Policy adapters

ChunkedActionPolicy is the minimal protocol RolloutRecorder drives. Any object with select_action(batch) and reset() satisfies it, which is the same shape a LeRobot policy exposes, so a LeRobot policy or your own wrapper around a vision-language-action checkpoint both work without an adapter class in this library.

from typing import Any, Protocol


class ChunkedActionPolicy(Protocol):
    def select_action(self, batch: dict[str, Any]) -> Any: ...
    def reset(self) -> None: ...

select_action receives a LeRobot-shaped batch with observation.state, observation.images.<name>, and task keys, and returns one six-element action in degrees. Policies that internally predict a chunk of future actions should cache that chunk and return one action per call, then clear the cache in reset().

End-to-End Usage

import gymnasium as gym
import numpy as np
from lerobot.datasets.lerobot_dataset import LeRobotDataset
from so101_nexus import SO101_JOINT_NAMES
from so101_nexus.policy_adapters import RolloutRecorder
from so101_nexus.teleop.dataset import FieldSelection, build_features
import so101_nexus.mujoco  # noqa: F401


class MyPolicy:
    """Replace with your own model. Returns absolute joint targets in degrees."""

    def select_action(self, batch):
        return np.zeros(len(SO101_JOINT_NAMES), dtype=np.float32)

    def reset(self):
        pass


policy = MyPolicy()

env = gym.make("MuJoCoPickLift-v1", render_mode="rgb_array", control_mode="pd_joint_pos")
action_features = {f"{name}.pos": float for name in SO101_JOINT_NAMES}
follower_features = {**action_features, "wrist": (480, 640, 3), "overhead": (480, 640, 3)}
features = build_features(FieldSelection(), follower_features, action_features)
dataset = LeRobotDataset.create(
    repo_id="local/policy-rollouts",
    fps=30,
    features=features,
    robot_type="sim_so_follower",
    use_videos=True,
)
recorder = RolloutRecorder(env, policy, dataset=dataset, max_steps_per_episode=160)
recorder.record_episodes(n=10, seed=0)
dataset.finalize()
env.close()

RolloutRecorder expects env observations with state, overhead_camera, and wrist_camera keys by default. The env state and env actions stay in radians. The recorder converts state to degrees for the policy and dataset, and converts the returned policy action back to radians for env.step(action).

Image Key Alignment

RolloutRecorder.camera_keys and the image keys your policy reads are two sides of the same contract. By default the recorder maps ("overhead_camera", "wrist_camera") to observation.images.overhead and observation.images.wrist. If you customize one side, customize the other side to match or the policy will fail at inference when a requested image key is missing.

Side Video Channel

Pass record_side_video=True to record the env's configured render view as an extra observation.images.side video channel alongside the rollout dataset. This requires the env to be constructed with render_mode="rgb_array"; pair it with RenderConfig(camera="side") for an angled tabletop bystander view. The side frame is visualization-only: it is rendered before each step (so it depicts the same state as the rest of the frame) and never enters policy batches, so the policy still sees exactly the overhead and wrist cameras it was trained on. Declare the channel in the dataset schema with FieldSelection(side_image=True) and a "side" entry in follower_features sized to the env's RenderConfig resolution.

Caveats

Action Semantics

The recorder treats policy actions as absolute joint positions in degrees, then converts those positions to radians and clips them before env.step(action). It does not integrate deltas. If a real-model rollout visibly drifts away from rest or immediately saturates the clip, validate the assumption by disabling your policy's internal chunk replay so every call re-queries the model, and printing state_deg next to action_deg for a few steps.

Calibration Gap

A checkpoint trained on physical demonstrations assumes that robot's homing pose. If your simulated rest pose sits far from the training distribution, for example a shoulder lift near -90 deg locally against a training corpus centered near +123 deg, zero-shot rollouts will look wrong even when the plumbing is correct. That is a hardware and dataset alignment issue. Rehome to a compatible setup or fine-tune on demonstrations recorded from your own setup.

Untrusted Checkpoints

Loading a policy from a remote checkpoint can execute code from that repository, not just deserialize weights. Loader flags such as trust_remote_code=True run the publisher's Python even when the weights themselves are a safe, weights-only format. Prefer weights-only formats, leave remote-code execution off unless you have read the code, pin the revision you load, and treat a checkpoint from an untrusted publisher the same way you would treat running that publisher's code.