Customize IC sampling¶

cfg.ic_sampler is an ICSpec composing two per-side samplers and a tuple of validators. The env runs validators in jax.lax.while_loop and sets EnvState.ic_valid=False if max_attempts is exhausted.

The protocol¶

orbitalgym.sampling.spec.ICSpec `dataclass` ¶

ICSpec(
    guard_sampler: Callable,
    bandit_sampler: Callable,
    validators: tuple[Callable, ...] = (),
    max_attempts: int = 100,
)

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.

orbitalgym.sampling.spec.SideSampler ¶

Bases: Protocol

Sample initial state for one side (guards or bandits).

Concrete implementations are frozen dataclasses with call. The env passes (n_vehicles, components, class_name) as keyword args so the same sampler instance can serve either side without baking side identity in.

orbitalgym.sampling.spec.Validator ¶

Bases: Protocol

Post-sample combined-state check. Returns scalar bool jax.Array.

The bundled samplers¶

orbitalgym.sampling.side.RelativeKeplerian `dataclass` ¶

RelativeKeplerian(
    mean_delta_sma_m: ndarray | float = 0.0,
    sigma_delta_sma_m: ndarray | float = 0.0,
    mean_delta_ecc: ndarray | float = 0.0,
    sigma_delta_ecc: ndarray | float = 0.0,
    mean_delta_inc_rad: ndarray | float = 0.0,
    sigma_delta_inc_rad: ndarray | float = 0.0,
    mean_delta_raan_rad: ndarray | float = 0.0,
    sigma_delta_raan_rad: ndarray | float = 0.0,
    mean_delta_argp_rad: ndarray | float = 0.0,
    sigma_delta_argp_rad: ndarray | float = 0.0,
    mean_delta_mean_anomaly_rad: ndarray | float = 0.0,
    sigma_delta_mean_anomaly_rad: ndarray | float = 0.0,
    mass_sampler: ConstantMass | UniformMass | None = None,
    attitude_sampler: IdentityAttitude
    | FixedAttitude
    | UniformAttitude
    | UniformBodyRates
    | UniformAttitudeAndRates
    | None = None,
)

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.

orbitalgym.sampling.side.RelativeEllipse `dataclass` ¶

RelativeEllipse(
    radial_ellipse_m: ndarray | float = 0.0,
    cross_track_m: ndarray | float = 0.0,
    along_track_offset_m: ndarray | float = 0.0,
    along_track_drift_m_per_orbit: ndarray | float = 0.0,
    sigma_radial_ellipse_m: ndarray | float = 0.0,
    sigma_cross_track_m: ndarray | float = 0.0,
    sigma_along_track_offset_m: ndarray | float = 0.0,
    phase_rad: ndarray | float | None = None,
    mass_sampler: ConstantMass | UniformMass | None = None,
    attitude_sampler: IdentityAttitude
    | FixedAttitude
    | UniformAttitude
    | UniformBodyRates
    | UniformAttitudeAndRates
    | None = None,
)

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.

The bundled validators¶

orbitalgym.sampling.validators.MinSeparation `dataclass` ¶

MinSeparation(
    distance_m: float,
    scope: SeparationScope = SeparationScope.ALL,
)

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

scope

ALL - every pair across both sides WITHIN_GUARDS - guard-guard pairs only WITHIN_BANDITS - bandit-bandit pairs only CROSS - guard-bandit pairs only

orbitalgym.sampling.validators.MaxRange `dataclass` ¶

MaxRange(
    distance_m: float, side: SideSelector = SideSelector.ALL
)

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

Worked example: tighter Gaussian for evaluation¶

For evaluation, you usually want a tighter spread around a fixed reference configuration so the metric isn't dominated by IC variance:

    import dataclasses

    import jax.numpy as jnp

    from orbitalgym import OrbitalGymEnv, SingleAgentView, make_pursuit_evasion
    from orbitalgym.sampling.side import RelativeEllipse
    from orbitalgym.sampling.spec import ICSpec

    tight_ic = ICSpec(
        guard_sampler=RelativeEllipse(
            radial_ellipse_m=1000.0,
            cross_track_m=0.0,
            along_track_offset_m=0.0,
            phase_rad=0.0,
            sigma_radial_ellipse_m=1.0,  # 10x tighter than default
        ),
        bandit_sampler=RelativeEllipse(
            radial_ellipse_m=500.0,
            cross_track_m=100.0,
            along_track_offset_m=200.0,
            phase_rad=jnp.pi / 2.0,
            sigma_radial_ellipse_m=1.0,
        ),
        validators=(),
        max_attempts=100,
    )

    cfg = make_pursuit_evasion(seed=0, max_horizon_s=200.0)
    cfg = dataclasses.replace(cfg, ic_sampler=tight_ic)
    env = OrbitalGymEnv(cfg)

Per-game applicability¶

Game	Default IC	What's reasonable to swap
Lady-Bandit-Guard	1km radial-ellipse co-orbit, guard@0, bandit@π	Tighter sigma for eval; Gaussian-around-RT for 2D bench
Pursuit-Evasion	Same shape, asymmetric phase	Wider initial separation for harder problems
Sun-Blocking	Sun-illuminated geometry	Twilight-band initial states for edge-case study
Observation-Blocking	Off-target initial geometry	Pre-aligned ICs for the simplest variant

Attitude samplers¶

Translational samplers (RelativeKeplerian, RelativeEllipse) accept an optional attitude_sampler field that initializes the per-vehicle quaternion and body rates at episode reset. When attitude_sampler=None (the default), initial attitude comes from Attitude.zeros() — identity quaternion (1, 0, 0, 0) and zero body rates.

Five samplers ship in orbitalgym.sampling.attitude:

