Learning

eden.learning holds learning-side code that is not a task: policy export, and algorithms that extend an RL library. A task selects an algorithm in its runner_options through an options class whose class_name names the implementation, so eden train --task <name> trains it and a saved config.json resumes through it. The training runs guide explains runners and custom algorithms; the runner options themselves are on the learning options page.

Policy export

eden export writes policy.onnx, policy.pt (TorchScript) and policy_spec.json for a runtime outside Eden. See the policy export guide.

Export a trained policy for a runtime outside Eden: an ONNX / TorchScript graph plus policy_spec.json.

eden export (:func:eden.utils.cli.tasks.export_task) builds the checkpoint's own env, loads the runner, and hands both to :func:export_policy, which writes

• policy.onnx — input obs (the actor's observation groups, concatenated, un-normalized), output actions (the raw policy output, before the action terms scale it), with a dynamic batch axis. A recurrent policy also takes h_in (and c_in for an LSTM), shaped (num_layers, batch, hidden), and returns h_out (c_out): the caller threads the state, starting from zeros;
• policy.pt — the same composition as TorchScript, batch 1 when recurrent: its hidden state lives in the module, and reset() zeroes it;
• policy_spec.json — the contract around the graph (:func:build_policy_spec): what fills obs, how actions becomes joint targets, the control rate, and where the graph came from.

The observation normalizer is part of both graphs. The graph comes from the policy model's own as_onnx() / as_jit() (rsl-rl-lib >= 4's export hooks), which is also how a learner overrides the default: RMA's :class:~eden.learning.wrappers.rsl_rl.rma.model.RmaActorModel defines both to export its deployment composition (estimator injected). An inherited hook can export something other than what the model runs, so every export is checked against the runner's inference policy on the env's own observations before anything is written (:func:verify_export). The files are written into a staging directory and moved into place only once checked, so a failed export leaves the output directory as it was.

The spec records the robot's nominal command: the joint-zero offset domain randomization draws per env is left out of the action offset, and an env whose default pose itself was randomized at startup is refused (:func:disable_contract_randomization strips those terms from a config before it is built).

Attributes

EXPORT_FORMATS module-attribute

EXPORT_FORMATS = ('onnx', 'jit')

EXPORT_FILENAMES module-attribute

EXPORT_FILENAMES = {'onnx': 'policy.onnx', 'jit': 'policy.pt'}

POLICY_SPEC_FILENAME module-attribute

POLICY_SPEC_FILENAME = 'policy_spec.json'

POLICY_SPEC_SCHEMA_VERSION module-attribute

POLICY_SPEC_SCHEMA_VERSION = 2

VERIFY_STEPS module-attribute

VERIFY_STEPS = 4

VERIFY_ATOL module-attribute

VERIFY_ATOL = 0.0001

VERIFY_RTOL module-attribute

VERIFY_RTOL = 0.0001

Classes

Functions:

export_policy

export_policy(runner: Any, env: RslRlVecEnvWrapper, out_dir: str, *, formats: tuple[str, ...] = EXPORT_FORMATS, task: str, checkpoint: str, verify: bool = True) -> dict[str, str]

Export runner's policy to out_dir and write policy_spec.json beside it.

PARAMETER	DESCRIPTION
runner	An rsl_rl runner with checkpoint loaded. TYPE: Any
env	The env the runner was built on: the observations the export is checked on, and what the spec describes. TYPE: RslRlVecEnvWrapper
out_dir	Directory to write into (created if missing). An earlier export there is replaced, and a graph of a format not exported this time is removed with it; on a failure the directory is left as it was. TYPE: str
formats	Any of :data:EXPORT_FORMATS. TYPE: tuple[str, ...] DEFAULT: EXPORT_FORMATS
task	The registered task name, recorded in the spec. TYPE: str
checkpoint	The checkpoint path, hashed into the spec. TYPE: str
verify	Check every exported graph against the runner's inference policy (:func:verify_export) before writing the spec. The ONNX check needs onnxruntime. TYPE: bool DEFAULT: True

RETURNS	DESCRIPTION
dict[str, str]	{format: path} for every file written, the spec under "spec".

Source code in eden/learning/export.py

def export_policy(
    runner: Any,
    env: RslRlVecEnvWrapper,
    out_dir: str,
    *,
    formats: tuple[str, ...] = EXPORT_FORMATS,
    task: str,
    checkpoint: str,
    verify: bool = True,
) -> dict[str, str]:
    """Export ``runner``'s policy to ``out_dir`` and write ``policy_spec.json`` beside it.

    Parameters
    ----------
    runner:
        An rsl_rl runner with ``checkpoint`` loaded.
    env:
        The env the runner was built on: the observations the export is checked on, and what the spec describes.
    out_dir:
        Directory to write into (created if missing). An earlier export there is replaced, and a graph of a format
        not exported this time is removed with it; on a failure the directory is left as it was.
    formats:
        Any of :data:`EXPORT_FORMATS`.
    task:
        The registered task name, recorded in the spec.
    checkpoint:
        The checkpoint path, hashed into the spec.
    verify:
        Check every exported graph against the runner's inference policy (:func:`verify_export`) before writing the
        spec. The ONNX check needs ``onnxruntime``.

    Returns
    -------
    dict[str, str]
        ``{format: path}`` for every file written, the spec under ``"spec"``.
    """
    unknown = sorted(set(formats) - set(EXPORT_FORMATS))
    if unknown or not formats:
        raise ValueError(
            f"Export formats {unknown or list(formats)} are not supported; choose one or more of {EXPORT_FORMATS}."
        )
    check_export_dependencies(formats, verify=verify)
    get_policy = getattr(getattr(runner, "alg", None), "get_policy", None)
    if get_policy is None:
        raise TypeError(
            f"{type(runner).__qualname__} has no `alg.get_policy()`, so there is no policy model to export. Export "
            f"reads the model through it, as rsl-rl-lib >= 4's algorithms provide it; under rsl-rl-lib 3, upgrade "
            f"(a 3.x checkpoint loads into a 4.x / 5.x runner)."
        )
    policy = get_policy()
    obs_groups = policy_obs_groups(policy)
    refuse_randomized_contract(env.unwrapped)
    os.makedirs(out_dir, exist_ok=True)
    # Staged, then moved in only once checked and described: a graph that failed its check, or a spec cut short, must
    # not land where a deployment could pick it up, nor beside the files of an earlier export.
    staging = tempfile.mkdtemp(prefix=".eden-export-", dir=out_dir)
    try:
        staged = {}
        for fmt in EXPORT_FORMATS:
            if fmt in formats:
                staged[fmt] = os.path.join(staging, EXPORT_FILENAMES[fmt])
                (_export_onnx if fmt == "onnx" else _export_jit)(policy, staged[fmt])
        if verify:
            verify_export(staged, runner.get_inference_policy(device=env.device), env, obs_groups)
        spec = build_policy_spec(env.unwrapped, policy, staged, task=task, checkpoint=checkpoint, runner_env=env)
        staged["spec"] = os.path.join(staging, POLICY_SPEC_FILENAME)
        with open(staged["spec"], "w") as f:
            json.dump(spec, f, indent=2)
            f.write("\n")
        files = {key: os.path.join(out_dir, os.path.basename(path)) for key, path in staged.items()}
        for key, path in staged.items():  # the spec last, so it never names a graph not yet in place
            os.replace(path, files[key])
    finally:
        shutil.rmtree(staging, ignore_errors=True)
    for fmt in EXPORT_FORMATS:
        stale = os.path.join(out_dir, EXPORT_FILENAMES[fmt])
        if fmt not in files and os.path.exists(stale):
            os.remove(stale)  # an earlier export's graph, which the new spec does not describe
    return files

check_export_dependencies

check_export_dependencies(formats: tuple[str, ...], *, verify: bool) -> None

Raise ImportError if a package formats needs is missing: onnx for the ONNX graph's spec, and onnxruntime to check it.

eden export calls this before it builds anything, so a missing extra costs no env build.

Source code in eden/learning/export.py

def check_export_dependencies(formats: tuple[str, ...], *, verify: bool) -> None:
    """Raise ``ImportError`` if a package ``formats`` needs is missing: ``onnx`` for the ONNX graph's spec, and
    ``onnxruntime`` to check it.

    ``eden export`` calls this before it builds anything, so a missing extra costs no env build.
    """
    if "onnx" not in formats:
        return
    missing = []
    for package in ("onnx", "onnxruntime") if verify else ("onnx",):
        try:
            importlib.import_module(package)
        except (ImportError, OSError) as exc:  # a broken onnxruntime-gpu wheel fails on its CUDA libraries
            missing.append(f"{package} ({' '.join(str(exc).split())})")
    if missing:
        skip = " onnxruntime only runs the parity check, which --no_verify (verify=False) skips." if verify else ""
        raise ImportError(
            f"Exporting the ONNX graph needs {' and '.join(missing)}. Install the extra with {install_hint('onnx')} "
            f"(for the GPU wheel: {install_hint('onnx-gpu')}). Or export TorchScript only with --format jit.{skip}"
        )

disable_contract_randomization

disable_contract_randomization(config: EdenRLConfig) -> list[str]

Disable config's event terms that would randomize what the spec records as nominal; return their names.

Call it before the env is built: eden export does. :func:export_policy refuses an env built with them.

Source code in eden/learning/export.py

def disable_contract_randomization(config: EdenRLConfig) -> list[str]:
    """Disable ``config``'s event terms that would randomize what the spec records as nominal; return their names.

    Call it before the env is built: ``eden export`` does. :func:`export_policy` refuses an env built with them.
    """
    from eden.managers.event_manager import EVENT_TERM_REGISTRY

    if config.event_options is None:
        return []
    disabled = []
    for name in config.event_options.term_keys():
        term_cls = EVENT_TERM_REGISTRY.get(getattr(config.event_options, name).name)
        if isinstance(term_cls, type) and issubclass(term_cls, _contract_randomization_terms()):
            config.event_options.disable_term(name)
            disabled.append(name)
    return disabled

refuse_randomized_contract

refuse_randomized_contract(env: RLEnvBase) -> None

Raise ValueError if env was built with an event term that randomizes what the spec records as nominal.

Source code in eden/learning/export.py

def refuse_randomized_contract(env: RLEnvBase) -> None:
    """Raise ``ValueError`` if ``env`` was built with an event term that randomizes what the spec records as nominal."""
    event_manager = getattr(env, "event_manager", None)
    if event_manager is None:
        return
    randomizing = {
        name: type(term).__name__
        for _, name, term in event_manager.iter_terms()
        if isinstance(term, _contract_randomization_terms())
    }
    if randomizing:
        listed = ", ".join(f"'{name}' ({cls})" for name, cls in randomizing.items())
        raise ValueError(
            f"The env was built with event term(s) {listed}, which bias each env's default joint pose at startup, so "
            f"policy_spec.json would record env 0's random draw as the robot's default pose and action offset. Build "
            f"the env from a config passed through eden.learning.export.disable_contract_randomization(config) "
            f"first, as eden export does; reading another env's row would only record another draw."
        )

policy_obs_groups

policy_obs_groups(policy: Any) -> list[str]

The observation groups policy concatenates into its input, in order — the layout of the graph's obs.

Source code in eden/learning/export.py

def policy_obs_groups(policy: Any) -> list[str]:
    """The observation groups ``policy`` concatenates into its input, in order — the layout of the graph's ``obs``."""
    groups = getattr(policy, "obs_groups", None)
    if not isinstance(groups, (list, tuple)) or not groups:
        raise TypeError(
            f"{type(policy).__qualname__} has no `obs_groups` list, so the observation layout of an export cannot be "
            f"known. Export reads the groups the model concatenates from that attribute, as rsl-rl-lib >= 4's "
            f"MLPModel sets it; a custom model must set it too."
        )
    return list(groups)

verify_export

verify_export(files: dict[str, str], inference_policy: Any, env: RslRlVecEnvWrapper, obs_groups: list[str]) -> None

Check every exported graph reproduces inference_policy over :data:VERIFY_STEPS real control steps.

The env is stepped by inference_policy from its current state, under torch.inference_mode as rsl_rl's rollout steps it; each graph then replays the same observations from a zero recurrent state. inference_policy is what eden inference runs, so a learner whose export hook exports anything else — an inherited as_onnx() that skips the model's own forward — fails here, before a spec is written next to it.

RAISES	DESCRIPTION
RuntimeError	Naming the format, the step and the largest deviation, when a graph disagrees.

Source code in eden/learning/export.py

def verify_export(files: dict[str, str], inference_policy: Any, env: RslRlVecEnvWrapper, obs_groups: list[str]) -> None:
    """Check every exported graph reproduces ``inference_policy`` over :data:`VERIFY_STEPS` real control steps.

    The env is stepped by ``inference_policy`` from its current state, under ``torch.inference_mode`` as rsl_rl's
    rollout steps it; each graph then replays the same observations from a zero recurrent state.
    ``inference_policy`` is what ``eden inference`` runs, so a learner whose export hook exports anything else — an
    inherited ``as_onnx()`` that skips the model's own forward — fails here, before a spec is written next to it.

    Raises
    ------
    RuntimeError
        Naming the format, the step and the largest deviation, when a graph disagrees.
    """
    reset = getattr(inference_policy, "reset", None)
    if reset is not None:
        reset()
    observations, expected = [], []
    obs = env.get_observations()
    # inference_mode, as rsl_rl's own rollout steps the env: after a `learn()` its buffers are inference tensors,
    # which only an inference-mode step may update in place.
    with torch.inference_mode():
        for _ in range(VERIFY_STEPS):
            actions = inference_policy(obs)
            observations.append(torch.cat([obs[group] for group in obs_groups], dim=-1).cpu())
            expected.append(actions.cpu())
            obs, *_ = env.step(actions)
    for fmt, path in files.items():
        produced = _run_onnx(path, observations) if fmt == "onnx" else _run_jit(path, observations)
        name = os.path.basename(path)
        for step, (got, want) in enumerate(zip(produced, expected, strict=True)):
            if got.shape != want.shape:
                mismatch = f"actions of shape {tuple(got.shape)} where the policy returned {tuple(want.shape)}"
            elif not torch.allclose(got, want, atol=VERIFY_ATOL, rtol=VERIFY_RTOL):
                mismatch = (
                    f"actions up to {(got - want).abs().max().item():.3e} away from the policy's "
                    f"(tolerance atol={VERIFY_ATOL}, rtol={VERIFY_RTOL})"
                )
            else:
                continue
            raise RuntimeError(
                f"The exported {fmt} graph ({name}) does not reproduce the policy: at step {step} it returned {mismatch}. "
                f"The model's export hook exports something other than what it runs; fix the hook rather than the "
                f"tolerance."
            )

build_policy_spec

build_policy_spec(env: RLEnvBase, policy: Any, files: dict[str, str], *, task: str, checkpoint: str, runner_env: RslRlVecEnvWrapper | None = None) -> dict[str, Any]

The policy_spec.json contract for policy as built on env.

