Publish native Metatron v1 G1 package, model card and artifact index

G1/MODEL_CARD.md

@@ -0,0 +1,155 @@
+# Metatron v1
+
+## License
+
+Model license: NVIDIA Open Model License
+
+The model weights, checkpoints, serialized model configurations and exported
+models are licensed under the [NVIDIA Open Model License](LICENSE.pdf).
+See [NOTICE.txt](NOTICE.txt) for attribution.
+
+- Deployment Geography: Global.
+- Use Case: Researchers and developers running G1 motion control, recovery and
+  interaction in compatible simulators or supervised real-G1 deployment workflows.
+- Release: Metatron v1; native package hosted in `nvidia/Metatron-v1/G1`.
+  Use the pinned revision selected by the ProtoMotions download helper. No public
+  ProtoMotions release tag is asserted by this model card.
+ Include both files when
+redistributing these artifacts.
+
+ProtoMotions source code remains Apache-2.0. Training datasets and robot assets
+have separate terms; this model license does not grant rights to those materials.
+
+## Overview
+
+Metatron v1 is a Unitree G1 policy for motion control, recovery and
+scene-conditioned interaction. It supports sparse motion goals, velocity
+commands and scene inputs.
+
+| Field | Value |
+| --- | --- |
+| Family | MaskedMimic spherical conditional VAE |
+| Robot | Unitree G1, 29 actions |
+| Motion corpus | BONES, AMASS/LAFAN, recovery and SAMP interaction data |
+| Primary checkpoint | `inference_last.ckpt` |
+
+## Intended Use
+
+- Control G1 motion from partial future-pose and velocity commands.
+- Perform recovery and scene-conditioned interaction.
+- Run in simulation or deploy on a real Unitree G1.
+
+## Training
+
+- Training simulator: **IsaacLab**.
+- Training method: supervised imitation from frozen motion-tracking, recovery
+  and scene-interaction experts.
+- Training data: BONES, AMASS/LAFAN, fall recovery and SAMP scene interaction.
+- Architecture: 64-dimensional spherical latent, six-layer, four-head,
+  1024-wide Transformer prior and conditioned action decoder.
+
+## Inputs and Outputs
+
+The policy consumes current and historical robot state and actions, masked
+global/local future-pose targets, style samples, planar velocity commands and up
+to two scene-box tokens. It outputs 29 raw actions, converted to motor commands
+by the configured G1 PD mapping. Inference uses the learned prior without
+external teachers.
+
+The exact body/joint order, history layout, normalization, model dimensions and
+action mapping are stored in `resolved_configs.json`.
+
+## Software Integration
+
+Native inference uses PyTorch and ProtoMotions, with simulator dependencies from
+the matching ProtoMotions checkout. Linux is the native reference environment;
+this is not a claim that every backend, OS or GPU has been qualified. Browser
+exports and their runtime requirements are documented in the repository's
+[main model card](../README.md).
+
+## Artifacts
+
+- `inference_last.ckpt`: policy weights and normalization statistics, including
+  the posterior; optimizer state is omitted.
+- `resolved_configs.json`: authoritative native inference configuration.
+- `experiment_config.py` and `config.yaml`: configuration reconstruction code
+  and arguments.
+- `provenance.json`: artifact integrity metadata.
+
+## Runtime Support
+
+- Training simulator: **IsaacLab**.
+- Simulator expectation: **Expected to transfer**.
+- Also works in other simulators, including **MuJoCo**, and on the **real G1**.
+- Preserve the trained observation, action, history and control-rate contracts.
+
+For native playback, install ProtoMotions and its MuJoCo requirements, then run
+from the repository root (the helper downloads a pinned package from this Hub):
+
+```bash
+python scripts/download_metatron.py
+python protomotions/inference_agent.py --simulator mujoco --num-envs 1 \
+  --checkpoint data/pretrained_models/metatron/G1/inference_last.ckpt \
+  --inference-config examples/inference/metatron.py
+```
+
+The bundled native preset supports empty scenes and cycles
+idle/location/steering/reaching with `X`.
+
+### Scene Support
+
+The **weights retain scene conditioning**. The PyTorch prior accepts
+`scene_box_corners` (two boxes, eight XYZ corners each, in the root-heading
+frame) and `scene_box_mask` (`true` means visible). Our separately pre-exported
+39-input ONNX includes the original scene preprocessing and also retains this
+capability. A scene-capable host must supply the matching object measurements.
+
+The **native preset and generic re-export disable scenes**: the bundled
+observation functions always produce zero corner tokens and false visibility
+masks. Running `deployment/export_to_onnx.py` on this checkpoint preserves the
+weights but builds that scene-disabled observation path. It does not reproduce
+the scene-capable 39-input ONNX. The preset rejects a furniture scene file
+rather than silently ignoring it.
+
+This distinction is intentional: these weights support scenes, but this
+package does not add a scene provider or a SAMP training recipe to ProtoMotions.
+
+## Limitations
+
+- This is an inference checkpoint, not a resumable training checkpoint.
+- Training datasets and furniture assets are not included.
+- The basic student recipes do not reproduce the full multi-expert training run.
+- Recovery and interaction can fail depending on goals, initialization and
+  physics. Hardware deployment requires robot-side safety limits, supervision
+  and an emergency-stop mechanism.
+
+## Provenance
+
+This package contains the SAMP-capable G1 policy used for the Metatron web demo,
+with its configuration adapted to the current ProtoMotions runtime.
+
+## Evaluation and Responsible Use
+
+The release tests exercise native configuration loading, finite actions through
+MuJoCo rollout/reset, frozen normalization parity and scene-input sensitivity.
+These are integration checks, not a comprehensive task-success or safety benchmark.
+No native throughput or task-success aggregate is published in this card.
+
+| Area | Scope and limitations |
+| --- | --- |
+| Bias | Coverage follows the motion corpora and expert rollouts; no demographic fairness evaluation is claimed. |
+| Explainability | A learned conditional policy maps state and partial goals to actions; attention and latent variables are not explanations or guarantees. |
+| Privacy | The package contains model tensors/configs, not raw participant records or motion datasets. No formal privacy guarantee or leakage audit is claimed. |
+| Safety & Security | Not a safety controller. Validate observations, limits and recovery on the intended platform; use supervision and emergency stops on hardware. Load only trusted, verified checkpoint/config packages. |
+
+Training motion sources are listed above. Labels are expert actions produced in
+simulation for supervised imitation. Corpus splits, collection protocols and
+demographic breakdowns are not included in this release; consult the source
+dataset documentation and licenses. Testing here refers to release integration
+checks, not an independently curated held-out benchmark. No training motion,
+recovery clip, furniture asset or dataset license is redistributed by this package.
+
+## References
+
+- [ProtoMotions source and citation](https://github.com/NVlabs/ProtoMotions)
+- [MaskedMimic](https://research.nvidia.com/labs/par/maskedmimic/)

G1/NOTICE.txt

@@ -0,0 +1,6 @@
+Licensed by NVIDIA Corporation under the NVIDIA Open Model License.
+Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+This notice applies to the model weights and serialized model configuration.
+Source code remains Apache-2.0. Robot assets, motion data, furniture and runtime
+dependencies retain their separate licenses.

G1/README.md

@@ -0,0 +1,30 @@
+# Metatron v1 — G1
+
+Native PyTorch inference package for the Unitree G1, including policy weights,
+portable resolved configuration and reconstruction recipe. This is the
+SAMP-capable G1 policy, not a SOMA policy or a resumable training checkpoint.
+
+See [MODEL_CARD.md](MODEL_CARD.md) for architecture, inputs, scene limitations,
+training provenance and safe use. [LICENSE.pdf](LICENSE.pdf) and
+[NOTICE.txt](NOTICE.txt) must accompany redistributed model artifacts; code is
+Apache-2.0 and third-party data/assets keep their own terms.
+
+From a compatible ProtoMotions checkout with MuJoCo dependencies installed:
+
+```bash
+python scripts/download_metatron.py
+python protomotions/inference_agent.py --simulator mujoco --num-envs 1 \
+  --checkpoint data/pretrained_models/metatron/G1/inference_last.ckpt \
+  --inference-config examples/inference/metatron.py
+```
+
+The helper pins the Hub revision and verifies every file. It includes the
+license and model card. `--configs-only` downloads the small sidecars for config
+tests, not a runnable policy. Avoid mixing files from different revisions.
+The current native preset cycles idle/location/steering/reaching with `X` and
+masks scene inputs. Use the separate `demo/model.onnx` and its paired application
+for the scene-capable browser interface; a generic native re-export is not the
+same graph. `racer/network.onnx` is the six-robot browser interface.
+
+No training data, recovery clips or robot meshes are included here. Robot and
+example motion assets are supplied separately by ProtoMotions under their own terms.

G1/config.yaml

@@ -0,0 +1,16 @@
+{
+  "robot_name": "g1",
+  "simulator": "mujoco",
+  "num_envs": 1,
+  "batch_size": 4,
+  "motion_file": "examples/data/g1_crouch_to_run.motion",
+  "experiment_path": "data/pretrained_models/metatron/G1/experiment_config.py",
+  "experiment_name": "metatron_v1_playback",
+  "scenes_file": null,
+  "checkpoint": null,
+  "headless": true,
+  "seed": 0,
+  "torch_deterministic": false,
+  "overrides": [],
+  "create_config_only": true
+}

G1/experiment_config.py

@@ -0,0 +1,546 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Reconstruct the scene-conditioned Metatron v1 network for flat-ground play.
+
+This recipe preserves the trained module order, normalization and five-body
+order. It is an inference qualification recipe, not a reproduction of the
+four-group training run. Scene tokens are explicitly hidden on this empty
+terrain; furniture interaction needs its own scene observations.
+"""
+
+from argparse import Namespace
+from typing import Any, List, Optional, Sequence
+
+from examples.experiments.masked_mimic.deployable_g1_observations import (
+    configure_observations,
+)
+from examples.inference.metatron import (
+    empty_scene_corners,
+    empty_scene_mask,
+)
+from protomotions.agents.common.config import (
+    MLPLayerConfig,
+    MLPWithConcatConfig,
+    MaskedMeanConfig,
+    ModuleContainerConfig,
+    ModuleOperationForwardConfig,
+    ModuleOperationReshapeConfig,
+    ObsProcessorConfig,
+    TransformerConfig,
+)
+from protomotions.agents.evaluators.config import MaskedMimicEvaluatorConfig
+from protomotions.agents.supervised.masked_mimic_config import (
+    KLDScheduleConfig,
+    LatentGeometry,
+    MaskedMimicModelConfig,
+    MaskedMimicSupervisedAgentConfig,
+    MaskedMimicVAEConfig,
+    TemporalLossConfig,
+    VAENoiseType,
+)
+from protomotions.components.motion_lib import MotionLibConfig
+from protomotions.components.scene_lib import SceneLibConfig
+from protomotions.components.terrains.config import TerrainConfig
+from protomotions.envs.action import make_bm_pd_action_config
+from protomotions.envs.base_env.config import EnvConfig
+from protomotions.envs.body_selection import BodySelection
+from protomotions.envs.context_views import EnvContext
+from protomotions.envs.control.masked_mimic_control import (
+    MaskedMimicControlConfig,
+    MaskedMimicStyleConfig,
+    MaskedMimicVelocityCommandConfig,
+)
+from protomotions.envs.mdp_component import MdpComponent
+from protomotions.envs.motion_manager.config import MimicMotionManagerConfig
+from protomotions.envs.obs.cast import to_float
+from protomotions.envs.obs.humanoid import compute_humanoid_max_coords_observations
+from protomotions.envs.obs.masked_mimic import (
+    compute_local_target_poses,
+    compute_target_masks_only,
+    compute_target_poses_only,
+    compute_target_time_offsets,
+)
+from protomotions.envs.obs.target_poses import build_max_coords_target_poses
+from protomotions.envs.sensors.config import (
+    IMUSensor,
+    JointEncoderSensor,
+    LocalizationSensor,
+    NoNoise,
+    SensorConfig,
+)
+from protomotions.robot_configs.base import RobotConfig
+from protomotions.simulator.base_simulator.config import SimulatorConfig
+from protomotions.utils.experiment_configs import ExperimentConfigs
+
+BODY_NAMES = (
+    "torso_link",
+    "right_ankle_roll_link",
+    "left_ankle_roll_link",
+    "left_rubber_hand",
+    "right_rubber_hand",
+)
+FRAMES = 5
+TOKEN_SIZE = 1024
+LATENT_SIZE = 64
+
+
+def configure_robot_and_simulator(
+    robot_cfg: RobotConfig,
+    simulator_cfg: SimulatorConfig,
+    args: Namespace,
+) -> None:
+    """Use G1's BM gains, pelvis IMU and 50 Hz policy without perturbations."""
+    if args.robot_name != "g1":
+        raise ValueError("This checkpoint architecture and body order require G1")
+    robot_cfg.update_fields(
+        trackable_bodies_subset=list(BODY_NAMES),
+        contact_bodies=[],
+        contact_partners=[],
+    )
+    robot_cfg.reset_noise = None
+    simulator_cfg.domain_randomization = None
+
+
+def terrain_config(args: Namespace) -> TerrainConfig:
+    """Flat ground, without procedural obstacles."""
+    return TerrainConfig(terrain_proportions=[0.0] * 7 + [1.0])
+
+
+def scene_lib_config(args: Namespace) -> SceneLibConfig:
+    """No perceived furniture; reject scenes rather than silently hiding them."""
+    if args.scenes_file is not None:
+        raise ValueError("Metatron native playback supports empty scenes only")
+    return SceneLibConfig(scene_file=None)
+
+
+def motion_lib_config(args: Namespace) -> MotionLibConfig:
+    """Reference clips for reset and optional motion tracking."""
+    return MotionLibConfig(motion_file=args.motion_file)
+
+
+def env_config(robot_cfg: RobotConfig, args: Namespace) -> EnvConfig:
+    """Bind the saved observation contract to current sensor/history owners."""
+    bodies = BodySelection.named(*BODY_NAMES)
+    cfg = EnvConfig(
+        motion_manager=MimicMotionManagerConfig(init_start_prob=1.0),
+        action_config=make_bm_pd_action_config(robot_cfg),
+        sensors=SensorConfig(
+            joint_encoders=JointEncoderSensor(
+                pos_noise=NoNoise("inference"),
+                vel_noise=NoNoise("inference"),
+            ),
+            imu=IMUSensor(
+                quat_noise=NoNoise("inference"),
+                gyro_noise=NoNoise("inference"),
+            ),
+            localization=LocalizationSensor(
+                body_name="pelvis",
+                noise=NoNoise("inference"),
+            ),
+        ),
+        control_components={
+            "masked_mimic": MaskedMimicControlConfig(
+                future_steps=(1, 2, 5, 7, 12, 18, 25),
+                reference_height_smoothing_seconds=0.1,
+                num_masked_future_steps=FRAMES,
+                sparse_target_condition_probability=0.5,
+                local_target_condition_probability=0.5,
+                anchor_pose_submasking=True,
+                style=MaskedMimicStyleConfig(),
+                velocity_command=MaskedMimicVelocityCommandConfig(),
+            )
+        },
+        observation_components={
+            "max_coords_obs": MdpComponent.from_kernel(
+                compute_humanoid_max_coords_observations,
+            ),
+            "mimic_target_poses": MdpComponent.from_kernel(
+                build_max_coords_target_poses,
+                with_velocities=True,
+                with_relative=True,
+            ),
+            "masked_mimic_target_poses": MdpComponent.from_kernel(
+                compute_target_poses_only,
+                conditionable_body_ids=bodies,
+                include_root_relative=False,
+                anchor_translation_mask=EnvContext.masked_mimic.anchor_translation_mask,
+                anchor_rotation_mask=EnvContext.masked_mimic.anchor_rotation_mask,
+            ),
+            "masked_mimic_local_target_poses": MdpComponent.from_kernel(
+                compute_local_target_poses,
+                conditionable_body_ids=bodies,
+            ),
+            "previous_actions": MdpComponent.from_kernel(
+                to_float,
+                x=EnvContext.current_action,
+            ),
+            "scene_box_corners": MdpComponent.from_kernel(empty_scene_corners),
+            "scene_box_mask": MdpComponent.from_kernel(empty_scene_mask),
+        },
+    )
+    configure_observations(robot_cfg, cfg)
+    for name, path in (
+        ("poses", EnvContext.masked_mimic.style_pose_samples),
+        ("velocities", EnvContext.masked_mimic.style_velocity_samples),
+        ("masks", EnvContext.masked_mimic.style_condition_masks),
+        ("frame_masks", EnvContext.masked_mimic.style_frame_masks),
+    ):
+        cfg.observation_components["masked_mimic_style_" + name] = (
+            MdpComponent.from_kernel(to_float, x=path)
+        )
+    for prefix, masks, times, frames in (
+        (
+            "masked_mimic",
+            EnvContext.masked_mimic.target_bodies_masks,
+            EnvContext.masked_mimic.time_offsets,
+            EnvContext.masked_mimic.target_poses_masks,
+        ),
+        (
+            "masked_mimic_local",
+            EnvContext.masked_mimic.local_target_bodies_masks,
+            EnvContext.masked_mimic.local_time_offsets,
+            EnvContext.masked_mimic.local_target_poses_masks,
+        ),
+    ):
+        submasks = {}
+        if prefix == "masked_mimic":
+            submasks = dict(
+                anchor_translation_mask=EnvContext.masked_mimic.anchor_translation_mask,
+                anchor_rotation_mask=EnvContext.masked_mimic.anchor_rotation_mask,
+            )
+        cfg.observation_components[prefix + "_target_masks"] = MdpComponent.from_kernel(
+            compute_target_masks_only,
+            conditionable_body_ids=bodies,
+            masked_mimic_target_bodies_masks=masks,
+            **submasks,
+        )
+        cfg.observation_components[prefix + "_target_times"] = MdpComponent.from_kernel(
+            compute_target_time_offsets,
+            masked_mimic_time_offsets=times,
+        )
+        cfg.observation_components[prefix + "_target_poses_masks"] = (
+            MdpComponent.from_kernel(to_float, x=frames)
+        )
+    return cfg
+
+
+def _obs(
+    keys: List[str],
+    outputs: List[str],
+    *,
+    frames: Optional[int] = None,
+    normalize: bool = True,
+) -> ObsProcessorConfig:
+    operations = []
+    if frames is not None:
+        operations.append(ModuleOperationReshapeConfig(["batch_size", frames, -1]))
+    if normalize:
+        operations.append(ModuleOperationForwardConfig())
+    return ObsProcessorConfig(
+        in_keys=keys,
+        out_keys=outputs,
+        normalize_obs=normalize,
+        module_operations=operations,
+    )
+
+
+def _mlp(
+    keys: List[str],
+    output: str,
+    size: int = TOKEN_SIZE,
+    *,
+    widths: Sequence[int] = (1024,) * 3,
+    frames: Optional[int] = None,
+    output_activation: Optional[str] = None,
+) -> MLPWithConcatConfig:
+    operations = [ModuleOperationForwardConfig()]
+    if frames is not None:
+        operations.insert(0, ModuleOperationReshapeConfig(["batch_size", frames, -1]))
+    return MLPWithConcatConfig(
+        in_keys=keys,
+        out_keys=[output],
+        num_out=size,
+        layers=[MLPLayerConfig(units=n, activation="silu") for n in widths],
+        module_operations=operations,
+        output_activation=output_activation,
+    )
+
+
+def _head(keys: List[str], output: str) -> MLPWithConcatConfig:
+    return _mlp(keys, output, LATENT_SIZE, widths=(512, 256, 128))
+
+
+def _container(models: List[Any], outputs: List[str]) -> ModuleContainerConfig:
+    produced = {key for model in models for key in model.out_keys}
+    inputs = list(
+        dict.fromkeys(
+            key for model in models for key in model.in_keys if key not in produced
+        )
+    )
+    return ModuleContainerConfig(models=models, in_keys=inputs, out_keys=outputs)
+
+
+def _prior() -> ModuleContainerConfig:
+    # Order matches the packaged state_dict. The two time processors contain
+    # identical copies of the original shared statistics: independent standard
+    # processors preserve inference without extending the common module API.
+    models = [
+        _obs(["scene_box_corners"], ["scene_box_norm"]),
+        _mlp(["scene_box_norm"], "scene_box_tokens", widths=(1024,)),
+        _obs(["previous_actions"], ["prior_previous_actions_norm"]),
+        _obs(["reduced_coords_obs"], ["prior_reduced_coords_norm"]),
+        _obs(["masked_mimic_local_target_poses"], ["local_target_poses_seq"], frames=5),
+        _obs(["fk_masked_mimic_target_poses"], ["global_target_poses_seq"], frames=5),
+        _obs(
+            ["masked_mimic_local_target_masks"],
+            ["local_target_masks_seq"],
+            frames=5,
+            normalize=False,
+        ),
+        _obs(
+            ["masked_mimic_target_masks"],
+            ["global_target_masks_seq"],
+            frames=5,
+            normalize=False,
+        ),
+        _obs(
+            ["masked_mimic_target_times"],
+            ["target_times_seq"],
+            frames=5,
+        ),
+        _obs(
+            ["masked_mimic_local_target_times"],
+            ["local_target_times_seq"],
+            frames=5,
+        ),
+        _obs(["historical_proprioception"], ["historical_pose_obs_seq"], frames=5),
+        _obs(["masked_mimic_style_poses"], ["style_poses_seq"], frames=5),
+        _obs(["masked_mimic_style_velocities"], ["style_velocities_seq"], frames=5),
+        _obs(
+            ["masked_mimic_style_masks"],
+            ["style_condition_masks_seq"],
+            frames=5,
+            normalize=False,
+        ),
+        _obs(
+            ["masked_mimic_style_frame_masks"],
+            ["style_frame_masks_seq"],
+            frames=5,
+            normalize=False,
+        ),
+        _obs(["fk_max_coords_obs"], ["prior_fk_max_coords_obs_norm"]),
+        _mlp(
+            [
+                "prior_fk_max_coords_obs_norm",
+                "prior_reduced_coords_norm",
+                "prior_previous_actions_norm",
+            ],
+            "current_state_token",
+            frames=1,
+        ),
+        _mlp(
+            ["masked_mimic_velocity_command", "masked_mimic_velocity_command_mask"],
+            "masked_mimic_velocity_command_token",
+            frames=1,
+        ),
+        _mlp(
+            ["style_poses_seq", "style_velocities_seq", "style_condition_masks_seq"],
+            "masked_mimic_style_frame_tokens",
+            frames=5,
+        ),
+        MaskedMeanConfig(
+            in_keys=["masked_mimic_style_frame_tokens", "style_frame_masks_seq"],
+            out_keys=["masked_mimic_style_token"],
+        ),
+        _mlp(
+            [
+                "local_target_poses_seq",
+                "local_target_masks_seq",
+                "local_target_times_seq",
+            ],
+            "masked_mimic_local_target_token",
+            frames=5,
+        ),
+        _mlp(
+            ["global_target_poses_seq", "global_target_masks_seq", "target_times_seq"],
+            "masked_mimic_global_target_token",
+            frames=5,
+        ),
+        _mlp(["historical_pose_obs_seq"], "historical_pose_obs_token", frames=5),
+        TransformerConfig(
+            in_keys=[
+                "current_state_token",
+                "masked_mimic_velocity_command_token",
+                "masked_mimic_local_target_token",
+                "masked_mimic_global_target_token",
+                "masked_mimic_style_token",
+                "historical_pose_obs_token",
+                "masked_mimic_velocity_command_mask",
+                "masked_mimic_local_target_poses_masks",
+                "masked_mimic_target_poses_masks",
+                "masked_mimic_style_frame_masks",
+                "scene_box_tokens",
+                "scene_box_mask",
+            ],
+            out_keys=["transformer_out"],
+            input_and_mask_mapping={
+                "masked_mimic_velocity_command_token": (
+                    "masked_mimic_velocity_command_mask"
+                ),
+                "masked_mimic_local_target_token": (
+                    "masked_mimic_local_target_poses_masks"
+                ),
+                "masked_mimic_global_target_token": "masked_mimic_target_poses_masks",
+                "masked_mimic_style_token": "masked_mimic_style_frame_masks",
+                "scene_box_tokens": "scene_box_mask",
+            },
+            transformer_token_size=1024,
+            latent_dim=1024,
+            num_heads=4,
+            ff_size=4096,
+            num_layers=6,
+            activation="silu",
+            output_activation="silu",
+            dropout=0.0,
+        ),
+        _head(["transformer_out"], "decoder_condition"),
+        _head(["transformer_out"], "prior_mu"),
+        _head(["transformer_out"], "prior_logvar"),
+    ]
+    return _container(models, ["prior_mu", "prior_logvar", "decoder_condition"])
+
+
+def _encoder() -> ModuleContainerConfig:
+    models = [
+        _obs(["max_coords_obs"], ["max_coords_obs_norm"]),
+        _obs(["mimic_target_poses"], ["mimic_target_poses_norm"]),
+        _obs(
+            ["masked_mimic_local_target_poses"],
+            ["masked_mimic_local_target_poses_norm"],
+        ),
+        _obs(["masked_mimic_target_poses"], ["masked_mimic_global_target_poses_norm"]),
+        _obs(
+            ["masked_mimic_target_times"],
+            ["masked_mimic_target_times_norm"],
+        ),
+        _obs(
+            ["masked_mimic_local_target_times"],
+            ["masked_mimic_local_target_times_norm"],
+        ),
+        _obs(["masked_mimic_style_poses"], ["masked_mimic_style_poses_norm"]),
+        _obs(["masked_mimic_style_velocities"], ["masked_mimic_style_velocities_norm"]),
+        ObsProcessorConfig(
+            in_keys=["masked_mimic_style_masks"],
+            out_keys=["masked_mimic_style_masks_flat"],
+            normalize_obs=False,
+            module_operations=[ModuleOperationReshapeConfig(["batch_size", -1])],
+        ),
+        _mlp(
+            [
+                "masked_mimic_local_target_poses_norm",
+                "masked_mimic_local_target_masks",
+                "masked_mimic_local_target_times_norm",
+                "masked_mimic_local_target_poses_masks",
+            ],
+            "encoder_local_target_features",
+            256,
+            widths=(512, 512),
+        ),
+        _mlp(
+            [
+                "masked_mimic_global_target_poses_norm",
+                "masked_mimic_target_masks",
+                "masked_mimic_target_times_norm",
+                "masked_mimic_target_poses_masks",
+            ],
+            "encoder_global_target_features",
+            256,
+            widths=(512, 512),
+        ),
+        _mlp(
+            [
+                "max_coords_obs_norm",
+                "mimic_target_poses_norm",
+                "encoder_local_target_features",
+                "encoder_global_target_features",
+                "masked_mimic_style_poses_norm",
+                "masked_mimic_style_velocities_norm",
+                "masked_mimic_style_masks_flat",
+                "masked_mimic_style_frame_masks",
+                "masked_mimic_velocity_command",
+                "masked_mimic_velocity_command_mask",
+            ],
+            "encoder_trunk_out",
+            512,
+            widths=(1024,) * 5,
+            output_activation="silu",
+        ),
+        _head(["encoder_trunk_out"], "encoder_mu"),
+        _head(["encoder_trunk_out"], "encoder_logvar"),
+    ]
+    return _container(models, ["encoder_mu", "encoder_logvar"])
+
+
+def _decoder() -> ModuleContainerConfig:
+    return _container(
+        [
+            _obs(["reduced_coords_obs"], ["decoder_reduced_coords_norm"]),
+            _obs(["fk_max_coords_obs"], ["fk_max_coords_obs_norm"]),
+            _obs(["previous_actions"], ["previous_actions_norm"]),
+            _head(["decoder_condition", "vae_latent"], "solution_condition_features"),
+            _head(["solution_condition_features"], "latent_dynamics_predicted"),
+            _mlp(
+                [
+                    "fk_max_coords_obs_norm",
+                    "previous_actions_norm",
+                    "solution_condition_features",
+                    "decoder_reduced_coords_norm",
+                ],
+                "actor_trunk_out",
+                29,
+            ),
+        ],
+        ["actor_trunk_out", "solution_condition_features", "latent_dynamics_predicted"],
+    )
+
+
+def agent_config(
+    robot_cfg: RobotConfig, env_cfg: EnvConfig, args: Namespace
+) -> MaskedMimicSupervisedAgentConfig:
+    """The learned spherical prior and conditional decoder, with no external teacher."""
+    return MaskedMimicSupervisedAgentConfig(
+        model=MaskedMimicModelConfig(
+            encoder=_encoder(),
+            prior=_prior(),
+            trunk=_decoder(),
+            vae=MaskedMimicVAEConfig(
+                vae_latent_dim=64,
+                latent_geometry=LatentGeometry.SPHERE,
+                vae_noise_type=VAENoiseType.ZEROS,
+                prior_mean_unit_norm_coeff=0.001,
+                kld_schedule=KLDScheduleConfig(
+                    start_epoch=500, end_epoch=2000, cyclical=True
+                ),
+            ),
+            temporal_losses=[
+                TemporalLossConfig("privileged_action", 0.05),
+                TemporalLossConfig("solution_condition_features", 0.001),
+                TemporalLossConfig(
+                    "decoder_condition",
+                    0.001,
+                    prediction_key="latent_dynamics_predicted",
+                    forward=True,
+                ),
+            ],
+        ),
+        evaluator=MaskedMimicEvaluatorConfig(),
+        batch_size=args.batch_size,
+        training_max_steps=args.training_max_steps,
+    )
+
+
+def apply_inference_overrides(configs: ExperimentConfigs, args: Namespace) -> None:
+    """Keep the saved sensor mounts and use deterministic prior inference."""
+    configs.simulator.domain_randomization = None
+    configs.agent.model.vae.vae_noise_type = VAENoiseType.ZEROS

G1/inference_last.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e5fc2e47a8f9ee9b96a674cabbf236b3c510f887427645a14bc1779b10f0b43a
+size 459455983

G1/provenance.json

@@ -0,0 +1,9 @@
+{
+  "name": "Metatron v1",
+  "robot": "g1",
+  "inference_checkpoint": "inference_last.ckpt",
+  "inference_checkpoint_sha256": "e5fc2e47a8f9ee9b96a674cabbf236b3c510f887427645a14bc1779b10f0b43a",
+  "inference_checkpoint_bytes": 459455983,
+  "model_tensor_count": 278,
+  "resumable": false
+}

G1/resolved_configs.json

@@ -0,0 +1,2601 @@
+{
+  "format_version": 1,
+  "protomotions": "3.1",
+  "configs": {
+    "robot": {
+      "$type": "protomotions.robot_configs.g1.G1RobotConfig",
+      "asset": {
+        "$type": "protomotions.robot_configs.base.RobotAssetConfig",
+        "asset_root": "protomotions/data/assets",
+        "self_collisions": true,
+        "asset_file_name": "mjcf/g1_holo_compat.xml",
+        "replace_cylinder_with_capsule": true,
+        "thickness": 0.01,
+        "max_angular_velocity": 1000.0,
+        "max_linear_velocity": 1000.0,
+        "density": 0.001,
+        "angular_damping": 0.0,
+        "linear_damping": 0.0,
+        "disable_gravity": null,
+        "fix_base_link": null
+      },
+      "common_naming_to_robot_body_names": {
+        "all_left_foot_bodies": [
+          "left_ankle_roll_link"
+        ],
+        "all_right_foot_bodies": [
+          "right_ankle_roll_link"
+        ],
+        "all_left_hand_bodies": [
+          "left_rubber_hand"
+        ],
+        "all_right_hand_bodies": [
+          "right_rubber_hand"
+        ],
+        "head_body_name": [
+          "head"
+        ],
+        "torso_body_name": [
+          "torso_link"
+        ]
+      },
+      "default_root_height": 0.8,
+      "semantic_forward_axis_xy": [
+        1.0,
+        0.0
+      ],
+      "default_dof_pos": {
+        "$tensor": [
+          -0.31200000643730164,
+          0.0,
+          0.0,
+          0.6690000295639038,
+          -0.3630000054836273,
+          0.0,
+          -0.31200000643730164,
+          0.0,
+          0.0,
+          0.6690000295639038,
+          -0.3630000054836273,
+          0.0,
+          0.0,
+          0.0,
+          0.0,
+          0.20000000298023224,
+          0.20000000298023224,
+          0.0,
+          0.6000000238418579,
+          0.0,
+          0.0,
+          0.0,
+          0.20000000298023224,
+          -0.20000000298023224,
+          0.0,
+          0.6000000238418579,
+          0.0,
+          0.0,
+          0.0
+        ],
+        "dtype": "float32"
+      },
+      "contact_bodies": [],
+      "contact_partners": [],
+      "trackable_bodies_subset": [
+        "torso_link",
+        "right_ankle_roll_link",
+        "left_ankle_roll_link",
+        "left_rubber_hand",
+        "right_rubber_hand"
+      ],
+      "non_termination_contact_bodies": "all",
+      "init_state": null,
+      "reset_noise": null,
+      "mimic_small_marker_bodies": null,
+      "anchor_body_name": "torso_link",
+      "imu_body_name": "pelvis",
+      "imu_mount_quat": [
+        0.0,
+        0.0,
+        0.0,
+        1.0
+      ],
+      "contact_pairs_multiplier": 16,
+      "control": {
+        "$type": "protomotions.robot_configs.base.ControlConfig",
+        "override_control_info": {
+          ".*_hip_(pitch|yaw)_joint": {
+            "$type": "protomotions.components.pose_lib.ControlInfo",
+            "stiffness": 40.17923847137318,
+            "damping": 2.5578897650279457,
+            "armature": 0.01017752,
+            "friction": null,
+            "effort_limit": 88,
+            "velocity_limit": 32
+          },
+          ".*_hip_roll_joint": {
+            "$type": "protomotions.components.pose_lib.ControlInfo",
+            "stiffness": 99.09842777666113,
+            "damping": 6.3088018534966395,
+            "armature": 0.025101925,
+            "friction": null,
+            "effort_limit": 139,
+            "velocity_limit": 20
+          },
+          ".*_knee_joint": {
+            "$type": "protomotions.components.pose_lib.ControlInfo",
+            "stiffness": 99.09842777666113,
+            "damping": 6.3088018534966395,
+            "armature": 0.025101925,
+            "friction": null,
+            "effort_limit": 139,
+            "velocity_limit": 20
+          },
+          ".*_ankle_.*": {
+            "$type": "protomotions.components.pose_lib.ControlInfo",
+            "stiffness": 28.50124619574858,
+            "damping": 1.814445686584846,
+            "armature": 0.00721945,
+            "friction": null,
+            "effort_limit": 50,
+            "velocity_limit": 37
+          },
+          "waist_yaw_joint": {
+            "$type": "protomotions.components.pose_lib.ControlInfo",
+            "stiffness": 40.17923847137318,
+            "damping": 2.5578897650279457,
+            "armature": 0.01017752,
+            "friction": null,
+            "effort_limit": 88,
+            "velocity_limit": 32
+          },
+          "waist_(roll|pitch)_joint": {
+            "$type": "protomotions.components.pose_lib.ControlInfo",
+            "stiffness": 28.50124619574858,
+            "damping": 1.814445686584846,
+            "armature": 0.00721945,
+            "friction": null,
+            "effort_limit": 50,
+            "velocity_limit": 37
+          },
+          ".*_(shoulder|elbow)_.*": {
+            "$type": "protomotions.components.pose_lib.ControlInfo",
+            "stiffness": 14.25062309787429,
+            "damping": 0.907222843292423,
+            "armature": 0.003609725,
+            "friction": null,
+            "effort_limit": 25,
+            "velocity_limit": 37
+          },
+          ".*_wrist_roll_joint": {
+            "$type": "protomotions.components.pose_lib.ControlInfo",
+            "stiffness": 14.25062309787429,
+            "damping": 0.907222843292423,
+            "armature": 0.003609725,
+            "friction": null,
+            "effort_limit": 25,
+            "velocity_limit": 37
+          },
+          ".*_wrist_pitch_joint": {
+            "$type": "protomotions.components.pose_lib.ControlInfo",
+            "stiffness": 16.77832748089279,
+            "damping": 1.06814150219,
+            "armature": 0.00425,
+            "friction": null,
+            "effort_limit": 5,
+            "velocity_limit": 22
+          },
+          ".*_wrist_yaw_joint": {
+            "$type": "protomotions.components.pose_lib.ControlInfo",
+            "stiffness": 16.77832748089279,
+            "damping": 1.06814150219,
+            "armature": 0.00425,
+            "friction": null,
+            "effort_limit": 5,
+            "velocity_limit": 22
+          }
+        },
+        "control_type": {
+          "$enum": "protomotions.robot_configs.base.ControlType",
+          "name": "BUILT_IN_PD"
+        },
+        "soft_pos_limit": 0.9,
+        "asset_default_joints": []
+      },
+      "control_hz": 50,
+      "simulation_params": {
+        "$type": "protomotions.robot_configs.base.SimulatorParams",
+        "isaacgym": {
+          "$type": "protomotions.simulator.isaacgym.config.IsaacGymSimParams",
+          "fps": 100,
+          "substeps": 2,
+          "physx": {
+            "$type": "protomotions.simulator.isaacgym.config.IsaacGymPhysXParams",
+            "num_threads": 4,
+            "solver_type": 1,
+            "num_position_iterations": 8,
+            "num_velocity_iterations": 4,
+            "contact_offset": 0.02,
+            "rest_offset": 0.0,
+            "bounce_threshold_velocity": 0.2,
+            "max_depenetration_velocity": 1,
+            "default_buffer_size_multiplier": 10.0
+          },
+          "flex": {
+            "$type": "protomotions.simulator.isaacgym.config.IsaacGymFlexParams",
+            "num_inner_iterations": 10,
+            "warm_start": 0.25
+          }
+        },
+        "isaaclab": {
+          "$type": "protomotions.simulator.isaaclab.config.IsaacLabSimParams",
+          "fps": 200,
+          "physx": {
+            "$type": "protomotions.simulator.isaaclab.config.IsaacLabPhysXParams",
+            "num_threads": 4,
+            "solver_type": 1,
+            "num_position_iterations": 8,
+            "num_velocity_iterations": 4,
+            "contact_offset": 0.02,
+            "rest_offset": 0.0,
+            "bounce_threshold_velocity": 0.2,
+            "max_depenetration_velocity": 1,
+            "default_buffer_size_multiplier": 10.0,
+            "gpu_found_lost_pairs_capacity": 2097152,
+            "gpu_max_rigid_contact_count": 8388608,
+            "gpu_found_lost_aggregate_pairs_capacity": 33554432,
+            "gpu_max_rigid_patch_count": 163840,
+            "gpu_collision_stack_size": 67108864
+          }
+        },
+        "newton": {
+          "$type": "protomotions.simulator.newton.config.NewtonSimParams",
+          "fps": 200,
+          "solver": "newton",
+          "integrator": "implicitfast",
+          "iterations": 100,
+          "ls_iterations": 50,
+          "ls_parallel": true,
+          "impratio": 10.0,
+          "njmax": 450,
+          "nconmax": 300,
+          "cone": "pyramidal",
+          "ccd_iterations": 200
+        },
+        "mujoco": {
+          "$type": "protomotions.simulator.mujoco.config.MujocoSimParams",
+          "fps": 1000,
+          "integrator": null
+        }
+      }
+    },
+    "simulator": {
+      "$type": "protomotions.simulator.mujoco.config.MujocoSimulatorConfig",
+      "_target_": "protomotions.simulator.mujoco.simulator.MujocoSimulator",
+      "w_last": false,
+      "default_robot_friction": 1.0,
+      "headless": true,
+      "num_envs": 1,
+      "sim": {
+        "$type": "protomotions.simulator.mujoco.config.MujocoSimParams",
+        "fps": 1000,
+        "integrator": null
+      },
+      "experiment_name": "metatron_v1_playback",
+      "camera": null,
+      "record_viewer": false,
+      "viewer_record_dir": "output/recordings/viewer",
+      "domain_randomization": null,
+      "projectile": {
+        "$type": "protomotions.simulator.base_simulator.config.ProjectileConfig",
+        "num_projectiles": 0,
+        "cube_half_size_range": [
+          0.05,
+          0.15
+        ],
+        "density": 500.0,
+        "speed_range": [
+          30.0,
+          40.0
+        ],
+        "spawn_distance_range": [
+          4.0,
+          5.0
+        ],
+        "spawn_height_range": [
+          -0.65,
+          1.1
+        ],
+        "direction_noise_std": 0.1,
+        "hide_delay": 2.0,
+        "hide_z": -2.0,
+        "hide_spacing": 4.0
+      },
+      "pd_target_max_accel": null,
+      "use_implicit_pd": true
+    },
+    "terrain": {
+      "$type": "protomotions.components.terrains.config.TerrainConfig",
+      "_target_": "protomotions.components.terrains.terrain.Terrain",
+      "map_length": 20.0,
+      "map_width": 20.0,
+      "border_size": 40.0,
+      "num_levels": 10,
+      "num_terrains": 10,
+      "terrain_proportions": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        1.0
+      ],
+      "slope_threshold": 0.9,
+      "num_samples_per_axis": 16,
+      "sample_width": 1.0,
+      "terrain_obs_num_samples": 256,
+      "horizontal_scale": 0.1,
+      "vertical_scale": 0.005,
+      "spacing_between_scenes": 10.0,
+      "minimal_humanoid_spacing": 1.0,
+      "terrain_path": null,
+      "load_terrain": false,
+      "save_terrain": false,
+      "sim_config": {
+        "$type": "protomotions.components.terrains.config.TerrainSimConfig",
+        "static_friction": 1.0,
+        "dynamic_friction": 1.0,
+        "restitution": 0.0,
+        "height_offset": 0.0,
+        "combine_mode": {
+          "$enum": "protomotions.components.terrains.config.CombineMode",
+          "name": "AVERAGE"
+        }
+      }
+    },
+    "scene_lib": {
+      "$type": "protomotions.components.scene_lib.SceneLibConfig",
+      "_target_": "protomotions.components.scene_lib.SceneLib",
+      "scene_file": null,
+      "asset_root": null,
+      "subset_method": {
+        "$enum": "protomotions.components.scene_lib.SubsetMethod",
+        "name": "FIRST"
+      },
+      "replicate_method": {
+        "$enum": "protomotions.components.scene_lib.ReplicationMethod",
+        "name": "WEIGHTED"
+      },
+      "pointcloud_samples_per_object": null,
+      "num_objects_per_env": null,
+      "mesh_collision_approximation": null,
+      "mesh_collision_max_convex_hulls": null,
+      "mesh_collision_hull_vertex_limit": null,
+      "mesh_collision_voxel_resolution": null,
+      "scene_indices": null,
+      "inline_scenes": null
+    },
+    "motion_lib": {
+      "$type": "protomotions.components.motion_lib.MotionLibConfig",
+      "_target_": "protomotions.components.motion_lib.MotionLib",
+      "motion_file": "examples/data/g1_crouch_to_run.motion",
+      "motion_file_shard_indices": null,
+      "motion_file_switch_mode": "fixed",
+      "get_motion_state_use_blend": true,
+      "max_seconds": null,
+      "clip_delta": 3.0,
+      "clip_seed": 42
+    },
+    "env": {
+      "$type": "protomotions.envs.base_env.config.EnvConfig",
+      "max_episode_length": 300,
+      "reset_grace_period": 5,
+      "_target_": "protomotions.envs.base_env.env.BaseEnv",
+      "motion_manager": {
+        "$type": "protomotions.envs.motion_manager.config.MimicMotionManagerConfig",
+        "_target_": "protomotions.envs.motion_manager.mimic_motion_manager.MimicMotionManager",
+        "init_start_prob": 1.0,
+        "sample_time_truncate_s": null,
+        "subset_method": null,
+        "exclude_motion_ids": null,
+        "exclude_motions_file": null,
+        "realign_motion_with_humanoid_on_each_step": false,
+        "resample_on_reset": true
+      },
+      "ref_respawn_offset": 0.05,
+      "ref_object_respawn_offset": 0.0,
+      "ref_contact_smooth_window": 0,
+      "recovery_reset": {
+        "$type": "protomotions.envs.base_env.config.RecoveryResetConfig",
+        "recovery_prob": 0.0,
+        "fall_sim_steps": 150
+      },
+      "skip_correct_terrain_height_on_flat": true,
+      "show_terrain_markers": false,
+      "save_dir": "",
+      "reward_components": {},
+      "control_components": {
+        "masked_mimic": {
+          "$type": "protomotions.envs.control.masked_mimic_control.MaskedMimicControlConfig",
+          "_target_": "protomotions.envs.control.masked_mimic_control.MaskedMimicControl",
+          "bootstrap_on_episode_end": true,
+          "future_steps": [
+            1,
+            2,
+            5,
+            7,
+            12,
+            18,
+            25
+          ],
+          "reference_height_smoothing_seconds": 0.1,
+          "low_pose_grace_seconds": 0.0,
+          "num_masked_future_steps": 5,
+          "target_profile_probabilities": [
+            0.2,
+            0.6,
+            0.2
+          ],
+          "target_gap_min_seconds": 0.1,
+          "target_gap_max_seconds": 2.0,
+          "visible_target_pose_prob": 0.8,
+          "fixed_conditioning": null,
+          "sparse_target_condition_probability": 0.5,
+          "local_target_condition_probability": 0.5,
+          "anchor_pose_submasking": true,
+          "style": {
+            "$type": "protomotions.envs.control.masked_mimic_control.MaskedMimicStyleConfig",
+            "window_seconds": 2.0,
+            "num_frames": 5,
+            "near_time_bias_power": 2.0,
+            "velocity_probability": 0.5,
+            "visible_frame_count_probabilities": [
+              0.5,
+              0.1,
+              0.1,
+              0.1,
+              0.1,
+              0.1
+            ]
+          },
+          "velocity_command": {
+            "$type": "protomotions.envs.control.masked_mimic_control.MaskedMimicVelocityCommandConfig",
+            "horizon_seconds": 0.5,
+            "probability": 0.5
+          }
+        }
+      },
+      "termination_components": {},
+      "observation_components": {
+        "max_coords_obs": {
+          "$kernel": "protomotions.envs.obs.humanoid.compute_humanoid_max_coords_observations",
+          "bindings": {
+            "body_pos": "current.rigid_body_pos",
+            "body_rot": "current.rigid_body_rot",
+            "body_vel": "current.rigid_body_vel",
+            "body_ang_vel": "current.rigid_body_ang_vel",
+            "ground_height": "current.ground_heights",
+            "anchor_body_index": "current.anchor_idx"
+          },
+          "params": {
+            "body_contacts": null,
+            "local_obs": true,
+            "root_height_obs": true,
+            "observe_contacts": false,
+            "w_last": true
+          }
+        },
+        "mimic_target_poses": {
+          "$kernel": "protomotions.envs.obs.target_poses.build_max_coords_target_poses",
+          "bindings": {
+            "current_state_body_pos": "current.rigid_body_pos",
+            "current_state_body_rot": "current.rigid_body_rot",
+            "current_state_body_vel": "current.rigid_body_vel",
+            "current_state_body_ang_vel": "current.rigid_body_ang_vel",
+            "mimic_ref_pos": "mimic.future_pos",
+            "mimic_ref_rot": "mimic.future_rot",
+            "mimic_ref_vel": "mimic.future_vel",
+            "mimic_ref_ang_vel": "mimic.future_ang_vel",
+            "anchor_body_index": "mimic.anchor_idx"
+          },
+          "params": {
+            "w_last": true,
+            "future_steps": null,
+            "body_indices": null,
+            "with_velocities": true,
+            "with_relative": true
+          }
+        },
+        "masked_mimic_target_poses": {
+          "$kernel": "protomotions.envs.obs.masked_mimic.compute_target_poses_only",
+          "bindings": {
+            "current_state_body_pos": "current.rigid_body_pos",
+            "current_state_body_rot": "current.rigid_body_rot",
+            "masked_mimic_ref_pos": "masked_mimic.ref_pos",
+            "masked_mimic_ref_rot": "masked_mimic.ref_rot",
+            "masked_mimic_target_bodies_masks": "masked_mimic.target_bodies_masks",
+            "anchor_body_index": "mimic.anchor_idx",
+            "anchor_translation_mask": "masked_mimic.anchor_translation_mask",
+            "anchor_rotation_mask": "masked_mimic.anchor_rotation_mask"
+          },
+          "params": {
+            "future_steps": null,
+            "ground_independently": false,
+            "conditionable_body_ids": {
+              "$type": "protomotions.envs.body_selection.BodySelection",
+              "names": [
+                "torso_link",
+                "right_ankle_roll_link",
+                "left_ankle_roll_link",
+                "left_rubber_hand",
+                "right_rubber_hand"
+              ],
+              "groups": [],
+              "pattern": null
+            },
+            "include_root_relative": false
+          }
+        },
+        "masked_mimic_local_target_poses": {
+          "$kernel": "protomotions.envs.obs.masked_mimic.compute_local_target_poses",
+          "bindings": {
+            "masked_mimic_ref_pos": "masked_mimic.local_ref_pos",
+            "masked_mimic_ref_rot": "masked_mimic.local_ref_rot",
+            "masked_mimic_target_bodies_masks": "masked_mimic.local_target_bodies_masks",
+            "masked_mimic_target_poses_masks": "masked_mimic.local_target_poses_masks",
+            "anchor_body_index": "mimic.anchor_idx"
+          },
+          "params": {
+            "future_steps": null,
+            "conditionable_body_ids": {
+              "$type": "protomotions.envs.body_selection.BodySelection",
+              "names": [
+                "torso_link",
+                "right_ankle_roll_link",
+                "left_ankle_roll_link",
+                "left_rubber_hand",
+                "right_rubber_hand"
+              ],
+              "groups": [],
+              "pattern": null
+            }
+          }
+        },
+        "previous_actions": {
+          "$kernel": "protomotions.envs.obs.cast.to_float",
+          "bindings": {
+            "x": "current_action"
+          },
+          "params": {}
+        },
+        "scene_box_corners": {
+          "$kernel": "examples.inference.metatron.empty_scene_corners",
+          "bindings": {
+            "dof_pos": "current.dof_pos"
+          },
+          "params": {}
+        },
+        "scene_box_mask": {
+          "$kernel": "examples.inference.metatron.empty_scene_mask",
+          "bindings": {
+            "dof_pos": "current.dof_pos"
+          },
+          "params": {}
+        },
+        "fk_max_coords_obs": {
+          "$kernel_object": "protomotions.utils.baked_fk.BakedAnchorFK",
+          "init": {
+            "anchor_idx": 16,
+            "trackable_body_indices": [
+              16,
+              13,
+              7,
+              24,
+              32
+            ]
+          },
+          "bindings": {
+            "dof_pos": "noisy.dof_pos",
+            "current_anchor_rot": "noisy.anchor_rot"
+          },
+          "params": {}
+        },
+        "reduced_coords_obs": {
+          "$kernel": "protomotions.envs.obs.humanoid.compute_humanoid_reduced_coords_observations",
+          "bindings": {
+            "dof_pos": "noisy.dof_pos",
+            "dof_vel": "noisy.dof_vel",
+            "anchor_rot": "noisy.anchor_rot",
+            "root_local_ang_vel": "noisy.root_local_ang_vel"
+          },
+          "params": {
+            "w_last": true,
+            "root_rot": null,
+            "root_pos": null,
+            "root_vel": null,
+            "ground_height": null,
+            "root_height_obs": false,
+            "root_vel_obs": false
+          }
+        },
+        "last_action": {
+          "$kernel": "protomotions.envs.obs.cast.to_float",
+          "bindings": {
+            "x": "current_action"
+          },
+          "params": {}
+        },
+        "historical_proprioception": {
+          "$kernel": "protomotions.envs.obs.humanoid_historical.compute_historical_proprioception",
+          "bindings": {
+            "historical_dof_pos": "noisy_historical.dof_pos",
+            "historical_dof_vel": "noisy_historical.dof_vel",
+            "historical_root_local_ang_vel": "noisy_historical.root_local_ang_vel",
+            "historical_anchor_rot": "noisy_historical.anchor_rot",
+            "historical_actions": "noisy_historical.raw_actions",
+            "dt": "dt"
+          },
+          "params": {
+            "time_origin": "current",
+            "fk": {
+              "$kernel_object": "protomotions.utils.baked_fk.BakedAnchorFK",
+              "init": {
+                "anchor_idx": 16,
+                "trackable_body_indices": [
+                  16,
+                  13,
+                  7,
+                  24,
+                  32
+                ]
+              }
+            },
+            "history_steps": [
+              8,
+              4,
+              3,
+              2,
+              1
+            ]
+          }
+        },
+        "fk_masked_mimic_target_poses": {
+          "$kernel": "protomotions.envs.obs.masked_mimic.compute_target_poses_only",
+          "bindings": {
+            "current_state_body_pos": "noisy.rigid_body_pos",
+            "current_state_body_rot": "noisy.rigid_body_rot",
+            "masked_mimic_ref_pos": "masked_mimic.ref_pos",
+            "masked_mimic_ref_rot": "masked_mimic.ref_rot",
+            "masked_mimic_target_bodies_masks": "masked_mimic.target_bodies_masks",
+            "anchor_body_index": "mimic.anchor_idx",
+            "anchor_translation_mask": "masked_mimic.anchor_translation_mask",
+            "anchor_rotation_mask": "masked_mimic.anchor_rotation_mask"
+          },
+          "params": {
+            "future_steps": null,
+            "conditionable_body_ids": {
+              "$tensor": [
+                16,
+                13,
+                7,
+                24,
+                32
+              ],
+              "dtype": "int64"
+            },
+            "ground_independently": true,
+            "include_root_relative": false
+          }
+        },
+        "masked_mimic_velocity_command": {
+          "$kernel": "protomotions.envs.obs.masked_mimic_velocity_command.compute_masked_mimic_velocity_command_obs",
+          "bindings": {
+            "anchor_rot": "noisy.anchor_rot",
+            "average_velocity": "masked_mimic.average_velocity",
+            "velocity_command_mask": "masked_mimic.velocity_command_mask"
+          },
+          "params": {}
+        },
+        "masked_mimic_velocity_command_mask": {
+          "$kernel": "protomotions.envs.obs.cast.to_float",
+          "bindings": {
+            "x": "masked_mimic.velocity_command_mask"
+          },
+          "params": {}
+        },
+        "masked_mimic_style_poses": {
+          "$kernel": "protomotions.envs.obs.cast.to_float",
+          "bindings": {
+            "x": "masked_mimic.style_pose_samples"
+          },
+          "params": {}
+        },
+        "masked_mimic_style_velocities": {
+          "$kernel": "protomotions.envs.obs.cast.to_float",
+          "bindings": {
+            "x": "masked_mimic.style_velocity_samples"
+          },
+          "params": {}
+        },
+        "masked_mimic_style_masks": {
+          "$kernel": "protomotions.envs.obs.cast.to_float",
+          "bindings": {
+            "x": "masked_mimic.style_condition_masks"
+          },
+          "params": {}
+        },
+        "masked_mimic_style_frame_masks": {
+          "$kernel": "protomotions.envs.obs.cast.to_float",
+          "bindings": {
+            "x": "masked_mimic.style_frame_masks"
+          },
+          "params": {}
+        },
+        "masked_mimic_target_masks": {
+          "$kernel": "protomotions.envs.obs.masked_mimic.compute_target_masks_only",
+          "bindings": {
+            "masked_mimic_target_bodies_masks": "masked_mimic.target_bodies_masks",
+            "anchor_translation_mask": "masked_mimic.anchor_translation_mask",
+            "anchor_rotation_mask": "masked_mimic.anchor_rotation_mask"
+          },
+          "params": {
+            "future_steps": null,
+            "conditionable_body_ids": {
+              "$type": "protomotions.envs.body_selection.BodySelection",
+              "names": [
+                "torso_link",
+                "right_ankle_roll_link",
+                "left_ankle_roll_link",
+                "left_rubber_hand",
+                "right_rubber_hand"
+              ],
+              "groups": [],
+              "pattern": null
+            }
+          }
+        },
+        "masked_mimic_target_times": {
+          "$kernel": "protomotions.envs.obs.masked_mimic.compute_target_time_offsets",
+          "bindings": {
+            "masked_mimic_time_offsets": "masked_mimic.time_offsets"
+          },
+          "params": {
+            "future_steps": null
+          }
+        },
+        "masked_mimic_target_poses_masks": {
+          "$kernel": "protomotions.envs.obs.cast.to_float",
+          "bindings": {
+            "x": "masked_mimic.target_poses_masks"
+          },
+          "params": {}
+        },
+        "masked_mimic_local_target_masks": {
+          "$kernel": "protomotions.envs.obs.masked_mimic.compute_target_masks_only",
+          "bindings": {
+            "masked_mimic_target_bodies_masks": "masked_mimic.local_target_bodies_masks"
+          },
+          "params": {
+            "future_steps": null,
+            "anchor_translation_mask": null,
+            "anchor_rotation_mask": null,
+            "conditionable_body_ids": {
+              "$type": "protomotions.envs.body_selection.BodySelection",
+              "names": [
+                "torso_link",
+                "right_ankle_roll_link",
+                "left_ankle_roll_link",
+                "left_rubber_hand",
+                "right_rubber_hand"
+              ],
+              "groups": [],
+              "pattern": null
+            }
+          }
+        },
+        "masked_mimic_local_target_times": {
+          "$kernel": "protomotions.envs.obs.masked_mimic.compute_target_time_offsets",
+          "bindings": {
+            "masked_mimic_time_offsets": "masked_mimic.local_time_offsets"
+          },
+          "params": {
+            "future_steps": null
+          }
+        },
+        "masked_mimic_local_target_poses_masks": {
+          "$kernel": "protomotions.envs.obs.cast.to_float",
+          "bindings": {
+            "x": "masked_mimic.local_target_poses_masks"
+          },
+          "params": {}
+        }
+      },
+      "action_config": {
+        "$type": "protomotions.envs.action.config.ActionConfig",
+        "fn": {
+          "$function": "protomotions.envs.action.action_functions.bm_pd_action"
+        },
+        "params": {
+          "pd_action_offset": {
+            "$tensor": [
+              -0.31200000643730164,
+              0.0,
+              0.0,
+              0.6690000295639038,
+              -0.3630000054836273,
+              0.0,
+              -0.31200000643730164,
+              0.0,
+              0.0,
+              0.6690000295639038,
+              -0.3630000054836273,
+              0.0,
+              0.0,
+              0.0,
+              0.0,
+              0.20000000298023224,
+              0.20000000298023224,
+              0.0,
+              0.6000000238418579,
+              0.0,
+              0.0,
+              0.0,
+              0.20000000298023224,
+              -0.20000000298023224,
+              0.0,
+              0.6000000238418579,
+              0.0,
+              0.0,
+              0.0
+            ],
+            "dtype": "float32"
+          },
+          "action_scale": {
+            "$tensor": [
+              0.5475465059280396,
+              0.3506614565849304,
+              0.5475465059280396,
+              0.3506614565849304,
+              0.4385773241519928,
+              0.4385773241519928,
+              0.5475465059280396,
+              0.3506614565849304,
+              0.5475465059280396,
+              0.3506614565849304,
+              0.4385773241519928,
+              0.4385773241519928,
+              0.5475465059280396,
+              0.4385773241519928,
+              0.4385773241519928,
+              0.4385773241519928,
+              0.4385773241519928,
+              0.4385773241519928,
+              0.4385773241519928,
+              0.4385773241519928,
+              0.07450086623430252,
+              0.07450086623430252,
+              0.4385773241519928,
+              0.4385773241519928,
+              0.4385773241519928,
+              0.4385773241519928,
+              0.4385773241519928,
+              0.07450086623430252,
+              0.07450086623430252
+            ],
+            "dtype": "float32"
+          },
+          "stiffness": {
+            "$tensor": [
+              40.179237365722656,
+              99.09842681884766,
+              40.179237365722656,
+              99.09842681884766,
+              28.501245498657227,
+              28.501245498657227,
+              40.179237365722656,
+              99.09842681884766,
+              40.179237365722656,
+              99.09842681884766,
+              28.501245498657227,
+              28.501245498657227,
+              40.179237365722656,
+              28.501245498657227,
+              28.501245498657227,
+              14.250622749328613,
+              14.250622749328613,
+              14.250622749328613,
+              14.250622749328613,
+              14.250622749328613,
+              16.77832794189453,
+              16.77832794189453,
+              14.250622749328613,
+              14.250622749328613,
+              14.250622749328613,
+              14.250622749328613,
+              14.250622749328613,
+              16.77832794189453,
+              16.77832794189453
+            ],
+            "dtype": "float32"
+          },
+          "damping": {
+            "$tensor": [
+              2.557889699935913,
+              6.308801651000977,
+              2.557889699935913,
+              6.308801651000977,
+              1.8144457340240479,
+              1.8144457340240479,
+              2.557889699935913,
+              6.308801651000977,
+              2.557889699935913,
+              6.308801651000977,
+              1.8144457340240479,
+              1.8144457340240479,
+              2.557889699935913,
+              1.8144457340240479,
+              1.8144457340240479,
+              0.9072228670120239,
+              0.9072228670120239,
+              0.9072228670120239,
+              0.9072228670120239,
+              0.9072228670120239,
+              1.0681414604187012,
+              1.0681414604187012,
+              0.9072228670120239,
+              0.9072228670120239,
+              0.9072228670120239,
+              0.9072228670120239,
+              0.9072228670120239,
+              1.0681414604187012,
+              1.0681414604187012
+            ],
+            "dtype": "float32"
+          }
+        }
+      },
+      "sensors": {
+        "$type": "protomotions.envs.sensors.config.SensorConfig",
+        "joint_encoders": {
+          "$type": "protomotions.envs.sensors.config.JointEncoderSensor",
+          "pos_noise": {
+            "$type": "protomotions.envs.sensors.config.NoNoise",
+            "reason": "inference"
+          },
+          "vel_noise": {
+            "$type": "protomotions.envs.sensors.config.NoNoise",
+            "reason": "inference"
+          }
+        },
+        "imu": {
+          "$type": "protomotions.envs.sensors.config.IMUSensor",
+          "quat_noise": {
+            "$type": "protomotions.envs.sensors.config.NoNoise",
+            "reason": "inference"
+          },
+          "gyro_noise": {
+            "$type": "protomotions.envs.sensors.config.NoNoise",
+            "reason": "inference"
+          }
+        },
+        "localization": {
+          "$type": "protomotions.envs.sensors.config.LocalizationSensor",
+          "body_name": "pelvis",
+          "noise": {
+            "$type": "protomotions.envs.sensors.config.NoNoise",
+            "reason": "inference"
+          },
+          "offset": [
+            0.0,
+            0.0,
+            0.0
+          ],
+          "drift": null
+        },
+        "base_velocity": null,
+        "height": null,
+        "seed": null
+      }
+    },
+    "agent": {
+      "$type": "protomotions.agents.supervised.masked_mimic_config.MaskedMimicSupervisedAgentConfig",
+      "batch_size": 4,
+      "training_max_steps": 10000000000000,
+      "training_max_iterations": null,
+      "_target_": "protomotions.agents.supervised.agent.SupervisedAgent",
+      "model": {
+        "$type": "protomotions.agents.supervised.masked_mimic_config.MaskedMimicModelConfig",
+        "_target_": "protomotions.agents.supervised.masked_mimic_model.MaskedMimicModel",
+        "in_keys": [],
+        "out_keys": [],
+        "encoder": {
+          "$type": "protomotions.agents.common.config.ModuleContainerConfig",
+          "models": [
+            {
+              "$type": "protomotions.agents.common.config.ObsProcessorConfig",
+              "normalize_obs": true,
+              "norm_clamp_value": 5.0,
+              "norm_ema_decay": 0.999,
+              "_target_": "protomotions.agents.common.common.ObsProcessor",
+              "in_keys": [
+                "max_coords_obs"
+              ],
+              "out_keys": [
+                "max_coords_obs_norm"
+              ],
+              "module_operations": [
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationForwardConfig"
+                }
+              ]
+            },
+            {
+              "$type": "protomotions.agents.common.config.ObsProcessorConfig",
+              "normalize_obs": true,
+              "norm_clamp_value": 5.0,
+              "norm_ema_decay": 0.999,
+              "_target_": "protomotions.agents.common.common.ObsProcessor",
+              "in_keys": [
+                "mimic_target_poses"
+              ],
+              "out_keys": [
+                "mimic_target_poses_norm"
+              ],
+              "module_operations": [
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationForwardConfig"
+                }
+              ]
+            },
+            {
+              "$type": "protomotions.agents.common.config.ObsProcessorConfig",
+              "normalize_obs": true,
+              "norm_clamp_value": 5.0,
+              "norm_ema_decay": 0.999,
+              "_target_": "protomotions.agents.common.common.ObsProcessor",
+              "in_keys": [
+                "masked_mimic_local_target_poses"
+              ],
+              "out_keys": [
+                "masked_mimic_local_target_poses_norm"
+              ],
+              "module_operations": [
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationForwardConfig"
+                }
+              ]
+            },
+            {
+              "$type": "protomotions.agents.common.config.ObsProcessorConfig",
+              "normalize_obs": true,
+              "norm_clamp_value": 5.0,
+              "norm_ema_decay": 0.999,
+              "_target_": "protomotions.agents.common.common.ObsProcessor",
+              "in_keys": [
+                "masked_mimic_target_poses"
+              ],
+              "out_keys": [
+                "masked_mimic_global_target_poses_norm"
+              ],
+              "module_operations": [
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationForwardConfig"
+                }
+              ]
+            },
+            {
+              "$type": "protomotions.agents.common.config.ObsProcessorConfig",
+              "normalize_obs": true,
+              "norm_clamp_value": 5.0,
+              "norm_ema_decay": 0.999,
+              "_target_": "protomotions.agents.common.common.ObsProcessor",
+              "in_keys": [
+                "masked_mimic_target_times"
+              ],
+              "out_keys": [
+                "masked_mimic_target_times_norm"
+              ],
+              "module_operations": [
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationForwardConfig"
+                }
+              ]
+            },
+            {
+              "$type": "protomotions.agents.common.config.ObsProcessorConfig",
+              "normalize_obs": true,
+              "norm_clamp_value": 5.0,
+              "norm_ema_decay": 0.999,
+              "_target_": "protomotions.agents.common.common.ObsProcessor",
+              "in_keys": [
+                "masked_mimic_local_target_times"
+              ],
+              "out_keys": [
+                "masked_mimic_local_target_times_norm"
+              ],
+              "module_operations": [
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationForwardConfig"
+                }
+              ]
+            },
+            {
+              "$type": "protomotions.agents.common.config.ObsProcessorConfig",
+              "normalize_obs": true,
+              "norm_clamp_value": 5.0,
+              "norm_ema_decay": 0.999,
+              "_target_": "protomotions.agents.common.common.ObsProcessor",
+              "in_keys": [
+                "masked_mimic_style_poses"
+              ],
+              "out_keys": [
+                "masked_mimic_style_poses_norm"
+              ],
+              "module_operations": [
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationForwardConfig"
+                }
+              ]
+            },
+            {
+              "$type": "protomotions.agents.common.config.ObsProcessorConfig",
+              "normalize_obs": true,
+              "norm_clamp_value": 5.0,
+              "norm_ema_decay": 0.999,
+              "_target_": "protomotions.agents.common.common.ObsProcessor",
+              "in_keys": [
+                "masked_mimic_style_velocities"
+              ],
+              "out_keys": [
+                "masked_mimic_style_velocities_norm"
+              ],
+              "module_operations": [
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationForwardConfig"
+                }
+              ]
+            },
+            {
+              "$type": "protomotions.agents.common.config.ObsProcessorConfig",
+              "normalize_obs": false,
+              "norm_clamp_value": 5.0,
+              "norm_ema_decay": 0.999,
+              "_target_": "protomotions.agents.common.common.ObsProcessor",
+              "in_keys": [
+                "masked_mimic_style_masks"
+              ],
+              "out_keys": [
+                "masked_mimic_style_masks_flat"
+              ],
+              "module_operations": [
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationReshapeConfig",
+                  "new_shape": [
+                    "batch_size",
+                    -1
+                  ]
+                }
+              ]
+            },
+            {
+              "$type": "protomotions.agents.common.config.MLPWithConcatConfig",
+              "normalize_obs": false,
+              "norm_clamp_value": 5.0,
+              "norm_ema_decay": 0.999,
+              "num_out": 256,
+              "layers": [
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 512,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                },
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 512,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                }
+              ],
+              "_target_": "protomotions.agents.common.mlp.MLPWithConcat",
+              "in_keys": [
+                "masked_mimic_local_target_poses_norm",
+                "masked_mimic_local_target_masks",
+                "masked_mimic_local_target_times_norm",
+                "masked_mimic_local_target_poses_masks"
+              ],
+              "out_keys": [
+                "encoder_local_target_features"
+              ],
+              "output_activation": null,
+              "module_operations": [
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationForwardConfig"
+                }
+              ]
+            },
+            {
+              "$type": "protomotions.agents.common.config.MLPWithConcatConfig",
+              "normalize_obs": false,
+              "norm_clamp_value": 5.0,
+              "norm_ema_decay": 0.999,
+              "num_out": 256,
+              "layers": [
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 512,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                },
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 512,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                }
+              ],
+              "_target_": "protomotions.agents.common.mlp.MLPWithConcat",
+              "in_keys": [
+                "masked_mimic_global_target_poses_norm",
+                "masked_mimic_target_masks",
+                "masked_mimic_target_times_norm",
+                "masked_mimic_target_poses_masks"
+              ],
+              "out_keys": [
+                "encoder_global_target_features"
+              ],
+              "output_activation": null,
+              "module_operations": [
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationForwardConfig"
+                }
+              ]
+            },
+            {
+              "$type": "protomotions.agents.common.config.MLPWithConcatConfig",
+              "normalize_obs": false,
+              "norm_clamp_value": 5.0,
+              "norm_ema_decay": 0.999,
+              "num_out": 512,
+              "layers": [
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 1024,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                },
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 1024,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                },
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 1024,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                },
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 1024,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                },
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 1024,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                }
+              ],
+              "_target_": "protomotions.agents.common.mlp.MLPWithConcat",
+              "in_keys": [
+                "max_coords_obs_norm",
+                "mimic_target_poses_norm",
+                "encoder_local_target_features",
+                "encoder_global_target_features",
+                "masked_mimic_style_poses_norm",
+                "masked_mimic_style_velocities_norm",
+                "masked_mimic_style_masks_flat",
+                "masked_mimic_style_frame_masks",
+                "masked_mimic_velocity_command",
+                "masked_mimic_velocity_command_mask"
+              ],
+              "out_keys": [
+                "encoder_trunk_out"
+              ],
+              "output_activation": "silu",
+              "module_operations": [
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationForwardConfig"
+                }
+              ]
+            },
+            {
+              "$type": "protomotions.agents.common.config.MLPWithConcatConfig",
+              "normalize_obs": false,
+              "norm_clamp_value": 5.0,
+              "norm_ema_decay": 0.999,
+              "num_out": 64,
+              "layers": [
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 512,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                },
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 256,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                },
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 128,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                }
+              ],
+              "_target_": "protomotions.agents.common.mlp.MLPWithConcat",
+              "in_keys": [
+                "encoder_trunk_out"
+              ],
+              "out_keys": [
+                "encoder_mu"
+              ],
+              "output_activation": null,
+              "module_operations": [
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationForwardConfig"
+                }
+              ]
+            },
+            {
+              "$type": "protomotions.agents.common.config.MLPWithConcatConfig",
+              "normalize_obs": false,
+              "norm_clamp_value": 5.0,
+              "norm_ema_decay": 0.999,
+              "num_out": 64,
+              "layers": [
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 512,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                },
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 256,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                },
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 128,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                }
+              ],
+              "_target_": "protomotions.agents.common.mlp.MLPWithConcat",
+              "in_keys": [
+                "encoder_trunk_out"
+              ],
+              "out_keys": [
+                "encoder_logvar"
+              ],
+              "output_activation": null,
+              "module_operations": [
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationForwardConfig"
+                }
+              ]
+            }
+          ],
+          "_target_": "protomotions.agents.common.common.ModuleContainer",
+          "in_keys": [
+            "max_coords_obs",
+            "mimic_target_poses",
+            "masked_mimic_local_target_poses",
+            "masked_mimic_target_poses",
+            "masked_mimic_target_times",
+            "masked_mimic_local_target_times",
+            "masked_mimic_style_poses",
+            "masked_mimic_style_velocities",
+            "masked_mimic_style_masks",
+            "masked_mimic_local_target_masks",
+            "masked_mimic_local_target_poses_masks",
+            "masked_mimic_target_masks",
+            "masked_mimic_target_poses_masks",
+            "masked_mimic_style_frame_masks",
+            "masked_mimic_velocity_command",
+            "masked_mimic_velocity_command_mask"
+          ],
+          "out_keys": [
+            "encoder_mu",
+            "encoder_logvar"
+          ]
+        },
+        "prior": {
+          "$type": "protomotions.agents.common.config.ModuleContainerConfig",
+          "models": [
+            {
+              "$type": "protomotions.agents.common.config.ObsProcessorConfig",
+              "normalize_obs": true,
+              "norm_clamp_value": 5.0,
+              "norm_ema_decay": 0.999,
+              "_target_": "protomotions.agents.common.common.ObsProcessor",
+              "in_keys": [
+                "scene_box_corners"
+              ],
+              "out_keys": [
+                "scene_box_norm"
+              ],
+              "module_operations": [
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationForwardConfig"
+                }
+              ]
+            },
+            {
+              "$type": "protomotions.agents.common.config.MLPWithConcatConfig",
+              "normalize_obs": false,
+              "norm_clamp_value": 5.0,
+              "norm_ema_decay": 0.999,
+              "num_out": 1024,
+              "layers": [
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 1024,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                }
+              ],
+              "_target_": "protomotions.agents.common.mlp.MLPWithConcat",
+              "in_keys": [
+                "scene_box_norm"
+              ],
+              "out_keys": [
+                "scene_box_tokens"
+              ],
+              "output_activation": null,
+              "module_operations": [
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationForwardConfig"
+                }
+              ]
+            },
+            {
+              "$type": "protomotions.agents.common.config.ObsProcessorConfig",
+              "normalize_obs": true,
+              "norm_clamp_value": 5.0,
+              "norm_ema_decay": 0.999,
+              "_target_": "protomotions.agents.common.common.ObsProcessor",
+              "in_keys": [
+                "previous_actions"
+              ],
+              "out_keys": [
+                "prior_previous_actions_norm"
+              ],
+              "module_operations": [
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationForwardConfig"
+                }
+              ]
+            },
+            {
+              "$type": "protomotions.agents.common.config.ObsProcessorConfig",
+              "normalize_obs": true,
+              "norm_clamp_value": 5.0,
+              "norm_ema_decay": 0.999,
+              "_target_": "protomotions.agents.common.common.ObsProcessor",
+              "in_keys": [
+                "reduced_coords_obs"
+              ],
+              "out_keys": [
+                "prior_reduced_coords_norm"
+              ],
+              "module_operations": [
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationForwardConfig"
+                }
+              ]
+            },
+            {
+              "$type": "protomotions.agents.common.config.ObsProcessorConfig",
+              "normalize_obs": true,
+              "norm_clamp_value": 5.0,
+              "norm_ema_decay": 0.999,
+              "_target_": "protomotions.agents.common.common.ObsProcessor",
+              "in_keys": [
+                "masked_mimic_local_target_poses"
+              ],
+              "out_keys": [
+                "local_target_poses_seq"
+              ],
+              "module_operations": [
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationReshapeConfig",
+                  "new_shape": [
+                    "batch_size",
+                    5,
+                    -1
+                  ]
+                },
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationForwardConfig"
+                }
+              ]
+            },
+            {
+              "$type": "protomotions.agents.common.config.ObsProcessorConfig",
+              "normalize_obs": true,
+              "norm_clamp_value": 5.0,
+              "norm_ema_decay": 0.999,
+              "_target_": "protomotions.agents.common.common.ObsProcessor",
+              "in_keys": [
+                "fk_masked_mimic_target_poses"
+              ],
+              "out_keys": [
+                "global_target_poses_seq"
+              ],
+              "module_operations": [
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationReshapeConfig",
+                  "new_shape": [
+                    "batch_size",
+                    5,
+                    -1
+                  ]
+                },
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationForwardConfig"
+                }
+              ]
+            },
+            {
+              "$type": "protomotions.agents.common.config.ObsProcessorConfig",
+              "normalize_obs": false,
+              "norm_clamp_value": 5.0,
+              "norm_ema_decay": 0.999,
+              "_target_": "protomotions.agents.common.common.ObsProcessor",
+              "in_keys": [
+                "masked_mimic_local_target_masks"
+              ],
+              "out_keys": [
+                "local_target_masks_seq"
+              ],
+              "module_operations": [
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationReshapeConfig",
+                  "new_shape": [
+                    "batch_size",
+                    5,
+                    -1
+                  ]
+                }
+              ]
+            },
+            {
+              "$type": "protomotions.agents.common.config.ObsProcessorConfig",
+              "normalize_obs": false,
+              "norm_clamp_value": 5.0,
+              "norm_ema_decay": 0.999,
+              "_target_": "protomotions.agents.common.common.ObsProcessor",
+              "in_keys": [
+                "masked_mimic_target_masks"
+              ],
+              "out_keys": [
+                "global_target_masks_seq"
+              ],
+              "module_operations": [
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationReshapeConfig",
+                  "new_shape": [
+                    "batch_size",
+                    5,
+                    -1
+                  ]
+                }
+              ]
+            },
+            {
+              "$type": "protomotions.agents.common.config.ObsProcessorConfig",
+              "normalize_obs": true,
+              "norm_clamp_value": 5.0,
+              "norm_ema_decay": 0.999,
+              "_target_": "protomotions.agents.common.common.ObsProcessor",
+              "in_keys": [
+                "masked_mimic_target_times"
+              ],
+              "out_keys": [
+                "target_times_seq"
+              ],
+              "module_operations": [
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationReshapeConfig",
+                  "new_shape": [
+                    "batch_size",
+                    5,
+                    -1
+                  ]
+                },
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationForwardConfig"
+                }
+              ]
+            },
+            {
+              "$type": "protomotions.agents.common.config.ObsProcessorConfig",
+              "normalize_obs": true,
+              "norm_clamp_value": 5.0,
+              "norm_ema_decay": 0.999,
+              "_target_": "protomotions.agents.common.common.ObsProcessor",
+              "in_keys": [
+                "masked_mimic_local_target_times"
+              ],
+              "out_keys": [
+                "local_target_times_seq"
+              ],
+              "module_operations": [
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationReshapeConfig",
+                  "new_shape": [
+                    "batch_size",
+                    5,
+                    -1
+                  ]
+                },
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationForwardConfig"
+                }
+              ]
+            },
+            {
+              "$type": "protomotions.agents.common.config.ObsProcessorConfig",
+              "normalize_obs": true,
+              "norm_clamp_value": 5.0,
+              "norm_ema_decay": 0.999,
+              "_target_": "protomotions.agents.common.common.ObsProcessor",
+              "in_keys": [
+                "historical_proprioception"
+              ],
+              "out_keys": [
+                "historical_pose_obs_seq"
+              ],
+              "module_operations": [
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationReshapeConfig",
+                  "new_shape": [
+                    "batch_size",
+                    5,
+                    -1
+                  ]
+                },
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationForwardConfig"
+                }
+              ]
+            },
+            {
+              "$type": "protomotions.agents.common.config.ObsProcessorConfig",
+              "normalize_obs": true,
+              "norm_clamp_value": 5.0,
+              "norm_ema_decay": 0.999,
+              "_target_": "protomotions.agents.common.common.ObsProcessor",
+              "in_keys": [
+                "masked_mimic_style_poses"
+              ],
+              "out_keys": [
+                "style_poses_seq"
+              ],
+              "module_operations": [
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationReshapeConfig",
+                  "new_shape": [
+                    "batch_size",
+                    5,
+                    -1
+                  ]
+                },
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationForwardConfig"
+                }
+              ]
+            },
+            {
+              "$type": "protomotions.agents.common.config.ObsProcessorConfig",
+              "normalize_obs": true,
+              "norm_clamp_value": 5.0,
+              "norm_ema_decay": 0.999,
+              "_target_": "protomotions.agents.common.common.ObsProcessor",
+              "in_keys": [
+                "masked_mimic_style_velocities"
+              ],
+              "out_keys": [
+                "style_velocities_seq"
+              ],
+              "module_operations": [
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationReshapeConfig",
+                  "new_shape": [
+                    "batch_size",
+                    5,
+                    -1
+                  ]
+                },
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationForwardConfig"
+                }
+              ]
+            },
+            {
+              "$type": "protomotions.agents.common.config.ObsProcessorConfig",
+              "normalize_obs": false,
+              "norm_clamp_value": 5.0,
+              "norm_ema_decay": 0.999,
+              "_target_": "protomotions.agents.common.common.ObsProcessor",
+              "in_keys": [
+                "masked_mimic_style_masks"
+              ],
+              "out_keys": [
+                "style_condition_masks_seq"
+              ],
+              "module_operations": [
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationReshapeConfig",
+                  "new_shape": [
+                    "batch_size",
+                    5,
+                    -1
+                  ]
+                }
+              ]
+            },
+            {
+              "$type": "protomotions.agents.common.config.ObsProcessorConfig",
+              "normalize_obs": false,
+              "norm_clamp_value": 5.0,
+              "norm_ema_decay": 0.999,
+              "_target_": "protomotions.agents.common.common.ObsProcessor",
+              "in_keys": [
+                "masked_mimic_style_frame_masks"
+              ],
+              "out_keys": [
+                "style_frame_masks_seq"
+              ],
+              "module_operations": [
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationReshapeConfig",
+                  "new_shape": [
+                    "batch_size",
+                    5,
+                    -1
+                  ]
+                }
+              ]
+            },
+            {
+              "$type": "protomotions.agents.common.config.ObsProcessorConfig",
+              "normalize_obs": true,
+              "norm_clamp_value": 5.0,
+              "norm_ema_decay": 0.999,
+              "_target_": "protomotions.agents.common.common.ObsProcessor",
+              "in_keys": [
+                "fk_max_coords_obs"
+              ],
+              "out_keys": [
+                "prior_fk_max_coords_obs_norm"
+              ],
+              "module_operations": [
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationForwardConfig"
+                }
+              ]
+            },
+            {
+              "$type": "protomotions.agents.common.config.MLPWithConcatConfig",
+              "normalize_obs": false,
+              "norm_clamp_value": 5.0,
+              "norm_ema_decay": 0.999,
+              "num_out": 1024,
+              "layers": [
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 1024,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                },
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 1024,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                },
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 1024,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                }
+              ],
+              "_target_": "protomotions.agents.common.mlp.MLPWithConcat",
+              "in_keys": [
+                "prior_fk_max_coords_obs_norm",
+                "prior_reduced_coords_norm",
+                "prior_previous_actions_norm"
+              ],
+              "out_keys": [
+                "current_state_token"
+              ],
+              "output_activation": null,
+              "module_operations": [
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationReshapeConfig",
+                  "new_shape": [
+                    "batch_size",
+                    1,
+                    -1
+                  ]
+                },
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationForwardConfig"
+                }
+              ]
+            },
+            {
+              "$type": "protomotions.agents.common.config.MLPWithConcatConfig",
+              "normalize_obs": false,
+              "norm_clamp_value": 5.0,
+              "norm_ema_decay": 0.999,
+              "num_out": 1024,
+              "layers": [
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 1024,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                },
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 1024,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                },
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 1024,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                }
+              ],
+              "_target_": "protomotions.agents.common.mlp.MLPWithConcat",
+              "in_keys": [
+                "masked_mimic_velocity_command",
+                "masked_mimic_velocity_command_mask"
+              ],
+              "out_keys": [
+                "masked_mimic_velocity_command_token"
+              ],
+              "output_activation": null,
+              "module_operations": [
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationReshapeConfig",
+                  "new_shape": [
+                    "batch_size",
+                    1,
+                    -1
+                  ]
+                },
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationForwardConfig"
+                }
+              ]
+            },
+            {
+              "$type": "protomotions.agents.common.config.MLPWithConcatConfig",
+              "normalize_obs": false,
+              "norm_clamp_value": 5.0,
+              "norm_ema_decay": 0.999,
+              "num_out": 1024,
+              "layers": [
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 1024,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                },
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 1024,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                },
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 1024,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                }
+              ],
+              "_target_": "protomotions.agents.common.mlp.MLPWithConcat",
+              "in_keys": [
+                "style_poses_seq",
+                "style_velocities_seq",
+                "style_condition_masks_seq"
+              ],
+              "out_keys": [
+                "masked_mimic_style_frame_tokens"
+              ],
+              "output_activation": null,
+              "module_operations": [
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationReshapeConfig",
+                  "new_shape": [
+                    "batch_size",
+                    5,
+                    -1
+                  ]
+                },
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationForwardConfig"
+                }
+              ]
+            },
+            {
+              "$type": "protomotions.agents.common.config.MaskedMeanConfig",
+              "_target_": "protomotions.agents.common.common.MaskedMean",
+              "in_keys": [
+                "masked_mimic_style_frame_tokens",
+                "style_frame_masks_seq"
+              ],
+              "out_keys": [
+                "masked_mimic_style_token"
+              ]
+            },
+            {
+              "$type": "protomotions.agents.common.config.MLPWithConcatConfig",
+              "normalize_obs": false,
+              "norm_clamp_value": 5.0,
+              "norm_ema_decay": 0.999,
+              "num_out": 1024,
+              "layers": [
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 1024,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                },
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 1024,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                },
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 1024,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                }
+              ],
+              "_target_": "protomotions.agents.common.mlp.MLPWithConcat",
+              "in_keys": [
+                "local_target_poses_seq",
+                "local_target_masks_seq",
+                "local_target_times_seq"
+              ],
+              "out_keys": [
+                "masked_mimic_local_target_token"
+              ],
+              "output_activation": null,
+              "module_operations": [
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationReshapeConfig",
+                  "new_shape": [
+                    "batch_size",
+                    5,
+                    -1
+                  ]
+                },
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationForwardConfig"
+                }
+              ]
+            },
+            {
+              "$type": "protomotions.agents.common.config.MLPWithConcatConfig",
+              "normalize_obs": false,
+              "norm_clamp_value": 5.0,
+              "norm_ema_decay": 0.999,
+              "num_out": 1024,
+              "layers": [
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 1024,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                },
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 1024,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                },
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 1024,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                }
+              ],
+              "_target_": "protomotions.agents.common.mlp.MLPWithConcat",
+              "in_keys": [
+                "global_target_poses_seq",
+                "global_target_masks_seq",
+                "target_times_seq"
+              ],
+              "out_keys": [
+                "masked_mimic_global_target_token"
+              ],
+              "output_activation": null,
+              "module_operations": [
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationReshapeConfig",
+                  "new_shape": [
+                    "batch_size",
+                    5,
+                    -1
+                  ]
+                },
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationForwardConfig"
+                }
+              ]
+            },
+            {
+              "$type": "protomotions.agents.common.config.MLPWithConcatConfig",
+              "normalize_obs": false,
+              "norm_clamp_value": 5.0,
+              "norm_ema_decay": 0.999,
+              "num_out": 1024,
+              "layers": [
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 1024,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                },
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 1024,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                },
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 1024,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                }
+              ],
+              "_target_": "protomotions.agents.common.mlp.MLPWithConcat",
+              "in_keys": [
+                "historical_pose_obs_seq"
+              ],
+              "out_keys": [
+                "historical_pose_obs_token"
+              ],
+              "output_activation": null,
+              "module_operations": [
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationReshapeConfig",
+                  "new_shape": [
+                    "batch_size",
+                    5,
+                    -1
+                  ]
+                },
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationForwardConfig"
+                }
+              ]
+            },
+            {
+              "$type": "protomotions.agents.common.config.TransformerConfig",
+              "_target_": "protomotions.agents.common.transformer.Transformer",
+              "in_keys": [
+                "current_state_token",
+                "masked_mimic_velocity_command_token",
+                "masked_mimic_local_target_token",
+                "masked_mimic_global_target_token",
+                "masked_mimic_style_token",
+                "historical_pose_obs_token",
+                "masked_mimic_velocity_command_mask",
+                "masked_mimic_local_target_poses_masks",
+                "masked_mimic_target_poses_masks",
+                "masked_mimic_style_frame_masks",
+                "scene_box_tokens",
+                "scene_box_mask"
+              ],
+              "out_keys": [
+                "transformer_out"
+              ],
+              "input_and_mask_mapping": {
+                "masked_mimic_velocity_command_token": "masked_mimic_velocity_command_mask",
+                "masked_mimic_local_target_token": "masked_mimic_local_target_poses_masks",
+                "masked_mimic_global_target_token": "masked_mimic_target_poses_masks",
+                "masked_mimic_style_token": "masked_mimic_style_frame_masks",
+                "scene_box_tokens": "scene_box_mask"
+              },
+              "transformer_token_size": 1024,
+              "latent_dim": 1024,
+              "num_heads": 4,
+              "ff_size": 4096,
+              "num_layers": 6,
+              "dropout": 0.0,
+              "activation": "silu",
+              "output_activation": "silu"
+            },
+            {
+              "$type": "protomotions.agents.common.config.MLPWithConcatConfig",
+              "normalize_obs": false,
+              "norm_clamp_value": 5.0,
+              "norm_ema_decay": 0.999,
+              "num_out": 64,
+              "layers": [
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 512,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                },
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 256,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                },
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 128,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                }
+              ],
+              "_target_": "protomotions.agents.common.mlp.MLPWithConcat",
+              "in_keys": [
+                "transformer_out"
+              ],
+              "out_keys": [
+                "decoder_condition"
+              ],
+              "output_activation": null,
+              "module_operations": [
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationForwardConfig"
+                }
+              ]
+            },
+            {
+              "$type": "protomotions.agents.common.config.MLPWithConcatConfig",
+              "normalize_obs": false,
+              "norm_clamp_value": 5.0,
+              "norm_ema_decay": 0.999,
+              "num_out": 64,
+              "layers": [
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 512,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                },
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 256,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                },
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 128,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                }
+              ],
+              "_target_": "protomotions.agents.common.mlp.MLPWithConcat",
+              "in_keys": [
+                "transformer_out"
+              ],
+              "out_keys": [
+                "prior_mu"
+              ],
+              "output_activation": null,
+              "module_operations": [
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationForwardConfig"
+                }
+              ]
+            },
+            {
+              "$type": "protomotions.agents.common.config.MLPWithConcatConfig",
+              "normalize_obs": false,
+              "norm_clamp_value": 5.0,
+              "norm_ema_decay": 0.999,
+              "num_out": 64,
+              "layers": [
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 512,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                },
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 256,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                },
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 128,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                }
+              ],
+              "_target_": "protomotions.agents.common.mlp.MLPWithConcat",
+              "in_keys": [
+                "transformer_out"
+              ],
+              "out_keys": [
+                "prior_logvar"
+              ],
+              "output_activation": null,
+              "module_operations": [
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationForwardConfig"
+                }
+              ]
+            }
+          ],
+          "_target_": "protomotions.agents.common.common.ModuleContainer",
+          "in_keys": [
+            "scene_box_corners",
+            "previous_actions",
+            "reduced_coords_obs",
+            "masked_mimic_local_target_poses",
+            "fk_masked_mimic_target_poses",
+            "masked_mimic_local_target_masks",
+            "masked_mimic_target_masks",
+            "masked_mimic_target_times",
+            "masked_mimic_local_target_times",
+            "historical_proprioception",
+            "masked_mimic_style_poses",
+            "masked_mimic_style_velocities",
+            "masked_mimic_style_masks",
+            "masked_mimic_style_frame_masks",
+            "fk_max_coords_obs",
+            "masked_mimic_velocity_command",
+            "masked_mimic_velocity_command_mask",
+            "masked_mimic_local_target_poses_masks",
+            "masked_mimic_target_poses_masks",
+            "scene_box_mask"
+          ],
+          "out_keys": [
+            "prior_mu",
+            "prior_logvar",
+            "decoder_condition"
+          ]
+        },
+        "trunk": {
+          "$type": "protomotions.agents.common.config.ModuleContainerConfig",
+          "models": [
+            {
+              "$type": "protomotions.agents.common.config.ObsProcessorConfig",
+              "normalize_obs": true,
+              "norm_clamp_value": 5.0,
+              "norm_ema_decay": 0.999,
+              "_target_": "protomotions.agents.common.common.ObsProcessor",
+              "in_keys": [
+                "reduced_coords_obs"
+              ],
+              "out_keys": [
+                "decoder_reduced_coords_norm"
+              ],
+              "module_operations": [
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationForwardConfig"
+                }
+              ]
+            },
+            {
+              "$type": "protomotions.agents.common.config.ObsProcessorConfig",
+              "normalize_obs": true,
+              "norm_clamp_value": 5.0,
+              "norm_ema_decay": 0.999,
+              "_target_": "protomotions.agents.common.common.ObsProcessor",
+              "in_keys": [
+                "fk_max_coords_obs"
+              ],
+              "out_keys": [
+                "fk_max_coords_obs_norm"
+              ],
+              "module_operations": [
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationForwardConfig"
+                }
+              ]
+            },
+            {
+              "$type": "protomotions.agents.common.config.ObsProcessorConfig",
+              "normalize_obs": true,
+              "norm_clamp_value": 5.0,
+              "norm_ema_decay": 0.999,
+              "_target_": "protomotions.agents.common.common.ObsProcessor",
+              "in_keys": [
+                "previous_actions"
+              ],
+              "out_keys": [
+                "previous_actions_norm"
+              ],
+              "module_operations": [
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationForwardConfig"
+                }
+              ]
+            },
+            {
+              "$type": "protomotions.agents.common.config.MLPWithConcatConfig",
+              "normalize_obs": false,
+              "norm_clamp_value": 5.0,
+              "norm_ema_decay": 0.999,
+              "num_out": 64,
+              "layers": [
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 512,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                },
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 256,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                },
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 128,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                }
+              ],
+              "_target_": "protomotions.agents.common.mlp.MLPWithConcat",
+              "in_keys": [
+                "decoder_condition",
+                "vae_latent"
+              ],
+              "out_keys": [
+                "solution_condition_features"
+              ],
+              "output_activation": null,
+              "module_operations": [
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationForwardConfig"
+                }
+              ]
+            },
+            {
+              "$type": "protomotions.agents.common.config.MLPWithConcatConfig",
+              "normalize_obs": false,
+              "norm_clamp_value": 5.0,
+              "norm_ema_decay": 0.999,
+              "num_out": 64,
+              "layers": [
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 512,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                },
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 256,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                },
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 128,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                }
+              ],
+              "_target_": "protomotions.agents.common.mlp.MLPWithConcat",
+              "in_keys": [
+                "solution_condition_features"
+              ],
+              "out_keys": [
+                "latent_dynamics_predicted"
+              ],
+              "output_activation": null,
+              "module_operations": [
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationForwardConfig"
+                }
+              ]
+            },
+            {
+              "$type": "protomotions.agents.common.config.MLPWithConcatConfig",
+              "normalize_obs": false,
+              "norm_clamp_value": 5.0,
+              "norm_ema_decay": 0.999,
+              "num_out": 29,
+              "layers": [
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 1024,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                },
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 1024,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                },
+                {
+                  "$type": "protomotions.agents.common.config.MLPLayerConfig",
+                  "units": 1024,
+                  "activation": "silu",
+                  "use_layer_norm": false
+                }
+              ],
+              "_target_": "protomotions.agents.common.mlp.MLPWithConcat",
+              "in_keys": [
+                "fk_max_coords_obs_norm",
+                "previous_actions_norm",
+                "solution_condition_features",
+                "decoder_reduced_coords_norm"
+              ],
+              "out_keys": [
+                "actor_trunk_out"
+              ],
+              "output_activation": null,
+              "module_operations": [
+                {
+                  "$type": "protomotions.agents.common.config.ModuleOperationForwardConfig"
+                }
+              ]
+            }
+          ],
+          "_target_": "protomotions.agents.common.common.ModuleContainer",
+          "in_keys": [
+            "reduced_coords_obs",
+            "fk_max_coords_obs",
+            "previous_actions",
+            "decoder_condition",
+            "vae_latent"
+          ],
+          "out_keys": [
+            "actor_trunk_out",
+            "solution_condition_features",
+            "latent_dynamics_predicted"
+          ]
+        },
+        "vae": {
+          "$type": "protomotions.agents.supervised.masked_mimic_config.MaskedMimicVAEConfig",
+          "kld_schedule": {
+            "$type": "protomotions.agents.supervised.masked_mimic_config.KLDScheduleConfig",
+            "init_kld_coeff": 0.0001,
+            "end_kld_coeff": 0.01,
+            "start_epoch": 500,
+            "end_epoch": 2000,
+            "cyclical": true
+          },
+          "vae_latent_dim": 64,
+          "vae_noise_type": {
+            "$enum": "protomotions.agents.supervised.masked_mimic_config.VAENoiseType",
+            "name": "ZEROS"
+          },
+          "latent_geometry": {
+            "$enum": "protomotions.agents.supervised.masked_mimic_config.LatentGeometry",
+            "name": "SPHERE"
+          },
+          "prior_mean_unit_norm_coeff": 0.001
+        },
+        "temporal_losses": [
+          {
+            "$type": "protomotions.agents.supervised.masked_mimic_config.TemporalLossConfig",
+            "key": "privileged_action",
+            "coefficient": 0.05,
+            "prediction_key": "privileged_action",
+            "forward": false
+          },
+          {
+            "$type": "protomotions.agents.supervised.masked_mimic_config.TemporalLossConfig",
+            "key": "solution_condition_features",
+            "coefficient": 0.001,
+            "prediction_key": "solution_condition_features",
+            "forward": false
+          },
+          {
+            "$type": "protomotions.agents.supervised.masked_mimic_config.TemporalLossConfig",
+            "key": "decoder_condition",
+            "coefficient": 0.001,
+            "prediction_key": "latent_dynamics_predicted",
+            "forward": true
+          }
+        ],
+        "optimizer": {
+          "$type": "protomotions.agents.base_agent.config.OptimizerConfig",
+          "_target_": "torch.optim.Adam",
+          "lr": 2e-05,
+          "weight_decay": 0.0,
+          "eps": 1e-08,
+          "betas": [
+            0.9,
+            0.999
+          ]
+        }
+      },
+      "num_steps": 32,
+      "gradient_clip_val": 0.0,
+      "fail_on_bad_grads": false,
+      "check_grad_mag": true,
+      "gamma": 0.99,
+      "bounds_loss_coef": 0.0,
+      "task_reward_w": 1.0,
+      "num_mini_epochs": 1,
+      "training_early_termination": null,
+      "save_epoch_checkpoint_every": 1000,
+      "save_last_checkpoint_every": 10,
+      "save_inference_checkpoint": false,
+      "max_episode_length_manager": null,
+      "evaluator": {
+        "$type": "protomotions.agents.evaluators.config.MaskedMimicEvaluatorConfig",
+        "_target_": "protomotions.agents.evaluators.masked_mimic_evaluator.MaskedMimicEvaluator",
+        "metric_components": {},
+        "failure_components": {},
+        "score_key": null,
+        "max_eval_steps": 600,
+        "eval_metrics_every": 200,
+        "save_predicted_motion_lib_every": 3,
+        "motion_weights_rules": {
+          "$type": "protomotions.agents.evaluators.config.MotionWeightsRulesConfig",
+          "motion_weights_update_success_discount": 0.999,
+          "motion_weights_update_failure_discount": 0.999,
+          "min_motion_weight": "1/num_motions"
+        },
+        "action_source": {
+          "$enum": "protomotions.agents.evaluators.config.MaskedMimicActionSource",
+          "name": "POSTERIOR"
+        }
+      },
+      "normalize_rewards": true,
+      "normalized_reward_clamp_value": 5.0,
+      "reward_norm_ema_decay": null,
+      "expert_model_path": null,
+      "rollout_actor": {
+        "$enum": "protomotions.agents.supervised.config.RolloutActor",
+        "name": "STUDENT"
+      },
+      "loss": {
+        "$type": "protomotions.agents.common.supervision.SupervisionLossConfig",
+        "loss_type": "mse",
+        "prediction_key": "privileged_action",
+        "target_key": "expert_actions",
+        "prediction_logvar_key": null,
+        "target_logvar_key": null,
+        "label_smoothing": 0.0,
+        "weight": 1.0,
+        "log_prefix": "masked_mimic",
+        "enabled": true,
+        "extra": {}
+      },
+      "expert_rollout_probability": 0.0
+    }
+  }
+}