IdentityAttitude() — identity quaternion (1, 0, 0, 0), zero rates. The default when no sampler is configured. Use this explicitly if you want to document intent.
FixedAttitude(quat_wxyz=(w,x,y,z), omega_rad_s=(wx,wy,wz)) — pinned quaternion and body rates broadcast to all N vehicles. Useful for deterministic demos and unit tests where every spacecraft starts with the same known orientation and spin rate.
UniformAttitude() — uniformly random quaternion (Marsaglia method: sample from N(0, I_4), normalize, canonicalize to w >= 0 hemisphere); zero body rates. Use when initial pointing should be isotropically random.
UniformBodyRates(omega_max_rad_s=(wx,wy,wz)) — identity quaternion; uniform random body rates in [-omega_max, +omega_max] per axis. Use when you want a known initial orientation with a random spin.
UniformAttitudeAndRates(omega_max_rad_s=(wx,wy,wz)) — both quaternion (Marsaglia) and body rates (uniform per axis) randomized. Use for the most adversarial initial conditions.

Plug them into RelativeKeplerian or RelativeEllipse via the attitude_sampler field:

from orbitalgym.sampling.attitude import FixedAttitude, UniformAttitude
from orbitalgym.sampling.side import RelativeKeplerian

# Guards start with a known slow spin (0.1 deg/s about body-z)
guard_sampler = RelativeKeplerian(
    sigma_delta_mean_anomaly_rad=0.1,
    attitude_sampler=FixedAttitude(
        quat_wxyz=(1.0, 0.0, 0.0, 0.0),
        omega_rad_s=(0.0, 0.0, 0.001745),   # 0.1 deg/s in rad/s
    ),
)

# Bandits start with a random orientation, zero spin
bandit_sampler = RelativeKeplerian(
    sigma_delta_mean_anomaly_rad=0.1,
    attitude_sampler=UniformAttitude(),
)

The attitude sampler is called only when ATTITUDE and BODY_RATES are in the side's components list. When they are absent, the sampler is ignored.

Worked example: asymmetric attitude ICs¶

Guards with a known initial slow spin, bandits tumbling randomly:

import math
import jax.numpy as jnp
from orbitalgym.config import ScenarioConfig, VehicleParamsSpec
from orbitalgym.dynamics.attitude import AttitudeParams
from orbitalgym.observations.conical import ConicalObservation
from orbitalgym.registry import (
    ActionComponentKey,
    AttitudeDynamicsKey,
    DynamicsKey,
    StateComponentKey,
)
from orbitalgym.sampling.attitude import FixedAttitude, UniformAttitudeAndRates
from orbitalgym.sampling.mass import ConstantMass
from orbitalgym.sampling.side import RelativeKeplerian
from orbitalgym.sampling.spec import ICSpec
from orbitalgym.sampling.validators import MinSeparation

guard_sampler = RelativeKeplerian(
    sigma_delta_mean_anomaly_rad=0.1,
    mass_sampler=ConstantMass(propellant_mass_kg=50.0),
    attitude_sampler=FixedAttitude(
        quat_wxyz=(1.0, 0.0, 0.0, 0.0),
        omega_rad_s=(0.0, 0.0, 0.001745),   # 0.1 deg/s about body-z
    ),
)

bandit_sampler = RelativeKeplerian(
    sigma_delta_mean_anomaly_rad=0.5,
    mass_sampler=ConstantMass(propellant_mass_kg=30.0),
    attitude_sampler=UniformAttitudeAndRates(
        omega_max_rad_s=(0.05, 0.05, 0.1),  # up to ~3–6 deg/s tumble
    ),
)

boresights = jnp.array([
    [ 1.0,  0.0, 0.0],
    [-1.0,  0.0, 0.0],
    [ 0.0,  1.0, 0.0],
    [ 0.0, -1.0, 0.0],
])

cfg = ScenarioConfig(
    n_guards=1,
    n_bandits=1,
    # ... (epoch, reference_orbit, dt, max_horizon_s, seed, guard_params, bandit_params)
    truth_dynamics=DynamicsKey.HCW_RTN,
    guard_components=(
        StateComponentKey.RTN,
        StateComponentKey.MASS,
        StateComponentKey.ATTITUDE,
        StateComponentKey.BODY_RATES,
    ),
    bandit_components=(
        StateComponentKey.RTN,
        StateComponentKey.MASS,
        StateComponentKey.ATTITUDE,
        StateComponentKey.BODY_RATES,
    ),
    guard_action_components=(
        ActionComponentKey.IMPULSIVE_MANEUVER,
        ActionComponentKey.ATTITUDE_CONTROL,
    ),
    bandit_action_components=(
        ActionComponentKey.IMPULSIVE_MANEUVER,
        ActionComponentKey.ATTITUDE_CONTROL,
    ),
    attitude_dynamics_key=AttitudeDynamicsKey.RIGID_BODY,
    guard_attitude_params=AttitudeParams(
        inertia_diag=jnp.array([10.0, 10.0, 5.0]),
        omega_max=jnp.array([0.5, 0.5, 1.0]),
    ),
    bandit_attitude_params=AttitudeParams(
        inertia_diag=jnp.array([8.0, 8.0, 4.0]),
        omega_max=jnp.array([0.5, 0.5, 0.5]),
    ),
    ic_sampler=ICSpec(
        guard_sampler=guard_sampler,
        bandit_sampler=bandit_sampler,
        validators=(MinSeparation(min_separation_m=50.0),),
    ),
)

The IC sampler is called at every env.reset. Each episode the guards start with the fixed orientation + spin while the bandits start tumbling randomly — creating an asymmetric observation challenge where the guard must point a moving sensor at an unpredictably-oriented bandit.