Skip to content

T4 — Plan with short-horizon search

Wire POMDPAdapter to a small random-shooting planner that simulates candidate action sequences forward through the env and picks the best by cumulative reward, ending with one planned step from a sampled initial state. Builds on T3 — Write an active-pursuer policy; reference implementation at examples/planners/receding_horizon.py.

What this is and isn't

This is random shooting — sample K action sequences, simulate each, take the first action of the best. It is not MCTS: there is no tree, no UCB, no expansion-by-uncertainty. The teaching value is the API patterns: how to call transition / reward, how to thread PRNG keys through a horizon loop, how the action vector is laid out, and how a pure adapter composes with jax.vmap and jax.lax.scan.

For real planning, swap the search loop for mctx or pomdp-py. The wiring stays the same.

The POMDP adapter surface

    from dataclasses import dataclass

    import jax
    import jax.numpy as jnp

    from orbitalgym import (
        OrbitalGymEnv,
        POMDPAdapter,
        Side,
        make_pursuit_evasion,
    )

    cfg = make_pursuit_evasion(seed=0)
    env = OrbitalGymEnv(cfg)
    adapter = POMDPAdapter(env)

    states_dim = adapter.states_dim
    action_dim_per_side = adapter.action_dim_per_side
    n_g = cfg.n_guards
    n_b = cfg.n_bandits
    full_action_dim = (n_g + n_b) * action_dim_per_side

POMDPAdapter exposes a POMDPs.jl-shaped interface over OrbitalGymEnv:

  • states_dim — flat-state dimensionality. The vector packs StateLayout.flatten (per-side guard/bandit truth), then two scalar tail entries, t and step, then the game's per-bandit dwell counters. The reference orbit is captured at __init__ time as a per-scenario constant, so it doesn't need to live in the flat vector.
  • action_dim_per_side — components of Δv per vehicle (3 for RTN dynamics, 2 for in-plane).
  • discount() → float — episode-bounded games return 1.0. Override by subclassing if your problem needs γ < 1.
  • initialstate(key) → s_flat — sample an initial state and return its flat vector.
  • transition(s_flat, a_flat, key) → s_flat' — apply one macro step: action_repeat env substeps of the same action, one by default.
  • observation(s, a, s', side) → obs — per-side observation at s'.
  • reward(s, a, s', side) → float — per-side scalar reward.

The action vector layout is guards-first then bandits. For n_g guards and n_b bandits, each with d = action_dim_per_side components, the flat action is

a_flat = [guard_0_dv(d), ..., guard_{n_g-1}_dv(d),
          bandit_0_dv(d), ..., bandit_{n_b-1}_dv(d)]

with total length (n_g + n_b) * d.

Sample an initial state

    s0 = adapter.initialstate(jax.random.PRNGKey(0))

s0 is a flat (states_dim,) array — the per-side truth, the two-scalar (t, step) tail, and the dwell counters. There is no hidden mutable bookkeeping; every call to transition reads its bookkeeping out of the input vector and writes the new bookkeeping into the output vector.

Build the planner

    @dataclass(frozen=True)
    class RecedingHorizonRollout:
        """Random-shooting planner over a flat-vector POMDPAdapter.

        Sample ``n_candidates`` action sequences of length ``horizon``,
        simulate each forward through ``adapter.transition``, score by
        cumulative reward to ``controlled_side``, and return the first
        action of the best-scoring sequence.

        ``POMDPAdapter`` is pure, so the entire search compiles to a
        single ``jax.vmap(jax.lax.scan)`` call: candidates run in
        parallel on whatever device JAX has, and each candidate's
        horizon rolls forward inside a scanned step.
        """

        adapter: POMDPAdapter
        controlled_side: Side
        n_candidates: int = 32
        horizon: int = 3
        dv_scale: float = 0.05

        def act(self, s_flat: jax.Array, key: jax.Array) -> jax.Array:
            n_g = self.adapter.env.config.n_guards
            n_b = self.adapter.env.config.n_bandits
            d = self.adapter.action_dim_per_side
            action_dim = (n_g + n_b) * d

            def score_sequence(seq_key):
                keys = jax.random.split(seq_key, self.horizon + 1)
                actions = jax.random.normal(keys[0], (self.horizon, action_dim)) * self.dv_scale

                def step(carry_state, idx):
                    a = actions[idx]
                    s_next = self.adapter.transition(carry_state, a, keys[idx + 1])
                    r = self.adapter.reward(carry_state, a, s_next, self.controlled_side)
                    return s_next, r

                _, rewards = jax.lax.scan(step, s_flat, jnp.arange(self.horizon))
                return rewards.sum(), actions[0]

            seq_keys = jax.random.split(key, self.n_candidates)
            scores, first_actions = jax.vmap(score_sequence)(seq_keys)
            best = jnp.argmax(scores)
            return first_actions[best]

The planner is a single jax.vmap(jax.lax.scan):

  1. Outer vmap over K candidates — each candidate gets its own PRNG sub-key. JAX runs them in parallel on whatever device is available.
  2. Inner lax.scan over H horizon steps — transition advances the flat state, reward scores it for the controlled side, and the scan accumulates the per-step rewards.
  3. argmax over the per-candidate cumulative rewards picks the best sequence; the planner returns its first action.

Because POMDPAdapter is pure, all of this compiles to a single JIT call. No Python loops, no per-candidate state snapshots — the horizon-rollout primitive is a JAX function from (seq_key) → (score, first_action) and vmap does the parallelism.

Plan one step

    planner = RecedingHorizonRollout(
        adapter=adapter,
        controlled_side=Side.BANDIT,
        n_candidates=8,
        horizon=3,
    )
    key_plan, key_step = jax.random.split(jax.random.PRNGKey(1))
    a_planner = planner.act(s0, key_plan)
    s1 = adapter.transition(s0, a_planner, key_step)
    r_bandit = adapter.reward(s0, a_planner, s1, Side.BANDIT)

a_planner has shape (full_action_dim,) = ((n_g + n_b) * d,) — the same flat layout the adapter consumes. Feeding it back into adapter.transition advances one real env step from s0, and adapter.reward(s0, a_planner, s1, Side.BANDIT) scores it from the bandit's perspective.

When to override discount

POMDPAdapter.discount() returns 1.0 by default because the bundled games (make_pursuit_evasion, make_lady_bandit_guard, make_sun_blocking) are episode-bounded — terminal conditions and max_horizon_s cap the trajectory, so undiscounted cumulative reward is well-defined. Pass POMDPAdapter(env, discount=γ) when you wire the adapter into an unbounded MDP or want to emphasise near-term reward; discount() then reports γ ** action_repeat, the per-macro-step discount. The planner above ignores discount; a serious search loop should multiply it into the inner-loop accumulator.

Going further

  • mctx — DeepMind's JAX-native MCTS library. Plugs into a pure dynamics function; POMDPAdapter.transition is exactly that, so it slots in as recurrent_fn directly.
  • pomdp-py — Python POMDPs framework with a similar interface to POMDPAdapter. The flat vector here is intentionally compatible with that ecosystem.

Next

→ T5 — Run on GPU / MPS moves the same rollouts off CPU and onto an accelerator.