Components¶

Bases: Protocol

Structural protocol for a role-agnostic policy.

Implementations are typically frozen dataclasses with a command_cls field (the per-side Command class) and an n_vehicles field (defaulting to None/0) that the env populates via dataclasses.replace at construction. The call returns a Command pytree — never a raw array.

Always emits the identity Command (zeros for every component).

Env-populated init knobs (filled by OrbitalGymEnv via dataclasses.replace): command_cls: the per-side Command class (built by build_command_class) n_vehicles: number of vehicles on this side

Replaces the previous (n_vehicles, action_dim) raw-array contract.

Uniform-random over a (A, dv_dim) discrete action grid, per vehicle.

Env-populated init knobs (filled by OrbitalGymEnv via dataclasses.replace): command_cls: per-side Command class n_vehicles: number of vehicles on this side

The action_grid row count is the discrete action-space size; rows are Δv vectors whose dimensionality is padded to the side's command dv_dim if needed (e.g. 2-D RT grid into a 3-D RTN command).

JAX-native Gumbel MuZero / PUCT search backed by mctx.

See module docstring for the joint vs independent coordination modes.

Sampling-based receding-horizon planner conforming to the Policy protocol.

See the module docstring for the joint and independent coordination modes.

Bandit (or guard) policy that closes on the opponent.

Assumes the obs is a FullObservation flat layout in 1v1 RTN: [own_rtn(6), opp_rtn(6)]. For other observation channels, wrap in a LeadInterceptPursuer-shaped adapter that reconstructs the same own/opp split.

Evades an approaching opponent by thrusting perpendicular to relative velocity.

Assumes obs layout [own_rtn(6), opp_rtn(6)]. Action lies in the orthogonal complement of the relative velocity vector.

Sun-line-aware heuristic blocker.

Assumes obs layout [own_rtn(6), opp_rtn(6)]. Thrusts along the Sun-direction line, oriented by the sign of the opponent's projection onto that line (so the thrust closes on the Sun-side of the opponent).

Wraps a base policy and adds PRNG-keyed Gaussian jitter to its action.

The jitter is drawn from N(0, sigma²I) with the same shape as the base action. Setting sigma=0.0 is a pure passthrough.

base is required — instantiate with JitteredPolicy(base=..., sigma=...).

Confidence-gated routing wrapper.

policy_state semantics: a scalar confidence in [0, 1]. The wrapped primary's own state is not threaded through this class — primaries that need state should be wrapped with their own state-handling.

Only the dv field is gated by confidence. Non-dv Command fields (e.g. Communicate.active, Communicate.payload) propagate from the primary policy unchanged — the fallback's non-dv fields are discarded. If the fallback should drive non-dv fields too, replace this wrapper with a custom routing policy.

Sums actions from base and offset. Both must produce same-shape actions.

Only the dv field is summed. Non-dv Command fields (e.g. Communicate.active, Communicate.payload) propagate from the base policy unchanged — the offset policy's non-dv fields are discarded. This matches the typical "scripted base + learned residual" usage where the residual is a Δv correction only.

Bases: Protocol

Structural protocol for an observation function.

Returns a tuple of Observation channels. Most implementations return one channel; composite functions return more.

