Adding a new navigation policy

May 30, 2026 · View on GitHub

This guide walks through plugging a new in-process navigation policy into wanderland-lab. The benchmark calls policies directly in-process — there is no client/server split. The only extension point is the Policy protocol at src/wanderland_lab/policies/base.py.

Worked references in the repo:

File	Role
`src/wanderland_lab/policies/base.py`	`Policy` protocol + `Action` + `PolicyCapabilities`
`src/wanderland_lab/policies/random.py`	Trivial reference (~60 lines)
`src/wanderland_lab/policies/citywalker.py`	Full example with batched forward
`configs/benchmark/policy/random.yaml`	Hydra wiring (minimal)
`configs/benchmark/policy/citywalker.yaml`	Hydra wiring (with nested model + config)

1. The `Policy` protocol

A policy is anything that satisfies this protocol:

class Policy(Protocol):
    @property
    def capabilities(self) -> PolicyCapabilities: ...
    def reset(self, session_id: str) -> None: ...
    def act(self, observation: Mapping[str, Any]) -> Action: ...
    def act_batch(
        self, observations: Sequence[Mapping[str, Any]]
    ) -> List[Action]: ...

`act(observation) -> Action`

The env guarantees observation is a Mapping[str, Any] containing at least these keys:

Key	Type	Notes
`images`	`np.ndarray (T, H, W, 3) uint8`	`T == capabilities.context_size`; oldest first, newest last
`ego_pos`	`np.ndarray (3,) float`	World-frame XYZ (Isaac is Z-up)
`ego_yaw`	`float`	World-frame yaw in radians
`goal`	`np.ndarray (3,) float`	World-frame goal XYZ
`dt`	`float`	Control timestep in seconds
`session_id`	`str`	Stable id for the current episode/env (see `reset`)

Optional keys, present only when capabilities request them:

Key	When	Type
`ref_path`	Always optional; passed when the dataset has a route	`Sequence[np.ndarray (3,)]` of world-frame waypoints
`goal_image`	When `capabilities.needs_goal_image is True`	`np.ndarray (H, W, 3) uint8`

Treat any other key as undefined — don't rely on it.

`act_batch(observations) -> List[Action]`

Vectorized variant. The env calls this when running num_envs > 1 so the policy can stack inputs into a single GPU forward. The default implementation should just loop over act; override only if real batching is a meaningful speedup. CityWalker stacks (B, T, 3, H, W) images and (B, k+1, 2) body coords, runs one model.forward, then slices outputs per-env.

`reset(session_id)`

Called once at the start of each episode (per-env). Use it to clear any per-session state — history ring buffers, per-episode RNGs, frame counters. The env disambiguates concurrent envs by passing distinct session_ids, so a single policy instance must support multiple parallel sessions if you want num_envs > 1. Pattern:

self._sessions.setdefault(session_id, _Session()).reset()

`capabilities`

Static traits the env reads once to set itself up. Don't change these mid-episode.

2. `PolicyCapabilities` fields

@dataclass(frozen=True)
class PolicyCapabilities:
    context_size: int
    holonomic: bool = False
    needs_goal_image: bool = False
    rgb_normalize: bool = True

Field	What it controls
`context_size`	The env keeps a ring buffer of the last `T == context_size` RGB frames and passes them as `obs["images"]`. Set to `1` for stateless single-frame policies.
`holonomic`	If `False`, the env zeros out `vy` in the executed action — non-holonomic agents only get forward + yaw. Set `True` only if your model genuinely predicts strafing.
`needs_goal_image`	If `True`, the env captures one RGB frame at the goal pose at episode start and includes it as `obs["goal_image"]`. Used by image-goal policies.
`rgb_normalize`	If `True`, your policy expects RGB in `[0, 1]` floats and the env / model wrapper handles ImageNet mean/std internally. If `False`, you take raw uint8 frames and do your own normalization.

3. The `Action` type

@dataclass(frozen=True)
class Action:
    vx: float        # body-frame forward velocity (m/s)
    vy: float        # body-frame lateral velocity (m/s); zeroed if not holonomic
    yaw_rate: float  # body-frame yaw rate (rad/s)

All actions are body-frame velocities. The env handles the body→world conversion before stepping physics. Sign conventions and the body→world mapping are documented in docs/coordinate_frames.md; the short version: at yaw=0, +vx moves toward −world X, +vy toward +world Y.

4. Walked example: a constant-forward policy

A minimal policy that drives the agent forward at a fixed speed, ignoring images and goals. Useful for end-to-end env smoke tests.

src/wanderland_lab/policies/constant_forward.py:

"""Constant-forward policy: ignores observations, drives forward at vx_const."""

from __future__ import annotations

from typing import Any, List, Mapping, Sequence

from .base import Action, Policy, PolicyCapabilities


