Policies

A policy extension packages a frozen policy so Eden can run it: a local actor (a checkpoint or an ONNX graph with its observation and action contract) or a client of a policy server (a VLA or world action model). Every policy is a PolicyBase subclass registered in POLICY_REGISTRY, run with eden inference --policy <name>. The policy zoo guide lists what each one needs, the remote policies guide explains server-backed clients, and extending policies shows how to add one.

from eden.extensions.policies.base import POLICY_REGISTRY, PolicyBase
from eden.options.extensions.policies import PolicyOptions, RemotePolicyOptions

Base class and registry

Base classes for policy extensions.

ClientMixin (server-backed policies) lives in :mod:eden.extensions.policies.client.

Attributes

POLICY_REGISTRY module-attribute

POLICY_REGISTRY = Registry('POLICY')

Classes

PolicyBase

PolicyBase(env: EnvBase, options: PolicyOptions)

Bases: ConfigurableMixin[PolicyOptions]

Base class for policies.

Any object that takes observations and returns actions are defined as a policy.

The *_options classmethods below describe the world the policy was trained in, and a task-less run (eden inference|deploy --policy) calls them for you. In :meth:action_options, :meth:observation_options and :meth:scene_entity_options, a parameter named after one of the policy's configurable attributes is filled from it (hand_name gets self.hand_name), so a name is declared once; entity_name is the caller's (--entity_name wins) and anything else keeps its own default. Parameters there therefore share a namespace with the class attributes — a knob that must keep its literal default needs a name no attribute uses. :meth:command_options is filled by the CLI instead (entity_name, motion_file, --clip_name), and :meth:robot_options / :meth:env_options take no arguments at all.

:attr:reads_env_observations says which side of the env boundary the policy's step() sits on, and a deploy is arranged around the answer.

Where a policy's files come from (registry, dataset, local_dir, assets_subdir) is fixed by the class, not configured per instance: a lookup that runs before any instance exists (a clip catalogue in :meth:command_options, a file an observation term loads) must read the same place __init__ does. To fetch from a mirror (an offline cluster, a shared read-only cache), set EDEN_ASSETS_DIR, or set local_dir on the class before anything reads it: MyPolicy.local_dir = "/data/mirror", as :meth:use_local_checkpoint redirects a checkpoint.

PARAMETER	DESCRIPTION
checkpoint_path	The path to the checkpoint file.

Source code in eden/extensions/policies/base.py

def __init__(self, env: EnvBase, options: PolicyOptions):
    super().__init__(options)
    self.env = env

    if self.checkpoint_path is not None:
        self.checkpoint_path = self.fetch_asset(self.checkpoint_path)
    self._verify_md5(self.checkpoint_path, self.md5_hash)

Attributes

reads_env_observations class-attribute instance-attribute

reads_env_observations: bool = False

checkpoint_path class-attribute instance-attribute

checkpoint_path: str | None = None

md5_hash class-attribute instance-attribute

md5_hash: str | None = None

registry class-attribute

registry: str | None = 'Kashu7100'

dataset class-attribute

dataset: str | None = 'eden_policies'

local_dir class-attribute

local_dir: str | None = None

assets_subdir class-attribute

assets_subdir: str = 'policies'

env instance-attribute

env = env

Methods:

resolve_local_dir classmethod

resolve_local_dir() -> str

Where this policy's files are cached: local_dir, else assets_subdir under the asset root.

Resolved on each call, so EDEN_ASSETS_DIR moves it. To fetch a file, call :meth:fetch_asset.

Source code in eden/extensions/policies/base.py

@classmethod
def resolve_local_dir(cls) -> str:
    """Where this policy's files are cached: ``local_dir``, else ``assets_subdir`` under the asset root.

    Resolved on each call, so ``EDEN_ASSETS_DIR`` moves it. To fetch a file, call :meth:`fetch_asset`.
    """
    return resolve_local_dir(cls.local_dir, cls.assets_subdir)

fetch_asset classmethod

fetch_asset(file: str) -> str

The local path of file from this policy's Hub dataset, downloaded if missing.

Fetches with the class's registry, dataset and :meth:resolve_local_dir, so it gives the same path on the class (a lookup before any env exists) as on an instance. Call this rather than passing local_dir (None by default) to get_asset_path yourself. Unlike EntityOptions.fetch_asset, file is required: a policy has no single file of its own.

Source code in eden/extensions/policies/base.py

@classmethod
def fetch_asset(cls, file: str) -> str:
    """The local path of ``file`` from this policy's Hub dataset, downloaded if missing.

    Fetches with the class's ``registry``, ``dataset`` and :meth:`resolve_local_dir`, so it gives the same
    path on the class (a lookup before any env exists) as on an instance. Call this rather than passing
    ``local_dir`` (``None`` by default) to ``get_asset_path`` yourself. Unlike ``EntityOptions.fetch_asset``,
    ``file`` is required: a policy has no single file of its own.
    """
    return get_asset_path(file, registry=cls.registry, dataset=cls.dataset, local_dir=cls.resolve_local_dir())

use_local_checkpoint classmethod

use_local_checkpoint(path: str, *, checkpoint_attr: str = 'checkpoint_path') -> None

Point a checkpoint at a local file and drop the hash pinned to the released one.

The two go together: md5_hash pins the released weights, so setting the path without clearing it makes _verify_md5 reject the very file the caller asked for. Use this for every --ckpt-path-style flag rather than assigning the path by hand.

PARAMETER	DESCRIPTION
path	Local checkpoint to load instead of the released one. Not hashed: an unreleased file has no pinned digest to check it against. TYPE: str
checkpoint_attr	Which checkpoint to redirect, for policies that load more than one (sonic's decoder_checkpoint_path, omnixtreme's residual_/fk_). The hash attribute is derived from it by the naming convention every policy follows — <prefix>checkpoint_path is pinned by <prefix>md5_hash — and both are required to exist, so a policy that breaks the convention fails here rather than silently keeping a stale pin. TYPE: str DEFAULT: 'checkpoint_path'

Source code in eden/extensions/policies/base.py

@classmethod
def use_local_checkpoint(cls, path: str, *, checkpoint_attr: str = "checkpoint_path") -> None:
    """Point a checkpoint at a local file and drop the hash pinned to the released one.

    The two go together: ``md5_hash`` pins the *released* weights, so setting the path
    without clearing it makes ``_verify_md5`` reject the very file the caller asked
    for. Use this for every ``--ckpt-path``-style flag rather than assigning the path
    by hand.

    Parameters
    ----------
    path: str
        Local checkpoint to load instead of the released one. Not hashed: an
        unreleased file has no pinned digest to check it against.
    checkpoint_attr: str
        Which checkpoint to redirect, for policies that load more than one (sonic's
        ``decoder_checkpoint_path``, omnixtreme's ``residual_``/``fk_``). The hash
        attribute is derived from it by the naming convention every policy follows —
        ``<prefix>checkpoint_path`` is pinned by ``<prefix>md5_hash`` — and both are
        required to exist, so a policy that breaks the convention fails here rather
        than silently keeping a stale pin.
    """
    if not checkpoint_attr.endswith("checkpoint_path"):
        raise ValueError(
            f"checkpoint_attr must end with 'checkpoint_path' so its md5 attribute can be derived; "
            f"got {checkpoint_attr!r} on {cls.__name__}."
        )
    hash_attr = f"{checkpoint_attr.removesuffix('checkpoint_path')}md5_hash"
    for attr in (checkpoint_attr, hash_attr):
        if not hasattr(cls, attr):
            raise AttributeError(
                f"{cls.__name__} has no `{attr}`. `use_local_checkpoint` redirects "
                f"`{checkpoint_attr}` and clears `{hash_attr}` together; both must exist."
            )
    setattr(cls, checkpoint_attr, path)
    setattr(cls, hash_attr, None)

require_entity

require_entity(name: str) -> Entity

Resolve a scene entity the policy reads, failing loud when the scene has no such entity.

Call this in __init__ for every entity step() touches, so a scene missing one raises at construction rather than deep in the first step. In step() read self.env.entities[name] instead of caching the handle: deployment hot-swaps it. Resolves against the mapping, like a term's resolve_entity.

Source code in eden/extensions/policies/base.py

def require_entity(self, name: str) -> Entity:
    """Resolve a scene entity the policy reads, failing loud when the scene has no such entity.

    Call this in ``__init__`` for every entity ``step()`` touches, so a scene missing one
    raises at construction rather than deep in the first step. In ``step()`` read
    ``self.env.entities[name]`` instead of caching the handle: deployment hot-swaps it.
    Resolves against the mapping, like a term's ``resolve_entity``.
    """
    return resolve_named_entity(
        self.env.entities,
        name,
        hint=(
            f"A task-less run builds the policy's own workspace from "
            f"{type(self).__name__}.scene_entity_options(); with --task, the task's scene "
            f"must declare '{name}'."
        ),
    )

robot_options classmethod

robot_options() -> RobotOptions | None

Return a robot configured for this policy, or None for the generic robot.

Frozen policies are usually trained against a specific default joint pose and PD-gain set (e.g. OpenHomie's). Task-less rollout/deploy builds the scene from this hook so the robot spawns in the pose and gains the policy expects; returning None falls back to the caller-chosen eden.robots class with its stock defaults.

Source code in eden/extensions/policies/base.py

@classmethod
def robot_options(cls) -> RobotOptions | None:
    """Return a robot configured for this policy, or ``None`` for the generic robot.

    Frozen policies are usually trained against a specific default joint
    pose and PD-gain set (e.g. OpenHomie's). Task-less rollout/deploy builds
    the scene from this hook so the robot spawns in the pose and gains the
    policy expects; returning ``None`` falls back to the caller-chosen
    ``eden.robots`` class with its stock defaults.
    """
    return None

scene_entity_options classmethod

scene_entity_options() -> dict[str, EntityOptionsLike]

Return the non-robot scene entities the policy reads, keyed by scene entity name (optional).

A policy that only reads its own robot needs none — the default. One whose observation is defined against something else (SimToolReal's tool and table, HUSKY's skateboard) cannot run on a robot-only scene, so a task-less run builds these alongside :meth:robot_options; with --task the task's scene must declare them, which :meth:require_entity checks at construction. A plane entity replaces the flat ground (e.g. with terrain) rather than adding one.

Source code in eden/extensions/policies/base.py

@classmethod
def scene_entity_options(cls) -> dict[str, EntityOptionsLike]:
    """Return the non-robot scene entities the policy reads, keyed by scene entity name (optional).

    A policy that only reads its own robot needs none — the default. One whose observation
    is defined against something else (SimToolReal's tool and table, HUSKY's skateboard)
    cannot run on a robot-only scene, so a task-less run builds these alongside
    :meth:`robot_options`; with ``--task`` the task's scene must declare them, which
    :meth:`require_entity` checks at construction. A ``plane`` entity replaces the flat
    ground (e.g. with terrain) rather than adding one.
    """
    return {}

env_options classmethod

env_options() -> dict[str, Any]

Set up required environment parameters for the policy, if any (optional).

Source code in eden/extensions/policies/base.py

@classmethod
def env_options(cls) -> dict[str, Any]:
    """Set up required environment parameters for the policy, if any (optional)."""
    return {}

action_options abstractmethod classmethod

action_options() -> dict[str, ActionTermOptions]

Set up required action group options for the policy (required).

Source code in eden/extensions/policies/base.py

@classmethod
@abstractmethod
def action_options(cls) -> dict[str, ActionTermOptions]:
    """Set up required action group options for the policy (required)."""
    raise NotImplementedError

command_options classmethod

command_options() -> dict[str, CommandTermOptions]

Set up required command group options for the policy, if any (optional).

Source code in eden/extensions/policies/base.py

@classmethod
def command_options(cls) -> dict[str, CommandTermOptions]:
    """Set up required command group options for the policy, if any (optional)."""
    return {}

observation_options abstractmethod classmethod

observation_options() -> dict[str, ObservationGroupOptions]

Set up required observation group options for the policy (required).

Source code in eden/extensions/policies/base.py

@classmethod
@abstractmethod
def observation_options(cls) -> dict[str, ObservationGroupOptions]:
    """Set up required observation group options for the policy (required)."""
    raise NotImplementedError

step abstractmethod

step(obs: dict[str, Tensor | dict[str, Tensor]]) -> dict[str, Tensor]

Get action from policy.

Source code in eden/extensions/policies/base.py

@abstractmethod
def step(self, obs: dict[str, torch.Tensor | dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]:
    """Get action from policy."""
    raise NotImplementedError

reset

reset(envs_idx: slice | Tensor | None = None) -> None

Reset per-env policy state at episode boundaries.

Override this in subclasses that hold per-env state across steps — e.g. heading offsets cached on the first step, last-action histories fed back as observations, attention KV caches, action-chunk queues. The default is a no-op, which is right for stateless policies.

When used as a teacher / motion generator in long-horizon rollouts the owner is expected to call policy.reset(envs_idx) whenever the underlying env resets the corresponding envs.

PARAMETER	DESCRIPTION
envs_idx	Indices of environments to reset. None means all environments. TYPE: slice | Tensor | None DEFAULT: None

Source code in eden/extensions/policies/base.py

def reset(self, envs_idx: (slice | torch.Tensor) | None = None) -> None:
    """Reset per-env policy state at episode boundaries.

    Override this in subclasses that hold per-env state across steps —
    e.g. heading offsets cached on the first step, last-action histories
    fed back as observations, attention KV caches, action-chunk queues.
    The default is a no-op, which is right for stateless policies.

    When used as a teacher / motion generator in long-horizon rollouts the
    owner is expected to call ``policy.reset(envs_idx)`` whenever the
    underlying env resets the corresponding envs.

    Parameters
    ----------
    envs_idx : slice | torch.Tensor | None, optional
        Indices of environments to reset. ``None`` means all environments.
    """
    return None

Functions:

Options

Configuration options for policy extensions (local and remote clients).

Attributes

ChunkMode module-attribute

ChunkMode = Literal['client', 'server', 'single']

ShortChunk module-attribute

ShortChunk = Literal['error', 'accept']

Classes

PolicyOptions

PolicyOptions(**data)

Bases: DynamicConfigurableOptions

Options for policy extensions.

PARAMETER	DESCRIPTION
checkpoint_path	The path to the checkpoint file.

Source code in eden/options/options.py

def __init__(self, **data):
    # Seed type metadata as construction data so BaseModel validates it into
    # __pydantic_extra__ (no-op for extra="forbid" classes, which use
    # the serialization envelope instead).
    data.update(option_metadata(self))
    # NOTE: all of the option modifications/validations should be done in the __init__ method
    try:
        BaseModel.__init__(self, **data)
    except ValidationError as exc:
        # A misspelled field on an extra="forbid" class should teach the field set, not just
        # name the typo. Only rewritten when every error is extra_forbidden.
        enriched = _enrich_extra_forbidden(type(self), exc)
        if enriched is None:
            raise
        raise enriched from None

Attributes

checkpoint_path class-attribute instance-attribute

checkpoint_path: str | None = None

ClientOptions

ClientOptions(**data)

Bases: DynamicConfigurableOptions

Options for policy client extensions (ClientMixin).

PARAMETER	DESCRIPTION
host	Policy server hostname, or a full URL (wss://example.com/ws, http://node:8000) whose scheme matches the transport.
port	Policy server port. None uses the port embedded in host (or the scheme's default when there is none); a port in host that disagrees with port raises.
request_timeout	Request timeout in seconds for HTTP transports.
set_task_endpoint	The server's set-task API (HTTP route); None (default) when the server has none and takes the task with each request.
reset_endpoint	The server's reset API: an HTTP route, a ZMQ endpoint name or a GalaxeaVLA reset message key; None (default) when the server has none, so a reset only clears local state.
act_endpoint	HTTP endpoint for requesting actions.
task_description	Language instruction sent on set_task / reset. Optional: policies without a language input leave it None.
set_task_payload_key	JSON key carrying the task description on set_task.
reset_payload_key	JSON key carrying the task description on reset.
transport	TRANSPORT_REGISTRY key selecting the wire protocol (http_json, http_json_numpy, http_multipart, openpi_websocket, galaxea_websocket, zmq_msgpack). Defaults to http_json.
serializer	SERIALIZER_REGISTRY key; None uses the transport's default.
chunk_mode	client (buffer the server's chunk and re-query when drained), server (the server caches the chunk and serves one row per call) or single (one action per call).
pred_action_horizon	Expected action-chunk length; None accepts any length.
exec_horizon	Rows executed open-loop before re-querying; None executes the whole chunk.
short_chunk	error raises when the server returns fewer than pred_action_horizon rows; accept queues what came back.
chunk_temporal_ensemble_steps	Temporal ensembling of the previous chunk's unused tail into the new head; zero keeps plain open-loop execution.
chunk_temporal_ensemble_dims	Restrict ensembling to the first N action dims (e.g. to spare a binary gripper).
response_action_key	Key holding the action chunk in a dict reply (action or actions).

Source code in eden/options/options.py

def __init__(self, **data):
    # Seed type metadata as construction data so BaseModel validates it into
    # __pydantic_extra__ (no-op for extra="forbid" classes, which use
    # the serialization envelope instead).
    data.update(option_metadata(self))
    # NOTE: all of the option modifications/validations should be done in the __init__ method
    try:
        BaseModel.__init__(self, **data)
    except ValidationError as exc:
        # A misspelled field on an extra="forbid" class should teach the field set, not just
        # name the typo. Only rewritten when every error is extra_forbidden.
        enriched = _enrich_extra_forbidden(type(self), exc)
        if enriched is None:
            raise
        raise enriched from None

Attributes

host class-attribute instance-attribute

host: str = 'localhost'

port class-attribute instance-attribute

port: int | None = 9010

request_timeout class-attribute instance-attribute

request_timeout: float = 30.0

set_task_endpoint class-attribute instance-attribute

set_task_endpoint: str | None = None

reset_endpoint class-attribute instance-attribute

reset_endpoint: str | None = None

act_endpoint class-attribute instance-attribute

act_endpoint: str = '/act'

task_description class-attribute instance-attribute

task_description: str | None = None

set_task_payload_key class-attribute instance-attribute

set_task_payload_key: str = 'task_description'

reset_payload_key class-attribute instance-attribute

reset_payload_key: str = 'task_description'

transport class-attribute instance-attribute

transport: str = 'http_json'

serializer class-attribute instance-attribute

serializer: str | None = None

chunk_mode class-attribute instance-attribute

chunk_mode: ChunkMode = 'client'

pred_action_horizon class-attribute instance-attribute

pred_action_horizon: int | None = None

exec_horizon class-attribute instance-attribute

exec_horizon: int | None = None

short_chunk class-attribute instance-attribute

short_chunk: ShortChunk = 'error'

chunk_temporal_ensemble_steps class-attribute instance-attribute

chunk_temporal_ensemble_steps: int = 0

chunk_temporal_ensemble_dims class-attribute instance-attribute

chunk_temporal_ensemble_dims: int | None = None

response_action_key class-attribute instance-attribute

response_action_key: str = 'action'

RemotePolicyOptions

RemotePolicyOptions(**data)

Bases: PolicyOptions, ClientOptions

Merged options for policies that combine ClientMixin and PolicyBase.

Use as the type argument when subclassing both: class FooPolicy( ClientMixin, PolicyBase, ConfigurableMixin[RemotePolicyOptions]):. Without this third base, __init_subclass__ would resolve _options_class_ to whichever parent appears first in MRO, hiding the other parent's fields.

Source code in eden/options/options.py

def __init__(self, **data):
    # Seed type metadata as construction data so BaseModel validates it into
    # __pydantic_extra__ (no-op for extra="forbid" classes, which use
    # the serialization envelope instead).
    data.update(option_metadata(self))
    # NOTE: all of the option modifications/validations should be done in the __init__ method
    try:
        BaseModel.__init__(self, **data)
    except ValidationError as exc:
        # A misspelled field on an extra="forbid" class should teach the field set, not just
        # name the typo. Only rewritten when every error is extra_forbidden.
        enriched = _enrich_extra_forbidden(type(self), exc)
        if enriched is None:
            raise
        raise enriched from None

ONNX runtime helpers

Shared ONNX-runtime infrastructure for policy extensions.

Provides the session layer (provider selection, thread policy, CPU fallback) and a generic torch.nn.Module wrapper so policies don't each re-implement session creation + inference. Single-input / single-output graphs use :class:OnnxModule directly (see the sonic / holosoma ports); multi-input or IO-binding graphs subclass :class:_OnnxWrapperBase (see omnixtreme).

Attributes

ort module-attribute

ort: Any = None

INTRA_OP_THREADS_ENV module-attribute

INTRA_OP_THREADS_ENV = 'ONNX_INTRA_OP_THREADS'

DEFAULT_INTRA_OP_THREADS module-attribute

DEFAULT_INTRA_OP_THREADS = 1

Classes

OnnxModule

OnnxModule(onnx_path: str, input_name: str, output_name: str, *, device: str | None = None, label: str = 'ONNX', expected_input_dim: int | None = None, expected_output_dim: int | None = None, enable_iobinding: bool = False, dynamic_batch: bool = False)

Bases: _OnnxWrapperBase

Generic single-input / single-output ONNX graph as a torch Module.

Covers the common policy shape: one named float32 input (e.g. actor_obs or obs_dict) and one named float32 output (e.g. action). forward accepts a tensor or array, runs the graph with automatic CPU fallback, and returns a float32 tensor on self.device.

PARAMETER	DESCRIPTION
onnx_path	Path to the ONNX file. TYPE: str
input_name	Names of the graph's single input / output tensor (validated against the loaded graph). TYPE: str
output_name	Names of the graph's single input / output tensor (validated against the loaded graph). TYPE: str
device	Device the returned tensor is placed on. TYPE: str | None DEFAULT: None
label	Human-readable tag used in log / fallback messages. TYPE: str DEFAULT: 'ONNX'
expected_input_dim	When set and the corresponding tensor has a static shape, the flattened dim is validated and a ValueError is raised on mismatch — a guard against a swapped / re-exported graph. TYPE: int | None DEFAULT: None
expected_output_dim	When set and the corresponding tensor has a static shape, the flattened dim is validated and a ValueError is raised on mismatch — a guard against a swapped / re-exported graph. TYPE: int | None DEFAULT: None
enable_iobinding	Use the GPU IO-binding fast path when the device is CUDA, a GPU provider is active and the output is float32 with static non-batch dims; otherwise the numpy path runs. TYPE: bool DEFAULT: False
dynamic_batch	Patch a static-batch-1 graph to a dynamic leading axis first (:func:dynamic_batch_onnx), so num_envs > 1 feeds are accepted. TYPE: bool DEFAULT: False

Source code in eden/extensions/policies/onnx.py

def __init__(
    self,
    onnx_path: str,
    input_name: str,
    output_name: str,
    *,
    device: str | None = None,
    label: str = "ONNX",
    expected_input_dim: int | None = None,
    expected_output_dim: int | None = None,
    enable_iobinding: bool = False,
    dynamic_batch: bool = False,
) -> None:
    super().__init__(onnx_path, device=device, label=label, dynamic_batch=dynamic_batch)
    if input_name not in self._input_names:
        raise ValueError(f"{label} at {onnx_path} has no input '{input_name}'; got {sorted(self._input_names)}")
    if output_name not in self._output_names:
        raise ValueError(f"{label} at {onnx_path} has no output '{output_name}'; got {sorted(self._output_names)}")
    self._input_name = input_name
    self._output_name = output_name
    # Per-row (feature) element count, used by the swapped-graph guard. Works
    # whether the leading axis is a literal 1 or the symbolic "batch" that
    # dynamic_batch leaves behind.
    self.input_dim = _nonbatch_numel(self.session, input_name)
    self.output_dim = _nonbatch_numel(self.session, output_name)
    self._validate_dim("input", self.input_dim, expected_input_dim)
    self._validate_dim("output", self.output_dim, expected_output_dim)

    # Per-row output shape (dims after the dynamic batch axis), used to
    # pre-allocate the GPU output for IO-binding. None if any non-batch dim
    # is dynamic — then IO-binding is disabled and we use the numpy path.
    out = next((o for o in self.session.get_outputs() if o.name == output_name), None)
    out_shape = out.shape if out is not None else None
    tail = list(out_shape[1:]) if out_shape else []
    self._out_tail = tuple(int(d) for d in tail) if tail and all(isinstance(d, int) for d in tail) else None
    # IO-binding pre-allocates a float32 output buffer, so it's only sound
    # when the graph output is float32 (matching the numpy path, which casts
    # to float32). A non-float32 output (e.g. float16 / int) falls back to the
    # numpy path rather than writing wrong-width bytes into a float32 buffer.
    out_is_float32 = out is not None and out.type == "tensor(float)"
    self._iobind = bool(enable_iobinding) and self._out_tail is not None and out_is_float32

Attributes

input_dim instance-attribute

input_dim = _nonbatch_numel(self.session, input_name)

output_dim instance-attribute

output_dim = _nonbatch_numel(self.session, output_name)

Methods:

forward

forward(obs: Tensor) -> Tensor

Source code in eden/extensions/policies/onnx.py

@torch.no_grad()
def forward(self, obs: torch.Tensor) -> torch.Tensor:
    if self._iobind and self._can_iobind():
        try:
            return self._forward_iobind(obs)
        except Exception as e:
            if _env_bool("ONNX_DEBUG"):
                warnings.warn(f"[{self._label}] IO-binding failed ({type(e).__name__}: {e}); using numpy path")
    feed = {self._input_name: _ensure_2d(to_numpy(obs, np.float32))}
    return self._run(feed, output_names=[self._output_name])

Functions:

dynamic_batch_onnx

dynamic_batch_onnx(src: str) -> str

Patch a static-batch-1 ONNX graph to accept an arbitrary leading batch.

Rewrites every Reshape shape target whose leading dim is a constant 1 to -1 and marks each graph input/output's dim 0 as the symbolic "batch". The patched graph is cached next to src as *.dynbatch.onnx (falling back to a temp dir if src's directory is read-only) and the cached path is returned; an existing cache is reused.

Needed because some released graphs (e.g. SONIC) are exported with a static batch of 1, which rejects num_envs > 1 feeds outright. The patch is a heuristic, so the freshly built graph is checked by :func:_verify_dynamic_batch (accepts a batched feed + permutation-equivariant rows) before it is used — a corrupt patch fails loudly here instead of silently returning wrong rows.

Source code in eden/extensions/policies/onnx.py

def dynamic_batch_onnx(src: str) -> str:
    """Patch a static-batch-1 ONNX graph to accept an arbitrary leading batch.

    Rewrites every ``Reshape`` shape target whose leading dim is a constant ``1``
    to ``-1`` and marks each graph input/output's dim 0 as the symbolic
    ``"batch"``. The patched graph is cached next to ``src`` as
    ``*.dynbatch.onnx`` (falling back to a temp dir if ``src``'s directory is
    read-only) and the cached path is returned; an existing cache is reused.

    Needed because some released graphs (e.g. SONIC) are exported with a static
    batch of 1, which rejects ``num_envs > 1`` feeds outright. The patch is a
    heuristic, so the freshly built graph is checked by :func:`_verify_dynamic_batch`
    (accepts a batched feed + permutation-equivariant rows) before it is used — a
    corrupt patch fails loudly here instead of silently returning wrong rows.
    """
    # ``splitext`` (not ``str.replace``) so a path with ``.onnx`` in a parent dir
    # isn't corrupted, and an extensionless path can't collapse ``dst`` onto
    # ``src`` and overwrite the verified original.
    dst = os.path.splitext(src)[0] + ".dynbatch.onnx"

    # Reuse the cache only if it is at least as new as the source: a re-downloaded
    # / updated checkpoint at the same path (newer mtime) forces a rebuild, so we
    # never run a stale patch derived from older weights.
    def _cache_fresh(path: str) -> bool:
        return os.path.exists(path) and os.path.getmtime(path) >= os.path.getmtime(src)

    if _cache_fresh(dst):
        return dst
    import onnx
    from onnx import numpy_helper

    model = onnx.load(src)
    # Names of the shape tensors feeding every Reshape (constants + initializers).
    shape_inputs = {n.input[1] for n in model.graph.node if n.op_type == "Reshape"}

    def _lead_one_to_minus_one(t) -> bool:
        arr = numpy_helper.to_array(t).copy()
        # Only a leading literal 1 with no other inferred (-1) dim: a Reshape
        # target may carry at most one -1, so never introduce a second one
        # (transformer attention reshapes already use -1 for the flattened tail).
        if arr.ndim == 1 and arr.size and arr[0] == 1 and -1 not in arr.tolist():
            arr[0] = -1
            t.CopyFrom(numpy_helper.from_array(arr, t.name))
            return True
        return False

    for node in model.graph.node:
        if node.op_type == "Constant" and node.output[0] in shape_inputs:
            for attr in node.attribute:
                if attr.name == "value":
                    _lead_one_to_minus_one(attr.t)
    for init in model.graph.initializer:
        if init.name in shape_inputs:
            _lead_one_to_minus_one(init)
    for vi in list(model.graph.input) + list(model.graph.output):
        vi.type.tensor_type.shape.dim[0].dim_param = "batch"

    try:
        onnx.save(model, dst)
    except OSError:
        import hashlib
        import tempfile

        cache = os.path.join(tempfile.gettempdir(), "eden_onnx_dynbatch")
        os.makedirs(cache, exist_ok=True)
        dst = os.path.join(cache, hashlib.md5(src.encode()).hexdigest() + ".dynbatch.onnx")
        if not _cache_fresh(dst):
            onnx.save(model, dst)
    # Just-built graph: confirm the patch is numerically equivalent to looped
    # batch-1 before any policy runs it. Cache hits above skip this (already
    # verified when first built).
    _verify_dynamic_batch(dst)
    return dst

Server-backed clients

A server-backed policy mixes in ClientMixin: it builds a request from the env's observations, sends it through a transport and serializer, and schedules the returned action chunk. The adapters turn a server's action format into the env's action terms.

Server-backed policies: ClientMixin and the per-embodiment setup machinery.

The counterpart of :mod:eden.extensions.policies.onnx (local actors): a client policy is class FooPolicy(ClientMixin, PolicyBase, ConfigurableMixin[RemotePolicyOptions]) whose inference happens on a server. The layer is policy-family neutral (VLA, world action model, remote RL controller, ...) and owns three things:

• a transport from :mod:eden.utils.policy_client, selected by key (transport="http_json", "openpi_websocket", "zmq_msgpack", ...),
• an :class:ActionChunkScheduler that turns server action chunks into one (num_envs, D) row per step,
• a uniform lifecycle: :meth:ClientMixin.set_task, :meth:ClientMixin.reset_server, :meth:ClientMixin.reset.

Every step runs the same pipeline::

obs -> observe(obs)                                   per-step hooks (frame history, measured pose)
    -> [if scheduler.needs_inference]
          payload = pack_request(obs)                 default: pack_observations(obs)
          reply   = transport.infer(payload, request_options())
          chunk   = decode_response(reply)            -> (num_envs, T, D)
          scheduler.accept(chunk)
    -> row = scheduler.pop()                          -> (num_envs, D)
    -> unpack_actions(row)                            -> {action_term: tensor}

step is also exposed as its stages observe (every step: frame histories, measured poses) / pack_request / infer_payload / accept_chunk / pop_action (plus :attr:ClientMixin.needs_inference). An evaluator can then run the blocking round-trip on a worker thread while it keeps stepping the simulation, so the robot moves on during the wait as much as the inference latency would let it. The evaluator must still call observe on every step. The contract is one request in flight: accept or discard its chunk before the next pack_request. Do not call reset / reset_server / set_task while a request is in flight, because nothing here locks or fences the transport.

Servers that serve several embodiments use :class:SetupClientMixin with a SETUP_REGISTRY of :class:PolicySetupBase strategies selected by policy_setup.

Attributes

Classes

ActionChunkScheduler

ActionChunkScheduler(*, num_envs: int, mode: ChunkMode = 'client', pred_action_horizon: int | None = None, exec_horizon: int | None = None, short_chunk: ShortChunk = 'error', ensemble_steps: int = 0, ensemble_dims: int | None = None)

Queue server action chunks and hand them out one row per env per step.

PARAMETER	DESCRIPTION
num_envs	Number of envs; chunks are (num_envs, T, D) and rows are (num_envs, D). TYPE: int
mode	client: this side owns the chunk and re-queries only when its queue drains. server: the server caches the chunk and serves one row per call, so every step is an inference call (GalaxeaVLA). single: the server returns one action per call (Octo). The latter two never buffer: a multi-row chunk raises. TYPE: ChunkMode DEFAULT: 'client'
pred_action_horizon	Expected chunk length T. None accepts any length. TYPE: int | None DEFAULT: None
exec_horizon	Rows executed before re-querying. None executes the whole chunk. Must not exceed pred_action_horizon when both are set. TYPE: int | None DEFAULT: None
short_chunk	What to do when the server returns fewer than pred_action_horizon rows: raise (the prediction is truncated relative to the trained horizon) or queue what came back. TYPE: ShortChunk DEFAULT: 'error'
ensemble_steps	Temporal ensembling of the previous chunk's unused tail into the new chunk's head: for K steps the old plan is weighted K/(K+1) … 1/(K+1). Zero disables it (bit-exact open-loop execution). Requires both horizons and ensemble_steps <= pred_action_horizon - exec_horizon. TYPE: int DEFAULT: 0
ensemble_dims	Only the first ensemble_dims action dims are ensembled (e.g. to keep a binary gripper channel untouched). None ensembles every dim. TYPE: int | None DEFAULT: None

Source code in eden/extensions/policies/client.py

def __init__(
    self,
    *,
    num_envs: int,
    mode: ChunkMode = "client",
    pred_action_horizon: int | None = None,
    exec_horizon: int | None = None,
    short_chunk: ShortChunk = "error",
    ensemble_steps: int = 0,
    ensemble_dims: int | None = None,
) -> None:
    if mode not in ("client", "server", "single"):
        raise ValueError(f"chunk_mode must be 'client', 'server' or 'single', got {mode!r}")
    if short_chunk not in ("error", "accept"):
        raise ValueError(f"short_chunk must be 'error' or 'accept', got {short_chunk!r}")
    if pred_action_horizon is not None and pred_action_horizon <= 0:
        raise ValueError(f"pred_action_horizon must be positive, got {pred_action_horizon}")
    if exec_horizon is not None and exec_horizon <= 0:
        raise ValueError(f"exec_horizon must be positive, got {exec_horizon}")
    if pred_action_horizon is not None and exec_horizon is not None and exec_horizon > pred_action_horizon:
        raise ValueError(f"exec_horizon ({exec_horizon}) must be <= pred_action_horizon ({pred_action_horizon})")
    if ensemble_steps < 0:
        raise ValueError("chunk_temporal_ensemble_steps must be non-negative")
    if ensemble_steps > 0:
        if pred_action_horizon is None or exec_horizon is None:
            raise ValueError("temporal ensembling requires both pred_action_horizon and exec_horizon")
        available_overlap = pred_action_horizon - exec_horizon
        if ensemble_steps > available_overlap:
            raise ValueError(
                f"chunk_temporal_ensemble_steps ({ensemble_steps}) exceeds the unused "
                f"prediction tail ({available_overlap})"
            )
    if ensemble_dims is not None and ensemble_dims <= 0:
        raise ValueError("chunk_temporal_ensemble_dims must be positive or None")

    self.num_envs = int(num_envs)
    self.mode = mode
    self.pred_action_horizon = pred_action_horizon
    self.exec_horizon = exec_horizon
    self.short_chunk = short_chunk
    self.ensemble_steps = int(ensemble_steps)
    self.ensemble_dims = ensemble_dims

    self.queues: list[collections.deque[np.ndarray]] = [collections.deque() for _ in range(self.num_envs)]
    self.raw_queues: list[collections.deque[np.ndarray]] = [collections.deque() for _ in range(self.num_envs)]
    self._previous_unused_tail: list[np.ndarray | None] = [None] * self.num_envs
    self._last_take = 0
    # Diagnostics
    self.chunk_index = -1
    self.chunk_action_index = -1
    self.last_prediction_chunk_length = 0
    self.last_temporal_ensemble_steps = 0
    self.last_raw_row: np.ndarray | None = None

Attributes

num_envs instance-attribute

num_envs = int(num_envs)

mode instance-attribute

mode = mode

pred_action_horizon instance-attribute

pred_action_horizon = pred_action_horizon

exec_horizon instance-attribute

exec_horizon = exec_horizon

short_chunk instance-attribute

short_chunk = short_chunk

ensemble_steps instance-attribute

ensemble_steps = int(ensemble_steps)

ensemble_dims instance-attribute

ensemble_dims = ensemble_dims

queues instance-attribute

queues: list[deque[ndarray]] = [collections.deque() for _ in range(self.num_envs)]

raw_queues instance-attribute

raw_queues: list[deque[ndarray]] = [collections.deque() for _ in range(self.num_envs)]

chunk_index instance-attribute

chunk_index = -1

chunk_action_index instance-attribute

chunk_action_index = -1

last_prediction_chunk_length instance-attribute

last_prediction_chunk_length = 0

last_temporal_ensemble_steps instance-attribute

last_temporal_ensemble_steps = 0

last_raw_row instance-attribute

last_raw_row: ndarray | None = None

needs_inference property

needs_inference: bool

True when the next :meth:pop needs a fresh server chunk first.

action_metadata property

action_metadata: dict[str, int]

Where the last popped row sat in its chunk (for recorders / evaluators).

Methods:

clear

clear() -> None

Drop buffered actions and ensembling state (task change / episode boundary).

Source code in eden/extensions/policies/client.py

def clear(self) -> None:
    """Drop buffered actions and ensembling state (task change / episode boundary)."""
    for q in self.queues:
        q.clear()
    for q in self.raw_queues:
        q.clear()
    self._previous_unused_tail = [None] * self.num_envs
    self._last_take = 0
    self.chunk_index = -1
    self.chunk_action_index = -1
    self.last_prediction_chunk_length = 0
    self.last_temporal_ensemble_steps = 0
    self.last_raw_row = None

accept

accept(chunk: ndarray, raw_chunk: ndarray | None = None) -> None

Queue the first exec_horizon rows of a (num_envs, T, D) chunk.

raw_chunk (same shape) is an optional untouched copy queued alongside for diagnostics, e.g. before a setup latched the gripper column.

Source code in eden/extensions/policies/client.py

def accept(self, chunk: np.ndarray, raw_chunk: np.ndarray | None = None) -> None:
    """Queue the first ``exec_horizon`` rows of a ``(num_envs, T, D)`` chunk.

    ``raw_chunk`` (same shape) is an optional untouched copy queued alongside for
    diagnostics, e.g. before a setup latched the gripper column.
    """
    chunk = np.asarray(chunk, dtype=np.float32)
    if chunk.ndim != 3 or chunk.shape[0] != self.num_envs:
        raise ValueError(f"Expected an action chunk of shape ({self.num_envs}, T, D), got {chunk.shape}")
    T = chunk.shape[1]
    if T == 0:
        raise ValueError("Policy server returned an empty action chunk.")
    if self.mode != "client" and T != 1:
        raise ValueError(
            f"chunk_mode={self.mode!r} expects one action per inference call, but the server returned "
            f"{T} rows. Use chunk_mode='client' if this server returns action chunks."
        )
    if raw_chunk is None:
        raw_chunk = chunk
    else:
        raw_chunk = np.asarray(raw_chunk, dtype=np.float32)
        if raw_chunk.shape != chunk.shape:
            raise ValueError(f"raw_chunk shape {raw_chunk.shape} does not match chunk shape {chunk.shape}")

    self.last_prediction_chunk_length = int(T)
    if self.pred_action_horizon is not None and T < self.pred_action_horizon and self.short_chunk == "error":
        raise RuntimeError(
            f"Policy server returned {T} actions but pred_action_horizon={self.pred_action_horizon}. "
            "The prediction chunk is truncated — check the server's action_horizon before "
            "open-loop execution."
        )
    take = T if self.exec_horizon is None else min(self.exec_horizon, T)
    if self.pred_action_horizon is not None:
        take = min(take, self.pred_action_horizon)

    self.last_temporal_ensemble_steps = 0
    for i in range(self.num_envs):
        env_chunk = chunk[i]
        previous_tail = self._previous_unused_tail[i]
        if self.ensemble_steps > 0 and previous_tail is not None:
            overlap = min(self.ensemble_steps, previous_tail.shape[0], take)
            dims = env_chunk.shape[1] if self.ensemble_dims is None else min(self.ensemble_dims, env_chunk.shape[1])
            if overlap > 0:
                env_chunk = env_chunk.copy()
                # For K steps, old-plan weights are K/(K+1), ..., 1/(K+1): the new
                # observation-conditioned prediction takes over completely after K rows.
                old_weights = np.arange(overlap, 0, -1, dtype=np.float32) / float(self.ensemble_steps + 1)
                env_chunk[:overlap, :dims] = (
                    old_weights[:, None] * previous_tail[:overlap, :dims]
                    + (1.0 - old_weights[:, None]) * env_chunk[:overlap, :dims]
                )
                self.last_temporal_ensemble_steps = overlap
        self._previous_unused_tail[i] = chunk[i][take:].copy()
        self.queues[i].extend(env_chunk[:take])
        self.raw_queues[i].extend(raw_chunk[i][:take])

    self._last_take = int(take)
    self.chunk_index += 1

pop

pop() -> ndarray

Return the next (num_envs, D) row; raises if nothing is queued.

Source code in eden/extensions/policies/client.py

def pop(self) -> np.ndarray:
    """Return the next ``(num_envs, D)`` row; raises if nothing is queued."""
    if not self.queues[0]:
        raise RuntimeError("No queued action: call accept() with a server chunk before pop().")
    self.chunk_action_index = self._last_take - len(self.queues[0])
    rows = np.stack([q.popleft() for q in self.queues], axis=0)
    self.last_raw_row = np.stack([q.popleft() for q in self.raw_queues], axis=0)
    return rows

ModalitySchema dataclass

ModalitySchema(video_keys: tuple[str, ...], state_keys: tuple[str, ...], action_keys: tuple[str, ...], language_key: str, action_horizon: int | None = None, video_horizon: int = 1, video_stride: int = 1)

What a setup expects the server's get_modality_config to say.

Isaac-GR00T and its derivatives select the checkpoint's embodiment server-side (--embodiment-tag) and advertise the resulting per-modality modality_keys and delta_indices. A setup declares the schema it packs and decodes; the two are compared once at connect time so a mismatch fails there, with both sides named, instead of as a KeyError on the server or a silently wrong video window.

Attributes

video_keys instance-attribute

video_keys: tuple[str, ...]

state_keys instance-attribute

state_keys: tuple[str, ...]

action_keys instance-attribute

action_keys: tuple[str, ...]

language_key instance-attribute

language_key: str

action_horizon class-attribute instance-attribute

action_horizon: int | None = None

video_horizon class-attribute instance-attribute

video_horizon: int = 1

video_stride class-attribute instance-attribute

video_stride: int = 1

Methods:

FrameHistory

FrameHistory(horizon: int, stride: int)

Per-step frame buffer that samples a strided temporal window the way delta_indices do.

Servers trained with a video window (GR00T delta_indices=[-15, 0], RLDX-1 [-6, -4, -2, 0]) expect horizon frames spaced stride control steps apart, ending at the current one. :meth:append records one frame per control step and :meth:window returns (horizon, ...) in chronological order; missing history is padded with the oldest available frame, as the reference MultiStepWrapper does by seeding its deque with copies of the first observation.

Built on :class:~eden.utils.buffers.circular_buffer.CircularBuffer (batch size 1, backfilled), so no per-key list bookkeeping is needed.

Source code in eden/extensions/policies/client.py

def __init__(self, horizon: int, stride: int) -> None:
    if horizon < 1:
        raise ValueError(f"horizon must be >= 1, got {horizon}")
    if stride < 1:
        raise ValueError(f"stride must be >= 1, got {stride}")
    self.horizon = int(horizon)
    self.stride = int(stride)
    self._max_len = (self.horizon - 1) * self.stride + 1
    self._buffer: CircularBuffer | None = None

Attributes

horizon instance-attribute

horizon = int(horizon)

stride instance-attribute

stride = int(stride)

Methods:

reset

reset() -> None

Source code in eden/extensions/policies/client.py

def reset(self) -> None:
    self._buffer = None

append

append(frame: ndarray | Tensor) -> None

Source code in eden/extensions/policies/client.py

def append(self, frame: np.ndarray | torch.Tensor) -> None:
    tensor = torch.as_tensor(frame)
    if self._buffer is None:
        self._buffer = CircularBuffer(self._max_len, batch_size=1, device=str(tensor.device))
    self._buffer.append(tensor.unsqueeze(0))

window

window() -> ndarray

(horizon, ...) frames at stride spacing, newest last.

Source code in eden/extensions/policies/client.py

def window(self) -> np.ndarray:
    """``(horizon, ...)`` frames at ``stride`` spacing, newest last."""
    if self._buffer is None:
        raise RuntimeError("No frames recorded: call append() once per control step before window().")
    history = self._buffer.buffer[0]  # (max_len, ...) oldest -> newest, backfilled
    offsets = [(self.horizon - 1 - i) * self.stride for i in range(self.horizon)]
    return np.stack([history[self._max_len - 1 - off].cpu().numpy() for off in offsets], axis=0)

ClientMixin

ClientMixin(env: EnvBase, options: PolicyOptions)

Bases: ConfigurableMixin[ClientOptions]

Mixin for policy extensions that delegate inference to a remote policy server.

Concrete subclasses implement :meth:pack_observations and :meth:unpack_actions (the serialization contract between Eden's observation/action format and the server's API) and pick a transport. Everything else — connection handling, chunk queuing, horizon rules, task/reset lifecycle — is shared. Override :meth:decode_response when the reply does not carry the chunk under a single key, :meth:observe for per-step state the policy needs even on steps that do not query the server, and :meth:request_options for per-request server options (sessions, memory flags).

Two requirements must both be met when combining with PolicyBase:

1. ClientMixin must appear before PolicyBase so that its concrete :meth:step takes precedence over the abstract one.
2. The subclass must re-parameterize ConfigurableMixin with the merged options type RemotePolicyOptions; otherwise __init_subclass__ inherits _options_class_ from whichever parent is first in MRO and :meth:configure rejects the other parent's fields.

PARAMETER	DESCRIPTION
host	Policy server address. host may be a full URL (wss://example.com/ws); use port=None to take the URL's port. See :func:eden.utils.policy_client.server_url.
port	Policy server address. host may be a full URL (wss://example.com/ws); use port=None to take the URL's port. See :func:eden.utils.policy_client.server_url.
request_timeout	Request timeout in seconds (HTTP transports).
set_task_endpoint	The server's set-task / reset API: an HTTP route, a ZMQ endpoint name or a GalaxeaVLA reset message key. None (the default) means the server has none, because it takes the task with each request, so :meth:set_task / :meth:reset_server only update local state. Declare one only when the official server exposes it.
reset_endpoint	The server's set-task / reset API: an HTTP route, a ZMQ endpoint name or a GalaxeaVLA reset message key. None (the default) means the server has none, because it takes the task with each request, so :meth:set_task / :meth:reset_server only update local state. Declare one only when the official server exposes it.
act_endpoint	HTTP route for inference.
task_description	Language instruction sent on set_task / reset; optional for policies with no language input.
set_task_payload_key	JSON keys carrying the task description.
reset_payload_key	JSON keys carrying the task description.
transport	TRANSPORT_REGISTRY key. Default "http_json".
serializer	SERIALIZER_REGISTRY key; None takes the transport's default.
chunk_mode	Who owns the action chunk: this client (buffer and re-query when drained), the server (it caches the chunk and serves one row per call), or nobody (one action per call).
pred_action_horizon	Expected chunk length from the server; None accepts any length.
exec_horizon	Rows executed open-loop before re-querying; None executes the whole chunk.
short_chunk	Whether a chunk shorter than pred_action_horizon is an error.
chunk_temporal_ensemble_steps	Temporal ensembling of overlapping chunks; zero disables.
chunk_temporal_ensemble_dims	Temporal ensembling of overlapping chunks; zero disables.
response_action_key	Key holding the action chunk in a dict reply ("action" or "actions").

Source code in eden/extensions/policies/client.py

def __init__(self, env: EnvBase, options: PolicyOptions):
    super().__init__(env, options)
    if not self.supports_batch and env.num_envs != 1:
        raise ValueError(
            f"{type(self).__name__} has supports_batch=False and only supports "
            f"num_envs=1 (got {env.num_envs}). Set supports_batch=True on the "
            f"subclass if the remote server accepts batched requests."
        )
    if env.num_envs > 1 and type(self).reset is ClientMixin.reset:
        # The env resets its envs one by one; the default reset can only reset the shared
        # server state of all of them, so it would raise at the first episode end.
        raise ValueError(
            f"{type(self).__name__} serves {env.num_envs} envs but keeps ClientMixin.reset, which "
            "refuses a partial reset because the server's episode state is shared by every env. Fix: "
            "override reset(envs_idx) with a per-env scheme (e.g. one server session per env), or "
            "run with num_envs=1."
        )
    self._num_envs: int = env.num_envs
    self._post_configure()
    self._scheduler = ActionChunkScheduler(
        num_envs=env.num_envs,
        mode=self.chunk_mode,
        pred_action_horizon=self.pred_action_horizon,
        exec_horizon=self.exec_horizon,
        short_chunk=self.short_chunk,
        ensemble_steps=self.chunk_temporal_ensemble_steps,
        ensemble_dims=self.chunk_temporal_ensemble_dims,
    )
    self.last_response_info: dict[str, Any] = {}
    self._transport: TransportBase = self.make_transport()
    self._post_connect()

Attributes

reads_env_observations class-attribute instance-attribute

reads_env_observations = True

host class-attribute instance-attribute

host: str = 'localhost'

port class-attribute instance-attribute

port: int | None = 9010

request_timeout class-attribute instance-attribute

request_timeout: float = 30.0

set_task_endpoint class-attribute instance-attribute

set_task_endpoint: str | None = None

reset_endpoint class-attribute instance-attribute

reset_endpoint: str | None = None

act_endpoint class-attribute instance-attribute

act_endpoint: str = '/act'

task_description class-attribute instance-attribute

task_description: str | None = None

set_task_payload_key class-attribute instance-attribute

set_task_payload_key: str = 'task_description'

reset_payload_key class-attribute instance-attribute

reset_payload_key: str = 'task_description'

transport class-attribute instance-attribute

transport: str = 'http_json'

serializer class-attribute instance-attribute

serializer: str | None = None

chunk_mode class-attribute instance-attribute

chunk_mode: ChunkMode = 'client'

pred_action_horizon class-attribute instance-attribute

pred_action_horizon: int | None = None

exec_horizon class-attribute instance-attribute

exec_horizon: int | None = None

short_chunk class-attribute instance-attribute

short_chunk: ShortChunk = 'error'

chunk_temporal_ensemble_steps class-attribute instance-attribute

chunk_temporal_ensemble_steps: int = 0

chunk_temporal_ensemble_dims class-attribute instance-attribute

chunk_temporal_ensemble_dims: int | None = None

response_action_key class-attribute instance-attribute

response_action_key: str = 'action'

supports_batch class-attribute

supports_batch: bool = False

action_dim class-attribute

action_dim: int | None = None

last_response_info instance-attribute

last_response_info: dict[str, Any] = {}

last_prediction_chunk_length property

last_prediction_chunk_length: int

Length T of the most recent server chunk.

last_temporal_ensemble_steps property

last_temporal_ensemble_steps: int

How many rows of the most recent chunk were temporally ensembled.

needs_inference property

needs_inference: bool

True when the next :meth:step will query the server.

action_metadata property

action_metadata: dict[str, int]

Chunk position of the last executed action.

Methods:

transport_config

transport_config() -> TransportConfig

Collect the connection settings and the selected transport's knobs from this policy.

Source code in eden/extensions/policies/client.py

def transport_config(self) -> TransportConfig:
    """Collect the connection settings and the selected transport's knobs from this policy."""
    if self.transport not in TRANSPORT_REGISTRY.keys():
        raise ValueError(
            f"{type(self).__name__}.transport={self.transport!r} is not a registered transport. "
            f"Available: {sorted(TRANSPORT_REGISTRY.keys())}"
        )
    transport_cls = TRANSPORT_REGISTRY.get(self.transport)
    knobs = {name: getattr(self, name) for name in transport_cls.knobs if hasattr(self, name)}
    return TransportConfig(
        host=self.host,
        port=self.port,
        request_timeout=self.request_timeout,
        act_endpoint=self.act_endpoint,
        set_task_endpoint=self.set_task_endpoint,
        reset_endpoint=self.reset_endpoint,
        set_task_payload_key=self.set_task_payload_key,
        reset_payload_key=self.reset_payload_key,
        task_description=self.task_description,
        serializer=self.serializer,
        knobs=knobs,
    )

make_transport

make_transport() -> TransportBase

Build the transport named by self.transport; override to inject a custom one.

Source code in eden/extensions/policies/client.py

def make_transport(self) -> TransportBase:
    """Build the transport named by ``self.transport``; override to inject a custom one."""
    return make_transport(self.transport, self.transport_config())

set_task

set_task(task_description: str) -> None

Set the task description locally and on the server, clearing pending actions.

Should be called after construction and before the first :meth:step.

Source code in eden/extensions/policies/client.py

def set_task(self, task_description: str) -> None:
    """Set the task description locally and on the server, clearing pending actions.

    Should be called after construction and before the first :meth:`step`.
    """
    self.task_description = task_description
    self._scheduler.clear()
    self._reset_local()
    self._transport.set_task(task_description)

reset_server

reset_server() -> None

Reset per-episode state locally, and on the server when it has a reset API (reset_endpoint).

Source code in eden/extensions/policies/client.py

def reset_server(self) -> None:
    """Reset per-episode state locally, and on the server when it has a reset API (``reset_endpoint``)."""
    self._scheduler.clear()
    self._reset_local()
    # The task text may have been assigned directly; reset routes that carry it must send the current one.
    self._transport.config.task_description = self.task_description
    self._transport.reset(self.reset_options())

close

close() -> None

Release the connection to the policy server (WebSocket / ZMQ socket, HTTP session).

Source code in eden/extensions/policies/client.py

def close(self) -> None:
    """Release the connection to the policy server (WebSocket / ZMQ socket, HTTP session)."""
    self._transport.close()

reset

reset(envs_idx: slice | Tensor | None = None) -> None

PolicyBase hook: reset the server once every env this client serves is being reset.

The server's episode state (history, chunk cache, recurrent memory) belongs to this client's connection / session as a whole, so a reset that covers every env maps to :meth:reset_server; with num_envs == 1 that is every non-empty reset. A partial reset of a batched client would wipe the other envs' state, so it raises until the client implements a per-env scheme (e.g. one server session per env).

Source code in eden/extensions/policies/client.py

def reset(self, envs_idx: (slice | torch.Tensor) | None = None) -> None:
    """``PolicyBase`` hook: reset the server once every env this client serves is being reset.

    The server's episode state (history, chunk cache, recurrent memory) belongs to this
    client's connection / session as a whole, so a reset that covers every env maps to
    :meth:`reset_server`; with ``num_envs == 1`` that is every non-empty reset. A partial
    reset of a batched client would wipe the other envs' state, so it raises until the
    client implements a per-env scheme (e.g. one server session per env).
    """
    covered = self._reset_coverage(envs_idx)
    if covered == 0:
        return
    if covered < self.env.num_envs:
        raise NotImplementedError(
            f"{type(self).__name__} cannot reset {covered} of {self.env.num_envs} envs: its server-side "
            "episode state is shared by all of them. Implement a per-env reset (e.g. one session per env)."
        )
    self.reset_server()

reset_options

reset_options() -> dict | None

Options sent with the server reset (e.g. session ids); None by default.

Source code in eden/extensions/policies/client.py

def reset_options(self) -> dict | None:
    """Options sent with the server ``reset`` (e.g. session ids); ``None`` by default."""
    return None

request_options

request_options() -> dict | None

Options sent with every inference request (e.g. memory flags); None by default.

Source code in eden/extensions/policies/client.py

def request_options(self) -> dict | None:
    """Options sent with every inference request (e.g. memory flags); ``None`` by default."""
    return None

observe

observe(obs: dict[str, Tensor | dict[str, Tensor]]) -> None

Per-step hook called before scheduling, on every step (default: no-op).

Use it for state that must track every control step regardless of when the server is queried: rolling frame histories, the measured pose a delta is applied to, ...

Source code in eden/extensions/policies/client.py

def observe(self, obs: dict[str, torch.Tensor | dict[str, torch.Tensor]]) -> None:
    """Per-step hook called before scheduling, on every step (default: no-op).

    Use it for state that must track every control step regardless of when the server
    is queried: rolling frame histories, the measured pose a delta is applied to, ...
    """

pack_request

pack_request(obs: dict[str, Tensor | dict[str, Tensor]]) -> Any

Build the server payload from an observation (main thread; reads sim state).

Source code in eden/extensions/policies/client.py

def pack_request(self, obs: dict[str, torch.Tensor | dict[str, torch.Tensor]]) -> Any:
    """Build the server payload from an observation (main thread; reads sim state)."""
    return self.pack_observations(obs)

infer_payload

infer_payload(payload: Any) -> ndarray

Blocking server round-trip -> (num_envs, T, D) chunk.

Safe on a worker thread while the main thread keeps stepping the simulation, as long as it is the only call using the transport, i.e. no reset / reset_server / set_task until it returns.

Source code in eden/extensions/policies/client.py

def infer_payload(self, payload: Any) -> np.ndarray:
    """Blocking server round-trip -> ``(num_envs, T, D)`` chunk.

    Safe on a worker thread while the main thread keeps stepping the simulation, as long
    as it is the only call using the transport, i.e. no ``reset`` / ``reset_server`` /
    ``set_task`` until it returns.
    """
    response = self._transport.infer(payload, self.request_options())
    self._record_response_info(response)
    return self.decode_response(response)

accept_chunk

accept_chunk(chunk: ndarray) -> None

Validate a predicted chunk and queue its executable prefix.

Source code in eden/extensions/policies/client.py

def accept_chunk(self, chunk: np.ndarray) -> None:
    """Validate a predicted chunk and queue its executable prefix."""
    chunk = self._normalize_action_shape(np.asarray(chunk, dtype=np.float32))
    self._check_action_width(chunk, self.action_dim, schema=type(self).__name__)
    self._scheduler.accept(chunk)

pop_action

pop_action() -> dict[str, Tensor]

Unpack and return the next queued (num_envs, D) row (main thread).

Source code in eden/extensions/policies/client.py

def pop_action(self) -> dict[str, torch.Tensor]:
    """Unpack and return the next queued ``(num_envs, D)`` row (main thread)."""
    return self.unpack_actions(self._scheduler.pop())

step

step(obs: dict[str, Tensor | dict[str, Tensor]]) -> dict[str, Tensor]

Query the policy server when needed and return the next action in Eden format.

Source code in eden/extensions/policies/client.py

def step(self, obs: dict[str, torch.Tensor | dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]:
    """Query the policy server when needed and return the next action in Eden format."""
    self.observe(obs)
    if self.needs_inference:
        self.accept_chunk(self.infer_payload(self.pack_request(obs)))
    return self.pop_action()

pack_observations abstractmethod

pack_observations(obs: dict[str, Tensor | dict[str, Tensor]]) -> Any

Convert Eden observations to the server payload.

When supports_batch=True, observation fields in the returned payload must carry a leading batch axis of size env.num_envs.

Source code in eden/extensions/policies/client.py

@abstractmethod
def pack_observations(self, obs: dict[str, torch.Tensor | dict[str, torch.Tensor]]) -> Any:
    """Convert Eden observations to the server payload.

    When ``supports_batch=True``, observation fields in the returned payload must
    carry a leading batch axis of size ``env.num_envs``.
    """
    raise NotImplementedError

decode_response

decode_response(response: Any) -> ndarray

Extract the action chunk from a server reply as (num_envs, T, D).

Default: response[response_action_key] normalized by :meth:_normalize_action_shape. Override for replies that spread the action over several keys.

Source code in eden/extensions/policies/client.py

def decode_response(self, response: Any) -> np.ndarray:
    """Extract the action chunk from a server reply as ``(num_envs, T, D)``.

    Default: ``response[response_action_key]`` normalized by
    :meth:`_normalize_action_shape`. Override for replies that spread the action over
    several keys.
    """
    if not isinstance(response, dict) or self.response_action_key not in response:
        raise KeyError(
            f"{type(self).__name__}: server reply has no {self.response_action_key!r} key "
            f"(got {sorted(response) if isinstance(response, dict) else type(response).__name__}). "
            "Set response_action_key or override decode_response()."
        )
    return self._normalize_action_shape(np.asarray(response[self.response_action_key], dtype=np.float32))

unpack_actions abstractmethod

unpack_actions(raw_action: ndarray) -> dict[str, Tensor]

Convert a (num_envs, D) action row from the server to an Eden action dict.

For non-batched subclasses this is always (1, D).

Source code in eden/extensions/policies/client.py

@abstractmethod
def unpack_actions(self, raw_action: np.ndarray) -> dict[str, torch.Tensor]:
    """Convert a ``(num_envs, D)`` action row from the server to an Eden action dict.

    For non-batched subclasses this is always ``(1, D)``.
    """
    raise NotImplementedError

PolicySetupBase

PolicySetupBase(policy: SetupClientMixin)

Bases: ABC

Strategy object for one server embodiment / schema.

Class-level defaults (pred_action_horizon, exec_horizon, default_image_size) fill the policy's pred_action_horizon / exec_horizon / image_size when the policy leaves them None, so each setup can carry the horizon and resolution its checkpoint was trained with.

Concrete setups own observation_options / action_options (classmethod factories so the env can be built without a live policy), pack_observations, and unpack_actions; they optionally override decode_response (reply -> (T, D) chunk), observe (per-step state) and reset.

Source code in eden/extensions/policies/client.py

def __init__(self, policy: SetupClientMixin) -> None:
    self.policy = policy

Attributes

pred_action_horizon class-attribute

pred_action_horizon: int | None = None

exec_horizon class-attribute

exec_horizon: int | None = None

default_image_size class-attribute

default_image_size: tuple[int, int] | None = None

action_dim class-attribute

action_dim: int | None = None

adapter class-attribute instance-attribute

adapter: ActionAdapter | None = None

last_raw_action_chunk class-attribute instance-attribute

last_raw_action_chunk: ndarray | None = None

policy instance-attribute

policy = policy

Methods:

reset

reset() -> None

Clear per-episode state (adapter latches, frame history, cached indices).

Source code in eden/extensions/policies/client.py

def reset(self) -> None:
    """Clear per-episode state (adapter latches, frame history, cached indices)."""
    self.last_raw_action_chunk = None
    if self.adapter is not None:
        self.adapter.reset()

observation_options abstractmethod classmethod

observation_options(**kwargs) -> dict[str, ObservationGroupOptions]

Observation group options to register with the Eden env.

Source code in eden/extensions/policies/client.py

@classmethod
@abstractmethod
def observation_options(cls, **kwargs) -> dict[str, ObservationGroupOptions]:
    """Observation group options to register with the Eden env."""

action_options abstractmethod classmethod

action_options(**kwargs) -> dict[str, ActionTermOptions]

Action term options to register with the Eden env.

Source code in eden/extensions/policies/client.py

@classmethod
@abstractmethod
def action_options(cls, **kwargs) -> dict[str, ActionTermOptions]:
    """Action term options to register with the Eden env."""

observe

observe(obs: dict) -> None

Called on every step before scheduling (default: nothing).

Source code in eden/extensions/policies/client.py

def observe(self, obs: dict) -> None:
    """Called on every step before scheduling (default: nothing)."""

validate_server_config

validate_server_config(config: dict) -> None

Check the server's advertised configuration against this setup's schema.

config is transport.server_config() — the handshake metadata (WebSocket) or the GR00T-family get_modality_config reply (ZMQ); {} when the server advertises nothing, in which case there is nothing to check. Raise ValueError naming the expected and served values on a mismatch.

Source code in eden/extensions/policies/client.py

def validate_server_config(self, config: dict) -> None:
    """Check the server's advertised configuration against this setup's schema.

    ``config`` is ``transport.server_config()`` — the handshake metadata (WebSocket) or
    the GR00T-family ``get_modality_config`` reply (ZMQ); ``{}`` when the server
    advertises nothing, in which case there is nothing to check. Raise ``ValueError``
    naming the expected and served values on a mismatch.
    """

pack_observations abstractmethod

pack_observations(obs: dict) -> Any

Convert Eden observation tensors to the server payload.

Source code in eden/extensions/policies/client.py

@abstractmethod
def pack_observations(self, obs: dict) -> Any:
    """Convert Eden observation tensors to the server payload."""

decode_response

decode_response(response: Any) -> ndarray | None

Reply -> (T, D) chunk, or None to use the policy's default key lookup.

Source code in eden/extensions/policies/client.py

def decode_response(self, response: Any) -> np.ndarray | None:
    """Reply -> ``(T, D)`` chunk, or ``None`` to use the policy's default key lookup."""
    return None

set_raw_server_action

set_raw_server_action(action: ndarray) -> None

Receive the untouched server row matching the row about to be unpacked (default: ignored).

Source code in eden/extensions/policies/client.py

def set_raw_server_action(self, action: np.ndarray) -> None:
    """Receive the untouched server row matching the row about to be unpacked (default: ignored)."""

unpack_actions

unpack_actions(raw_action: ndarray) -> dict[str, Tensor]

One (D,) action row -> {action_term: tensor}; delegates to adapter by default.

Source code in eden/extensions/policies/client.py

def unpack_actions(self, raw_action: np.ndarray) -> dict[str, torch.Tensor]:
    """One ``(D,)`` action row -> ``{action_term: tensor}``; delegates to ``adapter`` by default."""
    if self.adapter is None:
        raise NotImplementedError(f"{type(self).__name__} has no adapter set and does not override unpack_actions")
    return self.adapter(raw_action, env=self.policy.env)

SetupClientMixin

SetupClientMixin(env: EnvBase, options: PolicyOptions)

Bases: ClientMixin

ClientMixin whose serialization contract is delegated to a policy_setup.

Subclasses set SETUP_REGISTRY (name -> :class:PolicySetupBase subclass) and the default policy_setup. observation_options / action_options take policy_setup= and forward the remaining kwargs to the setup.

PARAMETER	DESCRIPTION
policy_setup	Key in SETUP_REGISTRY selecting the embodiment.
image_size	Image size sent to the server; None takes the setup's training default.

Source code in eden/extensions/policies/client.py

def __init__(self, env: EnvBase, options: PolicyOptions):
    super().__init__(env, options)
    if not self.supports_batch and env.num_envs != 1:
        raise ValueError(
            f"{type(self).__name__} has supports_batch=False and only supports "
            f"num_envs=1 (got {env.num_envs}). Set supports_batch=True on the "
            f"subclass if the remote server accepts batched requests."
        )
    if env.num_envs > 1 and type(self).reset is ClientMixin.reset:
        # The env resets its envs one by one; the default reset can only reset the shared
        # server state of all of them, so it would raise at the first episode end.
        raise ValueError(
            f"{type(self).__name__} serves {env.num_envs} envs but keeps ClientMixin.reset, which "
            "refuses a partial reset because the server's episode state is shared by every env. Fix: "
            "override reset(envs_idx) with a per-env scheme (e.g. one server session per env), or "
            "run with num_envs=1."
        )
    self._num_envs: int = env.num_envs
    self._post_configure()
    self._scheduler = ActionChunkScheduler(
        num_envs=env.num_envs,
        mode=self.chunk_mode,
        pred_action_horizon=self.pred_action_horizon,
        exec_horizon=self.exec_horizon,
        short_chunk=self.short_chunk,
        ensemble_steps=self.chunk_temporal_ensemble_steps,
        ensemble_dims=self.chunk_temporal_ensemble_dims,
    )
    self.last_response_info: dict[str, Any] = {}
    self._transport: TransportBase = self.make_transport()
    self._post_connect()

Attributes

SETUP_REGISTRY class-attribute

SETUP_REGISTRY: dict[str, type[PolicySetupBase]] = {}

policy_setup class-attribute instance-attribute

policy_setup: str = ''

image_size class-attribute instance-attribute

image_size: tuple[int, int] | None = None

Methods:

observation_options classmethod

observation_options(policy_setup: str | None = None, **kwargs) -> dict[str, ObservationGroupOptions]

Observation group options for policy_setup (default: the class default).

Source code in eden/extensions/policies/client.py

@classmethod
def observation_options(cls, policy_setup: str | None = None, **kwargs) -> dict[str, ObservationGroupOptions]:
    """Observation group options for ``policy_setup`` (default: the class default)."""
    return cls._setup_class(policy_setup or cls.policy_setup).observation_options(**kwargs)

action_options classmethod

action_options(policy_setup: str | None = None, **kwargs) -> dict[str, ActionTermOptions]

Action term options for policy_setup (default: the class default).

Source code in eden/extensions/policies/client.py

@classmethod
def action_options(cls, policy_setup: str | None = None, **kwargs) -> dict[str, ActionTermOptions]:
    """Action term options for ``policy_setup`` (default: the class default)."""
    return cls._setup_class(policy_setup or cls.policy_setup).action_options(**kwargs)

observe

observe(obs: dict[str, Tensor | dict[str, Tensor]]) -> None

Source code in eden/extensions/policies/client.py

def observe(self, obs: dict[str, torch.Tensor | dict[str, torch.Tensor]]) -> None:
    self._setup.observe(obs)

pack_observations

pack_observations(obs: dict[str, Tensor | dict[str, Tensor]]) -> Any

Source code in eden/extensions/policies/client.py

def pack_observations(self, obs: dict[str, torch.Tensor | dict[str, torch.Tensor]]) -> Any:
    return self._setup.pack_observations(obs)

decode_response

decode_response(response: Any) -> ndarray

Source code in eden/extensions/policies/client.py

def decode_response(self, response: Any) -> np.ndarray:
    chunk = self._setup.decode_response(response)
    if chunk is None:
        return super().decode_response(response)
    return self._normalize_action_shape(np.asarray(chunk, dtype=np.float32))

accept_chunk

accept_chunk(chunk: ndarray) -> None

Source code in eden/extensions/policies/client.py

def accept_chunk(self, chunk: np.ndarray) -> None:
    chunk = self._normalize_action_shape(np.asarray(chunk, dtype=np.float32))
    self._check_action_width(
        chunk, type(self._setup).action_dim, schema=f"{type(self).__name__} setup {self.policy_setup!r}"
    )
    raw, self._setup.last_raw_action_chunk = self._setup.last_raw_action_chunk, None
    raw_chunk = None if raw is None else self._normalize_action_shape(np.asarray(raw, dtype=np.float32))
    self._scheduler.accept(chunk, raw_chunk=raw_chunk)

pop_action

pop_action() -> dict[str, Tensor]

Source code in eden/extensions/policies/client.py

def pop_action(self) -> dict[str, torch.Tensor]:
    row = self._scheduler.pop()
    self._setup.set_raw_server_action(self._scheduler.last_raw_row[0])
    return self.unpack_actions(row)

unpack_actions

unpack_actions(raw_action: ndarray) -> dict[str, Tensor]

Source code in eden/extensions/policies/client.py

def unpack_actions(self, raw_action: np.ndarray) -> dict[str, torch.Tensor]:
    raw_action = np.asarray(raw_action)
    if raw_action.ndim == 2:
        if raw_action.shape[0] != 1:
            raise ValueError(
                f"{type(self).__name__} setups decode one env at a time; got a batched row {raw_action.shape}."
            )
        raw_action = raw_action[0]
    return self._setup.unpack_actions(raw_action)

Functions:

concat_action_streams

concat_action_streams(action_dict: dict[str, ndarray], keys: tuple[str, ...]) -> ndarray

Concatenate per-key (1, T, d_k) action streams in keys order into one (T, D) chunk.

Servers of the Isaac-GR00T family answer with one array per modality key. A (T, d) array (batch already squeezed) is accepted; a batch larger than one is rejected because the client serves a single env.

Source code in eden/extensions/policies/client.py

def concat_action_streams(action_dict: dict[str, np.ndarray], keys: tuple[str, ...]) -> np.ndarray:
    """Concatenate per-key ``(1, T, d_k)`` action streams in ``keys`` order into one ``(T, D)`` chunk.

    Servers of the Isaac-GR00T family answer with one array per modality key. A ``(T, d)``
    array (batch already squeezed) is accepted; a batch larger than one is rejected because
    the client serves a single env.
    """
    cols = []
    for key in keys:
        if key not in action_dict:
            raise KeyError(f"Server reply lacks action key {key!r}; got {sorted(action_dict)}")
        arr = np.asarray(action_dict[key], dtype=np.float32)
        if arr.ndim == 2:  # (T, d) — server already squeezed the batch axis
            arr = arr[np.newaxis]
        if arr.ndim != 3:
            raise ValueError(f"Action key {key!r} must be (B, T, d); got shape {arr.shape}")
        if arr.shape[0] != 1:
            raise ValueError(f"Remote clients serve num_envs=1; action key {key!r} has batch={arr.shape[0]}")
        cols.append(arr[0])  # (T, d)
    return np.concatenate(cols, axis=-1).astype(np.float32)

validate_modality_config

validate_modality_config(config: dict, schema: ModalitySchema, *, setup_name: str) -> None

Raise ValueError when config (a get_modality_config reply) contradicts schema.

An empty config means the server advertises nothing and passes. Keys are compared as sets for video / state (payloads are keyed dicts) and as an ordered list for action (the client concatenates streams in that order); horizons come from delta_indices.

Source code in eden/extensions/policies/client.py

def validate_modality_config(config: dict, schema: ModalitySchema, *, setup_name: str) -> None:
    """Raise ``ValueError`` when ``config`` (a ``get_modality_config`` reply) contradicts ``schema``.

    An empty ``config`` means the server advertises nothing and passes. Keys are compared
    as sets for video / state (payloads are keyed dicts) and as an ordered list for action
    (the client concatenates streams in that order); horizons come from ``delta_indices``.
    """
    if not config:
        return
    problems: list[str] = []

    def modality(name: str) -> dict | None:
        entry = config.get(name)
        if not isinstance(entry, dict) or "modality_keys" not in entry:
            problems.append(f"server config has no {name!r} modality with modality_keys (got {sorted(config)})")
            return None
        return entry

    video, state, action, language = (modality(k) for k in ("video", "state", "action", "language"))
    if video is not None:
        served = tuple(video["modality_keys"])
        if set(served) != set(schema.video_keys):
            problems.append(f"video keys: setup packs {sorted(schema.video_keys)}, server serves {sorted(served)}")
        delta = list(video.get("delta_indices", []))
        if delta and len(delta) != schema.video_horizon:
            problems.append(
                f"video horizon: setup sends {schema.video_horizon} frame(s), server expects {len(delta)} ({delta})"
            )
        elif len(delta) > 1:
            strides = {b - a for a, b in zip(delta[:-1], delta[1:])}
            if strides != {schema.video_stride}:
                problems.append(
                    f"video stride: setup samples every {schema.video_stride} step(s), server expects {delta}"
                )
    if state is not None:
        served = tuple(state["modality_keys"])
        if set(served) != set(schema.state_keys):
            problems.append(f"state keys: setup packs {sorted(schema.state_keys)}, server serves {sorted(served)}")
    if action is not None:
        served = tuple(action["modality_keys"])
        if served != tuple(schema.action_keys):
            problems.append(
                f"action keys (order matters): setup decodes {list(schema.action_keys)}, server serves {list(served)}"
            )
        delta = list(action.get("delta_indices", []))
        if schema.action_horizon is not None and delta and len(delta) != schema.action_horizon:
            problems.append(f"action horizon: setup expects {schema.action_horizon} rows, server predicts {len(delta)}")
    if language is not None and schema.language_key not in tuple(language["modality_keys"]):
        problems.append(
            f"language key: setup sends {schema.language_key!r}, server accepts {list(language['modality_keys'])}"
        )

    if problems:
        raise ValueError(
            f"Setup {setup_name!r} does not match the served checkpoint's modality config:\n  - "
            + "\n  - ".join(problems)
        )

Action adapters

Manipulation action adapters for remote policy clients.

An adapter converts a raw action array returned by a policy server into the {controller_name: torch.Tensor} dict that Eden's action manager expects. It is the mirror image of the server-specific obs packer: the packer serializes Eden's world state for the server; the adapter deserializes the server's action chunk back into Eden's controller contract.

The adapters here are transport-agnostic on purpose — they are composed into per-setup strategies (OpenPI pi/, GR00T, RLDX-1, Cosmos3, GalaxeaVLA) but nothing about them is wire-specific, so any client policy that receives a numpy action vector can reuse them. They are the manipulation family's decoders; other families add their own module beside this one. The module also carries the manipulation-side composition helpers shared by the setups (image observation groups, the WidowX EE-delta action pair, robot DOF layouts read from the robot option classes) and the construction checks every client runs against the env it is given (check_*: the control rate, the gripper decision term, the plain proprio group, joint terms that pass commands through, the DROID arm term and joint state, LIBERO's camera resolution).

Four adapter shapes cover the setups that use them:

• :class:EEDeltaAdapter — 7-D [ee_delta(6), gripper(1)] for the WidowX/Bridge and Google-Robot setups (OpenPI, GR00T, Octo). Euler or axis-angle rotation deltas; both reach the IK controller as a rotation vector.
• :class:JointVelocityAdapter — (n_arm + 1)-D joint velocity + gripper for OpenPI DROID. Emits true velocity commands by default, with an optional position-integration fallback for arms without a usable velocity controller.
• :class:JointPositionAdapter — (n_arm + 1)-D absolute joint position targets + gripper for DROID joint-position checkpoints (OpenPI, Cosmos3, GalaxeaVLA).
• :class:BimanualJointPositionAdapter — 14-D joint positions split into two arm controllers for Aloha.

None of them decides the gripper: each forwards the server's gripper scalar to the env's GripperDecisionController, which applies the setup's rule (threshold, hysteresis, sticky latch). See :func:check_env_gripper_term.

Attributes

WIDOWX_DOFS module-attribute

WIDOWX_DOFS: tuple[str, ...] = tuple(WidowX250.arm_dofs_name)

DROID_ARM_DOFS module-attribute

DROID_ARM_DOFS: tuple[str, ...] = tuple(FrankaResearch3.arm_dofs_name)

DROID_CONTROL_HZ module-attribute

DROID_CONTROL_HZ: float = 15.0

DROID_CONTROL_TIMING module-attribute

DROID_CONTROL_TIMING: tuple[str, int] = ('1/120', 8)

DROID_CONTROL_TIMING_NOTE module-attribute

DROID_CONTROL_TIMING_NOTE: str = " as PolaRiS does, with sim_substeps=9 to keep the Robotiq 2F85's substep under 1 ms (Eden's PolaRiS tasks step every 20 ms)"

ROBOTIQ_CLOSE_POS module-attribute

ROBOTIQ_CLOSE_POS: float = Robotiq2f85.close_dofs_pos['left_driver_joint']

DROID_FLANGE_FROM_FR3_ATTACHMENT_EULER module-attribute

DROID_FLANGE_FROM_FR3_ATTACHMENT_EULER: tuple[float, float, float] = (0.0, 0.0, -0.75 * np.pi)

BRIDGE_DEFAULT_ROT module-attribute

BRIDGE_DEFAULT_ROT = np.array([[0.0, 0.0, 1.0], [0.0, 1.0, 0.0], [-1.0, 0.0, 0.0]], dtype=np.float64)

LIBERO_CAMERA_RESOLUTION module-attribute

LIBERO_CAMERA_RESOLUTION: tuple[int, int] = (256, 256)

GRIPPER_DECISION_FIELDS module-attribute

GRIPPER_DECISION_FIELDS = ('input', 'scale', 'threshold', 'open_threshold', 'close_threshold', 'hold_steps', 'sticky_repeats', 'zero_holds', 'continuous', 'snap_threshold')

The :class:GripperDecisionController fields that decide what a gripper command does.

Classes

ActionAdapter

Bases: ABC

A server action row -> {action_term: tensor} decoder.

bind(env) runs once, when the setup that owns the adapter is built: it resolves and validates everything the adapter reads from the env (entity, DOF names), so a bad handle raises at construction instead of on the first control step. reset() clears per-episode state. Both default to no-ops.

Attributes

row_width instance-attribute

row_width: int

Methods:

bind

bind(env: 'EnvBase') -> None

Resolve and validate what this adapter reads from env (default: nothing).

Source code in eden/extensions/policies/adapters.py

def bind(self, env: "EnvBase") -> None:
    """Resolve and validate what this adapter reads from ``env`` (default: nothing)."""

reset

reset() -> None

Clear per-episode state (default: nothing).

Source code in eden/extensions/policies/adapters.py

def reset(self) -> None:
    """Clear per-episode state (default: nothing)."""

EEDeltaAdapter

EEDeltaAdapter(*, rotation: Literal['euler', 'axis_angle'] = 'euler')

Bases: ActionAdapter

Adapter for 7-D end-effector delta actions.

Input shape: (7,) — [dx, dy, dz, drx, dry, drz, gripper].

Output dict: * ee_controller — (num_envs, 6) [dxyz, rotation vector]: the task-space delta :class:DifferentialIKController solves J @ dq = delta for, whose angular rows are an angular displacement (axis * angle), not Euler angles. * gripper_controller — (num_envs, 1) the server's gripper scalar, unchanged; the env's GripperDecisionController interprets it (openness / sign / sticky latch, per setup).

PARAMETER

DESCRIPTION

rotation

Rotation convention of the input delta slice raw[3:6]. "axis_angle" is already a rotation vector and passes through unchanged. "euler" (extrinsic xyz roll / pitch / yaw) is converted to the rotation vector of the same rotation, as SimplerEnv does (euler2axangle) before it steps the arm. WidowX / Google-Robot servers output Euler (LIBERO's axis-angle deltas go to robosuite OSC_POSE, not through this adapter). TYPE: ('euler', 'axis_angle') DEFAULT: "euler"

Source code in eden/extensions/policies/adapters.py

def __init__(self, *, rotation: Literal["euler", "axis_angle"] = "euler"):
    if rotation not in ("euler", "axis_angle"):
        raise ValueError(f"rotation must be 'euler' or 'axis_angle', got {rotation!r}")
    self.rotation = rotation

Attributes

row_width class-attribute instance-attribute

row_width = 7

rotation instance-attribute

rotation = rotation

Methods:

JointVelocityAdapter

JointVelocityAdapter(*, n_arm: int = 7, entity_name: str = 'robot', arm_dofs_name: list[str] | None = None, dt: float | None = None, mode: Literal['velocity', 'velocity_integrate'] = 'velocity', action_scale: float | list[float] = 1.0, clip: float = 1.0, velocity_lag: dict[str, 'Sequence[float]'] | None = None, velocity_compensation: dict[str, 'Sequence[float]'] | None = None, velocity_lag_dt: float | None = None, gripper_threshold: float = 0.5, debug_gripper: bool = False, arm_key: str = 'ee_controller', gripper_key: str = 'gripper_controller')

Bases: _JointAdapterBase

Adapter for joint-velocity + gripper actions (DROID).

Pi0 DROID returns (T, 8) chunks where the first 7 dims are joint velocities in DROID's unit action space, not rad/s, and the last dim is a gripper command binarized at 0.5 on the robot side (here by the env's GripperDecisionController). openpi's examples/droid/main.py clips every dim element-wise to [-1, 1] before RobotEnv.step (which leaves DROID's own v / max|v| rescale nothing to do), and DROID's RobotIKSolver.joint_velocity_to_delta turns a unit velocity into max_joint_delta = 0.2 rad per 15 Hz control step, i.e. 3 rad/s per unit. That conversion is the caller's action_scale (DroidSetup passes 3.0 with clip=1.0 and dt=1/15 s); the default action_scale=1.0 is a pass-through and is not DROID's.

In mode="velocity" (default), the arm output is action_scale * clip(v), the rad/s velocity command consumed by :class:ImplicitVelocityController.

In mode="velocity_integrate", this adapter integrates velocities to joint-position targets on the client side::

target_pos = current_pos + dt * action_scale * clip(v, -1, 1)

and the integrated target is consumed by a standard :class:ImplicitPDController.

Limitation: client-side integration assumes the downstream position controller tracks the reference cleanly. Stiff Franka-like gains usually do, but over long horizons small integration errors accumulate; prefer mode="velocity" wherever the arm has an ImplicitVelocityController.

PARAMETER	DESCRIPTION
n_arm	Number of arm joints (length of the velocity slice in raw_action). TYPE: int DEFAULT: 7
entity_name	Eden entity name used to read current joint positions for integration. TYPE: str DEFAULT: 'robot'
arm_dofs_name	Ordered DOF names under entity_name, one per arm joint. Must match the dofs_name list on the downstream arm controller. Required by the integration modes (resolved on the env in :meth:bind); unused in mode="velocity". TYPE: list[str] | None DEFAULT: None
dt	Integration timestep (seconds). If None, uses env.dt. Only used by integration modes. TYPE: float | None DEFAULT: None
mode	velocity emits rad/s commands directly. velocity_integrate integrates them into position targets (the fallback above). TYPE: ('velocity', 'velocity_integrate') DEFAULT: "velocity"
action_scale	Rad/s per unit of the clipped arm output, applied before it is sent to the downstream controller (scalar, or one value per arm joint). DROID's is 3.0: 0.2 rad per 1/15 s control step (DROID_RAD_PER_S_PER_UNIT in the pi setups). The default 1.0 passes the clipped value through, which is right only for a server whose arm output is already rad/s. TYPE: float | list[float] DEFAULT: 1.0
clip	Each arm dim is clipped element-wise to [-clip, clip] before scaling, in the server's units (DROID: openpi's [-1, 1] unit-space clip, so at most 3 rad/s per joint after the 3.0 scale). TYPE: float DEFAULT: 1.0
velocity_lag	Optional {"tau_s": (n_arm,), "dc_gain": (n_arm,)} first-order lag applied to the commanded velocity in mode="velocity", modelling the real robot's controller bandwidth (Genesis's velocity actuator is ~10x faster than DROID's). None (default) is an exact no-op. See :class:_FirstOrderLag. TYPE: dict | None DEFAULT: None
velocity_compensation	Optional {"real_tau_s", "real_dc_gain", "sim_tau_s", "sim_dc_gain"} (each (n_arm,)): the H_real / H_sim command filter, in mode="velocity". Mutually exclusive with velocity_lag. TYPE: dict | None DEFAULT: None
velocity_lag_dt	The period between two adapter calls in seconds, required when velocity_lag or velocity_compensation is set. A :class:~eden.extensions.policies.client.ClientMixin client executes one action row per env step, so for it this is env.dt (DroidSetup pins that to DROID's 1/15 s and passes it). Only an evaluator that holds one action across several env steps has a longer one (e.g. 0.064 s on a 20 ms env), which is why it is passed rather than read from the env. TYPE: float | None DEFAULT: None
gripper_threshold	Closedness above which the diagnostics label the gripper intent "close". The scalar itself is forwarded unchanged; the env's GripperDecisionController decides. TYPE: float DEFAULT: 0.5
debug_gripper	Log each gripper scalar and its intent. TYPE: bool DEFAULT: False
arm_key	Keys in the returned dict — must match the controller names registered in the env's action manager. TYPE: str DEFAULT: 'ee_controller'

Source code in eden/extensions/policies/adapters.py

def __init__(
    self,
    *,
    n_arm: int = 7,
    entity_name: str = "robot",
    arm_dofs_name: list[str] | None = None,
    dt: float | None = None,
    mode: Literal["velocity", "velocity_integrate"] = "velocity",
    action_scale: float | list[float] = 1.0,
    clip: float = 1.0,
    velocity_lag: dict[str, "Sequence[float]"] | None = None,
    velocity_compensation: dict[str, "Sequence[float]"] | None = None,
    velocity_lag_dt: float | None = None,
    gripper_threshold: float = 0.5,
    debug_gripper: bool = False,
    arm_key: str = "ee_controller",
    gripper_key: str = "gripper_controller",
):
    super().__init__(
        n_arm=n_arm,
        gripper_threshold=gripper_threshold,
        debug_gripper=debug_gripper,
        arm_key=arm_key,
        gripper_key=gripper_key,
    )
    if mode not in ("velocity", "velocity_integrate"):
        raise ValueError(f"JointVelocityAdapter.mode must be 'velocity' or 'velocity_integrate', got {mode!r}")
    self.mode = mode
    self.entity_name = entity_name
    self.arm_dofs_name = None if arm_dofs_name is None else list(arm_dofs_name)
    if mode != "velocity":
        if self.arm_dofs_name is None:
            raise ValueError(
                f"JointVelocityAdapter(mode={mode!r}) integrates from the arm's current joint positions, so "
                f"it needs arm_dofs_name: the {n_arm} arm joint names on {entity_name!r}, in the arm "
                "controller's dofs_name order."
            )
        if len(self.arm_dofs_name) != n_arm:
            raise ValueError(
                f"JointVelocityAdapter: arm_dofs_name names {len(self.arm_dofs_name)} joints "
                f"({self.arm_dofs_name}), but n_arm={n_arm}. Fix: pass exactly one name per arm joint."
            )
    if dt is not None and (not np.isfinite(dt) or dt <= 0.0):
        raise ValueError(f"dt must be a positive finite action period, got {dt!r}")
    self.dt = None if dt is None else float(dt)

    # Shape () or (n_arm,): broadcasts over the arm velocities either way.
    self.action_scale = np.array(action_scale, dtype=np.float32)
    if self.action_scale.shape not in ((), (n_arm,)):
        raise ValueError(
            "action_scale must be a scalar or have one value per arm joint; "
            f"expected shape ({n_arm},), got {self.action_scale.shape}"
        )
    if not np.all(np.isfinite(self.action_scale)):
        raise ValueError("action_scale must contain only finite values")
    if not np.isfinite(clip) or clip <= 0.0:
        raise ValueError(f"clip must be positive and finite, got {clip!r}")
    self.clip = float(clip)

    self._velocity_filter = self._make_velocity_filter(velocity_lag, velocity_compensation, velocity_lag_dt)
    self.velocity_lag_dt = None if velocity_lag_dt is None else float(velocity_lag_dt)
    self._arm_dofs_idx_local: torch.Tensor | None = None

Attributes

mode instance-attribute

mode = mode

entity_name instance-attribute

entity_name = entity_name

arm_dofs_name instance-attribute

arm_dofs_name = None if arm_dofs_name is None else list(arm_dofs_name)

dt instance-attribute

dt = None if dt is None else float(dt)

action_scale instance-attribute

action_scale = np.array(action_scale, dtype=np.float32)

clip instance-attribute

clip = float(clip)

velocity_lag_dt instance-attribute

velocity_lag_dt = None if velocity_lag_dt is None else float(velocity_lag_dt)

Methods:

bind

bind(env: 'EnvBase') -> None

Resolve the arm joints the integration modes read, once, against env.

Only the derived indices are cached; the entity is looked up on every call, since deployment can swap env.entities[entity_name].

Source code in eden/extensions/policies/adapters.py

def bind(self, env: "EnvBase") -> None:
    """Resolve the arm joints the integration modes read, once, against ``env``.

    Only the derived indices are cached; the entity is looked up on every call, since
    deployment can swap ``env.entities[entity_name]``.
    """
    if self.mode == "velocity":
        return
    if self.entity_name not in env.entities:
        raise ValueError(
            f"JointVelocityAdapter.entity_name={self.entity_name!r} is not an entity of the env. "
            f"Available: {sorted(env.entities)}"
        )
    selection = env.entities[self.entity_name].find_named_dofs(self.arm_dofs_name)
    if len(selection.names) != self.n_arm:
        raise ValueError(
            f"JointVelocityAdapter: arm_dofs_name {self.arm_dofs_name} selects {len(selection.names)} joints "
            f"on {self.entity_name!r} ({list(selection.names)}), but n_arm={self.n_arm}. Fix: name each arm "
            "joint exactly (a glob can match more joints than intended)."
        )
    self._arm_dofs_idx_local = selection.dofs_idx_local

reset

reset() -> None

Source code in eden/extensions/policies/adapters.py

def reset(self) -> None:
    super().reset()
    # Stateful filter over the episode's command history - must not leak
    # across episode boundaries.
    if self._velocity_filter is not None:
        self._velocity_filter.reset()

JointPositionAdapter

JointPositionAdapter(*, n_arm: int = 7, gripper_threshold: float = 0.5, debug_gripper: bool = False, arm_key: str = 'ee_controller', gripper_key: str = 'gripper_controller')

Bases: _JointAdapterBase

Adapter for absolute joint-position + gripper actions.

Pi0.5 DROID joint-position checkpoints are trained in delta joint-position space, but OpenPI's output transform reconstructs absolute joint targets by adding the current state back before the client receives the chunk. This adapter therefore passes the first n_arm dimensions directly to an ImplicitPDController configured with reference_source=ZERO, and the gripper scalar to the env's GripperDecisionController. Shared by OpenPI's droid_jointpos, Cosmos3 and GalaxeaVLA's DROID-Franka setup.

Source code in eden/extensions/policies/adapters.py

def __init__(
    self,
    *,
    n_arm: int = 7,
    gripper_threshold: float = 0.5,
    debug_gripper: bool = False,
    arm_key: str = "ee_controller",
    gripper_key: str = "gripper_controller",
):
    if not isinstance(n_arm, int) or isinstance(n_arm, bool) or n_arm <= 0:
        raise ValueError(f"n_arm must be a positive integer, got {n_arm!r}")
    self.n_arm = n_arm
    self.row_width = n_arm + 1
    self.gripper_threshold = gripper_threshold
    self.debug_gripper = debug_gripper
    self.arm_key = arm_key
    self.gripper_key = gripper_key
    self._debug_step = 0
    self.last_diagnostics: dict[str, np.ndarray | float | str] | None = None

BimanualJointPositionAdapter

BimanualJointPositionAdapter(*, left_n: int = 7, right_n: int = 7, left_key: str = 'left_arm_controller', right_key: str = 'right_arm_controller')

Bases: ActionAdapter

Adapter for 14-D bimanual joint-position actions (Aloha).

Splits a (left_n + right_n)-D action vector into its two arms' slices and emits them under separate controller keys.

Source code in eden/extensions/policies/adapters.py

def __init__(
    self,
    *,
    left_n: int = 7,
    right_n: int = 7,
    left_key: str = "left_arm_controller",
    right_key: str = "right_arm_controller",
):
    self.left_n = left_n
    self.row_width = left_n + right_n
    self.left_key = left_key
    self.right_key = right_key

Attributes

left_n instance-attribute

left_n = left_n

row_width instance-attribute

row_width = left_n + right_n

left_key instance-attribute

left_key = left_key

right_key instance-attribute

right_key = right_key

Methods:

Functions:

ee_delta_to_world

ee_delta_to_world(delta_b: Tensor, base_quat_w: Tensor) -> Tensor

A robot-base-frame (num_envs, 6) [dxyz, rotation vector] delta in the world frame.

EE-delta servers (Bridge, Fractal, RobotArena) compose their deltas in the robot base frame, while :class:DifferentialIKController solves against the world-frame Jacobian: the two agree only for a base with identity rotation. Both parts are displacements, so the base's world orientation base_quat_w ((num_envs, 4) wxyz, the controlled entity's get_quat()) alone maps them.

Source code in eden/extensions/policies/adapters.py

def ee_delta_to_world(delta_b: torch.Tensor, base_quat_w: torch.Tensor) -> torch.Tensor:
    """A robot-base-frame ``(num_envs, 6)`` ``[dxyz, rotation vector]`` delta in the world frame.

    EE-delta servers (Bridge, Fractal, RobotArena) compose their deltas in the robot base frame, while
    :class:`DifferentialIKController` solves against the world-frame Jacobian: the two agree only for a
    base with identity rotation. Both parts are displacements, so the base's world orientation
    ``base_quat_w`` (``(num_envs, 4)`` wxyz, the controlled entity's ``get_quat()``) alone maps them.
    """
    return torch.cat(
        [transform_by_quat(delta_b[:, :3], base_quat_w), transform_by_quat(delta_b[:, 3:6], base_quat_w)], dim=-1
    )

bridge_rpy

bridge_rpy(quat_wxyz: ndarray) -> ndarray

Bridge / SimplerEnv WidowX rpy: the EE rotation relative to :data:BRIDGE_DEFAULT_ROT, xyz Euler.

R @ BRIDGE_DEFAULT_ROT.T as quaternions, decomposed into extrinsic xyz Euler angles.

Source code in eden/extensions/policies/adapters.py

def bridge_rpy(quat_wxyz: np.ndarray) -> np.ndarray:
    """Bridge / SimplerEnv WidowX ``rpy``: the EE rotation relative to :data:`BRIDGE_DEFAULT_ROT`, xyz Euler.

    ``R @ BRIDGE_DEFAULT_ROT.T`` as quaternions, decomposed into extrinsic xyz Euler angles.
    """
    quat = torch.as_tensor(quat_wxyz, dtype=torch.float64)
    return torch.stack(euler_xyz_from_quat(quat_mul(quat, _BRIDGE_DEFAULT_QUAT_INV))).numpy()

robosuite_axis_angle

robosuite_axis_angle(quat_wxyz: ndarray) -> ndarray

Robosuite's quat2axisangle for a wxyz quaternion, as LIBERO's official clients compute EE state.

Both OpenPI (examples/libero/main.py) and Isaac-GR00T (gr00t/eval/sim/LIBERO/libero_env.py) copy it. Kept robosuite-exact -- no sign canonicalisation (the angle may exceed pi when w < 0), hard zero at the singularity -- rather than routed through :func:eden.utils.geom.axis_angle_from_quat or Rotation.as_rotvec, because the servers' normalisation statistics were fit on robosuite's output.

Source code in eden/extensions/policies/adapters.py

def robosuite_axis_angle(quat_wxyz: np.ndarray) -> np.ndarray:
    """Robosuite's ``quat2axisangle`` for a wxyz quaternion, as LIBERO's official clients compute EE state.

    Both OpenPI (``examples/libero/main.py``) and Isaac-GR00T (``gr00t/eval/sim/LIBERO/libero_env.py``)
    copy it. Kept robosuite-exact -- no sign canonicalisation (the angle may exceed pi when w < 0), hard
    zero at the singularity -- rather than routed through :func:`eden.utils.geom.axis_angle_from_quat` or
    ``Rotation.as_rotvec``, because the servers' normalisation statistics were fit on robosuite's output.
    """
    w, x, y, z = quat_wxyz
    w = float(np.clip(w, -1.0, 1.0))
    den = math.sqrt(max(0.0, 1.0 - w * w))
    if math.isclose(den, 0.0):
        return np.zeros(3)
    angle = 2.0 * math.acos(w)
    return np.array([x, y, z]) * angle / den

libero_task_action_options

libero_task_action_options() -> dict[str, 'ActionTermOptions']

The stock libero task's action terms, from the task's own builder (no restated constants).

osc_pose (robosuite OSC_POSE on the grip site) and gripper (the stateful LiberoPandaGripperController, +1 closes / -1 opens): the controllers LIBERO's demonstrations, and so every LIBERO checkpoint (OpenPI, GR00T), were recorded with.

Source code in eden/extensions/policies/adapters.py

def libero_task_action_options() -> dict[str, "ActionTermOptions"]:
    """The stock ``libero`` task's action terms, from the task's own builder (no restated constants).

    ``osc_pose`` (robosuite ``OSC_POSE`` on the grip site) and ``gripper`` (the stateful
    ``LiberoPandaGripperController``, ``+1`` closes / ``-1`` opens): the controllers LIBERO's
    demonstrations, and so every LIBERO checkpoint (OpenPI, GR00T), were recorded with.
    """
    from eden.tasks.libero.config import _make_action_options

    options = _make_action_options()
    # term_keys(), not model_extra: the extras also carry the _option_module_ / _option_class_ markers.
    return {name: getattr(options, name) for name in options.term_keys()}

check_libero_camera_resolution

check_libero_camera_resolution(env, *, subject: str, trained_on: str, wrong_fix: str) -> None

Refuse an env whose image / wrist_image groups render at another size than LIBERO's 256x256.

The camera is the one each group's term names in params["camera_name"] (:func:make_image_groups); a camera term without one is judged by the image its group publishes. subject names the client and setup, trained_on the upstream render the checkpoints saw, and wrong_fix is the closing "Do not ..." sentence about the client's own image-size knob. Envs without an observation manager skip.

Source code in eden/extensions/policies/adapters.py

def check_libero_camera_resolution(env, *, subject: str, trained_on: str, wrong_fix: str) -> None:
    """Refuse an env whose ``image`` / ``wrist_image`` groups render at another size than LIBERO's 256x256.

    The camera is the one each group's term names in ``params["camera_name"]`` (:func:`make_image_groups`);
    a camera term without one is judged by the image its group publishes. ``subject`` names the client and
    setup, ``trained_on`` the upstream render the checkpoints saw, and ``wrong_fix`` is the closing "Do not
    ..." sentence about the client's own image-size knob. Envs without an observation manager skip.
    """
    manager = getattr(env, "observation_manager", None)
    if manager is None:
        return
    width, height = LIBERO_CAMERA_RESOLUTION
    for group, name, term in manager.iter_terms():
        if group not in ("image", "wrist_image"):
            continue
        camera = (getattr(term, "params", None) or {}).get("camera_name")
        if camera is None:
            rows, cols = manager.group_obs_dim[group][:2]
            resolution, where, camera = (
                (int(cols), int(rows)),
                f"the camera behind observation term {group}/{name}",
                "<camera>",
            )
        else:
            resolution = tuple(int(v) for v in env.cameras[camera].cam_resolution)
            where = f"camera {camera!r} (observation term {group}/{name})"
        if resolution == LIBERO_CAMERA_RESOLUTION:
            continue
        raise ValueError(
            f"{subject} sends {where}, which renders at cam_resolution={resolution}, but {trained_on}, so the frame "
            f"would be upsampled off that distribution. Fix: in the env's cameras_options set the {camera!r} camera's "
            f"cam_resolution=({width}, {height}), e.g. cfg.cameras_options.{camera}.cam_resolution = ({width}, "
            f"{height}). {wrong_fix}"
        )

libero_training_view

libero_training_view(image: Tensor | ndarray) -> ndarray

An upright Eden render as LIBERO's training images show it: mirrored left-right, (H, W, 3) uint8.

LIBERO's official clients send obs["agentview_image"][::-1, ::-1] and the same for robot0_eye_in_hand_image (OpenPI examples/libero/main.py lines 115-116, Isaac-GR00T gr00t/eval/sim/LIBERO/libero_env.py lines 148-149), as the LIBERO datasets were recorded. robosuite 1.4's raw frame is OpenGL bottom-up (the upright view flipped vertically), so that pair of flips leaves the upright view mirrored left-right. Eden renders upright, so only the left-right mirror reproduces it; porting [::-1, ::-1] literally would turn the image upside down.

Source code in eden/extensions/policies/adapters.py

def libero_training_view(image: torch.Tensor | np.ndarray) -> np.ndarray:
    """An upright Eden render as LIBERO's training images show it: mirrored left-right, ``(H, W, 3)`` uint8.

    LIBERO's official clients send ``obs["agentview_image"][::-1, ::-1]`` and the same for
    ``robot0_eye_in_hand_image`` (OpenPI ``examples/libero/main.py`` lines 115-116, Isaac-GR00T
    ``gr00t/eval/sim/LIBERO/libero_env.py`` lines 148-149), as the LIBERO datasets were recorded. robosuite
    1.4's raw frame is OpenGL bottom-up (the upright view flipped vertically), so that pair of flips leaves the
    upright view mirrored left-right. Eden renders upright, so only the left-right mirror reproduces it;
    porting ``[::-1, ::-1]`` literally would turn the image upside down.
    """
    return np.ascontiguousarray(to_uint8_hwc(image)[:, ::-1])

dofs_pos_joint_names

dofs_pos_joint_names(env: 'EnvBase', term) -> list[str]

The joints a built DofsPos term reads, in its column order.

Source code in eden/extensions/policies/adapters.py

def dofs_pos_joint_names(env: "EnvBase", term) -> list[str]:
    """The joints a built ``DofsPos`` term reads, in its column order."""
    entity = env.entities[term.entity_name]
    if term.dofs_name is not None:
        return list(entity.find_named_dofs(list(term.dofs_name)).names)
    if term.terms_idx is not None:
        return [entity.dofs_name[idx] for idx in term.terms_idx]
    return list(entity.dofs_name)

gripper_dofs_from_name

gripper_dofs_from_name(gripper_dof_name: str) -> list[str]

Expand a single Robotiq 2F-85 driver name to both actuated drivers; other names pass through.

Source code in eden/extensions/policies/adapters.py

def gripper_dofs_from_name(gripper_dof_name: str) -> list[str]:
    """Expand a single Robotiq 2F-85 driver name to both actuated drivers; other names pass through."""
    if gripper_dof_name in Robotiq2f85.actuated_dofs_name:
        return list(Robotiq2f85.actuated_dofs_name)
    return [gripper_dof_name]

make_image_groups

make_image_groups(camera_name: str, use_wrist_image: bool, wrist_camera_name: str) -> dict[str, ObservationGroupOptions]

Build image (and optionally wrist_image) observation groups, one rgb term each.

Source code in eden/extensions/policies/adapters.py

def make_image_groups(
    camera_name: str,
    use_wrist_image: bool,
    wrist_camera_name: str,
) -> dict[str, ObservationGroupOptions]:
    """Build ``image`` (and optionally ``wrist_image``) observation groups, one ``rgb`` term each."""
    from eden.managers.terms.observations.camera import rgb_image
    from eden.options.managers.observations import ObservationGroupOptions
    from eden.terms import ObsTerm

    groups: dict[str, ObservationGroupOptions] = {
        "image": ObservationGroupOptions(
            rgb=ObsTerm.configure(func=rgb_image, params={"camera_name": camera_name}),
        ),
    }
    if use_wrist_image:
        groups["wrist_image"] = ObservationGroupOptions(
            rgb=ObsTerm.configure(func=rgb_image, params={"camera_name": wrist_camera_name}),
        )
    return groups

make_ee_pose_group

make_ee_pose_group(entity_name: str, ee_link_name: str, *, gripper_entity_name: str | None, gripper_dof_name: str, open_gripper_pos: float = open_finger_pos, gripper_close_pos: float = ROBOTIQ_CLOSE_POS, offset_pos: tuple[float, float, float] = (0.0, 0.0, 0.0)) -> ObservationGroupOptions

[ee_pos(3), ee_quat_wxyz(4), gripper_openness(1)] proprio group (Bridge / Google-robot / RobotArena).

With the gripper on the arm entity, openness is pos / open_gripper_pos (WidowX fingers: fully open at open_gripper_pos, closed at 0). With a separately mounted gripper it is 1 - pos / gripper_close_pos (Robotiq driver: open at 0, closed at gripper_close_pos).

Source code in eden/extensions/policies/adapters.py

def make_ee_pose_group(
    entity_name: str,
    ee_link_name: str,
    *,
    gripper_entity_name: str | None,
    gripper_dof_name: str,
    open_gripper_pos: float = WidowX250.open_finger_pos,
    gripper_close_pos: float = ROBOTIQ_CLOSE_POS,
    offset_pos: tuple[float, float, float] = (0.0, 0.0, 0.0),
) -> ObservationGroupOptions:
    """``[ee_pos(3), ee_quat_wxyz(4), gripper_openness(1)]`` proprio group (Bridge / Google-robot / RobotArena).

    With the gripper on the arm entity, openness is ``pos / open_gripper_pos`` (WidowX fingers:
    fully open at ``open_gripper_pos``, closed at 0). With a separately mounted gripper it is
    ``1 - pos / gripper_close_pos`` (Robotiq driver: open at 0, closed at ``gripper_close_pos``).
    """
    from eden.managers.terms.observations.manipulation import EndEffectorPose, GripperState
    from eden.options.managers.observations import ObservationGroupOptions

    if gripper_entity_name is None:  # WidowX fingers on the arm: open at open_gripper_pos, closed at 0
        gripper_entity_name, open_pos, close_pos = entity_name, open_gripper_pos, 0.0
    else:  # a mounted Robotiq: open at 0, closed at gripper_close_pos
        open_pos, close_pos = 0.0, gripper_close_pos
    return ObservationGroupOptions(
        ee_pose=EndEffectorPose.configure(
            entity_name=entity_name, ee_link_name=ee_link_name, rotation="quat", offset_pos=tuple(offset_pos)
        ),
        gripper=GripperState.configure(
            entity_name=gripper_entity_name,
            dofs_name=[gripper_dof_name],
            normalize="openness",
            open_pos=open_pos,
            close_pos=close_pos,
        ),
    )

make_droid_joint_state_group

make_droid_joint_state_group(arm_entity_name: str, arm_dofs_name: list[str] | tuple[str, ...], *, gripper_entity_name: str | None = None, gripper_dof_name: str, gripper_close_pos: float = ROBOTIQ_CLOSE_POS) -> ObservationGroupOptions

[joint_pos(N), gripper_closedness(1)] proprio group for DROID-style joint-state servers.

The last value is DROID's robot_state["gripper_position"]: closedness pos / gripper_close_pos clipped to [0, 1] (a Robotiq driver can travel past its close target), read from gripper_dof_name on gripper_entity_name -- the arm entity when that is None.

Source code in eden/extensions/policies/adapters.py

def make_droid_joint_state_group(
    arm_entity_name: str,
    arm_dofs_name: list[str] | tuple[str, ...],
    *,
    gripper_entity_name: str | None = None,
    gripper_dof_name: str,
    gripper_close_pos: float = ROBOTIQ_CLOSE_POS,
) -> ObservationGroupOptions:
    """``[joint_pos(N), gripper_closedness(1)]`` proprio group for DROID-style joint-state servers.

    The last value is DROID's ``robot_state["gripper_position"]``: closedness ``pos / gripper_close_pos``
    clipped to ``[0, 1]`` (a Robotiq driver can travel past its close target), read from
    ``gripper_dof_name`` on ``gripper_entity_name`` -- the arm entity when that is ``None``.
    """
    from eden.managers.terms.observations.manipulation import GripperState
    from eden.managers.terms.observations.proprio import DofsPos
    from eden.options.managers.observations import ObservationGroupOptions

    return ObservationGroupOptions(
        joint_pos=DofsPos.configure(entity_name=arm_entity_name, dofs_name=list(arm_dofs_name)),
        gripper=GripperState.configure(
            entity_name=gripper_entity_name if gripper_entity_name is not None else arm_entity_name,
            dofs_name=[gripper_dof_name],
            normalize="closedness",
            open_pos=0.0,
            close_pos=gripper_close_pos,
            clip=(0.0, 1.0),
        ),
    )

make_widowx_ee_delta_actions

make_widowx_ee_delta_actions(entity_name: str, ee_link_name: str, dofs_name: list[str], scale: float, *, gripper_input: str = 'openness', gripper_threshold: float = 0.5, gripper_sticky_repeats: int = 0, gripper_zero_holds: bool = False) -> dict[str, ActionTermOptions]

ee_controller (differential IK) + gripper_controller for WidowX EE-delta policies.

The gripper term is a :class:~eden.managers.terms.actions.binary_actions.GripperDecisionController that receives the server's raw gripper scalar (Bridge servers emit openness, > 0.5 = open) and decides open/close itself; gripper_sticky_repeats enables the fractal latch the Google-robot checkpoints were trained with.

Source code in eden/extensions/policies/adapters.py

def make_widowx_ee_delta_actions(
    entity_name: str,
    ee_link_name: str,
    dofs_name: list[str],
    scale: float,
    *,
    gripper_input: str = "openness",
    gripper_threshold: float = 0.5,
    gripper_sticky_repeats: int = 0,
    gripper_zero_holds: bool = False,
) -> dict[str, ActionTermOptions]:
    """``ee_controller`` (differential IK) + ``gripper_controller`` for WidowX EE-delta policies.

    The gripper term is a :class:`~eden.managers.terms.actions.binary_actions.GripperDecisionController`
    that receives the server's raw gripper scalar (Bridge servers emit *openness*, ``> 0.5`` =
    open) and decides open/close itself; ``gripper_sticky_repeats`` enables the fractal latch
    the Google-robot checkpoints were trained with.
    """
    from eden.managers.terms.actions.binary_actions import GripperDecisionController
    from eden.managers.terms.actions.task_space_actions import DifferentialIKController

    return dict(
        ee_controller=DifferentialIKController.configure(
            entity_name=entity_name,
            ee_link_name=ee_link_name,
            dofs_name=dofs_name,
            position_scale=scale,
            orientation_scale=scale,
        ),
        gripper_controller=GripperDecisionController.configure(
            entity_name=entity_name,
            dofs_name=list(WidowX250.gripper_dofs_name),
            open_action={"left_finger": WidowX250.open_finger_pos, "right_finger": -WidowX250.open_finger_pos},
            close_action=0.0,
            input=gripper_input,
            threshold=gripper_threshold,
            sticky_repeats=gripper_sticky_repeats,
            zero_holds=gripper_zero_holds,
        ),
    )

make_droid_joint_position_actions

make_droid_joint_position_actions(arm_entity_name: str = 'robot', arm_dofs_name: list[str] | None = None, *, gripper_entity_name: str | None = None, gripper_dof_name: str = 'left_driver_joint', arm_key: str = 'ee_controller', **gripper_options) -> dict[str, ActionTermOptions]

Absolute DROID arm joint targets + the Robotiq gripper term, as joint-position DROID servers need.

The servers return absolute joint targets, so the arm is an ImplicitPDController with a zero reference. The gripper is :func:make_droid_gripper_action on gripper_entity_name (the arm entity when the Robotiq is part of it), configured by gripper_options.

Source code in eden/extensions/policies/adapters.py

def make_droid_joint_position_actions(
    arm_entity_name: str = "robot",
    arm_dofs_name: list[str] | None = None,
    *,
    gripper_entity_name: str | None = None,
    gripper_dof_name: str = "left_driver_joint",
    arm_key: str = "ee_controller",
    **gripper_options,
) -> dict[str, ActionTermOptions]:
    """Absolute DROID arm joint targets + the Robotiq gripper term, as joint-position DROID servers need.

    The servers return absolute joint targets, so the arm is an ``ImplicitPDController`` with a zero
    reference. The gripper is :func:`make_droid_gripper_action` on ``gripper_entity_name`` (the arm
    entity when the Robotiq is part of it), configured by ``gripper_options``.
    """
    from eden.constants import ReferenceSource
    from eden.managers.terms.actions.joint_actions import ImplicitPDController

    arm = ImplicitPDController.configure(
        entity_name=arm_entity_name,
        dofs_name=list(DROID_ARM_DOFS) if arm_dofs_name is None else list(arm_dofs_name),
        reference_source=ReferenceSource.ZERO,
    )
    gripper = make_droid_gripper_action(
        gripper_entity_name if gripper_entity_name is not None else arm_entity_name,
        gripper_dof_name,
        **gripper_options,
    )
    return {arm_key: arm, "gripper_controller": gripper}

make_droid_gripper_action

make_droid_gripper_action(gripper_entity_name: str, gripper_dof_name: str, *, gripper_close_pos: float = ROBOTIQ_CLOSE_POS, input: str = 'closedness', threshold: float = 0.5, open_threshold: float | None = None, close_threshold: float | None = None, hold_steps: int = 0, continuous: bool = False, snap_threshold: float | None = None, on_missing: str = 'hold') -> ActionTermOptions

gripper_controller for DROID-style servers: closedness in, [0, gripper_close_pos] targets out.

The open/close decision (threshold, hysteresis, hold, continuous) is made by the term, not the client. The single Robotiq driver name expands to both actuated drivers.

Source code in eden/extensions/policies/adapters.py

def make_droid_gripper_action(
    gripper_entity_name: str,
    gripper_dof_name: str,
    *,
    gripper_close_pos: float = ROBOTIQ_CLOSE_POS,
    input: str = "closedness",
    threshold: float = 0.5,
    open_threshold: float | None = None,
    close_threshold: float | None = None,
    hold_steps: int = 0,
    continuous: bool = False,
    snap_threshold: float | None = None,
    on_missing: str = "hold",
) -> ActionTermOptions:
    """``gripper_controller`` for DROID-style servers: closedness in, ``[0, gripper_close_pos]`` targets out.

    The open/close decision (threshold, hysteresis, hold, continuous) is made by the term, not
    the client. The single Robotiq driver name expands to both actuated drivers.
    """
    from eden.managers.terms.actions.binary_actions import GripperDecisionController

    return GripperDecisionController.configure(
        entity_name=gripper_entity_name,
        dofs_name=gripper_dofs_from_name(gripper_dof_name),
        open_action=0.0,
        close_action=gripper_close_pos,
        input=input,
        threshold=threshold,
        open_threshold=open_threshold,
        close_threshold=close_threshold,
        hold_steps=hold_steps,
        continuous=continuous,
        snap_threshold=snap_threshold,
        on_missing=on_missing,
    )

check_control_rate

check_control_rate(env: 'EnvBase', *, control_hz: float, subject: str, source: str, timing: tuple[str, int] | None = None, timing_note: str = '', wrong_fix: str | None = None) -> None

Refuse an env whose control period env.dt is not the upstream client's 1 / control_hz s.

A remote checkpoint's rows are one upstream control step each (an EE or joint delta, the next waypoint of a trajectory), and exec_horizon / a video window's stride count env steps, so at another rate every chunk executes at the wrong speed, silently. subject names the client and setup ("PiPolicy.policy_setup='droid'") and source where upstream sets the rate. timing is the (sim_dt as typed, decimation) the fix suggests, with timing_note appended (e.g. the sim_substeps a Robotiq 2F85 needs); None derives one from the env's sim_dt. wrong_fix replaces the closing "do not rescale or resample the actions" sentence for a client with a more plausible wrong fix.

Source code in eden/extensions/policies/adapters.py

def check_control_rate(
    env: "EnvBase",
    *,
    control_hz: float,
    subject: str,
    source: str,
    timing: tuple[str, int] | None = None,
    timing_note: str = "",
    wrong_fix: str | None = None,
) -> None:
    """Refuse an env whose control period ``env.dt`` is not the upstream client's ``1 / control_hz`` s.

    A remote checkpoint's rows are one upstream control step each (an EE or joint delta, the next waypoint
    of a trajectory), and ``exec_horizon`` / a video window's stride count env steps, so at another rate
    every chunk executes at the wrong speed, silently. ``subject`` names the client and setup
    (``"PiPolicy.policy_setup='droid'"``) and ``source`` where upstream sets the rate. ``timing`` is the
    ``(sim_dt as typed, decimation)`` the fix suggests, with ``timing_note`` appended (e.g. the
    ``sim_substeps`` a Robotiq 2F85 needs); ``None`` derives one from the env's ``sim_dt``. ``wrong_fix`` replaces
    the closing "do not rescale or resample the actions" sentence for a client with a more plausible wrong fix.
    """
    env_dt = float(env.dt)
    period = 1.0 / control_hz
    if math.isclose(env_dt, period, rel_tol=1e-6):
        return
    speed = f"{period / env_dt:.2f}x too fast" if env_dt < period else f"{env_dt / period:.2f}x too slow"
    if timing is None:
        decimation = max(1, round(period / env.sim_dt))
        exact = math.isclose(env.sim_dt * decimation, period, rel_tol=1e-6)
        timing = (f"{env.sim_dt!r}" if exact else f"1/{control_hz * decimation:g}", decimation)
    sim_dt, decimation = timing
    if wrong_fix is None:
        wrong_fix = (
            f"Do not rescale or resample the actions instead: the checkpoint's chunk is a trajectory at "
            f"{control_hz:g} Hz, and a faster env would also re-query the server that much more often."
        )
    raise ValueError(
        f"{subject} executes one action row per {period:.6f} s ({control_hz:g} Hz) control step, as {source}, but "
        f"the env steps every env.dt={env_dt:.6f} s (sim_dt * decimation), which would run each chunk {speed}. Fix: "
        f"build the env with sim_dt * decimation = 1/{control_hz:g} s, e.g. EnvOptions(sim_dt={sim_dt}, "
        f"decimation={decimation}){timing_note}. {wrong_fix}"
    )

check_env_gripper_term

check_env_gripper_term(env, term_name: str, **expected) -> None

Fail fast when the env's built gripper term disagrees with settings a policy also carries.

Some client policies expose gripper knobs (thresholds, hold, decode mode) that the env's :class:GripperDecisionController also carries. When the env is built (it always is before a policy is constructed in the drivers), the two are compared so a mismatch is reported at setup time rather than silently ignored. Envs without an action manager (unit tests) skip.

Every field of :data:GRIPPER_DECISION_FIELDS is compared: one the caller does not pass in expected must be at the controller's default, because hysteresis, a hold, a sticky latch or continuous mode each replace the plain threshold rule the caller relies on.

Clients forward the server's raw gripper scalar, so any other gripper term (e.g. a task config's own BinaryJointController) would silently misread it: a missing term, a term of another type, or an expected field the term does not have all raise.

Source code in eden/extensions/policies/adapters.py

def check_env_gripper_term(env, term_name: str, **expected) -> None:
    """Fail fast when the env's built gripper term disagrees with settings a policy also carries.

    Some client policies expose gripper knobs (thresholds, hold, decode mode) that the env's
    :class:`GripperDecisionController` also carries. When the env is built (it always is before a
    policy is constructed in the drivers), the two are compared so a mismatch is reported at
    setup time rather than silently ignored. Envs without an action manager (unit tests) skip.

    Every field of :data:`GRIPPER_DECISION_FIELDS` is compared: one the caller does not pass in
    ``expected`` must be at the controller's default, because hysteresis, a hold, a sticky latch or
    continuous mode each replace the plain threshold rule the caller relies on.

    Clients forward the server's raw gripper scalar, so any other gripper term (e.g. a task
    config's own ``BinaryJointController``) would silently misread it: a missing term, a term of
    another type, or an expected field the term does not have all raise.
    """
    from eden.managers.terms.actions.binary_actions import GripperDecisionController

    manager = getattr(env, "action_manager", None)
    if manager is None:
        return
    term = manager.terms.get(term_name)
    if term is None:
        raise ValueError(
            f"The env has no action term {term_name!r}; available: {sorted(manager.terms)}. Build the env "
            "with the policy's action_options() so its GripperDecisionController receives the gripper command."
        )
    if not isinstance(term, GripperDecisionController):
        raise TypeError(
            f"Action term {term_name!r} is a {type(term).__name__}, but this policy forwards the server's raw "
            "gripper scalar and needs a GripperDecisionController to interpret it. Build the env's gripper "
            "term from the policy's action_options()."
        )
    unknown = sorted(name for name in expected if not hasattr(term, name))
    if unknown:
        raise AttributeError(f"Gripper term {term_name!r} has no field(s) {unknown} to check against the policy.")
    undeclared = {
        name: getattr(GripperDecisionController, name) for name in GRIPPER_DECISION_FIELDS if name not in expected
    }
    mismatched = {
        name: (getattr(term, name), value)
        for name, value in {**expected, **undeclared}.items()
        if getattr(term, name) != value
    }
    if mismatched:
        details = ", ".join(
            f"{k}: env term has {a!r}, policy {'has' if k in expected else 'does not use it (default)'} {b!r}"
            for k, (a, b) in mismatched.items()
        )
        raise ValueError(
            f"Gripper term {term_name!r} was configured differently from the policy ({details}). "
            "The decision now lives in the env's GripperDecisionController: build the env's gripper term from "
            "the policy's action_options(), not from another client's or a task's own gripper term."
        )

check_plain_observation_group

check_plain_observation_group(env, group: str, *, width: int | None = None, scales: dict[str, Sequence[float]] | None = None, fix: str) -> None

Fail fast when an observation group a client slices by position is not the plain current reading.

Clients pack a group by column (proprio[:7] is the arm), so the group must be one flat (num_envs, width) tensor of this step's values: no history window (it puts older frames first), no scale/clip/modifiers (a delay or a rescale reaches the server unannounced), and no noise when the group corrupts. scales names the terms whose post-compute scale is part of the wire contract (term -> the exact per-column scale the factory sets, e.g. robosuite's second-finger sign flip): such a term must exist and carry exactly that scale, every other term a scale of 1. fix is appended to the message and should name the factory that builds the group. Envs without an observation manager (unit tests) skip.

Source code in eden/extensions/policies/adapters.py

def check_plain_observation_group(
    env, group: str, *, width: int | None = None, scales: dict[str, Sequence[float]] | None = None, fix: str
) -> None:
    """Fail fast when an observation group a client slices by position is not the plain current reading.

    Clients pack a group by column (``proprio[:7]`` is the arm), so the group must be one flat
    ``(num_envs, width)`` tensor of this step's values: no history window (it puts older frames
    first), no ``scale``/``clip``/``modifiers`` (a delay or a rescale reaches the server unannounced),
    and no noise when the group corrupts. ``scales`` names the terms whose post-compute ``scale`` is
    part of the wire contract (term -> the exact per-column scale the factory sets, e.g. robosuite's
    second-finger sign flip): such a term must exist and carry exactly that scale, every other term
    a scale of 1. ``fix`` is appended to the message and should name the factory that builds the
    group. Envs without an observation manager (unit tests) skip.
    """
    manager = getattr(env, "observation_manager", None)
    if manager is None:
        return
    groups = env.config.observation_options
    if group not in groups.term_keys():
        raise ValueError(
            f"The env has no observation group {group!r}; available: {sorted(groups.term_keys())}. Fix: {fix}"
        )
    group_options = getattr(groups, group)
    problems = []
    if not group_options.concatenate_terms:
        problems.append("concatenate_terms=False (the client needs one tensor)")
    if group_options.history_length is not None:
        problems.append(f"history_length={group_options.history_length}")
    from eden.managers.terms.observations.manipulation import GripperState

    built = {name: obs_term for g, name, obs_term in manager.iter_terms() if g == group}
    scales = scales or {}
    term_names = manager.group_term_names[group]
    problems.extend(
        f"no term {name!r} (it carries scale={tuple(scales[name])!r})" for name in scales if name not in term_names
    )
    for name in term_names:
        term = getattr(group_options, name)
        if term.history_length:
            problems.append(f"{name}.history_length={term.history_length}")
        scale = np.asarray(term.scale, dtype=float)
        if name in scales:
            if not np.array_equal(scale, np.asarray(scales[name], dtype=float)):
                problems.append(f"{name}.scale={term.scale!r} (expected {tuple(scales[name])!r})")
        elif not bool(np.all(scale == 1.0)):
            problems.append(f"{name}.scale={term.scale!r}")
        # A [0, 1] clip on a normalised gripper reading only bounds an overshoot past the stroke ends
        # (make_droid_joint_state_group sets it); any other clip changes values the server reads.
        unit_gripper_clip = (
            isinstance(built.get(name), GripperState)
            and built[name].normalize != "raw"
            and tuple(term.clip or ()) == (0.0, 1.0)
        )
        if term.clip is not None and not unit_gripper_clip:
            problems.append(f"{name}.clip={term.clip!r}")
        if term.modifiers:
            problems.append(f"{name}.modifiers={[type(m).__name__ for m in term.modifiers]}")
        if term.noise is not None and group_options.enable_corruption:
            problems.append(f"{name}.noise={type(term.noise).__name__}")
    actual = manager.group_obs_dim[group]
    if width is not None and not problems and tuple(actual) != (width,):
        problems.append(f"its terms {manager.group_term_names[group]} are {tuple(actual)} wide, not ({width},)")
    if problems:
        raise ValueError(
            f"Observation group {group!r} is not the plain current reading the client packs by column: "
            f"{'; '.join(problems)}. A history or a post-compute transform shifts or rescales what the server "
            f"reads. Fix: {fix}"
        )

check_gripper_state_matches_action

check_gripper_state_matches_action(env, *, group: str, term: str, action_term: str) -> None

Fail fast when a gripper state term and the gripper action term describe different grippers.

A client sends the state normalised over [open_pos, close_pos] and the env maps the server's command onto [open_action, close_action]; the two come from separate observation_options and action_options arguments, so a stroke set in one call only would skew every request. The state must read a DOF of the action term's entity, and its open / close positions must equal the action term's (per-DOF dict targets are compared for that DOF). Envs without managers skip.

Source code in eden/extensions/policies/adapters.py

def check_gripper_state_matches_action(env, *, group: str, term: str, action_term: str) -> None:
    """Fail fast when a gripper state term and the gripper action term describe different grippers.

    A client sends the state normalised over ``[open_pos, close_pos]`` and the env maps the server's
    command onto ``[open_action, close_action]``; the two come from separate ``observation_options``
    and ``action_options`` arguments, so a stroke set in one call only would skew every request. The
    state must read a DOF of the action term's entity, and its open / close positions must equal the
    action term's (per-DOF dict targets are compared for that DOF). Envs without managers skip.
    """
    obs_manager = getattr(env, "observation_manager", None)
    action_manager = getattr(env, "action_manager", None)
    if obs_manager is None or action_manager is None:
        return
    from eden.managers.terms.observations.manipulation import GripperState

    state = next((t for g, n, t in obs_manager.iter_terms() if (g, n) == (group, term)), None)
    action = action_manager.terms.get(action_term)
    if not isinstance(state, GripperState) or action is None:
        return  # The callers' own layout and gripper-term checks name a missing or foreign term.
    problems = []
    if state.entity_name != action.entity_name or not set(state.dofs_name) <= set(action.dofs_name):
        problems.append(
            f"it reads {state.entity_name!r} {list(state.dofs_name)} but {action_term!r} drives "
            f"{action.entity_name!r} {list(action.dofs_name)}"
        )
    dof = state.dofs_name[0] if state.dofs_name else None
    for side, state_pos, target in (
        ("open", state.open_pos, action.open_action),
        ("close", state.close_pos, action.close_action),
    ):
        target = target.get(dof) if isinstance(target, dict) else target
        if state_pos is not None and target is not None and not math.isclose(state_pos, target, abs_tol=1e-6):
            problems.append(f"{side}_pos={state_pos!r} but {action_term!r} has {side}_action={target!r}")
    if problems:
        raise ValueError(
            f"Gripper state {group}/{term} and action term {action_term!r} describe different grippers: "
            f"{'; '.join(problems)}. Fix: pass the same gripper entity and the same open/close positions to "
            "the policy's observation_options() and action_options()."
        )

check_joint_pass_through

check_joint_pass_through(term, *, term_name: str, subject: str, fix: str, commands: str = "the server's absolute joint targets", wrong_fix: str = _PASS_THROUGH_WRONG_FIX) -> None

Raise unless joint action term executes the client's commands untouched (see :func:_pass_through_problems).

subject names the client and setup ("PiPolicy.policy_setup='aloha'"); commands what it sends (absolute targets by default); fix the action_options(...) call that builds the right term; and wrong_fix the client-side workaround to rule out. The caller has already checked the term's type. Every problem is listed at once.

Source code in eden/extensions/policies/adapters.py

def check_joint_pass_through(
    term,
    *,
    term_name: str,
    subject: str,
    fix: str,
    commands: str = "the server's absolute joint targets",
    wrong_fix: str = _PASS_THROUGH_WRONG_FIX,
) -> None:
    """Raise unless joint action ``term`` executes the client's commands untouched (see :func:`_pass_through_problems`).

    ``subject`` names the client and setup (``"PiPolicy.policy_setup='aloha'"``); ``commands`` what it sends
    (absolute targets by default); ``fix`` the ``action_options(...)`` call that builds the right term; and
    ``wrong_fix`` the client-side workaround to rule out. The caller has already checked the term's type. Every
    problem is listed at once.
    """
    _refuse_term_problems(
        term,
        _pass_through_problems(term),
        term_name=term_name,
        subject=subject,
        fix=fix,
        commands=commands,
        wrong_fix=wrong_fix,
    )

check_droid_arm_term

check_droid_arm_term(arm, *, term_name: str, subject: str, fix: str, commands: str = "the server's absolute joint targets", wrong_fix: str = _DROID_ARM_WRONG_FIX) -> None

:func:check_joint_pass_through for a DROID arm term, which must also take DROID's joints in DROID's order.

That is one column per DROID arm joint, and DROID's own joint names in DROID's order (a rig with other names is checked by width only): the server reads row i as DROID's joint i.

Source code in eden/extensions/policies/adapters.py

def check_droid_arm_term(
    arm,
    *,
    term_name: str,
    subject: str,
    fix: str,
    commands: str = "the server's absolute joint targets",
    wrong_fix: str = _DROID_ARM_WRONG_FIX,
) -> None:
    """:func:`check_joint_pass_through` for a DROID arm term, which must also take DROID's joints in DROID's order.

    That is one column per DROID arm joint, and DROID's own joint names in DROID's order (a rig with other names is
    checked by width only): the server reads row ``i`` as DROID's joint ``i``.
    """
    problems = _pass_through_problems(arm)
    if arm.action_dim != len(DROID_ARM_DOFS):
        problems.append(f"width {arm.action_dim} (DROID's arm has {len(DROID_ARM_DOFS)} joints)")
    dofs = list(arm.dofs_name)
    if set(dofs) <= set(DROID_ARM_DOFS) and dofs != [name for name in DROID_ARM_DOFS if name in dofs]:
        problems.append(f"dofs_name={dofs} (needs {list(DROID_ARM_DOFS)}, in that order)")
    _refuse_term_problems(
        arm, problems, term_name=term_name, subject=subject, fix=fix, commands=commands, wrong_fix=wrong_fix
    )

check_droid_joint_state

check_droid_joint_state(env: 'EnvBase', arm, *, group: str, subject: str, fix: str) -> None

Raise unless group is [DofsPos of arm's joints in its order, GripperState closedness].

A DROID client sends group[:7] as joint_position and group[7] as gripper_position, and the server replies commands for those joints in that order, which the adapter hands column by column to arm's dofs_name. observation_options() and action_options() take the arm entity and DOF names separately, so the two can differ without either failing to build. The joint angles must be absolute (no offset_from_default); fix names the factory calls that build matching terms.

Source code in eden/extensions/policies/adapters.py

def check_droid_joint_state(env: "EnvBase", arm, *, group: str, subject: str, fix: str) -> None:
    """Raise unless ``group`` is ``[DofsPos of arm's joints in its order, GripperState closedness]``.

    A DROID client sends ``group[:7]`` as ``joint_position`` and ``group[7]`` as ``gripper_position``, and the
    server replies commands for those joints in that order, which the adapter hands column by column to ``arm``'s
    ``dofs_name``. ``observation_options()`` and ``action_options()`` take the arm entity and DOF names separately,
    so the two can differ without either failing to build. The joint angles must be absolute (no
    ``offset_from_default``); ``fix`` names the factory calls that build matching terms.
    """
    from eden.managers.terms.observations.manipulation import GripperState
    from eden.managers.terms.observations.proprio import DofsPos

    terms = {name: term for g, name, term in env.observation_manager.iter_terms() if g == group}
    arm_dofs = list(arm.dofs_name)
    problems = []
    if list(terms) != ["joint_pos", "gripper"]:
        problems.append(f"it holds the terms {list(terms)}, not ['joint_pos', 'gripper'] in that order")
    else:
        joints, gripper = terms["joint_pos"], terms["gripper"]
        if not isinstance(joints, DofsPos):
            problems.append(f"'joint_pos' is a {type(joints).__name__}, not a DofsPos")
        else:
            read = dofs_pos_joint_names(env, joints)
            if (joints.entity_name, read) != (arm.entity_name, arm_dofs):
                problems.append(
                    f"'joint_pos' reads {joints.entity_name!r} {read} but the arm term drives {arm.entity_name!r} "
                    f"{arm_dofs}, so each joint would be sent another joint's command"
                )
            if joints.offset_from_default:
                problems.append("'joint_pos' has offset_from_default=True (the server needs absolute joint angles)")
        if not isinstance(gripper, GripperState) or gripper.normalize != "closedness":
            problems.append(
                f"'gripper' is a {type(gripper).__name__} with normalize={getattr(gripper, 'normalize', None)!r}, "
                "not a GripperState normalized to 'closedness' (DROID's gripper_position, 1 = closed)"
            )
    if problems:
        raise ValueError(
            f"{subject} sends {group}[:{len(arm_dofs)}] as DROID's joint_position and {group}[{len(arm_dofs)}] as its "
            "gripper_position, and the server replies commands for those joints in that order, but the env's "
            f"{group!r} group does not match: {'; '.join(problems)}. Fix: {fix}. Do not reorder the server's rows in "
            "the client instead: the state it is sent and the commands it returns share one joint order."
        )

Transports, serializers and errors

Request/reply transports to an inference (or any request/reply) server.

A transport moves a payload to a server and returns the decoded reply. It is built from an explicit :class:TransportConfig (address, timeout, endpoints, serializer key and transport-specific knobs) and selected by TRANSPORT_REGISTRY key, so nothing here depends on who the caller is. Optional dependencies are imported lazily.

Every transport maps failures onto :mod:eden.utils.policy_client.errors: unreachable server or timeout -> :class:TransportError (a ConnectionError), server-side failure -> :class:ServerError (a RuntimeError).

==================== ============================================= ================== Key Wire Default serializer ==================== ============================================= ================== http_json requests POST, JSON body / JSON reply (strict JSON) http_json_numpy POST with a json_numpy text body json_numpy http_multipart POST multipart form; payload is (data, files) none openpi_websocket WebSocket, msgpack + OpenPI numpy hook msgpack_openpi galaxea_websocket as above, str-keyed records, __reset__ msg msgpack_galaxea zmq_msgpack ZMQ REQ/REP {"endpoint", "data"} envelope msgpack_numpy ==================== ============================================= ==================

Attributes

TRANSPORT_REGISTRY module-attribute

TRANSPORT_REGISTRY = Registry('TRANSPORT')

Classes

TransportConfig dataclass

TransportConfig(host: str = 'localhost', port: int | None = None, request_timeout: float = 30.0, act_endpoint: str = '/act', set_task_endpoint: str | None = None, reset_endpoint: str | None = None, set_task_payload_key: str = 'task_description', reset_payload_key: str = 'task_description', task_description: str | None = None, serializer: str | None = None, knobs: dict[str, Any] = dict())

Everything a transport needs to know, independent of the caller.

PARAMETER	DESCRIPTION
host	Server address, joined by :func:server_url: host is a bare hostname or a full URL (wss://example.com/ws, http://node:8000). A port embedded in the URL and port must agree; pass port=None to use the URL's. TYPE: str DEFAULT: 'localhost'
port	Server address, joined by :func:server_url: host is a bare hostname or a full URL (wss://example.com/ws, http://node:8000). A port embedded in the URL and port must agree; pass port=None to use the URL's. TYPE: str DEFAULT: 'localhost'
request_timeout	Request timeout in seconds (HTTP transports). TYPE: float DEFAULT: 30.0
act_endpoint	HTTP route for inference, starting with /. TYPE: str DEFAULT: '/act'
set_task_endpoint	The server's set-task / reset API: an HTTP route, a ZMQ endpoint name or a GalaxeaVLA reset message key. None (the default) means the server has none (it takes the task with each request), so set_task / reset send nothing. A transport whose protocol has no such message rejects a non-None value. TYPE: str | None DEFAULT: None
reset_endpoint	The server's set-task / reset API: an HTTP route, a ZMQ endpoint name or a GalaxeaVLA reset message key. None (the default) means the server has none (it takes the task with each request), so set_task / reset send nothing. A transport whose protocol has no such message rejects a non-None value. TYPE: str | None DEFAULT: None
set_task_payload_key	JSON keys carrying the task description (HTTP transports). TYPE: str DEFAULT: 'task_description'
reset_payload_key	JSON keys carrying the task description (HTTP transports). TYPE: str DEFAULT: 'task_description'
task_description	Initial task text; :meth:TransportBase.set_task updates it. TYPE: str | None DEFAULT: None
serializer	SERIALIZER_REGISTRY key; None takes the transport's default. TYPE: str | None DEFAULT: None
knobs	Transport-specific settings; each transport declares the ones it reads (with defaults) in its knobs class attribute. TYPE: dict[str, Any] DEFAULT: dict()
transport
ignoring

Attributes

host class-attribute instance-attribute

host: str = 'localhost'

port class-attribute instance-attribute

port: int | None = None

request_timeout class-attribute instance-attribute

request_timeout: float = 30.0

act_endpoint class-attribute instance-attribute

act_endpoint: str = '/act'

set_task_endpoint class-attribute instance-attribute

set_task_endpoint: str | None = None

reset_endpoint class-attribute instance-attribute

reset_endpoint: str | None = None

set_task_payload_key class-attribute instance-attribute

set_task_payload_key: str = 'task_description'

reset_payload_key class-attribute instance-attribute

reset_payload_key: str = 'task_description'

task_description class-attribute instance-attribute

task_description: str | None = None

serializer class-attribute instance-attribute

serializer: str | None = None

knobs class-attribute instance-attribute

knobs: dict[str, Any] = field(default_factory=dict)

Methods:

TransportBase

TransportBase(config: TransportConfig)

Bases: ABC

Request/reply channel to a server.

Subclasses declare in knobs the optional settings they read from config.knobs (with defaults), so a transport's configuration surface is visible in one place.

Source code in eden/utils/policy_client/transport.py

def __init__(self, config: TransportConfig) -> None:
    unknown = set(config.knobs) - set(self.knobs)
    if unknown:
        raise ValueError(
            f"{type(self).__name__} does not accept knobs {sorted(unknown)}; known: {sorted(self.knobs)}"
        )
    defaults = {f.name: f.default for f in fields(TransportConfig)}
    for name in sorted(_OPTIONAL_CONFIG_FIELDS - self.config_fields):
        value = getattr(config, name)
        if value != defaults[name]:
            raise ValueError(
                f"{type(self).__name__} does not read {name} (got {value!r}), so it would have no effect. "
                f"Leave it at its default {defaults[name]!r}; this transport's own settings are knobs "
                f"{sorted(self.knobs)}."
            )
    for name, supported in (
        ("reset_endpoint", self.supports_reset_endpoint),
        ("set_task_endpoint", self.supports_set_task_endpoint),
    ):
        if not supported and getattr(config, name) is not None:
            raise ValueError(
                f"{type(self).__name__} has no {name.removesuffix('_endpoint')} message; its servers take "
                f"the task with each request. Set {name}=None (got {getattr(config, name)!r})."
            )
    if self.default_serializer is None and config.serializer is not None:
        raise ValueError(
            f"{type(self).__name__} builds its own request body and takes no serializer. "
            f"Set serializer=None (got {config.serializer!r}), or pick a transport that encodes with one."
        )
    self.config = config
    key = config.serializer if config.serializer is not None else self.default_serializer
    self.serializer: SerializerBase | None = make_serializer(key) if key is not None else None
    self.metadata: dict = {}

Attributes

default_serializer class-attribute

default_serializer: str | None = None

knobs class-attribute

knobs: dict[str, Any] = {}

supports_reset_endpoint class-attribute

supports_reset_endpoint: bool = False

supports_set_task_endpoint class-attribute

supports_set_task_endpoint: bool = False

config_fields class-attribute

config_fields: frozenset[str] = frozenset()

config instance-attribute

config = config

serializer instance-attribute

serializer: SerializerBase | None = make_serializer(key) if key is not None else None

metadata instance-attribute

metadata: dict = {}

Methods:

knob

knob(name: str) -> Any

Read a declared knob, falling back to its default.

Source code in eden/utils/policy_client/transport.py

def knob(self, name: str) -> Any:
    """Read a declared knob, falling back to its default."""
    if name not in self.knobs:
        raise KeyError(f"{type(self).__name__} does not declare knob {name!r}")
    return self.config.knobs.get(name, self.knobs[name])

infer abstractmethod

infer(payload: Any, options: dict | None = None) -> Any

Send one payload and return the decoded reply.

Source code in eden/utils/policy_client/transport.py

@abstractmethod
def infer(self, payload: Any, options: dict | None = None) -> Any:
    """Send one payload and return the decoded reply."""

set_task

set_task(task_description: str) -> None

Record a new task text; servers with a set_task route also get told.

Source code in eden/utils/policy_client/transport.py

def set_task(self, task_description: str) -> None:
    """Record a new task text; servers with a set_task route also get told."""
    self.config.task_description = task_description

reset

reset(options: dict | None = None) -> None

Reset server-side episode state through config.reset_endpoint; default: no such API.

Source code in eden/utils/policy_client/transport.py

def reset(self, options: dict | None = None) -> None:
    """Reset server-side episode state through ``config.reset_endpoint``; default: no such API."""

server_config

server_config() -> dict

Whatever the server advertises about itself (handshake metadata by default).

Source code in eden/utils/policy_client/transport.py

def server_config(self) -> dict:
    """Whatever the server advertises about itself (handshake metadata by default)."""
    return dict(self.metadata)

close

close() -> None

Release the connection.

Source code in eden/utils/policy_client/transport.py

def close(self) -> None:
    """Release the connection."""

HttpTransportBase

HttpTransportBase(config: TransportConfig)

Bases: TransportBase

Shared requests.Session plumbing and error mapping for the HTTP transports.

knobs: set_task_as_query (send the task as a URL query parameter instead of a JSON body; Octo) and reset_with_payload (send {reset_payload_key: task_description} on reset, else a bare POST).

Source code in eden/utils/policy_client/transport.py

def __init__(self, config: TransportConfig) -> None:
    super().__init__(config)
    import requests

    for name in ("act_endpoint", "set_task_endpoint", "reset_endpoint"):
        route = getattr(config, name)
        if route is not None and not (isinstance(route, str) and route.startswith("/")):
            raise ValueError(
                f"{type(self).__name__} {name}={route!r} is not an HTTP route: it is appended to the "
                f"server URL, so it must start with '/'. Set {name}={'/' + str(route).lstrip('/')!r}."
            )
    self.session = requests.Session()
    self.base_url = server_url(config.host, config.port, schemes=("http", "https")).rstrip("/")

Attributes

knobs class-attribute instance-attribute

knobs = {'set_task_as_query': False, 'reset_with_payload': True}

supports_reset_endpoint class-attribute instance-attribute

supports_reset_endpoint = True

supports_set_task_endpoint class-attribute instance-attribute

supports_set_task_endpoint = True

config_fields class-attribute instance-attribute

config_fields = _OPTIONAL_CONFIG_FIELDS

session instance-attribute

session = requests.Session()

base_url instance-attribute

base_url = server_url(config.host, config.port, schemes=('http', 'https')).rstrip('/')

Methods:

post

post(endpoint: str, **kwargs)

POST to endpoint and return the raw response, mapping request failures.

Source code in eden/utils/policy_client/transport.py

def post(self, endpoint: str, **kwargs):
    """POST to ``endpoint`` and return the raw response, mapping request failures."""
    import requests

    url = f"{self.base_url}{endpoint}"
    timeout = self.config.request_timeout
    try:
        response = self.session.post(url, timeout=timeout, **kwargs)
        response.raise_for_status()
        return response
    except requests.ConnectTimeout as e:
        raise TransportTimeoutError(
            f"Connecting to policy server at {self.base_url} timed out after request_timeout={timeout} s. "
            "Is it running, and is that address reachable from this machine?"
        ) from e
    except requests.Timeout as e:  # connected, but no reply in time (ReadTimeout)
        raise TransportTimeoutError(
            f"Policy server at {url} did not answer within request_timeout={timeout} s. If the server is "
            "just slow on this request (e.g. compiling on its first call), raise request_timeout."
        ) from e
    except requests.ConnectionError as e:
        raise TransportError(f"Cannot reach policy server at {self.base_url}. Is it running?") from e
    except requests.HTTPError as e:
        body = e.response.text[:2000] if e.response is not None else ""
        raise ServerError(f"Policy server request to {url} failed ({e}):\n{body}") from e
    except requests.RequestException as e:
        raise TransportError(f"Policy server request to {url} failed: {e}") from e

post_json

post_json(endpoint: str, payload: dict | None = None) -> dict

POST a JSON body (none when payload is None) and return the parsed JSON reply.

Source code in eden/utils/policy_client/transport.py

def post_json(self, endpoint: str, payload: dict | None = None) -> dict:
    """POST a JSON body (none when ``payload`` is ``None``) and return the parsed JSON reply."""
    if payload is None:
        return self._json(self.post(endpoint))
    body = self._json_body(payload, endpoint)
    return self._json(self.post(endpoint, data=body, headers={"Content-Type": "application/json"}))

set_task

set_task(task_description: str) -> None

Source code in eden/utils/policy_client/transport.py

def set_task(self, task_description: str) -> None:
    super().set_task(task_description)
    endpoint = self.config.set_task_endpoint
    if endpoint is None:
        return
    if self.knob("set_task_as_query"):
        self.post(endpoint, params={self.config.set_task_payload_key: task_description})
    else:
        self.post_json(endpoint, {self.config.set_task_payload_key: task_description})

reset

reset(options: dict | None = None) -> None

Source code in eden/utils/policy_client/transport.py

def reset(self, options: dict | None = None) -> None:
    endpoint = self.config.reset_endpoint
    if endpoint is None:
        return
    if self.knob("reset_with_payload"):
        self.post_json(endpoint, {self.config.reset_payload_key: self.config.task_description})
    else:
        self.post(endpoint)

close

close() -> None

Source code in eden/utils/policy_client/transport.py

def close(self) -> None:
    self.session.close()

HttpJsonTransport

HttpJsonTransport(config: TransportConfig)

Bases: HttpTransportBase

JSON body in, JSON reply out (RobotArena-style servers: open-pi-zero, Octo).

Source code in eden/utils/policy_client/transport.py

def __init__(self, config: TransportConfig) -> None:
    super().__init__(config)
    import requests

    for name in ("act_endpoint", "set_task_endpoint", "reset_endpoint"):
        route = getattr(config, name)
        if route is not None and not (isinstance(route, str) and route.startswith("/")):
            raise ValueError(
                f"{type(self).__name__} {name}={route!r} is not an HTTP route: it is appended to the "
                f"server URL, so it must start with '/'. Set {name}={'/' + str(route).lstrip('/')!r}."
            )
    self.session = requests.Session()
    self.base_url = server_url(config.host, config.port, schemes=("http", "https")).rstrip("/")

Methods:

infer

infer(payload: Any, options: dict | None = None) -> Any

Source code in eden/utils/policy_client/transport.py

def infer(self, payload: Any, options: dict | None = None) -> Any:
    return self.post_json(self.config.act_endpoint, payload)

HttpJsonNumpyTransport

HttpJsonNumpyTransport(config: TransportConfig)

Bases: HttpTransportBase

json_numpy-encoded text body and reply (MolmoAct2's host_server_droid.py).

Source code in eden/utils/policy_client/transport.py

def __init__(self, config: TransportConfig) -> None:
    super().__init__(config)
    import requests

    for name in ("act_endpoint", "set_task_endpoint", "reset_endpoint"):
        route = getattr(config, name)
        if route is not None and not (isinstance(route, str) and route.startswith("/")):
            raise ValueError(
                f"{type(self).__name__} {name}={route!r} is not an HTTP route: it is appended to the "
                f"server URL, so it must start with '/'. Set {name}={'/' + str(route).lstrip('/')!r}."
            )
    self.session = requests.Session()
    self.base_url = server_url(config.host, config.port, schemes=("http", "https")).rstrip("/")

Attributes

default_serializer class-attribute instance-attribute

default_serializer = 'json_numpy'

Methods:

infer

infer(payload: Any, options: dict | None = None) -> Any

Source code in eden/utils/policy_client/transport.py

def infer(self, payload: Any, options: dict | None = None) -> Any:
    body = self.serializer.encode(payload)
    response = self.post(self.config.act_endpoint, data=body, headers={"Content-Type": "application/json"})
    result = self.serializer.decode(response.text)
    if isinstance(result, dict) and "error" in result:
        raise ServerError(f"Policy server returned error: {result['error']}")
    return result

HttpMultipartTransport

HttpMultipartTransport(config: TransportConfig)

Bases: HttpTransportBase

Multipart form upload; the payload is a (data, files) pair as requests expects.

Source code in eden/utils/policy_client/transport.py

def __init__(self, config: TransportConfig) -> None:
    super().__init__(config)
    import requests

    for name in ("act_endpoint", "set_task_endpoint", "reset_endpoint"):
        route = getattr(config, name)
        if route is not None and not (isinstance(route, str) and route.startswith("/")):
            raise ValueError(
                f"{type(self).__name__} {name}={route!r} is not an HTTP route: it is appended to the "
                f"server URL, so it must start with '/'. Set {name}={'/' + str(route).lstrip('/')!r}."
            )
    self.session = requests.Session()
    self.base_url = server_url(config.host, config.port, schemes=("http", "https")).rstrip("/")

Methods:

infer

infer(payload: Any, options: dict | None = None) -> Any

Source code in eden/utils/policy_client/transport.py

def infer(self, payload: Any, options: dict | None = None) -> Any:
    try:
        data, files = payload
    except (TypeError, ValueError) as e:
        raise TypeError("http_multipart expects the payload as a (data, files) tuple") from e
    return self._json(self.post(self.config.act_endpoint, data=data, files=files))

OpenPiWebsocketTransport

OpenPiWebsocketTransport(config: TransportConfig)

Bases: TransportBase

Port of openpi_client.websocket_client_policy.WebsocketClientPolicy (no openpi_client needed).

Connects to ws://host:port (host may be a full ws:///wss:// URL, see :func:server_url), reads the server's metadata handshake, then exchanges one msgpack message per inference. A str reply is the server reporting an error.

knobs:

• api_key: sent as an Authorization: Api-Key … header.
• ws_ping_timeout: keepalive pings every 20 s, each answered within this many seconds, as the OpenPI client does by default; a dead server then closes the connection instead of blocking :meth:infer (which has no timeout of its own) forever. None disables them: a server whose inference blocks its event loop (slow JIT compilation, diffusion sampling) cannot answer pings in time, and the client would drop a live connection.
• wait_for_server_s: a refused connection is retried every 5 s, as the OpenPI client waits for its server to come up, for at most this many seconds (None: forever).

Source code in eden/utils/policy_client/transport.py

def __init__(self, config: TransportConfig) -> None:
    super().__init__(config)
    require_module("websockets")  # a missing policy-client extra fails here, naming it
    from websockets import exceptions as ws_exc
    from websockets.sync import client as ws_sync

    self._ws_exc, self._ws_sync = ws_exc, ws_sync
    self.uri = server_url(config.host, config.port, schemes=("ws", "wss"))
    self._conn = None
    self.connect()

Attributes

default_serializer class-attribute instance-attribute

default_serializer = 'msgpack_openpi'

knobs class-attribute instance-attribute

knobs = {'api_key': None, 'ws_ping_timeout': 20.0, 'wait_for_server_s': 600.0}

wait_for_server class-attribute

wait_for_server: bool = True

uri instance-attribute

uri = server_url(config.host, config.port, schemes=('ws', 'wss'))

Methods:

connect

connect() -> None

Source code in eden/utils/policy_client/transport.py

def connect(self) -> None:
    ws_exc, ws_sync = self._ws_exc, self._ws_sync
    wait = self.wait_for_server
    started = time.monotonic()
    while True:
        try:
            conn = ws_sync.connect(
                self.uri,
                compression=None,
                max_size=None,
                additional_headers=self._headers(),
                ping_interval=None if self.knob("ws_ping_timeout") is None else 20,
                ping_timeout=self.knob("ws_ping_timeout"),
                open_timeout=self._open_timeout(),
            )
        except ConnectionRefusedError as e:
            if not wait:
                raise TransportError(f"Cannot reach policy server at {self.uri}. Is it running?") from e
            limit = self.knob("wait_for_server_s")
            waited = time.monotonic() - started
            if limit is not None and waited >= limit:
                raise TransportError(
                    f"Policy server at {self.uri} still refuses connections after {waited:.0f} s of waiting. "
                    "Start the server or fix host/port; if it is still loading, wait longer with "
                    "knobs={'wait_for_server_s': <seconds>} (None waits forever)."
                ) from e
            _LOG.info("Waiting for policy server at %s (%.0f s so far) ...", self.uri, waited)
            time.sleep(5 if limit is None else min(5.0, limit - waited))
            continue
        except TimeoutError as e:  # TCP connect or handshake open_timeout
            raise TransportTimeoutError(f"Connecting to policy server at {self.uri} timed out.") from e
        except OSError as e:
            raise TransportError(f"Cannot reach policy server at {self.uri}. Is it running?") from e
        except ws_exc.InvalidStatus as e:
            # The server answered the handshake with an HTTP error (e.g. a rejected api_key).
            raise ServerError(f"Policy server at {self.uri} rejected the connection: {e}") from e
        except ws_exc.WebSocketException as e:
            raise TransportError(f"WebSocket handshake with {self.uri} failed: {e}") from e
        self._conn = conn
        self.metadata = self._recv()
        return

infer

infer(payload: Any, options: dict | None = None) -> Any

Source code in eden/utils/policy_client/transport.py

def infer(self, payload: Any, options: dict | None = None) -> Any:
    self._send(payload)
    return self._recv()

close

close() -> None

Source code in eden/utils/policy_client/transport.py

def close(self) -> None:
    if self._conn is not None:
        try:
            self._conn.close()
        finally:
            self._conn = None

GalaxeaWebsocketTransport

GalaxeaWebsocketTransport(config: TransportConfig)

Bases: OpenPiWebsocketTransport

GalaxeaVLA scripts/serve_policy.py: OpenPI-style socket with str-keyed records and a reset message.

Always sends the Authorization header (empty key allowed), does not retry on refused connections, and clears the server's chunk cache with {reset_endpoint: True} (serve_policy.py's "__reset__") on both reset and set_task; with reset_endpoint=None neither sends anything. A dict reply carrying error is a server failure.

knobs: api_key, ws_ping_timeout, ws_timeout_s (connection open timeout).

Source code in eden/utils/policy_client/transport.py

def __init__(self, config: TransportConfig) -> None:
    super().__init__(config)
    require_module("websockets")  # a missing policy-client extra fails here, naming it
    from websockets import exceptions as ws_exc
    from websockets.sync import client as ws_sync

    self._ws_exc, self._ws_sync = ws_exc, ws_sync
    self.uri = server_url(config.host, config.port, schemes=("ws", "wss"))
    self._conn = None
    self.connect()

Attributes

default_serializer class-attribute instance-attribute

default_serializer = 'msgpack_galaxea'

knobs class-attribute instance-attribute

knobs = {'api_key': None, 'ws_ping_timeout': None, 'ws_timeout_s': 60.0}

wait_for_server class-attribute instance-attribute

wait_for_server = False

supports_reset_endpoint class-attribute instance-attribute

supports_reset_endpoint = True

Methods:

set_task

set_task(task_description: str) -> None

Source code in eden/utils/policy_client/transport.py

def set_task(self, task_description: str) -> None:
    TransportBase.set_task(self, task_description)
    self.reset()

reset

reset(options: dict | None = None) -> None

Source code in eden/utils/policy_client/transport.py

def reset(self, options: dict | None = None) -> None:
    key = self.config.reset_endpoint
    if key is None:
        return
    self._send({key: True})
    # serve_policy.py acknowledges with {"__reset__": True}; errors come back as {"error": ...}.
    reply = self._recv()
    if not (isinstance(reply, dict) and reply.get(key)):
        raise ServerError(f"GalaxeaVLA server did not acknowledge {key}; got {reply!r}")

ZmqMsgpackTransport

ZmqMsgpackTransport(config: TransportConfig)

Bases: TransportBase

ZMQ REQ/REP client for Isaac-GR00T's PolicyServer and its derivatives (RLDX-1).

Requests are {"endpoint": name, "data": {...}, "api_token"?: str}; get_action returns [action_dict, info] and reset a dict. A dict reply with error is a server failure. After a timeout the REQ socket is wedged, so it is recreated before the timeout is re-raised. Select the serializer per server: msgpack_numpy (GR00T, default) or msgpack_npy_blob (RLDX-1).

knobs: timeout_ms (send/receive timeout) and api_token. config.reset_endpoint names the server's reset endpoint ("reset" for GR00T / RLDX-1); None sends nothing.

Source code in eden/utils/policy_client/transport.py

def __init__(self, config: TransportConfig) -> None:
    super().__init__(config)
    self._zmq = require_module("zmq")
    self._context = self._zmq.Context.instance()
    self._socket = None
    self._init_socket()

Attributes

default_serializer class-attribute instance-attribute

default_serializer = 'msgpack_numpy'

knobs class-attribute instance-attribute

knobs = {'timeout_ms': 60000, 'api_token': None}

supports_reset_endpoint class-attribute instance-attribute

supports_reset_endpoint = True

address property

address: str

Methods:

call_endpoint

call_endpoint(endpoint: str, data: dict | None, requires_input: bool = True) -> Any

Round-trip one {"endpoint": …} request.

Source code in eden/utils/policy_client/transport.py

def call_endpoint(self, endpoint: str, data: dict | None, requires_input: bool = True) -> Any:
    """Round-trip one ``{"endpoint": …}`` request."""
    request: dict = {"endpoint": endpoint}
    if requires_input:
        request["data"] = data
    api_token = self.knob("api_token")
    if api_token:
        request["api_token"] = api_token
    try:
        self._socket.send(self.serializer.encode(request))
        message = self._socket.recv()
    except self._zmq.error.Again as e:
        # REQ socket is wedged after a timeout — recreate it before re-raising.
        self._socket.close(linger=0)
        self._init_socket()
        raise TransportTimeoutError(
            f"Policy server at {self.address} did not answer {endpoint!r} within timeout_ms="
            f"{self.knob('timeout_ms')}. Is it running at that address? If it is just slow on this request, "
            "raise it with knobs={'timeout_ms': <ms>}."
        ) from e
    response = self.serializer.decode(message)
    if isinstance(response, dict) and "error" in response:
        raise ServerError(f"Policy server error: {response['error']}")
    return response

infer

infer(payload: Any, options: dict | None = None) -> Any

Source code in eden/utils/policy_client/transport.py

def infer(self, payload: Any, options: dict | None = None) -> Any:
    response = self.call_endpoint("get_action", {"observation": payload, "options": options})
    # Server returns [action_dict, info_dict]; msgpack decodes to a list.
    if not (isinstance(response, (list, tuple)) and len(response) == 2 and isinstance(response[0], dict)):
        raise ServerError(
            f"Policy server at {self.address} get_action reply must be [action_dict, info]; "
            f"got {type(response).__name__} {str(response)[:200]}"
        )
    return tuple(response)

reset

reset(options: dict | None = None) -> None

Source code in eden/utils/policy_client/transport.py

def reset(self, options: dict | None = None) -> None:
    if self.config.reset_endpoint is None:
        return
    self.call_endpoint(self.config.reset_endpoint, {"options": options})

server_config

server_config() -> dict

The server's get_modality_config reply (GR00T family), or {} when it has none.

Source code in eden/utils/policy_client/transport.py

def server_config(self) -> dict:
    """The server's ``get_modality_config`` reply (GR00T family), or ``{}`` when it has none."""
    result = self.call_endpoint("get_modality_config", None, requires_input=False)
    return result if isinstance(result, dict) else {}

close

close() -> None

Source code in eden/utils/policy_client/transport.py

def close(self) -> None:
    if self._socket is not None:
        self._socket.close(linger=0)
        self._socket = None

Functions:

server_url

server_url(host: str, port: int | None, *, schemes: tuple[str, ...]) -> str

Join host and port into scheme://host[:port][/path].

host is either a bare hostname (localhost, 10.0.0.1, node:8000), which gets schemes[0], or a full URL whose scheme must be one of schemes. Its path is kept, so host="wss://example.com/ws", port=None addresses the server as given. A port in host and port must agree. Pass port=None to use the URL's port, or the scheme's default port when the URL has none. tcp (ZMQ) needs a port and no path.

Source code in eden/utils/policy_client/transport.py

def server_url(host: str, port: int | None, *, schemes: tuple[str, ...]) -> str:
    """Join ``host`` and ``port`` into ``scheme://host[:port][/path]``.

    ``host`` is either a bare hostname (``localhost``, ``10.0.0.1``, ``node:8000``), which gets
    ``schemes[0]``, or a full URL whose scheme must be one of ``schemes``. Its path is kept,
    so ``host="wss://example.com/ws", port=None`` addresses the server as given. A port in
    ``host`` and ``port`` must agree. Pass ``port=None`` to use the URL's port, or the
    scheme's default port when the URL has none. ``tcp`` (ZMQ) needs a port and no path.
    """
    parts = urlsplit(host if "://" in host else f"{schemes[0]}://{host}")
    scheme = parts.scheme.lower()
    if scheme not in schemes:
        raise ValueError(f"Policy server host {host!r} uses scheme {scheme!r}; this transport accepts {list(schemes)}.")
    if not parts.hostname:
        raise ValueError(f"Policy server host {host!r} has no hostname.")
    try:
        embedded = parts.port
    except ValueError as e:
        raise ValueError(f"Policy server host {host!r} has an invalid port.") from e
    if embedded is not None and port is not None and embedded != port:
        raise ValueError(
            f"Policy server host {host!r} embeds port {embedded}, but port={port}. Pass port=None to use the "
            "URL's port, or drop the port from the URL."
        )
    resolved = embedded if embedded is not None else port
    if scheme == "tcp" and (resolved is None or parts.path not in ("", "/")):
        raise ValueError(f"A tcp:// policy server address needs a port and no path; got host={host!r}, port={port}.")
    netloc = f"[{parts.hostname}]" if ":" in parts.hostname else parts.hostname
    if parts.username is not None:
        userinfo = parts.username if parts.password is None else f"{parts.username}:{parts.password}"
        netloc = f"{userinfo}@{netloc}"
    if resolved is not None:
        netloc = f"{netloc}:{resolved}"
    return urlunsplit((scheme, netloc, parts.path, parts.query, parts.fragment))

make_transport

make_transport(key: str, config: TransportConfig) -> TransportBase

Instantiate the transport registered under key; unknown keys list the available ones.

Source code in eden/utils/policy_client/transport.py

def make_transport(key: str, config: TransportConfig) -> TransportBase:
    """Instantiate the transport registered under ``key``; unknown keys list the available ones."""
    if key not in TRANSPORT_REGISTRY.keys():
        raise ValueError(f"Unknown transport {key!r}. Available: {sorted(TRANSPORT_REGISTRY.keys())}")
    return TRANSPORT_REGISTRY.get(key)(config)

Wire serialisers for request/reply clients.

A serialiser turns a payload dict into the bytes (or text) a server expects and decodes the reply. They are selected by SERIALIZER_REGISTRY key (TransportConfig.serializer) so the choice is a plain string that survives config serialisation. Each one imports its optional dependency lazily and reports the extra to install when it is missing.

===================== ========================================================== Key Wire format ===================== ========================================================== json json.dumps / json.loads json_numpy json_numpy (numpy arrays as base64 inside JSON text) msgpack_openpi msgpack with OpenPI's numpy hook (b"__ndarray__" records) msgpack_galaxea same record layout with str keys, as the GalaxeaVLA server speaks msgpack_numpy msgpack_numpy (Isaac-GR00T PolicyServer) msgpack_npy_blob msgpack with arrays tunnelled as np.save blobs (RLDX-1) ===================== ==========================================================

Attributes

SERIALIZER_REGISTRY module-attribute

SERIALIZER_REGISTRY = Registry('SERIALIZER')

Classes

SerializerBase

Bases: ABC

Encode a payload for the wire and decode a server reply.

Subclasses' encode / decode are wrapped so codec failures surface as :class:SerializationError.

Methods:

encode abstractmethod

encode(data: Any) -> bytes | str

Serialise data (dict of str / numbers / numpy arrays) for transmission.

Source code in eden/utils/policy_client/serializer.py

@abstractmethod
def encode(self, data: Any) -> bytes | str:
    """Serialise ``data`` (dict of str / numbers / numpy arrays) for transmission."""

decode abstractmethod

decode(raw: bytes | str) -> Any

Deserialise a server reply.

Source code in eden/utils/policy_client/serializer.py

@abstractmethod
def decode(self, raw: bytes | str) -> Any:
    """Deserialise a server reply."""

JsonSerializer

Bases: SerializerBase

Plain JSON text; numpy arrays must already be converted to lists by the policy.

Methods:

encode

encode(data: Any) -> str

Source code in eden/utils/policy_client/serializer.py

def encode(self, data: Any) -> str:
    return json.dumps(data)

decode

decode(raw: bytes | str) -> Any

Source code in eden/utils/policy_client/serializer.py

def decode(self, raw: bytes | str) -> Any:
    return json.loads(raw)

JsonNumpySerializer

JsonNumpySerializer()

Bases: SerializerBase

json_numpy text: numpy arrays are embedded as base64 records.

Source code in eden/utils/policy_client/serializer.py

def __init__(self) -> None:
    self._json_numpy = require_module("json_numpy")

Methods:

encode

encode(data: Any) -> str

Source code in eden/utils/policy_client/serializer.py

def encode(self, data: Any) -> str:
    return self._json_numpy.dumps(data)

decode

decode(raw: bytes | str) -> Any

Source code in eden/utils/policy_client/serializer.py

def decode(self, raw: bytes | str) -> Any:
    return self._json_numpy.loads(raw)

MsgpackOpenPiSerializer

MsgpackOpenPiSerializer()

Bases: SerializerBase

msgpack with OpenPI's numpy hook, byte-for-byte what openpi_client.msgpack_numpy produces.

Array records use bytes keys (b"__ndarray__") because that is what the OpenPI WebsocketPolicyServer looks for when unpacking. Bytes keys survive msgpack's raw=False decoding unchanged (they are packed as bin), so the same hook decodes the server's replies.

Source code in eden/utils/policy_client/serializer.py

def __init__(self) -> None:
    self._msgpack = require_module("msgpack")

Attributes

binary_keys class-attribute instance-attribute

binary_keys: bool = True

Methods:

encode

encode(data: Any) -> bytes

Source code in eden/utils/policy_client/serializer.py

def encode(self, data: Any) -> bytes:
    return self._msgpack.packb(data, default=self._pack_default)

decode

decode(raw: bytes | str) -> Any

Source code in eden/utils/policy_client/serializer.py

def decode(self, raw: bytes | str) -> Any:
    return self._msgpack.unpackb(raw, object_hook=self._object_hook)

MsgpackGalaxeaSerializer

MsgpackGalaxeaSerializer()

Bases: MsgpackOpenPiSerializer

OpenPI record layout with str keys, matching the GalaxeaVLA server's codec.

Replies are normalised recursively so nested dicts come back with str keys.

Source code in eden/utils/policy_client/serializer.py

def __init__(self) -> None:
    self._msgpack = require_module("msgpack")

Attributes

binary_keys class-attribute instance-attribute

binary_keys = False

MsgpackNumpySerializer

MsgpackNumpySerializer()

Bases: SerializerBase

msgpack_numpy as used by Isaac-GR00T's ZMQ PolicyServer.

GR00T's MsgSerializer additionally tunnels ModalityConfig objects as {"__ModalityConfig__": True, "as_json": ...} records (the get_modality_config reply), with as_json a dict or its JSON text. They are decoded to the plain dict so clients can read modality_keys / delta_indices without importing gr00t.

Source code in eden/utils/policy_client/serializer.py

def __init__(self) -> None:
    self._mnp = require_module("msgpack_numpy")

Methods:

encode

encode(data: Any) -> bytes

Source code in eden/utils/policy_client/serializer.py

def encode(self, data: Any) -> bytes:
    return self._mnp.packb(data)

decode

decode(raw: bytes | str) -> Any

Source code in eden/utils/policy_client/serializer.py

def decode(self, raw: bytes | str) -> Any:
    return self._mnp.unpackb(raw, object_hook=self._decode_hook, raw=False)

MsgpackNpyBlobSerializer

MsgpackNpyBlobSerializer()

Bases: SerializerBase

Client-side port of RLDX-1's MsgSerializer: arrays travel as np.save byte blobs.

Torch tensors are moved to CPU (half precision widened to float32) and encoded like arrays. RLDX tunnels ModalityConfig objects (the get_modality_config reply) as {"__ModalityConfig_class__": True, "as_json": ...} records; they are decoded to the plain dict (from a dict or its JSON text), as :class:MsgpackNumpySerializer does for GR00T's tag.

Source code in eden/utils/policy_client/serializer.py

def __init__(self) -> None:
    self._msgpack = require_module("msgpack")

Methods:

encode

encode(data: Any) -> bytes

Source code in eden/utils/policy_client/serializer.py

def encode(self, data: Any) -> bytes:
    return self._msgpack.packb(data, default=self._encode_default)

decode

decode(raw: bytes | str) -> Any

Source code in eden/utils/policy_client/serializer.py

def decode(self, raw: bytes | str) -> Any:
    return self._msgpack.unpackb(raw, object_hook=self._decode_hook, raw=False)

Functions:

require_module

require_module(module: str, extra: str = 'policy-client')

Import an optional client dependency, or raise naming the extra that provides it.

Source code in eden/utils/policy_client/serializer.py

def require_module(module: str, extra: str = "policy-client"):
    """Import an optional client dependency, or raise naming the extra that provides it."""
    try:
        return __import__(module)
    except ImportError as e:  # pragma: no cover - depends on the install
        raise ImportError(f"The '{module}' package is required here. Install it with {install_hint(extra)}.") from e

make_serializer

make_serializer(key: str) -> SerializerBase

Instantiate the registered serialiser key; unknown keys list the available ones.

Source code in eden/utils/policy_client/serializer.py

def make_serializer(key: str) -> SerializerBase:
    """Instantiate the registered serialiser ``key``; unknown keys list the available ones."""
    if key not in SERIALIZER_REGISTRY.keys():
        raise ValueError(f"Unknown serializer {key!r}. Available: {sorted(SERIALIZER_REGISTRY.keys())}")
    return SERIALIZER_REGISTRY.get(key)()

Error hierarchy for request/reply clients (eden.utils.policy_client).

TransportError also derives from :class:ConnectionError and every error derives from :class:RuntimeError, so callers that already distinguish "server unreachable" from "protocol failure" with those builtins keep working.

::

ClientError (RuntimeError)
├── TransportError (…, ConnectionError)   server unreachable / connection dropped
│   └── TransportTimeoutError             request timed out
├── ServerError                           server answered with an error (HTTP 4xx/5xx, {"error": …})
└── SerializationError                    payload or reply could not be (de)serialised

Classes

ClientError

Bases: RuntimeError

Base class for every remote-policy client failure.

TransportError

Bases: ClientError, ConnectionError

The policy server could not be reached or the connection was lost.

TransportTimeoutError

Bases: TransportError

A request to the policy server timed out.

ServerError

Bases: ClientError

The policy server answered, but with an error.

SerializationError

Bases: ClientError

A payload or reply could not be encoded or decoded.

Local policies

Locomotion

HolosomaLocomotionPolicy

HolosomaLocomotionPolicy(env: EnvBase, options: PolicyOptions)

Bases: PolicyBase

Holosoma velocity-command locomotion controller for G1 29-DOF (ONNX backend).

Reads proprioception directly from the env each step and drives the frozen ONNX policy with the current 3-D velocity command (see :meth:set_command). Only the Unitree G1 29-DOF robot is supported; G1Dofs29.dofs_name order matches Holosoma's joint order, so no DOF remapping is needed.

Load the ONNX graph, cache DOF indices, and allocate command/phase buffers.

Source code in eden/extensions/policies/holosoma/policy.py

def __init__(self, env: EnvBase, options: PolicyOptions) -> None:
    """Load the ONNX graph, cache DOF indices, and allocate command/phase buffers."""
    super().__init__(env, options)

    device = env.device
    self.device = device
    self._dtype = torch.float32
    N = env.num_envs

    # Single-IO ONNX graph (actor_obs -> action). OnnxModule validates the
    # static input/output dims against these (no-op for a dynamic axis),
    # raising on a swapped / re-exported graph.
    self._policy = OnnxModule(
        self.checkpoint_path, "actor_obs", "action", device=str(device),
        label="Holosoma Loco ONNX", expected_input_dim=_OBS_DIM, expected_output_dim=NUM_DOFS,
    )  # fmt: skip

    robot = self.require_entity(self.entity_name)
    # All 29 DOFs in Holosoma / G1Dofs29 declaration order.
    _, dofs_idx = robot.find_named_dofs_idx_local(G1_DOFS29_NAME, preserve_order=True)
    self._dofs_idx = torch.as_tensor(dofs_idx, dtype=torch.long, device=device)

    self._default_dof_pos = torch.tensor(G1_LOCO_DEFAULT_DOF_POS, dtype=self._dtype, device=device)

    # Current 3-D command [vel_x, vel_y, ang_vel_z]; zero → standing.
    self._commands = torch.zeros(N, 3, device=device, dtype=self._dtype)
    # Previous raw policy output, fed back as the `actions` observation.
    self._last_action = torch.zeros(N, NUM_DOFS, device=device, dtype=self._dtype)

    # Gait phase clock (left foot at 0, right at π) and per-env standing flag.
    self._phase = torch.zeros(N, 2, device=device, dtype=self._dtype)
    self._phase[:, 1] = torch.pi
    self._is_standing = torch.zeros(N, dtype=torch.bool, device=device)
    self._phase_dt = 2.0 * torch.pi / (self.rl_rate * self.gait_period)

    # World-frame "down" vector, expanded once to the env batch.
    self._gravity_world = torch.tensor([0.0, 0.0, -1.0], device=device, dtype=self._dtype).expand(N, 3)

Attributes

entity_name class-attribute instance-attribute

entity_name: str = 'robot'

checkpoint_path class-attribute instance-attribute

checkpoint_path: str | None = 'holosoma/ppo_g1_29dof.onnx'

md5_hash class-attribute instance-attribute

md5_hash: str | None = 'cfe906f8c5a87b5c4b367b1241e8b217'

policy_action_scale class-attribute instance-attribute

policy_action_scale: float = 0.25

rl_rate class-attribute instance-attribute

rl_rate: float = 50.0

gait_period class-attribute instance-attribute

gait_period: float = 1.0

device instance-attribute

device = device

Methods:

set_command

set_command(command: Tensor) -> None

Set the 3-D velocity command for all (or a slice of) environments.

PARAMETER	DESCRIPTION
command	Either (3,) (broadcast to all envs) or (num_envs, 3). Layout is [vel_x, vel_y, ang_vel_z]. A ~zero command makes the policy stand (gait phase pins to π). TYPE: Tensor

Source code in eden/extensions/policies/holosoma/policy.py

def set_command(self, command: torch.Tensor) -> None:
    """Set the 3-D velocity command for all (or a slice of) environments.

    Parameters
    ----------
    command : torch.Tensor
        Either ``(3,)`` (broadcast to all envs) or ``(num_envs, 3)``. Layout
        is ``[vel_x, vel_y, ang_vel_z]``. A ~zero command makes the policy
        stand (gait phase pins to π).
    """
    command = torch.as_tensor(command, dtype=self._dtype, device=self.device)
    if command.ndim == 1:
        command = command.unsqueeze(0).expand(self._commands.shape[0], -1)
    self._commands[:] = command

robot_options classmethod

robot_options() -> RobotOptions

Return a G1Dofs29 with Holosoma's training default joint pose.

Source code in eden/extensions/policies/holosoma/policy.py

@classmethod
def robot_options(cls) -> RobotOptions:
    """Return a ``G1Dofs29`` with Holosoma's training default joint pose."""
    return holosoma_g1_robot()

env_options classmethod

env_options() -> dict[str, Any]

Return env options matching Holosoma's training sim: 200 Hz physics, decimation 4 → 50 Hz control (== rl_rate). Holosoma trains all backends at fps=200, control_decimation=4, substeps=1 (holosoma.config_values.simulator), so reproduce that exact cadence rather than another (sim_dt, decimation) pair with the same 50 Hz product.

Source code in eden/extensions/policies/holosoma/policy.py

@classmethod
def env_options(cls) -> dict[str, Any]:
    """Return env options matching Holosoma's training sim: 200 Hz physics,
    decimation 4 → 50 Hz control (== ``rl_rate``). Holosoma trains all
    backends at ``fps=200, control_decimation=4, substeps=1``
    (``holosoma.config_values.simulator``), so reproduce that exact cadence
    rather than another (sim_dt, decimation) pair with the same 50 Hz product.
    """
    return dict(
        sim_dt=0.005,
        sim_substeps=1,
        decimation=4,
    )

action_options classmethod

action_options(entity_name: str = 'robot') -> dict[str, ActionTermOptions]

Configure a 29-DOF implicit PD controller consuming absolute targets.

:meth:step already produces absolute joint-position targets (default + scale * action), so the controller uses reference_source=ZERO and scale=1.0 (no extra offset / scaling).

Source code in eden/extensions/policies/holosoma/policy.py

@classmethod
def action_options(cls, entity_name: str = "robot") -> dict[str, ActionTermOptions]:
    """Configure a 29-DOF implicit PD controller consuming absolute targets.

    :meth:`step` already produces absolute joint-position targets
    (``default + scale * action``), so the controller uses
    ``reference_source=ZERO`` and ``scale=1.0`` (no extra offset / scaling).
    """
    return dict(
        dofs_pos_controller=actions.ImplicitPDController.configure(
            entity_name=entity_name,
            dofs_name=G1_DOFS29_NAME,
            scale=1.0,
            reference_source=ReferenceSource.ZERO,
        ),
    )

observation_options classmethod

observation_options(entity_name: str = 'robot') -> dict[str, ObservationGroupOptions]

No env-computed observations.

The policy builds its own 100-dim actor_obs from raw simulator state in :meth:step (the obs argument is ignored), so the env needs no observation terms — an empty group set builds and steps fine.

Source code in eden/extensions/policies/holosoma/policy.py

@classmethod
def observation_options(cls, entity_name: str = "robot") -> dict[str, ObservationGroupOptions]:
    """No env-computed observations.

    The policy builds its own 100-dim ``actor_obs`` from raw simulator state
    in :meth:`step` (the ``obs`` argument is ignored), so the env needs no
    observation terms — an empty group set builds and steps fine.
    """
    return {}

reset

reset(envs_idx: slice | Tensor | None = None) -> None

Reset prev-action and gait-phase state for the given environments.

PARAMETER	DESCRIPTION
envs_idx	Environments to reset; None resets all. TYPE: slice | Tensor | None DEFAULT: None

Source code in eden/extensions/policies/holosoma/policy.py

def reset(self, envs_idx: (slice | torch.Tensor) | None = None) -> None:
    """Reset prev-action and gait-phase state for the given environments.

    Parameters
    ----------
    envs_idx : slice | torch.Tensor | None, optional
        Environments to reset; ``None`` resets all.
    """
    if envs_idx is None:
        self._last_action.zero_()
        self._phase[:, 0] = 0.0
        self._phase[:, 1] = torch.pi
        self._is_standing.zero_()
    else:
        self._last_action[envs_idx] = 0.0
        self._phase[envs_idx, 0] = 0.0
        self._phase[envs_idx, 1] = torch.pi
        self._is_standing[envs_idx] = False

step

step(obs: dict[str, Tensor | dict[str, Tensor]]) -> dict[str, Tensor]

Advance the gait clock, run the policy, and return 29-DOF targets.

PARAMETER	DESCRIPTION
obs	Ignored — raw simulator state is read directly from the env. TYPE: dict

RETURNS	DESCRIPTION
dict[str, Tensor]	{"dofs_pos_controller": (N, 29)} absolute joint-position targets.

Source code in eden/extensions/policies/holosoma/policy.py

def step(self, obs: dict[str, torch.Tensor | dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]:
    """Advance the gait clock, run the policy, and return 29-DOF targets.

    Parameters
    ----------
    obs : dict
        Ignored — raw simulator state is read directly from the env.

    Returns
    -------
    dict[str, torch.Tensor]
        ``{"dofs_pos_controller": (N, 29)}`` absolute joint-position targets.
    """
    # Read off the entity, not an observation group: this policy rides along as an extra deploy
    # slot on another policy's or task's env, whose groups are not its own (see the
    # "Declare the observation" paragraph of docs/guides/policy-porting.md).
    del obs
    robot = self.env.entities[self.entity_name]

    # Phase advances once per control step, before observation construction
    # (matches LocomotionPolicy.run: update_phase_time() then policy_action()).
    self._update_phase()

    actor_obs = self._build_obs(robot)
    action = self._policy(actor_obs)  # raw policy output, (N, 29)
    self._last_action = action

    targets = self._default_dof_pos + self.policy_action_scale * action
    return {"dofs_pos_controller": targets}

HomiePolicy

HomiePolicy(env: EnvBase, options: PolicyOptions)

Bases: PolicyBase

OpenHomie frozen walk/stand tracker for G1 lower-body locomotion.

Reads proprioception directly from the env each step and drives the frozen tracker with the current 4-D command (see :meth:set_command). Only the Unitree G1 29-DOF robot is supported; G1Dofs29.dofs_name order matches HOMIE's joint order, so no DOF remapping is needed.

Build the tracker, cache DOF indices, and allocate command buffers.

Source code in eden/extensions/policies/homie/policy.py

def __init__(self, env: EnvBase, options: PolicyOptions) -> None:
    """Build the tracker, cache DOF indices, and allocate command buffers."""
    super().__init__(env, options)

    device = env.device
    self.device = device

    # Resolve the walk/stand checkpoints via the Hugging Face asset resolver
    # (downloads from Kashu7100/eden_policies on first use; passes through
    # existing local paths). Empty string disables that model.
    walk_path = self.fetch_asset(self.walk_model_path) if self.walk_model_path else ""
    stand_path = self.fetch_asset(self.stand_model_path) if self.stand_model_path else ""
    self._verify_md5(walk_path, self.walk_md5_hash) if walk_path else None
    self._verify_md5(stand_path, self.stand_md5_hash) if stand_path else None

    cfg = HomieConfig(
        walk_model_path=walk_path,
        stand_model_path=stand_path,
        command_scale=self.command_scale,
        walk_stand_threshold=self.walk_stand_threshold,
        clip_height_lower=self.clip_height_lower,
        clip_height_upper=self.clip_height_upper,
    )
    self._tracker = HomieTracker(config=cfg, num_envs=env.num_envs, device=device)

    robot = self.require_entity(self.entity_name)
    # All 29 DOFs in HOMIE / G1Dofs29 declaration order.
    _, dofs_idx = robot.find_named_dofs_idx_local(G1_DOFS29_NAME, preserve_order=True)
    self._dofs_idx = torch.as_tensor(dofs_idx, dtype=torch.long, device=device)

    # Arm joint targets are held at the HOMIE default pose (DOFs 15..28).
    self._arm_targets = (
        torch.tensor(HOMIE_DEFAULT_DOF_POS[HOMIE_NUM_ACTIONS:], dtype=torch.float32, device=device)
        .unsqueeze(0)
        .expand(env.num_envs, -1)
    )

    # Current 4-D command [vel_x, vel_y, ang_vel_z, height_delta]; settable
    # via set_command(). Defaults to zero → standing.
    self._commands = torch.zeros(env.num_envs, 4, device=device, dtype=torch.float32)

    # World-frame "down" vector, expanded to the env batch once.
    self._gravity_world = torch.tensor([0.0, 0.0, -1.0], device=device, dtype=torch.float32).expand(env.num_envs, 3)

Attributes

entity_name class-attribute instance-attribute

entity_name: str = 'robot'

checkpoint_path class-attribute instance-attribute

checkpoint_path: str | None = None

walk_model_path class-attribute instance-attribute

walk_model_path: str = 'homie/model_walk.pt'

stand_model_path class-attribute instance-attribute

stand_model_path: str = 'homie/model_stand.pt'

walk_md5_hash class-attribute instance-attribute

walk_md5_hash: str | None = '556dda04493332fdaf83385c93218169'

stand_md5_hash class-attribute instance-attribute

stand_md5_hash: str | None = '31edd14f666173d583218b5fac24e671'

command_scale class-attribute instance-attribute

command_scale: float = 1.0

walk_stand_threshold class-attribute instance-attribute

walk_stand_threshold: float = 0.1

clip_height_lower class-attribute instance-attribute

clip_height_lower: float = 0.24

clip_height_upper class-attribute instance-attribute

clip_height_upper: float = 0.74

device instance-attribute

device = device

Methods:

set_command

set_command(command: Tensor) -> None

Set the 4-D locomotion command for all (or a slice of) environments.

PARAMETER	DESCRIPTION
command	Either (4,) (broadcast to all envs) or (num_envs, 4). Layout is [vel_x, vel_y, ang_vel_z, height_delta] where height_delta is an offset about HOMIE_DEFAULT_HEIGHT (0 → 0.74 m). TYPE: Tensor

Source code in eden/extensions/policies/homie/policy.py

def set_command(self, command: torch.Tensor) -> None:
    """Set the 4-D locomotion command for all (or a slice of) environments.

    Parameters
    ----------
    command : torch.Tensor
        Either ``(4,)`` (broadcast to all envs) or ``(num_envs, 4)``. Layout
        is ``[vel_x, vel_y, ang_vel_z, height_delta]`` where ``height_delta``
        is an offset about ``HOMIE_DEFAULT_HEIGHT`` (0 → 0.74 m).
    """
    command = torch.as_tensor(command, dtype=torch.float32, device=self.device)
    if command.ndim == 1:
        command = command.unsqueeze(0).expand(self._commands.shape[0], -1)
    self._commands[:] = command

robot_options classmethod

robot_options() -> RobotOptions

Return a G1Dofs29 with HOMIE's default pose and PD gains baked in.

Source code in eden/extensions/policies/homie/policy.py

@classmethod
def robot_options(cls) -> RobotOptions:
    """Return a ``G1Dofs29`` with HOMIE's default pose and PD gains baked in."""
    return homie_g1_robot()

env_options classmethod

env_options() -> dict[str, Any]

Return env options for 50 Hz control (sim_dt=0.002, decimation=10).

Source code in eden/extensions/policies/homie/policy.py

@classmethod
def env_options(cls) -> dict[str, Any]:
    """Return env options for 50 Hz control (sim_dt=0.002, decimation=10)."""
    return dict(
        sim_dt=0.002,
        decimation=10,
    )

action_options classmethod

action_options(entity_name: str = 'robot') -> dict[str, ActionTermOptions]

Configure a 29-DOF implicit PD controller consuming absolute targets.

:meth:step already produces absolute joint-position targets (default + scale * action for the lower body, default for the arms), so the controller uses reference_source=ZERO and scale=1.0.

Source code in eden/extensions/policies/homie/policy.py

@classmethod
def action_options(cls, entity_name: str = "robot") -> dict[str, ActionTermOptions]:
    """Configure a 29-DOF implicit PD controller consuming absolute targets.

    :meth:`step` already produces absolute joint-position targets
    (default + scale * action for the lower body, default for the arms), so
    the controller uses ``reference_source=ZERO`` and ``scale=1.0``.
    """
    return dict(
        dofs_pos_controller=actions.ImplicitPDController.configure(
            entity_name=entity_name,
            dofs_name=G1_DOFS29_NAME,
            scale=1.0,
            reference_source=ReferenceSource.ZERO,
        ),
    )

observation_options classmethod

observation_options(entity_name: str = 'robot') -> dict[str, ObservationGroupOptions]

Return a minimal placeholder observation group.

HOMIE reads raw simulator state directly in :meth:step; Eden still requires at least one observation term to build the env.

Source code in eden/extensions/policies/homie/policy.py

@classmethod
def observation_options(cls, entity_name: str = "robot") -> dict[str, _OGO]:
    """Return a minimal placeholder observation group.

    HOMIE reads raw simulator state directly in :meth:`step`; Eden still
    requires at least one observation term to build the env.
    """
    return dict(
        policy=ObservationGroupOptions(
            concatenate_terms=True,
            enable_corruption=False,
            dofs_pos=observations.DofsPos.configure(entity_name=entity_name),
        ),
    )

reset

reset(envs_idx: slice | Tensor | None = None) -> None

Reset the tracker history/prev-actions for the given environments.

PARAMETER	DESCRIPTION
envs_idx	Environments to reset; None resets all. TYPE: slice | Tensor | None DEFAULT: None

Source code in eden/extensions/policies/homie/policy.py

def reset(self, envs_idx: (slice | torch.Tensor) | None = None) -> None:
    """Reset the tracker history/prev-actions for the given environments.

    Parameters
    ----------
    envs_idx : slice | torch.Tensor | None, optional
        Environments to reset; ``None`` resets all.
    """
    self._tracker.reset(envs_idx)

step

step(obs: dict[str, Tensor | dict[str, Tensor]]) -> dict[str, Tensor]

Read proprioception, run the tracker, and return 29-DOF targets.

PARAMETER	DESCRIPTION
obs	Ignored — raw simulator state is read directly from the env. TYPE: dict

RETURNS	DESCRIPTION
dict[str, Tensor]	{"dofs_pos_controller": (N, 29)} absolute joint-position targets.

Source code in eden/extensions/policies/homie/policy.py

def step(self, obs: dict[str, torch.Tensor | dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]:
    """Read proprioception, run the tracker, and return 29-DOF targets.

    Parameters
    ----------
    obs : dict
        Ignored — raw simulator state is read directly from the env.

    Returns
    -------
    dict[str, torch.Tensor]
        ``{"dofs_pos_controller": (N, 29)}`` absolute joint-position targets.
    """
    # Read off the entity, not an observation group: this policy rides along as an extra deploy
    # slot on another policy's or task's env, whose groups are not its own (see the
    # "Declare the observation" paragraph of docs/guides/policy-porting.md).
    del obs
    robot = self.env.entities[self.entity_name]

    # Proprioception in HOMIE joint order (== G1Dofs29 order).
    dof_pos = robot.get_dofs_pos(dofs_idx_local=self._dofs_idx)
    dof_vel = robot.get_dofs_vel(dofs_idx_local=self._dofs_idx)

    # Body-frame base angular velocity and projected gravity (root frame).
    base_ang_vel = robot.get_ang(frame="body")
    base_quat = robot.get_quat()
    projected_gravity = inv_transform_by_quat(self._gravity_world, base_quat)

    homie_actions = self._tracker.step(
        raw_commands=self._commands,
        base_ang_vel=base_ang_vel,
        projected_gravity=projected_gravity,
        dof_pos=dof_pos,
        dof_vel=dof_vel,
    )  # (N, 15)

    lower_targets = self._tracker.actions_to_targets(homie_actions)  # (N, 15)
    targets = torch.cat([lower_targets, self._arm_targets], dim=-1)  # (N, 29)
    return {"dofs_pos_controller": targets}

OliVelocityLocoPolicy

OliVelocityLocoPolicy(env: EnvBase, options: PolicyOptions)

Bases: PolicyBase

Frozen velocity-tracking loco/stand controller for LimX Oli (rsl_rl backend).

Source code in eden/extensions/policies/limx_oli_loco/policy.py

def __init__(self, env: EnvBase, options: PolicyOptions) -> None:
    super().__init__(env, options)
    if not self.checkpoint_path:
        raise ValueError(
            "OliVelocityLocoPolicy has no released checkpoint; set checkpoint_path to a local "
            "velocity_tracking_flat model_*.pt trained for Oli (eden deploy --policy_checkpoint)."
        )
    device = env.device
    self.device = device
    self._dtype = torch.float32
    n = env.num_envs

    self._actor, self._obs_mean, self._obs_std = _load_rsl_rl_actor(self.checkpoint_path, device)

    robot = self.require_entity(self.entity_name)
    _, dofs_idx = robot.find_named_dofs_idx_local(OLI_DOFS_NAME, preserve_order=True)
    self._dofs_idx = torch.as_tensor(dofs_idx, dtype=torch.long, device=device)

    defaults = Oli().default_dofs_pos
    self._default_dof_pos = torch.as_tensor(
        resolve_dofs_values(defaults, list(OLI_DOFS_NAME)), dtype=self._dtype, device=device
    )
    self._action_scale = torch.as_tensor(
        resolve_dofs_values(OLI_ACTION_SCALE, list(OLI_DOFS_NAME)), dtype=self._dtype, device=device
    )

    self._commands = torch.zeros(n, 3, device=device, dtype=self._dtype)
    self._last_action = torch.zeros(n, NUM_DOFS, device=device, dtype=self._dtype)
    self._gravity_world = torch.tensor([0.0, 0.0, -1.0], device=device, dtype=self._dtype).expand(n, 3)

Attributes

entity_name class-attribute instance-attribute

entity_name: str = 'robot'

checkpoint_path class-attribute instance-attribute

checkpoint_path: str | None = None

md5_hash class-attribute instance-attribute

md5_hash: str | None = None

registry class-attribute

registry: str | None = None

dataset class-attribute

dataset: str | None = None

device instance-attribute

device = device

Methods:

set_command

set_command(command: Tensor) -> None

Set the 3-D (vx, vy, wz) velocity command. All zeros (the default) is a stand.

Source code in eden/extensions/policies/limx_oli_loco/policy.py

def set_command(self, command: torch.Tensor) -> None:
    """Set the 3-D ``(vx, vy, wz)`` velocity command. All zeros (the default) is a stand."""
    self._commands[:] = torch.as_tensor(command, dtype=self._dtype, device=self.device)

reset

reset(envs_idx=None) -> None

Clear the action feedback at a handover so the policy re-enters from rest.

Source code in eden/extensions/policies/limx_oli_loco/policy.py

def reset(self, envs_idx=None) -> None:
    """Clear the action feedback at a handover so the policy re-enters from rest."""
    if envs_idx is None:
        self._last_action.zero_()
    else:
        self._last_action[envs_idx] = 0.0

robot_options classmethod

robot_options() -> Oli

Source code in eden/extensions/policies/limx_oli_loco/policy.py

@classmethod
def robot_options(cls) -> Oli:
    return Oli()

action_options classmethod

action_options(entity_name: str = 'robot') -> dict[str, ActionTermOptions]

An absolute-target controller: :meth:step has already applied OLI_ACTION_SCALE.

The deploy slot contract requires reference_source=ZERO and scale=1.0 so the returned targets are written to the robot unchanged; the training scale lives inside :meth:step, not here.

Source code in eden/extensions/policies/limx_oli_loco/policy.py

@classmethod
def action_options(cls, entity_name: str = "robot") -> dict[str, ActionTermOptions]:
    """An absolute-target controller: :meth:`step` has already applied ``OLI_ACTION_SCALE``.

    The deploy slot contract requires ``reference_source=ZERO`` and ``scale=1.0`` so the
    returned targets are written to the robot unchanged; the training scale lives inside
    :meth:`step`, not here.
    """
    # Through the term class, not a bare ``ActionTermOptions``: ``configure`` stamps the
    # registry key the deploy slot builder resolves the controller by.
    return {
        "dofs_pos_controller": actions.ImplicitPDController.configure(
            entity_name=entity_name,
            dofs_name=list(OLI_DOFS_NAME),
            reference_source=ReferenceSource.ZERO,
            scale=1.0,
            offset=0.0,
        )
    }

command_options classmethod

command_options(entity_name: str = 'robot') -> dict[str, CommandTermOptions]

The velocity_command term the observation group below reads, held at zero (a stand).

Declared so the env this policy owns (--policy limx_oli_loco alone, or named first) can build: generated_commands resolves its term on the env's command manager at construction, and a task-less env has no command manager unless a policy asks for one. :meth:step reads its own :meth:set_command buffer, not this term, so the operator's --velocity_command / --base_command reach the policy through the deploy loop as they do for Holosoma; the term here only has to exist and be quiet.

Source code in eden/extensions/policies/limx_oli_loco/policy.py

@classmethod
def command_options(cls, entity_name: str = "robot") -> dict[str, CommandTermOptions]:
    """The ``velocity_command`` term the observation group below reads, held at zero (a stand).

    Declared so the env this policy *owns* (``--policy limx_oli_loco`` alone, or named first)
    can build: ``generated_commands`` resolves its term on the env's command manager at
    construction, and a task-less env has no command manager unless a policy asks for one.
    :meth:`step` reads its own :meth:`set_command` buffer, not this term, so the operator's
    ``--velocity_command`` / ``--base_command`` reach the policy through the deploy loop as
    they do for Holosoma; the term here only has to exist and be quiet.
    """
    return {
        "velocity_command": ConstantVelocityCommand.configure(
            entity_name=entity_name, command_value=(0.0, 0.0, 0.0)
        )
    }

observation_options classmethod

observation_options(entity_name: str = 'robot') -> dict[str, ObservationGroupOptions]

The training group, in declaration order — the layout :meth:step reproduces.

Source code in eden/extensions/policies/limx_oli_loco/policy.py

@classmethod
def observation_options(cls, entity_name: str = "robot") -> dict[str, ObservationGroupOptions]:
    """The training group, in declaration order — the layout :meth:`step` reproduces."""
    return {
        "policy": ObservationGroupOptions(
            enable_corruption=False,
            base_lin_vel=ObservationTermOptions(
                name="base_lin_vel", params={"entity_name": entity_name, "frame": "body"}
            ),
            base_ang_vel=ObservationTermOptions(
                name="base_ang_vel", params={"entity_name": entity_name, "frame": "body"}
            ),
            projected_gravity=ObservationTermOptions(name="projected_gravity", entity_name=entity_name),
            dofs_pos=observations.DofsPos.configure(entity_name=entity_name, offset_from_default=True),
            dofs_vel=observations.DofsVel.configure(entity_name=entity_name),
            actions=ObservationTermOptions(name="last_action"),
            commands=ObservationTermOptions(name="generated_commands", params={"command_name": "velocity_command"}),
        )
    }

step

step(obs) -> dict[str, Tensor]

Run the frozen actor and return absolute 31-DOF joint targets.

obs is ignored: riding along on another task's env, this policy's observation group is not the one that env computes, so it reads robot state directly (the Holosoma pattern).

Source code in eden/extensions/policies/limx_oli_loco/policy.py

def step(self, obs) -> dict[str, torch.Tensor]:
    """Run the frozen actor and return absolute 31-DOF joint targets.

    ``obs`` is ignored: riding along on another task's env, this policy's observation group
    is not the one that env computes, so it reads robot state directly (the Holosoma pattern).
    """
    del obs
    robot = self.env.entities[self.entity_name]
    actor_obs = self._build_obs(robot)
    with torch.no_grad():
        action = self._actor((actor_obs - self._obs_mean) / (self._obs_std + 0.01))
    self._last_action = action
    return {"dofs_pos_controller": self._default_dof_pos + self._action_scale * action}

HuskyPolicy

HuskyPolicy(env: EnvBase, options: PolicyOptions)

Bases: PolicyBase

HUSKY skateboarding controller (G1 23-DOF + skateboard, ONNX backend).

Source code in eden/extensions/policies/husky/policy.py

def __init__(self, env: EnvBase, options: PolicyOptions) -> None:
    super().__init__(env, options)

    self.device = env.device

    robot = self.require_entity(self.entity_name)
    # The policy balances on the board without observing it, so a scene that lacks one runs
    # and falls: fail here instead.
    self.require_entity(self.skateboard_name)
    # The joint observations and the action term read the robot's DOFs in its own order.
    if list(robot.dofs_name) != list(POLICY_JOINT_NAMES):
        raise ValueError(
            f"HuskyPolicy drives the skater G1's 23 DOFs in {list(POLICY_JOINT_NAMES)} order; this robot's "
            f"dofs_name is {robot.dofs_name}. Build the robot from HuskyPolicy.robot_options()."
        )

    # The built robot must carry the training stance in every env: the joint-position
    # observation and the action term are both relative to its default pose.
    expected_default = torch.tensor(DEFAULT_DOFS_POS, device=env.device)
    built_default = robot.get_default_dofs_pos(POLICY_JOINT_NAMES).to(env.device)  # (num_envs, NUM_JOINTS)
    # ``allclose``'s tolerance, per env: the first env whose row is off is the one named.
    mismatched = ~torch.isclose(built_default, expected_default, atol=1e-6).all(dim=1)
    if mismatched.any().item():
        env_idx = int(mismatched.nonzero()[0].item())
        raise ValueError(
            f"Robot default_dofs_pos does not match the HUSKY training keyframe in env {env_idx}: "
            f"built={built_default[env_idx].tolist()} expected={expected_default.tolist()}"
        )

    # The env has to be the one this policy declares: a task env publishes its own
    # ``policy`` group and command terms, which the actor would read as garbage.
    group_dims = env.observation_manager.group_obs_dim
    if group_dims.get(OBS_GROUP) != (ACTOR_OBS_DIM,):
        raise ValueError(
            f"HuskyPolicy reads a {ACTOR_OBS_DIM}-dim '{OBS_GROUP}' observation group, but this env "
            f"publishes {group_dims}. Build the env from HuskyPolicy.observation_options() (a task-less "
            f"`eden inference --policy husky` does), not from another task's config."
        )
    self._skate_command()  # resolves the command term now, so a missing one fails at construction

    self._actor = HuskyActor(self.checkpoint_path, device=str(env.device))

Attributes

reads_env_observations class-attribute instance-attribute

reads_env_observations = True

entity_name class-attribute instance-attribute

entity_name: str = 'robot'

skateboard_name class-attribute instance-attribute

skateboard_name: str = 'skateboard'

checkpoint_path class-attribute instance-attribute

checkpoint_path: str = 'husky/skater_flat_g1.onnx'

md5_hash class-attribute instance-attribute

md5_hash: str | None = '69549f8d4a4fcc2629e4ba988666f841'

device instance-attribute

device = env.device

Methods:

robot_options classmethod

robot_options() -> HuskyG1

Skater G1 in the push-stance keyframe, placed to start on the deck.

Source code in eden/extensions/policies/husky/policy.py

@classmethod
def robot_options(cls) -> HuskyG1:
    """Skater G1 in the push-stance keyframe, placed to start on the deck."""
    return HuskyG1()

skateboard_options classmethod

skateboard_options() -> HuskySkateboard

Skateboard placed under the robot's left foot (training placement).

Source code in eden/extensions/policies/husky/policy.py

@classmethod
def skateboard_options(cls) -> HuskySkateboard:
    """Skateboard placed under the robot's left foot (training placement)."""
    return HuskySkateboard(default_root_pos=SKATEBOARD_DEFAULT_ROOT_POS)

scene_entity_options classmethod

scene_entity_options() -> dict[str, EntityOptionsLike]

The board the policy rides — without it the robot has nothing to balance on.

Source code in eden/extensions/policies/husky/policy.py

@classmethod
def scene_entity_options(cls) -> dict[str, EntityOptionsLike]:
    """The board the policy rides — without it the robot has nothing to balance on."""
    return {cls.skateboard_name: cls.skateboard_options()}

env_options classmethod

env_options() -> dict[str, Any]

Source code in eden/extensions/policies/husky/policy.py

@classmethod
def env_options(cls) -> dict[str, Any]:
    # 50 Hz control (decimation 4 x sim_dt 0.005) and the training solver
    # iterations (mjlab: iterations=10, ls_iterations=20).
    return dict(
        sim_dt=0.005,
        decimation=4,
        iterations=10,
        ls_iterations=20,
        max_collision_pairs=64,
        batch_dofs_info=False,
        enable_self_collision=False,
    )

action_options classmethod

action_options(entity_name: str = 'robot') -> dict[str, ActionTermOptions]

The training task's action term: step returns the actor's raw (normalized) action.

Source code in eden/extensions/policies/husky/policy.py

@classmethod
def action_options(cls, entity_name: str = "robot") -> dict[str, ActionTermOptions]:
    """The training task's action term: ``step`` returns the actor's raw (normalized) action."""
    from eden.tasks.husky.config import joint_position_action

    return dict(dofs_pos_controller=joint_position_action(entity_name))

command_options classmethod

command_options() -> dict[str, CommandTermOptions]

The held [lin_vel_x, relative_heading] command and the phase clock the actor observes.

Source code in eden/extensions/policies/husky/policy.py

@classmethod
def command_options(cls) -> dict[str, CommandTermOptions]:
    """The held ``[lin_vel_x, relative_heading]`` command and the phase clock the actor observes."""
    from eden.tasks.husky.custom_terms.commands import ConstantSkateCommand

    return {
        COMMAND_NAME: ConstantSkateCommand.configure(
            cycle_time_s=CYCLE_TIME_S, still_vx_threshold=STILL_VX_THRESHOLD
        )
    }

observation_options classmethod

observation_options(entity_name: str = 'robot') -> dict[str, ObservationGroupOptions]

The training task's actor group, with its 5-step history and without its training noise.

Source code in eden/extensions/policies/husky/policy.py

@classmethod
def observation_options(cls, entity_name: str = "robot") -> dict[str, ObservationGroupOptions]:
    """The training task's actor group, with its 5-step history and without its training noise."""
    from eden.tasks.husky.config import actor_observation_terms

    return {
        OBS_GROUP: ObservationGroupOptions(
            **actor_observation_terms(entity_name=entity_name, noise=False, history=True),
            enable_corruption=False,
        )
    }

set_command

set_command(command: Tensor) -> None

Set the skate command [lin_vel_x, relative_heading] (shared or per-env).

lin_vel_x is the target board speed in m/s (training range 0–1.5); below 0.1 the phase clock pins at 0 and the policy balances in place. relative_heading (rad, trained on ±pi/4) steers the board relative to its initial heading. Takes effect from the next observation the env computes.

Only an env built from :meth:command_options holds a command. On --task husky the training task's :class:~eden.tasks.husky.custom_terms.commands.SkateCommand samples its own, which is what the policy then observes.

Source code in eden/extensions/policies/husky/policy.py

def set_command(self, command: torch.Tensor) -> None:
    """Set the skate command ``[lin_vel_x, relative_heading]`` (shared or per-env).

    ``lin_vel_x`` is the target board speed in m/s (training range 0–1.5);
    below 0.1 the phase clock pins at 0 and the policy balances in place.
    ``relative_heading`` (rad, trained on ±pi/4) steers the board relative
    to its initial heading. Takes effect from the next observation the env computes.

    Only an env built from :meth:`command_options` holds a command. On ``--task husky`` the
    training task's :class:`~eden.tasks.husky.custom_terms.commands.SkateCommand` samples its
    own, which is what the policy then observes.
    """
    term = self._skate_command()
    if not hasattr(term, "set_command"):
        raise TypeError(
            f"The env's '{COMMAND_NAME}' command term is a {type(term).__name__}, which samples its own "
            f"command and holds none, so there is nothing to set. Build the env from "
            f"HuskyPolicy.command_options() (a ConstantSkateCommand) to hold one, or leave the command to "
            f"the {type(term).__name__}."
        )
    term.set_command(command)

step

step(obs: dict[str, Tensor | dict[str, Tensor]]) -> dict[str, Tensor]

Source code in eden/extensions/policies/husky/policy.py

def step(self, obs: dict[str, torch.Tensor | dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]:
    return {"dofs_pos_controller": self._actor(obs[OBS_GROUP])}

Whole-body control and motion tracking

AMOPolicy

AMOPolicy(env: EnvBase, options: PolicyOptions)

Bases: PolicyBase

AMO: Adaptive Motion Optimization for Hyper-Dexterous Humanoid Whole-Body Control.

Source code in eden/extensions/policies/amo/policy.py

def __init__(self, env: EnvBase, options: PolicyOptions):
    super().__init__(env, options)

    self._policy = torch.jit.load(self.checkpoint_path, map_location=env.device)
    self._policy.eval()

Attributes

reads_env_observations class-attribute instance-attribute

reads_env_observations = True

checkpoint_path class-attribute instance-attribute

checkpoint_path: str = 'amo_jit.pt'

md5_hash class-attribute instance-attribute

md5_hash: str | None = '3efbda329f247f41a7459f50f88f4bc5'

Methods:

robot_options classmethod

robot_options() -> G1Dofs29NoWristDof

Return a G1Dofs29NoWristDof with AMO's training default pose and PD gains.

Source code in eden/extensions/policies/amo/policy.py

@classmethod
def robot_options(cls) -> G1Dofs29NoWristDof:
    """Return a ``G1Dofs29NoWristDof`` with AMO's training default pose and PD gains."""
    return amo_g1_robot()

env_options classmethod

env_options() -> dict[str, Any]

Source code in eden/extensions/policies/amo/policy.py

@classmethod
def env_options(cls) -> dict[str, Any]:
    return dict(
        sim_dt=0.002,
        decimation=10,
    )

action_options classmethod

action_options(entity_name: str = 'robot') -> dict[str, ActionTermOptions]

Source code in eden/extensions/policies/amo/policy.py

@classmethod
def action_options(cls, entity_name: str = "robot") -> dict[str, ActionTermOptions]:
    return dict(
        dofs_pos_controller=actions.ImplicitPDController.configure(
            reference_source=ReferenceSource.DEFAULT,
            scale=0.25,
            entity_name=entity_name,
            dofs_name=list(G1_UCSD_NO_WRIST_DOFS_NAME),
        ),
    )

command_options classmethod

command_options(**kwargs) -> dict[str, CommandTermOptions]

Source code in eden/extensions/policies/amo/policy.py

@classmethod
def command_options(cls, **kwargs) -> dict[str, CommandTermOptions]:
    resampling_time_range = kwargs.pop("resampling_time_range", (4.0, 4.0))
    return dict(amo_command=AMOCommand.configure(resampling_time_range=resampling_time_range, **kwargs))

observation_options classmethod

observation_options(entity_name: str = 'robot') -> dict[str, ObservationGroupOptions]

Source code in eden/extensions/policies/amo/policy.py

@classmethod
def observation_options(cls, entity_name: str = "robot") -> dict[str, ObservationGroupOptions]:
    return dict(
        obs_prop=ObservationGroupOptions(
            concatenate_terms=True,
            history_length=26,  # Current + 25 past frames
            flatten_history_dim=True,
            backfill=False,  # Zero pad history
            # Terms order matches play_amo.py: ang_vel, rpy, dyaw, dofs_pos, dofs_vel, last_action, gait, adapter
            ang_vel=ObsTerm.configure(
                func=observations.base_ang_vel,
                scale=0.25,
                params={"entity_name": entity_name, "frame": "body"},
            ),
            rpy=ObsTerm.configure(
                func=observations.base_rpy,
                params={
                    "entity_name": entity_name,
                    "terms_idx": [0, 1],
                },  # roll, pitch
            ),
            dyaw=ObsTerm.configure(
                func=amo_dyaw_obs,
                params={"entity_name": entity_name, "command_name": "amo_command"},
            ),
            dofs_pos=observations.DofsPos.configure(scale=1.0, entity_name=entity_name, offset_from_default=True),
            dofs_vel=ObsTerm.configure(
                func=amo_dofs_vel,
                scale=0.05,
                params={"entity_name": entity_name},
            ),
            last_action=ObsTerm.configure(
                func=amo_last_action,
                scale=1.0,
                params={
                    "entity_name": entity_name,
                    "action_name": "dofs_pos_controller",
                    "command_name": "amo_command",
                },
            ),
            gait=ObsTerm.configure(
                func=amo_gait_obs,
                params={"entity_name": entity_name, "command_name": "amo_command"},
            ),
            adapter=ObsTerm.configure(
                func=amo_adapter_output,
                params={"entity_name": entity_name, "command_name": "amo_command"},
            ),
            terms_order=[
                "ang_vel",
                "rpy",
                "dyaw",
                "dofs_pos",
                "dofs_vel",
                "last_action",
                "gait",
                "adapter",
            ],
        ),
        demo_obs=ObservationGroupOptions(
            concatenate_terms=True,
            history_length=1,
            term=ObsTerm.configure(
                func=amo_demo_obs,
                params={"entity_name": entity_name, "command_name": "amo_command"},
            ),
        ),
    )

step

step(obs: dict[str, Tensor | dict[str, Tensor]]) -> dict[str, Tensor]

Source code in eden/extensions/policies/amo/policy.py

def step(self, obs: dict[str, torch.Tensor | dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]:
    with torch.no_grad():
        # obs_prop frame is 93 wide (3+2+2+23+23+23+2+15); history is oldest first: [t-25, ..., t-1, t].
        dim_prop = 93
        full_hist = obs["obs_prop"]  # (N, 26 * 93)
        current_prop = full_hist[..., -dim_prop:]
        obs_hist = full_hist[..., -dim_prop * 11 : -dim_prop]  # 10 frames
        extra_hist = full_hist[..., -dim_prop * 26 : -dim_prop]  # 25 frames

        obs_demo = obs["demo_obs"]
        obs_priv = torch.zeros(self.env.num_envs, 3, device=self.env.device)

        # Policy input: [prop, demo, priv, obs_hist]
        policy_in = torch.cat([current_prop, obs_demo, obs_priv, obs_hist], dim=1)

        raw_action = self._policy(policy_in, extra_hist)
        raw_action = torch.clamp(raw_action, -40.0, 40.0)

        # Combine with arm actions
        cmd_term = self.env.command_manager.get_term("amo_command")
        arm_target = cmd_term.get_arm_target()  # (N, 8) absolute

        default_arm = cmd_term.default_arm_pos
        arm_delta = (arm_target - default_arm) / 0.25

        action = torch.cat([raw_action, arm_delta], dim=1)

    return dict(dofs_pos_controller=action)

BeyondMimicMovesPolicy

BeyondMimicMovesPolicy(env: EnvBase, options: PolicyOptions)

Bases: PolicyBase

BeyondMimic G1-Moves: 60 motion-tracking policies for the 29-DOF Unitree G1.

Source dataset: https://huggingface.co/datasets/exptech/g1-moves

The clip is selected at configure time via the clip_name option (e.g. "B_DadDance"). The policy resolves the per-clip checkpoint and motion file from the dataset's manifest.json and downloads them on first use via huggingface_hub.

Notes

The expected env / observation / action setup mirrors examples/projects/mimic/beyondmimic (and the BeyondMimicConfig task), since the upstream g1-moves env.yaml is functionally identical:

• decimation=4, sim_dt=0.005, 29-DOF G1
• Observation (160 dims): [motion(58), motion_anchor_pos_b(3), motion_anchor_ori_b(6), base_lin_vel(3), base_ang_vel(3), dofs_pos_rel(29), dofs_vel(29), last_action(29)]
• Joint order is the MuJoCo XML declaration order (no permutation needed).

The robot entity default_dofs_pos MUST equal :data:G1_MOVES_DEFAULT_DOFS_POS, otherwise both dofs_pos_rel (obs) and the action target end up biased relative to the policy's training frame and the robot won't track. eden.robots.G1Dofs29's out-of-the-box defaults differ (e.g. hip_pitch=-0.1 vs. mjlab's -0.312, elbow=1.5 vs. 0.6), so build the robot with :func:g1_moves_robot (what :meth:robot_options returns).

Source code in eden/extensions/policies/beyondmimic_moves/policy.py

def __init__(self, env: EnvBase, options: PolicyOptions) -> None:
    # NOTE: PolicyBase.__init__ skips asset resolution while checkpoint_path is None.
    super().__init__(env, options)

    clip_dir = f"{self._clip_category(self.clip_name)}/{self.clip_name}"
    self.checkpoint_path = self.fetch_asset(f"{clip_dir}/policy/{self.clip_name}_policy.pt")
    self.motion_file = self.prepare_motion_file(self.clip_name)

    self._policy = self._build_actor(self.checkpoint_path, env.device)

    self._verify_robot_defaults()

Attributes

reads_env_observations class-attribute instance-attribute

reads_env_observations = True

registry class-attribute

registry: str = 'exptech'

dataset class-attribute

dataset: str = 'g1-moves'

clip_name class-attribute instance-attribute

clip_name: str = 'B_DadDance'

checkpoint_path class-attribute instance-attribute

checkpoint_path: str | None = self.fetch_asset(f'{clip_dir}/policy/{self.clip_name}_policy.pt')

motion_file instance-attribute

motion_file = self.prepare_motion_file(self.clip_name)

Methods:

prepare_motion_file classmethod

prepare_motion_file(clip_name: str) -> str

Resolve a clip's motion file, downloading and remapping body axes as needed.

Returns the path to a Genesis-BFS-ordered NPZ next to the original mjlab download. Safe to call before constructing the env (so the resolved path can be passed to :meth:command_options); subsequent calls hit the cached <clip>_eden.npz without re-reading the source.

Source code in eden/extensions/policies/beyondmimic_moves/policy.py

@classmethod
def prepare_motion_file(cls, clip_name: str) -> str:
    """Resolve a clip's motion file, downloading and remapping body axes as needed.

    Returns the path to a Genesis-BFS-ordered NPZ next to the original mjlab
    download. Safe to call before constructing the env (so the resolved path
    can be passed to :meth:`command_options`); subsequent calls hit the cached
    ``<clip>_eden.npz`` without re-reading the source.
    """
    category = cls._clip_category(clip_name)
    raw_npz = cls.fetch_asset(f"{category}/{clip_name}/training/{clip_name}.npz")
    return _remap_motion_npz_to_genesis_order(raw_npz)

robot_options classmethod

robot_options() -> G1Dofs29

Return a G1Dofs29 pinned to the g1-moves training default pose.

Source code in eden/extensions/policies/beyondmimic_moves/policy.py

@classmethod
def robot_options(cls) -> G1Dofs29:
    """Return a ``G1Dofs29`` pinned to the g1-moves training default pose."""
    return g1_moves_robot()

env_options classmethod

env_options() -> dict[str, Any]

Source code in eden/extensions/policies/beyondmimic_moves/policy.py

@classmethod
def env_options(cls) -> dict[str, Any]:
    # Mirror mjlab's SimulationCfg + the eden BeyondMimicConfig settings used
    # to validate sim2sim transfer. Eden's EnvOptions defaults
    # (iterations=50, ls_iterations=50, enable_self_collision=True,
    # enable_multi_contact=True, use_gjk_collision=True, max_collision_pairs=50)
    # diverge enough from MuJoCo's solver behaviour that physics drift makes
    # the policy fall.
    return dict(
        sim_dt=0.005,
        sim_substeps=1,
        decimation=4,
        iterations=10,
        ls_iterations=20,
        max_collision_pairs=20,
        batch_dofs_info=False,
        use_gjk_collision=False,
        enable_self_collision=False,
        enable_multi_contact=False,
    )

action_options classmethod

action_options(entity_name: str = 'robot') -> dict[str, ActionTermOptions]

Source code in eden/extensions/policies/beyondmimic_moves/policy.py

@classmethod
def action_options(cls, entity_name: str = "robot") -> dict[str, ActionTermOptions]:
    return dict(
        dofs_pos_controller=actions.ExplicitPDController.configure(
            entity_name=entity_name,
            dofs_name=G1Dofs29().dofs_name,
            scale=G1_ACTION_SCALE,
            reference_source=ReferenceSource.DEFAULT,
        ),
    )

command_options classmethod

command_options(clip_name: str | None = None, debug_vis: bool = False, entity_name: str = 'robot', ghost_entity_name: str = '', sampling_mode: str = 'start') -> dict[str, CommandTermOptions]

Source code in eden/extensions/policies/beyondmimic_moves/policy.py

@classmethod
def command_options(
    cls,
    clip_name: str | None = None,
    debug_vis: bool = False,
    entity_name: str = "robot",
    ghost_entity_name: str = "",
    sampling_mode: str = "start",
) -> dict[str, CommandTermOptions]:
    # This is a per-clip tracker, not a general one: the motion file is derived from the
    # clip, never passed in. command_options runs during env.build() (before the policy
    # instance exists, so self.motion_file is not available yet), so resolve it here via the
    # classmethod resolver — the same path __init__ uses. Defaults to the configured clip.
    return dict(
        motion=MotionCommand.configure(
            resampling_time_range=(1.0e9, 1.0e9),
            debug_vis=debug_vis,
            entity_name=entity_name,
            ghost_entity_name=ghost_entity_name,
            motion_file=cls.prepare_motion_file(clip_name or cls.clip_name),
            anchor_link_name=_ANCHOR_BODY,
            links_name=_BODY_NAMES,
            pose_range={},
            velocity_range={},
            dofs_position_range=(0.0, 0.0),
            sampling_mode=sampling_mode,
        ),
    )

observation_options classmethod

observation_options(entity_name: str = 'robot') -> dict[str, ObservationGroupOptions]

Source code in eden/extensions/policies/beyondmimic_moves/policy.py

@classmethod
def observation_options(cls, entity_name: str = "robot") -> dict[str, ObservationGroupOptions]:
    return dict(
        policy=ObservationGroupOptions(
            concatenate_terms=True,
            enable_corruption=False,
            motion=ObsTerm.configure(
                func=observations.generated_commands,
                params={"command_name": "motion"},
            ),
            motion_anchor_pos_b=ObsTerm.configure(
                func=g1_moves_motion_anchor_pos_b,
                params={"command_name": "motion"},
            ),
            motion_anchor_ori_b=ObsTerm.configure(
                func=g1_moves_motion_anchor_ori_b,
                params={"command_name": "motion"},
            ),
            base_lin_vel=ObsTerm.configure(
                func=g1_moves_imu_lin_vel_b,
                params={"entity_name": entity_name},
            ),
            base_ang_vel=ObsTerm.configure(
                func=observations.base_ang_vel,
                params={"entity_name": entity_name, "frame": "body"},
            ),
            dofs_pos=observations.DofsPos.configure(entity_name=entity_name, offset_from_default=True),
            dofs_vel=observations.DofsVel.configure(entity_name=entity_name),
            last_action=ObsTerm.configure(
                func=observations.last_action,
            ),
            terms_order=[
                "motion",
                "motion_anchor_pos_b",
                "motion_anchor_ori_b",
                "base_lin_vel",
                "base_ang_vel",
                "dofs_pos",
                "dofs_vel",
                "last_action",
            ],
        ),
    )

step

step(obs: dict[str, Tensor | dict[str, Tensor]]) -> dict[str, Tensor]

Source code in eden/extensions/policies/beyondmimic_moves/policy.py

def step(self, obs: dict[str, torch.Tensor | dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]:
    with torch.no_grad():
        raw_action = self._policy(obs["policy"])
    return dict(dofs_pos_controller=raw_action)

GMTPolicy

GMTPolicy(env: EnvBase, options: PolicyOptions)

Bases: PolicyBase

GMT: General Motion Tracking for Humanoid Whole-Body Control.

NOTE

This policy only supports 23 dofs G1 robot.

Source code in eden/extensions/policies/gmt/policy.py

def __init__(self, env: EnvBase, options: PolicyOptions):
    super().__init__(env, options)

    self._policy = torch.jit.load(self.checkpoint_path).to(env.device)
    self._policy.eval()

Attributes

reads_env_observations class-attribute instance-attribute

reads_env_observations = True

checkpoint_path class-attribute instance-attribute

checkpoint_path: str = 'gmt_g1_23dofs.pt'

md5_hash class-attribute instance-attribute

md5_hash: str | None = 'eb4cfef3de37d27c3b7eae0ac7baa8bf'

Methods:

robot_options classmethod

robot_options() -> G1Dofs29NoWristDof

Return a G1Dofs29NoWristDof matching GMT's training pose / gains / armature / damping.

Source code in eden/extensions/policies/gmt/policy.py

@classmethod
def robot_options(cls) -> G1Dofs29NoWristDof:
    """Return a ``G1Dofs29NoWristDof`` matching GMT's training pose / gains / armature / damping."""
    return gmt_g1_robot()

env_options classmethod

env_options() -> dict[str, Any]

Source code in eden/extensions/policies/gmt/policy.py

@classmethod
def env_options(cls) -> dict[str, Any]:
    return dict(
        sim_dt=0.001,
        decimation=20,
    )

action_options classmethod

action_options(entity_name: str = 'robot') -> dict[str, ActionTermOptions]

Source code in eden/extensions/policies/gmt/policy.py

@classmethod
def action_options(cls, entity_name: str = "robot") -> dict[str, ActionTermOptions]:
    return dict(
        dofs_pos_controller=actions.ImplicitPDController.configure(
            reference_source=ReferenceSource.DEFAULT,
            scale=0.5,
            entity_name=entity_name,
            # NOTE: only 23 dofs are used for the GMT
            dofs_name=list(G1_UCSD_NO_WRIST_DOFS_NAME),
        ),
    )

command_options classmethod

command_options(motion_file: str, debug_vis: bool = False) -> dict[str, CommandTermOptions]

Source code in eden/extensions/policies/gmt/policy.py

@classmethod
def command_options(cls, motion_file: str, debug_vis: bool = False) -> dict[str, CommandTermOptions]:
    return dict(
        motion=GMTMotionLibCommand.configure(
            resampling_time_range=(60.0, 60.0),
            debug_vis=debug_vis,
            motion_file=motion_file,
        )
    )

observation_options classmethod

observation_options(entity_name: str = 'robot') -> dict[str, ObservationGroupOptions]

Source code in eden/extensions/policies/gmt/policy.py

@classmethod
def observation_options(cls, entity_name: str = "robot") -> dict[str, ObservationGroupOptions]:
    return dict(
        mimic_obs=ObservationGroupOptions(
            command=ObsTerm.configure(
                func=observations.generated_commands,
                params={"command_name": "motion"},
            ),
        ),
        obs_prop=ObservationGroupOptions(
            concatenate_terms=True,
            history_length=21,
            flatten_history_dim=True,
            backfill=False,
            ang_vel=ObsTerm.configure(
                func=observations.base_ang_vel,
                scale=0.25,
                params={"entity_name": entity_name, "frame": "body"},
            ),
            rpy=ObsTerm.configure(
                func=observations.base_rpy,
                # NOTE: only roll and pitch are used for the GMT
                params={"entity_name": entity_name, "terms_idx": [0, 1]},
            ),
            dofs_pos=observations.DofsPos.configure(scale=1.0, entity_name=entity_name, offset_from_default=True),
            dofs_vel=ObsTerm.configure(
                func=dofs_vel_no_ankle_gmt,
                scale=0.05,
                params={"entity_name": entity_name},
            ),
            last_action=ObsTerm.configure(
                func=observations.last_action,
                scale=1.0,
                params={"action_name": "dofs_pos_controller"},
            ),
            terms_order=["ang_vel", "rpy", "dofs_pos", "dofs_vel", "last_action"],
        ),
    )

step

step(obs: dict[str, Tensor | dict[str, Tensor]]) -> dict[str, Tensor]

Get action from policy.

Source code in eden/extensions/policies/gmt/policy.py

def step(self, obs: dict[str, torch.Tensor | dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]:
    """Get action from policy."""
    with torch.no_grad():
        obs_prop = obs["obs_prop"][..., -74:]
        obs_hist = obs["obs_prop"][..., :-74]
        raw_action = self._policy(torch.cat([obs["mimic_obs"], obs_prop, obs_hist], dim=1))
        raw_action = torch.clamp(raw_action, -10, 10)
    return dict(dofs_pos_controller=raw_action)

HumanoidGptPolicy

HumanoidGptPolicy(env: EnvBase, options: PolicyOptions)

Bases: PolicyBase

Humanoid-GPT whole-body motion-tracking policy for the Unitree G1 (29 DOF).

Runs the released Transformer (default) or MLP (policy_type="mlp") ONNX checkpoint; both share the 136-dim observation / action contract, the transformer additionally consuming a K-frame obs history.

policy_type, transformer_history_len, ref_ema_alpha and ref_cvel_at_com shape the observation as well as the network, so pass the same values to :meth:observation_options and :meth:configure; an instance configured differently from the env it is handed is refused.

Source code in eden/extensions/policies/humanoid_gpt/policy.py

def __init__(self, env: EnvBase, options: PolicyOptions) -> None:
    # Select the checkpoint pair on ``options`` (not ``self``) BEFORE super():
    # ConfigurableMixin mirrors options -> self, clobbering a pre-super attr.
    if options.policy_type not in ("mlp", "transformer"):
        raise ValueError(f"Unknown policy_type={options.policy_type!r} (expected 'mlp' or 'transformer').")
    options.checkpoint_path, options.md5_hash = self._select_checkpoint(
        options.policy_type,
        options.checkpoint_path,
        options.md5_hash,
        options.mlp_checkpoint_path,
        options.mlp_md5_hash,
        options.transformer_checkpoint_path,
        options.transformer_md5_hash,
    )
    super().__init__(env, options)

    self.device = env.device
    robot = self.require_entity("robot")
    if list(robot.dofs_name) != G1_DOFS29_NAME:
        raise ValueError(
            f"HumanoidGptPolicy drives the G1's 29 DOFs in G1Dofs29 order; this robot's dofs_name is "
            f"{robot.dofs_name}. Run it with --robot G1Dofs29."
        )
    self._check_env_matches_configuration(env)

    # Patch the batch axis only for num_envs>1; num_envs=1 runs unpatched.
    self._model = HumanoidGptOnnx(self.checkpoint_path, device=str(env.device), dynamic_batch=(env.num_envs > 1))
    if self.policy_type == "transformer":
        self._validate_transformer_graph()

Attributes

reads_env_observations class-attribute instance-attribute

reads_env_observations = True

checkpoint_path class-attribute instance-attribute

checkpoint_path: str | None = None

md5_hash class-attribute instance-attribute

md5_hash: str | None = None

mlp_checkpoint_path class-attribute instance-attribute

mlp_checkpoint_path: str = _MLP_CHECKPOINT

mlp_md5_hash class-attribute instance-attribute

mlp_md5_hash: str | None = '7147c555a68432c76723c52468bbdd26'

transformer_checkpoint_path class-attribute instance-attribute

transformer_checkpoint_path: str = _TRANSFORMER_CHECKPOINT

transformer_md5_hash class-attribute instance-attribute

transformer_md5_hash: str | None = 'c2a884643074cee94ab57722bea116b4'

policy_type class-attribute instance-attribute

policy_type: str = 'transformer'

ref_cvel_at_com class-attribute instance-attribute

ref_cvel_at_com: bool = True

ref_ema_alpha class-attribute instance-attribute

ref_ema_alpha: float | None = 0.8

transformer_history_len class-attribute instance-attribute

transformer_history_len: int = 16

device instance-attribute

device = env.device

Methods:

env_options classmethod

env_options() -> dict[str, Any]

Source code in eden/extensions/policies/humanoid_gpt/policy.py

@classmethod
def env_options(cls) -> dict[str, Any]:
    # control_dt = sim_dt * decimation = 0.01 * 2 = 0.02s (50 Hz), matching
    # Humanoid-GPT's ctrl_dt and the 50 Hz reference motion.
    return dict(
        sim_dt=0.01,
        sim_substeps=2,
        decimation=2,
        iterations=20,
        ls_iterations=20,
        max_collision_pairs=20,
        batch_dofs_info=False,
        use_gjk_collision=False,
        enable_self_collision=False,
        enable_multi_contact=False,
    )

action_options classmethod

action_options(entity_name: str = 'robot') -> dict[str, ActionTermOptions]

The network's raw action, scaled per joint about Humanoid-GPT's nominal stance.

Source code in eden/extensions/policies/humanoid_gpt/policy.py

@classmethod
def action_options(cls, entity_name: str = "robot") -> dict[str, ActionTermOptions]:
    """The network's raw action, scaled per joint about Humanoid-GPT's nominal stance."""
    return dict(
        dofs_pos_controller=actions.ImplicitPDController.configure(
            entity_name=entity_name,
            dofs_name=list(G1_DOFS29_NAME),
            scale=dict(zip(G1_DOFS29_NAME, resolve_action_scale())),
            # The nominal stance, not the robot's default pose (the stock G1Dofs29 differs on six joints).
            offset=dict(zip(G1_DOFS29_NAME, HGPT_DEFAULT_DOFS_POS)),
            reference_source=ReferenceSource.ZERO,
        ),
    )

command_options classmethod

command_options(motion_file: str | list[str] = 'motion.npz', debug_vis: bool = False, entity_name: str = 'robot', ghost_entity_name: str = '', sampling_mode: str = 'start') -> dict[str, CommandTermOptions]

Source code in eden/extensions/policies/humanoid_gpt/policy.py

@classmethod
def command_options(
    cls,
    motion_file: str | list[str] = "motion.npz",
    debug_vis: bool = False,
    entity_name: str = "robot",
    ghost_entity_name: str = "",
    sampling_mode: str = "start",
) -> dict[str, CommandTermOptions]:
    # list of clips -> multi-clip term; single path -> single-clip; step() is
    # agnostic to which backs it.
    return motion_command_options(
        motion_file,
        anchor_link_name=_ANCHOR_BODY,
        links_name=_MOTION_BODY_NAMES,
        entity_name=entity_name,
        ghost_entity_name=ghost_entity_name,
        debug_vis=debug_vis,
        sampling_mode=sampling_mode,
    )

observation_options classmethod

observation_options(entity_name: str = 'robot', policy_type: str = _CLASS_SETTING, transformer_history_len: int = _CLASS_SETTING, ref_ema_alpha: float | None = _CLASS_SETTING, ref_cvel_at_com: bool = _CLASS_SETTING) -> dict[str, ObservationGroupOptions]

The 136-dim frame; for the Transformer, stacked into a (K, 136) history backfilled at reset.

Each knob not passed is the class's own setting, which is what a task-less eden inference|deploy --policy humanoid_gpt builds with. Pass the values the policy is configured with otherwise, e.g. observation_options(policy_type="mlp") alongside configure(policy_type="mlp").

Source code in eden/extensions/policies/humanoid_gpt/policy.py

@classmethod
def observation_options(
    cls,
    entity_name: str = "robot",
    policy_type: str = _CLASS_SETTING,
    transformer_history_len: int = _CLASS_SETTING,
    ref_ema_alpha: float | None = _CLASS_SETTING,
    ref_cvel_at_com: bool = _CLASS_SETTING,
) -> dict[str, ObservationGroupOptions]:
    """The 136-dim frame; for the Transformer, stacked into a ``(K, 136)`` history backfilled at reset.

    Each knob not passed is the class's own setting, which is what a task-less
    ``eden inference|deploy --policy humanoid_gpt`` builds with. Pass the values the policy is
    configured with otherwise, e.g. ``observation_options(policy_type="mlp")`` alongside
    ``configure(policy_type="mlp")``.
    """
    knobs = {
        "policy_type": policy_type,
        "transformer_history_len": transformer_history_len,
        "ref_ema_alpha": ref_ema_alpha,
        "ref_cvel_at_com": ref_cvel_at_com,
    }
    knobs = {name: getattr(cls, name) if value is _CLASS_SETTING else value for name, value in knobs.items()}
    if knobs["policy_type"] not in ("mlp", "transformer"):
        raise ValueError(f"Unknown policy_type={knobs['policy_type']!r} (expected 'mlp' or 'transformer').")
    entity = {"entity_name": entity_name}
    history = (
        {"history_length": knobs["transformer_history_len"], "flatten_history_dim": False}
        if knobs["policy_type"] == "transformer"
        else {}
    )
    return {
        OBS_GROUP: ObservationGroupOptions(
            enable_corruption=False,
            **history,
            gyro=ObsTerm.configure(func=observations.base_ang_vel, params={**entity, "frame": "body"}),
            gvec=observations.ProjectedGravity.configure(entity_name=entity_name),
            dofs_pos=ObsTerm.configure(func=hgpt_dofs_pos_rel, params=entity),
            dofs_vel=observations.DofsVel.configure(entity_name=entity_name),
            last_action=ObsTerm.configure(func=observations.last_action),
            reference=HgptReference.configure(
                root_link_name=_ROOT_BODY,
                entity_name=entity_name,
                ema_alpha=knobs["ref_ema_alpha"],
                cvel_at_com=knobs["ref_cvel_at_com"],
            ),
        )
    }

step

step(obs: dict[str, Tensor | dict[str, Tensor]]) -> dict[str, Tensor]

Source code in eden/extensions/policies/humanoid_gpt/policy.py

def step(self, obs: dict[str, torch.Tensor | dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]:
    # clone: infer() aliases the ONNX output buffer, which the next call overwrites.
    return {"dofs_pos_controller": self._model.infer(obs[OBS_GROUP]).clone()}

MetaMotivoPolicy

MetaMotivoPolicy(env: EnvBase, options: PolicyOptions)

Bases: PolicyBase

MetaMotivo (FBcpr) policy for SMPL humanoid control on Eden.

The task-context vector z determines the behavior. It can be set directly via :meth:set_context, inferred from a goal via :meth:compute_goal_context, resampled from the prior via :meth:sample_random_context, or left as the random sample drawn on the first :meth:step.

PARAMETER	DESCRIPTION
model_name	Hugging Face model id, loaded directly with from_pretrained (no Eden checkpoint). Variants: "facebook/metamotivo-S-1" (24.5M) and "facebook/metamotivo-M-1" (288M). TYPE: str

Source code in eden/extensions/policies/metamotivo/policy.py

def __init__(self, env: EnvBase, options: PolicyOptions):
    super().__init__(env, options)

    self._z: torch.Tensor | None = None
    self._step_count = 0
    self._model = None

    self._load_model()

Attributes

reads_env_observations class-attribute instance-attribute

reads_env_observations = True

checkpoint_path class-attribute instance-attribute

checkpoint_path: str | None = None

registry class-attribute

registry: str | None = None

dataset class-attribute

dataset: str | None = None

z_dim property

z_dim: int

Dimensionality of the task-context vector.

context property

context: Tensor | None

Current task-context vector z.

Methods:

robot_options classmethod

robot_options() -> SmplHumanoid

Return the HumEnv SMPL humanoid MetaMotivo was trained against.

Source code in eden/extensions/policies/metamotivo/policy.py

@classmethod
def robot_options(cls) -> SmplHumanoid:
    """Return the HumEnv SMPL humanoid MetaMotivo was trained against."""
    return metamotivo_smpl_robot()

env_options classmethod

env_options() -> dict[str, Any]

Source code in eden/extensions/policies/metamotivo/policy.py

@classmethod
def env_options(cls) -> dict[str, Any]:
    return dict(
        # HumEnv: sim_dt=1/450, action_repeat=15 -> ~30Hz control
        sim_dt=1.0 / 450.0,
        decimation=15,
    )

action_options classmethod

action_options(entity_name: str = 'robot') -> dict[str, ActionTermOptions]

Source code in eden/extensions/policies/metamotivo/policy.py

@classmethod
def action_options(cls, entity_name: str = "robot") -> dict[str, ActionTermOptions]:
    return dict(
        smpl_actuator=MujocoSMPLActuator.configure(
            entity_name=entity_name,
            dofs_name=SmplHumanoid().dofs_name,
        ),
    )

observation_options classmethod

observation_options(entity_name: str = 'robot') -> dict[str, ObservationGroupOptions]

Source code in eden/extensions/policies/metamotivo/policy.py

@classmethod
def observation_options(cls, entity_name: str = "robot") -> dict[str, ObservationGroupOptions]:
    return dict(
        proprio=ObservationGroupOptions(
            concatenate_terms=True,
            history_length=1,
            flatten_history_dim=True,
            proprio=ObsTerm.configure(
                func=humenv_proprio,
                params={"entity_name": entity_name},
            ),
        ),
    )

step

step(obs: dict[str, Tensor | dict[str, Tensor]]) -> dict[str, Tensor]

Compute action from observation using the MetaMotivo model.

PARAMETER	DESCRIPTION
obs	Observation dict from Eden's observation manager. Expected key: "proprio" with shape (num_envs, obs_dim). TYPE: dict

RETURNS	DESCRIPTION
dict[str, Tensor]	Action dict with key "smpl_actuator", shape (num_envs, 69).

Source code in eden/extensions/policies/metamotivo/policy.py

def step(self, obs: dict[str, torch.Tensor | dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]:
    """Compute action from observation using the MetaMotivo model.

    Parameters
    ----------
    obs : dict
        Observation dict from Eden's observation manager.
        Expected key: ``"proprio"`` with shape ``(num_envs, obs_dim)``.

    Returns
    -------
    dict[str, torch.Tensor]
        Action dict with key ``"smpl_actuator"``, shape ``(num_envs, 69)``.
    """
    with torch.no_grad():
        proprio = obs["proprio"]  # (num_envs, obs_dim)

        # Ensure context z is initialized
        if self._z is None or self._z.shape[0] != proprio.shape[0]:
            self._z = self._model.sample_z(proprio.shape[0], device=str(self.env.device))

        # Get action from the model
        action = self._model.act(
            obs=proprio.float(),
            z=self._z.float(),
            mean=True,
        )

        self._step_count += 1

    return dict(smpl_actuator=action)

set_context

set_context(z: Tensor) -> None

Set the task-context vector directly.

PARAMETER	DESCRIPTION
z	Shape (num_envs, z_dim); a 1-D (z_dim,) is unsqueezed to (1, z_dim). TYPE: Tensor

Source code in eden/extensions/policies/metamotivo/policy.py

def set_context(self, z: torch.Tensor) -> None:
    """Set the task-context vector directly.

    Parameters
    ----------
    z : torch.Tensor
        Shape ``(num_envs, z_dim)``; a 1-D ``(z_dim,)`` is unsqueezed to ``(1, z_dim)``.
    """
    if z.dim() == 1:
        z = z.unsqueeze(0)
    self._z = z.to(self.env.device)

sample_random_context

sample_random_context(batch_size: int | None = None) -> Tensor

Sample a random context from the model's z prior.

PARAMETER	DESCRIPTION
batch_size	If None, uses env.num_envs. TYPE: int | None DEFAULT: None

RETURNS	DESCRIPTION
Tensor	Shape (batch_size, z_dim).

Source code in eden/extensions/policies/metamotivo/policy.py

def sample_random_context(self, batch_size: int | None = None) -> torch.Tensor:
    """Sample a random context from the model's z prior.

    Parameters
    ----------
    batch_size : int | None
        If None, uses ``env.num_envs``.

    Returns
    -------
    torch.Tensor
        Shape ``(batch_size, z_dim)``.
    """
    if batch_size is None:
        batch_size = self.env.num_envs
    self._z = self._model.sample_z(batch_size, device=str(self.env.device))
    return self._z

compute_goal_context

compute_goal_context(goal_obs: Tensor) -> Tensor

Compute context from a goal observation via the backward map, and make it current.

PARAMETER	DESCRIPTION
goal_obs	Target proprioceptive observation, shape (1, obs_dim) or (obs_dim,). TYPE: Tensor

RETURNS	DESCRIPTION
Tensor	Context vector, shape (1, z_dim); the current context becomes its expansion to (num_envs, z_dim).

Source code in eden/extensions/policies/metamotivo/policy.py

def compute_goal_context(self, goal_obs: torch.Tensor) -> torch.Tensor:
    """Compute context from a goal observation via the backward map, and make it current.

    Parameters
    ----------
    goal_obs : torch.Tensor
        Target proprioceptive observation, shape ``(1, obs_dim)`` or ``(obs_dim,)``.

    Returns
    -------
    torch.Tensor
        Context vector, shape ``(1, z_dim)``; the current context becomes its
        expansion to ``(num_envs, z_dim)``.
    """
    if goal_obs.dim() == 1:
        goal_obs = goal_obs.unsqueeze(0)
    goal_obs = goal_obs.to(self.env.device).float()
    z = self._model.goal_inference(next_obs=goal_obs)
    self._z = z.expand(self.env.num_envs, -1)
    return z

compute_tracking_context

compute_tracking_context(target_obs: Tensor) -> Tensor

Compute tracking context from a sequence of target observations.

PARAMETER	DESCRIPTION
target_obs	Target observations, shape (seq_len, obs_dim). TYPE: Tensor

RETURNS	DESCRIPTION
Tensor	Context vectors, shape (seq_len, z_dim). The current context is not changed.

Source code in eden/extensions/policies/metamotivo/policy.py

def compute_tracking_context(self, target_obs: torch.Tensor) -> torch.Tensor:
    """Compute tracking context from a sequence of target observations.

    Parameters
    ----------
    target_obs : torch.Tensor
        Target observations, shape ``(seq_len, obs_dim)``.

    Returns
    -------
    torch.Tensor
        Context vectors, shape ``(seq_len, z_dim)``. The current context is not changed.
    """
    target_obs = target_obs.to(self.env.device).float()
    z = self._model.tracking_inference(next_obs=target_obs)
    return z

OmniXtremePolicy

OmniXtremePolicy(env: EnvBase, options: PolicyOptions)

Bases: PolicyBase

OmniXtreme: Dual-policy humanoid motion tracking (flow-matching base + residual correction).

Source code in eden/extensions/policies/omnixtreme/policy.py

def __init__(self, env: EnvBase, options: PolicyOptions):
    super().__init__(env, options)

    self._residual_path = self.fetch_asset(self.residual_checkpoint_path)
    self._fk_path = self.fetch_asset(self.fk_checkpoint_path)

    self._verify_md5(self._residual_path, self.residual_md5_hash)
    self._verify_md5(self._fk_path, self.fk_md5_hash)

    device_str = str(env.device)
    self._base_policy = OnnxBasePolicyWrapper(self.checkpoint_path, device=device_str)
    self._residual_policy = OnnxResidualPolicyWrapper(self._residual_path, device=device_str)

    # The anchor observation runs FK through the command term (see omnixtreme_anchor_ori_6d).
    cmd_term = env.command_manager.get_term("omnixtreme_motion")
    cmd_term.fk_session = _make_fk_onnx_session(self._fk_path)
    cmd_term._fk_onnx_path = self._fk_path

    # Residual gain, overridable with the RESIDUAL_GAIN environment variable.
    self._residual_gain = float(os.environ.get("RESIDUAL_GAIN", "1.0"))

    self.action_term = env.action_manager.get_term("dofs_pos_controller")

    # The residual policy reads and writes joints in BeyondMimic (Isaac Lab) order, not the
    # URDF order of the env's DOFs: ``x[:, _inv_perm]`` maps URDF -> BeyondMimic, ``x[:, _perm]`` back.
    # fmt: off
    perm_list = [
        0,3,6,9,13,17,1,4,7,10,14,18,2,5,8,11,15,19,21,23,25,27,12,16,20,22,24,26,28,
    ]
    inv_perm_list = [
        0,6,12,1,7,13,2,8,14,3,9,15,22,4,10,16,23,5,11,17,24,18,25,19,26,20,27,21,28,
    ]
    # fmt: on
    self._perm = torch.tensor(perm_list, dtype=torch.long, device=env.device)
    self._inv_perm = torch.tensor(inv_perm_list, dtype=torch.long, device=env.device)

Attributes

reads_env_observations class-attribute instance-attribute

reads_env_observations = True

entity_name class-attribute instance-attribute

entity_name: str = 'robot'

checkpoint_path class-attribute instance-attribute

checkpoint_path: str = 'omnixtreme/base_policy.onnx'

md5_hash class-attribute instance-attribute

md5_hash: str | None = '67b75756640bbb52720ee5920bbf0910'

residual_checkpoint_path class-attribute instance-attribute

residual_checkpoint_path: str = 'omnixtreme/residual_policy.onnx'

residual_md5_hash class-attribute instance-attribute

residual_md5_hash: str | None = '7ad3515080d32e51241f959984c0b568'

fk_checkpoint_path class-attribute instance-attribute

fk_checkpoint_path: str = 'omnixtreme/fk.onnx'

fk_md5_hash class-attribute instance-attribute

fk_md5_hash: str | None = '94c9896adda6939b0ddfa508ddc0d0be'

action_term instance-attribute

action_term = env.action_manager.get_term('dofs_pos_controller')

Methods:

robot_options classmethod

robot_options() -> OmniXtreme_G1

Return the OmniXtreme G1 (7520-22.5 hip-pitch motors, training default pose).

Source code in eden/extensions/policies/omnixtreme/policy.py

@classmethod
def robot_options(cls) -> OmniXtreme_G1:
    """Return the OmniXtreme G1 (7520-22.5 hip-pitch motors, training default pose)."""
    return OmniXtreme_G1()

env_options classmethod

env_options() -> dict[str, Any]

Source code in eden/extensions/policies/omnixtreme/policy.py

@classmethod
def env_options(cls) -> dict[str, Any]:
    return dict(sim_dt=0.004, decimation=5)

action_options classmethod

action_options(entity_name: str = 'robot') -> dict[str, ActionTermOptions]

Source code in eden/extensions/policies/omnixtreme/policy.py

@classmethod
def action_options(cls, entity_name: str = "robot") -> dict[str, ActionTermOptions]:
    return dict(
        dofs_pos_controller=actions.ImplicitPDController.configure(
            reference_source=ReferenceSource.DEFAULT,
            scale=ACTION_SCALE,
            entity_name=entity_name,
            dofs_name=en.robots.G1Dofs29NoHands().dofs_name,
            modifier=EnvelopeClip.configure(),
        ),
    )

command_options classmethod

command_options(motion_file: str, debug_vis: bool = False) -> dict[str, CommandTermOptions]

Source code in eden/extensions/policies/omnixtreme/policy.py

@classmethod
def command_options(cls, motion_file: str, debug_vis: bool = False) -> dict[str, CommandTermOptions]:
    return dict(
        omnixtreme_motion=OmniXtremeRefCommand.configure(
            resampling_time_range=(60.0, 60.0),
            motion_file=motion_file,
            debug_vis=debug_vis,
        )
    )

observation_options classmethod

observation_options(entity_name: str = 'robot') -> dict[str, ObservationGroupOptions]

Source code in eden/extensions/policies/omnixtreme/policy.py

@classmethod
def observation_options(cls, entity_name: str = "robot") -> dict[str, ObservationGroupOptions]:
    return dict(
        obs_prop=ObservationGroupOptions(
            concatenate_terms=True,
            history_length=15,
            flatten_history_dim=True,
            backfill=False,
            dofs_pos=observations.DofsPos.configure(scale=1.0, entity_name=entity_name, offset_from_default=True),
            dofs_vel=observations.DofsVel.configure(scale=1.0, entity_name=entity_name),
            ang_vel=ObsTerm.configure(
                func=observations.base_ang_vel,
                scale=1.0,
                params={"entity_name": entity_name, "frame": "body"},
            ),
            prev_base_action=ObsTerm.configure(
                func=omnixtreme_prev_base_action,
                scale=1.0,
                params={"command_name": "omnixtreme_motion"},
            ),
            terms_order=["dofs_pos", "dofs_vel", "ang_vel", "prev_base_action"],
        ),
        command_obs=ObservationGroupOptions(
            concatenate_terms=True,
            motion=ObsTerm.configure(
                func=observations.generated_commands,
                params={"command_name": "omnixtreme_motion"},
            ),
            anchor=ObsTerm.configure(
                func=omnixtreme_anchor_ori_6d,
                params={"entity_name": entity_name, "command_name": "omnixtreme_motion"},
            ),
            terms_order=["motion", "anchor"],
        ),
    )

step

step(obs: dict[str, Tensor | dict[str, Tensor]]) -> dict[str, Tensor]

Source code in eden/extensions/policies/omnixtreme/policy.py

def step(self, obs: dict[str, torch.Tensor | dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]:
    with torch.no_grad():
        # --- Extract observations ---
        obs_prop_flat = obs["obs_prop"]  # (B, 1350) = 15 * 90
        command_obs = obs["command_obs"]  # (B, 64)
        B = obs_prop_flat.shape[0]
        dim_per_frame = NUM_DOFS * 3 + 3  # dof_pos + dof_vel + ang_vel + prev_action = 90

        real_obs = obs_prop_flat[..., -dim_per_frame:]
        real_hist = obs_prop_flat.reshape(B, 15, dim_per_frame)

        # --- Run base policy ---
        obs_dict = {
            "real_obs": real_obs,
            "command_obs": command_obs,
            "real_historical_obs_raw": real_hist,
            "self_obs": torch.zeros(B, 1, device=self.env.device),
            "initial_noise": torch.randn(B, NUM_DOFS, device=self.env.device, dtype=torch.float32),
        }
        a_base = self._base_policy(obs_dict)

        # Store base action for next-step observation
        cmd = self.env.command_manager.get_term("omnixtreme_motion")
        cmd.prev_base_action[:] = a_base

        # --- Build residual observation (183 dims) ---
        dof_pos_rel = real_obs[..., :NUM_DOFS]
        dof_vel = real_obs[..., NUM_DOFS : NUM_DOFS * 2]
        ang_vel = real_obs[..., NUM_DOFS * 2 : NUM_DOFS * 2 + 3]

        term = self.action_term
        unscaled = (term.processed_action - term.effective_offset) / term.effective_scale
        prev_action_unscaled = unscaled[:, self._inv_perm]

        # Reorder to proto (BeyondMimic) order for residual policy
        real_reordered = torch.cat(
            [
                ang_vel,
                dof_pos_rel[:, self._inv_perm],
                dof_vel[:, self._inv_perm],
                prev_action_unscaled,
            ],
            dim=-1,
        )

        ref_pos = command_obs[..., :NUM_DOFS]
        ref_vel = command_obs[..., NUM_DOFS : NUM_DOFS * 2]
        anchor6d = command_obs[..., NUM_DOFS * 2 :]
        cmd_reordered = torch.cat(
            [
                ref_pos[:, self._inv_perm],
                ref_vel[:, self._inv_perm],
                anchor6d,
            ],
            dim=-1,
        )

        residual_obs = torch.cat([cmd_reordered, real_reordered, a_base[:, self._inv_perm]], dim=-1)

        # --- Run residual policy and combine ---
        a_res = self._residual_policy(residual_obs)[:, self._perm]  # reindex to URDF
        raw_actions = torch.clamp(a_base + self._residual_gain * a_res, -40.0, 40.0)

    return dict(dofs_pos_controller=raw_actions)

ProtoMotionV3Policy

ProtoMotionV3Policy(env: EnvBase, options: PolicyOptions)

Bases: PolicyBase

ProtoMotions v3 (g1-bones-deploy) motion-tracking policy on Unitree G1 (29 DOF).

Loads the AMP actor directly from last.ckpt and reproduces the obs + action transform pipeline in PyTorch, so checkpoints can be tested without re-exporting an ONNX bundle.

Source code in eden/extensions/policies/protomotionv3/policy.py

def __init__(self, env: EnvBase, options: PolicyOptions) -> None:
    super().__init__(env, options)

    if env.num_envs != 1:
        raise ValueError(
            f"{type(self).__name__} only supports num_envs=1 (got {env.num_envs}); "
            "matches the g1-bones-deploy deployment contract."
        )
    self.device = env.device
    robot = self.require_entity("robot")
    # The joint observations and the action term read the robot's DOFs in its own order.
    if list(robot.dofs_name) != G1_DOFS29_NAME:
        raise ValueError(
            f"ProtoMotionV3Policy drives the G1's 29 DOFs in G1Dofs29 order; this robot's dofs_name is "
            f"{robot.dofs_name}. Build the robot from ProtoMotionV3Policy.robot_options()."
        )

    # The env has to be the one this policy declares: a task env publishes its own
    # ``policy`` group, which the actor would read as garbage.
    group_dims = env.observation_manager.group_obs_dim
    if group_dims.get(OBS_GROUP) != (_OBS_DIM,):
        raise ValueError(
            f"ProtoMotionV3Policy reads a {_OBS_DIM}-dim '{OBS_GROUP}' observation group, but this env "
            f"publishes {group_dims}. Build the env from ProtoMotionV3Policy.observation_options() (a "
            f"task-less `eden inference --policy protomotionv3` does), not from another task's config."
        )

    self._actor = _CkptActor(self.checkpoint_path, env.device)

Attributes

reads_env_observations class-attribute instance-attribute

reads_env_observations = True

checkpoint_path class-attribute instance-attribute

checkpoint_path: str = 'protomotion_v3/g1_bones_deploy.ckpt'

md5_hash class-attribute instance-attribute

md5_hash: str | None = '205a22b9becc4c8b589da7509d940b5c'

device instance-attribute

device = env.device

Methods:

robot_options classmethod

robot_options() -> G1Dofs29

Return the stock G1Dofs29, whose gains match the training values.

Its default pose does not: the training pose is the action term's offset.

Source code in eden/extensions/policies/protomotionv3/policy.py

@classmethod
def robot_options(cls) -> G1Dofs29:
    """Return the stock ``G1Dofs29``, whose gains match the training values.

    Its default pose does not: the training pose is the action term's ``offset``.
    """
    return G1Dofs29()

env_options classmethod

env_options() -> dict[str, Any]

Source code in eden/extensions/policies/protomotionv3/policy.py

@classmethod
def env_options(cls) -> dict[str, Any]:
    # Matches `simulation_params.genesis` from the training resolved config:
    #   fps=100, decimation=2, substeps=2 → control_dt=0.02s (50 Hz).
    return dict(
        sim_dt=0.01,
        sim_substeps=2,
        decimation=2,
        iterations=20,
        ls_iterations=20,
        max_collision_pairs=20,
        batch_dofs_info=False,
        use_gjk_collision=False,
        enable_self_collision=False,
        enable_multi_contact=False,
    )

action_options classmethod

action_options(entity_name: str = 'robot') -> dict[str, ActionTermOptions]

bm_pd_action: the actor's mean action, scaled per joint about the training default pose.

Source code in eden/extensions/policies/protomotionv3/policy.py

@classmethod
def action_options(cls, entity_name: str = "robot") -> dict[str, ActionTermOptions]:
    """``bm_pd_action``: the actor's mean action, scaled per joint about the training default pose."""
    dofs_name = G1Dofs29().dofs_name
    return dict(
        dofs_pos_controller=actions.ImplicitPDController.configure(
            entity_name=entity_name,
            dofs_name=dofs_name,
            scale=dict(zip(dofs_name, _ACTION_SCALE)),
            # ProtoMotions' default pose, not the robot's (the stock G1Dofs29 differs on 12 joints).
            offset=dict(zip(dofs_name, _PD_ACTION_OFFSET)),
            reference_source=ReferenceSource.ZERO,
        ),
    )

command_options classmethod

command_options(motion_file: str = 'motion.npz', debug_vis: bool = False, entity_name: str = 'robot', ghost_entity_name: str = '', sampling_mode: str = 'start') -> dict[str, CommandTermOptions]

Source code in eden/extensions/policies/protomotionv3/policy.py

@classmethod
def command_options(
    cls,
    motion_file: str = "motion.npz",
    debug_vis: bool = False,
    entity_name: str = "robot",
    ghost_entity_name: str = "",
    sampling_mode: str = "start",
) -> dict[str, CommandTermOptions]:
    return dict(
        motion=MotionCommand.configure(
            resampling_time_range=(1.0e9, 1.0e9),
            debug_vis=debug_vis,
            entity_name=entity_name,
            ghost_entity_name=ghost_entity_name,
            motion_file=motion_file,
            anchor_link_name=_ANCHOR_BODY,
            links_name=_MOTION_BODY_NAMES,
            pose_range={},
            velocity_range={},
            dofs_position_range=(0.0, 0.0),
            sampling_mode=sampling_mode,
        ),
    )

observation_options classmethod

observation_options(entity_name: str = 'robot') -> dict[str, ObservationGroupOptions]

The 349-dim input: reduced coords (64), future target poses (256), previous joint targets (29).

Source code in eden/extensions/policies/protomotionv3/policy.py

@classmethod
def observation_options(cls, entity_name: str = "robot") -> dict[str, ObservationGroupOptions]:
    """The 349-dim input: reduced coords (64), future target poses (256), previous joint targets (29)."""
    entity = {"entity_name": entity_name}
    return {
        OBS_GROUP: ObservationGroupOptions(
            enable_corruption=False,
            # noisy_reduced_coords_obs (64)
            dofs_pos=observations.DofsPos.configure(entity_name=entity_name),
            dofs_vel=observations.DofsVel.configure(entity_name=entity_name),
            root_ang_vel=ObsTerm.configure(func=observations.base_ang_vel, params={**entity, "frame": "body"}),
            anchor_gravity=ProtomotionAnchorGravity.configure(
                entity_name=entity_name, anchor_link_name=_ANCHOR_BODY
            ),
            # noisy_mimic_reduced_coords_target_poses (256)
            target_poses=ProtomotionTargetPoses.configure(
                entity_name=entity_name, anchor_link_name=_ANCHOR_BODY, future_offsets=_FUTURE_STEP_OFFSETS
            ),
            # historical_previous_processed_actions (29)
            prev_targets=ProtomotionPrevTargets.configure(),
        )
    }

step

step(obs: dict[str, Tensor | dict[str, Tensor]]) -> dict[str, Tensor]

Source code in eden/extensions/policies/protomotionv3/policy.py

def step(self, obs: dict[str, torch.Tensor | dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]:
    return {"dofs_pos_controller": self._actor(obs[OBS_GROUP])}

SonicPolicy

SonicPolicy(env: EnvBase, options: PolicyOptions)

Bases: PolicyBase

GEAR-SONIC universal-token whole-body controller (g1 mode, ONNX backend).

The SONIC release is split into two ONNX graphs (encoder + decoder); both are downloaded from Kashu7100/eden_policies (mirror of NVIDIA's HF repo) on first use. To use the official NVIDIA bundle directly, set the checkpoint_path / decoder_checkpoint_path class attrs to a local path before instantiating, or pass them through configure(...).

Source code in eden/extensions/policies/sonic/policy.py

def __init__(self, env: EnvBase, options: PolicyOptions) -> None:
    super().__init__(env, options)

    # PolicyBase already resolved the encoder path. Resolve the decoder
    # path via the same Hugging Face asset resolver (mirrors omnixtreme).
    self._decoder_path = self.fetch_asset(self.decoder_checkpoint_path)
    self._verify_md5(self._decoder_path, self.decoder_md5_hash)

    device = env.device
    self.device = device
    self._dtype = torch.float32

    # OnnxModule validates the static input dim against these expectations
    # (no-op when the graph has a dynamic axis), raising on a swapped/re-exported file.
    self._encoder = SonicEncoder(self.checkpoint_path, device=str(device), expected_input_dim=_ENCODER_INPUT_DIM)
    self._decoder = SonicDecoder(self._decoder_path, device=str(device), expected_input_dim=_DECODER_INPUT_DIM)

    robot = self.require_entity(self.entity_name)
    # MJCF (Eden) DOF order — matches the motion file and Genesis kinematic state.
    _, dofs_idx = robot.find_named_dofs_idx_local(G1Dofs29().dofs_name, preserve_order=True)
    self._dofs_idx_mj = torch.as_tensor(dofs_idx, dtype=torch.long, device=device)
    # Pelvis (anchor) idx_local in Genesis kinematic-tree order (BFS in genesis 1.1.x,
    # DFS in 1.2.0). Name-resolved, so it indexes the robot's own get_links_* correctly
    # under either ordering (it is NOT used to slice the motion npz).
    _, pelvis_idx = robot.find_named_links_idx_local([_ANCHOR_BODY], preserve_order=True)
    self._pelvis_link_idx = torch.as_tensor(pelvis_idx, dtype=torch.long, device=device)

    # Permutations: MJCF → IsaacLab (gather index that picks each IL slot's
    # value out of an MJCF-ordered tensor) and the inverse.
    self._mj_to_il = torch.as_tensor(MJ_TO_IL, dtype=torch.long, device=device)
    self._il_to_mj = torch.as_tensor(IL_TO_MJ, dtype=torch.long, device=device)

    # Action-transform constants (MJCF / Eden order).
    self._default_angles_mj = torch.tensor(DEFAULT_ANGLES_MJ, device=device, dtype=self._dtype)
    self._action_scale_mj = torch.tensor(ACTION_SCALE_MJ, device=device, dtype=self._dtype)

    # Future motion frame offsets [0, step, 2*step, ..., (N-1)*step] — the
    # first future frame IS the current motion frame, matching both the C++
    # deploy (`target_frame = current_frame + frame_idx * step_size`,
    # frame_idx=0..N-1) and gear_sonic's `future_time_steps_init =
    # arange(num_future_frames) * frame_skips`.
    self._future_offsets = torch.arange(
        0, _FUTURE_LEN * _FUTURE_STEP, _FUTURE_STEP, dtype=torch.long, device=device
    )

    # encoder_mode_4: scalar mode_id at slot 0, zeros at 1..3 — NOT a
    # one-hot vector. The C++ deploy code
    # (`GatherEncoderMode(buf, off, fill_zeros_num=3)`) writes
    # ``buf[off] = static_cast<float>(GetEncodeMode())`` and zero-pads the
    # remaining three slots. For g1 (mode_id=0) this is `[0, 0, 0, 0]`.
    # A one-hot `[1, 0, 0, 0]` would tell the encoder we're in teleop mode
    # (mode_id=1) and feed the g1 motion obs to the wrong branch.
    self._encoder_mode_value = torch.zeros(NUM_MODES, device=device, dtype=self._dtype)
    self._encoder_mode_value[0] = float(MODE_G1)  # = 0.0 — kept explicit for clarity

    # Heading correction (apply_delta_heading in the C++ deploy). Captured
    # per env on the first step() after that env's reset: the yaw-only
    # rotation that aligns the motion's start heading with the robot's
    # start heading. Without this, the robot's world frame drifts away from
    # the motion's world frame and future-anchor obs degrade as the episode
    # progresses. Stored per env so partial resets / per-env motion
    # wrap-around never clobber the offsets of envs mid-episode.
    N = env.num_envs
    self._heading_offset_wxyz = torch.zeros(N, 4, device=device, dtype=self._dtype)
    self._heading_offset_wxyz[:, 0] = 1.0  # identity until first capture
    self._needs_heading_capture = torch.ones(N, dtype=torch.bool, device=device)
    # Previous-step motion frame index. We use this to auto-detect motion
    # wrap-around (`MotionCommand._update_command` resamples when
    # `time_steps >= time_step_total`) per env and re-capture that env's
    # heading offset against the new motion-start pose. Without this, long
    # episodes that outlive the motion silently keep the stale offset
    # across the wrap.
    self._prev_time_steps: torch.Tensor | None = None

    # ----- history ring buffers -----
    # All in IsaacLab order, newest-LAST (so the model gets the most recent
    # frame at index _HISTORY_LEN-1 when the buffer is flattened).
    N = env.num_envs
    self._hist_q = torch.zeros(N, _HISTORY_LEN, _NUM_DOFS, device=device, dtype=self._dtype)
    self._hist_qdot = torch.zeros(N, _HISTORY_LEN, _NUM_DOFS, device=device, dtype=self._dtype)
    self._hist_action = torch.zeros(N, _HISTORY_LEN, _NUM_DOFS, device=device, dtype=self._dtype)
    self._hist_base_ang = torch.zeros(N, _HISTORY_LEN, 3, device=device, dtype=self._dtype)
    self._hist_gravity = torch.zeros(N, _HISTORY_LEN, 3, device=device, dtype=self._dtype)
    self._hist_buffers = (
        self._hist_q,
        self._hist_qdot,
        self._hist_action,
        self._hist_base_ang,
        self._hist_gravity,
    )

    # Scratch buffers, refilled in-place each step.
    self._encoder_buf = torch.zeros(N, _ENCODER_INPUT_DIM, device=device, dtype=self._dtype)
    self._decoder_buf = torch.zeros(N, _DECODER_INPUT_DIM, device=device, dtype=self._dtype)

    # World-frame "down" vector, expanded once to the env-batch shape so
    # the per-step pelvis-frame projection can skip a `zeros_like` + scalar
    # assignment. Read-only on the hot path.
    self._gravity_world = torch.tensor([0.0, 0.0, -1.0], device=device, dtype=self._dtype).expand(N, 3)

Attributes

entity_name class-attribute instance-attribute

entity_name: str = 'robot'

checkpoint_path class-attribute instance-attribute

checkpoint_path: str = 'sonic/model_encoder.onnx'

md5_hash class-attribute instance-attribute

md5_hash: str | None = 'a21244b1cdceac9dc8625442a55d42e0'

decoder_checkpoint_path class-attribute instance-attribute

decoder_checkpoint_path: str = 'sonic/model_decoder.onnx'

decoder_md5_hash class-attribute instance-attribute

decoder_md5_hash: str | None = '1d4391ad404d0ff36281abec349b23fd'

device instance-attribute

device = device

Methods:

robot_options classmethod

robot_options() -> G1Dofs29

Return a G1Dofs29 carrying the standing pose and PD gains the checkpoint expects.

This is the robot a task-less eden deploy/inference --policy sonic builds its scene from, so it decides both where the robot stands before the policy takes over and — via :class:~eden.extensions.deployment.base.DeploymentBase, which resolves default_dofs_kp / default_dofs_kd once into the Kp/Kd sent on the wire — what PD loop tracks its targets.

Pose is :data:~eden.extensions.policies.sonic.joint_maps.SONIC_DOFS_POS, the default_angles[] the action transform is centred on. The stock G1's shallower stance would leave the robot 0.37 rad off at the knees and 0.9 rad at the elbows when SONIC's first target lands.

Gains differ from the stock G1 in the hip-pitch row alone: SONIC derives them from the same motor table by the same Raibert heuristic Eden does, and hip pitch is the one row where the tables disagree — SONIC's is a 7520-22, Eden's stock spec a 7520-14. The stock gains would run the hips at 0.41x the released checkpoint's stiffness (40.18 vs 99.10 N·m/rad) under the 7520-22 action scale, i.e. 2.4x less torque per unit action than trained.

Armature and effort limits stay on the stock spec: those describe the physical actuator, not the controller, and which motor a given G1 revision carries is a hardware question this port cannot settle. Note the consequence — Genesis builds hip pitch with the 7520-14 armature (0.0102 vs SONIC's 0.0251 kg·m²) and clamps its torque at 88 rather than 139 N·m, so a Genesis rollout is still not SONIC's training dynamics even with the gains fixed.

Source code in eden/extensions/policies/sonic/policy.py

@classmethod
def robot_options(cls) -> G1Dofs29:
    """Return a ``G1Dofs29`` carrying the standing pose and PD gains the checkpoint expects.

    This is the robot a task-less ``eden deploy/inference --policy sonic`` builds its
    scene from, so it decides both where the robot stands before the policy takes over
    and — via :class:`~eden.extensions.deployment.base.DeploymentBase`, which resolves
    ``default_dofs_kp`` / ``default_dofs_kd`` once into the Kp/Kd sent on the wire —
    what PD loop tracks its targets.

    **Pose** is :data:`~eden.extensions.policies.sonic.joint_maps.SONIC_DOFS_POS`, the
    ``default_angles[]`` the action transform is centred on. The stock G1's shallower
    stance would leave the robot 0.37 rad off at the knees and 0.9 rad at the elbows
    when SONIC's first target lands.

    **Gains** differ from the stock G1 in the hip-pitch row alone: SONIC derives them
    from the same motor table by the same Raibert heuristic Eden does, and hip pitch is
    the one row where the tables disagree — SONIC's is a 7520-22, Eden's stock spec a
    7520-14. The stock gains would run the hips at **0.41x** the released checkpoint's
    stiffness (40.18 vs 99.10 N·m/rad) under the 7520-22 action scale, i.e. 2.4x less
    torque per unit action than trained.

    Armature and effort limits stay on the stock spec: those describe the physical
    actuator, not the controller, and which motor a given G1 revision carries is a
    hardware question this port cannot settle. Note the consequence — Genesis builds
    hip pitch with the 7520-14 armature (0.0102 vs SONIC's 0.0251 kg·m²) and clamps its
    torque at 88 rather than 139 N·m, so a Genesis rollout is still not SONIC's
    training dynamics even with the gains fixed.
    """
    return G1Dofs29(default_dofs_pos=SONIC_DOFS_POS, default_dofs_kp=SONIC_DOFS_KP, default_dofs_kd=SONIC_DOFS_KD)

env_options classmethod

env_options() -> dict[str, Any]

Source code in eden/extensions/policies/sonic/policy.py

@classmethod
def env_options(cls) -> dict[str, Any]:
    # 50 Hz control to match the released SONIC ONNX (motion target_fps=50,
    # decimation=2 → control_dt=0.02 s). The deploy stack also runs at 50 Hz.
    return dict(
        sim_dt=0.01,
        sim_substeps=2,
        decimation=2,
        iterations=20,
        ls_iterations=20,
        max_collision_pairs=20,
        batch_dofs_info=False,
        use_gjk_collision=False,
        enable_self_collision=False,
        enable_multi_contact=False,
    )

action_options classmethod

action_options(entity_name: str = 'robot') -> dict[str, ActionTermOptions]

Source code in eden/extensions/policies/sonic/policy.py

@classmethod
def action_options(cls, entity_name: str = "robot") -> dict[str, ActionTermOptions]:
    # `SonicPolicy.step` already converts the policy output to absolute
    # joint-position targets (default_angles + scale * action) in MJCF
    # order, so the action term just hands those targets to Genesis's
    # built-in PD as-is.
    return dict(
        dofs_pos_controller=actions.ImplicitPDController.configure(
            entity_name=entity_name,
            dofs_name=G1Dofs29().dofs_name,
            scale=1.0,
            reference_source=ReferenceSource.ZERO,
        ),
    )

command_options classmethod

command_options(motion_file: str | list[str] = 'motion.npz', debug_vis: bool = False, entity_name: str = 'robot', ghost_entity_name: str = '', sampling_mode: str = 'start') -> dict[str, CommandTermOptions]

Source code in eden/extensions/policies/sonic/policy.py

@classmethod
def command_options(
    cls,
    motion_file: str | list[str] = "motion.npz",
    debug_vis: bool = False,
    entity_name: str = "robot",
    ghost_entity_name: str = "",
    sampling_mode: str = "start",
) -> dict[str, CommandTermOptions]:
    # A list of clips -> per-env multi-clip term; a single path -> single-clip.
    # SonicPolicy.step reads motion through the shared accessor interface, so it
    # is unaware which term backs it.
    return motion_command_options(
        motion_file,
        anchor_link_name=_ANCHOR_BODY,
        links_name=_MOTION_BODY_NAMES,
        entity_name=entity_name,
        ghost_entity_name=ghost_entity_name,
        debug_vis=debug_vis,
        sampling_mode=sampling_mode,
    )

observation_options classmethod

observation_options(entity_name: str = 'robot') -> dict[str, ObservationGroupOptions]

Source code in eden/extensions/policies/sonic/policy.py

@classmethod
def observation_options(cls, entity_name: str = "robot") -> dict[str, ObservationGroupOptions]:
    # The policy reads raw simulator state directly in `step()`. Eden still
    # requires at least one observation term to build the env, so we expose
    # a minimal placeholder.
    return dict(
        policy=ObservationGroupOptions(
            concatenate_terms=True,
            enable_corruption=False,
            dofs_pos=observations.DofsPos.configure(entity_name=entity_name),
        ),
    )

reset

reset(envs_idx: slice | Tensor | None = None) -> None

Clear per-episode state at episode boundaries.

Resets: * History ring buffers (q, qdot, action, base_ang, gravity). Without clearing, the first decoder call after a reset sees the previous episode's last 10 frames and produces a discontinuous initial action. * The heading offset — re-captured per env on the next step() against whatever pose the motion was just resampled to. Critical for sampling_mode != "start" (where the motion frame index changes across episodes) and on motion wrap-around. Only the affected rows are marked for re-capture, so partial resets in batched rollouts leave the other envs' offsets intact.

Source code in eden/extensions/policies/sonic/policy.py

def reset(self, envs_idx: (slice | torch.Tensor) | None = None) -> None:
    """Clear per-episode state at episode boundaries.

    Resets:
    * History ring buffers (``q``, ``qdot``, ``action``, ``base_ang``,
      ``gravity``). Without clearing, the first decoder call after a reset
      sees the previous episode's last 10 frames and produces a
      discontinuous initial action.
    * The heading offset — re-captured per env on the next ``step()``
      against whatever pose the motion was just resampled to. Critical for
      ``sampling_mode != "start"`` (where the motion frame index changes
      across episodes) and on motion wrap-around. Only the affected rows
      are marked for re-capture, so partial resets in batched rollouts
      leave the other envs' offsets intact.
    """
    for buf in self._hist_buffers:
        if envs_idx is None:
            buf.zero_()
        else:
            buf[envs_idx] = 0.0
    # Mark the affected envs for heading re-capture on their next step
    # against the fresh post-reset state.
    if envs_idx is None:
        self._needs_heading_capture.fill_(True)
        self._prev_time_steps = None
    else:
        self._needs_heading_capture[envs_idx] = True

step

step(obs: dict[str, Tensor | dict[str, Tensor]]) -> dict[str, Tensor]

Source code in eden/extensions/policies/sonic/policy.py

def step(self, obs: dict[str, torch.Tensor | dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]:
    # Read off the entity, not an observation group: this policy rides along as an extra deploy
    # slot on another policy's or task's env, whose groups are not its own (see the
    # "Declare the observation" section of docs/guides/policy-porting.md).
    del obs  # raw simulator state is read directly from the env
    robot = self.env.entities[self.entity_name]
    motion_cmd = self.env.command_manager.get_term("motion")

    pelvis_quat_wxyz = self._update_state_history(robot)

    # ---- motion future references (g1 mode encoder input) ----
    time_steps = motion_cmd.time_steps
    # Per-env motion length (multi-clip envs each track their own clip), so
    # the end-of-motion clamp is per-env. Single-clip: a constant vector.
    last_frame = motion_cmd.frame_total_per_env() - 1  # (N,)
    # NOTE: end-of-motion behavior is `clamp` (hold the last frame), not
    # `wrap` — matches the C++ deploy. If the motion wraps via
    # `_update_command`, the wrap-detection block below re-captures the
    # heading offset against the new motion-start pose.
    # `.minimum` (not `.clamp_max`) because the bound is a per-env tensor.
    future_t = (time_steps[:, None] + self._future_offsets[None, :]).minimum(last_frame[:, None])

    # Motion joint pos / vel are fed RAW (no default-angle subtraction) —
    # this matches both gear_sonic's `command_multi_future` (returns
    # `motion_lib.get_dof_pos(...)` directly) and the C++ deploy's
    # `GatherMotionJointPositionsMultiFrame` (raw memcpy from
    # `motion->JointPositions(frame)`). Only the robot's history obs
    # subtracts the default — see the q_rel computation in _update_state_history.
    # Motion data is stored in MJCF order (Eden's convention); reorder to IL.
    future_q_mj = motion_cmd.gather_motion_frames("dofs_pos", future_t)  # (N, F, 29)
    future_qdot_mj = motion_cmd.gather_motion_frames("dofs_vel", future_t)  # (N, F, 29)
    future_q_il = future_q_mj[..., self._mj_to_il]
    future_qdot_il = future_qdot_mj[..., self._mj_to_il]

    # ---- heading alignment (yaw-only) ----
    # `apply_delta_heading` in the C++ deploy =
    #   R(robot_init_yaw) · R(motion_init_yaw)^{-1}
    # Captured per env on the first step after that env's reset AND on that
    # env's motion wrap-around (when `MotionCommand._update_command`
    # resamples and its `time_steps` regresses). The same yaw offset is
    # applied to all future motion anchor quats until the env's next
    # reset/wrap. Without this, the robot's world heading drifts away from
    # the motion's heading and the anchor-orientation obs degrades.
    motion_pelvis_now_wxyz = motion_cmd.gather_motion_frames("links_quat_w", time_steps[:, None])[:, 0, 0]  # (N, 4)
    capture = self._needs_heading_capture.clone()
    if self._prev_time_steps is not None:
        capture |= time_steps < self._prev_time_steps  # per-env wrap detection
    # Compute the offset full-batch and write only the rows that need (re)capture.
    # No `bool(capture.any())` guard: that scalar read forces a per-step GPU→CPU
    # sync on the inference hot path. The masked `torch.where` is already a no-op
    # for the (usual) all-false case, matching CommandManager's unconditional-
    # masked-op convention.
    # `apply_delta_heading = robot_yaw * inv(motion_yaw)`, left-multiplied onto
    # the raw motion anchor quat.
    robot_yaw = yaw_quat(pelvis_quat_wxyz)
    motion_yaw = yaw_quat(motion_pelvis_now_wxyz)
    offset = transform_quat_by_quat(inv_quat(motion_yaw), robot_yaw)
    self._heading_offset_wxyz = torch.where(capture[:, None], offset, self._heading_offset_wxyz)
    self._needs_heading_capture &= ~capture
    self._prev_time_steps = time_steps.clone()

    # Future anchor (pelvis) orientation as 6D rotation in robot pelvis frame:
    #   corrected = apply_delta_heading · raw_ref
    #   rel       = inv(R(robot_pelvis_now)) · corrected
    # then convert to the 6D rep used by gear_sonic's motion_anchor_ori_b_mf.
    #
    # Both operands are (N, F, 4) — `transform_quat_by_quat` broadcasts on
    # the leading axes, so we skip the explicit `.expand(...).reshape(...)`
    # round-trips and only flatten once before the final `_quat_to_6d`.
    ref_pelvis_quat = motion_cmd.gather_motion_frames("links_quat_w", future_t)[..., 0, :]  # (N, F, 4)
    heading = self._heading_offset_wxyz[:, None, :]  # (N, 1, 4)
    cur_pelvis_inv = inv_quat(pelvis_quat_wxyz)[:, None, :]  # (N, 1, 4)
    corrected_ref = transform_quat_by_quat(ref_pelvis_quat, heading)  # (N, F, 4)
    rel_quat = transform_quat_by_quat(corrected_ref, cur_pelvis_inv)  # (N, F, 4)
    rel_anchor_6d = _quat_to_6d(rel_quat).reshape(-1, _FUTURE_LEN * 6)

    # ---- build encoder input (1762 dim, mostly zero for g1 mode) ----
    N = self._encoder_buf.shape[0]
    self._encoder_buf.zero_()
    self._write_slot(
        self._encoder_buf,
        _ENCODER_OFFSETS,
        "encoder_mode_4",
        self._encoder_mode_value.expand(N, -1),
    )
    self._write_slot(
        self._encoder_buf,
        _ENCODER_OFFSETS,
        "motion_joint_positions_10frame_step5",
        future_q_il.reshape(N, _FUTURE_LEN * _NUM_DOFS),
    )
    self._write_slot(
        self._encoder_buf,
        _ENCODER_OFFSETS,
        "motion_joint_velocities_10frame_step5",
        future_qdot_il.reshape(N, _FUTURE_LEN * _NUM_DOFS),
    )
    self._write_slot(self._encoder_buf, _ENCODER_OFFSETS, "motion_anchor_orientation_10frame_step5", rel_anchor_6d)
    # All other slots stay zero — they're for teleop / SMPL modes.

    # ---- encoder → token_state → decoder → action (IL order) ----
    action_il = self._decode(self._encode())
    self._push_action_history(action_il)

    # ---- isaaclab→mjcf permutation + action transform ----
    action_mj = action_il[..., self._il_to_mj]
    joint_pos_targets = self._default_angles_mj + self._action_scale_mj * action_mj
    return {"dofs_pos_controller": joint_pos_targets}

Twist2Policy

Twist2Policy(env: EnvBase, options: PolicyOptions)

Bases: PolicyBase

TWIST2 whole-body motion-tracking controller (G1 29-DOF, ONNX backend).

Source code in eden/extensions/policies/twist2/policy.py

def __init__(self, env: EnvBase, options: PolicyOptions) -> None:
    super().__init__(env, options)

    self.device = env.device
    robot = self.require_entity(self.entity_name)
    # The joint observations and the action term read the robot's DOFs in its own order.
    if list(robot.dofs_name) != list(POLICY_JOINT_NAMES):
        raise ValueError(
            f"Twist2Policy drives the G1's 29 DOFs in {list(POLICY_JOINT_NAMES)} order; this robot's "
            f"dofs_name is {robot.dofs_name}. Build the robot from Twist2Policy.robot_options()."
        )

    # The built robot must carry the knees-bent keyframe in every env: the joint-position
    # observation and the action term are both relative to its default pose.
    expected_default = torch.tensor(DEFAULT_POS, device=env.device)
    built_default = robot.get_default_dofs_pos(POLICY_JOINT_NAMES).to(env.device)  # (num_envs, NUM_JOINTS)
    # ``allclose``'s tolerance, per env: the first env whose row is off is the one named.
    mismatched = ~torch.isclose(built_default, expected_default, atol=1e-6).all(dim=1)
    if mismatched.any().item():
        env_idx = int(mismatched.nonzero()[0].item())
        raise ValueError(
            f"Robot default_dofs_pos does not match TWIST2's knees-bent keyframe in env {env_idx}: "
            f"built={built_default[env_idx].tolist()} expected={expected_default.tolist()}. "
            "Build the robot from Twist2Policy.robot_options()."
        )

    # The env has to be the one this policy declares: a task env publishes its own
    # ``policy`` group, which the actor would read as garbage.
    expected_dim = ACTOR_HISTORY_LENGTH * ACTOR_CURRENT_DIM
    group_dims = env.observation_manager.group_obs_dim
    if group_dims.get(OBS_GROUP) != (expected_dim,):
        raise ValueError(
            f"Twist2Policy reads a {expected_dim}-dim '{OBS_GROUP}' observation group, but this env "
            f"publishes {group_dims}. Build the env from Twist2Policy.observation_options() (a task-less "
            f"`eden inference --policy twist2` does), not from another task's config."
        )

    self._actor = Twist2Actor(self.checkpoint_path, device=str(env.device))

Attributes

reads_env_observations class-attribute instance-attribute

reads_env_observations = True

entity_name class-attribute instance-attribute

entity_name: str = 'robot'

checkpoint_path class-attribute instance-attribute

checkpoint_path: str = 'twist2/pretrained.onnx'

md5_hash class-attribute instance-attribute

md5_hash: str | None = 'a045b72bd49ce0993aa145f8958a5564'

device instance-attribute

device = env.device

Methods:

robot_options classmethod

robot_options() -> G1Dofs29

G1Dofs29 with the knees-bent default pose TWIST2 trains against.

The stock gains / action scale already match mjlab; only the default joint pose differs (Eden ships a home-ish pose), so we override it.

Source code in eden/extensions/policies/twist2/policy.py

@classmethod
def robot_options(cls) -> G1Dofs29:
    """``G1Dofs29`` with the knees-bent default pose TWIST2 trains against.

    The stock gains / action scale already match mjlab; only the default
    joint pose differs (Eden ships a home-ish pose), so we override it.
    """
    return G1Dofs29(default_dofs_pos=dict(KNEES_BENT_DOFS_POS), default_root_pos=(0.0, 0.0, 0.76))

env_options classmethod

env_options() -> dict[str, Any]

Source code in eden/extensions/policies/twist2/policy.py

@classmethod
def env_options(cls) -> dict[str, Any]:
    # 50 Hz control (decimation 2 × sim_dt 0.01) to match the deploy stack.
    return dict(
        sim_dt=0.01,
        sim_substeps=2,
        decimation=2,
        iterations=20,
        ls_iterations=20,
        max_collision_pairs=20,
        batch_dofs_info=False,
        use_gjk_collision=False,
        enable_self_collision=False,
        enable_multi_contact=False,
    )

action_options classmethod

action_options(entity_name: str = 'robot') -> dict[str, ActionTermOptions]

The actor's raw (normalized) action, scaled per joint onto the knees-bent default pose.

Source code in eden/extensions/policies/twist2/policy.py

@classmethod
def action_options(cls, entity_name: str = "robot") -> dict[str, ActionTermOptions]:
    """The actor's raw (normalized) action, scaled per joint onto the knees-bent default pose."""
    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,
        ),
    )

command_options classmethod

command_options(motion_file: str = 'motion.npz', debug_vis: bool = False, entity_name: str = 'robot', ghost_entity_name: str = '', sampling_mode: str = 'start') -> dict[str, CommandTermOptions]

Source code in eden/extensions/policies/twist2/policy.py

@classmethod
def command_options(
    cls,
    motion_file: str = "motion.npz",
    debug_vis: bool = False,
    entity_name: str = "robot",
    ghost_entity_name: str = "",
    sampling_mode: str = "start",
) -> dict[str, CommandTermOptions]:
    return dict(
        motion=MotionCommand.configure(
            resampling_time_range=(1.0e9, 1.0e9),
            debug_vis=debug_vis,
            entity_name=entity_name,
            ghost_entity_name=ghost_entity_name,
            motion_file=motion_file,
            anchor_link_name=_ANCHOR_BODY,
            links_name=_MOTION_BODY_NAMES,
            pose_range={},
            velocity_range={},
            dofs_position_range=(0.0, 0.0),
            sampling_mode=sampling_mode,
        ),
    )

observation_options classmethod

observation_options(entity_name: str = 'robot') -> dict[str, ObservationGroupOptions]

One 127-dim frame, mimic then proprio, stacked into an 11-frame history (backfilled at reset).

Source code in eden/extensions/policies/twist2/policy.py

@classmethod
def observation_options(cls, entity_name: str = "robot") -> dict[str, ObservationGroupOptions]:
    """One 127-dim frame, mimic then proprio, stacked into an 11-frame history (backfilled at reset)."""
    entity = {"entity_name": entity_name}
    return {
        OBS_GROUP: ObservationGroupOptions(
            enable_corruption=False,
            history_length=ACTOR_HISTORY_LENGTH,
            flatten_history_dim=True,
            # mimic (35)
            motion_root=ObsTerm.configure(func=twist2_motion_root),
            motion_dofs_pos=ObsTerm.configure(func=twist2_motion_dofs_pos),
            # proprio (92)
            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),
        )
    }

step

step(obs: dict[str, Tensor | dict[str, Tensor]]) -> dict[str, Tensor]

Source code in eden/extensions/policies/twist2/policy.py

def step(self, obs: dict[str, torch.Tensor | dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]:
    history = obs[OBS_GROUP]  # (N, 11 * 127), oldest first, so the newest frame is the current one
    actor_obs = torch.cat([history[:, -ACTOR_CURRENT_DIM:], history], dim=-1)
    return {"dofs_pos_controller": self._actor(actor_obs)}

Manipulation

OmniResetPolicy

OmniResetPolicy(env: EnvBase, options: PolicyOptions)

Bases: PolicyBase

UWLab's state-based OmniReset expert for one object pair (pair), acting at 10 Hz.

Source code in eden/extensions/policies/omnireset/policy.py

def __init__(self, env: EnvBase, options: PolicyOptions) -> None:
    super().__init__(env, options)
    if self.pair not in EXPERT_CHECKPOINTS:
        raise ValueError(
            f"OmniResetPolicy: no released expert for pair {self.pair!r}; pairs: {list(EXPERT_CHECKPOINTS)}."
        )
    if self.checkpoint_path is None:
        from huggingface_hub import hf_hub_download

        checkpoint, md5 = EXPERT_CHECKPOINTS[self.pair]
        self.checkpoint_path = hf_hub_download(
            C.UWLAB_REPO, checkpoint, repo_type="dataset", revision=C.DATASETS_REVISION
        )
        self._verify_md5(self.checkpoint_path, md5)
    self._actor = OmniResetActor(self.checkpoint_path, device=env.device)
    for name in (ROBOT, INSERTIVE, RECEPTIVE):
        self.require_entity(name)

Attributes

reads_env_observations class-attribute instance-attribute

reads_env_observations = True

pair class-attribute instance-attribute

pair: str = 'peg'

checkpoint_path class-attribute instance-attribute

checkpoint_path: str | None = None

md5_hash class-attribute instance-attribute

md5_hash: str | None = None

Methods:

robot_options classmethod

robot_options() -> RobotOptions

Source code in eden/extensions/policies/omnireset/policy.py

@classmethod
def robot_options(cls) -> RobotOptions:
    return scene_entity_options(C.PAIRS[cls.pair])[ROBOT]  # the same robot for every pair

scene_entity_options classmethod

scene_entity_options(pair: str = 'peg') -> dict[str, EntityOptionsLike]

The workspace and pair's objects (pair is filled from the policy's own).

Source code in eden/extensions/policies/omnireset/policy.py

@classmethod
def scene_entity_options(cls, pair: str = "peg") -> dict[str, EntityOptionsLike]:
    """The workspace and ``pair``'s objects (``pair`` is filled from the policy's own)."""
    if pair not in C.PAIRS:
        raise ValueError(f"OmniResetPolicy: no object pair {pair!r}; pairs: {list(C.PAIRS)}.")
    entities = scene_entity_options(C.PAIRS[pair])
    del entities[ROBOT]
    return entities

env_options classmethod

env_options() -> dict[str, Any]

Source code in eden/extensions/policies/omnireset/policy.py

@classmethod
def env_options(cls) -> dict[str, Any]:
    return dict(sim_dt=C.SIM_DT, sim_substeps=RIGS["uwlab"].sim_substeps, decimation=C.DECIMATION)

action_options classmethod

action_options() -> dict[str, ActionTermOptions]

Source code in eden/extensions/policies/omnireset/policy.py

@classmethod
def action_options(cls) -> dict[str, ActionTermOptions]:
    return action_terms(RIGS["uwlab"])

observation_options classmethod

observation_options() -> dict[str, ObservationGroupOptions]

Return the actor's 43-D group, term for term in POLICY_OBSERVATION_LAYOUT order.

See :data:~eden.extensions.policies.omnireset.constants.POLICY_OBSERVATION_LAYOUT.

Source code in eden/extensions/policies/omnireset/policy.py

@classmethod
def observation_options(cls) -> dict[str, ObservationGroupOptions]:
    """Return the actor's 43-D group, term for term in ``POLICY_OBSERVATION_LAYOUT`` order.

    See :data:`~eden.extensions.policies.omnireset.constants.POLICY_OBSERVATION_LAYOUT`.
    """
    return dict(policy=observation_group(RIGS["uwlab"]))

step

step(obs: dict[str, Tensor | dict[str, Tensor]]) -> dict[str, Tensor]

Source code in eden/extensions/policies/omnireset/policy.py

def step(self, obs: dict[str, torch.Tensor | dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]:
    action = self._actor(obs["policy"])
    return {"arm": action[:, :6], "gripper": action[:, 6:7]}

SimToolRealPolicy

SimToolRealPolicy(env: EnvBase, options: PolicyOptions)

Bases: PolicyBase

Object-centric zero-shot tool-manipulation controller (29-DOF, SAPG-LSTM).

Source code in eden/extensions/policies/simtoolreal/policy.py

def __init__(self, env: EnvBase, options: PolicyOptions) -> None:
    super().__init__(env, options)
    device = env.device
    self.device = device
    self._dtype = torch.float32
    N = env.num_envs

    self._actor = SimToolRealActor(self.checkpoint_path, device=str(device))
    self._state = self._actor.init_state(N)

    arm = self.require_entity(self.entity_name)
    hand = self.require_entity(self.hand_name)
    # Touched, not cached: ``step()`` re-resolves it (deployment hot-swaps entities), but a
    # scene without it must fail here, not one step into the rollout.
    self.require_entity(self.object_name)

    def _idx(values):
        return torch.as_tensor(values, dtype=torch.long, device=device)

    # Resolved at construction so a renamed joint/link raises here, not as a scrambled
    # obs vector on the first step.
    self._arm_dofs_idx = _idx(arm.find_named_dofs_idx_local(list(C.ARM_DOFS), preserve_order=True)[1])
    self._hand_dofs_idx = _idx(hand.find_named_dofs_idx_local(list(C.HAND_DOFS), preserve_order=True)[1])
    self._palm_link_idx = _idx(arm.find_named_links_idx_local([C.PALM_LINK], preserve_order=True)[1])
    self._ft_link_idx = _idx(hand.find_named_links_idx_local(list(C.FINGERTIP_LINKS), preserve_order=True)[1])

    def _t(x):
        return torch.as_tensor(x, dtype=self._dtype, device=device)

    self._q_lower = _t(C.Q_LOWER_LIMITS)
    self._q_upper = _t(C.Q_UPPER_LIMITS)
    self._palm_offset = _t(C.PALM_CENTER_OFFSET)
    self._ft_offset = _t(C.FINGERTIP_OFFSET)  # (3,) same for all 5
    self._kp_corners = _t(C.KEYPOINT_CORNERS)  # (4, 3)
    self._kp_scale = C.OBJECT_BASE_SIZE * C.KEYPOINT_SCALE / 2.0

    # Absolute target buffer (canonical order); prev_action_targets obs.
    self._prev_targets = torch.zeros(N, C.NUM_JOINTS, device=device, dtype=self._dtype)
    self._initialized = False

    # Per-env buffers; override via the setters or a ``goal_pose`` command term.
    self._object_scales = torch.tensor(self.object_scales, device=device, dtype=self._dtype).expand(N, 3).clone()
    self._goal_pos = torch.tensor(C.GOAL_START_POS, device=device, dtype=self._dtype).expand(N, 3).clone()
    self._goal_quat = torch.tensor(C.GOAL_START_QUAT_WXYZ, device=device, dtype=self._dtype).expand(N, 4).clone()

Attributes

entity_name class-attribute instance-attribute

entity_name: str = 'robot'

hand_name class-attribute instance-attribute

hand_name: str = 'hand'

object_name class-attribute instance-attribute

object_name: str = 'object'

goal_command class-attribute instance-attribute

goal_command: str = 'goal_pose'

object_asset class-attribute instance-attribute

object_asset: str = 'long_screwdriver'

object_scales class-attribute instance-attribute

object_scales: tuple[float, float, float] = C.DEFAULT_OBJECT_SCALES

checkpoint_path class-attribute instance-attribute

checkpoint_path: str | None = 'simtoolreal/model.pth'

md5_hash class-attribute instance-attribute

md5_hash: str | None = 'c9f9207125f56b4d4ebaea1e816ab7a8'

device instance-attribute

device = device

Methods:

set_goal

set_goal(pos: Tensor, quat_wxyz: Tensor) -> None

Set the per-env object goal pose (world frame, wxyz).

Source code in eden/extensions/policies/simtoolreal/policy.py

def set_goal(self, pos: torch.Tensor, quat_wxyz: torch.Tensor) -> None:
    """Set the per-env object goal pose (world frame, wxyz)."""
    self._goal_pos = pos.to(self.device, self._dtype)
    self._goal_quat = quat_wxyz.to(self.device, self._dtype)

set_object_scales

set_object_scales(scales: Tensor) -> None

Set the per-env object grasp-bbox scale (the object_scales obs).

Source code in eden/extensions/policies/simtoolreal/policy.py

def set_object_scales(self, scales: torch.Tensor) -> None:
    """Set the per-env object grasp-bbox scale (the object_scales obs)."""
    self._object_scales = scales.to(self.device, self._dtype)

robot_options classmethod

robot_options() -> KukaSharpa

KukaSharpa at SimToolReal's base placement, fingertips grippier than the palm.

The composite already defaults to this port's home pose, gains, armature and gravity compensation. The fingertip friction reproduces the original's modify_asset_frictions (robot 0.5, fingertips 1.5); the whole hand keeps colliding, since the fingers must contact an object all over to wrap it.

Source code in eden/extensions/policies/simtoolreal/policy.py

@classmethod
def robot_options(cls) -> KukaSharpa:
    """``KukaSharpa`` at SimToolReal's base placement, fingertips grippier than the palm.

    The composite already defaults to this port's home pose, gains, armature and gravity
    compensation. The fingertip friction reproduces the original's
    ``modify_asset_frictions`` (robot 0.5, fingertips 1.5); the whole hand keeps colliding,
    since the fingers must contact an object all over to wrap it.
    """
    rig = KukaSharpa(base_pos=C.ROBOT_BASE_POS)
    rig.hand.default_geoms_friction = {link: C.FINGER_TIP_FRICTION for link in C.FINGERTIP_LINKS}
    return rig

scene_entity_options classmethod

scene_entity_options(*, object_asset: str | None = None, vis_mode: str = 'visual') -> dict[str, EntityOptionsLike]

The workspace the policy's observations are defined against: table, tool, goal ghost.

The observation is object-centric (tool keypoints relative to the palm and the goal), so a robot-only scene cannot host this policy. Frictions and placements are SimToolReal's; see constants.

PARAMETER	DESCRIPTION
object_asset	:data:~eden.options.objects.dextoolbench.DEXTOOLBENCH_REGISTRY key to spawn. None reads :attr:object_asset — the base hook must stay callable with no arguments, and restating the literal here could drift from the attribute. TYPE: str | None DEFAULT: None
vis_mode	"collision" renders collision geometry instead of the visual meshes. TYPE: str DEFAULT: 'visual'

Source code in eden/extensions/policies/simtoolreal/policy.py

@classmethod
def scene_entity_options(
    cls, *, object_asset: str | None = None, vis_mode: str = "visual"
) -> dict[str, EntityOptionsLike]:
    """The workspace the policy's observations are defined against: table, tool, goal ghost.

    The observation is object-centric (tool keypoints relative to the palm and the goal),
    so a robot-only scene cannot host this policy. Frictions and placements are
    SimToolReal's; see ``constants``.

    Parameters
    ----------
    object_asset : str | None
        :data:`~eden.options.objects.dextoolbench.DEXTOOLBENCH_REGISTRY` key to spawn.
        ``None`` reads :attr:`object_asset` — the base hook must stay callable with no
        arguments, and restating the literal here could drift from the attribute.
    vis_mode : str
        ``"collision"`` renders collision geometry instead of the visual meshes.
    """
    tool_cls = DEXTOOLBENCH_REGISTRY.get(object_asset or cls.object_asset)
    return {
        "table": BoxOptions(
            size=C.TABLE_SIZE,
            default_root_pos=C.TABLE_POS,
            fixed=True,
            vis_mode=vis_mode,
            material=RigidMaterialOptions(friction=C.TABLE_FRICTION),
        ),
        cls.object_name: tool_cls(
            default_root_pos=C.OBJECT_START_POS,
            default_root_quat=C.OBJECT_START_QUAT_WXYZ,
            vis_mode=vis_mode,
            material=RigidMaterialOptions(friction=C.OBJECT_FRICTION),
        ),
        # Collision-free and not physics-driven: a marker for the goal pose, re-posed each
        # step by whoever moves the goal (the example's sampler; a fixed goal stays put).
        "goal_ghost": tool_cls(
            default_root_pos=C.GOAL_START_POS,
            default_root_quat=C.GOAL_START_QUAT_WXYZ,
            is_fixed_base=True,
            collision=False,
            **dict(GHOST_DEFAULTS),
        ),
    }

env_options classmethod

env_options() -> dict[str, Any]

Source code in eden/extensions/policies/simtoolreal/policy.py

@classmethod
def env_options(cls) -> dict[str, Any]:
    # 60 Hz control (120 Hz sim, decimation 2); softer constraints for contact-rich grasps.
    return dict(sim_dt=C.SIM_DT, decimation=C.DECIMATION, constraint_timeconst=0.02)

action_options classmethod

action_options(entity_name: str = 'robot', hand_name: str = 'hand') -> dict[str, ActionTermOptions]

Source code in eden/extensions/policies/simtoolreal/policy.py

@classmethod
def action_options(cls, entity_name: str = "robot", hand_name: str = "hand") -> dict[str, ActionTermOptions]:
    # One term per entity: a controller drives one entity, so the canonical 29-vector is
    # split here. The PD passes the absolute targets through unchanged.
    def _pd(name: str, dofs: list[str]):
        return actions.ImplicitPDController.configure(
            entity_name=name,
            dofs_name=dofs,
            scale=1.0,
            reference_source=ReferenceSource.ZERO,
        )

    return dict(
        dofs_pos_controller=_pd(entity_name, list(C.ARM_DOFS)),
        hand_dofs_pos_controller=_pd(hand_name, list(C.HAND_DOFS)),
    )

observation_options classmethod

observation_options(entity_name: str = 'robot') -> dict[str, ObservationGroupOptions]

Source code in eden/extensions/policies/simtoolreal/policy.py

@classmethod
def observation_options(cls, entity_name: str = "robot") -> dict[str, ObservationGroupOptions]:
    # The policy reads raw sim state in step(); Eden just needs >=1 obs term to build.
    return dict(
        policy=ObservationGroupOptions(
            concatenate_terms=True,
            enable_corruption=False,
            dofs_pos=observations.DofsPos.configure(entity_name=entity_name),
        ),
    )

reset

reset(envs_idx: slice | Tensor | None = None) -> None

Source code in eden/extensions/policies/simtoolreal/policy.py

def reset(self, envs_idx: (slice | torch.Tensor) | None = None) -> None:
    h, c = self._state
    if envs_idx is None:
        h.zero_()
        c.zero_()
        self._prev_targets.zero_()
        self._initialized = False
    else:
        h[:, envs_idx] = 0.0
        c[:, envs_idx] = 0.0
        self._prev_targets[envs_idx] = 0.0

step

step(obs: dict[str, Tensor | dict[str, Tensor]]) -> dict[str, Tensor]

Source code in eden/extensions/policies/simtoolreal/policy.py

def step(self, obs: dict[str, torch.Tensor | dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]:
    # Read off the entity, not an observation group: this policy rides along as an extra deploy
    # slot on another policy's or task's env, whose groups are not its own (see the
    # "Declare the observation" section of docs/guides/policy-porting.md).
    del obs  # raw sim state + object/goal pose read directly
    env = self.env
    arm = env.entities[self.entity_name]
    hand = env.entities[self.hand_name]
    obj = env.entities[self.object_name]

    # LOAD-BEARING ``relative=True``: the default solver frame re-origins a rigid body
    # onto its inertial frame, 3-15 cm down these tools' handles. IsaacGym and MuJoCo
    # report the URDF-root pose, so that is what the policy trained on. See
    # docs/guides/entities.md.
    object_pos = obj.get_pos(relative=True)
    object_quat = obj.get_quat(relative=True)

    # Prefer a ``goal_pose`` command term if the task has one, else the ``set_goal`` buffer.
    try:
        goal_cmd = env.command_manager.get_term(self.goal_command)
        self._goal_pos, self._goal_quat = goal_cmd.goal_pos, goal_cmd.goal_quat
    except Exception:
        pass

    # Seed the EMA from where the robot actually is.
    if not self._initialized:
        self._prev_targets = self._read_q(arm, hand)[0].clone()
        self._initialized = True

    obs_vec = self._build_obs(
        arm, hand, object_pos, object_quat, self._goal_pos, self._goal_quat, self._object_scales
    )
    action, self._state = self._actor(obs_vec, self._state)
    targets = self._apply_action(action)
    # Split the canonical 29-vector back out to the two controllers.
    return {
        "dofs_pos_controller": targets[:, : C.NUM_ARM_JOINTS],
        "hand_dofs_pos_controller": targets[:, C.NUM_ARM_JOINTS :],
    }

Remote policy clients

PiPolicy

PiPolicy(env: EnvBase, options: PolicyOptions)

Bases: SetupClientMixin, PolicyBase, ConfigurableMixin[RemotePolicyOptions]

OpenPI (pi0 / pi0-FAST) remote inference client.

All class attributes are configurable via PiPolicy.configure(...).

Source code in eden/extensions/policies/client.py

def __init__(self, env: EnvBase, options: PolicyOptions):
    super().__init__(env, options)
    if not self.supports_batch and env.num_envs != 1:
        raise ValueError(
            f"{type(self).__name__} has supports_batch=False and only supports "
            f"num_envs=1 (got {env.num_envs}). Set supports_batch=True on the "
            f"subclass if the remote server accepts batched requests."
        )
    if env.num_envs > 1 and type(self).reset is ClientMixin.reset:
        # The env resets its envs one by one; the default reset can only reset the shared
        # server state of all of them, so it would raise at the first episode end.
        raise ValueError(
            f"{type(self).__name__} serves {env.num_envs} envs but keeps ClientMixin.reset, which "
            "refuses a partial reset because the server's episode state is shared by every env. Fix: "
            "override reset(envs_idx) with a per-env scheme (e.g. one server session per env), or "
            "run with num_envs=1."
        )
    self._num_envs: int = env.num_envs
    self._post_configure()
    self._scheduler = ActionChunkScheduler(
        num_envs=env.num_envs,
        mode=self.chunk_mode,
        pred_action_horizon=self.pred_action_horizon,
        exec_horizon=self.exec_horizon,
        short_chunk=self.short_chunk,
        ensemble_steps=self.chunk_temporal_ensemble_steps,
        ensemble_dims=self.chunk_temporal_ensemble_dims,
    )
    self.last_response_info: dict[str, Any] = {}
    self._transport: TransportBase = self.make_transport()
    self._post_connect()

Attributes

SETUP_REGISTRY class-attribute

SETUP_REGISTRY: dict[str, type[PolicySetupBase]] = SETUP_REGISTRY

host class-attribute instance-attribute

host: str = 'localhost'

port class-attribute instance-attribute

port: int | None = 8000

api_key class-attribute instance-attribute

api_key: str | None = None

ws_ping_timeout class-attribute instance-attribute

ws_ping_timeout: float | None = None

transport class-attribute instance-attribute

transport: str = 'openpi_websocket'

response_action_key class-attribute instance-attribute

response_action_key: str = 'actions'

policy_setup class-attribute instance-attribute

policy_setup: str = 'widowx_bridge'

task_description class-attribute instance-attribute

task_description: str | None = None

pred_action_horizon class-attribute instance-attribute

pred_action_horizon: int | None = None

exec_horizon class-attribute instance-attribute

exec_horizon: int | None = None

chunk_temporal_ensemble_steps class-attribute instance-attribute

chunk_temporal_ensemble_steps: int = 0

chunk_temporal_ensemble_dims class-attribute instance-attribute

chunk_temporal_ensemble_dims: int | None = None

image_size class-attribute instance-attribute

image_size: tuple[int, int] = (224, 224)

use_wrist_image class-attribute instance-attribute

use_wrist_image: bool | None = None

aloha_left_camera_name class-attribute instance-attribute

aloha_left_camera_name: str = 'cam_left_wrist'

droid_arm_controller_kind class-attribute instance-attribute

droid_arm_controller_kind: str = 'velocity'

droid_velocity_lag class-attribute instance-attribute

droid_velocity_lag: dict | None = None

droid_velocity_compensation class-attribute instance-attribute

droid_velocity_compensation: dict | None = None

droid_gripper_threshold class-attribute instance-attribute

droid_gripper_threshold: float = 0.5

droid_debug_gripper class-attribute instance-attribute

droid_debug_gripper: bool = False

Gr00tPolicy

Gr00tPolicy(env: EnvBase, options: PolicyOptions)

Bases: SetupClientMixin, PolicyBase, ConfigurableMixin[RemotePolicyOptions]

NVIDIA Isaac-GR00T remote-inference client.

All class attributes are configurable via :meth:Gr00tPolicy.configure.

Source code in eden/extensions/policies/client.py

def __init__(self, env: EnvBase, options: PolicyOptions):
    super().__init__(env, options)
    if not self.supports_batch and env.num_envs != 1:
        raise ValueError(
            f"{type(self).__name__} has supports_batch=False and only supports "
            f"num_envs=1 (got {env.num_envs}). Set supports_batch=True on the "
            f"subclass if the remote server accepts batched requests."
        )
    if env.num_envs > 1 and type(self).reset is ClientMixin.reset:
        # The env resets its envs one by one; the default reset can only reset the shared
        # server state of all of them, so it would raise at the first episode end.
        raise ValueError(
            f"{type(self).__name__} serves {env.num_envs} envs but keeps ClientMixin.reset, which "
            "refuses a partial reset because the server's episode state is shared by every env. Fix: "
            "override reset(envs_idx) with a per-env scheme (e.g. one server session per env), or "
            "run with num_envs=1."
        )
    self._num_envs: int = env.num_envs
    self._post_configure()
    self._scheduler = ActionChunkScheduler(
        num_envs=env.num_envs,
        mode=self.chunk_mode,
        pred_action_horizon=self.pred_action_horizon,
        exec_horizon=self.exec_horizon,
        short_chunk=self.short_chunk,
        ensemble_steps=self.chunk_temporal_ensemble_steps,
        ensemble_dims=self.chunk_temporal_ensemble_dims,
    )
    self.last_response_info: dict[str, Any] = {}
    self._transport: TransportBase = self.make_transport()
    self._post_connect()

Attributes

SETUP_REGISTRY class-attribute

SETUP_REGISTRY: dict[str, type[PolicySetupBase]] = SETUP_REGISTRY

host class-attribute instance-attribute

host: str = 'localhost'

port class-attribute instance-attribute

port: int | None = 5555

timeout_ms class-attribute instance-attribute

timeout_ms: int = 60000

api_token class-attribute instance-attribute

api_token: str | None = None

transport class-attribute instance-attribute

transport: str = 'zmq_msgpack'

reset_endpoint class-attribute instance-attribute

reset_endpoint: str | None = 'reset'

short_chunk class-attribute instance-attribute

short_chunk: str = 'accept'

policy_setup class-attribute instance-attribute

policy_setup: str = 'simpler_env_widowx'

use_sim_policy_wrapper class-attribute instance-attribute

use_sim_policy_wrapper: bool | None = None

task_description class-attribute instance-attribute

task_description: str | None = None

pred_action_horizon class-attribute instance-attribute

pred_action_horizon: int | None = None

exec_horizon class-attribute instance-attribute

exec_horizon: int | None = None

image_size class-attribute instance-attribute

image_size: tuple[int, int] | None = None

video_horizon class-attribute instance-attribute

video_horizon: int | None = None

video_stride class-attribute instance-attribute

video_stride: int | None = None

debug_gripper class-attribute instance-attribute

debug_gripper: bool = False

Methods:

observe

observe(obs: dict) -> None

Refuse an observation dict without a group the setup reads, then run the setup's per-step hook.

Construction checks the env publishes every group in observation_groups; this catches a dict the caller built by hand (a dropped wrist_image, a renamed group) on the step it arrives, naming the fix.

Source code in eden/extensions/policies/gr00t/policy.py

def observe(self, obs: dict) -> None:
    """Refuse an observation dict without a group the setup reads, then run the setup's per-step hook.

    Construction checks the env publishes every group in ``observation_groups``; this catches a dict the caller
    built by hand (a dropped ``wrist_image``, a renamed group) on the step it arrives, naming the fix.
    """
    missing = [group for group in self._setup.observation_groups if group not in obs]
    if missing:
        raise KeyError(
            f"Gr00tPolicy.step() got observation groups {sorted(obs)}, without {missing}, which "
            f"policy_setup={self.policy_setup!r} reads. Fix: pass the dict env.reset() / env.step() return for an "
            f"env built from Gr00tPolicy.observation_options(policy_setup={self.policy_setup!r}). Do not fill a "
            "missing view with a copy of another camera: the checkpoint's schema has a distinct view per key."
        )
    super().observe(obs)

pack_observations

pack_observations(obs: dict) -> dict

Source code in eden/extensions/policies/gr00t/policy.py

def pack_observations(self, obs: dict) -> dict:
    payload = super().pack_observations(obs)
    return to_sim_wrapper_request(payload) if self.use_sim_policy_wrapper else payload

infer_payload

infer_payload(payload: Any) -> ndarray

Round-trip as the base does; a server rejecting the request's format gets the flag to flip named.

Source code in eden/extensions/policies/gr00t/policy.py

def infer_payload(self, payload: Any) -> np.ndarray:
    """Round-trip as the base does; a server rejecting the request's format gets the flag to flip named."""
    try:
        return super().infer_payload(payload)
    except ServerError as e:
        hint = wire_format_hint(str(e), use_sim_policy_wrapper=self.use_sim_policy_wrapper)
        if hint is None:
            raise
        raise ServerError(f"{e}\n{hint}") from e

Cosmos3Policy

Cosmos3Policy(env: EnvBase, options: PolicyOptions)

Bases: ClientMixin, PolicyBase, ConfigurableMixin[RemotePolicyOptions]

Cosmos3-Nano-Policy-DROID remote inference client (WebSocket).

All class attributes are configurable via Cosmos3Policy.configure(...).

Source code in eden/extensions/policies/client.py

def __init__(self, env: EnvBase, options: PolicyOptions):
    super().__init__(env, options)
    if not self.supports_batch and env.num_envs != 1:
        raise ValueError(
            f"{type(self).__name__} has supports_batch=False and only supports "
            f"num_envs=1 (got {env.num_envs}). Set supports_batch=True on the "
            f"subclass if the remote server accepts batched requests."
        )
    if env.num_envs > 1 and type(self).reset is ClientMixin.reset:
        # The env resets its envs one by one; the default reset can only reset the shared
        # server state of all of them, so it would raise at the first episode end.
        raise ValueError(
            f"{type(self).__name__} serves {env.num_envs} envs but keeps ClientMixin.reset, which "
            "refuses a partial reset because the server's episode state is shared by every env. Fix: "
            "override reset(envs_idx) with a per-env scheme (e.g. one server session per env), or "
            "run with num_envs=1."
        )
    self._num_envs: int = env.num_envs
    self._post_configure()
    self._scheduler = ActionChunkScheduler(
        num_envs=env.num_envs,
        mode=self.chunk_mode,
        pred_action_horizon=self.pred_action_horizon,
        exec_horizon=self.exec_horizon,
        short_chunk=self.short_chunk,
        ensemble_steps=self.chunk_temporal_ensemble_steps,
        ensemble_dims=self.chunk_temporal_ensemble_dims,
    )
    self.last_response_info: dict[str, Any] = {}
    self._transport: TransportBase = self.make_transport()
    self._post_connect()

Attributes

action_dim class-attribute

action_dim: int | None = 8

host class-attribute instance-attribute

host: str = 'localhost'

port class-attribute instance-attribute

port: int | None = 8000

api_key class-attribute instance-attribute

api_key: str | None = None

ws_ping_timeout class-attribute instance-attribute

ws_ping_timeout: float | None = None

transport class-attribute instance-attribute

transport: str = 'openpi_websocket'

response_action_key class-attribute instance-attribute

response_action_key: str = 'action'

pred_action_horizon class-attribute instance-attribute

pred_action_horizon: int | None = 32

exec_horizon class-attribute instance-attribute

exec_horizon: int | None = 32

image_size class-attribute instance-attribute

image_size: tuple[int, int] = (360, 640)

secondary_camera_name class-attribute instance-attribute

secondary_camera_name: str | None = None

duplicate_exterior_view class-attribute instance-attribute

duplicate_exterior_view: bool = False

gripper_threshold class-attribute instance-attribute

gripper_threshold: float = 0.5

debug_gripper class-attribute instance-attribute

debug_gripper: bool = False

Methods:

observation_options classmethod

observation_options(camera_name: str = 'camera', wrist_camera_name: str = 'wrist_camera', secondary_camera_name: str | None = None, arm_entity_name: str = 'robot', arm_dofs_name: list[str] | None = None, gripper_entity_name: str | None = None, gripper_dof_name: str = 'left_driver_joint', gripper_close_pos: float = 0.8) -> dict[str, ObservationGroupOptions]

Wire the DROID joint-state proprio + exterior/wrist camera groups.

Mirrors OpenPI's droid_jointpos observation contract: an 8-D proprio vector [joint_pos(7), gripper(1)] plus an image (exterior) and wrist_image group. When secondary_camera_name is given, an image_2 group carries the second exterior view; pass the same value to Cosmos3Policy.configure() (the policy checks the two agree when it is constructed).

Source code in eden/extensions/policies/cosmos3/policy.py

@classmethod
def observation_options(
    cls,
    camera_name: str = "camera",
    wrist_camera_name: str = "wrist_camera",
    secondary_camera_name: str | None = None,
    arm_entity_name: str = "robot",
    arm_dofs_name: list[str] | None = None,
    gripper_entity_name: str | None = None,
    gripper_dof_name: str = "left_driver_joint",
    gripper_close_pos: float = 0.8,
) -> dict[str, ObservationGroupOptions]:
    """Wire the DROID joint-state proprio + exterior/wrist camera groups.

    Mirrors OpenPI's ``droid_jointpos`` observation contract: an 8-D
    ``proprio`` vector ``[joint_pos(7), gripper(1)]`` plus an ``image``
    (exterior) and ``wrist_image`` group.  When ``secondary_camera_name``
    is given, an ``image_2`` group carries the second exterior view; pass
    the same value to ``Cosmos3Policy.configure()`` (the policy checks the
    two agree when it is constructed).
    """
    from eden.managers.terms.observations.camera import rgb_image
    from eden.options.managers.observations import ObservationGroupOptions
    from eden.terms import ObsTerm

    if arm_dofs_name is None:
        arm_dofs_name = list(DROID_ARM_DOFS)

    groups = make_image_groups(camera_name, use_wrist_image=True, wrist_camera_name=wrist_camera_name)
    if secondary_camera_name is not None:
        groups["image_2"] = ObservationGroupOptions(
            rgb=ObsTerm.configure(func=rgb_image, params={"camera_name": secondary_camera_name}),
        )

    groups["proprio"] = make_droid_joint_state_group(
        arm_entity_name,
        arm_dofs_name,
        gripper_entity_name=gripper_entity_name,
        gripper_dof_name=gripper_dof_name,
        gripper_close_pos=gripper_close_pos,
    )
    return groups

action_options classmethod

action_options(arm_entity_name: str = 'robot', arm_dofs_name: list[str] | None = None, gripper_entity_name: str | None = None, gripper_dof_name: str = 'left_driver_joint', gripper_close_pos: float = 0.8, gripper_threshold: float = 0.5) -> dict[str, ActionTermOptions]

Absolute joint-position arm controller + GripperDecisionController.

The server returns absolute arm joint targets, so the arm controller is an ImplicitPDController with a zero reference (targets are applied directly). The gripper term receives the server's closedness and binarizes it at gripper_threshold (> threshold closes).

Source code in eden/extensions/policies/cosmos3/policy.py

@classmethod
def action_options(
    cls,
    arm_entity_name: str = "robot",
    arm_dofs_name: list[str] | None = None,
    gripper_entity_name: str | None = None,
    gripper_dof_name: str = "left_driver_joint",
    gripper_close_pos: float = 0.8,
    gripper_threshold: float = 0.5,
) -> dict[str, ActionTermOptions]:
    """Absolute joint-position arm controller + ``GripperDecisionController``.

    The server returns absolute arm joint targets, so the arm controller is
    an ``ImplicitPDController`` with a zero reference (targets are applied
    directly). The gripper term receives the server's closedness and binarizes
    it at ``gripper_threshold`` (``> threshold`` closes).
    """
    return make_droid_joint_position_actions(
        arm_entity_name,
        arm_dofs_name,
        gripper_entity_name=gripper_entity_name,
        gripper_dof_name=gripper_dof_name,
        arm_key=_ARM_TERM,
        gripper_close_pos=gripper_close_pos,
        threshold=gripper_threshold,
    )

pack_observations

pack_observations(obs: dict) -> dict

Convert Eden obs tensors to the Cosmos3 RoboLab server payload dict.

Source code in eden/extensions/policies/cosmos3/policy.py

def pack_observations(self, obs: dict) -> dict:
    """Convert Eden obs tensors to the Cosmos3 RoboLab server payload dict."""
    if self.task_description is None:
        raise ValueError(
            "Cosmos3Policy sends 'prompt' with every request, but Cosmos3Policy.task_description is None. Fix: "
            "policy.set_task('<instruction>') before the first step, or "
            "Cosmos3Policy.configure(task_description='<instruction>'). Do not fill in a generic instruction: "
            "the checkpoint is language-conditioned and acts on whatever it is told."
        )
    proprio = to_numpy_unbatched(obs["proprio"])
    joint_pos = proprio[:7].astype(np.float32)
    gripper_closedness = proprio[7:8].astype(np.float32)

    payload: dict = {
        "prompt": self.task_description,
        "observation/joint_position": joint_pos,
        "observation/gripper_position": gripper_closedness,
    }

    # RoboArena concat view: wrist (top) + two exterior views (bottom), composed by the
    # server. The groups were checked against this policy at construction.
    exterior_img = prepare_image_padded(obs["image"], self.image_size)
    if self.secondary_camera_name is not None:
        exterior2_img = prepare_image_padded(obs["image_2"], self.image_size)
    else:  # duplicate_exterior_view=True
        exterior2_img = exterior_img
    payload["observation/wrist_image_left"] = prepare_image_padded(obs["wrist_image"], self.image_size)
    payload["observation/exterior_image_1_left"] = exterior_img
    payload["observation/exterior_image_2_left"] = exterior2_img

    return payload

unpack_actions

unpack_actions(raw_action: ndarray) -> dict[str, Tensor]

Decode one (1, 8) action row → Eden controller dict.

Layout: [joint_pos(7), gripper(1)]. Gripper convention: the server already re-applies its 1 - x transform, so the scalar is closedness (> 0.5 ⇒ CLOSE). Delegated to JointPositionAdapter.

Source code in eden/extensions/policies/cosmos3/policy.py

def unpack_actions(self, raw_action: np.ndarray) -> dict[str, torch.Tensor]:
    """Decode one ``(1, 8)`` action row → Eden controller dict.

    Layout: ``[joint_pos(7), gripper(1)]``. Gripper convention: the server
    already re-applies its ``1 - x`` transform, so the scalar is closedness
    (``> 0.5`` ⇒ CLOSE). Delegated to ``JointPositionAdapter``.
    """
    return self._adapter(np.asarray(raw_action, dtype=np.float32).reshape(-1), env=self.env)

GalaxeaVlaPolicy

GalaxeaVlaPolicy(env: EnvBase, options: PolicyOptions)

Bases: SetupClientMixin, PolicyBase, ConfigurableMixin[RemotePolicyOptions]

GalaxeaVLA (G0.5) remote WebSocket inference client.

All class attributes are configurable via GalaxeaVlaPolicy.configure(...).

Source code in eden/extensions/policies/client.py

def __init__(self, env: EnvBase, options: PolicyOptions):
    super().__init__(env, options)
    if not self.supports_batch and env.num_envs != 1:
        raise ValueError(
            f"{type(self).__name__} has supports_batch=False and only supports "
            f"num_envs=1 (got {env.num_envs}). Set supports_batch=True on the "
            f"subclass if the remote server accepts batched requests."
        )
    if env.num_envs > 1 and type(self).reset is ClientMixin.reset:
        # The env resets its envs one by one; the default reset can only reset the shared
        # server state of all of them, so it would raise at the first episode end.
        raise ValueError(
            f"{type(self).__name__} serves {env.num_envs} envs but keeps ClientMixin.reset, which "
            "refuses a partial reset because the server's episode state is shared by every env. Fix: "
            "override reset(envs_idx) with a per-env scheme (e.g. one server session per env), or "
            "run with num_envs=1."
        )
    self._num_envs: int = env.num_envs
    self._post_configure()
    self._scheduler = ActionChunkScheduler(
        num_envs=env.num_envs,
        mode=self.chunk_mode,
        pred_action_horizon=self.pred_action_horizon,
        exec_horizon=self.exec_horizon,
        short_chunk=self.short_chunk,
        ensemble_steps=self.chunk_temporal_ensemble_steps,
        ensemble_dims=self.chunk_temporal_ensemble_dims,
    )
    self.last_response_info: dict[str, Any] = {}
    self._transport: TransportBase = self.make_transport()
    self._post_connect()

Attributes

SETUP_REGISTRY class-attribute

SETUP_REGISTRY: dict[str, type[PolicySetupBase]] = SETUP_REGISTRY

host class-attribute instance-attribute

host: str = 'localhost'

port class-attribute instance-attribute

port: int | None = 8765

api_key class-attribute instance-attribute

api_key: str | None = None

ws_timeout_s class-attribute instance-attribute

ws_timeout_s: float = 60.0

transport class-attribute instance-attribute

transport: str = 'galaxea_websocket'

reset_endpoint class-attribute instance-attribute

reset_endpoint: str | None = '__reset__'

chunk_mode class-attribute instance-attribute

chunk_mode: str = 'server'

short_chunk class-attribute instance-attribute

short_chunk: str = 'accept'

policy_setup class-attribute instance-attribute

policy_setup: str = 'galaxea_r1lite'

task_description class-attribute instance-attribute

task_description: str = 'pick up the object'

control_frequency class-attribute instance-attribute

control_frequency: float = 15.0

gripper_close_threshold class-attribute instance-attribute

gripper_close_threshold: float = DROID_GRIPPER_THRESHOLD_DEFAULTS['gripper_close_threshold']

gripper_decode_mode class-attribute instance-attribute

gripper_decode_mode: str = 'continuous-snap'

gripper_snap_threshold class-attribute instance-attribute

gripper_snap_threshold: float = DROID_GRIPPER_THRESHOLD_DEFAULTS['gripper_snap_threshold']

debug_gripper class-attribute instance-attribute

debug_gripper: bool = False

MolmoAct2Policy

MolmoAct2Policy(env: EnvBase, options: PolicyOptions)

Bases: ClientMixin, PolicyBase, ConfigurableMixin[RemotePolicyOptions]

Client for the MolmoAct2 DROID inference server (absolute joint control).

Source code in eden/extensions/policies/client.py

def __init__(self, env: EnvBase, options: PolicyOptions):
    super().__init__(env, options)
    if not self.supports_batch and env.num_envs != 1:
        raise ValueError(
            f"{type(self).__name__} has supports_batch=False and only supports "
            f"num_envs=1 (got {env.num_envs}). Set supports_batch=True on the "
            f"subclass if the remote server accepts batched requests."
        )
    if env.num_envs > 1 and type(self).reset is ClientMixin.reset:
        # The env resets its envs one by one; the default reset can only reset the shared
        # server state of all of them, so it would raise at the first episode end.
        raise ValueError(
            f"{type(self).__name__} serves {env.num_envs} envs but keeps ClientMixin.reset, which "
            "refuses a partial reset because the server's episode state is shared by every env. Fix: "
            "override reset(envs_idx) with a per-env scheme (e.g. one server session per env), or "
            "run with num_envs=1."
        )
    self._num_envs: int = env.num_envs
    self._post_configure()
    self._scheduler = ActionChunkScheduler(
        num_envs=env.num_envs,
        mode=self.chunk_mode,
        pred_action_horizon=self.pred_action_horizon,
        exec_horizon=self.exec_horizon,
        short_chunk=self.short_chunk,
        ensemble_steps=self.chunk_temporal_ensemble_steps,
        ensemble_dims=self.chunk_temporal_ensemble_dims,
    )
    self.last_response_info: dict[str, Any] = {}
    self._transport: TransportBase = self.make_transport()
    self._post_connect()

Attributes

action_dim class-attribute

action_dim: int | None = 8

host class-attribute instance-attribute

host: str = 'localhost'

port class-attribute instance-attribute

port: int | None = 8000

act_endpoint class-attribute instance-attribute

act_endpoint: str = '/act'

set_task_endpoint class-attribute instance-attribute

set_task_endpoint: str | None = None

reset_endpoint class-attribute instance-attribute

reset_endpoint: str | None = None

transport class-attribute instance-attribute

transport: str = 'http_json_numpy'

response_action_key class-attribute instance-attribute

response_action_key: str = 'actions'

pred_action_horizon class-attribute instance-attribute

pred_action_horizon: int | None = 15

request_timeout class-attribute instance-attribute

request_timeout: float = 120.0

num_steps class-attribute instance-attribute

num_steps: int = 10

enable_cuda_graph class-attribute instance-attribute

enable_cuda_graph: bool = False

binarize_gripper class-attribute instance-attribute

binarize_gripper: bool = True

gripper_threshold class-attribute instance-attribute

gripper_threshold: float = 0.5

flip_gripper class-attribute instance-attribute

flip_gripper: bool = False

Methods:

pack_observations

pack_observations(obs: dict[str, Tensor | dict[str, Tensor]]) -> dict

Source code in eden/extensions/policies/molmoact2/policy.py

def pack_observations(
    self,
    obs: dict[str, torch.Tensor | dict[str, torch.Tensor]],
) -> dict:
    if self.task_description is None:
        raise ValueError(
            "MolmoAct2Policy sends 'instruction' with every request, but MolmoAct2Policy.task_description is "
            "None. Fix: policy.set_task('<instruction>') before the first step, or "
            "MolmoAct2Policy.configure(task_description='<instruction>'). Do not fill in a generic "
            "instruction: the checkpoint is language-conditioned and acts on whatever it is told."
        )
    missing = [group for group in ("image", "proprio") if group not in obs]
    if missing:
        raise KeyError(
            f"MolmoAct2Policy reads the observation groups ['image', 'proprio'], but obs has no {missing} (its "
            f"groups are {sorted(obs)}). Fix: pass the observations of an env built from "
            "MolmoAct2Policy.observation_options(...) as they are. Do not fill a missing group with zeros: the "
            "server would act on a blank view or a zero joint state."
        )
    image_group = obs["image"]
    if not isinstance(image_group, dict):
        raise TypeError(
            f"MolmoAct2Policy expects obs['image'] to be a dict with 'external' and 'wrist' keys, but it is a "
            f"{type(image_group).__name__}. Fix: build the env's 'image' group from "
            "MolmoAct2Policy.observation_options(...), which keeps the two cameras apart (concatenate_terms=False)."
        )
    if "external" not in image_group or "wrist" not in image_group:
        raise KeyError(
            "MolmoAct2Policy sends obs['image']['external'] and obs['image']['wrist'] as its two views, but "
            f"obs['image'] holds {sorted(image_group)}. Fix: build the env's 'image' group from "
            "MolmoAct2Policy.observation_options(external_camera=..., wrist_camera=...), which names the env's "
            "cameras. Do not send one camera as both views: the checkpoint was trained on two."
        )

    proprio = to_numpy_unbatched(obs["proprio"]).astype(np.float32)
    if proprio.shape != (8,):
        raise ValueError(f"MolmoAct2 expects 8-D proprio [q1..q7, gripper], got {proprio.shape}.")

    payload = {
        "external_cam": to_uint8_hwc(image_group["external"]),
        "wrist_cam": to_uint8_hwc(image_group["wrist"]),
        "instruction": self.task_description,
        "state": proprio,
        "num_steps": int(self.num_steps),
    }
    if self.enable_cuda_graph:
        payload["enable_cuda_graph"] = True
    return payload

accept_chunk

accept_chunk(chunk: ndarray) -> None

Refuse a chunk with any non-finite arm target before its first row executes, then queue it.

Every row of the reply is scanned, including any past the 15 that execute: a non-finite value anywhere means the server's output is broken, so the rows before it are not trusted either.

Source code in eden/extensions/policies/molmoact2/policy.py

def accept_chunk(self, chunk: np.ndarray) -> None:
    """Refuse a chunk with any non-finite arm target before its first row executes, then queue it.

    Every row of the reply is scanned, including any past the 15 that execute: a non-finite value anywhere
    means the server's output is broken, so the rows before it are not trusted either.
    """
    chunk = self._normalize_action_shape(np.asarray(chunk, dtype=np.float32))
    self._check_action_width(chunk, self.action_dim, schema=type(self).__name__)
    # Arm columns only, on the whole (num_envs, T, 8) chunk: a NaN gripper alone holds the gripper.
    bad = ~np.isfinite(chunk[:, :, :7]).all(axis=2)
    if bad.any():
        env_idx, row_idx = (int(i) for i in np.argwhere(bad)[0])
        raise ValueError(
            f"MolmoAct2 server returned non-finite arm joint targets for env {env_idx} in row {row_idx} of its "
            f"{chunk.shape[1]}-row chunk: {[round(float(v), 6) for v in chunk[env_idx, row_idx]]} ([q1..q7, "
            f"gripper], rad). The whole chunk is refused before any row reaches the arm term {_ARM_TERM!r} "
            "(control_dofs_pos). Fix: check the served checkpoint and its norm_stats.json (franka_droid) on the "
            "server; a NaN gripper alone is allowed (it holds the gripper). Do not replace the bad values with "
            "the current or a default pose, or execute the rows before it: the whole chunk was predicted from "
            "the same broken output."
        )
    super().accept_chunk(chunk)

unpack_actions

unpack_actions(raw_action: ndarray) -> dict[str, Tensor]

Decode an (num_envs, 8) MolmoAct2-DROID action row.

raw_action[:, :7] are absolute Franka joint targets in rad and are passed straight through to the arm ImplicitPDController (scale=1). raw_action[:, 7] is gripper closedness in [0, 1] (0=open, 1=closed), forwarded raw to the env's GripperDecisionController, which binarizes it at gripper_threshold (or follows it with binarize_gripper=False).

Stores self.last_diagnostics with the raw gripper scalar, threshold, and the intent the env's gripper term derives from it, so callers can record/plot gripper traces.

Source code in eden/extensions/policies/molmoact2/policy.py

def unpack_actions(self, raw_action: np.ndarray) -> dict[str, torch.Tensor]:
    """Decode an ``(num_envs, 8)`` MolmoAct2-DROID action row.

    ``raw_action[:, :7]`` are absolute Franka joint targets in rad and are
    passed straight through to the arm ``ImplicitPDController`` (scale=1).
    ``raw_action[:, 7]`` is gripper **closedness in ``[0, 1]``** (0=open,
    1=closed), forwarded raw to the env's ``GripperDecisionController``, which
    binarizes it at ``gripper_threshold`` (or follows it with ``binarize_gripper=False``).

    Stores ``self.last_diagnostics`` with the raw gripper scalar, threshold, and the intent
    the env's gripper term derives from it, so callers can record/plot gripper traces.
    """
    device = self.env.device
    arm = torch.as_tensor(raw_action[:, :7], dtype=torch.float32, device=device)
    # The raw closedness goes to the env's GripperDecisionController, configured by
    # action_options() with the same flip / threshold / binarize semantics.
    gripper = torch.as_tensor(raw_action[:, 7:8], dtype=torch.float32, device=device)

    raw_g = float(raw_action[0, 7])
    self.last_diagnostics = {
        "raw_gripper_action": raw_g,
        "gripper_threshold": float(self.gripper_threshold),
        "gripper_intent": self._gripper_intent(raw_g),
        "gripper_command": raw_g,
        "binarize_gripper": bool(self.binarize_gripper),
        "flip_gripper": bool(self.flip_gripper),
    }
    return {_ARM_TERM: arm, "gripper_controller": gripper}

action_options classmethod

action_options(arm_entity_name: str = 'robot', arm_dofs_name: list[str] | None = None, gripper_entity_name: str | None = None, gripper_dof_name: str = 'left_driver_joint', gripper_close_pos: float = 0.8, gripper_threshold: float = 0.5, binarize_gripper: bool = True, flip_gripper: bool = False) -> dict[str, ActionTermOptions]

Absolute joint arm controller + GripperDecisionController fed the server's closedness.

flip_gripper treats the scalar as openness instead; binarize_gripper=False drives the fingers continuously by the (clipped) closedness. Keep these equal to the policy's attributes of the same name (checked at policy construction).

Source code in eden/extensions/policies/molmoact2/policy.py

@classmethod
def action_options(
    cls,
    arm_entity_name: str = "robot",
    arm_dofs_name: list[str] | None = None,
    gripper_entity_name: str | None = None,
    gripper_dof_name: str = "left_driver_joint",
    gripper_close_pos: float = 0.8,
    gripper_threshold: float = 0.5,
    binarize_gripper: bool = True,
    flip_gripper: bool = False,
) -> dict[str, ActionTermOptions]:
    """Absolute joint arm controller + ``GripperDecisionController`` fed the server's closedness.

    ``flip_gripper`` treats the scalar as openness instead; ``binarize_gripper=False``
    drives the fingers continuously by the (clipped) closedness. Keep these equal to the
    policy's attributes of the same name (checked at policy construction).
    """
    return make_droid_joint_position_actions(
        arm_entity_name,
        arm_dofs_name,
        gripper_entity_name=gripper_entity_name,
        gripper_dof_name=gripper_dof_name,
        arm_key=_ARM_TERM,
        gripper_close_pos=gripper_close_pos,
        input="openness" if flip_gripper else "closedness",
        threshold=gripper_threshold,
        continuous=not binarize_gripper,
    )

observation_options classmethod

observation_options(external_camera: str = 'camera', wrist_camera: str = 'wrist_camera', arm_entity_name: str = 'robot', arm_dofs_name: list[str] | None = None, gripper_entity_name: str | None = None, gripper_dof_name: str = 'left_driver_joint', gripper_close_pos: float = 0.8) -> dict[str, ObservationGroupOptions]

Source code in eden/extensions/policies/molmoact2/policy.py

@classmethod
def observation_options(
    cls,
    external_camera: str = "camera",
    wrist_camera: str = "wrist_camera",
    arm_entity_name: str = "robot",
    arm_dofs_name: list[str] | None = None,
    gripper_entity_name: str | None = None,
    gripper_dof_name: str = "left_driver_joint",
    gripper_close_pos: float = 0.8,
) -> dict[str, ObservationGroupOptions]:
    if arm_dofs_name is None:
        arm_dofs_name = list(DROID_ARM_DOFS)

    image_group = ObservationGroupOptions(
        concatenate_terms=False,
        external=ObsTerm.configure(
            func=rgb_image,
            params={"camera_name": external_camera},
        ),
        wrist=ObsTerm.configure(
            func=rgb_image,
            params={"camera_name": wrist_camera},
        ),
    )

    return dict(
        image=image_group,
        proprio=make_droid_joint_state_group(
            arm_entity_name,
            arm_dofs_name,
            gripper_entity_name=gripper_entity_name,
            gripper_dof_name=gripper_dof_name,
            gripper_close_pos=gripper_close_pos,
        ),
    )

OctoPolicy

OctoPolicy(env: EnvBase, options: PolicyOptions)

Bases: ClientMixin, PolicyBase, ConfigurableMixin[RemotePolicyOptions]

Octo: A transformer-based robot manipulation policy.

Communicates with a remote Octo policy server (RobotArena/src/server/server_octo.py):

POST /set_task?task_description= POST /reset POST /act {"instruction": str, "image": list} -> {"raw_action": {"world_vector": list, "rotation_delta": list, "open_gripper": list}, "action": {"world_vector": list, "rot_axangle": list, "gripper": list, "terminate_episode": list}}

Action layout sent to Eden: [world_vector (3), rotation_delta (3), gripper (1)] - [:6] -> ee_controller (DifferentialIKController; the Euler delta becomes a rotation vector) - [6] -> gripper_controller (from action["gripper"], server-binarized +1=open/-1=close, forwarded raw)

Gripper convention: server +1=open / -1=close, forwarded raw; the env's GripperDecisionController (input="open_sign") turns the sign into open/close.

Source code in eden/extensions/policies/client.py

def __init__(self, env: EnvBase, options: PolicyOptions):
    super().__init__(env, options)
    if not self.supports_batch and env.num_envs != 1:
        raise ValueError(
            f"{type(self).__name__} has supports_batch=False and only supports "
            f"num_envs=1 (got {env.num_envs}). Set supports_batch=True on the "
            f"subclass if the remote server accepts batched requests."
        )
    if env.num_envs > 1 and type(self).reset is ClientMixin.reset:
        # The env resets its envs one by one; the default reset can only reset the shared
        # server state of all of them, so it would raise at the first episode end.
        raise ValueError(
            f"{type(self).__name__} serves {env.num_envs} envs but keeps ClientMixin.reset, which "
            "refuses a partial reset because the server's episode state is shared by every env. Fix: "
            "override reset(envs_idx) with a per-env scheme (e.g. one server session per env), or "
            "run with num_envs=1."
        )
    self._num_envs: int = env.num_envs
    self._post_configure()
    self._scheduler = ActionChunkScheduler(
        num_envs=env.num_envs,
        mode=self.chunk_mode,
        pred_action_horizon=self.pred_action_horizon,
        exec_horizon=self.exec_horizon,
        short_chunk=self.short_chunk,
        ensemble_steps=self.chunk_temporal_ensemble_steps,
        ensemble_dims=self.chunk_temporal_ensemble_dims,
    )
    self.last_response_info: dict[str, Any] = {}
    self._transport: TransportBase = self.make_transport()
    self._post_connect()

Attributes

action_dim class-attribute

action_dim: int | None = 7

host class-attribute instance-attribute

host: str = 'localhost'

port class-attribute instance-attribute

port: int | None = 9010

camera_name class-attribute instance-attribute

camera_name: str = 'camera'

chunk_mode class-attribute instance-attribute

chunk_mode: str = 'single'

set_task_endpoint class-attribute instance-attribute

set_task_endpoint: str | None = '/set_task'

reset_endpoint class-attribute instance-attribute

reset_endpoint: str | None = '/reset'

set_task_as_query class-attribute instance-attribute

set_task_as_query: bool = True

reset_with_payload class-attribute instance-attribute

reset_with_payload: bool = False

Methods:

decode_response

decode_response(response: dict) -> ndarray

Flatten Octo's two-dict reply into one (1, 1, 7) action row.

Uses raw_action["world_vector"] + raw_action["rotation_delta"] as the 6-D EE delta. RobotArena's Genesis default_test.py also uses the raw Euler rotation_delta for its target-pose transform, while action["gripper"] carries the processed WidowX open/close convention. The processed action["world_vector"] / action["rot_axangle"] are not a substitute: the server scales both by its action_scale, and rot_axangle is axis-angle.

Source code in eden/extensions/policies/octo/policy.py

def decode_response(self, response: dict) -> np.ndarray:
    """Flatten Octo's two-dict reply into one ``(1, 1, 7)`` action row.

    Uses raw_action["world_vector"] + raw_action["rotation_delta"] as the
    6-D EE delta. RobotArena's Genesis default_test.py also uses the raw
    Euler rotation_delta for its target-pose transform, while action["gripper"]
    carries the processed WidowX open/close convention. The processed
    action["world_vector"] / action["rot_axangle"] are not a substitute: the
    server scales both by its action_scale, and rot_axangle is axis-angle.
    """
    try:
        raw_action = response["raw_action"]
        action = response["action"]
        world_vector = raw_action["world_vector"]
        rotation_delta = raw_action["rotation_delta"]
        gripper = action["gripper"]
    except (KeyError, TypeError) as e:
        if isinstance(response, dict):
            got = {k: list(v) if isinstance(v, dict) else type(v).__name__ for k, v in response.items()}
        else:
            got = type(response).__name__
        raise KeyError(
            "Octo reply must carry raw_action{world_vector, rotation_delta} and action{gripper}, as "
            f"RobotArena's server_octo.py /act sends them; got {got}. Do not fall back to "
            "action['world_vector'] / action['rot_axangle']: the server scales both by its action_scale "
            "and rot_axangle is axis-angle, not the raw Euler rotation_delta this client converts."
        ) from e

    flat_action = np.concatenate(
        [
            np.asarray(world_vector, dtype=np.float32).reshape(-1),  # (3,) translation delta
            np.asarray(rotation_delta, dtype=np.float32).reshape(-1),  # (3,) Euler rotation delta
            np.asarray(gripper, dtype=np.float32).reshape(-1)[:1],  # (1,) binarized +1=open/-1=close
        ]
    )  # (7,)
    if flat_action.shape != (7,):
        raise ValueError(f"Expected Octo action shape (7,), got {flat_action.shape}: {flat_action}")
    return self._normalize_action_shape(flat_action)

action_options classmethod

action_options(entity_name: str = 'robot') -> dict[str, ActionTermOptions]

Source code in eden/extensions/policies/octo/policy.py

@classmethod
def action_options(cls, entity_name: str = "robot") -> dict[str, ActionTermOptions]:
    # The WidowX configuration Octo is trained on. Octo emits +1 = open / -1 = close; the
    # GripperDecisionController reads that sign directly.
    return make_widowx_ee_delta_actions(
        entity_name, "gripper_link", list(WIDOWX_DOFS), 0.5, gripper_input="open_sign"
    )

observation_options classmethod

observation_options(camera_name: str = 'camera', **kwargs) -> dict[str, ObservationGroupOptions]

Source code in eden/extensions/policies/octo/policy.py

@classmethod
def observation_options(cls, camera_name: str = "camera", **kwargs) -> dict[str, ObservationGroupOptions]:
    # Extra kwargs are accepted for polymorphic call sites and ignored.
    return dict(
        image=ObservationGroupOptions(
            concatenate_terms=False,
            rgb=ObsTerm.configure(
                func=rgb_image,
                params={"camera_name": camera_name},
            ),
        ),
    )

pack_observations

pack_observations(obs: dict[str, Tensor | dict[str, Tensor]]) -> dict

Convert Eden observations to the Octo server payload.

Source code in eden/extensions/policies/octo/policy.py

def pack_observations(
    self,
    obs: dict[str, torch.Tensor | dict[str, torch.Tensor]],
) -> dict:
    """Convert Eden observations to the Octo server payload."""
    image_obs = obs["image"]
    image: torch.Tensor = image_obs["rgb"] if isinstance(image_obs, dict) else image_obs
    return {
        "instruction": self.task_description,
        # (H, W, 3) uint8: drops the num_envs=1 batch axis and an RGBA render's alpha; a float image raises.
        "image": to_uint8_hwc(image).tolist(),
    }

unpack_actions

unpack_actions(raw_action: ndarray) -> dict[str, Tensor]

Convert a (1, 7) (or (7,)) action [world_vector (3), rotation_delta (3), gripper (1)].

Returns tensors with a leading batch axis of 1. OctoPolicy keeps ClientMixin.supports_batch at the default False because the upstream Octo HTTP server accepts a single image per /act request.

The Euler rotation_delta reaches the IK controller as the rotation vector of the same rotation (:class:EEDeltaAdapter). The gripper sign (+1 open / -1 close) is forwarded raw; the env's GripperDecisionController(input="open_sign") interprets it.

Source code in eden/extensions/policies/octo/policy.py

def unpack_actions(self, raw_action: np.ndarray) -> dict[str, torch.Tensor]:
    """Convert a ``(1, 7)`` (or ``(7,)``) action ``[world_vector (3), rotation_delta (3), gripper (1)]``.

    Returns tensors with a leading batch axis of 1. ``OctoPolicy`` keeps
    ``ClientMixin.supports_batch`` at the default ``False`` because the
    upstream Octo HTTP server accepts a single image per ``/act`` request.

    The Euler rotation_delta reaches the IK controller as the rotation vector of the
    same rotation (:class:`EEDeltaAdapter`). The gripper sign (+1 open / -1 close) is
    forwarded raw; the env's ``GripperDecisionController(input="open_sign")`` interprets it.
    """
    return self._adapter(np.asarray(raw_action, dtype=np.float32).reshape(-1), env=self.env)

See also

• Policy zoo guide
• Remote policies guide
• Extending policies