Skip to content

T2 — Build an RTN scenario with visualisations

Build a 2-guard / 1-bandit Pursuit-Evasion scenario in 3D RTN, run a multi-agent rollout, and render trajectory and reward plots. Builds on T1 — First rollout (RT 2D); a companion notebook lives at examples/workflow_3d_rtn.ipynb.

Imports

    import jax
    import jax.numpy as jnp

    from orbitalgym import OrbitalGymEnv, make_pursuit_evasion
    from orbitalgym.env.types import BySide
    from orbitalgym.policies import ZeroControl
    from orbitalgym.rollout import episode_mask, rollout

Build the scenario

    cfg = make_pursuit_evasion(
        n_guards=2,
        n_bandits=1,
        dt=10.0,
        max_horizon_s=2000.0,
        seed=0,
    )

make_pursuit_evasion defaults to HCW-RTN, so we get the 6-state form [r, θ, n, ṙ, θ̇, ṅ] — radial / along-track / cross-track position and velocity in the rotating reference-orbit frame. Two guards, one bandit, 200 steps at 10 s.

Roll the episode out (multi-agent)

    env = OrbitalGymEnv(cfg)
    policies = BySide(
        guard=ZeroControl(command_cls=env.guard_command_cls, n_vehicles=cfg.n_guards),
        bandit=ZeroControl(command_cls=env.bandit_command_cls, n_vehicles=cfg.n_bandits),
    )
    init_ps = BySide(
        guard=lambda c, s, k: None,
        bandit=lambda c, s, k: None,
    )
    traj = rollout(
        env,
        policies,
        init_ps,
        jax.random.PRNGKey(cfg.seed),
        n_steps=cfg.max_steps,
    )

This is rollout, not rollout_single_agent — both sides have an explicit policy and both sides' actions are recorded. The symmetric core takes a BySide[Policy] for the policies and a BySide[Callable] for the per-side policy-state initialisers; the lambda c, s, k: None callables are no-op initialisers because ZeroControl is stateless.

Read off the shapes

    guard_rtn = traj.env_state.guards.rtn  # (T, N_g, 6)
    bandit_rtn = traj.env_state.bandits.rtn  # (T, N_b, 6)
    valid = episode_mask(traj)  # (T,) bool
  • traj.env_state.guards.rtn — shape (T, N_g, 6). Outer-to-inner: time, vehicle, feature.
  • traj.env_state.bandits.rtn — shape (T, N_b, 6). Same convention.
  • episode_mask(traj) — (T,) bool, True for steps before and including termination.

To grab the i-th guard's full time series of position: traj.env_state.guards.rtn[:, i, :3].

To grab the t-th step's full state across guards: traj.env_state.guards.rtn[t].

Compute capture distance over time

    # Pairwise distance: guard 0 vs. bandit 0 over time.
    pos_guard0 = guard_rtn[:, 0, :3]  # (T, 3)
    pos_bandit0 = bandit_rtn[:, 0, :3]  # (T, 3)
    dist_g0_b0 = jnp.linalg.norm(pos_guard0 - pos_bandit0, axis=-1)

Pairwise distance between guard 0 and bandit 0. With zero control on both sides, this oscillates (relative ellipses precess against each other) without falling below cfg.game.capture_distance_m.

Render the trajectory in 3D

    ax3d = plot_rtn_3d(traj.env_state.guards.rtn)
    ax3d.figure.savefig(tmp / "guards_rtn_3d.png", dpi=120)
    plt.close(ax3d.figure)

plot_rtn_3d accepts a (T, N, 6) array and overlays each vehicle's ellipse on the same axes. The reference orbit appears as the origin.

Render the reward curve

    fig, ax = plt.subplots()
    plot_reward_curve(traj.sides.guard.reward, ax=ax)
    fig.savefig(tmp / "guard_reward.png", dpi=120)
    plt.close(fig)

For Pursuit-Evasion (zero-sum) the bandit reward is the negation of the guard reward — plot one and the other is implied.

Next

→ T3 — Write an active-pursuer policy replaces ZeroControl on the bandit with a HCW-predict-and-thrust policy that actually closes on the guard.