Skip to content

Core API

Top-level environment, configuration, and rollout primitives.

OrbitalGymEnv

orbitalgym.env.core.OrbitalGymEnv 

OrbitalGymEnv(config: ScenarioConfig)

Composed environment. reset / step are pure functions under the hood.

reset 

reset(key: Array) -> tuple[EnvState, BySide]

Reset returns (env_state, BySide(guard=SideOutput, bandit=SideOutput)).

reset_from_state 

reset_from_state(
    state: EnvState, key: Array
) -> tuple[EnvState, BySide]

Start an episode from a stored state instead of sampling one.

Returns the state with the clock zeroed, derived frame views re-materialized from the truth frame, and the initial observations computed with key. The reference orbit carried by state is kept so a bank sampled under one epoch replays under the same one.

step 

step(
    key: Array, state: EnvState, actions: Actions
) -> StepOutput

Symmetric step. Both sides run through identical machinery.

EnvState

orbitalgym.env.core.EnvState

SingleAgentView

orbitalgym.env.single_agent.SingleAgentView 

SingleAgentView(env: OrbitalGymEnv)

Projects the symmetric OrbitalGymEnv core to a single-agent view.

Caller provides actions for the controlled side only. The opponent is driven by the policy declared in cfg.{side}_policy, which ScenarioConfig.__post_init__ defaults to ZeroControl() when the builder leaves it unset.

reset 

reset(key: Array)

Returns (env_state, obs_controlled, opponent_policy_state).

step 

step(
    key: Array,
    state: EnvState,
    controlled_action: Array,
    opp_policy_state: Any,
)

