Dynamics¶

OrbitalGymEnv composes per-side action components (commands → post-component side state) with the configured dynamics. The bundled ImpulsiveManeuver action component reads command.dv, converts to the truth frame, optionally deducts propellant, and writes the result to state.applied_dv (a transient state field). The env's step() then explicitly invokes the configured truth_dynamics with applied_dv as the control input, propagating the truth-frame state, and zeros applied_dv after consumption. Both action components and dynamics are swappable.

What ships¶

Both HCW dynamics ignore J2, drag, and finite-burn duration. They are exact for instantaneous impulses applied to a circular reference orbit.

Step functions vs typed-instance dynamics¶

Two dynamics shapes ship side by side:

• Step functions (HCW_RT, HCW_RTN, J2_ECI) are stateless callables registered against their key. The env resolves the key to the function and calls it directly each step. No per-instance state.
• Typed-instance classes (KeplerianEciDynamics, AstrojaxOrbitDynamics) have per-instance fields (e.g. epoch_mjd_utc, force_model) that drive how the dynamics is configured. Their Epoch and integrator RHS are built once in __post_init__ (Python time, before any JIT trace) and stored on the instance — this avoids module-level state and lets each scenario carry its own reference epoch. When the env is given a bare DynamicsKey.KEPLERIAN_ECI, it instantiates the class with defaults (J2000 epoch); to bind your scenario's epoch, pass the typed instance directly: truth_dynamics=KeplerianEciDynamics(epoch_mjd_utc=58849.0). ScenarioConfig.__post_init__ already does this for you for reference_orbit_dynamics, threading cfg.epoch_mjd_utc through.

The dynamics protocol¶

dynamics: (state, dv, params, dt) → next_state

• state — raw per-side dynamics array, shape (N_side, 4) for HCW_RT or (N_side, 6) for HCW_RTN. One row per vehicle.
• dv — applied impulse, shape (N_side, action_dim). Sourced from the per-side ImpulsiveManeuver action component's command.
• params — VehicleParams (mass, max-thrust, dry-mass, mean motion).
• dt — scalar timestep in seconds.

OrbitalGymEnv exposes the configured dynamics callables as env.truth_dynamics, env.policy_dynamics, and env.belief_dynamics. You can call them directly:

    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))

    # `env.truth_dynamics` is the bound dynamics callable. It takes the raw
    # per-side dynamics array (shape (N_side, 6) for HCW-RTN), a Δv impulse
    # of shape (N_side, action_dim), the per-side `VehicleParams`, and a
    # scalar `dt`. It returns the next-step raw array of the same shape.
    guard_rtn = state.guards.rtn  # (N_g, 6)
    zero_dv = jnp.zeros((cfg.n_guards, 3))
    next_guard_rtn = env.truth_dynamics(guard_rtn, zero_dv, env.guard_params, cfg.dt)
    assert next_guard_rtn.shape == guard_rtn.shape

    # The same callable handles the bandit side — symmetric API.
    bandit_rtn = state.bandits.rtn  # (N_b, 6)
    next_bandit_rtn = env.truth_dynamics(
        bandit_rtn,
        jnp.zeros((cfg.n_bandits, 3)),
        env.bandit_params,
        cfg.dt,
    )
    assert next_bandit_rtn.shape == bandit_rtn.shape

The dynamics callable operates on the raw rtn array (or rt for the 2D variant), not the assembled side-state pytree. The env's step method handles the wrapping/unwrapping; calling the dynamics directly gives you the same physics with no observation, reward, or termination machinery in the way.

Attitude dynamics layer¶

Translational dynamics (always runs) and attitude dynamics (runs iff cfg.attitude_dynamics_key is set) are orthogonal. Every step:

1. Action components write applied_dv and/or applied_torque onto the side state.
2. The translational dynamics block propagates the orbital state using applied_dv, then zeros it.
3. The attitude dynamics block integrates (quat, omega) using applied_torque, then zeros it.

The two layers share no state. Scenarios with no attitude_dynamics_key run only the translational block; the attitude fields (quat, omega) are untouched by env.step and stay at whatever the IC sampler produced.

See Attitude dynamics & conical sensors for the full story — rigid_body_attitude_step, AttitudeParams, reaction-wheel saturation, and the ConicalObservation model that makes attitude decision-relevant.

Composition with the actuator¶

OrbitalGymEnv.step runs the actuator first, then the dynamics:

    import jax

    from orbitalgym import Actions, BySide, 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))

    # `env.step` runs the actuator (action → applied Δv) and then the
    # dynamics step (state + Δv → next state) for each side.
    actions = Actions(
        sides=BySide(
            guard=env.guard_command_cls.zeros(cfg.n_guards),
            bandit=env.bandit_command_cls.zeros(cfg.n_bandits),
        )
    )
    step_output = env.step(jax.random.PRNGKey(1), state, actions)

    # The dynamics step preserves the per-side state shape.
    assert step_output.state.guards.rtn.shape == state.guards.rtn.shape
    assert step_output.state.bandits.rtn.shape == state.bandits.rtn.shape

The actuator owns the action shape semantics (e.g. "action is a desired Δv, clipped to max_dv_mps"). The dynamics owns the physics (e.g. "given a Δv, propagate the HCW state by dt"). Swap either independently.

Truth, policy, and belief dynamics¶

ScenarioConfig carries three dynamics roles:

• truth_dynamics — what the env actually integrates each step. The ground truth.
• policy_dynamics — what POMDPAdapter's transition uses for the model-based solver. May differ from truth (e.g. simpler model for a faster planner; or pretend-J2 for robustness studies).
• belief_dynamics — what the belief updater uses to propagate the belief between observations. Defaults to policy_dynamics when unset (read the resolved value as cfg.belief_dynamics_resolved / env.belief_dynamics).

Bundled games default all three to the same key.

Writing a custom dynamics module¶

See Extending → Customize dynamics for the full extension surface — registered vs typed-instance patterns, frame/kind metadata, and the headline mixed-frame recipe (ECI truth + HCW belief/policy via AstrojaxOrbitDynamics).

What's not yet swappable¶

The actuator surface assumes impulsive dynamics (one Δv per timestep). Finite-burn or thrust-vs-time actuators would need a richer protocol — out of scope for the bundled framework, addressable by extending the Actuator protocol.

Precision and dtype¶

The package configures astrojax for float64 at import time, which also sets JAX's jax_enable_x64=True. This is required for sub-mm KOE↔ECI round-trip residuals on Earth-orbit scenarios. Some hardware (notably Apple Metal / MPS) is float32-only; flip with the package's helper before constructing any env:

import jax.numpy as jnp
import orbitalgym

orbitalgym.set_precision(jnp.float32)   # for MPS / GPU throughput
# ... build ScenarioConfig and OrbitalGymEnv after this ...

set_precision flips both astrojax's internal dtype and jax_enable_x64 together — calling astrojax.config.set_dtype(jnp.float32) alone does not toggle jax_enable_x64 back to False and leaves the package in an inconsistent state. Existing typed-instance dynamics captured the old dtype in their __post_init__; build new envs after the precision flip to pick up the new dtype.

Where to next¶

• Attitude dynamics & conical sensors — the attitude dynamics layer, rigid_body_attitude_step, AttitudeParams, and the ConicalObservation model.
• State layout & adding a Power component — what the dynamics step reads and writes.
• API reference → Components — Dynamics and Actuator protocols.