mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-06 09:37:06 +00:00
Compare commits
3 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 Martino Russi
|
f442c21e46 | ||
|
Martino Russi
|
ba89c73b67 | ||
|
Steven Palma
|
7957d4e2dc |
@@ -87,7 +87,7 @@ Learn more about it in the [LeRobotDataset Documentation](https://huggingface.co
|
||||
|
||||
## SoTA Models
|
||||
|
||||
LeRobot implements state-of-the-art policies in pure PyTorch, covering Imitation Learning, Reinforcement Learning, and Vision-Language-Action (VLA) models, with more coming soon. It also provides you with the tools to instrument and inspect your training process.
|
||||
LeRobot implements state-of-the-art policies in pure PyTorch, covering Imitation Learning, Reinforcement Learning, Vision-Language-Action (VLA) models, World Models, and Reward Models, with more coming soon. It also provides you with the tools to instrument and inspect your training process.
|
||||
|
||||
<p align="center">
|
||||
<img alt="Gr00t Architecture" src="./media/readme/VLA_architecture.jpg" width="640px">
|
||||
@@ -101,13 +101,13 @@ lerobot-train \
|
||||
--dataset.repo_id=lerobot/aloha_mobile_cabinet
|
||||
```
|
||||
|
||||
| Category | Models |
|
||||
| -------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
|
||||
| **Imitation Learning** | [ACT](./docs/source/policy_act_README.md), [Diffusion](./docs/source/policy_diffusion_README.md), [VQ-BeT](./docs/source/policy_vqbet_README.md), [Multitask DiT Policy](./docs/source/policy_multi_task_dit_README.md) |
|
||||
| **Reinforcement Learning** | [HIL-SERL](./docs/source/hilserl.mdx), [TDMPC](./docs/source/policy_tdmpc_README.md) & QC-FQL (coming soon) |
|
||||
| **VLAs Models** | [Pi0](./docs/source/pi0.mdx), [Pi0Fast](./docs/source/pi0fast.mdx), [Pi0.5](./docs/source/pi05.mdx), [GR00T N1.7](./docs/source/policy_groot_README.md), [SmolVLA](./docs/source/policy_smolvla_README.md), [XVLA](./docs/source/xvla.mdx), [EO-1](./docs/source/eo1.mdx), [MolmoAct2](./docs/source/molmoact2.mdx), [WALL-OSS](./docs/source/walloss.mdx) |
|
||||
| **World Models** | [VLA-JEPA](./docs/source/vla_jepa.mdx) (more coming soon) |
|
||||
| **Reward Models** | [SARM](./docs/source/sarm.mdx), [TOPReward](./docs/source/topreward.mdx), [Robometer](./docs/source/robometer.mdx) |
|
||||
| Category | Models |
|
||||
| -------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ |
|
||||
| **Imitation Learning** | [ACT](./docs/source/policy_act_README.md), [Diffusion](./docs/source/policy_diffusion_README.md), [VQ-BeT](./docs/source/policy_vqbet_README.md), [Multitask DiT Policy](./docs/source/policy_multi_task_dit_README.md) |
|
||||
| **Reinforcement Learning** | [HIL-SERL](./docs/source/hilserl.mdx), [TDMPC](./docs/source/policy_tdmpc_README.md) & QC-FQL (coming soon) |
|
||||
| **VLAs Models** | [Pi0](./docs/source/pi0.mdx), [Pi0Fast](./docs/source/pi0fast.mdx), [Pi0.5](./docs/source/pi05.mdx), [GR00T N1.7](./docs/source/policy_groot_README.md), [SmolVLA](./docs/source/policy_smolvla_README.md), [XVLA](./docs/source/xvla.mdx), [EO-1](./docs/source/eo1.mdx), [MolmoAct2](./docs/source/molmoact2.mdx), [WALL-OSS](./docs/source/walloss.mdx), [EVO1](./docs/source/evo1.mdx) |
|
||||
| **World Models** | [VLA-JEPA](./docs/source/vla_jepa.mdx), [LingBot-VA](./docs/source/lingbot_va.mdx), [FastWAM](./docs/source/fastwam.mdx) |
|
||||
| **Reward Models** | [SARM](./docs/source/sarm.mdx), [TOPReward](./docs/source/topreward.mdx), [Robometer](./docs/source/robometer.mdx) |
|
||||
|
||||
Similarly to the hardware, you can easily implement your own policy & leverage LeRobot's data collection, training, and visualization tools, and share your model to the HF Hub
|
||||
|
||||
|
||||
@@ -295,11 +295,12 @@ The file names are load-bearing: the factory does lazy imports by name, and the
|
||||
|
||||
### Wiring
|
||||
|
||||
Three places need to know about your policy. All by name.
|
||||
Four places need to know about your policy. All by name.
|
||||
|
||||
1. **`policies/__init__.py`** — re-export `MyPolicyConfig` and add it to `__all__`. **Don't** re-export the modeling class; it loads lazily through the factory (so `import lerobot` stays fast).
|
||||
2. **`factory.py:get_policy_class`** — add a branch returning `MyPolicy` from a lazy import.
|
||||
3. **`factory.py:make_policy_config`** and **`factory.py:make_pre_post_processors`** — same idea, two more branches.
|
||||
4. **`templates/lerobot_modelcard_template.md` and the root `README.md`** — the template is what `push_model_to_hub` renders into the model card of every checkpoint trained with your policy: add a one-line description of your policy in the `model_name` branches, map it in `policy_docs` so cards link to your MDX guide, and optionally add an architecture image to `diagrams`. Then add your policy to the models table in the root `README.md`, under the right category, linking to your doc page.
|
||||
|
||||
Mirror an existing policy that's structurally similar to yours; the diff is small.
|
||||
|
||||
@@ -371,6 +372,8 @@ The general expectations are in [`CONTRIBUTING.md`](https://github.com/huggingfa
|
||||
- [ ] Optional deps live behind a `[project.optional-dependencies]` extra and the `TYPE_CHECKING + require_package` guard.
|
||||
- [ ] `tests/policies/` updated; backward-compat artifact committed & policy-specific tests.
|
||||
- [ ] `src/lerobot/policies/<name>/README.md` symlinked into `docs/source/policy_<name>_README.md`; user-facing `docs/source/<name>.mdx` written and added to `_toctree.yml`.
|
||||
- [ ] `templates/lerobot_modelcard_template.md` has a description entry and a `policy_docs` link for your policy.
|
||||
- [ ] The models table in the root `README.md` lists your policy in the right category, linking to your doc page.
|
||||
- [ ] At least one reproducible benchmark eval in the policy MDX with a published checkpoint (sim benchmark, or real-robot dataset + checkpoint).
|
||||
|
||||
The fastest way to get a clean PR is to copy the directory of the existing policy closest to yours, rename, and replace contents method by method. Don't wait until everything is polished — open a draft PR early and iterate with us; reviewers would much rather give feedback on a half-finished branch than a fully-merged one.
|
||||
|
||||
@@ -160,13 +160,13 @@ This will follow the recipe found [here](https://github.com/NVIDIA/Isaac-GR00T/b
|
||||
|
||||
Preliminary LeRobot integration results (GR00T-LeRobot, `eval.n_episodes >= 50` per suite):
|
||||
|
||||
| Suite | Success rate |
|
||||
| ---------------- | -----------: |
|
||||
| LIBERO Spatial | 94% |
|
||||
| LIBERO Object | 98% |
|
||||
| LIBERO Goal | 93% |
|
||||
| LIBERO 10 (Long) | 90% |
|
||||
| **Average** | **93.75%** |
|
||||
| Suite | Success rate | Checkpoint |
|
||||
| ---------------- | -----------: | ------------------------------------------------------------------------------------------------------------- |
|
||||
| LIBERO Spatial | 91% | [nvidia/gr00t17-lerobot-libero_spatial-640](https://huggingface.co/nvidia/gr00t17-lerobot-libero_spatial-640) |
|
||||
| LIBERO Object | 81% | [nvidia/gr00t17-lerobot-libero_object-640](https://huggingface.co/nvidia/gr00t17-lerobot-libero_object-640) |
|
||||
| LIBERO Goal | 97% | [nvidia/gr00t17-lerobot-libero_goal-640](https://huggingface.co/nvidia/gr00t17-lerobot-libero_goal-640) |
|
||||
| LIBERO 10 (Long) | 84% | [nvidia/gr00t17-lerobot-libero_10-640](https://huggingface.co/nvidia/gr00t17-lerobot-libero_10-640) |
|
||||
| **Average** | **88.25%** | |
|
||||
|
||||
```bash
|
||||
export MODEL_ID=your_trained_model_on_huggingface
|
||||
|
||||
@@ -113,13 +113,11 @@ class DynamixelMotorsBus(SerialMotorsBus):
|
||||
port: str,
|
||||
motors: dict[str, Motor],
|
||||
calibration: dict[str, MotorCalibration] | None = None,
|
||||
protocol_version: int = PROTOCOL_VERSION,
|
||||
):
|
||||
require_package("dynamixel-sdk", extra="dynamixel", import_name="dynamixel_sdk")
|
||||
super().__init__(port, motors, calibration)
|
||||
self.port_handler = dxl.PortHandler(self.port)
|
||||
self.packet_handler = dxl.PacketHandler(protocol_version)
|
||||
print(f"Using protocol version {protocol_version}")
|
||||
self.packet_handler = dxl.PacketHandler(PROTOCOL_VERSION)
|
||||
self.sync_reader = dxl.GroupSyncRead(self.port_handler, self.packet_handler, 0, 0)
|
||||
self.sync_writer = dxl.GroupSyncWrite(self.port_handler, self.packet_handler, 0, 0)
|
||||
self._comm_success = dxl.COMM_SUCCESS
|
||||
|
||||
@@ -33,58 +33,6 @@
|
||||
# 2. We can change the value of the MyControlTableKey enums without impacting the client code
|
||||
|
||||
|
||||
# {data_name: (address, size_byte)}
|
||||
# https://emanual.robotis.com/docs/en/dxl/ax/{MODEL}/#control-table
|
||||
AX_SERIES_CONTROL_TABLE = {
|
||||
# EEPROM Area
|
||||
"Model_Number": (0, 2),
|
||||
"Firmware_Version": (2, 1),
|
||||
"ID": (3, 1),
|
||||
"Baud_Rate": (4, 1),
|
||||
"Return_Delay_Time": (5, 1),
|
||||
"CW_Angle_Limit": (6, 2),
|
||||
"CCW_Angle_Limit": (8, 2),
|
||||
"Temperature_Limit": (11, 1),
|
||||
"Min_Voltage_Limit": (12, 1),
|
||||
"Max_Voltage_Limit": (13, 1),
|
||||
"Max_Torque": (14, 2),
|
||||
"Status_Return_Level": (16, 1),
|
||||
"Alarm_LED": (17, 1),
|
||||
"Shutdown": (18, 1),
|
||||
# RAM Area
|
||||
"Torque_Enable": (24, 1),
|
||||
"LED": (25, 1),
|
||||
"CW_Compliance_Margin": (26, 1),
|
||||
"CCW_Compliance_Margin": (27, 1),
|
||||
"CW_Compliance_Slope": (28, 1),
|
||||
"CCW_Compliance_Slope": (29, 1),
|
||||
"Goal_Position": (30, 2),
|
||||
"Moving_Speed": (32, 2),
|
||||
"Torque_Limit": (34, 2),
|
||||
"Present_Position": (36, 2),
|
||||
"Present_Speed": (38, 2),
|
||||
"Present_Load": (40, 2),
|
||||
"Present_Voltage": (42, 1),
|
||||
"Present_Temperature": (43, 1),
|
||||
"Registered": (44, 1),
|
||||
"Moving": (46, 1),
|
||||
"Lock": (47, 1),
|
||||
"Punch": (48, 2),
|
||||
}
|
||||
|
||||
# https://emanual.robotis.com/docs/en/dxl/ax/{MODEL}/#baud-rate4
|
||||
AX_SERIES_BAUDRATE_TABLE = {
|
||||
9_600: 207,
|
||||
19_200: 103,
|
||||
57_600: 34,
|
||||
115_200: 16,
|
||||
200_000: 9,
|
||||
250_000: 7,
|
||||
400_000: 4,
|
||||
500_000: 3,
|
||||
1_000_000: 1,
|
||||
}
|
||||
|
||||
# {data_name: (address, size_byte)}
|
||||
# https://emanual.robotis.com/docs/en/dxl/x/{MODEL}/#control-table
|
||||
X_SERIES_CONTROL_TABLE = {
|
||||
@@ -166,14 +114,6 @@ X_SERIES_ENCODINGS_TABLE = {
|
||||
"Present_Velocity": X_SERIES_CONTROL_TABLE["Present_Velocity"][1],
|
||||
}
|
||||
|
||||
# {data_name: size_byte}
|
||||
AX_SERIES_ENCODINGS_TABLE = {
|
||||
"Goal_Position": AX_SERIES_CONTROL_TABLE["Goal_Position"][1],
|
||||
"Moving_Speed": AX_SERIES_CONTROL_TABLE["Moving_Speed"][1],
|
||||
"Present_Position": AX_SERIES_CONTROL_TABLE["Present_Position"][1],
|
||||
"Present_Speed": AX_SERIES_CONTROL_TABLE["Present_Speed"][1],
|
||||
}
|
||||
|
||||
MODEL_ENCODING_TABLE = {
|
||||
"x_series": X_SERIES_ENCODINGS_TABLE,
|
||||
"xl330-m077": X_SERIES_ENCODINGS_TABLE,
|
||||
@@ -182,8 +122,6 @@ MODEL_ENCODING_TABLE = {
|
||||
"xm430-w350": X_SERIES_ENCODINGS_TABLE,
|
||||
"xm540-w270": X_SERIES_ENCODINGS_TABLE,
|
||||
"xc430-w150": X_SERIES_ENCODINGS_TABLE,
|
||||
"ax_series": AX_SERIES_ENCODINGS_TABLE,
|
||||
"ax-12a": AX_SERIES_ENCODINGS_TABLE,
|
||||
}
|
||||
|
||||
# {model: model_resolution}
|
||||
@@ -196,8 +134,6 @@ MODEL_RESOLUTION = {
|
||||
"xm430-w350": 4096,
|
||||
"xm540-w270": 4096,
|
||||
"xc430-w150": 4096,
|
||||
"ax_series": 1024,
|
||||
"ax-12a": 1024,
|
||||
}
|
||||
|
||||
# {model: model_number}
|
||||
@@ -209,7 +145,6 @@ MODEL_NUMBER_TABLE = {
|
||||
"xm430-w350": 1020,
|
||||
"xm540-w270": 1120,
|
||||
"xc430-w150": 1070,
|
||||
"ax-12a": 12,
|
||||
}
|
||||
|
||||
# {model: available_operating_modes}
|
||||
@@ -231,8 +166,6 @@ MODEL_CONTROL_TABLE = {
|
||||
"xm430-w350": X_SERIES_CONTROL_TABLE,
|
||||
"xm540-w270": X_SERIES_CONTROL_TABLE,
|
||||
"xc430-w150": X_SERIES_CONTROL_TABLE,
|
||||
"ax_series": AX_SERIES_CONTROL_TABLE,
|
||||
"ax-12a": AX_SERIES_CONTROL_TABLE,
|
||||
}
|
||||
|
||||
MODEL_BAUDRATE_TABLE = {
|
||||
@@ -243,8 +176,6 @@ MODEL_BAUDRATE_TABLE = {
|
||||
"xm430-w350": X_SERIES_BAUDRATE_TABLE,
|
||||
"xm540-w270": X_SERIES_BAUDRATE_TABLE,
|
||||
"xc430-w150": X_SERIES_BAUDRATE_TABLE,
|
||||
"ax_series": AX_SERIES_BAUDRATE_TABLE,
|
||||
"ax-12a": AX_SERIES_BAUDRATE_TABLE,
|
||||
}
|
||||
|
||||
AVAILABLE_BAUDRATES = [
|
||||
|
||||
@@ -302,6 +302,33 @@ def _pad_evo1_stats(
|
||||
return padded_stats
|
||||
|
||||
|
||||
def _refresh_evo1_normalization_steps(
|
||||
config: Evo1Config,
|
||||
preprocessor: PolicyProcessorPipeline,
|
||||
postprocessor: PolicyProcessorPipeline,
|
||||
) -> None:
|
||||
"""Re-pad checkpoint-loaded (un)normalizer stats/features to EVO1's fixed widths.
|
||||
|
||||
Loading a checkpoint injects the raw dataset stats (unpadded to max_state_dim/max_action_dim)
|
||||
into the (un)normalizer via the generic override path in make_pre_post_processors. Those stats
|
||||
and their declared features must be re-padded/reshaped to EVO1's fixed widths, otherwise
|
||||
normalization fails against the padded state/action tensors (e.g. state padded to 24 vs. 8-dim
|
||||
LIBERO stats). Padding is a no-op when stats are already at the target width.
|
||||
"""
|
||||
normalization_features = _evo1_normalization_features(config)
|
||||
action_features = _evo1_action_features(config)
|
||||
for step in preprocessor.steps:
|
||||
if isinstance(step, NormalizerProcessorStep):
|
||||
step.features = normalization_features
|
||||
step.stats = _pad_evo1_stats(config, step.stats)
|
||||
step.to(device=step.device, dtype=step.dtype)
|
||||
for step in postprocessor.steps:
|
||||
if isinstance(step, UnnormalizerProcessorStep):
|
||||
step.features = action_features
|
||||
step.stats = _pad_evo1_stats(config, step.stats)
|
||||
step.to(device=step.device, dtype=step.dtype)
|
||||
|
||||
|
||||
def reconcile_evo1_processors(
|
||||
config: Evo1Config,
|
||||
preprocessor: PolicyProcessorPipeline,
|
||||
@@ -309,16 +336,19 @@ def reconcile_evo1_processors(
|
||||
) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]:
|
||||
"""Reconcile checkpoint-loaded pipelines with the current EVO1 config.
|
||||
|
||||
Two things cannot be restored from a serialized pipeline alone: the EVO1 batch converter
|
||||
(converters are plain functions and are never serialized), and eval-time CLI overrides of the
|
||||
action postprocessing flags (`postprocess_action_dim`, `binarize_gripper`, `gripper_*`). This
|
||||
restores the converter and rebuilds the action step from the current config so those overrides
|
||||
take effect.
|
||||
Three things cannot be restored from a serialized pipeline alone: the EVO1 batch converter
|
||||
(converters are plain functions and are never serialized), eval-time CLI overrides of the
|
||||
action postprocessing flags (`postprocess_action_dim`, `binarize_gripper`, `gripper_*`), and the
|
||||
(un)normalizer stats/features when the generic override path injects raw, unpadded dataset
|
||||
stats. This restores the converter, re-pads the normalization stats to EVO1's fixed widths, and
|
||||
rebuilds the action step from the current config so those overrides take effect.
|
||||
"""
|
||||
# Pipelines reloaded from a checkpoint come back with the default batch converter, which drops
|
||||
# non-observation extras (embodiment_id, state_mask, custom task fields) needed by EVO1.
|
||||
preprocessor.to_transition = evo1_batch_to_transition
|
||||
|
||||
_refresh_evo1_normalization_steps(config, preprocessor, postprocessor)
|
||||
|
||||
action_step = Evo1ActionProcessorStep(
|
||||
action_dim=_evo1_action_dim(config),
|
||||
binarize_gripper=config.binarize_gripper,
|
||||
|
||||
@@ -30,13 +30,19 @@ This is a Gaussian Actor policy (Gaussian policy with a tanh squash) — the pol
|
||||
{% elif model_name == "eo1" %}
|
||||
[EO-1](https://huggingface.co/papers/2508.21112) is a Vision-Language-Action model for general robot control. It pairs a Qwen2.5-VL backbone for vision-language understanding with a continuous flow-matching action head that denoises action chunks.
|
||||
{% elif model_name == "groot" %}
|
||||
[GR00T N1.5](https://github.com/NVIDIA/Isaac-GR00T) is an open, cross-embodiment foundation model from NVIDIA for generalized humanoid robot reasoning and skills. It takes language and images as input and uses a flow-matching action transformer to predict actions conditioned on vision, language, and proprioception.
|
||||
[GR00T N1.7](https://github.com/NVIDIA/Isaac-GR00T) is an open, cross-embodiment foundation model from NVIDIA for generalized humanoid robot reasoning and skills. It uses a Cosmos-Reason2/Qwen3-VL backbone and a flow-matching action transformer to predict actions conditioned on vision, language, and proprioception.
|
||||
{% elif model_name == "multi_task_dit" %}
|
||||
[Multi-Task Diffusion Transformer (DiT)](https://huggingface.co/papers/2507.05331) extends Diffusion Policy with a large Diffusion Transformer and text + vision conditioning for multi-task robot learning. It supports both diffusion and flow-matching objectives and reaches high dexterity with only ~450M parameters.
|
||||
{% elif model_name == "wall_x" %}
|
||||
[WALL-OSS](https://huggingface.co/papers/2509.11766) is an open-source foundation model for embodied intelligence from XSquare Robot. Built on Qwen2.5-VL, it uses a tightly-coupled multimodal architecture with flow matching to unify semantic reasoning and high-frequency action generation for cross-embodiment control.
|
||||
{% elif model_name == "xvla" %}
|
||||
[X-VLA](https://huggingface.co/papers/2510.10274) is a soft-prompted, flow-matching Vision-Language-Action framework that treats each robot or hardware setup as a "task" encoded with a small set of learnable Soft Prompt embeddings, letting a single model reconcile diverse robot morphologies, sensors, and action spaces.
|
||||
{% elif model_name == "evo1" %}
|
||||
[EVO1](https://github.com/MINT-SJTU/Evo-1) is a Vision-Language-Action policy built around an InternVL3 backbone and a continuous flow-matching action head. It embeds camera images and the language instruction with InternVL3 and predicts future action chunks via flow matching.
|
||||
{% elif model_name == "fastwam" %}
|
||||
[FastWAM](https://arxiv.org/abs/2603.16666) is a World Action Model policy that keeps video world-modeling during training but predicts actions directly at inference time, initializing its visual world-model components from the Wan2.2 video-diffusion stack.
|
||||
{% elif model_name == "lingbot_va" %}
|
||||
[LingBot-VA](https://github.com/Robbyant/lingbot-va) is an autoregressive video-action world-model policy built on the Wan2.2 video-diffusion stack. It interleaves the prediction of future video latents and robot actions in a single autoregressive sequence, feeding observed keyframes back into its KV cache for closed-loop world modeling.
|
||||
{% else %}
|
||||
This is a **{{ model_name }}** policy trained with [LeRobot](https://github.com/huggingface/lerobot).
|
||||
{% endif %}
|
||||
@@ -75,7 +81,10 @@ This policy has been trained and pushed to the Hub using [LeRobot](https://githu
|
||||
"groot": "groot",
|
||||
"xvla": "xvla",
|
||||
"multi_task_dit": "multi_task_dit",
|
||||
"wall_x": "walloss"
|
||||
"wall_x": "walloss",
|
||||
"evo1": "evo1",
|
||||
"fastwam": "fastwam",
|
||||
"lingbot_va": "lingbot_va"
|
||||
} %}
|
||||
{% if policy_docs.get(model_name) %}Learn how to train and run it in the [LeRobot {{ model_name }} guide](https://huggingface.co/docs/lerobot/main/en/{{ policy_docs[model_name] }}), or browse the [full documentation](https://huggingface.co/docs/lerobot/index).
|
||||
{% else %}See the [full LeRobot documentation](https://huggingface.co/docs/lerobot/index).
|
||||
|
||||
@@ -496,6 +496,60 @@ def test_evo1_processor_save_load_round_trip_applies_config_overrides(tmp_path):
|
||||
assert "embodiment_id" in processed
|
||||
|
||||
|
||||
def test_reconcile_evo1_processors_repads_overridden_stats(tmp_path):
|
||||
"""Loading a checkpoint and injecting raw (unpadded) dataset stats must be re-padded.
|
||||
|
||||
Regression test: lerobot-train passes the raw dataset stats as normalizer/unnormalizer
|
||||
overrides when resuming from a checkpoint (e.g. stage2 from a stage1 checkpoint). Those stats
|
||||
are at the dataset dims (e.g. LIBERO state=8/action=7), but EVO1 pads state/action to
|
||||
max_state_dim/max_action_dim before normalization, so reconcile_evo1_processors must re-pad the
|
||||
stats or normalization crashes with a shape mismatch.
|
||||
"""
|
||||
config = make_config()
|
||||
preprocessor, postprocessor = make_evo1_pre_post_processors(config, dataset_stats=make_stats())
|
||||
preprocessor.save_pretrained(tmp_path)
|
||||
postprocessor.save_pretrained(tmp_path)
|
||||
|
||||
# Reload with the generic override path injecting raw, unpadded dataset stats.
|
||||
raw_stats = make_stats()
|
||||
loaded_pre = PolicyProcessorPipeline.from_pretrained(
|
||||
tmp_path,
|
||||
config_filename=f"{POLICY_PREPROCESSOR_DEFAULT_NAME}.json",
|
||||
overrides={"normalizer_processor": {"stats": raw_stats}},
|
||||
to_transition=batch_to_transition,
|
||||
to_output=transition_to_batch,
|
||||
)
|
||||
loaded_post = PolicyProcessorPipeline.from_pretrained(
|
||||
tmp_path,
|
||||
config_filename=f"{POLICY_POSTPROCESSOR_DEFAULT_NAME}.json",
|
||||
overrides={"unnormalizer_processor": {"stats": raw_stats}},
|
||||
to_transition=policy_action_to_transition,
|
||||
to_output=transition_to_policy_action,
|
||||
)
|
||||
|
||||
# Sanity: the override really injected unpadded stats before reconciliation.
|
||||
normalizer = next(step for step in loaded_pre.steps if isinstance(step, NormalizerProcessorStep))
|
||||
assert normalizer._tensor_stats[OBS_STATE]["min"].shape == (STATE_DIM,)
|
||||
|
||||
loaded_pre, loaded_post = reconcile_evo1_processors(config, loaded_pre, loaded_post)
|
||||
|
||||
normalizer = next(step for step in loaded_pre.steps if isinstance(step, NormalizerProcessorStep))
|
||||
unnormalizer = next(step for step in loaded_post.steps if isinstance(step, UnnormalizerProcessorStep))
|
||||
assert normalizer._tensor_stats[OBS_STATE]["min"].shape == (MAX_STATE_DIM,)
|
||||
assert normalizer._tensor_stats[ACTION]["min"].shape == (MAX_ACTION_DIM,)
|
||||
assert unnormalizer._tensor_stats[ACTION]["min"].shape == (MAX_ACTION_DIM,)
|
||||
|
||||
# Normalizing a padded state must not raise (this is the exact runtime path that crashed).
|
||||
processed = loaded_pre(
|
||||
{
|
||||
"task": "pick the block",
|
||||
OBS_STATE: torch.zeros(STATE_DIM),
|
||||
f"{OBS_IMAGES}.front": torch.rand(3, 16, 16),
|
||||
}
|
||||
)
|
||||
assert processed[OBS_STATE].shape == (1, MAX_STATE_DIM)
|
||||
|
||||
|
||||
def test_evo1_policy_forward_and_inference_use_batched_embedding(monkeypatch):
|
||||
monkeypatch.setattr(modeling_evo1, "Evo1Model", DummyEvo1Model)
|
||||
policy = modeling_evo1.Evo1Policy(make_config())
|
||||
|
||||
Reference in New Issue
Block a user