actions is the full Actions pytree (both sides). When the observation fn is called pre-step (env.reset, single_agent's pre-decision opp obs), actions is the identity Command for both sides — no information leakage from a hypothetical step. When called post-step (env.step), actions is the action that produced next_state.

One channel covering every tracked entity, no noise, always visible.

Layout requirements: must expose n_guards, n_bandits, dynamics_state_dim attributes.

One channel: per-pair 3-D position measurement gated by distance.

H selects the first 3 components (position) from the d-dim dynamics state. Layout requirements: n_guards, n_bandits, dynamics_state_dim.

One channel: per-pair full-state measurement gated by body-frame cone.

Parameters¶

layout: Scenario layout object exposing n_guards, n_bandits, and dynamics_state_dim. sensor_boresights_body: Shape (k, 3) array of unit vectors in the observer's body frame defining sensor pointing directions. half_angle_rad: Half-angle (radians) of each cone — either a scalar (same for all sensors) or a (k,) array / tuple of per-sensor half-angles. sigma_floor: Measurement-noise standard deviation at zero range. With sigma_range_frac=0.0 this is the noise std for every pair. Set both to 0.0 for noiseless (ground-truth) measurements. sigma_range_frac: Growth of the noise std per metre of observer-target range: sigma_ij = sigma_floor + sigma_range_frac * range_ij. The same sigma_ij scales every measurement row, so velocity rows are noised in proportion to range exactly as position rows are.

Bases: Protocol

Structural protocol for a per-step reward function.

Called as reward_fn(prev_state, action, next_state, side, params, t). action is a full Actions container (both sides) so multi-side rewards (e.g. zero-sum penalties using both sides' actions) can be expressed naturally. Returns shape per scope: PER_VEHICLE: (N_side,) PER_SIDE: () (scalar)

Bases: StrEnum

Negative sum of guard-to-reference-orbit-origin distances.

Returns 0.0 for the bandit side (preserving the bootstrap's single-reward-on-guard-only behavior). Future zero-sum games negate explicitly inside their own reward classes.

No-op reward — returns 0.0 for any side.

The default reward for NoGame. Use as a transparent placeholder when the user is expected to supply a custom reward via the reward_fn= constructor kwarg on ScenarioConfig.

Potential-shaped reward with zero-sum terminal events for Lady-Bandit-Guard.

The terminal events are mirrored; the shaping is not, because each side gets a potential over the distance it is trying to close. The shaping gives each side a dense gradient toward the right behavior without moving the game's equilibria, while the terminal events dominate the cumulative return when triggered. The potential is zero at an absorbing state, so the step that ends the episode pays no shaping beyond returning the potential the side had.

shaping_scale_m is the distance that makes a potential unity, shaping_gain scales the shaping term (zero for a terminal-only game), shaping_discount must match the planner's per-step discount for a side's shaped value to be exactly its unshaped value less its own potential, and home_weight prices a guard's distance from the lady inside the guard's potential.

catch_dwell_steps and breach_dwell_steps hold back the matching terminal bonus until a bandit has spent that many consecutive steps inside the radius, matching the termination condition so the bonus is paid on the step the episode ends. Zero pays it on the first step inside the radius.

guard_separation_m, lady_keepout_m and separation_cost set the guard-side crowding charge; dv_cost is the price of fuel in reward units per m/s of delta-v. Both are real per-side costs rather than shaping: each side pays only its own, so raising either makes that side more careful without handing the other an advantage.

potential(state, side: Side) -> jax.Array

Bases: Protocol

Episode-level termination — returns scalar bool.

Per-side done in SideOutput is this scalar broadcast (kept for shape uniformity). Per-side termination (one side gives up, the other does not) is a Phase-N concern; not needed by the four bootstrap games.

Episode ends when state.step >= params.max_steps. No spatial event.

The default termination for NoGame. Reads max_steps from the cfg at call time, so it can be constructed without a cfg in scope — which is what lets each Game subclass build its default termination cleanly inside default_termination_fn.

Terminate on max_steps OR breach OR catch OR every bandit repelled.

The catch event is optional: set catch_radius_m=0 to disable, leaving only the max-steps and breach gates. catch_dwell_steps and breach_dwell_steps require the corresponding radius to be held for that many consecutive steps; zero, the default, fires on the first step inside it. Both repel gates are off by default: escape_radius_m=0 disables the distance gate and repel_on_empty_tank=False the propellant gate. max_steps, the step duration and the reference orbit are read from params at call time (single source of truth on the cfg).

Bases: Protocol

Structural protocol for any belief representation.

Implementations must expose a mean: jax.Array attribute holding the per-pair state estimate. Shape conventions follow the concrete belief classes ((N_obs, N_total, d) for KF/EKF).

Bases: Protocol

Initialize a belief state from the env's ground-truth state.

Bases: Protocol

Update a belief given a new observation.

Linear Kalman filter, per-(observer, target) block-diagonal.

stm, control_matrix, process_noise are uniform across all tracked entities. Per-pair correction matrices H and R come from each Observation channel (different sensors may have different m, H, R).

Initial belief: per-observer mean = truth of every tracked entity.

Layout requirements: must expose n_guards, n_bandits, dynamics_state_dim attributes.

Uninformed initial belief: every (observer, tracked) pair gets the same prior.

Extended Kalman filter with autodiff Jacobians.

dynamics_fn(x, u, dt) -> x_next is the per-target nonlinear propagator. Per-pair Jacobian F_ik = jax.jacfwd(dynamics_fn)(mean[i,k]) at the current mean. For Observation channels with obs_fn set, H_ik = jax.jacfwd(obs_fn)(predicted_mean[i,k]).

Particle-cloud belief, per (observer, target) pair.

particles and log_weights carry the same leading (N_obs, N_total, K) indices; n_eff, weight_entropy, and resampled carry (N_obs, N_total). log_weights is stored already log-softmax-normalized along the particle (last) axis so taking softmax produces a valid probability vector at any moment.

The three pair-level metrics are pre-resample — they describe the posterior produced by the most recent measurement update. After a resample fires for a pair, weights are reset to uniform; storing the pre-resample values is what makes the collapse curve visible across a belief history (softmax(log_weights) after the resample is uniform by construction).

mean: Array

Sequential-importance-resampling particle filter.

Parameters¶

dynamics_fn: Per-vehicle propagator with signature (x, u, dt) -> x_next. Same shape contract as EKFBeliefUpdater.dynamics_fn. process_noise: (d, d) covariance Q used to draw additive Gaussian process noise during predict (and resample jitter). Pass a tiny positive-definite matrix when you want near-deterministic propagation; pure zeros are allowed but make resample-jitter a no-op. dt: Step size in seconds passed to dynamics_fn. n_eff_threshold: Resample fires per-pair when N_eff / K falls below this fraction. 0.5 is the standard default; 1.0 resamples every step, 0.0 never resamples. resample_jitter_scale: After resampling, add scale * chol(Q) * randn to particles so duplicate samples differ slightly. Default 0.0 (no jitter). Useful when the measurement is sharp and resamples produce many copies of the same particle.

All particles start at truth, plus tiny isotropic jitter.

Use when the belief should begin as a near-delta on the true state — e.g. the guard knows its own state exactly and starts with a perfect cheat-prior on the bandit. Jitter prevents singular weight updates downstream.

Uniform-on-bandit-ring prior for the LBG ring-intercept scenario.

For each cross-pair (observer i, target k where k targets the opposing side), sample K particles uniformly along the 2:1 RT-plane natural-motion ring at radius ring_radius_m. Velocity components are set from the closed-form natural-motion solution so every particle is a valid free-drift trajectory under HCW dynamics — the cloud stays on the ring as it propagates.

Self-pairs and same-side pairs are initialized to truth + tiny jitter so the observer's own-state estimate matches reality. The only ring-distributed pairs are (own-side observer, opposing-side target).

Layout requirements: n_guards, n_bandits, dynamics_state_dim (must be 4 for RT plane or 6 for RTN). RTN places zero cross-track position and rate, creating an in-plane ring.

Parameters¶

layout: cfg.layout from a built ScenarioConfig. ring_radius_m (scalar, or one radius per opposing vehicle): Radial-ellipse semi-major axis (matches the bandit's RelativeEllipse.radial_ellipse_m). Along-track amplitude is 2× this under HCW 2:1 motion. along_track_offset_m (scalar, or one offset per opposing vehicle): Along-track centre of the ring, matching the bandit's RelativeEllipse.along_track_offset_m. Added to the along-track position of every cross-pair particle; velocities are untouched, so the cloud is still a drift-free ring, centred this far ahead of the lady. mean_motion_rad_s: Reference orbit mean motion (use orbitalgym.reference_orbit.mean_motion(cfg.reference_orbit)). n_particles: Number of particles per (observer, target) pair. truth_jitter_scale: Jitter on self-pair / same-side particles (where particles are anchored to truth). Cross-pair particles do not get this jitter — their spread comes from the ring distribution itself.

Stale-track prior: a Gaussian cloud around the opponent's truth.

Models an observer that has been tracking the opposing side and holds a track that is right but not sharp. For each cross-pair (observer i, opposing-side target k), particles are the target's true state plus independent Gaussian noise — sigma_pos_m on every position component and sigma_vel_mps on every velocity component. Self-pairs and same-side pairs are truth plus tiny jitter, as in the other particle initializers.

Layout requirements: dynamics_state_dim must be 4 (RT plane, components r, t, r_dot, t_dot) or 6 (RTN, with the cross-track position and rate drawn at the same sigmas).

Parameters¶

layout: cfg.layout from a built ScenarioConfig. n_particles: Number of particles per (observer, target) pair. sigma_pos_m: Standard deviation on each position component of a cross-pair particle. sigma_vel_mps: Standard deviation on each velocity component of a cross-pair particle. Pairing it with mean_motion * sigma_pos_m keeps the position and velocity uncertainty consistent under HCW relative motion. truth_jitter_scale: Jitter on the truth-anchored self-pair and same-side particles.

Bridge between the per-side pytree state and a single flat vector.

Bundles four things that derive together from n_guards, n_bandits, and the per-side state-component tuples (cfg.guard_components, cfg.bandit_components):

• guard_state_cls, bandit_state_cls — the dynamically-assembled per-side dataclass types. These contain one field per state component (e.g. rtn for an RTN dynamics frame, mass, power, ...). The set of fields is determined by the cfg.{side}_components tuple.
• n_guards, n_bandits — vehicle counts.
• flat_dim — the size of the full flat state vector. For a 1v1 RTN scenario with masses, this is (6 + 1) + (6 + 1) = 14 (six dims of state per vehicle from RTN, one extra from MASS, times two sides). Custom component sets change this.
• _flatten_fn, _unflatten_fn — JIT-compatible bijections between (guards_pytree, bandits_pytree) and a flat (flat_dim,) vector. POMDP planners use these to operate on the flat representation.

Built once per cfg in ScenarioConfig.__post_init__. Read it as cfg.layout.

dynamics_state_dim: int

bandit(bandit_state: Any, j: int) -> Any

build(
    guard_state_cls: type,
    bandit_state_cls: type,
    n_guards: int,
    n_bandits: int,
) -> StateLayout

guard(guard_state: Any, i: int) -> Any

2D in-plane relative state in the RT (radial + along-track) sub-frame.

rt layout: [r, t, r_dot, t_dot] — 4-dim per vehicle.

Full 3D relative state in the RTN frame.

rtn layout: [r, t, n, r_dot, t_dot, n_dot] — 6-dim per vehicle.

Absolute 6D inertial state (Earth-Centered Inertial).

eci layout: [rx, ry, rz, vx, vy, vz] — position [m] then velocity [m/s].

Propellant-mass tracking.

Only propellant_mass is state; wet_mass is derived on demand as params.dry_mass_kg + state.propellant_mass.

Scalar state-of-charge. No-op dynamics in bootstrap (stub).

Quaternion attitude (w, x, y, z). Protocol-level stub in bootstrap.

Body-frame angular rates (ωx, ωy, ωz). Protocol-level stub in bootstrap.

Per-step applied translational Δv (m/s) in the truth frame.

Transient: written by ImpulsiveManeuver.apply (or any other action component that produces Δv), consumed by the translational dynamics block in env.step, then zeroed.

Width is always 3 on state for shape uniformity. The dynamics step slices [:, :2] for Frame.RT, [:, :3] for RTN/ECI.

Per-step applied body-frame torque (N·m).

Transient: written by AttitudeControl.apply (or any other torque- producing action component), consumed by rigid_body_attitude_step in env.step, then zeroed.

Bases: Protocol

Composable per-side action-pytree component.

Mirrors StateComponent for the schema half (fields, zeros). Adds apply(command, side_state, side_params, dt, ref_eci6, key) -> side_state which is invoked by OrbitalGymEnv.step once per side, in component registration order. Components are responsible for any state mutation (including invoking dynamics — see ImpulsiveManeuver); the env loop only folds.

fields and zeros are instance methods so frame-aware components (e.g. ImpulsiveManeuver, where self.action_frame.dim decides the dv shape) can vary their schema. Components without per-instance state can simply ignore self.

apply(
    command: Any,
    side_state: Any,
    side_params: Any,
    dt: float,
    ref_eci6: Array,
    key: Array,
) -> Any

fields() -> Mapping[str, tuple[int, ...]]

zeros(n: int) -> Mapping[str, jax.Array]

Impulsive Δv producer.

Reads command.dv (in action_frame), converts to truth_frame, optionally deducts propellant via the rocket equation, and writes the result onto side_state.applied_dv (always width 3, zero-padded for RT).

Does NOT integrate the dynamics — env.step owns the translational dynamics call (see the explicit propagation block in OrbitalGymEnv.step).

Per-step delta-v magnitude is limited to max_thrust_n * dt / m_wet and to the delta-v obtainable from remaining propellant. Direction is preserved when clipping. An empty tank produces no delta-v.

The Δv field shape follows action_frame.dim — 2 for RT (radial, along-track) and 3 for RTN/ECI. fields() and zeros() are instance methods (not staticmethods) so the assembled Command pytree carries the correct shape per scenario; build_command_class is given component instances so it can read this.

Init knobs (env wires these from ScenarioConfig): action_frame: frame the policy emits Δv in truth_frame: frame the truth dynamics integrates in track_mass: whether to deduct propellant via the rocket equation

Decentralized-communication action component.

Apply is state-identity: communication's effects are visible to other agents via the observation pipeline (which receives the full Actions pytree under the Phase 3 protocol change). The reward function reads active to charge a per-broadcast cost.

Payload shape is fixed at (6,) (RTN position+velocity) — JAX requires static field shapes for jit-stability. Scenarios with different payload dim must register a subclass with overridden fields()/zeros().

Torque-commanding attitude control.

Reads command.torque (length rotation_dim — 1 in RT, 3 in RTN), lifts to a 3-vector in 2D (zero x,y), optionally clips per-axis to ±torque_max, writes the result onto side_state.applied_torque.

Does NOT integrate dynamics — env.step invokes rigid_body_attitude_step after the action fold. Mirror of the unified action↔dynamics pattern set by ImpulsiveManeuver: action produces, env consumes.

Per-side rigid-body parameters.

inertia_diag: principal-axis inertias I_xx, I_yy, I_zz (kg m^2). omega_max: per-axis angular-rate saturation (rad/s).

Composable initial-condition specification.

Used as ScenarioConfig.ic_sampler. The env reads this once at reset, runs validators in a jax.lax.while_loop capped at max_attempts, and sets EnvState.ic_valid=False if the cap is hit.

Sample Keplerian element differences around the reference orbit's osculating elements.

Each delta (delta_a, delta_e, delta_i, delta_RAAN, delta_omega, delta_M) is Gaussian: mean_* (default 0) and sigma_* (default 0). Scalars broadcast to all N vehicles; (N,) arrays specify per-vehicle parameters. M = mean anomaly (astrojax convention).

Singular at reference orbit e=0, i=0 (RAAN/omega/M ill-defined). Use a slightly inclined or eccentric reference orbit for clean sampling.

Precision note: this sampler relies on astrojax for KOE<->ECI conversions. The package configures astrojax for float64 at import time (orbitalgym/__init__.py), giving sub-mm round-trip residuals on Earth-orbit scenarios.

Sample bounded relative orbits via NSROE with ellipse extents in meters.

The non-singular relative orbital elements (Schaub & Junkins; Lyu & D'Amico) map to physical extents as: a * |delta_e| = radial-ellipse semi-major axis (m) a * |delta_i| = cross-track amplitude (m) a * delta_lambda = mean along-track separation (m) a * delta_a = secular drift coefficient (1.5ndelta_a per radian)

radial_ellipse_m / cross_track_m: scalar broadcasts; (N,) per-vehicle. along_track_offset_m: same. along_track_drift_m_per_orbit: optional; default 0 -> bounded. sigma_*: optional Gaussian noise added to the corresponding extent. phase_rad: per-vehicle initial phase along the relative ellipse. None -> uniform random per vehicle. scalar / (N,) -> deterministic per vehicle.

Same propellant mass for all N vehicles.

Independent uniform [low, high] propellant mass per vehicle.

Reject if the configured pair scope has any pair closer than distance_m.

Reject if any selected vehicle is farther than distance_m from the reference orbit origin.

Since RTN positions are reference-orbit-relative, the reference orbit origin is (0,0,0) in this frame. side: ALL - both guards and bandits GUARDS - guards only BANDITS - bandits only

Identity quaternion (1, 0, 0, 0); zero body rates. Default behavior.

Pinned quaternion + body rates. Both broadcast to all N vehicles.

Useful for demonstrations and deterministic test scenarios where the user wants every spacecraft to start with the same known orientation and slow spin.

Uniformly random quaternion (Marsaglia method); zero body rates.

Samples q from N(0, I_4)^N and normalizes. Canonicalizes to the hemisphere with w >= 0 (since q and -q represent the same rotation).

Identity quaternion; uniform random body rates per axis.

omega ~ Uniform(-omega_max, omega_max) per axis, per vehicle.

Both quaternion and body rates randomized.