Skip to content

Symmetric core & data shapes

The OrbitalGymEnv core is symmetric: the same step function and the same data shapes work whether you're learning the guard, learning the bandit, training both simultaneously, or planning over beliefs. This page is the spine of the documentation — every other in-depth page assumes you've internalised the types and indexing patterns described here.

Sides

from orbitalgym import Side
Side.GUARD       # "guard"
Side.BANDIT      # "bandit"
Side.GUARD.opposite()   # Side.BANDIT

Side is a Python-level StrEnum, never a traced JAX array. That distinction matters: side identity is a static structural property of the computation graph, not data flowing through it. You can if side is Side.GUARD: ... inside a jit-traced function and JAX will specialise the trace per-side.

BySide[T] — the universal one-per-side container

Anywhere a value is naturally per-side — actions, observations, rewards, beliefs, scripted policies — the type is BySide[T].

    import jax
    import jax.numpy as jnp

    from orbitalgym import BySide, Side

    bs = BySide(guard=42.0, bandit=-1.0)
    assert bs.get(Side.GUARD) == 42.0
    assert bs.map(lambda x: x * 2).bandit == -2.0

BySide is registered as a JAX pytree, so vmap, lax.scan, and tree.map traverse it transparently:

    # BySide is a JAX pytree — vmap and tree.map traverse it.
    bs_arr = BySide(guard=jnp.array([1.0, 2.0]), bandit=jnp.array([3.0, 4.0]))
    bs_squared = jax.tree.map(lambda x: x**2, bs_arr)
    assert bool(jnp.all(bs_squared.bandit == jnp.array([9.0, 16.0])))

The big four step types

env.step takes Actions(sides=BySide(guard, bandit)) and returns StepOutput(state, outputs, episode_done, info), where:

  • state — the next EnvState.
  • outputs.guard, outputs.bandit — SideOutput(obs, reward, done).
  • episode_done — scalar bool. Per-side done fields are this scalar broadcast for shape uniformity across sides with different vehicle counts.
  • info["applied_dv"] — BySide of (N_side, 3) truth-frame Δv the step actually imparted, zero-padded to width 3. Thrust and remaining propellant both clip a command, so anything that predicts the ego motion, accounts for Δv spent, or draws a burn must read this rather than the commanded dv. rollout stores it on traj.applied_dv, where the episode metrics and the thrust quivers in RolloutScene pick it up.

Axis ordering: BATCH → TIME → VEHICLE → FEATURE

Outer-to-inner. Read every shape spec on this principle:

Axis When present
B (batch) Only under jax.vmap (e.g. seed batch).
T (time) Only on Trajectory leaves (added by lax.scan in rollout).
N_side (vehicle) Always present per side, even if N=1.
feature Innermost. action_dim=3 for HCW-RTN, 6 for RTN state. 
    from orbitalgym import OrbitalGymEnv, SingleAgentView, make_pursuit_evasion
    from orbitalgym.policies import ZeroControl
    from orbitalgym.rollout import rollout_single_agent

    cfg = make_pursuit_evasion(n_guards=2, n_bandits=1, seed=0, max_horizon_s=1000.0)
    env = OrbitalGymEnv(cfg)
    view = SingleAgentView(env)
    guard_policy = ZeroControl(command_cls=env.guard_command_cls, n_vehicles=cfg.n_guards)
    traj = rollout_single_agent(
        view,
        guard_policy,
        lambda c, s, k: None,
        jax.random.PRNGKey(0),
        n_steps=cfg.max_steps,
    )

    # T=time, N=vehicle, feature
    assert traj.sides.guard.action.dv.shape == (cfg.max_steps, cfg.n_guards, 3)
    assert traj.sides.guard.reward.shape == (cfg.max_steps,)
    assert traj.env_state.guards.rtn.shape == (cfg.max_steps, cfg.n_guards, 6)

Indexing recipes

    # i-th guard's reward time series — well, reward is PER_SIDE for
    # the bundled scope, so it's just `traj.sides.guard.reward`. For
    # PER_VEHICLE rewards it would be `traj.sides.guard.reward[:, i]`.
    side_reward_t = traj.sides.guard.reward  # (T,)

    # i-th guard's RTN trajectory:
    guard0_rtn = traj.env_state.guards.rtn[:, 0]  # (T, 6)

    # Step t's full state across all guards:
    step5_state = traj.env_state.guards.rtn[5]  # (N_g, 6)

The recipe pattern: pick the axes you want as colons, slice or index the rest.

Episode termination and the mask

Trajectory.episode_done is (T,) bool, latched after termination — once True, all subsequent entries are True. To restrict analysis to valid steps:

    from orbitalgym.rollout import episode_mask

    valid = episode_mask(traj)  # (T,) bool, latched
    valid_rewards = side_reward_t[valid]  # only steps before termination

Under vmap

Wrap a rollout in jax.vmap and every leaf gains a leading batch dimension. The full B → T → N → feature order applies:

    # Under vmap over seeds, every shape gains a leading B dimension:
    def run_one(seed):
        return rollout_single_agent(
            view,
            guard_policy,
            lambda c, s, k: None,
            seed,
            n_steps=cfg.max_steps,
        )

    seeds = jax.vmap(jax.random.PRNGKey)(jnp.arange(4))
    batched = jax.jit(jax.vmap(run_one))(seeds)
    assert batched.sides.guard.action.dv.shape == (4, cfg.max_steps, cfg.n_guards, 3)
    # B=batch, T=time, N=vehicle, feature

This is the workhorse pattern for seed-parallel evaluation; see How-to → vmap rollouts for a more detailed recipe and T5 — GPU / MPS for the acceleration story.

Single-agent vs multi-agent framing

The same env supports both. In multi-agent framing, the caller passes both sides' actions and reads both sides' outputs (this is what the PettingZoo adapter and the symmetric rollout use).

In single-agent framing, SingleAgentView reads cfg.controlled_side (default Side.GUARD) and runs the opposite side's cfg.<side>_policy internally. The view's step takes only the controlled side's action and returns only its reward and observation. The Gymnasium adapter wraps this view.

Why per-side policies live on the config

ScenarioConfig.guard_policy and bandit_policy are populated by the make_ builders to ZeroControl by default. SingleAgentView reads the opposite side's policy at construction. Multi-agent adapters ignore both fields. This means a swap from "guard learning vs heuristic bandit" to "bandit learning vs heuristic guard" is a single dataclasses.replace(cfg, controlled_side=Side.BANDIT) call — see How-to → Switch the controlled side.

Where to next