Sections: model (the files and their tensors), observations (what fills obs, slice by slice), actions (how actions becomes a command), entities (each commanded entity's DOFs in column order, with default positions and PD gains), rate and provenance. runner_env is the rsl_rl wrapper, for its action clip. env must pass :func:refuse_randomized_contract, as :func:export_policy checks first.

Source code in eden/learning/export.py

def build_policy_spec(
    env: RLEnvBase,
    policy: Any,
    files: dict[str, str],
    *,
    task: str,
    checkpoint: str,
    runner_env: RslRlVecEnvWrapper | None = None,
) -> dict[str, Any]:
    """The ``policy_spec.json`` contract for ``policy`` as built on ``env``.

    Sections: ``model`` (the files and their tensors), ``observations`` (what fills ``obs``, slice by slice),
    ``actions`` (how ``actions`` becomes a command), ``entities`` (each commanded entity's DOFs in column order,
    with default positions and PD gains), ``rate`` and ``provenance``. ``runner_env`` is the rsl_rl wrapper, for
    its action clip. ``env`` must pass :func:`refuse_randomized_contract`, as :func:`export_policy` checks first.
    """
    obs_groups = policy_obs_groups(policy)
    clip_actions = getattr(runner_env, "clip_actions", None)
    action_terms = _action_spec(env)
    return {
        "schema_version": POLICY_SPEC_SCHEMA_VERSION,
        "model": _model_spec(files),
        "observations": _observation_spec(env, obs_groups),
        "actions": {
            "dim": env.action_manager.total_action_dim,
            "clip": None if clip_actions is None else [-float(clip_actions), float(clip_actions)],
            "terms": action_terms,
        },
        "entities": _entity_spec(env, {term["entity_name"] for term in action_terms if term.get("dofs_name")}),
        "rate": {
            "control_dt": float(env.dt),
            "decimation": int(env.env_options.decimation),
            "sim_dt": float(env.env_options.sim_dt),
        },
        "provenance": _provenance(env, task, checkpoint),
    }

rsl-rl-lib algorithms

eden.learning.wrappers.rsl_rl holds one subpackage per algorithm. Each has an options module (what a task configures) next to the implementation.

Models, algorithms and runners for rsl-rl-lib 5, selected through their options' class_name.

• :mod:.simba — a SimBaV2 actor / critic behind one encoder per observation group.
• :mod:.muon_ppo — PPO whose optimizer is Muon for the 2-D weights and AdamW for the rest.
• :mod:.him — HIMLoco (OpenHomie's HIMPPO): a self-supervised history estimator feeding the actor, plus rollout-time symmetry and mirror losses.
• :mod:.amp_ppo — PPO plus an adversarial-motion-prior style reward (HUSKY).
• :mod:.rma — the RMA teacher (actor with estimator + scan / privileged / history encoders, PPO with privileged-latent regularization and DAgger) and its depth-student distillation runner (extreme parkour).

Each subpackage ships an options class subclassing the rsl_rl option it replaces (RslRlModelOptions / RslRlPpoAlgorithmOptions / RslRlBaseRunnerOptions) whose class_name defaults to its implementation, so a task selects it with actor=<Options>(...), algorithm=<Options>(...) or runner_options=<Options>(...).

Attributes

MIN_RSL_RL_VERSION module-attribute

MIN_RSL_RL_VERSION = (5, 5, 1)

SimBaV2 actor-critic

Actor / critic options for :class:~eden.learning.wrappers.rsl_rl.simba.model.SimBaModel.

Declarative only — importing this module does not pull in rsl_rl's model API, so a task config that uses it still imports under any rsl_rl Eden supports. The model its class_name names needs rsl-rl-lib >= 5 and says so when imported without it. The defaults are Weave's G1-Inspire-HOI networks.

Classes

RslRlSimBaCriticOptions

RslRlSimBaCriticOptions(**data)

Bases: _SimBaFields

Critic: a deterministic SimBaV2 value network (Weave: 2 blocks).

PARAMETER	DESCRIPTION
num_blocks	Number of SimBaV2 residual blocks. TYPE: int

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

num_blocks class-attribute instance-attribute

num_blocks: int = 2

RslRlSimBaActorOptions

RslRlSimBaActorOptions(**data)

Bases: _SimBaFields, RslRlModelWithDistributionOptions

Actor: a SimBaV2 policy with a state-independent Gaussian std (Weave: 1 block, scalar std 1.0).

PARAMETER	DESCRIPTION
num_blocks	Number of SimBaV2 residual blocks. TYPE: int
distribution_cfg	Output distribution; Weave's is a learned per-action scalar std initialised at 1.0. TYPE: RslRlGaussianDistributionOptions

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

num_blocks class-attribute instance-attribute

num_blocks: int = 1

distribution_cfg class-attribute instance-attribute

distribution_cfg: RslRlGaussianDistributionOptions = RslRlGaussianDistributionOptions(init_std=1.0, std_type='scalar')

SimBaModel

SimBaModel(obs: TensorDict, obs_groups: dict[str, list[str]], obs_set: str, output_dim: int, hidden_dim: int = 512, num_blocks: int = 2, expansion: int = 4, latent_dim: int = 128, encoder_hidden_dims: dict[str, list[int]] | None = None, default_encoder_hidden_dims: list[int] | tuple[int, ...] = (256,), c_shift: float = 3.0, compile: bool = False, hidden_dims: list[int] | tuple[int, ...] = (), activation: str = 'swish', obs_normalization: bool = True, distribution_cfg: dict | None = None)

Bases: MLPModel

Per-group MLP encoders feeding a SimBaV2 backbone, as an rsl_rl >= 5 model.

PARAMETER	DESCRIPTION
obs	A batch of observations; only the group widths are read. TYPE: TensorDict
obs_groups	Observation set -> groups, as resolved by rsl_rl. TYPE: dict[str, list[str]]
obs_set	The set this model reads ("actor" or "critic"). TYPE: str
output_dim	Output width (action dimension, or 1 for the critic). TYPE: int
hidden_dim	Width of the SimBaV2 hypersphere. TYPE: int DEFAULT: 512
num_blocks	Number of SimBaV2 residual blocks. TYPE: int DEFAULT: 2
expansion	Inner-layer width multiplier of each block. TYPE: int DEFAULT: 4
latent_dim	Output width of every group encoder. TYPE: int DEFAULT: 128
encoder_hidden_dims	Per-group hidden widths of the group encoders. A key must name an observation group the env provides; a group of this set without a key uses default_encoder_hidden_dims. TYPE: dict[str, list[int]] DEFAULT: None
default_encoder_hidden_dims	Hidden widths for groups not in encoder_hidden_dims. TYPE: list[int] DEFAULT: (256,)
c_shift	Input shift of the SimBaV2 embedder. TYPE: float DEFAULT: 3.0
compile	torch.compile the trunk's forward (inductor, fullgraph, static shapes), as Weave does in training. The compiled function is kept outside the module state, so state_dict keys and export are unchanged. TYPE: bool DEFAULT: False
hidden_dims	Inherited from rsl_rl's MLPModel; SimBa has no plain MLP trunk, so it must be empty. TYPE: list[int] | tuple[int, ...] DEFAULT: ()
activation	Activation between the layers of the group encoders (Weave: SiLU, which rsl_rl names "swish"). TYPE: str DEFAULT: 'swish'
obs_normalization	Normalise the concatenated observation with EmpiricalNormalization. TYPE: bool DEFAULT: True
distribution_cfg	Output distribution (actor), or None for a deterministic output (critic). TYPE: dict | None DEFAULT: None

Source code in eden/learning/wrappers/rsl_rl/simba/model.py

def __init__(
    self,
    obs: TensorDict,
    obs_groups: dict[str, list[str]],
    obs_set: str,
    output_dim: int,
    hidden_dim: int = 512,
    num_blocks: int = 2,
    expansion: int = 4,
    latent_dim: int = 128,
    encoder_hidden_dims: dict[str, list[int]] | None = None,
    default_encoder_hidden_dims: list[int] | tuple[int, ...] = (256,),
    c_shift: float = 3.0,
    compile: bool = False,
    hidden_dims: list[int] | tuple[int, ...] = (),
    activation: str = "swish",
    obs_normalization: bool = True,
    distribution_cfg: dict | None = None,
) -> None:
    if len(hidden_dims) != 0:
        raise ValueError(
            f"SimBaModel got hidden_dims={list(hidden_dims)}, but SimBa has no plain MLP trunk to size. Set the "
            "backbone with hidden_dim / num_blocks / expansion and the per-group encoders with "
            "encoder_hidden_dims, and leave hidden_dims empty."
        )
    encoder_hidden_dims = dict(encoder_hidden_dims or {})
    unknown = sorted(set(encoder_hidden_dims) - set(obs.keys()))
    if unknown:
        raise ValueError(
            f"SimBaModel encoder_hidden_dims names observation groups {unknown} that the env does not provide; "
            f"available groups: {sorted(obs.keys())}."
        )
    # rsl_rl's MLPModel builds the normaliser, the distribution and a plain MLP trunk; the trunk is replaced
    # below, so build it as small as possible (a 1-wide hidden layer; rsl_rl's MLP needs at least one) rather
    # than at the MLPModel default size.
    super().__init__(
        obs,
        obs_groups,
        obs_set,
        output_dim,
        hidden_dims=[1],
        activation=activation,
        obs_normalization=obs_normalization,
        distribution_cfg=distribution_cfg,
    )
    trunk_output_dim = self.distribution.input_dim if self.distribution is not None else output_dim
    group_dims = [obs[group].shape[-1] for group in self.obs_groups]
    specs = [list(encoder_hidden_dims.get(group, default_encoder_hidden_dims)) for group in self.obs_groups]
    encoder = GroupEncoder(group_dims, specs, latent_dim, resolve_nn_activation(activation))
    self.mlp = nn.Sequential(
        encoder, SimBa(encoder.out_features, trunk_output_dim, hidden_dim, num_blocks, expansion, c_shift)
    )
    if self.distribution is not None:
        self.distribution.init_mlp_weights(self.mlp)
    self._compile_trunk = compile
    self._compiled_trunk = None

Attributes

mlp instance-attribute

mlp = nn.Sequential(encoder, SimBa(encoder.out_features, trunk_output_dim, hidden_dim, num_blocks, expansion, c_shift))

Methods:

forward

forward(obs: TensorDict, masks: Tensor | None = None, hidden_state: Any = None, stochastic_output: bool = False) -> Tensor

Source code in eden/learning/wrappers/rsl_rl/simba/model.py

def forward(
    self,
    obs: TensorDict,
    masks: torch.Tensor | None = None,
    hidden_state: Any = None,
    stochastic_output: bool = False,
) -> torch.Tensor:
    if not self._compile_trunk:
        return super().forward(obs, masks, hidden_state, stochastic_output)
    if self._compiled_trunk is None:
        self._compiled_trunk = torch.compile(
            self.mlp.forward, backend="inductor", mode="default", fullgraph=True, dynamic=False
        )
    obs = unpad_trajectories(obs, masks) if masks is not None and not self.is_recurrent else obs
    latent = self.get_latent(obs, masks, hidden_state)
    trunk_output = self._compiled_trunk(latent)
    if self.distribution is not None:
        if stochastic_output:
            self.distribution.update(trunk_output)
            return self.distribution.sample()
        return self.distribution.deterministic_output(trunk_output)
    return trunk_output

PPO with Muon

Algorithm options for :class:~eden.learning.wrappers.rsl_rl.muon_ppo.ppo.MuonPPO.

Declarative only — importing this module pulls in neither rsl_rl's algorithm API nor torch.optim.Muon, so a task config that uses it still imports under any rsl_rl Eden supports. The algorithm its class_name names needs rsl-rl-lib >= 5 and torch >= 2.10, and says so when imported without them.

Classes

RslRlMuonPpoAlgorithmOptions

RslRlMuonPpoAlgorithmOptions(**data)

Bases: RslRlPpoAlgorithmOptions

PPO with Muon (2-D weights) + AdamW (the rest), Weave's hyperparameters as defaults.

PARAMETER	DESCRIPTION
class_name	The algorithm class rsl_rl resolves; leave at the default. TYPE: str
learning_rate	Initial learning rate of both optimizers (then adapted to desired_kl). TYPE: float
entropy_coef	Entropy bonus coefficient. TYPE: float
weight_decay	AdamW weight decay. TYPE: float
muon_weight_decay	Muon weight decay (Weave leaves torch.optim.Muon's default, 0.1). TYPE: float

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

class_name class-attribute instance-attribute

class_name: str = _MUON_PPO_CLASS_NAME

learning_rate class-attribute instance-attribute

learning_rate: float = 0.0001

entropy_coef class-attribute instance-attribute

entropy_coef: float = 0.0005

weight_decay class-attribute instance-attribute

weight_decay: float = 0.01

muon_weight_decay class-attribute instance-attribute

muon_weight_decay: float = 0.1

MuonPPO

MuonPPO(*args: Any, weight_decay: float = 0.01, muon_weight_decay: float = 0.1, **kwargs: Any)

Bases: PPO

rsl_rl PPO whose Adam is replaced by :class:MuonAdamW over the actor and critic.

PARAMETER	DESCRIPTION
*args	Forwarded to rsl_rl.algorithms.PPO. TYPE: Any DEFAULT: ()
weight_decay	AdamW weight decay (Weave: 0.01). TYPE: float DEFAULT: 0.01
muon_weight_decay	Muon weight decay (Weave: torch's default, 0.1). TYPE: float DEFAULT: 0.1
**kwargs	Forwarded to rsl_rl.algorithms.PPO. TYPE: Any DEFAULT: {}

Source code in eden/learning/wrappers/rsl_rl/muon_ppo/ppo.py

def __init__(self, *args: Any, weight_decay: float = 0.01, muon_weight_decay: float = 0.1, **kwargs: Any) -> None:
    super().__init__(*args, **kwargs)
    self.optimizer = MuonAdamW(
        [self.actor, self.critic],
        lr=self.learning_rate,
        weight_decay=weight_decay,
        muon_weight_decay=muon_weight_decay,
    )

Attributes

optimizer instance-attribute

optimizer = MuonAdamW([self.actor, self.critic], lr=self.learning_rate, weight_decay=weight_decay, muon_weight_decay=muon_weight_decay)

Methods:

update

update() -> dict[str, float]

Run the stock PPO update, and log the value function's explained variance (Weave's extra metric).

RETURNS	DESCRIPTION
dict[str, float]	rsl_rl's loss dict plus explained_variance.

Source code in eden/learning/wrappers/rsl_rl/muon_ppo/ppo.py

def update(self) -> dict[str, float]:
    """Run the stock PPO update, and log the value function's explained variance (Weave's extra metric).

    Returns
    -------
    dict[str, float]
        rsl_rl's loss dict plus ``explained_variance``.
    """
    with torch.no_grad():
        values = self.storage.values.flatten()
        returns = self.storage.returns.flatten()
        explained_variance = (1.0 - (returns - values).var() / (returns.var() + 1e-8)).item()
    loss_dict = super().update()
    loss_dict["explained_variance"] = explained_variance
    return loss_dict

PPO with an adversarial motion prior

Algorithm options for :class:~eden.learning.wrappers.rsl_rl.amp_ppo.ppo.AmpPPO.

Declarative only — importing this module pulls in neither rsl_rl's algorithm API nor the discriminator, so a task config that uses it still imports under any rsl_rl Eden supports. The algorithm its class_name names needs rsl-rl-lib >= 5.5.1 and says so when imported without it.

Classes

RslRlAmpPpoAlgorithmOptions

RslRlAmpPpoAlgorithmOptions(**data)

Bases: RslRlPpoAlgorithmOptions

PPO plus an adversarial motion prior (AMP) style reward; defaults are the HUSKY release's.

The env must publish the observation group amp_obs_group as [joint positions x amp_num_frames (oldest first), style gate]: the style reward is added only where the gate is set (> 0.5), and every window is banked as a policy sample for the discriminator.

PARAMETER	DESCRIPTION
class_name	The algorithm class rsl_rl resolves; leave at the default. TYPE: str
amp_motion_file	Asset path (resolved by eden.utils.assets.get_asset_path) of the expert .npz bundle: motion_<i> arrays of shape (T, frame_dim) plus fps and num_motions. frame_dim must equal the joint-position width of the AMP observation group. TYPE: str
amp_obs_group	Observation group carrying the window and the gate. TYPE: str
amp_num_frames	Frames per window; must equal the group's joint-position history_length. TYPE: int
amp_reward_coef	Scale inside the LSGAN style reward coef * clamp(1 - 0.25 (d - 1)^2, 0). TYPE: float
amp_reward_scale	Extra multiplier on the style reward before it is added to the env reward. TYPE: float
amp_replay_size	Capacity of the policy-window replay buffer. TYPE: int
amp_discr_hidden_dims	Hidden widths of the spectral-norm discriminator MLP. TYPE: list[int]
amp_grad_pen_scale	Gradient-penalty weight at the expert samples. TYPE: float
amp_trunk_weight_decay	Weight decay of the discriminator trunk. TYPE: float
amp_head_weight_decay	Weight decay of the discriminator head. TYPE: float
amp_warm_start	Allow loading a checkpoint without AMP state (one trained as plain PPO): its actor-critic loads and the discriminator starts fresh. Off, such a checkpoint is refused rather than resumed under a changed reward. TYPE: bool

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

class_name class-attribute instance-attribute

class_name: str = _AMP_PPO_CLASS_NAME

amp_motion_file instance-attribute

amp_motion_file: str

amp_obs_group class-attribute instance-attribute

amp_obs_group: str = 'amp'

amp_num_frames class-attribute instance-attribute

amp_num_frames: int = 5

amp_reward_coef class-attribute instance-attribute

amp_reward_coef: float = 5.0

amp_reward_scale class-attribute instance-attribute

amp_reward_scale: float = 0.02

amp_replay_size class-attribute instance-attribute

amp_replay_size: int = 100000

amp_discr_hidden_dims class-attribute instance-attribute

amp_discr_hidden_dims: list[int] = [256, 256]

amp_grad_pen_scale class-attribute instance-attribute

amp_grad_pen_scale: float = 5.0

amp_trunk_weight_decay class-attribute instance-attribute

amp_trunk_weight_decay: float = 0.001

amp_head_weight_decay class-attribute instance-attribute

amp_head_weight_decay: float = 0.1

amp_warm_start class-attribute instance-attribute

amp_warm_start: bool = False

Methods:

for_inference

for_inference() -> RslRlPpoAlgorithmOptions

Stock PPO with the same PPO fields — what replays an AMP checkpoint (the AMP state only shapes training).

Source code in eden/learning/wrappers/rsl_rl/amp_ppo/options.py

def for_inference(self) -> RslRlPpoAlgorithmOptions:
    """Stock PPO with the same PPO fields — what replays an AMP checkpoint (the AMP state only shapes training)."""
    fields = {name: getattr(self, name) for name in RslRlPpoAlgorithmOptions.model_fields if name != "class_name"}
    return RslRlPpoAlgorithmOptions(**fields)

preflight

preflight() -> None

Load the expert clips (eden train calls this before building the env), so a bad file fails first.

Source code in eden/learning/wrappers/rsl_rl/amp_ppo/options.py

def preflight(self) -> None:
    """Load the expert clips (``eden train`` calls this before building the env), so a bad file fails first."""
    from eden.learning.wrappers.rsl_rl.amp_ppo.amp import AmpMotionSampler

    AmpMotionSampler(self.amp_motion_file, self.amp_num_frames, "cpu")

AmpPPO

AmpPPO(actor, critic, storage, *, amp_motion_file: str, amp_obs_group: str = 'amp', amp_num_frames: int = 5, amp_reward_coef: float = 5.0, amp_reward_scale: float = 0.02, amp_replay_size: int = 100000, amp_discr_hidden_dims: tuple[int, ...] = (256, 256), amp_grad_pen_scale: float = 5.0, amp_trunk_weight_decay: float = 0.001, amp_head_weight_decay: float = 0.1, amp_warm_start: bool = False, **ppo_kwargs)

Bases: PPO

PPO + an adversarial motion prior on gated windows of joint positions.

The amp_* keyword arguments are documented on :class:~eden.learning.wrappers.rsl_rl.amp_ppo.options.RslRlAmpPpoAlgorithmOptions; the rest go to rsl_rl's PPO.

Source code in eden/learning/wrappers/rsl_rl/amp_ppo/ppo.py

def __init__(
    self,
    actor,
    critic,
    storage,
    *,
    amp_motion_file: str,
    amp_obs_group: str = "amp",
    amp_num_frames: int = 5,
    amp_reward_coef: float = 5.0,
    amp_reward_scale: float = 0.02,
    amp_replay_size: int = 100_000,
    amp_discr_hidden_dims: tuple[int, ...] = (256, 256),
    amp_grad_pen_scale: float = 5.0,
    amp_trunk_weight_decay: float = 1e-3,
    amp_head_weight_decay: float = 1e-1,
    amp_warm_start: bool = False,
    **ppo_kwargs,
) -> None:
    super().__init__(actor, critic, storage, **ppo_kwargs)
    self.amp_obs_group = amp_obs_group
    self.amp_warm_start = amp_warm_start
    self.amp_num_frames = amp_num_frames
    self.amp_reward_coef = amp_reward_coef
    self.amp_reward_scale = amp_reward_scale
    self.amp_grad_pen_scale = amp_grad_pen_scale

    self.amp_expert = AmpMotionSampler(amp_motion_file, amp_num_frames, self.device)
    window_dim = self.amp_expert.frame_dim * amp_num_frames
    self.discriminator = AmpDiscriminator(window_dim, tuple(amp_discr_hidden_dims)).to(self.device)
    self.amp_normalizer = FrameNormalizer(self.amp_expert.frame_dim).to(self.device)
    self.amp_replay = AmpReplayBuffer(window_dim, amp_replay_size, self.device)
    self.amp_optimizer = torch.optim.Adam(
        [
            {"params": self.discriminator.trunk.parameters(), "weight_decay": amp_trunk_weight_decay},
            {"params": self.discriminator.head.parameters(), "weight_decay": amp_head_weight_decay},
        ],
        lr=self.learning_rate,
    )
    self._window_dim = window_dim
    # Accumulated as a tensor to avoid a GPU sync inside the rollout loop.
    self._style_reward_sum = torch.zeros((), device=self.device)
    self._style_reward_steps = 0

Attributes

amp_obs_group instance-attribute

amp_obs_group = amp_obs_group

amp_warm_start instance-attribute

amp_warm_start = amp_warm_start

amp_num_frames instance-attribute

amp_num_frames = amp_num_frames

amp_reward_coef instance-attribute

amp_reward_coef = amp_reward_coef

amp_reward_scale instance-attribute

amp_reward_scale = amp_reward_scale

amp_grad_pen_scale instance-attribute

amp_grad_pen_scale = amp_grad_pen_scale

amp_expert instance-attribute

amp_expert = AmpMotionSampler(amp_motion_file, amp_num_frames, self.device)

discriminator instance-attribute

discriminator = AmpDiscriminator(window_dim, tuple(amp_discr_hidden_dims)).to(self.device)

amp_normalizer instance-attribute

amp_normalizer = FrameNormalizer(self.amp_expert.frame_dim).to(self.device)

amp_replay instance-attribute

amp_replay = AmpReplayBuffer(window_dim, amp_replay_size, self.device)

amp_optimizer instance-attribute

amp_optimizer = torch.optim.Adam([{'params': self.discriminator.trunk.parameters(), 'weight_decay': amp_trunk_weight_decay}, {'params': self.discriminator.head.parameters(), 'weight_decay': amp_head_weight_decay}], lr=self.learning_rate)

Methods:

construct_algorithm staticmethod

construct_algorithm(obs: TensorDict, env, cfg: dict, device: str) -> AmpPPO

rsl_rl's PPO construction, then the AMP group checked against the env's first observation.

Source code in eden/learning/wrappers/rsl_rl/amp_ppo/ppo.py

@staticmethod
def construct_algorithm(obs: TensorDict, env, cfg: dict, device: str) -> AmpPPO:
    """rsl_rl's PPO construction, then the AMP group checked against the env's first observation."""
    alg = PPO.construct_algorithm(obs, env, cfg, device)
    if alg.amp_obs_group not in obs.keys():
        raise ValueError(
            f"AmpPPO reads the observation group amp_obs_group={alg.amp_obs_group!r}, which the env does not "
            f"provide (groups: {sorted(obs.keys())}). Add it to the task's observation_options, or point "
            "amp_obs_group at the group that carries the AMP window and gate."
        )
    alg._split_amp_obs(obs)  # raises, naming the expected layout, on a mis-sized group
    return alg

process_env_step

process_env_step(obs: TensorDict, rewards: Tensor, dones: Tensor, extras: dict[str, Tensor]) -> None

Add the gated style reward and bank the policy window, then defer to PPO.

Source code in eden/learning/wrappers/rsl_rl/amp_ppo/ppo.py

def process_env_step(
    self, obs: TensorDict, rewards: torch.Tensor, dones: torch.Tensor, extras: dict[str, torch.Tensor]
) -> None:
    """Add the gated style reward and bank the policy window, then defer to PPO."""
    raw_windows, gate = self._split_amp_obs(obs)
    frames = raw_windows.view(-1, self.amp_num_frames, self.amp_expert.frame_dim)
    style = self.discriminator.style_reward(self.amp_normalizer.normalize(frames).flatten(1), self.amp_reward_coef)
    style = style * gate.float() * self.amp_reward_scale
    rewards = rewards + style

    self.amp_replay.insert(raw_windows)
    self._style_reward_sum = self._style_reward_sum + style.mean()
    self._style_reward_steps += 1
    super().process_env_step(obs, rewards, dones, extras)

update

update() -> dict[str, float]

Run the PPO update, then the discriminator's LSGAN updates.

Source code in eden/learning/wrappers/rsl_rl/amp_ppo/ppo.py

def update(self) -> dict[str, float]:
    """Run the PPO update, then the discriminator's LSGAN updates."""
    batch_size = self.storage.num_envs * self.storage.num_transitions_per_env // self.num_mini_batches
    loss_dict = super().update()

    # The source trained the discriminator inside the PPO optimizer, so its
    # learning rate followed the KL-adaptive schedule; mirror that here.
    for group in self.amp_optimizer.param_groups:
        group["lr"] = self.learning_rate

    num_updates = self.num_learning_epochs * self.num_mini_batches
    mean_amp_loss = mean_grad_pen = mean_expert_d = mean_policy_d = 0.0
    for _ in range(num_updates):
        policy_windows = self.amp_replay.sample(batch_size).view(-1, self.amp_num_frames, self.amp_expert.frame_dim)
        expert_windows = self.amp_expert.sample(batch_size)
        self.amp_normalizer.update(self._gather(policy_windows[:, -1]))
        self.amp_normalizer.update(self._gather(expert_windows[:, -1]))
        policy_in = self.amp_normalizer.normalize(policy_windows).flatten(1)
        expert_in = self.amp_normalizer.normalize(expert_windows).flatten(1)

        policy_d = self.discriminator(policy_in)
        expert_d = self.discriminator(expert_in)
        expert_loss = torch.nn.functional.mse_loss(expert_d, torch.ones_like(expert_d))
        policy_loss = torch.nn.functional.mse_loss(policy_d, -torch.ones_like(policy_d))
        amp_loss = 0.5 * (expert_loss + policy_loss)
        grad_pen = self.discriminator.gradient_penalty(expert_in, self.amp_grad_pen_scale)

        self.amp_optimizer.zero_grad()
        (amp_loss + grad_pen).backward()
        if self.is_multi_gpu:
            reduce_gradients_in_buckets(
                self.discriminator.parameters(), self.gpu_world_size, self.grad_reduce_bucket_mb
            )
        self.amp_optimizer.step()

        mean_amp_loss += amp_loss.item()
        mean_grad_pen += grad_pen.item()
        mean_expert_d += expert_d.mean().item()
        mean_policy_d += policy_d.mean().item()

    loss_dict["amp"] = mean_amp_loss / num_updates
    loss_dict["amp_grad_pen"] = mean_grad_pen / num_updates
    loss_dict["amp_expert_d"] = mean_expert_d / num_updates
    loss_dict["amp_policy_d"] = mean_policy_d / num_updates
    loss_dict["amp_style_reward"] = float(self._style_reward_sum) / max(self._style_reward_steps, 1)
    self._style_reward_sum = torch.zeros((), device=self.device)
    self._style_reward_steps = 0
    return loss_dict

broadcast_parameters

broadcast_parameters() -> None

PPO's broadcast, plus the discriminator (spectral-norm vectors included) and the frame normalizer.

Source code in eden/learning/wrappers/rsl_rl/amp_ppo/ppo.py

def broadcast_parameters(self) -> None:
    """PPO's broadcast, plus the discriminator (spectral-norm vectors included) and the frame normalizer."""
    super().broadcast_parameters()
    for module in (self.discriminator, self.amp_normalizer):
        for tensor in module.state_dict().values():
            torch.distributed.broadcast(tensor, src=0)

save

save() -> dict

Source code in eden/learning/wrappers/rsl_rl/amp_ppo/ppo.py

def save(self) -> dict:
    saved = super().save()
    saved["amp_discriminator_state_dict"] = self.discriminator.state_dict()
    saved["amp_optimizer_state_dict"] = self.amp_optimizer.state_dict()
    saved["amp_normalizer_state_dict"] = self.amp_normalizer.state_dict()
    return saved

load

load(loaded_dict: dict, load_cfg: dict | None, strict: bool) -> bool

PPO's load, plus the discriminator, its optimizer and the frame normalizer.

A checkpoint without them was trained as plain PPO. Continuing it under AMP changes the reward under a value function trained without the style term, so it is refused unless amp_warm_start says that is the intent (the discriminator then starts fresh).

Source code in eden/learning/wrappers/rsl_rl/amp_ppo/ppo.py

def load(self, loaded_dict: dict, load_cfg: dict | None, strict: bool) -> bool:
    """PPO's load, plus the discriminator, its optimizer and the frame normalizer.

    A checkpoint without them was trained as plain PPO. Continuing it under AMP changes the reward under a value
    function trained without the style term, so it is refused unless ``amp_warm_start`` says that is the intent
    (the discriminator then starts fresh).
    """
    missing = [key for key in _AMP_CHECKPOINT_KEYS if key not in loaded_dict]
    if missing and not self.amp_warm_start:
        raise ValueError(
            f"This checkpoint has no AMP state (missing {missing}; it has {sorted(loaded_dict)}), so it was "
            "trained as plain PPO. Loading it into AmpPPO would add the style reward to a value function trained "
            "without it. To start AMP from its weights with a fresh discriminator, set amp_warm_start=True "
            "(--override 'runner_options: {algorithm: {amp_warm_start: true}}'); to replay it, use "
            "`eden inference`, which loads it under stock PPO."
        )
    resumed = super().load(loaded_dict, load_cfg, strict)
    if not missing:
        self.discriminator.load_state_dict(loaded_dict["amp_discriminator_state_dict"], strict=strict)
        if load_cfg is None or load_cfg.get("optimizer"):
            self.amp_optimizer.load_state_dict(loaded_dict["amp_optimizer_state_dict"])
        self.amp_normalizer.load_state_dict(loaded_dict["amp_normalizer_state_dict"])
    return resumed

train_mode

train_mode() -> None

Source code in eden/learning/wrappers/rsl_rl/amp_ppo/ppo.py

def train_mode(self) -> None:
    super().train_mode()
    self.discriminator.train()

eval_mode

eval_mode() -> None

Source code in eden/learning/wrappers/rsl_rl/amp_ppo/ppo.py

def eval_mode(self) -> None:
    super().eval_mode()
    self.discriminator.eval()

Adversarial-motion-prior components used by :class:~eden.learning.wrappers.rsl_rl.amp_ppo.ppo.AmpPPO.

Ports the adversarial-motion-prior machinery of TeleHuman/humanoid_skateboarding (its vendored rsl_rl fork: DiscriminatorMulti, G1_AMPLoader, ReplayBufferMulti, Normalizer):

• :class:AmpMotionSampler draws windows of num_frames consecutive frames (linear interpolation between recorded frames) from expert motion clips. A frame is one joint-position vector, in the order the env's AMP observation group emits it (HUSKY: the G1's 23 no-wrist joints in MJCF order).
• :class:AmpDiscriminator is the LSGAN discriminator (spectral-norm MLP, expert target +1, policy target -1); its raw output d maps to the style reward coef * clamp(1 - 0.25 (d - 1)^2, min=0).
• :class:AmpReplayBuffer is a flat ring buffer of policy windows.
• :class:FrameNormalizer is a running mean/std over single frames, applied per-frame to whole windows before the discriminator.

Deviation from the source: the source updated its normalizer with already-normalized batches (a bug that decays the statistics toward identity); this one updates from raw frames.

Classes

AmpMotionSampler

AmpMotionSampler(motion_file: str, num_frames: int, device: device | str)

Samples expert windows of num_frames consecutive frames from motion clips.

PARAMETER	DESCRIPTION
motion_file	Asset path of the .npz bundle (motion_<i> arrays of shape (T, frame_dim), plus fps and num_motions). TYPE: str
num_frames	Frames per window. The window spans [t - (num_frames - 2) dt, t + dt] around the sampled time t, matching the source loader's offsets. TYPE: int
device	Device the sampled windows live on. TYPE: device | str

Source code in eden/learning/wrappers/rsl_rl/amp_ppo/amp.py

def __init__(self, motion_file: str, num_frames: int, device: torch.device | str) -> None:
    data = np.load(get_asset_path(motion_file))
    self.device = device
    self.num_frames = num_frames
    self.dt = 1.0 / float(data["fps"])
    self.motions = [
        torch.as_tensor(data[f"motion_{i}"], dtype=torch.float32, device=device)
        for i in range(int(data["num_motions"]))
    ]
    self.frame_dim = int(self.motions[0].shape[1])
    for i, motion in enumerate(self.motions):
        if motion.shape[0] < num_frames + 2:
            raise ValueError(f"motion_{i} in {motion_file} is too short ({motion.shape[0]} frames)")

Attributes

device instance-attribute

device = device

num_frames instance-attribute

num_frames = num_frames

dt instance-attribute

dt = 1.0 / float(data['fps'])

motions instance-attribute

motions = [torch.as_tensor(data[f'motion_{i}'], dtype=torch.float32, device=device) for i in range(int(data['num_motions']))]

frame_dim instance-attribute

frame_dim = int(self.motions[0].shape[1])

Methods:

sample

sample(batch_size: int) -> Tensor

Return expert windows, shape (batch_size, num_frames, frame_dim), oldest frame first.

Source code in eden/learning/wrappers/rsl_rl/amp_ppo/amp.py

def sample(self, batch_size: int) -> torch.Tensor:
    """Return expert windows, shape (batch_size, num_frames, frame_dim), oldest frame first."""
    counts = torch.multinomial(
        torch.ones(len(self.motions), device=self.device), batch_size, replacement=True
    ).bincount(minlength=len(self.motions))
    windows = []
    offsets = torch.arange(-(self.num_frames - 2), 2, device=self.device, dtype=torch.float32)
    for motion, count in zip(self.motions, counts.tolist()):
        if count == 0:
            continue
        num_rows = motion.shape[0]
        # Continuous frame index such that every offset stays in [0, num_rows - 1].
        base = torch.empty(count, device=self.device).uniform_(self.num_frames - 2, num_rows - 2)
        idx = base.unsqueeze(1) + offsets.unsqueeze(0)  # (count, num_frames)
        low = idx.floor().long()
        frac = (idx - low.float()).unsqueeze(-1)
        frames = motion[low] * (1.0 - frac) + motion[low + 1] * frac
        windows.append(frames)
    return torch.cat(windows, dim=0)

AmpReplayBuffer

AmpReplayBuffer(window_dim: int, capacity: int, device: device | str)

Flat ring buffer of policy windows (one row = one flattened window).

Source code in eden/learning/wrappers/rsl_rl/amp_ppo/amp.py

def __init__(self, window_dim: int, capacity: int, device: torch.device | str) -> None:
    self._buffer = torch.zeros(capacity, window_dim, device=device)
    self._capacity = capacity
    self._next = 0
    self._size = 0

Methods:

insert

insert(windows: Tensor) -> None

Insert a batch of flattened windows, overwriting the oldest entries.

Source code in eden/learning/wrappers/rsl_rl/amp_ppo/amp.py

def insert(self, windows: torch.Tensor) -> None:
    """Insert a batch of flattened windows, overwriting the oldest entries."""
    n = windows.shape[0]
    idx = (self._next + torch.arange(n, device=windows.device)) % self._capacity
    self._buffer[idx] = windows.detach().clone()
    self._next = int((self._next + n) % self._capacity)
    self._size = min(self._size + n, self._capacity)

sample

sample(batch_size: int) -> Tensor

Source code in eden/learning/wrappers/rsl_rl/amp_ppo/amp.py

def sample(self, batch_size: int) -> torch.Tensor:
    if self._size == 0:
        raise ValueError("AmpReplayBuffer is empty — insert policy windows before sampling.")
    idx = torch.randint(self._size, (batch_size,), device=self._buffer.device)
    return self._buffer[idx]

FrameNormalizer

FrameNormalizer(frame_dim: int, clip: float = 10.0, eps: float = 0.01)

Bases: Module

Running mean/std over single frames, applied per-frame to windows.

Source code in eden/learning/wrappers/rsl_rl/amp_ppo/amp.py

def __init__(self, frame_dim: int, clip: float = 10.0, eps: float = 1e-2) -> None:
    super().__init__()
    self.clip = clip
    self.eps = eps
    self.register_buffer("mean", torch.zeros(frame_dim))
    self.register_buffer("var", torch.ones(frame_dim))
    self.register_buffer("count", torch.tensor(1e-4))

Attributes

clip instance-attribute

clip = clip

eps instance-attribute

eps = eps

Methods:

update

update(frames: Tensor) -> None

Fold a batch of raw frames (..., frame_dim) into the running statistics.

Source code in eden/learning/wrappers/rsl_rl/amp_ppo/amp.py

@torch.no_grad()
def update(self, frames: torch.Tensor) -> None:
    """Fold a batch of raw frames (..., frame_dim) into the running statistics."""
    flat = frames.reshape(-1, frames.shape[-1])
    batch_count = flat.shape[0]
    if batch_count == 0:
        return
    batch_mean = flat.mean(dim=0)
    batch_var = flat.var(dim=0, unbiased=False)
    delta = batch_mean - self.mean
    total = self.count + batch_count
    self.mean += delta * batch_count / total
    m_a = self.var * self.count
    m_b = batch_var * batch_count
    self.var.copy_((m_a + m_b + delta.square() * self.count * batch_count / total) / total)
    self.count.copy_(total)

normalize

normalize(frames: Tensor) -> Tensor

Normalize frames or windows of frames (..., frame_dim).

Source code in eden/learning/wrappers/rsl_rl/amp_ppo/amp.py

def normalize(self, frames: torch.Tensor) -> torch.Tensor:
    """Normalize frames or windows of frames (..., frame_dim)."""
    std = (self.var + self.eps**2).sqrt()
    return ((frames - self.mean) / std).clamp(-self.clip, self.clip)

AmpDiscriminator

AmpDiscriminator(window_dim: int, hidden_dims: tuple[int, ...] = (256, 256))

Bases: Module

LSGAN discriminator over flattened windows (spectral-norm MLP + linear head).

Source code in eden/learning/wrappers/rsl_rl/amp_ppo/amp.py

def __init__(self, window_dim: int, hidden_dims: tuple[int, ...] = (256, 256)) -> None:
    super().__init__()
    layers: list[nn.Module] = []
    in_dim = window_dim
    for hidden in hidden_dims:
        layers += [spectral_norm(nn.Linear(in_dim, hidden)), nn.ReLU()]
        in_dim = hidden
    self.trunk = nn.Sequential(*layers)
    self.head = spectral_norm(nn.Linear(in_dim, 1))

Attributes

trunk instance-attribute

trunk = nn.Sequential(*layers)

head instance-attribute

head = spectral_norm(nn.Linear(in_dim, 1))

Methods:

forward

forward(windows: Tensor) -> Tensor

Source code in eden/learning/wrappers/rsl_rl/amp_ppo/amp.py

def forward(self, windows: torch.Tensor) -> torch.Tensor:
    return self.head(self.trunk(windows))

style_reward

style_reward(windows: Tensor, coef: float) -> Tensor

Per-window style reward coef * clamp(1 - 0.25 (d - 1)^2, min=0), shape (N,).

Source code in eden/learning/wrappers/rsl_rl/amp_ppo/amp.py

def style_reward(self, windows: torch.Tensor, coef: float) -> torch.Tensor:
    """Per-window style reward ``coef * clamp(1 - 0.25 (d - 1)^2, min=0)``, shape (N,)."""
    with torch.no_grad():
        d = self.forward(windows).squeeze(-1)
    return coef * torch.clamp(1.0 - 0.25 * torch.square(d - 1.0), min=0.0)

gradient_penalty

gradient_penalty(expert_windows: Tensor, scale: float) -> Tensor

Penalize the discriminator's gradient norm at the expert samples.

Source code in eden/learning/wrappers/rsl_rl/amp_ppo/amp.py

def gradient_penalty(self, expert_windows: torch.Tensor, scale: float) -> torch.Tensor:
    """Penalize the discriminator's gradient norm at the expert samples."""
    expert = expert_windows.detach().clone().requires_grad_(True)
    d = self.head(self.trunk(expert))
    grad = torch.autograd.grad(
        outputs=d, inputs=expert, grad_outputs=torch.ones_like(d), create_graph=True, retain_graph=True
    )[0]
    return scale * grad.norm(2, dim=1).square().mean()

Functions:

RMA teacher and depth-student distillation

Options for the RMA teacher (actor model + PPO) and its depth-student distillation.

Declarative only — importing this module pulls in neither rsl_rl's model/algorithm API nor torch modules, so a task config that uses it still imports under any rsl_rl Eden supports. The classes their class_name fields name need rsl-rl-lib >= 5.5.1 and say so when imported without it. Defaults are the extreme-parkour teacher's (Cheng et al. 2023).

Classes

RslRlRmaActorOptions

RslRlRmaActorOptions(**data)

Bases: RslRlModelWithDistributionOptions

The RMA actor: estimator + scan / privileged / history encoders + backbone over one flat observation.

The actor observation set must concatenate to [prop | scan | priv_explicit | priv_latent | num_hist x prop] with the widths below (checked when the model is built).

PARAMETER	DESCRIPTION
class_name	The model class rsl_rl resolves; leave at the default. TYPE: str
num_prop	Chunk widths of the flat observation (num_hist frames of the num_prop-d proprioception). TYPE: int
num_scan	Chunk widths of the flat observation (num_hist frames of the num_prop-d proprioception). TYPE: int
num_priv_explicit	Chunk widths of the flat observation (num_hist frames of the num_prop-d proprioception). TYPE: int
num_priv_latent	Chunk widths of the flat observation (num_hist frames of the num_prop-d proprioception). TYPE: int
num_hist	Chunk widths of the flat observation (num_hist frames of the num_prop-d proprioception). TYPE: int
scan_encoder_dims	Scan encoder widths; the last is the scan latent width (Tanh-bounded). TYPE: list[int]
priv_encoder_dims	Privileged encoder widths; the last is the adaptation latent width, which the history encoder matches. TYPE: list[int]
hidden_dims	Actor backbone hidden widths. TYPE: list[int]
activation	Activation of every encoder and the backbone. TYPE: str
estimator_hidden_dims	Hidden widths of the base-velocity estimator (proprioception -> privileged-explicit chunk). TYPE: list[int]
tanh_backbone_output	Squash the action mean with Tanh. TYPE: bool
obs_normalization	Must stay False: the chunks are sliced by position and carry the task's fixed scales. TYPE: bool
distribution_cfg	A state-independent Gaussian (std_type="scalar" is the teacher's parameterization). TYPE: RslRlGaussianDistributionOptions

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

class_name class-attribute instance-attribute

class_name: str = _RMA_ACTOR_MODEL_CLASS_NAME

num_prop class-attribute instance-attribute

num_prop: int = 53

num_scan class-attribute instance-attribute

num_scan: int = 132

num_priv_explicit class-attribute instance-attribute

num_priv_explicit: int = 9

num_priv_latent class-attribute instance-attribute

num_priv_latent: int = 29

num_hist class-attribute instance-attribute

num_hist: int = 10

scan_encoder_dims class-attribute instance-attribute

scan_encoder_dims: list[int] = [128, 64, 32]

priv_encoder_dims class-attribute instance-attribute

priv_encoder_dims: list[int] = [64, 20]

hidden_dims class-attribute instance-attribute

hidden_dims: list[int] = [512, 256, 128]

activation class-attribute instance-attribute

activation: str = 'elu'

estimator_hidden_dims class-attribute instance-attribute

estimator_hidden_dims: list[int] = [128, 64]

tanh_backbone_output class-attribute instance-attribute

tanh_backbone_output: bool = False

obs_normalization class-attribute instance-attribute

obs_normalization: bool = False

distribution_cfg class-attribute instance-attribute

distribution_cfg: RslRlGaussianDistributionOptions = RslRlGaussianDistributionOptions(init_std=1.0, std_type='scalar')

RslRlRmaPpoAlgorithmOptions

RslRlRmaPpoAlgorithmOptions(**data)

Bases: RslRlPpoAlgorithmOptions

PPO for an RMA actor: concurrent estimator, privileged-latent regularization, DAgger, optional mirror loss.

Needs actor=RslRlRmaActorOptions(...). The PPO defaults are the extreme-parkour teacher's.

PARAMETER	DESCRIPTION
class_name	The algorithm class rsl_rl resolves; leave at the default. TYPE: str
estimator_learning_rate	Learning rate of the estimator's own Adam. TYPE: float
train_with_estimated_states	Act on the estimator's prediction of the privileged-explicit chunk during rollouts (storage keeps the truth). TYPE: bool
dagger_update_freq	Every this-many learning iterations, roll out with the history latent and regress the history encoder onto the privileged latent afterwards. TYPE: int
priv_reg_coef_schedule	(start, end, ramp_start, ramp_length) of the privileged-regularization coefficient, ramped linearly over update counts (PPO and DAgger updates both count). TYPE: list[float]
symmetry_coef	Weight of the mirror-equivariance loss (0 disables it). TYPE: float
mirror_class_name	"module.path:Class" of the task's left-right mirror, built as Class(dims, device) with mirror_obs(flat) / mirror_action(actions). Required when symmetry_coef > 0. TYPE: str | None

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

class_name class-attribute instance-attribute

class_name: str = _RMA_PPO_CLASS_NAME

learning_rate class-attribute instance-attribute

learning_rate: float = 0.0002

entropy_coef class-attribute instance-attribute

entropy_coef: float = 0.01

estimator_learning_rate class-attribute instance-attribute

estimator_learning_rate: float = 0.0001

train_with_estimated_states class-attribute instance-attribute

train_with_estimated_states: bool = True

dagger_update_freq class-attribute instance-attribute

dagger_update_freq: int = 20

priv_reg_coef_schedule class-attribute instance-attribute

priv_reg_coef_schedule: list[float] = [0.0, 0.1, 2000, 3000]

symmetry_coef class-attribute instance-attribute

symmetry_coef: float = 0.0

mirror_class_name class-attribute instance-attribute

mirror_class_name: str | None = None

RslRlRmaDepthDistillationAlgorithmOptions

RslRlRmaDepthDistillationAlgorithmOptions(**data)

Bases: ConfigurableOptions

DAgger distillation of an RMA teacher into a depth student (recurrent depth encoder + copied actor net).

PARAMETER	DESCRIPTION
class_name	The algorithm class the runner resolves; leave at the default. TYPE: str
learning_rate	Adam learning rate of the student (depth actor + depth encoder). TYPE: float
max_grad_norm	Gradient-norm clip on the depth actor. TYPE: float
gru_hidden_dim	Hidden width of the depth encoder's GRU. TYPE: int
update_interval	Control steps between depth-encoder passes; must equal the env's depth refresh cadence. TYPE: int
yaw_scale	Stretch of the Tanh-bounded yaw head. TYPE: float
yaw_slots	(start, stop) of the two delta-yaw slots inside the proprioception chunk, which the student predicts. TYPE: list[int]
num_pretrain_iters	Iterations driven by the teacher's actions before the student takes over. TYPE: int
depth_obs_group	Observation group carrying the processed depth frame. TYPE: str
yaw_gate_obs_group	Observation group whose flag (> 0.5) lets the student see its own yaw prediction. TYPE: str

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

class_name class-attribute instance-attribute

class_name: str = _RMA_DEPTH_DISTILLATION_CLASS_NAME

learning_rate class-attribute instance-attribute

learning_rate: float = 0.001

max_grad_norm class-attribute instance-attribute

max_grad_norm: float = 1.0

gru_hidden_dim class-attribute instance-attribute

gru_hidden_dim: int = 512

update_interval class-attribute instance-attribute

update_interval: int = 5

yaw_scale class-attribute instance-attribute

yaw_scale: float = 1.5

yaw_slots class-attribute instance-attribute

yaw_slots: list[int] = [6, 8]

num_pretrain_iters class-attribute instance-attribute

num_pretrain_iters: int = 0

depth_obs_group class-attribute instance-attribute

depth_obs_group: str = 'depth_camera'

yaw_gate_obs_group class-attribute instance-attribute

yaw_gate_obs_group: str = 'delta_yaw_ok'

RslRlRmaDepthDistillationRunnerOptions

RslRlRmaDepthDistillationRunnerOptions(**data)

Bases: RslRlBaseRunnerOptions

Runner options selecting the depth-student DAgger loop; train with --checkpoint <teacher checkpoint>.

PARAMETER	DESCRIPTION
class_name	The runner eden train builds; leave at the default. TYPE: str
actor	The teacher's architecture — it must match the teacher checkpoint being distilled. TYPE: RslRlRmaActorOptions
algorithm	The distillation's knobs. TYPE: RslRlRmaDepthDistillationAlgorithmOptions

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

class_name class-attribute instance-attribute

class_name: str = _RMA_DEPTH_DISTILLATION_RUNNER_CLASS_NAME

actor class-attribute instance-attribute

actor: RslRlRmaActorOptions = RslRlRmaActorOptions()

algorithm class-attribute instance-attribute

algorithm: RslRlRmaDepthDistillationAlgorithmOptions = RslRlRmaDepthDistillationAlgorithmOptions()

RmaActorModel

RmaActorModel(obs: TensorDict, obs_groups: dict[str, list[str]], obs_set: str, output_dim: int, num_prop: int = 53, num_scan: int = 132, num_priv_explicit: int = 9, num_priv_latent: int = 29, num_hist: int = 10, scan_encoder_dims: tuple[int, ...] | list[int] = (128, 64, 32), priv_encoder_dims: tuple[int, ...] | list[int] = (64, 20), hidden_dims: tuple[int, ...] | list[int] = (512, 256, 128), activation: str = 'elu', estimator_hidden_dims: tuple[int, ...] | list[int] = (128, 64), tanh_backbone_output: bool = False, obs_normalization: bool = False, distribution_cfg: dict | None = None)

Bases: MLPModel

RMA actor over one flat [prop | scan | priv_explicit | priv_latent | history] observation.

Subclasses rsl_rl's MLPModel for its distribution plumbing (output_*, log-prob, KL, entropy) and its recurrent no-ops, but replaces the trunk: there is no self.mlp, and observation normalization is refused — the task's terms carry hand-tuned fixed scales, and the chunks are sliced by position.

PARAMETER	DESCRIPTION
obs	rsl_rl's model arguments; the groups of obs_set concatenate to the flat layout. TYPE: TensorDict
obs_groups	rsl_rl's model arguments; the groups of obs_set concatenate to the flat layout. TYPE: TensorDict
obs_set	rsl_rl's model arguments; the groups of obs_set concatenate to the flat layout. TYPE: TensorDict
output_dim	rsl_rl's model arguments; the groups of obs_set concatenate to the flat layout. TYPE: TensorDict
num_prop	Chunk widths (:class:~eden.learning.wrappers.rsl_rl.rma.dims.RmaObsDims). TYPE: int DEFAULT: 53
num_scan	Chunk widths (:class:~eden.learning.wrappers.rsl_rl.rma.dims.RmaObsDims). TYPE: int DEFAULT: 53
num_priv_explicit	Chunk widths (:class:~eden.learning.wrappers.rsl_rl.rma.dims.RmaObsDims). TYPE: int DEFAULT: 53
num_priv_latent	Chunk widths (:class:~eden.learning.wrappers.rsl_rl.rma.dims.RmaObsDims). TYPE: int DEFAULT: 53
num_hist	Chunk widths (:class:~eden.learning.wrappers.rsl_rl.rma.dims.RmaObsDims). TYPE: int DEFAULT: 53
scan_encoder_dims	Encoder widths; the last entry of each is its latent width. TYPE: list[int] DEFAULT: (128, 64, 32)
priv_encoder_dims	Encoder widths; the last entry of each is its latent width. TYPE: list[int] DEFAULT: (128, 64, 32)
hidden_dims	Actor backbone hidden widths. TYPE: list[int] DEFAULT: (512, 256, 128)
activation	Activation everywhere but the scan latent (Tanh) and the estimator output (linear). TYPE: str DEFAULT: 'elu'
estimator_hidden_dims	Hidden widths of the base-velocity estimator. TYPE: list[int] DEFAULT: (128, 64)
tanh_backbone_output	Squash the action mean with Tanh. TYPE: bool DEFAULT: False
obs_normalization	Must be False. TYPE: bool DEFAULT: False
distribution_cfg	The output distribution (a state-independent Gaussian for the RMA teacher). TYPE: dict DEFAULT: None

Source code in eden/learning/wrappers/rsl_rl/rma/model.py

def __init__(
    self,
    obs: TensorDict,
    obs_groups: dict[str, list[str]],
    obs_set: str,
    output_dim: int,
    num_prop: int = 53,
    num_scan: int = 132,
    num_priv_explicit: int = 9,
    num_priv_latent: int = 29,
    num_hist: int = 10,
    scan_encoder_dims: tuple[int, ...] | list[int] = (128, 64, 32),
    priv_encoder_dims: tuple[int, ...] | list[int] = (64, 20),
    hidden_dims: tuple[int, ...] | list[int] = (512, 256, 128),
    activation: str = "elu",
    estimator_hidden_dims: tuple[int, ...] | list[int] = (128, 64),
    tanh_backbone_output: bool = False,
    obs_normalization: bool = False,
    distribution_cfg: dict | None = None,
) -> None:
    nn.Module.__init__(self)  # not MLPModel.__init__: it would build the plain MLP trunk this model replaces
    if obs_normalization:
        raise ValueError(
            "RmaActorModel(obs_normalization=True): the RMA actor slices its observation by position and trains "
            "on the task's fixed term scales; set obs_normalization=False on the actor options."
        )
    if distribution_cfg is None:
        raise ValueError(
            "RmaActorModel needs a distribution_cfg (the RMA teacher's actor is stochastic); pass "
            "RslRlGaussianDistributionOptions(init_std=1.0, std_type='scalar')."
        )
    self.obs_groups, self.obs_dim = self._get_obs_dim(obs, obs_groups, obs_set)
    self.dims = RmaObsDims(num_prop, num_scan, num_priv_explicit, num_priv_latent, num_hist)
    if self.obs_dim != self.dims.total:
        raise ValueError(
            f"The '{obs_set}' observation set {self.obs_groups} is {self.obs_dim}-d, but the RMA actor options "
            f"describe {self.dims.layout()}. Make the task's {obs_set} groups emit that layout, or set the "
            "actor options' num_prop / num_scan / num_priv_explicit / num_priv_latent / num_hist to the task's."
        )
    self.obs_normalization = False
    self.obs_normalizer = nn.Identity()

    dist_class, dist_cfg = resolve_class(distribution_cfg)
    self.distribution = dist_class(output_dim, **dist_cfg)
    if self.distribution.input_dim != output_dim:
        raise ValueError(
            f"RmaActorModel's backbone emits the {output_dim}-d action mean only; {dist_class.__name__} needs "
            f"{self.distribution.input_dim} inputs. Use a state-independent Gaussian distribution."
        )
    act = resolve_nn_activation(activation)
    self.net = RmaActor(
        self.dims,
        output_dim,
        list(scan_encoder_dims),
        list(priv_encoder_dims),
        list(hidden_dims),
        act,
        tanh_backbone_output=tanh_backbone_output,
    )
    self.estimator = Estimator(num_prop, num_priv_explicit, list(estimator_hidden_dims), act)

Attributes

dims instance-attribute

dims = RmaObsDims(num_prop, num_scan, num_priv_explicit, num_priv_latent, num_hist)

obs_normalization instance-attribute

obs_normalization = False

obs_normalizer instance-attribute

obs_normalizer = nn.Identity()

distribution instance-attribute

distribution = dist_class(output_dim, **dist_cfg)

net instance-attribute

net = RmaActor(self.dims, output_dim, list(scan_encoder_dims), list(priv_encoder_dims), list(hidden_dims), act, tanh_backbone_output=tanh_backbone_output)

estimator instance-attribute

estimator = Estimator(num_prop, num_priv_explicit, list(estimator_hidden_dims), act)

Methods:

forward

forward(obs: TensorDict, masks: Tensor | None = None, hidden_state=None, stochastic_output: bool = False, *, hist_encoding: bool = True, estimate: bool = True) -> Tensor

Action (mean, or a sample with stochastic_output); the switches are described in the module docstring.

Source code in eden/learning/wrappers/rsl_rl/rma/model.py

def forward(
    self,
    obs: TensorDict,
    masks: torch.Tensor | None = None,
    hidden_state=None,
    stochastic_output: bool = False,
    *,
    hist_encoding: bool = True,
    estimate: bool = True,
) -> torch.Tensor:
    """Action (mean, or a sample with ``stochastic_output``); the switches are described in the module docstring."""
    mean = self.mean(self.get_flat_obs(obs), hist_encoding=hist_encoding, estimate=estimate)
    if stochastic_output:
        self.distribution.update(mean)
        return self.distribution.sample()
    return self.distribution.deterministic_output(mean)

mean

mean(flat: Tensor, *, hist_encoding: bool, estimate: bool) -> Tensor

Action mean from the flat observation (no distribution update).

Source code in eden/learning/wrappers/rsl_rl/rma/model.py

def mean(self, flat: torch.Tensor, *, hist_encoding: bool, estimate: bool) -> torch.Tensor:
    """Action mean from the flat observation (no distribution update)."""
    if estimate:
        flat = self.with_estimate(flat)
    return self.net(flat, hist_encoding)

with_estimate

with_estimate(flat: Tensor) -> Tensor

A copy of flat whose privileged-explicit chunk is the estimator's prediction (detached).

Detached: the policy loss never trains the estimator — it has its own regression loss and optimizer.

Source code in eden/learning/wrappers/rsl_rl/rma/model.py

def with_estimate(self, flat: torch.Tensor) -> torch.Tensor:
    """A copy of ``flat`` whose privileged-explicit chunk is the estimator's prediction (detached).

    Detached: the policy loss never trains the estimator — it has its own regression loss and optimizer.
    """
    d = self.dims
    flat = flat.clone()
    flat[:, d.priv_explicit_slice] = self.estimator(flat[:, d.prop_slice]).detach()
    return flat

get_flat_obs

get_flat_obs(obs: TensorDict) -> Tensor

The concatenated observation groups of this model's set.

Source code in eden/learning/wrappers/rsl_rl/rma/model.py

def get_flat_obs(self, obs: TensorDict) -> torch.Tensor:
    """The concatenated observation groups of this model's set."""
    return torch.cat([obs[group] for group in self.obs_groups], dim=-1)

policy_parameters

policy_parameters()

Every parameter but the estimator's — what the PPO optimizer and its gradient clip cover.

Source code in eden/learning/wrappers/rsl_rl/rma/model.py

def policy_parameters(self):
    """Every parameter but the estimator's — what the PPO optimizer and its gradient clip cover."""
    return (p for name, p in self.named_parameters() if not name.startswith("estimator."))

get_latent

get_latent(obs, masks=None, hidden_state=None) -> Tensor

Source code in eden/learning/wrappers/rsl_rl/rma/model.py

def get_latent(self, obs, masks=None, hidden_state=None) -> torch.Tensor:
    return self.get_flat_obs(obs)

as_jit

as_jit() -> Module

A TorchScript-compatible deployment policy over the flat observation.

Source code in eden/learning/wrappers/rsl_rl/rma/model.py

def as_jit(self) -> nn.Module:
    """A TorchScript-compatible deployment policy over the flat observation."""
    return _ExportedRmaActor(self)

as_onnx

as_onnx(verbose: bool = False) -> Module

An ONNX-exportable deployment policy over the flat observation.

verbose is unused and optional: rsl-rl-lib 5.5.2 calls as_onnx(), earlier 5.x as_onnx(verbose=...).

Source code in eden/learning/wrappers/rsl_rl/rma/model.py

def as_onnx(self, verbose: bool = False) -> nn.Module:
    """An ONNX-exportable deployment policy over the flat observation.

    ``verbose`` is unused and optional: rsl-rl-lib 5.5.2 calls ``as_onnx()``, earlier 5.x ``as_onnx(verbose=...)``.
    """
    return _ExportedRmaActor(self)

RmaPPO

RmaPPO(actor: RmaActorModel, critic, storage, *, estimator_learning_rate: float = 0.0001, train_with_estimated_states: bool = True, dagger_update_freq: int = 20, priv_reg_coef_schedule: tuple[float, float, int, int] | list = (0.0, 0.1, 2000, 3000), symmetry_coef: float = 0.0, mirror_class_name: str | None = None, optimizer: str = 'adam', **ppo_kwargs)

Bases: PPO

PPO + estimator + privileged-latent regularization + DAgger (+ mirror loss) for an :class:RmaActorModel.

The extra keyword arguments are documented on :class:~eden.learning.wrappers.rsl_rl.rma.options.RslRlRmaPpoAlgorithmOptions; the rest go to rsl_rl's PPO.

Source code in eden/learning/wrappers/rsl_rl/rma/ppo.py

def __init__(
    self,
    actor: RmaActorModel,
    critic,
    storage,
    *,
    estimator_learning_rate: float = 1.0e-4,
    train_with_estimated_states: bool = True,
    dagger_update_freq: int = 20,
    priv_reg_coef_schedule: tuple[float, float, int, int] | list = (0.0, 0.1, 2000, 3000),
    symmetry_coef: float = 0.0,
    mirror_class_name: str | None = None,
    optimizer: str = "adam",
    **ppo_kwargs,
) -> None:
    if not isinstance(actor, RmaActorModel):
        raise TypeError(
            f"RmaPPO trains an RmaActorModel, got {type(actor).__name__}; set the runner's actor to "
            "RslRlRmaActorOptions(...) (eden.learning.wrappers.rsl_rl.rma)."
        )
    for key in ("rnd_cfg", "symmetry_cfg", "multi_gpu_cfg"):
        if ppo_kwargs.get(key) is not None:
            raise ValueError(
                f"RmaPPO does not support {key} (got {ppo_kwargs[key]!r}): leave it None. Its mirror loss is "
                "symmetry_coef + mirror_class_name, and it trains on one GPU — launch without torchrun."
            )
    if ppo_kwargs.get("use_mixed_precision"):
        raise ValueError("RmaPPO does not support use_mixed_precision=True; leave it False.")
    if critic.is_recurrent:
        raise ValueError(
            f"RmaPPO needs a feed-forward critic, got the recurrent {type(critic).__name__}: its update draws flat "
            "mini-batches without hidden states. Use RslRlCriticOptions (an MLP) for the critic."
        )
    if dagger_update_freq < 1:
        raise ValueError(f"RmaPPO(dagger_update_freq={dagger_update_freq}) must be >= 1.")
    if symmetry_coef > 0.0 and mirror_class_name is None:
        raise ValueError(
            f"RmaPPO(symmetry_coef={symmetry_coef}) needs mirror_class_name='module.path:Class' — a class built "
            "as Class(dims, device) with mirror_obs(flat) / mirror_action(actions) for the task's layout "
            "(e.g. 'eden.tasks.extreme_parkour.symmetry:ParkourMirror'); or set symmetry_coef=0.0."
        )
    super().__init__(actor, critic, storage, optimizer=optimizer, **ppo_kwargs)
    # PPO's optimizer without the estimator, which trains on its own regression loss (and optimizer) below.
    self.optimizer = resolve_optimizer(optimizer)(
        chain(self.actor.policy_parameters(), self.critic.parameters()), lr=self.learning_rate
    )
    self.estimator_optimizer = torch.optim.Adam(self.actor.estimator.parameters(), lr=estimator_learning_rate)
    self.hist_encoder_optimizer = torch.optim.Adam(
        self.actor.net.history_encoder.parameters(), lr=self.learning_rate
    )
    self.train_with_estimated_states = train_with_estimated_states
    self.dagger_update_freq = dagger_update_freq
    self.priv_reg_coef_schedule = tuple(priv_reg_coef_schedule)
    self.symmetry_coef = symmetry_coef
    self.mirror = resolve_callable(mirror_class_name)(self.actor.dims, self.device) if symmetry_coef > 0.0 else None
    # Drives the priv-reg ramp; advances on every PPO update and every DAgger update.
    self.counter = 0
    # The learning iteration whose rollout is being collected (the runner's `it`); picks DAgger iterations.
    self.iteration = 0

Attributes

actor instance-attribute

actor: RmaActorModel

distributed_unsupported class-attribute instance-attribute

distributed_unsupported = 'RmaPPO trains on one GPU: its estimator and history-encoder optimizers have no gradient reduction, so a distributed launch would train them independently per rank. Launch the RMA teacher as one process (without torchrun).'

optimizer instance-attribute

optimizer = resolve_optimizer(optimizer)(chain(self.actor.policy_parameters(), self.critic.parameters()), lr=self.learning_rate)

estimator_optimizer instance-attribute

estimator_optimizer = torch.optim.Adam(self.actor.estimator.parameters(), lr=estimator_learning_rate)

hist_encoder_optimizer instance-attribute

hist_encoder_optimizer = torch.optim.Adam(self.actor.net.history_encoder.parameters(), lr=self.learning_rate)

train_with_estimated_states instance-attribute

train_with_estimated_states = train_with_estimated_states

dagger_update_freq instance-attribute

dagger_update_freq = dagger_update_freq

priv_reg_coef_schedule instance-attribute

priv_reg_coef_schedule = tuple(priv_reg_coef_schedule)

symmetry_coef instance-attribute

symmetry_coef = symmetry_coef

mirror instance-attribute

mirror = resolve_callable(mirror_class_name)(self.actor.dims, self.device) if symmetry_coef > 0.0 else None

counter instance-attribute

counter = 0

iteration instance-attribute

iteration = 0

is_dagger_iteration property

is_dagger_iteration: bool

Whether the current rollout uses the history latent and is followed by a DAgger update.

Methods:

act

act(obs: TensorDict) -> Tensor

Sample actions — through the estimator and, on DAgger iterations, the history latent — and record them.

Source code in eden/learning/wrappers/rsl_rl/rma/ppo.py

def act(self, obs: TensorDict) -> torch.Tensor:
    """Sample actions — through the estimator and, on DAgger iterations, the history latent — and record them."""
    self.transition.hidden_states = (self.actor.get_hidden_state(), self.critic.get_hidden_state())
    self.transition.actions = self.actor(
        obs,
        stochastic_output=True,
        hist_encoding=self.is_dagger_iteration,
        estimate=self.train_with_estimated_states,
    ).detach()
    self.transition.values = self.critic(obs).detach()
    self.transition.actions_log_prob = self.actor.get_output_log_prob(self.transition.actions).detach()
    self.transition.distribution_params = tuple(p.detach() for p in self.actor.output_distribution_params)
    self.transition.observations = obs  # ground truth: the estimator's target, the critic's input
    return self.transition.actions

update

update() -> dict[str, float]

PPO epochs with the estimator, priv-reg and mirror terms; then DAgger on DAgger iterations.

Source code in eden/learning/wrappers/rsl_rl/rma/ppo.py

def update(self) -> dict[str, float]:
    """PPO epochs with the estimator, priv-reg and mirror terms; then DAgger on DAgger iterations."""
    actor, dims = self.actor, self.actor.dims
    mean_value_loss = mean_surrogate_loss = mean_entropy = 0.0
    mean_priv_reg_loss = mean_estimator_loss = mean_symmetry_loss = 0.0
    priv_reg_coef = 0.0
    start, end, ramp_start, ramp_len = self.priv_reg_coef_schedule

    for batch in self.storage.mini_batch_generator(self.num_mini_batches, self.num_learning_epochs):
        if self.normalize_advantage_per_mini_batch:  # else compute_returns normalized them over the rollout
            with torch.no_grad():
                batch.advantages = (batch.advantages - batch.advantages.mean()) / (batch.advantages.std() + 1e-8)
        # PPO always trains through the privileged path on the stored ground truth; the history encoder trains
        # in the DAgger update only.
        actor(batch.observations, stochastic_output=True, hist_encoding=False, estimate=False)
        actions_log_prob = actor.get_output_log_prob(batch.actions)
        values = self.critic(batch.observations)
        distribution_params = actor.output_distribution_params
        mu = actor.output_mean
        entropy = actor.output_entropy
        flat = actor.get_flat_obs(batch.observations)

        # Privileged regularization toward the frozen history latent, coefficient ramped by the counter.
        priv_latent = actor.net.infer_priv_latent(flat)
        with torch.no_grad():
            hist_latent = actor.net.infer_hist_latent(flat)
        priv_reg_loss = (priv_latent - hist_latent).norm(p=2, dim=1).mean()
        priv_reg_coef = start + min(max(self.counter - ramp_start, 0) / ramp_len, 1.0) * (end - start)

        # Estimator step on its own optimizer: the stored obs carry the ground-truth priv_explicit.
        predicted = actor.estimator(flat[:, dims.prop_slice])
        estimator_loss = (predicted - flat[:, dims.priv_explicit_slice]).pow(2).mean()
        self.estimator_optimizer.zero_grad()
        estimator_loss.backward()
        nn.utils.clip_grad_norm_(actor.estimator.parameters(), self.max_grad_norm)
        self.estimator_optimizer.step()

        if self.desired_kl is not None and self.schedule == "adaptive":
            with torch.inference_mode():
                kl_mean = actor.get_kl_divergence(batch.old_distribution_params, distribution_params).mean()
                if kl_mean > self.desired_kl * 2.0:
                    self.learning_rate = max(1e-5, self.learning_rate / 1.5)
                elif kl_mean < self.desired_kl / 2.0 and kl_mean > 0.0:
                    self.learning_rate = min(1e-2, self.learning_rate * 1.5)
                for param_group in self.optimizer.param_groups:
                    param_group["lr"] = self.learning_rate

        ratio = torch.exp(actions_log_prob - torch.squeeze(batch.old_actions_log_prob))
        surrogate = -torch.squeeze(batch.advantages) * ratio
        surrogate_clipped = -torch.squeeze(batch.advantages) * torch.clamp(
            ratio, 1.0 - self.clip_param, 1.0 + self.clip_param
        )
        surrogate_loss = torch.max(surrogate, surrogate_clipped).mean()

        if self.use_clipped_value_loss:
            value_clipped = batch.values + (values - batch.values).clamp(-self.clip_param, self.clip_param)
            value_loss = torch.max((values - batch.returns).pow(2), (value_clipped - batch.returns).pow(2)).mean()
        else:
            value_loss = (batch.returns - values).pow(2).mean()

        # Mirror equivariance on the action mean; both sides carry gradient.
        symmetry_loss = torch.zeros((), device=self.device)
        if self.mirror is not None:
            mu_mirrored = actor.mean(self.mirror.mirror_obs(flat), hist_encoding=False, estimate=False)
            symmetry_loss = (mu_mirrored - self.mirror.mirror_action(mu)).pow(2).mean()

        loss = (
            surrogate_loss
            + self.value_loss_coef * value_loss
            - self.entropy_coef * entropy.mean()
            + priv_reg_coef * priv_reg_loss
            + self.symmetry_coef * symmetry_loss
        )
        self.optimizer.zero_grad()
        loss.backward()
        nn.utils.clip_grad_norm_(chain(actor.policy_parameters(), self.critic.parameters()), self.max_grad_norm)
        self.optimizer.step()

        mean_value_loss += value_loss.item()
        mean_surrogate_loss += surrogate_loss.item()
        mean_entropy += entropy.mean().item()
        mean_priv_reg_loss += priv_reg_loss.item()
        mean_estimator_loss += estimator_loss.item()
        mean_symmetry_loss += symmetry_loss.item()

    self.critic.update_normalization(self.storage.observations.flatten(0, 1))
    self.counter += 1
    num_updates = self.num_learning_epochs * self.num_mini_batches
    loss_dict = {
        "value": mean_value_loss / num_updates,
        "surrogate": mean_surrogate_loss / num_updates,
        "entropy": mean_entropy / num_updates,
        "priv_reg": mean_priv_reg_loss / num_updates,
        "priv_reg_coef": priv_reg_coef,
        "estimator": mean_estimator_loss / num_updates,
    }
    if self.mirror is not None:
        loss_dict["symmetry"] = mean_symmetry_loss / num_updates
    if self.is_dagger_iteration:
        loss_dict["hist_latent"] = self.update_dagger()
    self.storage.clear()
    self.iteration += 1
    return loss_dict

update_dagger

update_dagger() -> float

Regress the history (adaptation) encoder onto the privileged latent over the stored rollout.

Source code in eden/learning/wrappers/rsl_rl/rma/ppo.py

def update_dagger(self) -> float:
    """Regress the history (adaptation) encoder onto the privileged latent over the stored rollout."""
    actor = self.actor
    mean_loss = 0.0
    for batch in self.storage.mini_batch_generator(self.num_mini_batches, self.num_learning_epochs):
        flat = actor.get_flat_obs(batch.observations)
        with torch.no_grad():
            priv_latent = actor.net.infer_priv_latent(flat)
        hist_latent_loss = (priv_latent - actor.net.infer_hist_latent(flat)).norm(p=2, dim=1).mean()
        self.hist_encoder_optimizer.zero_grad()
        hist_latent_loss.backward()
        nn.utils.clip_grad_norm_(actor.net.history_encoder.parameters(), self.max_grad_norm)
        self.hist_encoder_optimizer.step()
        mean_loss += hist_latent_loss.item()
    self.counter += 1
    return mean_loss / (self.num_learning_epochs * self.num_mini_batches)

save

save() -> dict

Source code in eden/learning/wrappers/rsl_rl/rma/ppo.py

def save(self) -> dict:
    saved = super().save()
    saved["estimator_optimizer_state_dict"] = self.estimator_optimizer.state_dict()
    saved["hist_encoder_optimizer_state_dict"] = self.hist_encoder_optimizer.state_dict()
    saved["rma_counter"] = self.counter
    return saved

load

load(loaded_dict: dict, load_cfg: dict | None, strict: bool) -> bool

PPO's load — with the actor checked key for key whatever strict says — plus the RMA state.

Eden loads with strict=False to tolerate drift, but a partially loaded RMA actor (another model's checkpoint, e.g. a plain MLP actor's mlp.*) would run with random encoders or estimator.

Source code in eden/learning/wrappers/rsl_rl/rma/ppo.py

def load(self, loaded_dict: dict, load_cfg: dict | None, strict: bool) -> bool:
    """PPO's load — with the actor checked key for key whatever ``strict`` says — plus the RMA state.

    Eden loads with ``strict=False`` to tolerate drift, but a partially loaded RMA actor (another model's
    checkpoint, e.g. a plain MLP actor's ``mlp.*``) would run with random encoders or estimator.
    """
    if (load_cfg is None or load_cfg.get("actor")) and "actor_state_dict" in loaded_dict:
        expected, found = set(self._raw_actor.state_dict()), set(loaded_dict["actor_state_dict"])
        if expected != found:
            raise ValueError(
                "This checkpoint's actor_state_dict is not an RmaActorModel's: missing "
                f"{sorted(expected - found)[:6]}{' ...' if len(expected - found) > 6 else ''}, unexpected "
                f"{sorted(found - expected)[:6]}{' ...' if len(found - expected) > 6 else ''}. It was saved by "
                "another actor model (or an RMA actor with other layers); load it with the task that trained it."
            )
    load_iteration = super().load(loaded_dict, load_cfg, strict)
    if load_cfg is None or load_cfg.get("optimizer"):
        for key, opt in (
            ("estimator_optimizer_state_dict", self.estimator_optimizer),
            ("hist_encoder_optimizer_state_dict", self.hist_encoder_optimizer),
        ):
            if key in loaded_dict:  # absent from a checkpoint another algorithm saved: keep the fresh state
                opt.load_state_dict(loaded_dict[key])
    if load_iteration:
        # The runner resumes at the saved iteration number, so the DAgger cadence and the ramp continue.
        self.iteration = loaded_dict["iter"]
        self.counter = loaded_dict.get("rma_counter", self.counter)
    return load_iteration

DAgger depth distillation of an RMA teacher into a depth-camera student (rsl-rl-lib >= 5.5.1).

Port of the vision distillation loop from "Extreme Parkour with Legged Robots" (Cheng et al. 2023). The frozen teacher (:class:~eden.learning.wrappers.rsl_rl.rma.model.RmaActorModel, history-latent path, true scandots, ground truth privileged-explicit chunk) labels every step; the student — depth_actor (a copy of the teacher's actor net) fed by a recurrent depth_encoder whose latent replaces the scan latent — acts and drives the environment from its depth latent and its own prediction of the delta-yaw slots; one optimizer step per iteration regresses student -> teacher actions plus the yaw prediction, backpropagating through the whole iteration's graph.

That loop fits neither rsl_rl's Distillation (which stores observations and trains on mini-batches) nor its runner loop (which collects under torch.inference_mode), so this module ships its own runner, :class:RmaDepthDistillationRunner — what :class:~eden.learning.wrappers.rsl_rl.rma.options.RslRlRmaDepthDistillationRunnerOptions's class_name names, so eden train --task extreme_parkour_student --checkpoint <teacher model_<it>.pt> runs it.

Cadence contract: the depth encoder runs only on control steps where the env refreshed the depth observation (common_step_counter % update_interval == 0); between refreshes the cached latent/yaw are reused. Every observation term of the student's groups that declares an update_interval must run at the same cadence (checked against the live env at construction), and num_steps_per_env must be a positive multiple of it so an iteration boundary never splits a cadence window (the update would otherwise backprop into the previous iteration's freed graph).

Checkpoints carry the teacher (actor_state_dict, estimator included — the same key a teacher checkpoint has) plus depth_actor_state_dict / depth_encoder_state_dict / depth_optimizer_state_dict. load() accepts either: a teacher checkpoint starts a fresh distillation (student actor re-copied from the teacher), a student checkpoint resumes one. Only a student checkpoint (or a student trained in this process) replays: :meth:RmaDepthDistillationRunner.get_inference_policy refuses an untrained depth encoder.

Like the source, the student reads the privileged-explicit chunk (base linear velocity) as the env provides it — in sim, the ground truth — in distillation and in replay alike. It is an in-sim policy: eden deploy refuses this runner, and a hardware port would have to substitute the frozen teacher's estimator for that chunk (in distillation too, so the student trains on what it would see).

Attributes

RMA_DEPTH_DISTILLATION_RUNNER_CLASS_NAME module-attribute

RMA_DEPTH_DISTILLATION_RUNNER_CLASS_NAME = 'eden.learning.wrappers.rsl_rl.rma.distillation:RmaDepthDistillationRunner'

RMA_DEPTH_DISTILLATION_CLASS_NAME module-attribute

RMA_DEPTH_DISTILLATION_CLASS_NAME = 'eden.learning.wrappers.rsl_rl.rma.distillation:RmaDepthDistillation'

Classes

RmaDepthDistillation

RmaDepthDistillation(teacher: RmaActorModel, *, depth_obs_group: str = _STUDENT_GROUP_KEYS['depth_obs_group'], yaw_gate_obs_group: str = _STUDENT_GROUP_KEYS['yaw_gate_obs_group'], yaw_slots: tuple[int, int] | list[int] = (6, 8), update_interval: int = 5, num_pretrain_iters: int = 0, gru_hidden_dim: int = 512, learning_rate: float = 0.001, max_grad_norm: float = 1.0, yaw_scale: float = 1.5, device: str = 'cpu', multi_gpu_cfg: dict | None = None)

The student stage's modules and optimizer; :meth:update applies one DAgger regression.

Holds the frozen teacher (an :class:RmaActorModel, kept whole — estimator included — so a student checkpoint stays self-contained), the trainable depth_encoder and the trainable depth_actor (initialized as a copy of the teacher's actor net). The rollout loop lives in :class:RmaDepthDistillationRunner. The keyword arguments are documented on :class:~eden.learning.wrappers.rsl_rl.rma.options.RslRlRmaDepthDistillationAlgorithmOptions.

Source code in eden/learning/wrappers/rsl_rl/rma/distillation.py

def __init__(
    self,
    teacher: RmaActorModel,
    *,
    depth_obs_group: str = _STUDENT_GROUP_KEYS["depth_obs_group"],
    yaw_gate_obs_group: str = _STUDENT_GROUP_KEYS["yaw_gate_obs_group"],
    yaw_slots: tuple[int, int] | list[int] = (6, 8),
    update_interval: int = 5,
    num_pretrain_iters: int = 0,
    gru_hidden_dim: int = 512,
    learning_rate: float = 1.0e-3,
    max_grad_norm: float = 1.0,
    yaw_scale: float = 1.5,
    device: str = "cpu",
    multi_gpu_cfg: dict | None = None,
) -> None:
    if multi_gpu_cfg is not None:
        raise ValueError("RmaDepthDistillation is single-GPU; launch without torchrun.")
    start, stop = yaw_slots
    if not (0 <= start < stop <= teacher.dims.num_prop and stop - start == 2):
        raise ValueError(
            f"yaw_slots={tuple(yaw_slots)} must name the 2 (delta_yaw, delta_next_yaw) slots inside the "
            f"{teacher.dims.num_prop}-d proprioception chunk; the yaw head predicts exactly 2 values."
        )
    self.device = device
    self.dims = teacher.dims
    self.depth_obs_group = depth_obs_group
    self.yaw_gate_obs_group = yaw_gate_obs_group
    self.yaw_slice = slice(start, stop)
    self.update_interval = update_interval
    self.num_pretrain_iters = num_pretrain_iters
    self.yaw_scale = yaw_scale
    self.max_grad_norm = max_grad_norm
    self.teacher = teacher.to(device).requires_grad_(False)
    self.teacher_loaded = False
    # Whether the depth encoder / actor hold a trained student (loaded, or updated in this process).
    self.student_ready = False
    # The depth latent substitutes for the teacher's scan latent, so its width is the scan latent's width.
    self.depth_encoder = RecurrentDepthEncoder(
        num_prop=self.dims.num_prop, latent_dim=teacher.net.scan_latent_dim, gru_hidden_dim=gru_hidden_dim
    ).to(device)
    self.depth_actor = copy.deepcopy(self.teacher.net).requires_grad_(True).to(device)
    self.optimizer = torch.optim.Adam(
        [*self.depth_actor.parameters(), *self.depth_encoder.parameters()], lr=learning_rate
    )

Attributes

device instance-attribute

device = device

dims instance-attribute

dims = teacher.dims

depth_obs_group instance-attribute

depth_obs_group = depth_obs_group

yaw_gate_obs_group instance-attribute

yaw_gate_obs_group = yaw_gate_obs_group

yaw_slice instance-attribute

yaw_slice = slice(start, stop)

update_interval instance-attribute

update_interval = update_interval

num_pretrain_iters instance-attribute

num_pretrain_iters = num_pretrain_iters

yaw_scale instance-attribute

yaw_scale = yaw_scale

max_grad_norm instance-attribute

max_grad_norm = max_grad_norm

teacher instance-attribute

teacher = teacher.to(device).requires_grad_(False)

teacher_loaded instance-attribute

teacher_loaded = False

student_ready instance-attribute

student_ready = False

depth_encoder instance-attribute

depth_encoder = RecurrentDepthEncoder(num_prop=self.dims.num_prop, latent_dim=teacher.net.scan_latent_dim, gru_hidden_dim=gru_hidden_dim).to(device)

depth_actor instance-attribute

depth_actor = copy.deepcopy(self.teacher.net).requires_grad_(True).to(device)

optimizer instance-attribute

optimizer = torch.optim.Adam([*self.depth_actor.parameters(), *self.depth_encoder.parameters()], lr=learning_rate)

Methods:

construct_algorithm staticmethod

construct_algorithm(obs: TensorDict, env, cfg: dict, device: str) -> RmaDepthDistillation

Build the teacher from cfg["actor"] (it must match the teacher checkpoint) and the student around it.

Source code in eden/learning/wrappers/rsl_rl/rma/distillation.py

@staticmethod
def construct_algorithm(obs: TensorDict, env, cfg: dict, device: str) -> RmaDepthDistillation:
    """Build the teacher from ``cfg["actor"]`` (it must match the teacher checkpoint) and the student around it."""
    alg_class, alg_cfg = resolve_class(cfg["algorithm"])
    actor_class, actor_cfg = resolve_class(cfg["actor"])
    cfg["obs_groups"] = resolve_obs_groups(obs, cfg["obs_groups"], ["actor"])
    for key, default in _STUDENT_GROUP_KEYS.items():
        group = alg_cfg.get(key, default)
        if group not in obs.keys():
            raise ValueError(
                f"Depth distillation reads the observation group algorithm.{key}={group!r}, which the env does "
                f"not provide (groups: {sorted(obs.keys())}). Is the env the depth-student task?"
            )
    teacher = actor_class(obs, cfg["obs_groups"], "actor", env.num_actions, **actor_cfg)
    if not isinstance(teacher, RmaActorModel):
        raise TypeError(
            "Depth distillation needs an RmaActorModel teacher (actor=RslRlRmaActorOptions(...)), got "
            f"{type(teacher).__name__}."
        )
    return alg_class(teacher, device=device, **alg_cfg, multi_gpu_cfg=cfg["multi_gpu"])

encode

encode(depth_frame: Tensor, flat: Tensor) -> tuple[Tensor, Tensor]

One depth-encoder pass: (depth latent, scaled 2-d yaw prediction).

The proprioception chunk is fed with the yaw slots hidden (the student must recover heading from depth); the Tanh-bounded yaw head is stretched by yaw_scale. Shared by the training loop and the inference policy so the encoding contract cannot drift between them.

Source code in eden/learning/wrappers/rsl_rl/rma/distillation.py

def encode(self, depth_frame: torch.Tensor, flat: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """One depth-encoder pass: ``(depth latent, scaled 2-d yaw prediction)``.

    The proprioception chunk is fed with the yaw slots hidden (the student must recover heading from depth); the
    Tanh-bounded yaw head is stretched by ``yaw_scale``. Shared by the training loop and the inference policy so
    the encoding contract cannot drift between them.
    """
    prop = flat[:, self.dims.prop_slice].clone()
    prop[:, self.yaw_slice] = 0.0
    latent_and_yaw = self.depth_encoder(depth_frame, prop)
    return latent_and_yaw[:, :-2], self.yaw_scale * latent_and_yaw[:, -2:]

teacher_actions

teacher_actions(obs: TensorDict) -> Tensor

The teacher's label: history latent, true scandots, ground-truth privileged-explicit chunk (no estimator).

Source code in eden/learning/wrappers/rsl_rl/rma/distillation.py

def teacher_actions(self, obs: TensorDict) -> torch.Tensor:
    """The teacher's label: history latent, true scandots, ground-truth privileged-explicit chunk (no estimator)."""
    with torch.no_grad():
        return self.teacher(obs, hist_encoding=True, estimate=False)

student_actions

student_actions(flat: Tensor, depth_latent: Tensor) -> Tensor

The student's action from the (yaw-substituted) flat observation and the depth latent.

Source code in eden/learning/wrappers/rsl_rl/rma/distillation.py

def student_actions(self, flat: torch.Tensor, depth_latent: torch.Tensor) -> torch.Tensor:
    """The student's action from the (yaw-substituted) flat observation and the depth latent."""
    return self.depth_actor(flat, hist_encoding=True, scandots_latent=depth_latent)

update

update(action_diffs: Tensor, yaw_errors: Tensor) -> dict[str, float]

One optimizer step on the iteration's teacher - student actions and true - predicted yaw errors.

Both losses are mean per-sample L2 norms; the gradient clip covers the actor only (faithful to the source).

Source code in eden/learning/wrappers/rsl_rl/rma/distillation.py

def update(self, action_diffs: torch.Tensor, yaw_errors: torch.Tensor) -> dict[str, float]:
    """One optimizer step on the iteration's ``teacher - student`` actions and ``true - predicted`` yaw errors.

    Both losses are mean per-sample L2 norms; the gradient clip covers the actor only (faithful to the source).
    """
    action_loss = action_diffs.norm(p=2, dim=1).mean()
    yaw_loss = yaw_errors.norm(p=2, dim=1).mean()
    loss = action_loss + yaw_loss
    self.optimizer.zero_grad()
    loss.backward()
    nn.utils.clip_grad_norm_(self.depth_actor.parameters(), self.max_grad_norm)
    self.optimizer.step()
    self.student_ready = True
    return {"depth_actor": action_loss.item(), "yaw": yaw_loss.item(), "total": loss.item()}

train_mode

train_mode() -> None

Source code in eden/learning/wrappers/rsl_rl/rma/distillation.py

def train_mode(self) -> None:
    self.teacher.eval()  # frozen
    self.depth_actor.train()
    self.depth_encoder.train()

eval_mode

eval_mode() -> None

Source code in eden/learning/wrappers/rsl_rl/rma/distillation.py

def eval_mode(self) -> None:
    self.teacher.eval()
    self.depth_actor.eval()
    self.depth_encoder.eval()

save

save() -> dict

Source code in eden/learning/wrappers/rsl_rl/rma/distillation.py

def save(self) -> dict:
    return {
        "actor_state_dict": self.teacher.state_dict(),
        "depth_actor_state_dict": self.depth_actor.state_dict(),
        "depth_encoder_state_dict": self.depth_encoder.state_dict(),
        "depth_optimizer_state_dict": self.optimizer.state_dict(),
    }

load

load(loaded_dict: dict, *, load_optimizer: bool = True) -> bool

Load the teacher, and the student halves when present; True when it resumed a student.

Source code in eden/learning/wrappers/rsl_rl/rma/distillation.py

def load(self, loaded_dict: dict, *, load_optimizer: bool = True) -> bool:
    """Load the teacher, and the student halves when present; ``True`` when it resumed a student."""
    self.teacher.load_state_dict(loaded_dict["actor_state_dict"])
    self.teacher_loaded = True
    if "depth_actor_state_dict" not in loaded_dict:
        self.depth_actor.load_state_dict(self.teacher.net.state_dict())  # a fresh distillation from the teacher
        return False
    self.depth_actor.load_state_dict(loaded_dict["depth_actor_state_dict"])
    self.depth_encoder.load_state_dict(loaded_dict["depth_encoder_state_dict"])
    if load_optimizer:
        self.optimizer.load_state_dict(loaded_dict["depth_optimizer_state_dict"])
    self.student_ready = True
    return True

RmaDepthDistillationRunner

RmaDepthDistillationRunner(env, train_cfg: dict, log_dir: str | None = None, device: str = 'cpu')

Bases: ProgressFileMixin, AgentIterationLogMixin, CheckpointRetentionMixin, OnPolicyRunner

Runs :class:RmaDepthDistillation's DAgger loop; see the module docstring for the loop and its contracts.

Source code in eden/learning/wrappers/rsl_rl/rma/distillation.py

def __init__(self, env, train_cfg: dict, log_dir: str | None = None, device: str = "cpu") -> None:
    if eden_dist.is_distributed():
        raise RuntimeError(
            "Depth distillation is single-GPU: it has no parameter broadcast or gradient reduction, so a "
            "distributed launch would train an independent student per rank. Launch without torchrun."
        )
    super().__init__(env, train_cfg, log_dir=log_dir, device=device)
    interval, steps = self.alg.update_interval, self.cfg["num_steps_per_env"]
    if steps <= 0 or steps % interval != 0:
        raise ValueError(
            f"num_steps_per_env ({steps}) must be a positive multiple of algorithm.update_interval ({interval}), "
            "so an iteration boundary never splits a depth-refresh window."
        )
    self._validate_env_cadence()

Attributes

alg instance-attribute

alg: RmaDepthDistillation

checkpoint_required class-attribute instance-attribute

checkpoint_required = "Depth distillation needs a teacher: pass --checkpoint <the teacher run's model_<it>.pt> (or --resume / --checkpoint <a student checkpoint> to continue a distillation)."

Methods:

learn

learn(num_learning_iterations: int, init_at_random_ep_len: bool = False) -> None

The vision DAgger loop; init_at_random_ep_len spreads the envs' episode clocks, as rsl_rl's loop does.

Source code in eden/learning/wrappers/rsl_rl/rma/distillation.py

def learn(self, num_learning_iterations: int, init_at_random_ep_len: bool = False) -> None:
    """The vision DAgger loop; ``init_at_random_ep_len`` spreads the envs' episode clocks, as rsl_rl's loop does."""
    if not self.alg.teacher_loaded:
        raise ValueError(self.checkpoint_required)
    alg, env_unwrapped = self.alg, self.env.unwrapped
    interval = alg.update_interval
    check_for_nan = self.cfg.get("check_for_nan", True)
    if init_at_random_ep_len:
        self.env.episode_length_buf = torch.randint_like(
            self.env.episode_length_buf, high=int(self.env.max_episode_length)
        )
    obs = self.env.get_observations().to(self.device)
    alg.train_mode()
    self.logger.init_logging_writer()
    depth_latent: torch.Tensor | None = None
    yaw_pred: torch.Tensor | None = None

    start_it = self.current_learning_iteration
    total_it = start_it + num_learning_iterations
    for it in range(start_it, total_it):
        start = time.time()
        action_diffs: list[torch.Tensor] = []
        yaw_errors: list[torch.Tensor] = []
        gate_fracs: list[torch.Tensor] = []  # 0-d tensors; one host sync per iteration
        # NO torch.inference_mode (unlike PPO): the student's actions carry the graph the update backpropagates.
        for _ in range(self.cfg["num_steps_per_env"]):
            # Clone: the yaw substitution must not leak into the env-owned observation buffer.
            flat = alg.teacher.get_flat_obs(obs).clone()
            gate = obs[alg.yaw_gate_obs_group].squeeze(-1) > 0.5
            if depth_latent is None or env_unwrapped.common_step_counter % interval == 0:
                depth_latent, yaw_pred = alg.encode(obs[alg.depth_obs_group], flat)
                yaw_errors.append(flat[:, alg.yaw_slice] - yaw_pred)
            # The label comes BEFORE the yaw substitution: the teacher sees the true delta-yaw.
            teacher_actions = alg.teacher_actions(obs)
            gate_fracs.append(gate.float().mean())
            # The student sees its own (detached) yaw prediction wherever the gate allows; the yaw head trains
            # only through the auxiliary loss, as in the source.
            flat[gate, alg.yaw_slice] = yaw_pred.detach()[gate]
            student_actions = alg.student_actions(flat, depth_latent)
            action_diffs.append(teacher_actions - student_actions)

            stepped = teacher_actions if it < alg.num_pretrain_iters else student_actions
            obs, rewards, dones, extras = self.env.step(stepped.detach().to(self.env.device))
            if check_for_nan:
                check_nan(obs, rewards, dones)
            obs, rewards, dones = obs.to(self.device), rewards.to(self.device), dones.to(self.device)
            self.logger.process_env_step(rewards, dones, extras, None)

        stop = time.time()
        collect_time, start = stop - start, stop

        loss_dict = alg.update(torch.cat(action_diffs), torch.cat(yaw_errors))
        loss_dict["delta_yaw_ok_frac"] = torch.stack(gate_fracs).mean().item()
        alg.depth_encoder.detach_hidden()
        # Force a fresh encoder pass at the next iteration's first step: reusing the latent across the update
        # would backprop into the freed graph.
        depth_latent = yaw_pred = None

        learn_time = time.time() - start
        self.current_learning_iteration = it
        self.logger.log(
            it=it,
            start_it=start_it,
            total_it=total_it,
            collect_time=collect_time,
            learn_time=learn_time,
            loss_dict=loss_dict,
            learning_rate=alg.optimizer.param_groups[0]["lr"],
            # Distillation is deterministic — no exploration noise to report.
            action_std=torch.zeros(self.env.num_actions, device=self.device),
            rnd_weight=None,
        )
        if self.logger.writer is not None and self._save_due(it, start_it):
            self.save(os.path.join(self.logger.log_dir, f"model_{it}.pt"))

    if self.logger.writer is not None:
        self.save(os.path.join(self.logger.log_dir, f"model_{self.current_learning_iteration}.pt"))
        self.logger.stop_logging_writer()

load

load(path: str, load_optimizer: bool = True, map_location: str | None = None) -> dict | None

Load a teacher checkpoint (fresh distillation) or a student checkpoint (resume).

Source code in eden/learning/wrappers/rsl_rl/rma/distillation.py

def load(self, path: str, load_optimizer: bool = True, map_location: str | None = None) -> dict | None:
    """Load a teacher checkpoint (fresh distillation) or a student checkpoint (resume)."""
    loaded_dict = torch.load(path, weights_only=False, map_location=map_location)
    if "actor_state_dict" not in loaded_dict:
        raise ValueError(
            f"'{path}' has no 'actor_state_dict' (keys: {sorted(loaded_dict)}), so it is not an RMA teacher or "
            "student checkpoint of the rsl-rl-lib 5 layout. An rsl_rl 3 checkpoint ('model_state_dict') does not "
            "load here: retrain the teacher with `eden train --task extreme_parkour_teacher`."
        )
    if self.alg.load(loaded_dict, load_optimizer=load_optimizer):
        self.current_learning_iteration = loaded_dict["iter"]
    return loaded_dict.get("infos")

get_inference_policy

get_inference_policy(device: str | None = None) -> Callable[[TensorDict], Tensor]

The deployed student: depth encoder every update_interval steps, its yaw in the yaw slots.

The yaw prediction replaces the yaw slots unconditionally (the source's play behavior — the gate is a training-time aid), and the privileged-explicit chunk stays as the env provides it (in-sim ground truth; a hardware deployment would substitute the teacher's estimator here).

Source code in eden/learning/wrappers/rsl_rl/rma/distillation.py

def get_inference_policy(self, device: str | None = None) -> Callable[[TensorDict], torch.Tensor]:
    """The deployed student: depth encoder every ``update_interval`` steps, its yaw in the yaw slots.

    The yaw prediction replaces the yaw slots unconditionally (the source's play behavior — the gate is a
    training-time aid), and the privileged-explicit chunk stays as the env provides it (in-sim ground truth; a
    hardware deployment would substitute the teacher's estimator here).
    """
    alg, env_unwrapped = self.alg, self.env.unwrapped
    if not alg.student_ready:
        raise RuntimeError(
            "The depth student has not been trained: this runner loaded a teacher checkpoint (no "
            "depth_actor_state_dict / depth_encoder_state_dict), so its depth encoder is randomly initialized and "
            "would drive the robot from a random latent and a random yaw. Replay a student checkpoint (a model_<it>.pt "
            "written by `eden train --task extreme_parkour_student`), or replay the teacher itself with "
            "--task extreme_parkour_teacher."
        )
    alg.eval_mode()
    if device is not None:
        alg.depth_actor.to(device)
        alg.depth_encoder.to(device)
    alg.depth_encoder.reset_hidden()
    state: dict[str, torch.Tensor | None] = {"latent": None, "yaw": None}

    def act(obs: TensorDict) -> torch.Tensor:
        with torch.no_grad():
            flat = alg.teacher.get_flat_obs(obs).clone()
            if state["latent"] is None or env_unwrapped.common_step_counter % alg.update_interval == 0:
                state["latent"], state["yaw"] = alg.encode(obs[alg.depth_obs_group], flat)
            flat[:, alg.yaw_slice] = state["yaw"]
            return alg.student_actions(flat, state["latent"])

    return act

HIMLoco

The :class:~eden.learning.wrappers.rsl_rl.him.ppo.HIMPPO algorithm options.

Classes

HIMPPOAlgorithmOptions

HIMPPOAlgorithmOptions(**data)

Bases: RslRlPpoAlgorithmOptions

PPO options for OpenHomie's HIMLoco learner.

Built by rsl-rl-lib's own OnPolicyRunner on any major Eden supports (3, 4, 5), so eden train and eden inference need nothing algorithm-specific; :meth:adapt_runner_dict routes the construction for the installed major. The task supplies the two things that tie HIMPPO to an observation layout:

• the observation groups: obs_groups["actor"] one group of stacked frames, obs_groups["critic"] one group of a single frame + the true base linear velocity (3, last);
• the mirror, data_augmentation_func.

PARAMETER	DESCRIPTION
data_augmentation_func	"module:function" mirroring the task's observations and actions, with rsl-rl-lib's symmetry signature f(env=, obs=TensorDict | None, actions=Tensor | None) -> (obs_aug, actions_aug), each the input stacked with its mirror along the batch axis — the same function stock PPO's symmetry extension takes. Required: every transition is stored with its mirror.
use_flip	Also train the estimator on mirrored observations and add the actor and critic mirror losses (OpenHomie use_flip). Symmetry lives here, so symmetry_cfg must stay None.
symmetry_scale	Weight of the two mirror losses.
target_skip_dim	Leading entries of the next critic observation the SwAV target encoder leaves out (OpenHomie drops the three velocity commands a frame starts with).

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

class_name class-attribute instance-attribute

class_name: str = 'eden.learning.wrappers.rsl_rl.him.ppo:HIMPPO'

entropy_coef class-attribute instance-attribute

entropy_coef: float = 0.01

data_augmentation_func instance-attribute

data_augmentation_func: str

use_flip class-attribute instance-attribute

use_flip: bool = True

symmetry_scale class-attribute instance-attribute

symmetry_scale: float = 1.0

target_skip_dim class-attribute instance-attribute

target_skip_dim: int = 3

Methods:

adapt_runner_dict

adapt_runner_dict(runner_dict: dict, major: int) -> dict

Point the version-translated runner dict at HIM's constructors for rsl-rl-lib major.

4 and 5 build the algorithm through class_name → construct_algorithm, which reads the translated dict as it is. 3 builds the model from policy.class_name and calls the algorithm class with it, so both go to HIM's v3 factories, and the model gets target_skip_dim.

PARAMETER	DESCRIPTION
runner_dict	The dict :func:~eden.options.learning.rsl_rl.translate.translate_runner_dict produced for major. TYPE: dict
major	The installed rsl-rl-lib major version. TYPE: int

RETURNS	DESCRIPTION
dict	runner_dict, adapted in place.

Source code in eden/learning/wrappers/rsl_rl/him/options.py

def adapt_runner_dict(self, runner_dict: dict, major: int) -> dict:
    """Point the version-translated runner dict at HIM's constructors for rsl-rl-lib ``major``.

    4 and 5 build the algorithm through ``class_name`` → ``construct_algorithm``, which reads the
    translated dict as it is. 3 builds the model from ``policy.class_name`` and calls the algorithm
    class with it, so both go to HIM's v3 factories, and the model gets ``target_skip_dim``.

    Parameters
    ----------
    runner_dict: dict
        The dict :func:`~eden.options.learning.rsl_rl.translate.translate_runner_dict` produced for ``major``.
    major: int
        The installed rsl-rl-lib major version.

    Returns
    -------
    dict
        ``runner_dict``, adapted in place.
    """
    if major == 3:
        module = "eden.learning.wrappers.rsl_rl.him"
        runner_dict["policy"]["class_name"] = f"{module}.actor_critic:HIMActorCritic.from_rsl_rl_v3"
        runner_dict["policy"]["target_skip_dim"] = self.target_skip_dim
        runner_dict["algorithm"]["class_name"] = f"{module}.ppo:HIMPPO.from_rsl_rl_v3"
    return runner_dict

HIMPPO

HIMPPO(actor_critic: HIMActorCritic, use_flip: bool = True, num_learning_epochs: int = 5, num_mini_batches: int = 4, clip_param: float = 0.2, gamma: float = 0.99, lam: float = 0.95, value_loss_coef: float = 1.0, entropy_coef: float = 0.01, learning_rate: float = 0.001, max_grad_norm: float = 1.0, use_clipped_value_loss: bool = True, schedule: str = 'adaptive', desired_kl: float = 0.01, device: str = 'cpu', symmetry_scale: float = 1.0, env=None, data_augmentation_func: Callable | None = None)

Source code in eden/learning/wrappers/rsl_rl/him/ppo.py

def __init__(
    self,
    actor_critic: HIMActorCritic,
    use_flip: bool = True,
    num_learning_epochs: int = 5,
    num_mini_batches: int = 4,
    clip_param: float = 0.2,
    gamma: float = 0.99,
    lam: float = 0.95,
    value_loss_coef: float = 1.0,
    entropy_coef: float = 0.01,
    learning_rate: float = 1e-3,
    max_grad_norm: float = 1.0,
    use_clipped_value_loss: bool = True,
    schedule: str = "adaptive",
    desired_kl: float = 0.01,
    device: str = "cpu",
    symmetry_scale: float = 1.0,
    env=None,
    data_augmentation_func: Callable | None = None,
):
    if data_augmentation_func is None:
        raise ValueError(
            "HIMPPO stores every transition's mirror, so it needs the task's mirror: set the algorithm options' "
            "data_augmentation_func to a 'module:function' with rsl-rl's symmetry signature."
        )
    self.env = env
    self.data_augmentation_func = data_augmentation_func
    self.device = device
    self.use_flip = use_flip
    self.desired_kl = desired_kl
    self.schedule = schedule
    self.learning_rate = learning_rate

    self.actor_critic = actor_critic
    self.actor_critic.to(self.device)
    self.storage: HIMRolloutStorage | None = None
    self.optimizer = HIMOptimizer(
        optim.Adam(self.actor_critic.parameters(), lr=learning_rate), self.actor_critic.estimator.optimizer
    )
    self.transition = HIMRolloutStorage.Transition()
    self.transition_sym = HIMRolloutStorage.Transition()
    self.symmetry_scale = symmetry_scale

    self.clip_param = clip_param
    self.num_learning_epochs = num_learning_epochs
    self.num_mini_batches = num_mini_batches
    self.value_loss_coef = value_loss_coef
    self.entropy_coef = entropy_coef
    self.gamma = gamma
    self.lam = lam
    self.max_grad_norm = max_grad_norm
    self.use_clipped_value_loss = use_clipped_value_loss

    # Read by rsl-rl-lib's runners: RND is logged when an algorithm has it, and v3 reaches the
    # model as ``alg.policy`` (v4/v5 through ``get_policy()``).
    self.rnd = None
    self.intrinsic_rewards = None
    self.policy = actor_critic
    self.actor_obs_group = actor_critic.obs_group
    self.critic_obs_group = actor_critic.critic_obs_group

Attributes

env instance-attribute

env = env

data_augmentation_func instance-attribute

data_augmentation_func = data_augmentation_func

device instance-attribute

device = device

use_flip instance-attribute

use_flip = use_flip

desired_kl instance-attribute

desired_kl = desired_kl

schedule instance-attribute

schedule = schedule

learning_rate instance-attribute

learning_rate = learning_rate

actor_critic instance-attribute

actor_critic: HIMActorCritic = actor_critic

storage instance-attribute

storage: HIMRolloutStorage | None = None

optimizer instance-attribute

optimizer = HIMOptimizer(optim.Adam(self.actor_critic.parameters(), lr=learning_rate), self.actor_critic.estimator.optimizer)

transition instance-attribute

transition = HIMRolloutStorage.Transition()

transition_sym instance-attribute

transition_sym = HIMRolloutStorage.Transition()

symmetry_scale instance-attribute

symmetry_scale = symmetry_scale

clip_param instance-attribute

clip_param = clip_param

num_learning_epochs instance-attribute

num_learning_epochs = num_learning_epochs

num_mini_batches instance-attribute

num_mini_batches = num_mini_batches

value_loss_coef instance-attribute

value_loss_coef = value_loss_coef

entropy_coef instance-attribute

entropy_coef = entropy_coef

gamma instance-attribute

gamma = gamma

lam instance-attribute

lam = lam

max_grad_norm instance-attribute

max_grad_norm = max_grad_norm

use_clipped_value_loss instance-attribute

use_clipped_value_loss = use_clipped_value_loss

rnd instance-attribute

rnd = None

intrinsic_rewards instance-attribute

intrinsic_rewards = None

policy instance-attribute

policy = actor_critic

actor_obs_group instance-attribute

actor_obs_group = actor_critic.obs_group

critic_obs_group instance-attribute

critic_obs_group = actor_critic.critic_obs_group

Methods:

construct_algorithm classmethod

construct_algorithm(obs, env, cfg: dict, device: str) -> HIMPPO

rsl-rl-lib 4 / 5's construction hook: build the actor-critic, the storage and the algorithm.

The mirror is run once on these observations' widths, so a layout it does not cover fails here, before a rollout.

Source code in eden/learning/wrappers/rsl_rl/him/ppo.py

@classmethod
def construct_algorithm(cls, obs, env, cfg: dict, device: str) -> HIMPPO:
    """rsl-rl-lib 4 / 5's construction hook: build the actor-critic, the storage and the algorithm.

    The mirror is run once on these observations' widths, so a layout it does not cover fails here,
    before a rollout.
    """
    if cfg.get("multi_gpu"):
        raise NotImplementedError(_NOT_DISTRIBUTED)
    cfg["obs_groups"] = resolve_obs_groups(obs, cfg["obs_groups"], ["actor", "critic"])
    actor_cfg, critic_cfg = cfg["actor"], cfg["critic"]
    if actor_cfg.get("obs_normalization") or critic_cfg.get("obs_normalization"):
        raise ValueError(_NO_OBS_NORMALIZATION)
    init_std = check_model_cfg(actor_cfg, critic_cfg)
    # Without the pre-reset critic observations every rollout would fail on its first step; refuse here.
    env_options = getattr(getattr(env, "unwrapped", env), "env_options", None)
    if env_options is not None and not getattr(env_options, "record_final_observations", True):
        raise ValueError(_NEEDS_FINAL_OBSERVATIONS)
    actor_critic = HIMActorCritic.from_groups(
        obs,
        single_group(cfg["obs_groups"], "actor"),
        single_group(cfg["obs_groups"], "critic"),
        env.num_actions,
        actor_hidden_dims=actor_cfg["hidden_dims"],
        critic_hidden_dims=critic_cfg["hidden_dims"],
        activation=actor_cfg["activation"],
        init_noise_std=init_std,
        target_skip_dim=cfg["algorithm"]["target_skip_dim"],
    )
    return cls._build(actor_critic, cfg["algorithm"], device, env, env.num_envs, cfg["num_steps_per_env"])

from_rsl_rl_v3 classmethod

from_rsl_rl_v3(actor_critic: HIMActorCritic, storage, device: str, multi_gpu_cfg=None, **alg_cfg) -> HIMPPO

rsl-rl-lib 3's construction call, alg_class(actor_critic, storage, device=, **alg_cfg, multi_gpu_cfg=).

Reached through the algorithm.class_name that HIMPPOAlgorithmOptions.adapt_runner_dict writes; the runner has already built actor_critic (HIMActorCritic.from_rsl_rl_v3). Its stock rollout storage is only read for its size. v3 hands the algorithm no env, so the mirror gets env=None.

Source code in eden/learning/wrappers/rsl_rl/him/ppo.py

@classmethod
def from_rsl_rl_v3(
    cls, actor_critic: HIMActorCritic, storage, device: str, multi_gpu_cfg=None, **alg_cfg
) -> HIMPPO:
    """rsl-rl-lib 3's construction call, ``alg_class(actor_critic, storage, device=, **alg_cfg, multi_gpu_cfg=)``.

    Reached through the ``algorithm.class_name`` that ``HIMPPOAlgorithmOptions.adapt_runner_dict`` writes;
    the runner has already built ``actor_critic`` (``HIMActorCritic.from_rsl_rl_v3``). Its stock rollout
    storage is only read for its size. v3 hands the algorithm no env, so the mirror gets ``env=None``.
    """
    if multi_gpu_cfg:
        raise NotImplementedError(_NOT_DISTRIBUTED)
    return cls._build(actor_critic, alg_cfg, device, None, storage.num_envs, storage.num_transitions_per_env)

init_storage

init_storage(num_envs, num_transitions_per_env, actor_obs_shape, critic_obs_shape, action_shape)

Source code in eden/learning/wrappers/rsl_rl/him/ppo.py

def init_storage(self, num_envs, num_transitions_per_env, actor_obs_shape, critic_obs_shape, action_shape):
    self.storage = HIMRolloutStorage(
        num_envs, num_transitions_per_env, actor_obs_shape, critic_obs_shape, action_shape, self.device
    )

train_mode

train_mode()

Source code in eden/learning/wrappers/rsl_rl/him/ppo.py

def train_mode(self):
    self.actor_critic.train()

eval_mode

eval_mode()

Source code in eden/learning/wrappers/rsl_rl/him/ppo.py

def eval_mode(self):
    self.actor_critic.eval()

get_policy

get_policy()

Source code in eden/learning/wrappers/rsl_rl/him/ppo.py

def get_policy(self):
    return self.actor_critic

save

save() -> dict

rsl-rl-lib 4 / 5's checkpoint contents (v3's runner saves policy and optimizer itself).

Source code in eden/learning/wrappers/rsl_rl/him/ppo.py

def save(self) -> dict:
    """rsl-rl-lib 4 / 5's checkpoint contents (v3's runner saves ``policy`` and ``optimizer`` itself)."""
    return {
        "actor_critic_state_dict": self.actor_critic.state_dict(),
        "estimator_optimizer_state_dict": self.optimizer.estimator.state_dict(),
        "optimizer_state_dict": self.optimizer.ppo.state_dict(),
    }

load

load(loaded_dict: dict, load_cfg: dict | None = None, strict: bool = True) -> bool

Load what load_cfg selects (everything when None); return whether to restore the iteration.

The weights always load strictly, whatever strict says. Eden's checkpoint loader passes strict=False so stock PPO can take older rsl-rl layouts, but a HIM checkpoint has one layout: a missing or extra key means the wrong checkpoint, and loading past it would leave those weights at their random init. ValueError, not KeyError, so the loader does not reword it as an rsl-rl version mismatch.

Source code in eden/learning/wrappers/rsl_rl/him/ppo.py

def load(self, loaded_dict: dict, load_cfg: dict | None = None, strict: bool = True) -> bool:
    """Load what ``load_cfg`` selects (everything when ``None``); return whether to restore the iteration.

    The weights always load strictly, whatever ``strict`` says. Eden's checkpoint loader passes
    ``strict=False`` so stock PPO can take older rsl-rl layouts, but a HIM checkpoint has one layout: a
    missing or extra key means the wrong checkpoint, and loading past it would leave those weights at
    their random init. ``ValueError``, not ``KeyError``, so the loader does not reword it as an
    rsl-rl version mismatch.
    """
    if load_cfg is None:
        load_cfg = {"actor": True, "critic": True, "optimizer": True, "iteration": True}
    if "actor_critic_state_dict" not in loaded_dict:
        raise ValueError(
            f"Not a HIMPPO checkpoint from rsl-rl-lib 4/5: no 'actor_critic_state_dict' among "
            f"{sorted(loaded_dict)}. A stock-PPO or rsl-rl-lib 3 checkpoint cannot initialize HIMPPO; "
            f"pass a checkpoint this algorithm saved."
        )
    if load_cfg.get("actor") or load_cfg.get("critic"):
        state = loaded_dict["actor_critic_state_dict"]
        expected = set(self.actor_critic.state_dict())
        missing, unexpected = sorted(expected - set(state)), sorted(set(state) - expected)
        if missing or unexpected:
            raise ValueError(
                f"HIMPPO checkpoint does not match this actor-critic: missing {missing}, unexpected "
                f"{unexpected}. It was saved for a different observation layout or network; train or "
                f"pick a checkpoint from this task's config."
            )
        self.actor_critic.load_state_dict(state, strict=True)
    if load_cfg.get("optimizer"):
        self.optimizer.ppo.load_state_dict(loaded_dict["optimizer_state_dict"])
        self.optimizer.estimator.load_state_dict(loaded_dict["estimator_optimizer_state_dict"])
        self.learning_rate = self.optimizer.param_groups[0]["lr"]
    return bool(load_cfg.get("iteration", False))

act

act(obs) -> Tensor

Source code in eden/learning/wrappers/rsl_rl/him/ppo.py

def act(self, obs) -> torch.Tensor:
    return self._act(obs[self.actor_obs_group], obs[self.critic_obs_group])

process_env_step

process_env_step(obs, rewards, dones, extras) -> None

Record the step; obs is the observation after it, post-reset for the envs that ended.

The estimator's target is the true next critic observation, so the ended envs take their pre-reset one from extras["final_observations"] — OpenHomie's next_critic_obs[termination_ids] = termination_privileged_obs.

Source code in eden/learning/wrappers/rsl_rl/him/ppo.py

def process_env_step(self, obs, rewards, dones, extras) -> None:
    """Record the step; ``obs`` is the observation *after* it, post-reset for the envs that ended.

    The estimator's target is the true next critic observation, so the ended envs take their
    pre-reset one from ``extras["final_observations"]`` — OpenHomie's
    ``next_critic_obs[termination_ids] = termination_privileged_obs``.
    """
    final = extras.get("final_observations")
    if final is None:
        raise KeyError(_NEEDS_FINAL_OBSERVATIONS)
    next_critic_obs = torch.where(
        dones.bool().unsqueeze(-1), final[self.critic_obs_group].to(self.device), obs[self.critic_obs_group]
    )
    self._process_env_step(rewards, dones, extras, next_critic_obs)

compute_returns

compute_returns(obs) -> None

Source code in eden/learning/wrappers/rsl_rl/him/ppo.py

def compute_returns(self, obs) -> None:
    last_values = self.actor_critic.evaluate(obs[self.critic_obs_group]).detach()
    self.storage.compute_returns(last_values, self.gamma, self.lam)

update

update()

Source code in eden/learning/wrappers/rsl_rl/him/ppo.py

def update(self):
    mean_value_loss = 0.0
    mean_surrogate_loss = 0.0
    mean_estimation_loss = 0.0
    mean_swap_loss = 0.0
    mean_actor_sym_loss = 0.0
    mean_critic_sym_loss = 0.0

    generator = self.storage.mini_batch_generator(self.num_mini_batches, self.num_learning_epochs)
    for (
        obs_batch,
        critic_obs_batch,
        actions_batch,
        next_critic_obs_batch,
        target_values_batch,
        advantages_batch,
        returns_batch,
        old_actions_log_prob_batch,
        old_mu_batch,
        old_sigma_batch,
    ) in generator:
        self.actor_critic.act(obs_batch)
        actions_log_prob_batch = self.actor_critic.get_actions_log_prob(actions_batch)
        value_batch = self.actor_critic.evaluate(critic_obs_batch)
        mu_batch = self.actor_critic.action_mean
        sigma_batch = self.actor_critic.action_std
        entropy_batch = self.actor_critic.entropy

        # Adaptive-KL learning-rate schedule
        if self.desired_kl is not None and self.schedule == "adaptive":
            with torch.inference_mode():
                kl = torch.sum(
                    torch.log(sigma_batch / old_sigma_batch + 1.0e-5)
                    + (torch.square(old_sigma_batch) + torch.square(old_mu_batch - mu_batch))
                    / (2.0 * torch.square(sigma_batch))
                    - 0.5,
                    axis=-1,
                )
                kl_mean = torch.mean(kl)
                if kl_mean > self.desired_kl * 2.0:
                    self.learning_rate = max(1e-5, self.learning_rate / 1.5)
                elif kl_mean < self.desired_kl / 2.0 and kl_mean > 0.0:
                    self.learning_rate = min(1e-2, self.learning_rate * 1.5)
                for param_group in self.optimizer.param_groups:
                    param_group["lr"] = self.learning_rate

        # Estimator update (on real ⊕ mirrored obs)
        if self.use_flip:
            estimator_obs = torch.cat((obs_batch, self._mirror_obs(self.actor_obs_group, obs_batch)), dim=0)
            estimator_next_critic_obs = torch.cat(
                (next_critic_obs_batch, self._mirror_obs(self.critic_obs_group, next_critic_obs_batch)), dim=0
            )
        else:
            estimator_obs = obs_batch
            estimator_next_critic_obs = next_critic_obs_batch
        estimation_loss, swap_loss = self.actor_critic.update_estimator(
            estimator_obs, estimator_next_critic_obs, lr=self.learning_rate
        )

        # Surrogate loss
        ratio = torch.exp(actions_log_prob_batch - torch.squeeze(old_actions_log_prob_batch))
        surrogate = -torch.squeeze(advantages_batch) * ratio
        surrogate_clipped = -torch.squeeze(advantages_batch) * torch.clamp(
            ratio, 1.0 - self.clip_param, 1.0 + self.clip_param
        )
        surrogate_loss = torch.max(surrogate, surrogate_clipped).mean()

        # Value loss
        if self.use_clipped_value_loss:
            value_clipped = target_values_batch + (value_batch - target_values_batch).clamp(
                -self.clip_param, self.clip_param
            )
            value_losses = (value_batch - returns_batch).pow(2)
            value_losses_clipped = (value_clipped - returns_batch).pow(2)
            value_loss = torch.max(value_losses, value_losses_clipped).mean()
        else:
            value_loss = (returns_batch - value_batch).pow(2).mean()

        # Symmetry (mirror) losses
        if self.use_flip:
            flipped_critic_obs_batch = self._mirror_obs(self.critic_obs_group, critic_obs_batch)
            actor_sym_loss = self._actor_symmetry_loss(obs_batch)
            critic_sym_loss = self.symmetry_scale * torch.mean(
                torch.square(
                    self.actor_critic.evaluate(flipped_critic_obs_batch)
                    - self.actor_critic.evaluate(critic_obs_batch).detach()
                )
            )
            loss = (
                surrogate_loss
                + self.value_loss_coef * value_loss
                - self.entropy_coef * entropy_batch.mean()
                + actor_sym_loss
                + critic_sym_loss
            )
        else:
            loss = surrogate_loss + self.value_loss_coef * value_loss - self.entropy_coef * entropy_batch.mean()

        self.optimizer.zero_grad()
        loss.backward()
        nn.utils.clip_grad_norm_(self.actor_critic.parameters(), self.max_grad_norm)
        self.optimizer.step()

        mean_value_loss += value_loss.item()
        mean_surrogate_loss += surrogate_loss.item()
        mean_estimation_loss += estimation_loss
        mean_swap_loss += swap_loss
        if self.use_flip:
            mean_actor_sym_loss += actor_sym_loss.item()
            mean_critic_sym_loss += critic_sym_loss.item()

    num_updates = self.num_learning_epochs * self.num_mini_batches
    self.storage.clear()
    return {
        "value_loss": mean_value_loss / num_updates,
        "surrogate_loss": mean_surrogate_loss / num_updates,
        "estimation_loss": mean_estimation_loss / num_updates,
        "swap_loss": mean_swap_loss / num_updates,
        "actor_sym_loss": mean_actor_sym_loss / num_updates,
        "critic_sym_loss": mean_critic_sym_loss / num_updates,
    }

HIMActorCritic

HIMActorCritic(num_actor_obs: int, num_critic_obs: int, num_one_step_obs: int, actor_history_length: int, num_actions: int, actor_hidden_dims=(512, 256, 256), critic_hidden_dims=(512, 256, 256), activation: str = 'elu', init_noise_std: float = 1.0, dynamic_latent_dim: int = 32, obs_group: str = 'policy', target_skip_dim: int = 3, critic_obs_group: str = 'critic')

Bases: Module

Source code in eden/learning/wrappers/rsl_rl/him/actor_critic.py

def __init__(
    self,
    num_actor_obs: int,
    num_critic_obs: int,
    num_one_step_obs: int,
    actor_history_length: int,
    num_actions: int,
    actor_hidden_dims=(512, 256, 256),
    critic_hidden_dims=(512, 256, 256),
    activation: str = "elu",
    init_noise_std: float = 1.0,
    dynamic_latent_dim: int = 32,
    obs_group: str = "policy",
    target_skip_dim: int = 3,
    critic_obs_group: str = "critic",
):
    super().__init__()
    act = _get_activation(activation)

    self.num_actor_obs = num_actor_obs
    self.num_critic_obs = num_critic_obs
    self.num_one_step_obs = num_one_step_obs
    self.actor_history_length = actor_history_length
    self.actor_proprioceptive_obs_length = actor_history_length * num_one_step_obs
    self.num_actions = num_actions
    self.dynamic_latent_dim = dynamic_latent_dim
    self.obs_group = obs_group
    self.critic_obs_group = critic_obs_group

    # Estimator: history -> velocity(3) + latent(dynamic_latent_dim)
    self.estimator = HIMEstimator(
        temporal_steps=actor_history_length,
        num_one_step_obs=num_one_step_obs,
        num_height_points=0,
        latent_dim=dynamic_latent_dim,
        target_skip_dim=target_skip_dim,
    )

    # Actor: current_obs(num_one_step_obs) ⊕ vel(3) ⊕ latent -> action
    mlp_input_dim_a = num_one_step_obs + VELOCITY_DIM + dynamic_latent_dim
    actor_layers: list[nn.Module] = [nn.Linear(mlp_input_dim_a, actor_hidden_dims[0]), act]
    for i in range(len(actor_hidden_dims)):
        if i == len(actor_hidden_dims) - 1:
            actor_layers.append(nn.Linear(actor_hidden_dims[i], num_actions))
        else:
            actor_layers += [nn.Linear(actor_hidden_dims[i], actor_hidden_dims[i + 1]), act]
    self.actor = nn.Sequential(*actor_layers)

    # Critic: privileged obs -> value
    critic_layers: list[nn.Module] = [nn.Linear(num_critic_obs, critic_hidden_dims[0]), act]
    for i in range(len(critic_hidden_dims)):
        if i == len(critic_hidden_dims) - 1:
            critic_layers.append(nn.Linear(critic_hidden_dims[i], 1))
        else:
            critic_layers += [nn.Linear(critic_hidden_dims[i], critic_hidden_dims[i + 1]), act]
    self.critic = nn.Sequential(*critic_layers)

    # State-independent action noise
    self.std = nn.Parameter(init_noise_std * torch.ones(num_actions))
    self.distribution: Normal | None = None
    # A call, as rsl-rl's own GaussianDistribution makes it. OpenHomie *assigns* False here, which
    # leaves validation on and replaces torch's method with a bool, so every rsl-rl distribution
    # built later in the process fails with "'bool' object is not callable".
    Normal.set_default_validate_args(False)

Attributes

is_recurrent class-attribute instance-attribute

is_recurrent = False

num_actor_obs instance-attribute

num_actor_obs = num_actor_obs

num_critic_obs instance-attribute

num_critic_obs = num_critic_obs

num_one_step_obs instance-attribute

num_one_step_obs = num_one_step_obs

actor_history_length instance-attribute

actor_history_length = actor_history_length

actor_proprioceptive_obs_length instance-attribute

actor_proprioceptive_obs_length = actor_history_length * num_one_step_obs

num_actions instance-attribute

num_actions = num_actions

dynamic_latent_dim instance-attribute

dynamic_latent_dim = dynamic_latent_dim

obs_group instance-attribute

obs_group = obs_group

critic_obs_group instance-attribute

critic_obs_group = critic_obs_group

estimator instance-attribute

estimator = HIMEstimator(temporal_steps=actor_history_length, num_one_step_obs=num_one_step_obs, num_height_points=0, latent_dim=dynamic_latent_dim, target_skip_dim=target_skip_dim)

actor instance-attribute

actor = nn.Sequential(*actor_layers)

critic instance-attribute

critic = nn.Sequential(*critic_layers)

std instance-attribute

std = nn.Parameter(init_noise_std * torch.ones(num_actions))

distribution instance-attribute

distribution: Normal | None = None

action_mean property

action_mean: Tensor

action_std property

action_std: Tensor

entropy property

entropy: Tensor

output_std property

output_std: Tensor

Current action-noise std (used for logging by the runner).

Methods:

from_groups classmethod

from_groups(obs, actor_group: str, critic_group: str, num_actions: int, *, actor_hidden_dims, critic_hidden_dims, activation: str, init_noise_std: float, target_skip_dim: int) -> HIMActorCritic

Size the networks off the observation TensorDict, checking it has the layout HIM assumes.

The actor_group holds a whole number of stacked frames, the critic_group one frame followed by the true base linear velocity (VELOCITY_DIM).

Source code in eden/learning/wrappers/rsl_rl/him/actor_critic.py

@classmethod
def from_groups(
    cls,
    obs,
    actor_group: str,
    critic_group: str,
    num_actions: int,
    *,
    actor_hidden_dims,
    critic_hidden_dims,
    activation: str,
    init_noise_std: float,
    target_skip_dim: int,
) -> HIMActorCritic:
    """Size the networks off the observation TensorDict, checking it has the layout HIM assumes.

    The ``actor_group`` holds a whole number of stacked frames, the ``critic_group`` one frame followed
    by the true base linear velocity (``VELOCITY_DIM``).
    """
    num_actor_obs = int(obs[actor_group].shape[-1])
    num_critic_obs = int(obs[critic_group].shape[-1])
    frame = num_critic_obs - VELOCITY_DIM
    if frame <= 0 or num_actor_obs % frame:
        raise ValueError(
            f"HIM needs the {critic_group!r} group to be one actor frame + the {VELOCITY_DIM}-d base linear "
            f"velocity, and the {actor_group!r} group a whole number of such frames; got actor width "
            f"{num_actor_obs} and critic width {num_critic_obs} (frame {frame})."
        )
    return cls(
        num_actor_obs=num_actor_obs,
        num_critic_obs=num_critic_obs,
        num_one_step_obs=frame,
        actor_history_length=num_actor_obs // frame,
        num_actions=num_actions,
        actor_hidden_dims=tuple(actor_hidden_dims),
        critic_hidden_dims=tuple(critic_hidden_dims),
        activation=activation,
        init_noise_std=float(init_noise_std),
        obs_group=actor_group,
        critic_obs_group=critic_group,
        target_skip_dim=target_skip_dim,
    )

from_rsl_rl_v3 classmethod

from_rsl_rl_v3(obs, obs_groups: dict, num_actions: int, **policy_cfg) -> HIMActorCritic

rsl-rl-lib 3's construction call, policy_class(obs, obs_groups, num_actions, **policy_cfg).

Reached through the policy.class_name that HIMPPOAlgorithmOptions.adapt_runner_dict writes.

Source code in eden/learning/wrappers/rsl_rl/him/actor_critic.py

@classmethod
def from_rsl_rl_v3(cls, obs, obs_groups: dict, num_actions: int, **policy_cfg) -> HIMActorCritic:
    """rsl-rl-lib 3's construction call, ``policy_class(obs, obs_groups, num_actions, **policy_cfg)``.

    Reached through the ``policy.class_name`` that ``HIMPPOAlgorithmOptions.adapt_runner_dict`` writes.
    """
    if policy_cfg.get("actor_obs_normalization") or policy_cfg.get("critic_obs_normalization"):
        raise ValueError(_NO_OBS_NORMALIZATION)
    std_type = policy_cfg.get("noise_std_type", _HIM_STD_TYPE)
    if std_type != _HIM_STD_TYPE:
        raise _unsupported(
            "policy.noise_std_type", std_type, _HIM_STD_TYPE, "HIM learns the std itself, not its log"
        )
    return cls.from_groups(
        obs,
        single_group(obs_groups, "policy"),
        single_group(obs_groups, "critic"),
        num_actions,
        actor_hidden_dims=policy_cfg["actor_hidden_dims"],
        critic_hidden_dims=policy_cfg["critic_hidden_dims"],
        activation=policy_cfg["activation"],
        init_noise_std=policy_cfg["init_noise_std"],
        target_skip_dim=policy_cfg["target_skip_dim"],
    )

reset

reset(dones=None) -> None

Source code in eden/learning/wrappers/rsl_rl/him/actor_critic.py

def reset(self, dones=None) -> None:
    pass

forward

forward(obs) -> Tensor

The deterministic action for an observation TensorDict: what rsl-rl-lib 4/5's get_inference_policy serves.

Source code in eden/learning/wrappers/rsl_rl/him/actor_critic.py

def forward(self, obs) -> torch.Tensor:
    """The deterministic action for an observation TensorDict: what rsl-rl-lib 4/5's ``get_inference_policy``
    serves.
    """
    return self.act_inference(obs)

update_distribution

update_distribution(obs_history: Tensor) -> None

Source code in eden/learning/wrappers/rsl_rl/him/actor_critic.py

def update_distribution(self, obs_history: torch.Tensor) -> None:
    action_mean = self.actor(self._actor_input(obs_history))
    self.distribution = Normal(action_mean, action_mean * 0.0 + self.std)

act

act(obs_history: Tensor, **kwargs) -> Tensor

Source code in eden/learning/wrappers/rsl_rl/him/actor_critic.py

def act(self, obs_history: torch.Tensor, **kwargs) -> torch.Tensor:
    self.update_distribution(obs_history)
    return self.distribution.sample()

get_actions_log_prob

get_actions_log_prob(actions: Tensor) -> Tensor

Source code in eden/learning/wrappers/rsl_rl/him/actor_critic.py

def get_actions_log_prob(self, actions: torch.Tensor) -> torch.Tensor:
    return self.distribution.log_prob(actions).sum(dim=-1)

act_inference

act_inference(obs_history, **kwargs) -> Tensor

The deterministic action, from the history tensor or the whole observation TensorDict.

rsl-rl-lib 3's get_inference_policy hands out this method itself, called on the TensorDict.

Source code in eden/learning/wrappers/rsl_rl/him/actor_critic.py

def act_inference(self, obs_history, **kwargs) -> torch.Tensor:
    """The deterministic action, from the history tensor or the whole observation TensorDict.

    rsl-rl-lib 3's ``get_inference_policy`` hands out this method itself, called on the TensorDict.
    """
    if not isinstance(obs_history, torch.Tensor):
        obs_history = obs_history[self.obs_group]
    return self.actor(self._actor_input(obs_history))

evaluate

evaluate(critic_observations: Tensor, **kwargs) -> Tensor

Source code in eden/learning/wrappers/rsl_rl/him/actor_critic.py

def evaluate(self, critic_observations: torch.Tensor, **kwargs) -> torch.Tensor:
    return self.critic(critic_observations)

update_estimator

update_estimator(obs_history: Tensor, next_critic_obs: Tensor, lr: float | None = None)

Source code in eden/learning/wrappers/rsl_rl/him/actor_critic.py

def update_estimator(self, obs_history: torch.Tensor, next_critic_obs: torch.Tensor, lr: float | None = None):
    return self.estimator.update(obs_history[:, : self.actor_proprioceptive_obs_length], next_critic_obs, lr)

FlashSAC

Off-policy SAC from the external rsl-rl-flashsac package, which is not on PyPI: install it with uv sync --inexact --extra flashsac --group flashsac-git in a checkout.

Runner options for FlashSAC (github.com/Holiday-Robot/rsl_rl_flashsac): off-policy SAC with SimBa-style residual networks, a distributional critic and a large GPU replay buffer.

Declarative only — importing this module pulls in neither rsl_rl_flashsac nor rsl_rl, so a task config that uses it imports anywhere. Training needs the flashsac extra (uv sync --inexact --extra flashsac --group flashsac-git; rsl-rl-lib >= 5.4). The class names are the ones rsl_rl_flashsac resolves (its own modules first, then rsl_rl's).

Classes

RslRlFlashSacActorOptions

RslRlFlashSacActorOptions(**data)

Bases: ConfigurableOptions

The FlashSAC actor: a tanh-squashed Gaussian on residual blocks.

PARAMETER	DESCRIPTION
class_name	The model class rsl_rl_flashsac resolves. TYPE: str
num_blocks	Residual blocks. TYPE: int
hidden_dim	Block width. TYPE: int
log_std_min	Bounds of the log standard deviation. TYPE: float
log_std_max	Bounds of the log standard deviation. TYPE: float

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

class_name class-attribute instance-attribute

class_name: str = 'FlashSACActor'

num_blocks class-attribute instance-attribute

num_blocks: int = 2

hidden_dim class-attribute instance-attribute

hidden_dim: int = 128

log_std_min class-attribute instance-attribute

log_std_min: float = -10.0

log_std_max class-attribute instance-attribute

log_std_max: float = 2.0

RslRlFlashSacCriticOptions

RslRlFlashSacCriticOptions(**data)

Bases: ConfigurableOptions

The FlashSAC critic: a distributional (categorical) double Q on residual blocks.

PARAMETER	DESCRIPTION
class_name	The model class rsl_rl_flashsac resolves. TYPE: str
num_blocks	Residual blocks. TYPE: int
hidden_dim	Block width. TYPE: int
num_bins	Return bins. TYPE: int
min_v	Support of the (normalised) return. TYPE: float
max_v	Support of the (normalised) return. TYPE: float

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

class_name class-attribute instance-attribute

class_name: str = 'FlashSACCritic'

num_blocks class-attribute instance-attribute

num_blocks: int = 2

hidden_dim class-attribute instance-attribute

hidden_dim: int = 256

num_bins class-attribute instance-attribute

num_bins: int = 101

min_v class-attribute instance-attribute

min_v: float = -5.0

max_v class-attribute instance-attribute

max_v: float = 5.0

RslRlFlashSacAlgorithmOptions

RslRlFlashSacAlgorithmOptions(**data)

Bases: ConfigurableOptions

FlashSAC's hyperparameters (rsl_rl_flashsac.algorithms.FlashSAC).

PARAMETER	DESCRIPTION
class_name	The algorithm class rsl_rl_flashsac resolves. TYPE: str
replay_buffer_size	Transitions kept, over all envs. TYPE: int
buffer_min_length	Transitions before the first update (random actions until then). TYPE: int
buffer_optimize_memory_usage	Store each observation once (next observations by index). TYPE: bool
buffer_device	Where the buffer lives; None the training device. TYPE: str | None
buffer_obs_dtype	Observation storage dtype (e.g. "float16"); None keeps float32. TYPE: str | None
num_learning_epochs	Gradient updates per runner iteration = their product. TYPE: int
num_mini_batches	Gradient updates per runner iteration = their product. TYPE: int
mini_batch_size	Transitions per update. TYPE: int
learning_rate_init	Warm-up start, peak and final learning rate. TYPE: float
learning_rate_peak	Warm-up start, peak and final learning rate. TYPE: float
learning_rate_end	Warm-up start, peak and final learning rate. TYPE: float
learning_rate_warmup_steps	Updates of linear warm-up. TYPE: int
learning_rate_decay_steps	Updates of decay; None spans the run (max_iterations x updates per iteration). TYPE: int | None
actor_bc_alpha	Behaviour-cloning weight on the actor loss. TYPE: float
actor_noise_zeta_mu	Exploration-noise schedule. TYPE: float
actor_noise_zeta_max	Exploration-noise schedule. TYPE: float
actor_update_period	Critic updates per actor update. TYPE: int
critic_target_update_tau	Polyak rate of the target critic. TYPE: float
temp_initial_value	Entropy temperature start and target action standard deviation. TYPE: float
temp_target_sigma	Entropy temperature start and target action standard deviation. TYPE: float
temp_target_entropy	Explicit target entropy; None derives it from temp_target_sigma. TYPE: float | None
gamma	Discount. TYPE: float
n_steps	n-step return length. TYPE: int
normalize_reward	Normalise rewards by a running return scale. TYPE: bool
normalized_G_max	Return bound after normalisation (should match the critic's support). TYPE: float
use_compile	torch.compile the networks; train in mixed precision. TYPE: bool
use_amp	torch.compile the networks; train in mixed precision. TYPE: bool
compile_mode	torch.compile mode. TYPE: str

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

class_name class-attribute instance-attribute

class_name: str = 'FlashSAC'

replay_buffer_size class-attribute instance-attribute

replay_buffer_size: int = 10000000

buffer_min_length class-attribute instance-attribute

buffer_min_length: int = 100000

buffer_optimize_memory_usage class-attribute instance-attribute

buffer_optimize_memory_usage: bool = True

buffer_device class-attribute instance-attribute

buffer_device: str | None = None

buffer_obs_dtype class-attribute instance-attribute

buffer_obs_dtype: str | None = None

num_learning_epochs class-attribute instance-attribute

num_learning_epochs: int = 1

num_mini_batches class-attribute instance-attribute

num_mini_batches: int = 2

mini_batch_size class-attribute instance-attribute

mini_batch_size: int = 2048

learning_rate_init class-attribute instance-attribute

learning_rate_init: float = 0.0003

learning_rate_peak class-attribute instance-attribute

learning_rate_peak: float = 0.0003

learning_rate_end class-attribute instance-attribute

learning_rate_end: float = 0.00015

learning_rate_warmup_steps class-attribute instance-attribute

learning_rate_warmup_steps: int = 0

learning_rate_decay_steps class-attribute instance-attribute

learning_rate_decay_steps: int | None = None

actor_bc_alpha class-attribute instance-attribute

actor_bc_alpha: float = 0.0

actor_noise_zeta_mu class-attribute instance-attribute

actor_noise_zeta_mu: float = 2.0

actor_noise_zeta_max class-attribute instance-attribute

actor_noise_zeta_max: int = 16

actor_update_period class-attribute instance-attribute

actor_update_period: int = 2

critic_target_update_tau class-attribute instance-attribute

critic_target_update_tau: float = 0.01

temp_initial_value class-attribute instance-attribute

temp_initial_value: float = 0.01

temp_target_sigma class-attribute instance-attribute

temp_target_sigma: float = 0.15

temp_target_entropy class-attribute instance-attribute

temp_target_entropy: float | None = None

gamma class-attribute instance-attribute

gamma: float = 0.99

n_steps class-attribute instance-attribute

n_steps: int = 3

normalize_reward class-attribute instance-attribute

normalize_reward: bool = True

normalized_G_max class-attribute instance-attribute

normalized_G_max: float = 5.0

use_compile class-attribute instance-attribute

use_compile: bool = True

compile_mode class-attribute instance-attribute

compile_mode: str = 'auto'

use_amp class-attribute instance-attribute

use_amp: bool = True

RslRlFlashSacRunnerOptions

RslRlFlashSacRunnerOptions(**data)

Bases: RslRlBaseRunnerOptions

rsl_rl_flashsac's OffPolicyRunner: each iteration steps every env num_steps_per_env times, then runs num_learning_epochs x num_mini_batches updates; eden train selects it from this class.

PARAMETER	DESCRIPTION
class_name	The runner class. TYPE: str
max_iterations	The run's length in iterations; None when :attr:num_env_steps sets it. Exactly one of the two is set. TYPE: int | None
num_env_steps	The run's length in env steps, over all envs of every rank: max_iterations is derived from it once the env count and the torchrun world size are known (:meth:for_num_envs, which :class:FlashSacVecEnvWrapper applies), so the run follows --num_envs and the GPU count. Exactly one of the two is set. TYPE: int | None
log_interval	Iterations per log line (one iteration is one env step by default). TYPE: int
start_training	Iterations before the first update. TYPE: int
action_bound_scale	The policy's tanh output is scaled to [-action_bound_scale, action_bound_scale]. TYPE: float
actor	The actor. TYPE: RslRlFlashSacActorOptions
critic	The critic. TYPE: RslRlFlashSacCriticOptions
algorithm	The algorithm. TYPE: RslRlFlashSacAlgorithmOptions

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

class_name class-attribute instance-attribute

class_name: Literal['OffPolicyRunner'] = 'OffPolicyRunner'

num_steps_per_env class-attribute instance-attribute

num_steps_per_env: int = 1

max_iterations class-attribute instance-attribute

max_iterations: int | None = 50000

num_env_steps class-attribute instance-attribute

num_env_steps: int | None = None

save_interval class-attribute instance-attribute

save_interval: int = 10000

log_interval class-attribute instance-attribute

log_interval: int = 20

start_training class-attribute instance-attribute

start_training: int = 0

action_bound_scale class-attribute instance-attribute

action_bound_scale: float = 1.0

actor class-attribute instance-attribute

actor: RslRlFlashSacActorOptions = RslRlFlashSacActorOptions()

critic class-attribute instance-attribute

critic: RslRlFlashSacCriticOptions = RslRlFlashSacCriticOptions()

algorithm class-attribute instance-attribute

algorithm: RslRlFlashSacAlgorithmOptions = RslRlFlashSacAlgorithmOptions()

Methods:

for_num_envs

for_num_envs(num_envs: int, *, world_size: int) -> RslRlFlashSacRunnerOptions

These options with max_iterations sized for num_envs envs on each of world_size ranks when :attr:num_env_steps sets the run.

Every rank runs max_iterations iterations of its own num_envs envs, so the run's env steps are split across the ranks rather than repeated on each. With algorithm.learning_rate_decay_steps=None the learning-rate decay spans the derived run too.

Source code in eden/learning/wrappers/rsl_rl/flashsac/options.py

def for_num_envs(self, num_envs: int, *, world_size: int) -> RslRlFlashSacRunnerOptions:
    """These options with ``max_iterations`` sized for ``num_envs`` envs on each of ``world_size`` ranks when
    :attr:`num_env_steps` sets the run.

    Every rank runs ``max_iterations`` iterations of its own ``num_envs`` envs, so the run's env steps are split
    across the ranks rather than repeated on each. With ``algorithm.learning_rate_decay_steps=None`` the
    learning-rate decay spans the derived run too.
    """
    if self.num_env_steps is None:
        return self
    per_iteration = num_envs * self.num_steps_per_env * world_size
    return self.model_copy(
        update={"max_iterations": max(1, self.num_env_steps // per_iteration), "num_env_steps": None}
    )

The Eden side of FlashSAC: the VecEnv adapter and the runner factory eden train and eden inference use for :class:~eden.learning.wrappers.rsl_rl.flashsac.options.RslRlFlashSacRunnerOptions.

Classes

FlashSacVecEnvWrapper

FlashSacVecEnvWrapper(env: RLEnvBase, options: RslRlFlashSacRunnerOptions)

Bases: RslRlVecEnvWrapper

:class:RslRlVecEnvWrapper plus what FlashSAC needs of the env.

• The action bounds FlashSAC reads (action_bias / action_scale), applied here: the policy emits tanh actions in [-1, 1] and the env receives action_bias + action_scale * a, i.e. [-action_bound_scale, action_bound_scale].
• On time-outs, the observation before the reset (extras["time_outs_obs"]), so a truncated episode bootstraps from its true last state. The env must record it (EnvOptions.record_final_observations), or every truncation would bootstrap from the next episode's first observation.
• The run sized for this env count over every rank (:meth:RslRlFlashSacRunnerOptions.for_num_envs): read runner_options.max_iterations off the wrapper, not the config.

For training, a mimic env whose motion terms mine hard motions (adaptive_motion_weights) is refused: FlashSAC's runner does not step that curriculum.

Source code in eden/learning/wrappers/rsl_rl/flashsac/runner.py

def __init__(self, env: RLEnvBase, options: RslRlFlashSacRunnerOptions) -> None:
    if not isinstance(options, RslRlFlashSacRunnerOptions):
        raise TypeError(
            f"FlashSacVecEnvWrapper takes RslRlFlashSacRunnerOptions, got {type(options).__name__}; build a PPO "
            "runner's env with RslRlVecEnvWrapper."
        )
    _refuse_unsupported(env.config, training=env.training)
    super().__init__(env, options=options.for_num_envs(env.num_envs, world_size=eden_dist.get_world_size()))
    self.action_bias = torch.zeros(self.num_actions, device=self.device)
    self.action_scale = torch.full((self.num_actions,), float(options.action_bound_scale), device=self.device)

Attributes

action_bias instance-attribute

action_bias = torch.zeros(self.num_actions, device=self.device)

action_scale instance-attribute

action_scale = torch.full((self.num_actions,), float(options.action_bound_scale), device=self.device)

Methods:

from_config classmethod

from_config(cfg, *, show_viewer: bool = False, eval_mode: bool = False, apply_seed: bool = True) -> FlashSacVecEnvWrapper

:meth:RslRlVecEnvWrapper.from_config, refusing what FlashSAC cannot do before the env is built.

The task's own :meth:~eden.utils.configs.EdenConfig.check_bindings runs first: it can name the variant flag that sets these fields (flashdexretarget: --learner flashsac).

Source code in eden/learning/wrappers/rsl_rl/flashsac/runner.py

@classmethod
def from_config(
    cls, cfg, *, show_viewer: bool = False, eval_mode: bool = False, apply_seed: bool = True
) -> FlashSacVecEnvWrapper:
    """:meth:`RslRlVecEnvWrapper.from_config`, refusing what FlashSAC cannot do before the env is built.

    The task's own :meth:`~eden.utils.configs.EdenConfig.check_bindings` runs first: it can name the variant flag
    that sets these fields (flashdexretarget: ``--learner flashsac``).
    """
    cfg.check_bindings()
    _refuse_unsupported(cfg, training=not eval_mode)
    return super().from_config(cfg, show_viewer=show_viewer, eval_mode=eval_mode, apply_seed=apply_seed)

step

step(actions: Tensor) -> tuple[TensorDict, Tensor, Tensor, dict]

Source code in eden/learning/wrappers/rsl_rl/flashsac/runner.py

def step(self, actions: torch.Tensor) -> tuple[TensorDict, torch.Tensor, torch.Tensor, dict]:
    obs, rew, dones, extras = super().step(self.action_bias + self.action_scale * actions)
    final = extras.get("final_observations")
    if final is not None:
        extras["time_outs_obs"] = TensorDict(dict(final), batch_size=[self.num_envs])
    return obs, rew, dones, extras

Functions:

make_flashsac_runner

make_flashsac_runner(env: FlashSacVecEnvWrapper, checkpoint: str | None = None, log_dir: str | None = None, *, inference: bool = False)

Build rsl_rl_flashsac's OffPolicyRunner, with Eden's runner features (:func:_runner_class), on env, optionally loading checkpoint.

PARAMETER	DESCRIPTION
env	The env. TYPE: FlashSacVecEnvWrapper
checkpoint	A FlashSAC checkpoint (model_<iter>.pt) to load, strictly. TYPE: str | None DEFAULT: None
log_dir	The run directory; en.log_dir by default. TYPE: str | None DEFAULT: None
inference	Only the policy will run (eden inference): the replay buffer is allocated at one row per env instead of replay_buffer_size (50M rows, ~154 GB, for flashdexretarget), and nothing is compiled. TYPE: bool DEFAULT: False

Source code in eden/learning/wrappers/rsl_rl/flashsac/runner.py

def make_flashsac_runner(
    env: FlashSacVecEnvWrapper,
    checkpoint: str | None = None,
    log_dir: str | None = None,
    *,
    inference: bool = False,
):
    """Build ``rsl_rl_flashsac``'s ``OffPolicyRunner``, with Eden's runner features (:func:`_runner_class`), on
    ``env``, optionally loading ``checkpoint``.

    Parameters
    ----------
    env : FlashSacVecEnvWrapper
        The env.
    checkpoint : str | None
        A FlashSAC checkpoint (``model_<iter>.pt``) to load, strictly.
    log_dir : str | None
        The run directory; ``en.log_dir`` by default.
    inference : bool
        Only the policy will run (``eden inference``): the replay buffer is allocated at one row per env instead
        of ``replay_buffer_size`` (50M rows, ~154 GB, for flashdexretarget), and nothing is compiled.
    """
    runner_cls = _runner_class(eden_dist.is_distributed())
    cfg = env.runner_dict  # a fresh copy, which constructing the runner consumes (FlashSAC pops the class names)
    if inference:
        cfg["algorithm"].update(replay_buffer_size=env.num_envs, buffer_min_length=env.num_envs, use_compile=False)
    device = cfg["device"]
    runner = runner_cls(env, cfg, log_dir=log_dir or en.log_dir, device=device)
    if checkpoint:
        # FlashSAC's own load, strict: the rsl_rl PPO path's legacy conversion and key audit are for its layouts.
        runner.load(checkpoint, map_location=device, strict=True)
    return runner

ConTrack CMDP-PPO

Constrained-MDP PPO trained through the external rsl_rl_contrack fork of rsl_rl.

Runner/algorithm options for the external rsl_rl_contrack CMDP-PPO backend.

Declarative only — importing this module needs neither the backend nor rsl_rl's runner API, so the contrack task config that selects it imports on any install. The runner its class_name names (:class:~eden.learning.wrappers.rsl_rl_contrack.runner.ContrackOnPolicyRunner) needs the backend and raises with the install command when it is missing, so eden train --task contrack is the training entry point.

ConTrack (https://github.com/EmptyBlueBox/ConTrack) trains tracking policies as a constrained MDP: rewards are split into task (rg), style (rs), and penalty (rp) groups, the critic outputs two values for the (rg+rp, rs+rp) channels, and a dual variable mixes the two advantage streams with w_task = sigmoid(λ) updated toward a target task ratio alpha. The reward groups are declared on the env config (RewardGroupOptions); these options carry only the learning hyperparameters. Defaults match the ConTrack release.

Attributes

Classes

RslRlContrackAlgorithmOptions

RslRlContrackAlgorithmOptions(**data)

Bases: RslRlPpoAlgorithmOptions

PPO options extended with ConTrack's dual-controller (CMDP) parameters.

PARAMETER	DESCRIPTION
cmdp_alpha	Target task ratio alpha — the constraint J_g >= alpha * J_g*.
cmdp_lambda_lr	Dual step size for λ.
cmdp_warmup_iters	Updates with w_task = 1 (pure task reward) before the dual controller takes over.
cmdp_vg_window	Episodes in the running mean of the task return Ĵ_g.
cmdp_vg_star_window	Updates in the running max defining the normalizer J_g*.
cmdp_lambda_init	Initial λ (0 gives w_task = 0.5).

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

entropy_coef class-attribute instance-attribute

entropy_coef: float = 0.001

num_learning_epochs class-attribute instance-attribute

num_learning_epochs: int = 10

num_mini_batches class-attribute instance-attribute

num_mini_batches: int = 16

learning_rate class-attribute instance-attribute

learning_rate: float = 0.0005

cmdp_alpha class-attribute instance-attribute

cmdp_alpha: float = 0.9

cmdp_lambda_lr class-attribute instance-attribute

cmdp_lambda_lr: float = 0.05

cmdp_warmup_iters class-attribute instance-attribute

cmdp_warmup_iters: int = 0

cmdp_vg_window class-attribute instance-attribute

cmdp_vg_window: int = 100

cmdp_vg_star_window class-attribute instance-attribute

cmdp_vg_star_window: int = 20

cmdp_lambda_init class-attribute instance-attribute

cmdp_lambda_init: float = 0.0

RslRlContrackRunnerOptions

RslRlContrackRunnerOptions(**data)

Bases: RslRlOnPolicyRunnerOptions

On-policy runner options that select the rsl_rl_contrack backend.

The actor / critic model options are consumed by rsl_rl_contrack.modules.ActorCritic (MLP-only; the critic automatically outputs the two CMDP value channels). Defaults mirror the ConTrack release hyperparameters, so RslRlContrackRunnerOptions() is training-ready.

PARAMETER	DESCRIPTION
class_name	The runner eden train / eden inference build; leave at the default. TYPE: str

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

class_name class-attribute instance-attribute

class_name: str = _CONTRACK_RUNNER_CLASS_NAME

num_steps_per_env class-attribute instance-attribute

num_steps_per_env: int = 16

max_iterations class-attribute instance-attribute

max_iterations: int = 5000

save_interval class-attribute instance-attribute

save_interval: int = 200

clip_actions class-attribute instance-attribute

clip_actions: ClipActions = 1.0

experiment_name class-attribute instance-attribute

experiment_name: str = 'contrack'

obs_groups class-attribute instance-attribute

obs_groups: dict = {'actor': ['policy'], 'critic': ['critic']}

actor class-attribute instance-attribute

actor: RslRlModelWithDistributionOptions = RslRlModelWithDistributionOptions(hidden_dims=[256, 256, 128], distribution_cfg=RslRlGaussianDistributionOptions(init_std=1.0, std_type='scalar'))

critic class-attribute instance-attribute

critic: RslRlModelOptions = RslRlModelOptions(hidden_dims=[256, 256, 128])

algorithm class-attribute instance-attribute

algorithm: RslRlContrackAlgorithmOptions = RslRlContrackAlgorithmOptions()

Runner for the external rsl_rl_contrack backend (ConTrack's CMDP-PPO fork of rsl_rl).

rsl_rl_contrack is a standalone RL library installed separately (it is not an Eden extra). :class:ContrackOnPolicyRunner is what :class:~eden.learning.wrappers.rsl_rl_contrack.options.RslRlContrackRunnerOptions's class_name names, so eden train / eden inference build it. Importing this module raises with the install command when the backend is missing — eden train resolves the runner before building the env, so that happens in seconds.

Attributes

CONTRACK_RUNNER_CLASS_NAME module-attribute

CONTRACK_RUNNER_CLASS_NAME = 'eden.learning.wrappers.rsl_rl_contrack.runner:ContrackOnPolicyRunner'

Classes

ContrackOnPolicyRunner

ContrackOnPolicyRunner(env: RslRlVecEnvWrapper, train_cfg: dict, log_dir: str | None = None, device: str = 'cpu')

Bases: CheckpointRetentionMixin, OnPolicyRunner

The fork's OnPolicyRunner behind Eden's runner signature, with Eden's checkpoint retention.

eden train hands it a plain :class:~eden.envs.wrappers.rsl_rl_env.RslRlVecEnvWrapper and the version-translated runner dict; the fork needs neither. It re-wraps the same built env in :class:ContrackVecEnvWrapper (which adds the CMDP reward-group extras the fork's PPO reads) and rebuilds the fork's train_cfg from the env's runner options. Checkpoint retention matters because the fork saves model_<it>.pt the same way rsl_rl does, and keep_last_checkpoints would otherwise be silently ignored.

Source code in eden/learning/wrappers/rsl_rl_contrack/runner.py

def __init__(self, env: RslRlVecEnvWrapper, train_cfg: dict, log_dir: str | None = None, device: str = "cpu"):
    if not isinstance(env, ContrackVecEnvWrapper):
        env = ContrackVecEnvWrapper(env.unwrapped, env.runner_options)
    super().__init__(env, contrack_train_cfg(env.runner_options), log_dir=log_dir, device=device)

Methods:

Functions:

contrack_train_cfg

contrack_train_cfg(options: RslRlContrackRunnerOptions) -> dict

The fork's train_cfg dict, assembled directly from the runner options.

No rsl-rl-lib version translation — the fork has a single, known API.

Source code in eden/learning/wrappers/rsl_rl_contrack/runner.py

def contrack_train_cfg(options: RslRlContrackRunnerOptions) -> dict:
    """The fork's ``train_cfg`` dict, assembled directly from the runner options.

    No rsl-rl-lib version translation — the fork has a single, known API.
    """
    # Eden's obs_groups convention names the actor set "actor"; the contrack
    # runner (rsl_rl v5 lineage) calls it "policy".
    obs_groups = {("policy" if k == "actor" else k): list(v) for k, v in options.obs_groups.items()}
    if not obs_groups:
        obs_groups = {"policy": ["policy"], "critic": ["critic"]}

    distribution = options.actor.distribution_cfg
    algorithm = options.algorithm
    return {
        "num_steps_per_env": options.num_steps_per_env,
        "save_interval": options.save_interval,
        "keep_last_checkpoints": options.keep_last_checkpoints,
        "obs_groups": obs_groups,
        "logger": options.logger,
        "wandb_project": options.wandb_project,
        "neptune_project": options.neptune_project,
        "experiment_name": options.experiment_name or "contrack",
        "run_name": options.run_name,
        "policy": {
            "class_name": "ActorCritic",
            "actor_hidden_dims": list(options.actor.hidden_dims),
            "critic_hidden_dims": list(options.critic.hidden_dims),
            "activation": options.actor.activation,
            "actor_obs_normalization": options.actor.obs_normalization,
            "critic_obs_normalization": options.critic.obs_normalization,
            "init_noise_std": distribution.init_std if distribution is not None else 1.0,
            "noise_std_type": distribution.std_type if distribution is not None else "scalar",
        },
        "algorithm": {
            "class_name": "PPO",
            "num_learning_epochs": algorithm.num_learning_epochs,
            "num_mini_batches": algorithm.num_mini_batches,
            "clip_param": algorithm.clip_param,
            "gamma": algorithm.gamma,
            "lam": algorithm.lam,
            "value_loss_coef": algorithm.value_loss_coef,
            "entropy_coef": algorithm.entropy_coef,
            "learning_rate": algorithm.learning_rate,
            "max_grad_norm": algorithm.max_grad_norm,
            "use_clipped_value_loss": algorithm.use_clipped_value_loss,
            "schedule": algorithm.schedule,
            "desired_kl": algorithm.desired_kl,
            "normalize_advantage_per_mini_batch": algorithm.normalize_advantage_per_mini_batch,
            "cmdp_alpha": algorithm.cmdp_alpha,
            "cmdp_lambda_lr": algorithm.cmdp_lambda_lr,
            "cmdp_warmup_iters": algorithm.cmdp_warmup_iters,
            "cmdp_vg_window": algorithm.cmdp_vg_window,
            "cmdp_vg_star_window": algorithm.cmdp_vg_star_window,
            "cmdp_lambda_init": algorithm.cmdp_lambda_init,
        },
    }

VecEnv wrapper that feeds ConTrack's CMDP-PPO (rsl_rl_contrack) its per-step reward-group extras.

ConTrack's CMDP-PPO consumes three per-step reward-group scalars from extras — rg (task), rs (style), rp (penalty) — and trains a two-channel critic on (rg + rp, rs + rp). :class:ContrackVecEnvWrapper forwards :meth:RewardManager.get_group_reward for the task's declared reward groups (see RewardGroupOptions in the task config); the per-episode group returns are logged by the reward manager itself (Episode_Reward_Group/<g> in the reset extras).

It also applies runner_options.check_for_nan (the fork's runner has no NaN check of its own), and maintains the object-tracking-error running sums the ConTrack OnPolicyRunner logs unconditionally (it_obj_pos_err_sum / it_obj_rot_err_sum / it_step_count).

Needs no backend: :class:~eden.learning.wrappers.rsl_rl_contrack.runner.ContrackOnPolicyRunner wraps the env in it.

Attributes

CMDP_GROUPS module-attribute

CMDP_GROUPS = ('rg', 'rs', 'rp')

Classes

ContrackVecEnvWrapper

ContrackVecEnvWrapper(env: RLEnvBase, options: RslRlContrackRunnerOptions)

Bases: RslRlVecEnvWrapper

RslRl wrapper that injects the per-step CMDP reward-group extras.

Requires the wrapped env's reward terms to be declared in exactly the CMDP groups (rg / rs / rp via RewardGroupOptions) — enforced at construction so a term outside the three channels (which the CMDP critic would never see) fails before training starts.

Source code in eden/learning/wrappers/rsl_rl_contrack/env.py

def __init__(self, env: RLEnvBase, options: RslRlContrackRunnerOptions):
    super().__init__(env, options)
    groups = set(self.unwrapped.reward_manager.group_term_names)
    if groups != set(CMDP_GROUPS):
        raise ValueError(
            f"ConTrack CMDP-PPO requires the reward terms to be partitioned into exactly the groups "
            f"{CMDP_GROUPS}; the env declares {sorted(groups)}. Declare every reward term inside "
            f"rg/rs/rp RewardGroupOptions (a term outside them would be invisible to the CMDP critic)."
        )

    # Read by rsl_rl_contrack.OnPolicyRunner.log() — zeroed there each iteration.
    self.it_obj_pos_err_sum = torch.zeros((), device=self.device)
    self.it_obj_rot_err_sum = torch.zeros((), device=self.device)
    self.it_step_count = 0

Attributes

object_entity_name class-attribute instance-attribute

object_entity_name: str = 'object'

it_obj_pos_err_sum instance-attribute

it_obj_pos_err_sum = torch.zeros((), device=self.device)

it_obj_rot_err_sum instance-attribute

it_obj_rot_err_sum = torch.zeros((), device=self.device)

it_step_count instance-attribute

it_step_count = 0

Methods:

step

step(actions: Tensor) -> tuple[TensorDict, Tensor, Tensor, dict]

Source code in eden/learning/wrappers/rsl_rl_contrack/env.py

def step(self, actions: torch.Tensor) -> tuple[TensorDict, torch.Tensor, torch.Tensor, dict]:
    obs, rew, dones, extras = super().step(actions)
    reward_manager = self.unwrapped.reward_manager
    for group in CMDP_GROUPS:
        extras[group] = reward_manager.get_group_reward(group)
    if self.runner_options.check_for_nan:
        _raise_on_nan(obs, rew, {group: extras[group] for group in CMDP_GROUPS})
    self._accumulate_object_tracking_error(self.unwrapped)
    return obs, rew, dones, extras

Functions:

See also

• Training runs guide
• Learning options
• Policy export guide