Customize observations¶

cfg.guard_observation_fn and cfg.bandit_observation_fn are per-side ObservationFn callables. Each returns a tuple of Observation channels — one per sensor modality. Belief updaters fold each channel sequentially.

For the conceptual story (what's in a channel, scope rules, the four bundled channels), read In depth → Observations. This page is the practical "how to write one" track.

The protocol¶

orbitalgym.observations.base.ObservationFn ¶

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.

orbitalgym.observations.types.Observation ¶

One sensor channel's measurement payload.

Shape conventions

obs: (N_obs, N_total, m) per-pair measurement visible: (N_obs, N_total) bool per-pair visibility mask obs_matrix: (m, d) linear measurement matrix H obs_noise: (m, m) or (N_obs, N_total, m, m) measurement noise R

obs_noise is either shared across every pair, shape (m, m), or conditioned on the pair, shape (N_obs, N_total, m, m). Use noise_for(i, j) to read the (m, m) block for one pair regardless of which form a channel supplies.

obs_fn (optional): nonlinear measurement function h(x) -> z. When set, the EKF updater computes H = jax.jacfwd(obs_fn)(predicted_mean) per pair and uses obs_fn directly for the innovation. The KF updater rejects channels with obs_fn set.

visibility_score_fn (optional): closure (particles: (N_obs, N_total, K, d)) -> (N_obs, N_total, K) returning a signed-distance score to the channel's visibility boundary in native units — positive inside, negative outside. Consumed by the particle filter for negative-information updates on non-detections. Non-gated channels (e.g. GPS) leave it None.

A gated channel must return OUT_OF_SCOPE_SCORE for every pair its visible mask can never set — self pairs and same-side pairs. Those pairs are outside the sensor's reach rather than inside its gate and undetected, so a non-detection says nothing about them; scoring them normally would let the negative-information update crush the observer's own-state particles.

noise_for ¶

noise_for(i, j) -> jax.Array

Measurement noise (m, m) for the (observer i, entity j) pair.

Accepts traced indices, so belief updaters can call it from inside a vmap over the observer and entity axes.

The four bundled channels as a vocabulary¶

Channel	Module	Use when
`FullObservation`	`observations/reference.py`	Bootstrap; observer sees full truth (own + opp).
`OnboardGPSObservation`	`observations/onboard_gps.py`	Realistic single-side GPS-style noise on own state.
`RangeLimitedObservation`	`observations/range_limited.py`	Pair visibility gated on range; multi-vehicle realism.
`CompositeObservation`	`observations/composite.py`	Glue: combine several channels into one tuple.

Worked example: a position-only channel¶

A custom channel that drops velocity and exposes only position:

    import dataclasses
    from dataclasses import dataclass
    from typing import Any

    import jax.numpy as jnp

    from orbitalgym import OrbitalGymEnv, SingleAgentView, make_pursuit_evasion
    from orbitalgym.observations.types import Observation

    @dataclass(frozen=True)
    class PositionOnlyObservation:
        """One channel: own + opponent position only (drops velocity).

        Returns shape `(N_self, N_total, 3)` for RTN dynamics. The
        obs_matrix projects the full `d`-dim state to its first 3
        position components.
        """

        layout: Any  # carries n_guards, n_bandits, dynamics_state_dim

        def __call__(self, env_state, actions, side, params, key, t):
            del actions, params, key, t
            from orbitalgym.belief._common import _truth_arrays_for_side

            own_truth, opp_truth = _truth_arrays_for_side(env_state, side.value)
            n_self = own_truth.shape[0]
            n_tgt = opp_truth.shape[0]
            n_total = n_self + n_tgt
            d = self.layout.dynamics_state_dim
            m = 3  # position-only

            stacked = jnp.concatenate([own_truth, opp_truth], axis=0)
            positions = stacked[:, :m]  # (N_total, 3)
            obs = jnp.broadcast_to(positions[None, :, :], (n_self, n_total, m))
            visible = jnp.ones((n_self, n_total), dtype=bool)
            H = jnp.zeros((m, d)).at[:m, :m].set(jnp.eye(m))  # noqa: N806
            R = jnp.eye(m) * 1e-6  # noqa: N806
            return (Observation(obs=obs, visible=visible, obs_matrix=H, obs_noise=R),)

Wire it onto the cfg:

    cfg = make_pursuit_evasion(seed=0, max_horizon_s=200.0)
    cfg = dataclasses.replace(
        cfg,
        guard_observation_fn=PositionOnlyObservation(layout=cfg.layout),
        bandit_observation_fn=PositionOnlyObservation(layout=cfg.layout),
    )
    env = OrbitalGymEnv(cfg)

cfg.layout is the canonical StateLayout for the scenario — a small struct bundling the per-side state dataclass types, vehicle counts, flat-vector dimension, and pytree↔flat conversion functions. Pass it to your custom channel (as above) so dynamics_state_dim, n_guards, and n_bandits line up with the rest of the scenario. See API → Components → State for the full class reference.

Composing channels¶

To compose, return a length-2+ tuple from your ObservationFn, or use CompositeObservation from the reference module. The belief updater folds the channels in tuple order.