Returns (next_state, obs_controlled, reward_controlled, episode_done, opp_ps', info).

ScenarioConfig

orbitalgym.config.ScenarioConfig dataclass 

ScenarioConfig(
    n_guards: int,
    n_bandits: int,
    epoch_mjd_utc: float,
    reference_orbit: ReferenceOrbitState,
    guard_components: tuple[StateComponentKey, ...],
    bandit_components: tuple[StateComponentKey, ...],
    guard_params: VehicleParamsSpec,
    bandit_params: VehicleParamsSpec,
    ic_sampler: ICSpec,
    dt: float,
    max_horizon_s: float,
    seed: int,
    truth_dynamics: Any = DynamicsKey.HCW_RTN,
    policy_dynamics: Any = DynamicsKey.HCW_RTN,
    belief_dynamics: Any = None,
    belief_dynamics_resolved: Any = DynamicsKey.HCW_RTN,
    reference_orbit_dynamics: Any = None,
    reference_orbit_dynamics_resolved: Any = DynamicsKey.KEPLERIAN_ECI,
    guard_action_components: tuple[
        ActionComponentKey, ...
    ] = (ActionComponentKey.IMPULSIVE_MANEUVER,),
    bandit_action_components: tuple[
        ActionComponentKey, ...
    ] = (ActionComponentKey.IMPULSIVE_MANEUVER,),
    action_frame: Frame | None = None,
    guard_components_extended: tuple[
        StateComponentKey, ...
    ] = (),
    bandit_components_extended: tuple[
        StateComponentKey, ...
    ] = (),
    guard_observation_fn: Any = None,
    bandit_observation_fn: Any = None,
    guard_ground_station_network: Any = None,
    bandit_ground_station_network: Any = None,
    reward_fn: Any = None,
    termination_fn: Any = None,
    guard_belief_initializer: Any = None,
    guard_belief_updater: Any = None,
    bandit_belief_initializer: Any = None,
    bandit_belief_updater: Any = None,
    controlled_side: Any = "guard",
    guard_policy: Any = None,
    bandit_policy: Any = None,
    game: Any = None,
    attitude_dynamics_key: AttitudeDynamicsKey
    | None = None,
    guard_attitude_params: Any = None,
    bandit_attitude_params: Any = None,
    attitude_control_torque_max: tuple[float, float, float]
    | None = None,
    pointing_boresight_body: tuple[float, float, float] = (
        1.0,
        0.0,
        0.0,
    ),
    layout: Any = None,
)

max_steps property 

max_steps: int

Derived step count: ceil(max_horizon_s / dt).

VehicleParamsSpec

orbitalgym.config.VehicleParamsSpec dataclass 

VehicleParamsSpec(
    dry_mass_kg: float,
    isp_s: float,
    max_thrust_n: float,
    slew_rate_rad_s: float = 0.0,
)

ReferenceOrbitState

orbitalgym.reference_orbit.ReferenceOrbitState

Reference orbit 6D state in Earth-Centered Inertial frame at the scenario epoch.

The dynamical anchor that HCW linearizes about. Canonical representation; osculating elements derived on demand via astrojax helpers.

from_keplerian classmethod 

from_keplerian(
    semi_major_axis_m: float,
    eccentricity: float,
    inclination: float,
    raan: float,
    argument_of_perigee: float,
    mean_anomaly: float,
    as_degrees: bool = True,
) -> ReferenceOrbitState

Build a ReferenceOrbitState from classical Keplerian elements.

Element convention matches astrojax / sampling.side.RelativeKeplerian: [a, e, i, RAAN, omega, M] with M the mean anomaly. Semi-major axis is in meters; angle units are controlled by as_degrees. Conversion delegates to astrojax.state_koe_to_eci (Montenbruck & Gill Satellite Orbits, Eq. 2.43–2.44).

Parameters:

Name Type Description Default
semi_major_axis_m float

Semi-major axis a in meters. Must be > 0.

 required
eccentricity float

Eccentricity e. Must satisfy 0 <= e < 1 (elliptical / circular only; parabolic and hyperbolic rejected).

 required
inclination float

Inclination i.

 required
raan float

Right ascension of ascending node Ω.

 required
argument_of_perigee float

Argument of perigee ω.

 required
mean_anomaly float

Mean anomaly M.

 required
as_degrees bool

If True (default), interpret all angle inputs as degrees. If False, interpret them as radians.

 True

Returns:

Type Description
ReferenceOrbitState

ReferenceOrbitState with position_eci (3,) meters and
ReferenceOrbitState

velocity_eci (3,) m/s.

Raises:

Type Description
ValueError

If semi_major_axis_m <= 0 or eccentricity is outside [0, 1).

Examples:

from orbitalgym.reference_orbit import ReferenceOrbitState

# 500 km circular equatorial orbit, periapsis on the +x axis.
ref = ReferenceOrbitState.from_keplerian(
    semi_major_axis_m=6878e3,
    eccentricity=0.0,
    inclination=0.0,
    raan=0.0,
    argument_of_perigee=0.0,
    mean_anomaly=0.0,
)

The canonical representation is Cartesian ECI (position_eci, velocity_eci). You can construct it either directly with those arrays, or from classical Keplerian elements via ReferenceOrbitState.from_keplerian(...) — angles default to degrees; pass as_degrees=False for radians. The Keplerian constructor delegates to astrojax.state_koe_to_eci and validates that semi_major_axis_m > 0 and 0 ≤ eccentricity < 1.

from orbitalgym.reference_orbit import ReferenceOrbitState

# Cartesian (canonical):
ref_cart = ReferenceOrbitState(
    position_eci=jnp.array([7000e3, 0.0, 0.0]),
    velocity_eci=jnp.array([0.0, 7546.05, 0.0]),
)

# Keplerian (degrees by default):
ref_kep = ReferenceOrbitState.from_keplerian(
    semi_major_axis_m=7000e3,
    eccentricity=0.001,
    inclination=51.6,
    raan=0.0,
    argument_of_perigee=0.0,
    mean_anomaly=0.0,
)

Rollout

Two drivers: rollout is obs-only (agent_view = obs); belief_rollout threads per-side beliefs (agent_view = belief). Both compile to a single lax.scan-based JIT region.

orbitalgym.rollout.rollout 

rollout(
    env,
    policies: BySide,
    init_policy_state_fns: BySide,
    key: Array,
    n_steps: int,
    *,
    initial_state: Any = None,
) -> Trajectory

Symmetric obs-only rollout. Both sides driven by their own policies.

agent_view passed to each policy is the flattened per-side observation.

Returns a Trajectory with sides.guard / sides.bandit each carrying (T, N_side, ...) leading-axis arrays, plus episode_done (T,) latched.

initial_state starts the episode from a stored EnvState (see orbitalgym.eval.bank) instead of sampling.

orbitalgym.rollout.rollout_single_agent 

rollout_single_agent(
    view,
    controlled_policy: Callable,
    init_controlled_ps_fn: Callable,
    key: Array,
    n_steps: int,
) -> Trajectory

Convenience wrapper for SingleAgentView-style runs.

Pulls the opponent's scripted policy from view.opponent_policy and threads only the controlled side's policy state externally.

orbitalgym.rollout.belief_rollout 

belief_rollout(
    env,
    policies: BySide,
    init_policy_state_fns: BySide,
    belief_initializers: BySide,
    belief_updaters: BySide,
    key: Array,
    n_steps: int,
    *,
    guard_ground_station_network: Any = None,
    bandit_ground_station_network: Any = None,
    team_sync_fns: Any = None,
    guard_link: Any = None,
    bandit_link: Any = None,
    initial_state: Any = None,
) -> tuple[Trajectory, BySide]

Belief-aware rollout: agent_view = post-update Belief.

Each tick:

  1. For each side, compute a per-tick contact mask of shape (n_side,) from the side's link predicate when given, else from its ground-station network. When neither is given the mask is all-False.
  2. Optionally fuse the side's belief across in-contact teammates via team_sync_fns.X(belief_X, contact_X).
  3. Wrap the (possibly synced) belief as a :class:~orbitalgym.belief.contact_aware.ContactAwareBelief and pass it to policies.X(ps_X, agent_view, key, t).
  4. env.step advances the world.
  5. Each side's belief_updater folds the new per-side observation channels into the synced belief (so team fusion persists into the next tick's prior). The control it predicts under is the Δv the env actually imparted, from StepOutput.info["applied_dv"], not the commanded Δv: a command clipped by thrust or an empty tank would otherwise leave the own-state belief drifting off truth.
  6. Freeze-on-done as in :func:rollout.