README.md

@@ -1,7 +1,7 @@
 ---
 license: other
 license_name: nvidia-open-model-license
-license_link: https://huggingface.co/nvidia/metatron-browser/blob/main/LICENSE.pdf
+license_link: https://huggingface.co/nvidia/Metatron-v1/blob/main/LICENSE.pdf
 language:
 - en
 library_name: onnx
@@ -13,7 +13,42 @@ tags:
 - metatron
 ---
 
-# Metatron v1 — G1 Browser Exports
+# Metatron v1
+
+## Choose an artifact
+
+| Package | Use | Documentation |
+| --- | --- | --- |
+| [G1/](G1/) | Native PyTorch G1 inference: checkpoint, resolved config, reconstruction recipe and legal files | [G1 model card](G1/MODEL_CARD.md), [download and run](G1/README.md) |
+| `demo/model.onnx` | Single-robot browser scene/recovery pipeline | Browser export card below |
+| `racer/network.onnx` | Six-robot batched G1 Racer network | Browser export card below |
+
+The native checkpoint is 459,455,983 bytes (SHA256
+`e5fc2e47a8f9ee9b96a674cabbf236b3c510f887427645a14bc1779b10f0b43a`).
+It retains the posterior and normalization statistics; the deterministic browser
+graphs export the prior/decoder path. Interfaces are not interchangeable.
+The native recipe masks scene inputs; see its card before re-exporting.
+
+`config.json` is a machine-readable artifact index with file sizes and SHA256
+digests, consumed by ProtoMotions' pinned native downloader. Model files are
+hosted here, not served by GitHub Pages. The existing browser paths are unchanged.
+
+## License
+
+Model license: NVIDIA Open Model License. The [full agreement](LICENSE.pdf) and
+[attribution notice](NOTICE.txt) govern model artifacts, including native weights
+and serialized configurations. Source code remains Apache-2.0. Dataset and asset
+rights are separate; this repository does not redistribute training motion clips.
+
+- Deployment Geography: Global.
+- Use Case: Researchers and developers running interactive humanoid control in a
+  browser; the separate native G1 package also supports simulator and supervised
+  real-G1 deployment workflows.
+- Release: Metatron v1, hosted in this repository and integrated with
+  [ProtoMotions](https://github.com/NVlabs/ProtoMotions). Pin a Hub revision; no
+  public ProtoMotions release tag is asserted here.
+
+# Browser export model card
 
 ## Overview
 
@@ -129,3 +164,13 @@ of the associated application bundles, motion clips and other assets.
 - [ProtoMotions](https://github.com/NVlabs/ProtoMotions)
 - [ProtoMotions documentation](https://protomotions.github.io/)
 - [NVIDIA Open Model License](https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-open-model-license/)
+
+## Download statistics
+
+Hub download totals are server-side query-file request counts, not unique users
+or game sessions. The native downloader reads the functional `config.json`
+artifact index; it does not send a separate analytics request. Locally cached
+loads do not request model files again. The existing ONNX loaders remain pinned
+to their own artifacts; adding the native package does not retroactively measure
+browser usage. Library-specific counting can differ; see
+[Hugging Face download statistics](https://huggingface.co/docs/hub/models-download-stats).

config.json

@@ -0,0 +1,68 @@
+{
+  "schema_version": 1,
+  "model": "Metatron v1",
+  "repo_id": "nvidia/Metatron-v1",
+  "packages": {
+    "G1": {
+      "files": [
+        {
+          "path": "G1/LICENSE.pdf",
+          "size": 79005,
+          "sha256": "4d2fb590aa9b30c47f2058bff17291df7fa2aa0c1bd775a20703da9bb267cfab"
+        },
+        {
+          "path": "G1/MODEL_CARD.md",
+          "size": 7125,
+          "sha256": "7b209dc7d24e34de9d34e002a43ef86f94b0676e973191e059cf0a5028ad57e5"
+        },
+        {
+          "path": "G1/NOTICE.txt",
+          "size": 345,
+          "sha256": "156bc21cf2ef9959c172cc903db26391f4c2c189244bbb17fd71a11f025d3899"
+        },
+        {
+          "path": "G1/README.md",
+          "size": 1562,
+          "sha256": "3a642332b2c8be0aad48a047a1022a76ae6aa5e6272d788791bb7b1387a5ad97"
+        },
+        {
+          "path": "G1/config.yaml",
+          "size": 428,
+          "sha256": "af1ed1dbbc9b92520e108bea585cf94036941e63ded51618d7dd8c8632194745"
+        },
+        {
+          "path": "G1/experiment_config.py",
+          "size": 20409,
+          "sha256": "90c0a6411f7717e8e699ebb419fabb66000010445c92464c7c05c60fc160bb06"
+        },
+        {
+          "path": "G1/inference_last.ckpt",
+          "size": 459455983,
+          "sha256": "e5fc2e47a8f9ee9b96a674cabbf236b3c510f887427645a14bc1779b10f0b43a"
+        },
+        {
+          "path": "G1/provenance.json",
+          "size": 289,
+          "sha256": "084f3a1e87514ba15fa03428fd2c851efef1b6f32dc7571e3624a4a1e5f408dd"
+        },
+        {
+          "path": "G1/resolved_configs.json",
+          "size": 91427,
+          "sha256": "2665afcf1d789fec1b979ce235413943e65e12ba67baa759d0bbfa7dc62f18cc"
+        }
+      ]
+    }
+  },
+  "browser_exports": [
+    {
+      "path": "demo/model.onnx",
+      "size": 403375789,
+      "sha256": "7180e82cae7ec04ff7766c262d430f57573d93678e025d8a71916885072a0cd4"
+    },
+    {
+      "path": "racer/network.onnx",
+      "size": 403036698,
+      "sha256": "130a22718a34c0da488e2157553b40efedfefcf6ca6cb6b97cfdd67f8dfbeaf3"
+    }
+  ]
+}