class ConstantForwardPolicy(Policy):
    def __init__(self, vx_const: float = 0.5, context_size: int = 1):
        self._capabilities = PolicyCapabilities(
            context_size=int(context_size),
            holonomic=False,        # only forward + yaw
            needs_goal_image=False,
            rgb_normalize=True,
        )
        self.vx_const = float(vx_const)

    @property
    def capabilities(self) -> PolicyCapabilities:
        return self._capabilities

    def reset(self, session_id: str) -> None:
        return None  # stateless

    def act(self, observation: Mapping[str, Any]) -> Action:
        return Action(vx=self.vx_const, vy=0.0, yaw_rate=0.0)

    def act_batch(
        self, observations: Sequence[Mapping[str, Any]]
    ) -> List[Action]:
        return [self.act(obs) for obs in observations]


__all__ = ["ConstantForwardPolicy"]

That's it — ~30 lines, runnable.

5. Wiring it into the benchmark CLI

Policies are instantiated by Hydra. Drop a YAML at configs/benchmark/policy/<name>.yaml and the benchmark CLI will pick it up via policy=<name>.

configs/benchmark/policy/constant_forward.yaml:

# Constant-forward smoke-test policy.
_target_: wanderland_lab.policies.constant_forward.ConstantForwardPolicy
vx_const: 0.5
context_size: 1

Run it on a GPU host (the benchmark CLI needs Isaac Sim). Use absolute paths — Hydra switches the process to a per-run working directory:

DATA=$(pwd)/data/wanderland

OMNI_KIT_ACCEPT_EULA=YES uv run wanderland-lab benchmark \
    policy=constant_forward \
    scene_id=<scene_id> \
    env.usd_path=${DATA}/<scene_id>/scene.usdz \
    env.episodes_json=${DATA}/<scene_id>/episodes.json \
    output.dir=$(pwd)/outputs/constant_forward

The wanderland-lab benchmark console script wraps python -m wanderland_lab.benchmark.cli (and sets the LD_PRELOAD=libcarb.so bootstrap Isaac Sim needs). Hydra calls ConstantForwardPolicy(vx_const=0.5, context_size=1) and hands the result to the env.

Nested kwargs

For policies with non-trivial constructor args (a model object, a config dataclass, ...), nest them with their own _target_. CityWalker is the canonical example:

_target_: wanderland_lab.policies.citywalker.CityWalkerPolicy
model:
  _target_: transformers.AutoModel.from_pretrained
  pretrained_model_name_or_path: ai4ce/citywalker
  trust_remote_code: true
policy_config:
  _target_: wanderland_lab.policies.citywalker.CityWalkerPolicyConfig
  vx_max: 0.8
  append_goal_row: true     # feed the goal as the last coords row
  normalize_coords: true
  target_horizon_m: 5.0
  step_scale: 1.0
device: cuda

Hydra recursively instantiates model and policy_config first, then passes them to CityWalkerPolicy.__init__. (See reproducing_baselines.md for what the goal-row fields do — they are what makes CityWalker actually move toward the goal.)

Override at the CLI

Any field can be overridden without editing the file:

OMNI_KIT_ACCEPT_EULA=YES uv run wanderland-lab benchmark \
    policy=constant_forward \
    policy.vx_const=0.8 \
    scene_id=<scene_id> \
    env.usd_path=... env.episodes_json=...

6. Optional: implementing `act_batch` for vectorized speedup

For num_envs == 1, the default act_batch (loop over act) is fine. If you plan to run multiple envs in parallel and your model takes a real GPU forward, override act_batch to do one batched forward instead of B sequential ones.

Pattern (from CityWalkerPolicy):

@torch.inference_mode()
def act_batch(self, observations):
    if not observations:
        return []
    # 1) Per-env preprocessing: history update, build per-env tensors.
    per_env = [self._prepare_inputs(obs) for obs in observations]

    # 2) Stack to (B, ...) and run ONE forward.
    imgs_t   = torch.cat([p.imgs_t   for p in per_env], dim=0)
    coords_t = torch.cat([p.coords_t for p in per_env], dim=0)
    out = self.model(imgs_t, coords_t)

    # 3) Per-env postprocessing: slice outputs into Actions.
    return [self._build_action(p, out[i]) for i, p in enumerate(per_env)]

def act(self, observation):
    # Keep single-env path consistent with batched: just batch-of-1.
    return self.act_batch([observation])[0]

Two non-obvious things about this pattern:

Preserve per-session state under batching. Each env has its own session_id; preprocessing must update the right _Session (history buffer etc.) before tensor stacking, not after.
Make act a thin wrapper over act_batch([obs])[0]. Otherwise the single-env and batched code paths drift, and parity with prior single-env runs breaks.

If your model is small / CPU-only / allocates per call, skip this — the default loop is faster than the bookkeeping overhead.