The synced belief — not the wrapped ContactAwareBelief — is what flows forward in the carry: wrapping is a per-tick view, fusion is a persistent state update.

Returns (Trajectory, BySide[belief_history]). belief_history is a BySide whose leaves are belief leaves with a leading time axis — what feeds :class:orbitalgym.viz.animation.RolloutScene's belief_history parameter for animated 2σ ellipsoids.

initial_state starts the episode from a stored EnvState (see orbitalgym.eval.bank) instead of sampling.

orbitalgym.rollout.episode_mask 

episode_mask(traj: Trajectory) -> jax.Array

Boolean mask True up to and including the terminating step.

Same semantics as Phase 0's episode_mask but reads from the new Trajectory shape's episode_done field.

Precision

orbitalgym.set_precision 

set_precision(dtype) -> None

Set both astrojax's internal dtype and JAX's jax_enable_x64 flag.

astrojax.config.set_dtype(jnp.float64) enables jax_enable_x64 as a side effect, but set_dtype(jnp.float32) does NOT toggle it back. This helper keeps the two settings consistent so callers can flip precision with a single call. Call set_precision(jnp.float64) for sub-mm orbit precision (required by KOE/ECI conversions inside KEPLERIAN_ECI and full-force astrojax dynamics); the package default is jnp.float32 for backend portability.

No package-level cache invalidation is needed: typed-instance dynamics (KeplerianEciDynamics, AstrojaxOrbitDynamics) build their Epoch / RHS state in __post_init__, so a fresh dtype is picked up the next time a new env is constructed. Existing instances retain their original dtype — flip the dtype before building the env you want to run on the new precision.

Logging

orbitalgym.logging.writer.save_run 

save_run(
    path: Path | str,
    config: ScenarioConfig,
    trajectory: Any,
) -> None

Save (config, trajectory) to a single HDF5 file.

Parameters:

Name Type Description Default
path Path | str

destination file path (created or overwritten)

 required
config ScenarioConfig

the ScenarioConfig used to produce the trajectory (round-tripped via JSON)

 required
trajectory Any

any pytree (typically a Trajectory) whose leaves are jax.Array

 required

The config is stored as a JSON string in a root-level attribute so an inspector (e.g., h5py.File(path).attrs['config']) can see the scenario without loading the arrays.

orbitalgym.logging.reader.load_run 

load_run(
    path: Path | str,
) -> tuple[ScenarioConfig, dict[str, Any]]

Load a saved run from HDF5.

Returns:

Type Description
ScenarioConfig

(config, trajectory_dict) where trajectory_dict maps dotted-path strings
dict[str, Any]

to jax.Arrays (e.g., "reward", "action", "env_state.guards.rtn").

Raises ValueError if the file is missing required structure (config attr, schema_version attr, /trajectory group) or if schema_version doesn't match what this reader understands.