Skip to content

State layout & adding a Power component

EnvState carries one side-state per side. Each side-state is a flax pytree composed of state components — RTN for relative-orbital state, Mass for propellant tracking, and so on. State components are swappable: pick the set you want when building the side-state class.

This page documents how to read the existing fields, then walks through adding a new Power component end-to-end.

Reading the bundled components

    import jax
    import jax.numpy as jnp

    from orbitalgym import OrbitalGymEnv, make_lady_bandit_guard

    cfg = make_lady_bandit_guard(n_guards=1, n_bandits=1, seed=0)
    env = OrbitalGymEnv(cfg)
    state, _ = env.reset(jax.random.PRNGKey(0))

    # state.guards.rtn has shape (N_g, 6) — one row per guard, six columns:
    # [r, theta, n, r_dot, theta_dot, n_dot]  (RTN position then velocity)
    rtn = state.guards.rtn
    assert rtn.shape == (cfg.n_guards, 6)

    # Position columns:
    pos = rtn[:, 0:3]  # (N_g, 3)
    vel = rtn[:, 3:6]  # (N_g, 3)

The columns of state.guards.rtn are [r, θ, n, ṙ, θ̇, ṅ] in the rotating reference-orbit frame. For HCW-RT (planar), the analogous attribute is state.guards.rt with 4 columns [r, θ, ṙ, θ̇].

If the side has the Mass component, propellant mass is exposed directly:

    # If the side carries the Mass component, propellant_mass has shape (N_side,).
    if hasattr(state.guards, "propellant_mass"):
        m_g = state.guards.propellant_mass
        assert m_g.shape == (cfg.n_guards,)

Flat layout

POMDPAdapter and any other consumer that wants a 1-D state vector uses env.layout.flatten and env.layout.unflatten:

    # POMDPAdapter flattens state into a 1-D vector for solver consumption.
    flat = env.layout.flatten(state.guards, state.bandits)
    rebuilt_g, rebuilt_b = env.layout.unflatten(flat)
    assert bool(jnp.allclose(rebuilt_g.rtn, state.guards.rtn))

The StateLayout-level flat vector excludes the scalar time / step counter, the game's dwell counters, and the reference orbit. POMDPAdapter packs (t, step) and the per-bandit dwell counters onto the tail of its own flat vector so transition can stay pure (see Use the POMDP adapter); the reference orbit is treated as a per-scenario constant.

Worked example: adding a Power component

Suppose your scenario needs a per-vehicle battery state-of-charge that decreases proportionally to commanded thrust magnitude. End-to-end:

1. The component dataclass

import flax.struct
import jax

@flax.struct.dataclass
class Power:
    soc: jax.Array      # (N_side,) state of charge in [0, 1]

2. Register it with the side-state assembler

Side-state classes are built dynamically via build_state_class (see orbitalgym.state.assemble). Register Power so the assembler knows how to allocate space for it in the flat layout.

3. Sampler hook

Initial SoC is sampled at episode reset. Add a sampler that draws soc ~ Uniform(0.7, 1.0) per vehicle and writes it into the Power component of the freshly-sampled side-state.

4. Dynamics-side update

Each step, drain SoC by α · |dv| for each vehicle, where dv is the realised impulse magnitude that step. The dynamics step gets (state, dv, params, dt); you read dv, decrement state.power.soc, and return the new state.

5. Plot hook (optional)

Add viz.summaries.plot_power_curve(state.guards.power.soc) analogous to the existing plot_mass_curve. Same (T, N) → matplotlib axis shape.

What this is not

The Power component is illustrative — it is not part of the merged library. The point is the pattern: a frozen dataclass + the assembler hook + a sampler + a dynamics-side update + an optional viz hook. Apply the same template to sensor health or any other per-vehicle quantity you want side-states to carry.

Two state components track rigid-body rotation:

  • Attitude — quaternion (n, 4) in wxyz convention. Initialized to the identity quaternion (1, 0, 0, 0). Represents body-to-reference rotation.
  • BodyRates — body-frame angular rate (n, 3) in rad/s. Zero-initialized.

Add them to guard_components or bandit_components together (they must appear as a pair; having one without the other raises in __post_init__):

from orbitalgym.registry import StateComponentKey

guard_components=(
    StateComponentKey.RTN,
    StateComponentKey.MASS,
    StateComponentKey.ATTITUDE,
    StateComponentKey.BODY_RATES,
)

Transient control components

Two additional components carry "control input for this tick" from action components to the dynamics block. Both are auto-extended onto the per-side components by ScenarioConfig.__post_init__ — users do not list them explicitly:

  • AppliedDV — translational Δv (n, 3) in m/s (truth frame). Added automatically whenever the side has any spatial component (RT, RTN, or ECI). Written by ImpulsiveManeuver.apply; read and zeroed by env.step's translational dynamics block.
  • AppliedTorque — body-frame torque (n, 3) in N·m. Added automatically whenever the side has both ATTITUDE and BODY_RATES. Written by AttitudeControl.apply; read and zeroed by env.step's attitude dynamics block.

Both are zeroed each step after dynamics consumes them, so they always read zero outside the dynamics block. They appear in the flat state layout (since the flat layout covers the full assembled pytree) but are not meaningful as persistent quantities — treat them as implementation detail of the action-to-dynamics handoff.

AppliedDV is always width 3 for shape uniformity. The dynamics block slices [:, :2] for Frame.RT and [:, :3] for RTN/ECI.

Where to next