Adding a Policy to the Zoo¶
A policy extension packages a frozen controller (an ONNX graph, a TorchScript module, a torch network, or a client for
a remote server) so that eden inference --policy <name> and eden deploy --policy <name> can run it. This page
covers the package layout, the PolicyBase contract, how weights are fetched, and the mistakes that produce a policy
that builds and then acts on the wrong input. Faithfully reproducing a policy from another stack is its own checklist:
Porting a Pretrained Policy. Server-backed policies add a transport layer on top:
Remote Policy Clients.
What a policy is¶
A policy is any object that takes the env's observations and returns actions. In Eden every policy:
- subclasses
PolicyBase(eden.extensions.policies.base) and is registered with@POLICY_REGISTRY.register(); - declares the world it was trained in through classmethods: the robot, the other scene entities, the env options, and the observation, action and command terms;
- implements
step(obs) -> {action_term_name: tensor}, andreset(envs_idx)if it holds per-env state.
The registry key is to_snake_case(ClassName) (Twist2Policy → twist2_policy). --policy also accepts the key
without _policy and the package name, as long as the package registers one policy. Policy packages are imported on
demand: the CLI imports eden.extensions.policies.<pkg>.policy when it resolves a name, so a package with optional
dependencies does not slow down or break the others.
Package layout¶
One package per policy under eden/extensions/policies/<name>/:
| File | Contents |
|---|---|
policy.py |
The registered PolicyBase subclass. Required: the CLI discovers a package by this file. |
constants.py |
The frozen I/O contract: joint-name order, observation scales, default pose, gains. |
onnx_wrapper.py / network.py / modules.py |
The actor, when it is not a plain ONNX graph. |
custom_terms/ |
Observation, action or command terms only this policy needs. |
setups.py |
Per-embodiment setups, for a server-backed client that serves several robots. |
Keep the package to the policy, its network and its contract. When Eden also has a training task for the policy, reuse
that task's terms and scene from eden/tasks/<name>/ rather than copying them: husky builds its observation and
action from eden/tasks/husky/config.py, and omnireset its scene from eden/tasks/omnireset/. A released checkpoint
and a policy trained in Eden then share one layout.
Add a runnable example under examples/policies/<name>/main.py with the metadata header described in
examples/AGENTS.md (tier: policy, its capabilities, hardware, optional_deps and run command), and tests
under tests/extensions/policies/<name>/.
A worked example: TWIST2¶
eden/extensions/policies/twist2/policy.py is a compact, complete policy that reads the env's observations. Its
skeleton, with the class constants and the __init__ checks trimmed:
from eden import actions, observations
from eden.constants import ReferenceSource
from eden.extensions.policies.base import POLICY_REGISTRY, PolicyBase, PolicyOptions
from eden.managers.terms.commands.mimic import MotionCommand
from eden.options.managers.observations import ObservationGroupOptions
from eden.options.robots.unitree_g1 import G1Dofs29
from eden.terms import ObsTerm
# Importing the custom terms registers them before any env is built.
from eden.extensions.policies.twist2.custom_terms.observations import (
twist2_imu_roll_pitch, twist2_motion_dofs_pos, twist2_motion_root,
)
from eden.extensions.policies.twist2.onnx_wrapper import Twist2Actor
# POLICY_JOINT_NAMES, KNEES_BENT_DOFS_POS, JOINT_SCALES, ... come from eden.extensions.policies.twist2.constants
@POLICY_REGISTRY.register()
class Twist2Policy(PolicyBase):
"""TWIST2 whole-body motion-tracking controller (G1 29-DOF, ONNX backend)."""
reads_env_observations = True # step() reads the observation dict the env publishes
entity_name: str = "robot"
checkpoint_path: str = "twist2/pretrained.onnx"
md5_hash: str | None = "a045b72bd49ce0993aa145f8958a5564"
def __init__(self, env, options: PolicyOptions) -> None:
super().__init__(env, options) # resolves and verifies checkpoint_path
robot = self.require_entity(self.entity_name)
# ... check robot.dofs_name, the built default pose and the observation width here ...
self._actor = Twist2Actor(self.checkpoint_path, device=str(env.device))
@classmethod
def robot_options(cls) -> G1Dofs29:
return G1Dofs29(default_dofs_pos=dict(KNEES_BENT_DOFS_POS), default_root_pos=(0.0, 0.0, 0.76))
@classmethod
def env_options(cls) -> dict:
return dict(sim_dt=0.01, sim_substeps=2, decimation=2, ...) # 50 Hz control
@classmethod
def action_options(cls, entity_name: str = "robot") -> dict:
return dict(
dofs_pos_controller=actions.ImplicitPDController.configure(
entity_name=entity_name,
dofs_name=list(POLICY_JOINT_NAMES),
scale=dict(zip(POLICY_JOINT_NAMES, JOINT_SCALES)),
reference_source=ReferenceSource.DEFAULT,
),
)
@classmethod
def command_options(cls, motion_file: str = "motion.npz", entity_name: str = "robot", ...) -> dict:
return dict(motion=MotionCommand.configure(motion_file=motion_file, entity_name=entity_name, ...))
@classmethod
def observation_options(cls, entity_name: str = "robot") -> dict:
entity = {"entity_name": entity_name}
return {
"policy": ObservationGroupOptions(
enable_corruption=False,
history_length=ACTOR_HISTORY_LENGTH,
flatten_history_dim=True,
motion_root=ObsTerm.configure(func=twist2_motion_root),
motion_dofs_pos=ObsTerm.configure(func=twist2_motion_dofs_pos),
base_ang_vel=ObsTerm.configure(
func=observations.base_ang_vel, params={**entity, "frame": "body"}, scale=BASE_ANG_VEL_SCALE
),
imu_roll_pitch=ObsTerm.configure(func=twist2_imu_roll_pitch, params=entity),
dofs_pos=observations.DofsPos.configure(entity_name=entity_name, offset_from_default=True),
dofs_vel=observations.DofsVel.configure(entity_name=entity_name, scale=list(JOINT_VEL_SCALES)),
last_action=ObsTerm.configure(func=observations.last_action),
)
}
def step(self, obs):
history = obs["policy"] # (N, 11 * 127), oldest frame first
actor_obs = torch.cat([history[:, -ACTOR_CURRENT_DIM:], history], dim=-1)
return {"dofs_pos_controller": self._actor(actor_obs)}
The PolicyBase contract¶
Class attributes¶
| Attribute | Default | Meaning |
|---|---|---|
checkpoint_path |
None |
The main checkpoint, relative to the policy's Hugging Face dataset. __init__ replaces it with the local path, downloading on first use. |
md5_hash |
None |
The MD5 the downloaded file must match; None skips the check. |
registry |
Eden's policy repository owner | The Hugging Face owner of the dataset (a class constant). |
dataset |
"eden_policies" |
The Hugging Face dataset holding the weights (a class constant). |
local_dir |
None |
The cache directory; None means assets_subdir under the asset root (a class constant). |
assets_subdir |
"policies" |
The subdirectory of the asset root an unset local_dir means. |
reads_env_observations |
False |
Whether step() consumes the env's observation dict (see below). |
Any other annotated class attribute (entity_name, hand_name, pair, …) is a configurable parameter:
MyPolicy.configure(entity_name="g1") sets it, and the spec classmethods can read it (next section).
Classmethods: the world the policy was trained in¶
| Classmethod | Required | Returns |
|---|---|---|
robot_options() |
no | The robot options with the training pose, gains and armature, or None to use --robot. Takes no arguments. |
scene_entity_options() |
no | {name: entity options} for every non-robot entity step() or the observations read. A plane entity replaces the flat ground. |
env_options() |
no | EnvOptions fields: at least sim_dt and decimation. Takes no arguments. |
action_options() |
yes | {term_name: action term options}. |
observation_options() |
yes | {group_name: ObservationGroupOptions}. May be {}. |
command_options() |
no | {term_name: command term options}. |
In action_options, observation_options and scene_entity_options, a parameter named after one of the policy's
configurable attributes is filled from that attribute: declare hand_name: str = "hand" on the class and
action_options(cls, entity_name="robot", hand_name="hand") receives cls.hand_name. So a name is declared once,
instead of restated as a literal that can drift from it. entity_name is the caller's (--entity_name wins), and any
other parameter keeps its default, which means a knob that must keep its literal default needs a name no attribute
uses.
command_options is filled by the CLI instead: entity_name, motion_file (from --motion_file, and required
whenever the signature takes it), and clip_name (from --clip_name, when the signature takes it).
env_options: match the control rate¶
sim_dt × decimation is the time per policy step, and it must match the training setup. A policy that does not
override env_options() runs at the env's defaults and will very likely perform poorly. The shipped values:
| Policy | sim_dt |
decimation |
Control rate |
|---|---|---|---|
amo |
0.002 s | 10 | 50 Hz |
gmt |
0.001 s | 20 | 50 Hz |
holosoma, husky |
0.005 s | 4 | 50 Hz |
sonic, twist2, protomotionv3, humanoid_gpt |
0.01 s | 2 | 50 Hz |
omnixtreme |
0.004 s | 5 | 50 Hz |
metamotivo |
1/450 s | 15 | 30 Hz |
Solver settings (iterations, substeps, contact options) belong here too when the training simulator's differ.
step() and reset()¶
obskeys are the group names fromobservation_options(). A group withconcatenate_terms=Falsearrives as a dict of its terms (obs["wrist_image"]["image"]), which is how camera images are delivered.- The returned dict's keys must be the term names from
action_options(). An unknown key raises aKeyError. A term the dict leaves out is not an error: it keeps its previous action, silently. Return every term on every step. - Return tensors on
self.env.device, and run inference undertorch.no_grad()(ortorch.inference_mode()). - Override
reset(envs_idx)when the policy holds per-env state across steps (a recurrent state, a last-action history it builds itself, an action-chunk queue).eden inferencecalls it with the envs that just ended an episode;envs_idx=Nonemeans all envs.
Observations: declare them, or read raw state¶
reads_env_observations says which side of the env boundary step() sits on, and eden deploy arranges a run around
the answer, so treat it as a contract:
True: the policy owns the env. The env is built from its observation groups and action terms, and its action goes through the env's action manager, which keeps env-computed observations such aslast_actionin step with what the robot was commanded. This is the preferred form: the layout is explicit, and the observation manager handles history, scale, NaN policy and resets. Only one such policy can be loaded at a time, and not alongside a task, since one env publishes one observation layout.False: the policy reads raw state.step()reads the robot's state off the entity and builds the input itself (holosoma,homie,sonic). Such a policy can ride along as an extra slot on another policy's env, converting its own actions; itsobservation_options()may return{}.
Rules that keep an observation correct:
- The term order is the layout. Terms concatenate in declaration order (or
terms_order, when set), and the network only sees numbers: a wrong order is dimensionally valid and silently wrong. - History is oldest first. With
history_length=Handflatten_history_dim=True, a group of frame widthDarrives as(N, H * D)with the newest frame last (obs[g][:, -D:]). Withflatten_history_dim=Falseit is(N, H, D). At an episode start the window is filled with the first frame whenbackfill=True(the default), or with zeros whenbackfill=False. - Turn training noise off. Pass
enable_corruption=Falseon a group whose term definitions carry noise (for example terms shared with a training task), ashuskyandtwist2do. - A declared group needs at least one term. No groups at all builds fine; a
concatenate_terms=Truegroup with no terms fails at construction withUnable to concatenate observation terms in group `<name>`.
Entities and commands¶
- Resolve every entity in
__init__withself.require_entity(name). A scene without it then fails at construction with a message that says where the entity should come from (the policy'sscene_entity_options()in a task-less run, the task's scene with--task). Instep(), readself.env.entities[name]instead of caching the handle: deployment swaps the entity for a live-robot stand-in. - Check the built robot against the contract in
__init__. The shipped ports refuse a robot whosedofs_nameis not the policy's joint order, or whose built default pose differs from the training keyframe in any env (robot.get_default_dofs_pos(POLICY_JOINT_NAMES)is(num_envs, n)in that order), and an env whose observation group has the wrong width. Copy those checks; they turn a silent wrong-input rollout into an error. - A command-driven policy defines
set_command(command)and a_commandsbuffer.--base_command VX VY WZis padded with zeros or truncated to the width ofpolicy._commands(HOMIE's is 4-D, HUSKY's 2-D) beforeset_commandis called. A policy whose command lives on a command term exposes_commandsas a property that reads the term, asHuskyPolicydoes. - A command term a frozen policy observes must tick in
pre_physics_advance, never inpost_physics_refresh, which a deployment does not run.
Custom terms¶
A term only the policy needs goes in custom_terms/, registered like any other term (see Writing Terms):
# eden/extensions/policies/mypolicy/custom_terms/observations.py
import torch
from eden.envs.base import EnvBase
from eden.managers.observation_manager import OBSERVATION_TERM_REGISTRY
@OBSERVATION_TERM_REGISTRY.register()
def mypolicy_root_height(env: EnvBase, *, entity_name: str = "robot") -> torch.Tensor:
return env.entities[entity_name].get_pos()[:, 2:3]
The built-in term registries find their terms by importing the modules they list, and a policy's custom_terms/ is
not among them, so import every custom term module in policy.py (at module level). The CLI imports policy.py
before it builds the env, and the import registers the terms. Prefix a policy-only term's name with the policy's
(twist2_imu_roll_pitch): term keys share one namespace across Eden.
A custom action term is needed when the robot's actuators differ from Eden's controllers (MetaMotivo's SMPL humanoid
uses MuJoCo-style affine actuators, MujocoSMPLActuator); a custom command term when the policy consumes a signal no
built-in command produces (AMOCommand, GMT's GMTMotionLibCommand).
Weights¶
PolicyBase.__init__ fetches checkpoint_path with the class's fetch_asset(), which downloads from the
registry/dataset Hugging Face dataset into resolve_local_dir() (local_dir, else assets_subdir under the asset
root), and checks md5_hash. Fetch every further file the same way, with self.fetch_asset(path):
- Never call
get_asset_path(..., local_dir=self.local_dir)yourself. An unsetlocal_dirisNoneand means "assets_subdirunder the asset root", resolved at load time; passing the rawNonefetches into the asset root instead. registry,dataset,local_dirandassets_subdirare class constants, not options: passing one toconfigure()raises. That keeps a lookup made before any instance exists (a clip catalog read bycommand_options(), a file an observation term loads) on the same path as__init__: callMyPolicy.fetch_asset(path)on the class.
| Scenario | What to do |
|---|---|
One .pt / .onnx file |
Set checkpoint_path and md5_hash. |
| Several files | Keep the main one in checkpoint_path; name the others <prefix>checkpoint_path / <prefix>md5_hash (decoder_checkpoint_path, residual_checkpoint_path) and fetch and verify them in __init__. The naming lets use_local_checkpoint(path, checkpoint_attr=...) redirect any of them. |
A model loaded from Hugging Face by its own library (metamotivo) |
Set checkpoint_path = None, registry = None, dataset = None, and load in __init__. |
Weights from another dataset (beyondmimic_moves) |
Override registry and dataset on the class. |
No released weights (limx_oli_loco) |
checkpoint_path = None, registry = None, dataset = None, and raise in __init__ when no local file was given. |
To run local weights, call MyPolicy.use_local_checkpoint(path) before constructing the policy: it sets the path
and clears the MD5 pinned to the released file together. To fetch from a mirror, export EDEN_ASSETS_DIR=/data/mirror
(policies then cache under /data/mirror/policies/), or set MyPolicy.local_dir = "/data/mirror" on the class before
anything reads it.
Run it¶
eden inference --policy mypolicy --vis # task-less: the policy's own world
eden inference --policy mypolicy --async --vis # deployment-style timing
eden deploy --policy mypolicy --mujoco --vis # MuJoCo sim-to-sim (see Deployment)
From Python, build the env from the policy's spec the way a task-less run does:
import eden as en
import genesis as gs
from eden.envs.base import RLEnvBase
from eden.extensions.policies.twist2.policy import Twist2Policy
from eden.options.envs import EnvOptions
from eden.options.managers.actions import ActionManagerOptions
from eden.options.managers.commands import CommandManagerOptions
from eden.options.managers.observations import ObservationManagerOptions
from eden.utils.configs import EdenRLConfig
en.init(backend=gs.gpu)
cfg = EdenRLConfig(
env_options=EnvOptions(num_envs=1, **Twist2Policy.env_options()),
scene_options=en.scenes.Empty(robot=Twist2Policy.robot_options(), **Twist2Policy.scene_entity_options()),
action_options=ActionManagerOptions(**Twist2Policy.action_options()),
observation_options=ObservationManagerOptions(**Twist2Policy.observation_options()),
command_options=CommandManagerOptions(**Twist2Policy.command_options(motion_file="motion.npz")),
)
env = RLEnvBase.from_config(cfg, eval_mode=True)
env.build()
policy = Twist2Policy(env=env, options=Twist2Policy.configure())
obs, _ = env.reset()
policy.reset()
for _ in range(1000):
obs, *_ = env.step(policy.step(obs))
The CLI's version of this is _build_extension_env in eden/utils/cli/tasks.py, which the policy tests use too.
Calling the spec classmethods directly skips the CLI's attribute filling, so pass any parameter that names a
configured attribute yourself.
Checklist¶
-
PolicyBasesubclass ineden/extensions/policies/<name>/policy.py, decorated with@POLICY_REGISTRY.register(). -
checkpoint_pathandmd5_hashset, or the weight source handled as in Weights. -
robot_options()with the training pose, gains and armature;scene_entity_options()for every other entity the policy reads. -
env_options()with the trainingsim_dtanddecimation(and solver settings that differ). -
action_options()keys equal the keysstep()returns, on every step. - Observation groups in the training order, history oldest first, training noise off.
-
reads_env_observationsset to match whatstep()reads. -
__init__resolves every entity withrequire_entityand checks the built robot and the observation width. - Every custom term module imported in
policy.py. - An example in
examples/policies/<name>/main.pywith its metadata header, and tests intests/extensions/policies/<name>/that pin the observation and action contract. - Verified layer by layer as in Porting a Pretrained Policy.
See also¶
- The Policy Zoo: every shipped policy
- Porting a Pretrained Policy: the fidelity checklist
- Remote Policy Clients: server-backed policies
- Writing Terms: term registration and lifecycle
- Deployment: policy slots on hardware