Skip to content

Controlled-side policy cookbook

Patterns for the side under your control. Each pattern is a recipe, not a single fixed implementation — the gallery (linked from each section) ships canonical implementations of three of these.

Pattern: scripted-with-fallback

A learned primary policy backed by a known-safe fallback that takes over when a confidence signal drops below threshold. The confidence signal is threaded through policy_state (typically: belief covariance trace, critic-head value, or any other uncertainty proxy).

    from orbitalgym.actions.assemble import build_command_class
    from orbitalgym.actions.components import ImpulsiveManeuver
    from orbitalgym.policies import ZeroControl
    from orbitalgym.policies.controlled import (
        CompositeActionPolicy,
        HeuristicWithFallbackPolicy,
    )
    from orbitalgym.registry import Frame

    @dataclass(frozen=True)
    class _Primary:
        n_vehicles: int = 1
        command_cls: Any = None

        def __call__(self, ps, agent_view, key, t):
            dv = jnp.broadcast_to(jnp.array([0.05, 0.0, 0.0]), (self.n_vehicles, 3))
            return self.command_cls.zeros(self.n_vehicles).replace(dv=dv), ps

    primary = _Primary(command_cls=cmd_cls)
    fallback = ZeroControl(command_cls=cmd_cls, n_vehicles=1)
    policy = HeuristicWithFallbackPolicy(
        primary=primary,
        fallback=fallback,
        threshold=0.5,
        n_vehicles=1,
        command_cls=cmd_cls,
    )
    # Pass confidence through policy_state. High confidence → primary.
    confident = jnp.asarray(0.9)
    cmd_high, _ = policy(confident, jnp.zeros(12), jax.random.PRNGKey(0), jnp.asarray(0))
    # Low confidence → fallback (ZeroControl).
    unconfident = jnp.asarray(0.1)
    cmd_low, _ = policy(unconfident, jnp.zeros(12), jax.random.PRNGKey(0), jnp.asarray(0))

Both branches are evaluated every step (this is JAX-native — jnp.where is the right primitive). The trade-off is intentional: tracing-friendly routing always pays both costs.

Gallery: HeuristicWithFallbackPolicy.

Pattern: belief-aware

A policy that consumes the per-side Belief directly. With the unified protocol, belief threading is the rollout's job — drive via belief_rollout (see Core API → Rollout) and the policy receives the post-update belief as agent_view. Read agent_view.mean (or richer belief statistics on future Belief types) to drive the action. No wrapper class is needed.

    # Belief-aware policies receive the per-side Belief as `agent_view` when
    # wired through `belief_rollout`. They read `belief.mean` directly — no
    # `BeliefConditionedPolicy` wrapper needed.
    @dataclass(frozen=True)
    class _MeanReadingPolicy:
        n_vehicles: int = 1
        command_cls: Any = None

        def __call__(self, ps, agent_view, key, t):
            # `agent_view.mean` is the belief's state estimate; here we shape
            # a small Δv from the first three components for illustration.
            mean3 = agent_view.mean[:3]
            dv = jnp.broadcast_to(mean3 * 0.001, (self.n_vehicles, 3))
            return self.command_cls.zeros(self.n_vehicles).replace(dv=dv), ps

    @dataclass(frozen=True)
    class _StubBelief:
        mean: jax.Array

    pol = _MeanReadingPolicy(command_cls=cmd_cls)
    bel = _StubBelief(mean=jnp.array([10.0, 20.0, 30.0]))
    bc_cmd, _ = pol(None, bel, jax.random.PRNGKey(0), jnp.asarray(0))

The same Policy works for obs-only rollout (where agent_view is the flat obs) and belief-aware belief_rollout (where it's the Belief). Memoryless reactive policies treat agent_view as the obs vector; belief-consumers read .mean. The dispatch is duck-typed.

Pattern: composite action

A learned offset policy summed onto a scripted base. Lets you add a small NN correction to a known-reasonable controller — usually faster to train than learning the entire control law from scratch.

    @dataclass(frozen=True)
    class _LearnedOffset:
        n_vehicles: int = 1
        command_cls: Any = None

        def __call__(self, ps, agent_view, key, t):
            # In real use: a NN forward pass.
            dv = jnp.broadcast_to(jnp.array([0.01, 0.0, 0.0]), (self.n_vehicles, 3))
            return self.command_cls.zeros(self.n_vehicles).replace(dv=dv), ps

    base = ZeroControl(command_cls=cmd_cls, n_vehicles=1)
    composite = CompositeActionPolicy(
        base=base, offset=_LearnedOffset(command_cls=cmd_cls), n_vehicles=1, command_cls=cmd_cls
    )
    cmd, _ = composite(None, jnp.zeros(12), jax.random.PRNGKey(0), jnp.asarray(0))

Gallery: CompositeActionPolicy.

Pattern: A/B harness

To compare two candidate policies head-to-head against the same heuristic opponent, run them in parallel under jax.vmap over a seed batch and compute the metric difference. This is a recipe, not a class — see How-to → vmap rollouts for the underlying mechanics.

The pattern: vmap the rollout over a jax.random.split(key, n_seeds), swap the controlled-side policy across two batches, and reduce the per-seed metrics with jnp.mean / jnp.std.

See also