Compare commits

..

1 Commits

Author SHA1 Message Date
Maxime Ellerbach a3755f07c5 fix(peft): allow fresh LoRA fine-tuning from a base-model checkpoint 2026-07-24 13:37:03 +00:00
31 changed files with 1137 additions and 668 deletions
+14 -16
View File
@@ -150,14 +150,14 @@ class MyPolicy(PreTrainedPolicy):
The methods called by the train/eval loops:
| Method | Used by | What it does |
| ----------------------------------------------------------------- | ----------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
| `reset() -> None` | `lerobot-eval` | Clear per-episode state at the start of each episode. |
| `select_action(batch, **kwargs) -> Tensor` | `lerobot-eval` | Return the next action `(B, action_dim)`. Called every step. |
| `predict_action_chunk(batch, **kwargs) -> Tensor` | the policy itself | Return an action chunk `(B, chunk_size, action_dim)`. Currently abstract on the base class — raise `NotImplementedError` if your policy doesn't chunk. |
| `forward(batch, reduction="mean") -> tuple[Tensor, dict \| None]` | `lerobot-train` | Return `(loss, output_dict)`. Accept `reduction="none"` if you want to support per-sample weighting. |
| `get_optim_params() -> dict` | the optimizer | Return `self.parameters()` for simple policies; return a named parameter dict for multi-optimizer policies (see `get_optim_params` in [`modeling_act.py`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/act/modeling_act.py) for a per-group learning-rate example). |
| `update() -> None` _(optional)_ | `lerobot-train` | Called after each optimizer step _if defined_. Use for EMA, target nets, replay buffers (TDMPC uses this). |
| Method | Used by | What it does |
| ----------------------------------------------------------------- | ----------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
| `reset() -> None` | `lerobot-eval` | Clear per-episode state at the start of each episode. |
| `select_action(batch, **kwargs) -> Tensor` | `lerobot-eval` | Return the next action `(B, action_dim)`. Called every step. |
| `predict_action_chunk(batch, **kwargs) -> Tensor` | the policy itself | Return an action chunk `(B, chunk_size, action_dim)`. Currently abstract on the base class — raise `NotImplementedError` if your policy doesn't chunk. |
| `forward(batch, reduction="mean") -> tuple[Tensor, dict \| None]` | `lerobot-train` | Return `(loss, output_dict)`. Accept `reduction="none"` if you want to support per-sample weighting. |
| `get_optim_params() -> dict` | the optimizer | Return `self.parameters()` for simple policies; return a named parameter dict for [multi-optimizer policies](https://github.com/huggingface/lerobot/blob/ecd38c50d7d15b4184cf42649ff1185ee2e11eeb/src/lerobot/policies/sac/modeling_sac.py#L61-L73). |
| `update() -> None` _(optional)_ | `lerobot-train` | Called after each optimizer step _if defined_. Use for EMA, target nets, replay buffers (TDMPC uses this). |
Batches are flat dictionaries keyed by the constants in [`lerobot.utils.constants`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/utils/constants.py): `OBS_STATE` (`observation.state.<motor>`), `OBS_IMAGES` (`observation.images.<camera>`), `OBS_LANGUAGE`, `ACTION`, etc. Reuse the constants — don't invent new prefixes.
@@ -295,10 +295,12 @@ The file names are load-bearing: the factory does lazy imports by name, and the
### Wiring
Two 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__`. This import is what registers your policy: `@PreTrainedConfig.register_subclass("my_policy")` runs, and from then on the factory resolves everything by convention. **Don't** re-export the modeling class; it loads lazily through the factory (so `import lerobot` stays fast).
2. **`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.
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.
@@ -330,10 +332,6 @@ This way:
Add a matching extra to [`pyproject.toml`](https://github.com/huggingface/lerobot/blob/main/pyproject.toml) `[project.optional-dependencies]` and include it in the `all` extra so `pip install 'lerobot[all]'` keeps installing everything.
### Avoid copying a modeling file — subclass it
If your policy needs to modify a backbone that already exists in `transformers` (custom conditioning, extra inputs, a swapped sub-module), **do not vendor a copy of its `modeling_*.py`**. Instead, subclass the smallest upstream unit and override only what changes. [`pi_gemma.py`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/pi_gemma.py) is the canonical reference: it injects AdaRMS conditioning into PaliGemma/Gemma in ~370 lines by subclassing `GemmaModel`/`PaliGemmaModel` and overriding the decoder-layer forward, instead of forking the ~2,000-line modeling file. Model surgery on a _loaded_ native model is also fine (layer truncation, tokenizer expansion, hidden-state capture — see `evo1/internvl3_embedder.py`, `eo1/modeling_eo1.py`, `groot/groot_n1_7.py` for working examples). Reviewers will ask for this pattern when a PR arrives with a copied modeling file; the only accepted exception is a model that does not exist in `transformers` at all.
### Benchmarks and a published checkpoint
A new policy is much easier to review — and far more useful — when it ships with a working checkpoint and at least one number you can reproduce.
@@ -369,7 +367,7 @@ If your policy is real-robot-only and no sim benchmark applies, swap the sim eva
The general expectations are in [`CONTRIBUTING.md`](https://github.com/huggingface/lerobot/blob/main/CONTRIBUTING.md) and the [PR template](https://github.com/huggingface/lerobot/blob/main/.github/PULL_REQUEST_TEMPLATE.md). On top of those, reviewers will look for:
- [ ] `MyPolicy` and `MyPolicyConfig` cover the surface above; `__init_subclass__` accepts the class.
- [ ] `policies/__init__.py` re-exports the config (this registers the policy; the factory resolves modeling/processor by naming convention).
- [ ] `factory.py` and `policies/__init__.py` are wired (lazy imports for modeling).
- [ ] `make_my_policy_pre_post_processors` follows the naming convention.
- [ ] 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.
+12 -3
View File
@@ -111,17 +111,26 @@ Requirements:
- The block-causal masks use PyTorch **flex-attention**, so build the policy with
`--policy.attn_mode=flex` for training (the default `torch` SDPA is inference-only).
- The full 5B DiT does not fit a single 2432 GB GPU under AdamW; fine-tune with **LoRA**
(`--policy.use_peft=true`) and/or optimizer offload. `get_optim_params` returns only the
trainable (e.g. adapter) parameters; the VAE + UMT5 text encoder stay frozen.
(`--peft.method_type=LORA`) and/or optimizer offload. `get_optim_params` returns only the
trainable (e.g. adapter) parameters; the VAE + UMT5 text encoder stay frozen. Install the
`lerobot[peft]` extra to enable PEFT support (see the [PEFT training guide](./peft_training)).
```bash
lerobot-train \
--policy.path=lerobot/lingbot_va_libero_long --policy.attn_mode=flex \
--policy.use_peft=true \
--peft.method_type=LORA --peft.r=32 --peft.lora_alpha=32 \
--peft.target_modules='transformer\.blocks\.\d+\.attn[12]\.(to_q|to_v)' \
--dataset.repo_id=<your LeRobot-format dataset> \
--batch_size=1 --steps=... --output_dir=outputs/train/lingbot_va
```
Unlike SmolVLA / π₀, LingBot-VA does not ship built-in default LoRA targets, so you must pass
`--peft.target_modules` explicitly. Only `self.transformer` (the dual-stream Wan transformer) is
trainable; the example above adapts the query/value projections of both its self-attention
(`attn1`) and cross-attention (`attn2`) blocks — the standard LoRA target set. Broaden it (e.g.
add `to_k`/`to_out`, or the `ffn` layers) if you need a higher-capacity adapter. Passing
`--peft.method_type` implies PEFT, so `--policy.use_peft=true` is not required.
The dataset must provide camera clips (a temporal window per camera, VAE-encoded to
`frame_chunk_size` latent frames) and `frame_chunk_size * action_per_frame` action steps per item.
+7
View File
@@ -62,3 +62,10 @@ to the `--peft.full_training_modules` parameter:
The learning rate and the scheduled target learning rate can usually be scaled by a factor of 10 compared to the
learning rate used for full fine-tuning (e.g., 1e-4 normal, so 1e-3 using LoRA).
## Other policies
The same `--peft.*` flags work for any pre-trained policy. Some policies (SmolVLA, π₀, π₀.₅) ship
built-in default `target_modules`, so `--peft.method_type=LORA` is enough. Others do not, and will
ask you to pass `--peft.target_modules` explicitly — for example LingBot-VA, whose recommended
targets are documented in its [dedicated guide](./lingbot_va#training--fine-tuning).
+9 -15
View File
@@ -205,30 +205,24 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC): # type: igno
f"{CONFIG_NAME} not found on the HuggingFace Hub in {model_id}"
) from e
# HACK: Parse the original config to get the config subclass, so that we can
# apply cli overrides.
# This is very ugly, ideally we'd like to be able to do that natively with draccus
# something like --policy.path (in addition to --policy.type)
with draccus.config_type("json"):
orig_config = draccus.parse(cls, config_file, args=[])
if config_file is None:
raise FileNotFoundError(f"{CONFIG_NAME} not found in {model_id}")
with open(config_file) as f:
config = json.load(f)
# Resolve the concrete config subclass from the serialized "type" tag, then parse
# the config (with CLI overrides) directly for that class. The "type" key is
# stripped because draccus only consumes it when parsing the registry base class.
policy_type = config.pop("type", None)
if policy_type is None:
raise ValueError(f"Missing 'type' field in {CONFIG_NAME} of {model_id}")
try:
config_cls = cls.get_choice_class(policy_type)
except Exception as e:
raise ValueError(
f"Policy type '{policy_type}' (from {CONFIG_NAME} of {model_id}) is not registered. "
f"Available policy types: {cls.get_known_choices()}"
) from e
config.pop("type")
with tempfile.NamedTemporaryFile("w+", delete=False, suffix=".json") as f:
json.dump(config, f)
config_file = f.name
cli_overrides = policy_kwargs.pop("cli_overrides", [])
with draccus.config_type("json"):
return draccus.parse(config_cls, config_file, args=cli_overrides)
return draccus.parse(orig_config.__class__, config_file, args=cli_overrides)
+5
View File
@@ -221,6 +221,11 @@ class TrainPipelineConfig(HubMixin):
)
active_cfg = self.trainable_config
# Keep the policy-level `use_peft` flag in sync with the presence of a `--peft.*` config.
if self.peft is not None and self.policy is not None:
self.policy.use_peft = True
if self.rename_map and active_cfg.pretrained_path is None:
raise ValueError(
"`rename_map` requires a pretrained policy checkpoint. "
-2
View File
@@ -32,7 +32,6 @@ from .pretrained import PreTrainedPolicy as PreTrainedPolicy
from .smolvla.configuration_smolvla import SmolVLAConfig as SmolVLAConfig
from .tdmpc.configuration_tdmpc import TDMPCConfig as TDMPCConfig
from .utils import make_robot_action, prepare_observation_for_inference
from .vla_jepa.configuration_vla_jepa import VLAJEPAConfig as VLAJEPAConfig
from .vqbet.configuration_vqbet import VQBeTConfig as VQBeTConfig
from .wall_x.configuration_wall_x import WallXConfig as WallXConfig
from .xvla.configuration_xvla import XVLAConfig as XVLAConfig
@@ -58,7 +57,6 @@ __all__ = [
"PI05Config",
"SmolVLAConfig",
"TDMPCConfig",
"VLAJEPAConfig",
"VQBeTConfig",
"WallXConfig",
"XVLAConfig",
+39 -2
View File
@@ -18,10 +18,17 @@ from typing import Any
import torch
from lerobot.processor import (
AddBatchDimensionProcessorStep,
DeviceProcessorStep,
NormalizerProcessorStep,
PolicyAction,
PolicyProcessorPipeline,
make_default_pre_post_processors,
RenameObservationsProcessorStep,
UnnormalizerProcessorStep,
policy_action_to_transition,
transition_to_policy_action,
)
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
from .configuration_act import ACTConfig
@@ -47,4 +54,34 @@ def make_act_pre_post_processors(
tuple[PolicyProcessorPipeline[dict[str, Any], dict[str, Any]], PolicyProcessorPipeline[PolicyAction, PolicyAction]]: A tuple containing the
pre-processor pipeline and the post-processor pipeline.
"""
return make_default_pre_post_processors(config, dataset_stats, normalizer_device=config.device)
input_steps = [
RenameObservationsProcessorStep(rename_map={}),
AddBatchDimensionProcessorStep(),
DeviceProcessorStep(device=config.device),
NormalizerProcessorStep(
features={**config.input_features, **config.output_features},
norm_map=config.normalization_mapping,
stats=dataset_stats,
device=config.device,
),
]
output_steps = [
UnnormalizerProcessorStep(
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
),
DeviceProcessorStep(device="cpu"),
]
return (
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
steps=input_steps,
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
),
PolicyProcessorPipeline[PolicyAction, PolicyAction](
steps=output_steps,
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
to_transition=policy_action_to_transition,
to_output=transition_to_policy_action,
),
)
@@ -19,10 +19,17 @@ from typing import Any
import torch
from lerobot.processor import (
AddBatchDimensionProcessorStep,
DeviceProcessorStep,
NormalizerProcessorStep,
PolicyAction,
PolicyProcessorPipeline,
make_default_pre_post_processors,
RenameObservationsProcessorStep,
UnnormalizerProcessorStep,
policy_action_to_transition,
transition_to_policy_action,
)
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
from .configuration_diffusion import DiffusionConfig
@@ -56,4 +63,32 @@ def make_diffusion_pre_post_processors(
Returns:
A tuple containing the configured pre-processor and post-processor pipelines.
"""
return make_default_pre_post_processors(config, dataset_stats)
input_steps = [
RenameObservationsProcessorStep(rename_map={}),
AddBatchDimensionProcessorStep(),
DeviceProcessorStep(device=config.device),
NormalizerProcessorStep(
features={**config.input_features, **config.output_features},
norm_map=config.normalization_mapping,
stats=dataset_stats,
),
]
output_steps = [
UnnormalizerProcessorStep(
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
),
DeviceProcessorStep(device="cpu"),
]
return (
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
steps=input_steps,
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
),
PolicyProcessorPipeline[PolicyAction, PolicyAction](
steps=output_steps,
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
to_transition=policy_action_to_transition,
to_output=transition_to_policy_action,
),
)
+37 -12
View File
@@ -23,16 +23,24 @@ import torch
from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
from lerobot.processor import (
AddBatchDimensionProcessorStep,
ComplementaryDataProcessorStep,
DeviceProcessorStep,
NormalizerProcessorStep,
PolicyAction,
PolicyProcessorPipeline,
ProcessorStep,
ProcessorStepRegistry,
make_default_policy_processor_steps,
make_policy_processor_pipelines,
RenameObservationsProcessorStep,
UnnormalizerProcessorStep,
)
from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
from lerobot.types import TransitionKey
from lerobot.utils.constants import OBS_STATE
from lerobot.utils.constants import (
OBS_STATE,
POLICY_POSTPROCESSOR_DEFAULT_NAME,
POLICY_PREPROCESSOR_DEFAULT_NAME,
)
from lerobot.utils.import_utils import _transformers_available, require_package
from .configuration_eo1 import EO1Config
@@ -234,12 +242,14 @@ def make_eo1_pre_post_processors(
]:
"""Build pre/post processor pipelines for EO1."""
steps = make_default_policy_processor_steps(config, dataset_stats)
input_steps: list[ProcessorStep] = [
steps.rename_observations,
steps.add_batch_dim,
steps.normalize,
RenameObservationsProcessorStep(rename_map={}),
AddBatchDimensionProcessorStep(),
NormalizerProcessorStep(
features={**config.input_features, **config.output_features},
norm_map=config.normalization_mapping,
stats=dataset_stats,
),
EO1ConversationTemplateStep(input_features=config.input_features, chunk_size=config.chunk_size),
EO1QwenProcessorStep(
processor_name=config.vlm_base,
@@ -247,12 +257,27 @@ def make_eo1_pre_post_processors(
image_max_pixels=config.image_max_pixels,
use_fast_processor=config.use_fast_processor,
),
steps.to_device,
DeviceProcessorStep(device=config.device),
]
output_steps: list[ProcessorStep] = [
steps.unnormalize,
steps.to_cpu,
UnnormalizerProcessorStep(
features=config.output_features,
norm_map=config.normalization_mapping,
stats=dataset_stats,
),
DeviceProcessorStep(device="cpu"),
]
return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
return (
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
steps=input_steps,
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
),
PolicyProcessorPipeline[PolicyAction, PolicyAction](
steps=output_steps,
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
to_transition=policy_action_to_transition,
to_output=transition_to_policy_action,
),
)
+363 -69
View File
@@ -17,7 +17,6 @@
from __future__ import annotations
import importlib
import inspect
import logging
from typing import TYPE_CHECKING, Any, TypedDict, Unpack
@@ -45,10 +44,26 @@ from lerobot.utils.constants import (
)
from lerobot.utils.feature_utils import dataset_to_policy_features
from .act.configuration_act import ACTConfig
from .diffusion.configuration_diffusion import DiffusionConfig
from .eo1.configuration_eo1 import EO1Config
from .evo1.configuration_evo1 import Evo1Config
from .fastwam.configuration_fastwam import FastWAMConfig
from .gaussian_actor.configuration_gaussian_actor import GaussianActorConfig
from .groot.configuration_groot import GrootConfig
from .lingbot_va.configuration_lingbot_va import LingBotVAConfig
from .molmoact2.configuration_molmoact2 import MolmoAct2Config
from .multi_task_dit.configuration_multi_task_dit import MultiTaskDiTConfig
from .pi0.configuration_pi0 import PI0Config
from .pi05.configuration_pi05 import PI05Config
from .pretrained import PreTrainedPolicy
from .smolvla.configuration_smolvla import SmolVLAConfig
from .tdmpc.configuration_tdmpc import TDMPCConfig
from .utils import validate_visual_features_consistency
from .vla_jepa.configuration_vla_jepa import VLAJEPAConfig
from .vqbet.configuration_vqbet import VQBeTConfig
from .wall_x.configuration_wall_x import WallXConfig
from .xvla.configuration_xvla import XVLAConfig
def _reconnect_relative_absolute_steps(
@@ -73,23 +88,100 @@ def get_policy_class(name: str) -> type[PreTrainedPolicy]:
"""
Retrieves a policy class by its registered name.
Resolution is convention-based: the draccus-registered config class of ``name`` is
looked up, its ``configuration_*`` module path is rewritten to ``modeling_*``, and
the ``<X>Policy`` class is imported from there. The modeling module is only imported
at call time, keeping heavy optional dependencies lazy. This works for both built-in
policies and third-party lerobot plugins (anything registered via
``@PreTrainedConfig.register_subclass``).
This function uses dynamic imports to avoid loading all policy classes into memory
at once, improving startup time and reducing dependencies.
Args:
name: The registered name of the policy (e.g. "act", "diffusion", "pi0").
name: The name of the policy. Supported names are "tdmpc", "diffusion", "act",
"multi_task_dit", "vqbet", "pi0", "pi05", "gaussian_actor", "smolvla", "wall_x",
"molmoact2", "eo1", "evo1".
Returns:
The policy class corresponding to the given name.
Raises:
ValueError: If the policy name is not registered.
ImportError: If the policy's optional dependencies are not installed.
NotImplementedError: If the policy name is not recognized.
"""
return _get_policy_cls_from_policy_name(name=name)
if name == "tdmpc":
from .tdmpc.modeling_tdmpc import TDMPCPolicy
return TDMPCPolicy
elif name == "diffusion":
from .diffusion.modeling_diffusion import DiffusionPolicy
return DiffusionPolicy
elif name == "act":
from .act.modeling_act import ACTPolicy
return ACTPolicy
elif name == "multi_task_dit":
from .multi_task_dit.modeling_multi_task_dit import MultiTaskDiTPolicy
return MultiTaskDiTPolicy
elif name == "vqbet":
from .vqbet.modeling_vqbet import VQBeTPolicy
return VQBeTPolicy
elif name == "pi0":
from .pi0.modeling_pi0 import PI0Policy
return PI0Policy
elif name == "pi0_fast":
from .pi0_fast.modeling_pi0_fast import PI0FastPolicy
return PI0FastPolicy
elif name == "pi05":
from .pi05.modeling_pi05 import PI05Policy
return PI05Policy
elif name == "gaussian_actor":
from .gaussian_actor.modeling_gaussian_actor import GaussianActorPolicy
return GaussianActorPolicy
elif name == "smolvla":
from .smolvla.modeling_smolvla import SmolVLAPolicy
return SmolVLAPolicy
elif name == "groot":
from .groot.modeling_groot import GrootPolicy
return GrootPolicy
elif name == "xvla":
from .xvla.modeling_xvla import XVLAPolicy
return XVLAPolicy
elif name == "wall_x":
from .wall_x.modeling_wall_x import WallXPolicy
return WallXPolicy
elif name == "eo1":
from .eo1.modeling_eo1 import EO1Policy
return EO1Policy
elif name == "molmoact2":
from .molmoact2.modeling_molmoact2 import MolmoAct2Policy
return MolmoAct2Policy
elif name == "vla_jepa":
from .vla_jepa.modeling_vla_jepa import VLAJEPAPolicy
return VLAJEPAPolicy
elif name == "lingbot_va":
from .lingbot_va.modeling_lingbot_va import LingBotVAPolicy
return LingBotVAPolicy
elif name == "fastwam":
from .fastwam.modeling_fastwam import FastWAMPolicy
return FastWAMPolicy
elif name == "evo1":
from .evo1.modeling_evo1 import Evo1Policy
return Evo1Policy
else:
try:
return _get_policy_cls_from_policy_name(name=name)
except Exception as e:
raise ValueError(f"Policy type '{name}' is not available.") from e
def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
@@ -100,8 +192,9 @@ def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
mapping a string identifier to the corresponding config class.
Args:
policy_type: The registered type of the policy (any name registered via
``@PreTrainedConfig.register_subclass``, e.g. "act", "diffusion", "pi0").
policy_type: The type of the policy. Supported types include "tdmpc",
"multi_task_dit", "diffusion", "act", "vqbet", "pi0", "pi05", "gaussian_actor",
"smolvla", "wall_x", "molmoact2", "eo1", "evo1".
**kwargs: Keyword arguments to be passed to the configuration class constructor.
Returns:
@@ -110,11 +203,48 @@ def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
Raises:
ValueError: If the `policy_type` is not recognized.
"""
try:
config_cls = PreTrainedConfig.get_choice_class(policy_type)
except Exception as e:
raise ValueError(f"Policy type '{policy_type}' is not available.") from e
return config_cls(**kwargs)
if policy_type == "tdmpc":
return TDMPCConfig(**kwargs)
elif policy_type == "diffusion":
return DiffusionConfig(**kwargs)
elif policy_type == "act":
return ACTConfig(**kwargs)
elif policy_type == "multi_task_dit":
return MultiTaskDiTConfig(**kwargs)
elif policy_type == "vqbet":
return VQBeTConfig(**kwargs)
elif policy_type == "pi0":
return PI0Config(**kwargs)
elif policy_type == "pi05":
return PI05Config(**kwargs)
elif policy_type == "gaussian_actor":
return GaussianActorConfig(**kwargs)
elif policy_type == "smolvla":
return SmolVLAConfig(**kwargs)
elif policy_type == "groot":
return GrootConfig(**kwargs)
elif policy_type == "xvla":
return XVLAConfig(**kwargs)
elif policy_type == "wall_x":
return WallXConfig(**kwargs)
elif policy_type == "eo1":
return EO1Config(**kwargs)
elif policy_type == "molmoact2":
return MolmoAct2Config(**kwargs)
elif policy_type == "vla_jepa":
return VLAJEPAConfig(**kwargs)
elif policy_type == "lingbot_va":
return LingBotVAConfig(**kwargs)
elif policy_type == "fastwam":
return FastWAMConfig(**kwargs)
elif policy_type == "evo1":
return Evo1Config(**kwargs)
else:
try:
config_cls = PreTrainedConfig.get_choice_class(policy_type)
return config_cls(**kwargs)
except Exception as e:
raise ValueError(f"Policy type '{policy_type}' is not available.") from e
class ProcessorConfigKwargs(TypedDict, total=False):
@@ -168,7 +298,8 @@ def make_pre_post_processors(
A tuple containing the input (pre-processor) and output (post-processor) pipelines.
Raises:
ValueError: If no processor factory exists for the given policy configuration type.
NotImplementedError: If a processor factory is not implemented for the given
policy configuration type.
"""
if pretrained_path:
if isinstance(policy_cfg, GrootConfig):
@@ -220,13 +351,197 @@ def make_pre_post_processors(
)
return preprocessor, postprocessor
# Create new processors from the policy config, resolving the per-policy factory
# function by naming convention (lazy import keeps optional dependencies optional).
return _make_processors_from_policy_config(
config=policy_cfg,
dataset_stats=kwargs.get("dataset_stats"),
dataset_meta=kwargs.get("dataset_meta"),
)
# Create a new processor based on policy type
if isinstance(policy_cfg, TDMPCConfig):
from .tdmpc.processor_tdmpc import make_tdmpc_pre_post_processors
processors = make_tdmpc_pre_post_processors(
config=policy_cfg,
dataset_stats=kwargs.get("dataset_stats"),
)
elif isinstance(policy_cfg, DiffusionConfig):
from .diffusion.processor_diffusion import make_diffusion_pre_post_processors
processors = make_diffusion_pre_post_processors(
config=policy_cfg,
dataset_stats=kwargs.get("dataset_stats"),
)
elif isinstance(policy_cfg, ACTConfig):
from .act.processor_act import make_act_pre_post_processors
processors = make_act_pre_post_processors(
config=policy_cfg,
dataset_stats=kwargs.get("dataset_stats"),
)
elif isinstance(policy_cfg, MultiTaskDiTConfig):
from .multi_task_dit.processor_multi_task_dit import (
make_multi_task_dit_pre_post_processors,
)
processors = make_multi_task_dit_pre_post_processors(
config=policy_cfg,
dataset_stats=kwargs.get("dataset_stats"),
)
elif isinstance(policy_cfg, VQBeTConfig):
from .vqbet.processor_vqbet import make_vqbet_pre_post_processors
processors = make_vqbet_pre_post_processors(
config=policy_cfg,
dataset_stats=kwargs.get("dataset_stats"),
)
elif isinstance(policy_cfg, PI0Config):
from .pi0.processor_pi0 import make_pi0_pre_post_processors
processors = make_pi0_pre_post_processors(
config=policy_cfg,
dataset_stats=kwargs.get("dataset_stats"),
)
elif isinstance(policy_cfg, PI05Config):
from .pi05.processor_pi05 import make_pi05_pre_post_processors
processors = make_pi05_pre_post_processors(
config=policy_cfg,
dataset_stats=kwargs.get("dataset_stats"),
)
elif isinstance(policy_cfg, GaussianActorConfig):
from .gaussian_actor.processor_gaussian_actor import make_gaussian_actor_pre_post_processors
processors = make_gaussian_actor_pre_post_processors(
config=policy_cfg,
dataset_stats=kwargs.get("dataset_stats"),
)
elif isinstance(policy_cfg, SmolVLAConfig):
from .smolvla.processor_smolvla import make_smolvla_pre_post_processors
processors = make_smolvla_pre_post_processors(
config=policy_cfg,
dataset_stats=kwargs.get("dataset_stats"),
)
elif isinstance(policy_cfg, GrootConfig):
from .groot.processor_groot import make_groot_pre_post_processors
processors = make_groot_pre_post_processors(
config=policy_cfg,
dataset_stats=kwargs.get("dataset_stats"),
dataset_meta=kwargs.get("dataset_meta"),
)
elif isinstance(policy_cfg, XVLAConfig):
from .xvla.processor_xvla import (
make_xvla_pre_post_processors,
)
processors = make_xvla_pre_post_processors(
config=policy_cfg,
dataset_stats=kwargs.get("dataset_stats"),
)
elif isinstance(policy_cfg, WallXConfig):
from .wall_x.processor_wall_x import make_wall_x_pre_post_processors
processors = make_wall_x_pre_post_processors(
config=policy_cfg,
dataset_stats=kwargs.get("dataset_stats"),
)
elif isinstance(policy_cfg, EO1Config):
from .eo1.processor_eo1 import make_eo1_pre_post_processors
processors = make_eo1_pre_post_processors(
config=policy_cfg,
dataset_stats=kwargs.get("dataset_stats"),
)
elif isinstance(policy_cfg, Evo1Config):
from .evo1.processor_evo1 import make_evo1_pre_post_processors
processors = make_evo1_pre_post_processors(
config=policy_cfg,
dataset_stats=kwargs.get("dataset_stats"),
)
elif isinstance(policy_cfg, MolmoAct2Config):
from .molmoact2.processor_molmoact2 import make_molmoact2_pre_post_processors
processors = make_molmoact2_pre_post_processors(
config=policy_cfg,
dataset_stats=kwargs.get("dataset_stats"),
dataset_meta=kwargs.get("dataset_meta"),
)
elif isinstance(policy_cfg, VLAJEPAConfig):
from .vla_jepa.processor_vla_jepa import make_vla_jepa_pre_post_processors
processors = make_vla_jepa_pre_post_processors(
config=policy_cfg,
dataset_stats=kwargs.get("dataset_stats"),
)
elif isinstance(policy_cfg, LingBotVAConfig):
from .lingbot_va.processor_lingbot_va import make_lingbot_va_pre_post_processors
processors = make_lingbot_va_pre_post_processors(
config=policy_cfg,
dataset_stats=kwargs.get("dataset_stats"),
)
elif isinstance(policy_cfg, FastWAMConfig):
from .fastwam.processor_fastwam import make_fastwam_pre_post_processors
processors = make_fastwam_pre_post_processors(
config=policy_cfg,
dataset_stats=kwargs.get("dataset_stats"),
)
else:
try:
processors = _make_processors_from_policy_config(
config=policy_cfg,
dataset_stats=kwargs.get("dataset_stats"),
)
except Exception as e:
raise ValueError(f"Processor for policy type '{policy_cfg.type}' is not implemented.") from e
return processors
def _has_peft_adapter_config(
pretrained_path: str,
revision: str | None = None,
) -> bool:
"""Return whether ``pretrained_path`` points to an existing PEFT adapter.
A PEFT adapter checkpoint always ships an ``adapter_config.json``.
A plain base-model checkpoint does not. This distinction lets us tell apart
two very different ``use_peft=True`` scenarios that both set ``pretrained_path``:
* loading/resuming a previously trained adapter (config lives at ``pretrained_path``)
* starting a *fresh* PEFT fine-tune on top of a base model
Works for both local directories and Hub repo ids.
"""
import os
adapter_config_name = "adapter_config.json"
if os.path.isdir(pretrained_path):
return os.path.isfile(os.path.join(pretrained_path, adapter_config_name))
from huggingface_hub import file_exists
from huggingface_hub.errors import HfHubHTTPError
try:
return file_exists(pretrained_path, adapter_config_name, revision=revision)
except (HfHubHTTPError, OSError):
return False
def make_policy(
@@ -323,13 +638,24 @@ def make_policy(
"the PEFT config parameters to be set. For training with PEFT, see `lerobot_train.py` on how to do that."
)
if cfg.pretrained_path and not cfg.use_peft:
# When `use_peft=True` and a checkpoint is given, the checkpoint can be one of two things:
# 1. A base model checkpoint (e.g., a pretrained policy) on which we want to start a fresh PEFT fine-tune.
# 2. A PEFT adapter checkpoint (e.g., a previously trained PEFT adapter)
# We distinguish between these two cases
load_existing_adapter = (
cfg.pretrained_path
and cfg.use_peft
and _has_peft_adapter_config(str(cfg.pretrained_path), cfg.pretrained_revision)
)
if cfg.pretrained_path and not load_existing_adapter:
# Load a pretrained policy and override the config if needed (for example, if there are inference-time
# hyperparameters that we want to vary).
# hyperparameters that we want to vary). This also covers starting a fresh PEFT fine-tune on top of a
# base model: the base weights are loaded here and `wrap_with_peft` builds the adapter afterwards.
kwargs["pretrained_name_or_path"] = cfg.pretrained_path
kwargs["revision"] = cfg.pretrained_revision
policy = policy_cls.from_pretrained(**kwargs)
elif cfg.pretrained_path and cfg.use_peft:
elif load_existing_adapter:
# Load a pretrained PEFT model on top of the policy. The pretrained path points to the folder/repo
# of the adapter and the adapter's config contains the path to the base policy. So we need the
# adapter config first, then load the correct policy and then apply PEFT.
@@ -370,12 +696,10 @@ def make_policy(
return policy
def _get_policy_cls_from_policy_name(name: str) -> type[PreTrainedPolicy]:
def _get_policy_cls_from_policy_name(name: str) -> type[PreTrainedConfig]:
"""Get policy class from its registered name using dynamic imports.
Works for built-in policies and 3rd party lerobot plugins alike: the config class
registered under ``name`` is resolved via the draccus ChoiceRegistry, and the policy
class is imported from the sibling ``modeling_*`` module by naming convention.
This is used as a helper function to import policies from 3rd party lerobot plugins.
Args:
name: The name of the policy.
@@ -401,39 +725,22 @@ def _get_policy_cls_from_policy_name(name: str) -> type[PreTrainedPolicy]:
"configuration_", "modeling_"
) # e.g., configuration_diffusion -> modeling_diffusion
try:
module = importlib.import_module(module_path)
except ModuleNotFoundError as e:
if e.name == module_path:
# The modeling_* module itself does not exist for this policy type. A missing
# optional dependency inside an existing module propagates unchanged instead,
# so its actionable install hint stays visible.
raise ValueError(f"Policy class for '{name}' is not implemented.") from e
raise
policy_cls = getattr(module, cls_name, None)
if policy_cls is None:
raise ValueError(
f"Policy class '{cls_name}' not found in '{module_path}'. "
f"Policies must expose '<Name>Policy' in the sibling 'modeling_*' module by naming convention."
)
module = importlib.import_module(module_path)
policy_cls = getattr(module, cls_name)
return policy_cls
def _make_processors_from_policy_config(
config: PreTrainedConfig,
dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
dataset_meta: Any | None = None,
) -> tuple[Any, Any]:
"""Create pre- and post-processors from a policy configuration using dynamic imports.
Resolves ``make_{type}_pre_post_processors`` from the policy's ``processor_*`` module
by naming convention. Works for built-in policies and 3rd party lerobot plugins.
This is used as a helper function to import processor factories from 3rd party lerobot plugins.
Args:
config: The policy configuration object.
dataset_stats: Dataset statistics for normalization.
dataset_meta: Dataset metadata, forwarded only to factories that declare a
``dataset_meta`` parameter (e.g. groot, molmoact2).
Returns:
A tuple containing the input (pre-processor) and output (post-processor) pipelines.
"""
@@ -446,19 +753,6 @@ def _make_processors_from_policy_config(
logging.debug(
f"Instantiating pre/post processors using function '{function_name}' from module '{module_path}'"
)
try:
module = importlib.import_module(module_path)
except ModuleNotFoundError as e:
if e.name == module_path:
# The processor_* module itself does not exist for this policy type. A missing
# optional dependency inside an existing module propagates unchanged instead,
# so its actionable install hint stays visible.
raise ValueError(f"Processor for policy type '{policy_type}' is not implemented.") from e
raise
function = getattr(module, function_name, None)
if function is None:
raise ValueError(f"Processor for policy type '{policy_type}' is not implemented.")
call_kwargs: dict[str, Any] = {"dataset_stats": dataset_stats}
if "dataset_meta" in inspect.signature(function).parameters:
call_kwargs["dataset_meta"] = dataset_meta
return function(config, **call_kwargs)
module = importlib.import_module(module_path)
function = getattr(module, function_name)
return function(config, dataset_stats=dataset_stats)
@@ -22,11 +22,20 @@ import torch
from lerobot.configs import PipelineFeatureType, PolicyFeature
from lerobot.processor import (
ActionProcessorStep,
AddBatchDimensionProcessorStep,
DeviceProcessorStep,
NormalizerProcessorStep,
PolicyAction,
PolicyProcessorPipeline,
ProcessorStepRegistry,
make_default_policy_processor_steps,
make_policy_processor_pipelines,
RenameObservationsProcessorStep,
UnnormalizerProcessorStep,
policy_action_to_transition,
transition_to_policy_action,
)
from lerobot.utils.constants import (
POLICY_POSTPROCESSOR_DEFAULT_NAME,
POLICY_PREPROCESSOR_DEFAULT_NAME,
)
from .configuration_fastwam import FastWAMConfig
@@ -96,20 +105,38 @@ def make_fastwam_pre_post_processors(
# anyway) and unsafe across fine-tuning: its `resize_size` would be inherited from the base
# checkpoint's camera geometry, not this dataset's, making the concatenation N_cameras x too wide.
steps = make_default_policy_processor_steps(config, normalization_stats, normalizer_device=config.device)
input_steps = [
steps.rename_observations,
steps.add_batch_dim,
steps.to_device,
steps.normalize,
RenameObservationsProcessorStep(rename_map={}),
AddBatchDimensionProcessorStep(),
DeviceProcessorStep(device=config.device),
NormalizerProcessorStep(
features={**config.input_features, **config.output_features},
norm_map=config.normalization_mapping,
stats=normalization_stats,
device=config.device,
),
]
output_steps = [
steps.unnormalize,
UnnormalizerProcessorStep(
features=config.output_features,
norm_map=config.normalization_mapping,
stats=normalization_stats,
),
]
if config.toggle_action_dimensions:
output_steps.append(
FastWAMActionToggleProcessorStep(toggle_dimensions=config.toggle_action_dimensions)
)
output_steps.append(steps.to_cpu)
return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
output_steps.append(DeviceProcessorStep(device="cpu"))
return (
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
steps=input_steps,
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
),
PolicyProcessorPipeline[PolicyAction, PolicyAction](
steps=output_steps,
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
to_transition=policy_action_to_transition,
to_output=transition_to_policy_action,
),
)
@@ -20,10 +20,17 @@ from typing import Any
import torch
from lerobot.processor import (
AddBatchDimensionProcessorStep,
DeviceProcessorStep,
NormalizerProcessorStep,
PolicyAction,
PolicyProcessorPipeline,
make_default_pre_post_processors,
RenameObservationsProcessorStep,
UnnormalizerProcessorStep,
policy_action_to_transition,
transition_to_policy_action,
)
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
from .configuration_gaussian_actor import GaussianActorConfig
@@ -55,4 +62,33 @@ def make_gaussian_actor_pre_post_processors(
Returns:
A tuple containing the configured pre-processor and post-processor pipelines.
"""
return make_default_pre_post_processors(config, dataset_stats)
# Add remaining processors
input_steps = [
RenameObservationsProcessorStep(rename_map={}),
AddBatchDimensionProcessorStep(),
DeviceProcessorStep(device=config.device),
NormalizerProcessorStep(
features={**config.input_features, **config.output_features},
norm_map=config.normalization_mapping,
stats=dataset_stats,
),
]
output_steps = [
UnnormalizerProcessorStep(
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
),
DeviceProcessorStep(device="cpu"),
]
return (
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
steps=input_steps,
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
),
PolicyProcessorPipeline[PolicyAction, PolicyAction](
steps=output_steps,
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
to_transition=policy_action_to_transition,
to_output=transition_to_policy_action,
),
)
@@ -25,12 +25,19 @@ import torch
from lerobot.configs.types import FeatureType, NormalizationMode
from lerobot.processor import (
AddBatchDimensionProcessorStep,
DeviceProcessorStep,
NormalizerProcessorStep,
PolicyAction,
PolicyProcessorPipeline,
ProcessorStep,
RenameObservationsProcessorStep,
UnnormalizerProcessorStep,
make_default_policy_processor_steps,
make_policy_processor_pipelines,
)
from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
from lerobot.utils.constants import (
POLICY_POSTPROCESSOR_DEFAULT_NAME,
POLICY_PREPROCESSOR_DEFAULT_NAME,
)
from .configuration_lingbot_va import LingBotVAConfig
@@ -45,13 +52,15 @@ def make_lingbot_va_pre_post_processors(
]:
"""Build the pre/post processor pipelines for LingBot-VA."""
steps = make_default_policy_processor_steps(config, dataset_stats)
input_steps: list[ProcessorStep] = [
steps.rename_observations,
steps.add_batch_dim,
steps.normalize,
steps.to_device,
RenameObservationsProcessorStep(rename_map={}),
AddBatchDimensionProcessorStep(),
NormalizerProcessorStep(
features={**config.input_features, **config.output_features},
norm_map=config.normalization_mapping,
stats=dataset_stats,
),
DeviceProcessorStep(device=config.device),
]
# Unnormalize actions from [-1, 1] to physical units (QUANTILES) using q01/q99 restored from the checkpoint.
@@ -61,7 +70,18 @@ def make_lingbot_va_pre_post_processors(
norm_map={FeatureType.ACTION: NormalizationMode.QUANTILES},
stats=dataset_stats,
),
steps.to_cpu,
DeviceProcessorStep(device="cpu"),
]
return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
return (
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
steps=input_steps,
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
),
PolicyProcessorPipeline[PolicyAction, PolicyAction](
steps=output_steps,
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
to_transition=policy_action_to_transition,
to_output=transition_to_policy_action,
),
)
@@ -19,12 +19,18 @@ from typing import Any
import torch
from lerobot.processor import (
AddBatchDimensionProcessorStep,
DeviceProcessorStep,
NormalizerProcessorStep,
PolicyAction,
PolicyProcessorPipeline,
RenameObservationsProcessorStep,
TokenizerProcessorStep,
make_default_policy_processor_steps,
make_policy_processor_pipelines,
UnnormalizerProcessorStep,
policy_action_to_transition,
transition_to_policy_action,
)
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
from .configuration_multi_task_dit import MultiTaskDiTConfig
@@ -60,11 +66,9 @@ def make_multi_task_dit_pre_post_processors(
A tuple containing the configured pre-processor and post-processor pipelines.
"""
steps = make_default_policy_processor_steps(config, dataset_stats, normalizer_device=config.device)
input_steps = [
steps.rename_observations,
steps.add_batch_dim,
RenameObservationsProcessorStep(rename_map={}),
AddBatchDimensionProcessorStep(),
TokenizerProcessorStep(
tokenizer_name=config.text_encoder_name,
padding=config.tokenizer_padding,
@@ -72,12 +76,32 @@ def make_multi_task_dit_pre_post_processors(
max_length=config.tokenizer_max_length,
truncation=config.tokenizer_truncation,
),
steps.to_device,
steps.normalize,
DeviceProcessorStep(device=config.device),
NormalizerProcessorStep(
features={**config.input_features, **config.output_features},
norm_map=config.normalization_mapping,
stats=dataset_stats,
device=config.device,
),
]
output_steps = [
steps.unnormalize,
steps.to_cpu,
UnnormalizerProcessorStep(
features=config.output_features,
norm_map=config.normalization_mapping,
stats=dataset_stats,
),
DeviceProcessorStep(device="cpu"),
]
return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
return (
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
steps=input_steps,
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
),
PolicyProcessorPipeline[PolicyAction, PolicyAction](
steps=output_steps,
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
to_transition=policy_action_to_transition,
to_output=transition_to_policy_action,
),
)
+32 -11
View File
@@ -21,16 +21,22 @@ import torch
from lerobot.configs import PipelineFeatureType, PolicyFeature
from lerobot.processor import (
AbsoluteActionsProcessorStep,
AddBatchDimensionProcessorStep,
ComplementaryDataProcessorStep,
DeviceProcessorStep,
NormalizerProcessorStep,
PolicyAction,
PolicyProcessorPipeline,
ProcessorStep,
ProcessorStepRegistry,
RelativeActionsProcessorStep,
RenameObservationsProcessorStep,
TokenizerProcessorStep,
make_default_policy_processor_steps,
make_policy_processor_pipelines,
UnnormalizerProcessorStep,
policy_action_to_transition,
transition_to_policy_action,
)
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
from .configuration_pi0 import PI0Config
@@ -130,12 +136,10 @@ def make_pi0_pre_post_processors(
action_names=getattr(config, "action_feature_names", None),
)
steps = make_default_policy_processor_steps(config, dataset_stats)
# OpenPI order: raw → relative → normalize → model → unnormalize → absolute
input_steps: list[ProcessorStep] = [
steps.rename_observations, # To mimic the same processor as pretrained one
steps.add_batch_dim,
RenameObservationsProcessorStep(rename_map={}), # To mimic the same processor as pretrained one
AddBatchDimensionProcessorStep(),
Pi0NewLineProcessor(), # Add newlines before tokenization for PaliGemma
TokenizerProcessorStep(
tokenizer_name="google/paligemma-3b-pt-224",
@@ -143,15 +147,32 @@ def make_pi0_pre_post_processors(
padding_side="right",
padding="max_length",
),
steps.to_device,
DeviceProcessorStep(device=config.device),
relative_step,
steps.normalize,
NormalizerProcessorStep(
features={**config.input_features, **config.output_features},
norm_map=config.normalization_mapping,
stats=dataset_stats,
),
]
output_steps: list[ProcessorStep] = [
steps.unnormalize,
UnnormalizerProcessorStep(
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
),
AbsoluteActionsProcessorStep(enabled=config.use_relative_actions, relative_step=relative_step),
steps.to_cpu,
DeviceProcessorStep(device="cpu"),
]
return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
return (
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
steps=input_steps,
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
),
PolicyProcessorPipeline[PolicyAction, PolicyAction](
steps=output_steps,
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
to_transition=policy_action_to_transition,
to_output=transition_to_policy_action,
),
)
+36 -12
View File
@@ -24,17 +24,26 @@ import torch
from lerobot.configs import PipelineFeatureType, PolicyFeature
from lerobot.processor import (
AbsoluteActionsProcessorStep,
AddBatchDimensionProcessorStep,
DeviceProcessorStep,
NormalizerProcessorStep,
PolicyAction,
PolicyProcessorPipeline,
ProcessorStep,
ProcessorStepRegistry,
RelativeActionsProcessorStep,
RenameObservationsProcessorStep,
TokenizerProcessorStep,
make_default_policy_processor_steps,
make_policy_processor_pipelines,
UnnormalizerProcessorStep,
policy_action_to_transition,
transition_to_policy_action,
)
from lerobot.types import EnvTransition, TransitionKey
from lerobot.utils.constants import OBS_STATE
from lerobot.utils.constants import (
OBS_STATE,
POLICY_POSTPROCESSOR_DEFAULT_NAME,
POLICY_PREPROCESSOR_DEFAULT_NAME,
)
from .configuration_pi05 import PI05Config
@@ -126,16 +135,18 @@ def make_pi05_pre_post_processors(
action_names=getattr(config, "action_feature_names", None),
)
steps = make_default_policy_processor_steps(config, dataset_stats)
# OpenPI order: raw → relative → normalize → model → unnormalize → absolute
input_steps: list[ProcessorStep] = [
steps.rename_observations, # To mimic the same processor as pretrained one
steps.add_batch_dim,
RenameObservationsProcessorStep(rename_map={}), # To mimic the same processor as pretrained one
AddBatchDimensionProcessorStep(),
relative_step,
# NOTE: NormalizerProcessorStep MUST come before Pi05PrepareStateTokenizerProcessorStep
# because the tokenizer step expects normalized state in [-1, 1] range for discretization
steps.normalize,
NormalizerProcessorStep(
features={**config.input_features, **config.output_features},
norm_map=config.normalization_mapping,
stats=dataset_stats,
),
Pi05PrepareStateTokenizerProcessorStep(max_state_dim=config.max_state_dim),
TokenizerProcessorStep(
tokenizer_name="google/paligemma-3b-pt-224",
@@ -143,13 +154,26 @@ def make_pi05_pre_post_processors(
padding_side="right",
padding="max_length",
),
steps.to_device,
DeviceProcessorStep(device=config.device),
]
output_steps: list[ProcessorStep] = [
steps.unnormalize,
UnnormalizerProcessorStep(
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
),
AbsoluteActionsProcessorStep(enabled=config.use_relative_actions, relative_step=relative_step),
steps.to_cpu,
DeviceProcessorStep(device="cpu"),
]
return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
return (
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
steps=input_steps,
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
),
PolicyProcessorPipeline[PolicyAction, PolicyAction](
steps=output_steps,
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
to_transition=policy_action_to_transition,
to_output=transition_to_policy_action,
),
)
@@ -25,17 +25,26 @@ from lerobot.configs import PipelineFeatureType, PolicyFeature
from lerobot.processor import (
AbsoluteActionsProcessorStep,
ActionTokenizerProcessorStep,
AddBatchDimensionProcessorStep,
DeviceProcessorStep,
NormalizerProcessorStep,
PolicyAction,
PolicyProcessorPipeline,
ProcessorStep,
ProcessorStepRegistry,
RelativeActionsProcessorStep,
RenameObservationsProcessorStep,
TokenizerProcessorStep,
make_default_policy_processor_steps,
make_policy_processor_pipelines,
UnnormalizerProcessorStep,
policy_action_to_transition,
transition_to_policy_action,
)
from lerobot.types import EnvTransition, TransitionKey
from lerobot.utils.constants import OBS_STATE
from lerobot.utils.constants import (
OBS_STATE,
POLICY_POSTPROCESSOR_DEFAULT_NAME,
POLICY_PREPROCESSOR_DEFAULT_NAME,
)
from .configuration_pi0_fast import PI0FastConfig
@@ -126,8 +135,6 @@ def make_pi0_fast_pre_post_processors(
action_names=getattr(config, "action_feature_names", None),
)
steps = make_default_policy_processor_steps(config, dataset_stats)
# Pi0Fast order: relative → normalize → tokenize → model → unnormalize → absolute
# This matches pi0/pi0.5: RelativeActionsProcessorStep runs first on raw absolute actions,
# caching the raw state. NormalizerProcessorStep then normalizes the raw relative actions,
@@ -137,10 +144,14 @@ def make_pi0_fast_pre_post_processors(
# before Pi0FastPrepareStateAndLanguageTokenizerProcessorStep, so the state tokenizer
# continues to receive normalized state in [-1, 1] as expected.
input_steps: list[ProcessorStep] = [
steps.rename_observations, # To mimic the same processor as pretrained one
steps.add_batch_dim,
RenameObservationsProcessorStep(rename_map={}), # To mimic the same processor as pretrained one
AddBatchDimensionProcessorStep(),
relative_step,
steps.normalize,
NormalizerProcessorStep(
features={**config.input_features, **config.output_features},
norm_map=config.normalization_mapping,
stats=dataset_stats,
),
Pi0FastPrepareStateAndLanguageTokenizerProcessorStep(max_state_dim=config.max_state_dim),
TokenizerProcessorStep(
tokenizer_name=config.text_tokenizer_name,
@@ -154,13 +165,26 @@ def make_pi0_fast_pre_post_processors(
fast_skip_tokens=config.fast_skip_tokens,
paligemma_tokenizer_name=config.text_tokenizer_name,
),
steps.to_device,
DeviceProcessorStep(device=config.device),
]
output_steps: list[ProcessorStep] = [
steps.unnormalize,
UnnormalizerProcessorStep(
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
),
AbsoluteActionsProcessorStep(enabled=config.use_relative_actions, relative_step=relative_step),
steps.to_cpu,
DeviceProcessorStep(device="cpu"),
]
return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
return (
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
steps=input_steps,
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
),
PolicyProcessorPipeline[PolicyAction, PolicyAction](
steps=output_steps,
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
to_transition=policy_action_to_transition,
to_output=transition_to_policy_action,
),
)
@@ -19,13 +19,19 @@ from typing import Any
import torch
from lerobot.processor import (
AddBatchDimensionProcessorStep,
DeviceProcessorStep,
NewLineTaskProcessorStep,
NormalizerProcessorStep,
PolicyAction,
PolicyProcessorPipeline,
RenameObservationsProcessorStep,
TokenizerProcessorStep,
make_default_policy_processor_steps,
make_policy_processor_pipelines,
UnnormalizerProcessorStep,
policy_action_to_transition,
transition_to_policy_action,
)
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
from .configuration_smolvla import SmolVLAConfig
@@ -60,11 +66,9 @@ def make_smolvla_pre_post_processors(
A tuple containing the configured pre-processor and post-processor pipelines.
"""
steps = make_default_policy_processor_steps(config, dataset_stats)
input_steps = [
steps.rename_observations, # To mimic the same processor as pretrained one
steps.add_batch_dim,
RenameObservationsProcessorStep(rename_map={}), # To mimic the same processor as pretrained one
AddBatchDimensionProcessorStep(),
NewLineTaskProcessorStep(),
TokenizerProcessorStep(
tokenizer_name=config.vlm_model_name,
@@ -72,11 +76,28 @@ def make_smolvla_pre_post_processors(
padding_side="right",
max_length=config.tokenizer_max_length,
),
steps.to_device,
steps.normalize,
DeviceProcessorStep(device=config.device),
NormalizerProcessorStep(
features={**config.input_features, **config.output_features},
norm_map=config.normalization_mapping,
stats=dataset_stats,
),
]
output_steps = [
steps.unnormalize,
steps.to_cpu,
UnnormalizerProcessorStep(
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
),
DeviceProcessorStep(device="cpu"),
]
return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
return (
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
steps=input_steps,
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
),
PolicyProcessorPipeline[PolicyAction, PolicyAction](
steps=output_steps,
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
to_transition=policy_action_to_transition,
to_output=transition_to_policy_action,
),
)
+37 -2
View File
@@ -19,10 +19,17 @@ from typing import Any
import torch
from lerobot.processor import (
AddBatchDimensionProcessorStep,
DeviceProcessorStep,
NormalizerProcessorStep,
PolicyAction,
PolicyProcessorPipeline,
make_default_pre_post_processors,
RenameObservationsProcessorStep,
UnnormalizerProcessorStep,
policy_action_to_transition,
transition_to_policy_action,
)
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
from .configuration_tdmpc import TDMPCConfig
@@ -54,4 +61,32 @@ def make_tdmpc_pre_post_processors(
Returns:
A tuple containing the configured pre-processor and post-processor pipelines.
"""
return make_default_pre_post_processors(config, dataset_stats)
input_steps = [
RenameObservationsProcessorStep(rename_map={}),
AddBatchDimensionProcessorStep(),
DeviceProcessorStep(device=config.device),
NormalizerProcessorStep(
features={**config.input_features, **config.output_features},
norm_map=config.normalization_mapping,
stats=dataset_stats,
),
]
output_steps = [
UnnormalizerProcessorStep(
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
),
DeviceProcessorStep(device="cpu"),
]
return (
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
steps=input_steps,
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
),
PolicyProcessorPipeline[PolicyAction, PolicyAction](
steps=output_steps,
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
to_transition=policy_action_to_transition,
to_output=transition_to_policy_action,
),
)
@@ -20,16 +20,20 @@ import torch
from lerobot.policies.vla_jepa.configuration_vla_jepa import VLAJEPAConfig
from lerobot.processor import (
AddBatchDimensionProcessorStep,
DeviceProcessorStep,
EnvTransition,
NormalizerProcessorStep,
PolicyAction,
PolicyProcessorPipeline,
ProcessorStep,
ProcessorStepRegistry,
RenameObservationsProcessorStep,
TransitionKey,
UnnormalizerProcessorStep,
make_default_policy_processor_steps,
make_policy_processor_pipelines,
)
from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
@ProcessorStepRegistry.register(name="vla_jepa_clip_actions")
@@ -108,12 +112,15 @@ def make_vla_jepa_pre_post_processors(
PolicyProcessorPipeline[PolicyAction, PolicyAction],
]:
features = {**config.input_features, **config.output_features}
steps = make_default_policy_processor_steps(config, dataset_stats)
input_steps = [
steps.rename_observations,
steps.add_batch_dim,
steps.to_device,
steps.normalize,
RenameObservationsProcessorStep(rename_map={}),
AddBatchDimensionProcessorStep(),
DeviceProcessorStep(device=config.device),
NormalizerProcessorStep(
features=features,
norm_map=config.normalization_mapping,
stats=dataset_stats,
),
]
output_steps: list[ProcessorStep] = []
if config.clip_normalized_actions:
@@ -122,8 +129,6 @@ def make_vla_jepa_pre_post_processors(
output_steps.append(
PreSnapGripperProcessorStep(gripper_dim=config.gripper_dim, threshold=config.gripper_threshold)
)
# NOTE: unlike the default policy unnormalizer (output features only), VLA-JEPA
# unnormalizes over BOTH input and output features.
output_steps.append(
UnnormalizerProcessorStep(
features=features,
@@ -135,5 +140,16 @@ def make_vla_jepa_pre_post_processors(
output_steps.append(
BinarizeGripperProcessorStep(gripper_dim=config.gripper_dim, threshold=config.gripper_threshold)
)
output_steps.append(steps.to_cpu)
return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
output_steps.append(DeviceProcessorStep(device="cpu"))
return (
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
steps=input_steps,
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
),
PolicyProcessorPipeline[PolicyAction, PolicyAction](
steps=output_steps,
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
to_transition=policy_action_to_transition,
to_output=transition_to_policy_action,
),
)
+37 -2
View File
@@ -20,10 +20,17 @@ from typing import Any
import torch
from lerobot.processor import (
AddBatchDimensionProcessorStep,
DeviceProcessorStep,
NormalizerProcessorStep,
PolicyAction,
PolicyProcessorPipeline,
make_default_pre_post_processors,
RenameObservationsProcessorStep,
UnnormalizerProcessorStep,
policy_action_to_transition,
transition_to_policy_action,
)
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
from .configuration_vqbet import VQBeTConfig
@@ -55,4 +62,32 @@ def make_vqbet_pre_post_processors(
Returns:
A tuple containing the configured pre-processor and post-processor pipelines.
"""
return make_default_pre_post_processors(config, dataset_stats)
input_steps = [
RenameObservationsProcessorStep(rename_map={}), # Let the possibility to the user to rename the keys
AddBatchDimensionProcessorStep(),
DeviceProcessorStep(device=config.device),
NormalizerProcessorStep(
features={**config.input_features, **config.output_features},
norm_map=config.normalization_mapping,
stats=dataset_stats,
),
]
output_steps = [
UnnormalizerProcessorStep(
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
),
DeviceProcessorStep(device="cpu"),
]
return (
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
steps=input_steps,
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
),
PolicyProcessorPipeline[PolicyAction, PolicyAction](
steps=output_steps,
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
to_transition=policy_action_to_transition,
to_output=transition_to_policy_action,
),
)
+32 -11
View File
@@ -20,13 +20,19 @@ import torch
from lerobot.configs import PipelineFeatureType, PolicyFeature
from lerobot.processor import (
AddBatchDimensionProcessorStep,
ComplementaryDataProcessorStep,
DeviceProcessorStep,
NormalizerProcessorStep,
PolicyAction,
PolicyProcessorPipeline,
ProcessorStepRegistry,
make_default_policy_processor_steps,
make_policy_processor_pipelines,
RenameObservationsProcessorStep,
UnnormalizerProcessorStep,
policy_action_to_transition,
transition_to_policy_action,
)
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
from .configuration_wall_x import WallXConfig
@@ -59,22 +65,37 @@ def make_wall_x_pre_post_processors(
A tuple containing the configured pre-processor and post-processor pipelines
"""
steps = make_default_policy_processor_steps(config, dataset_stats)
input_steps = [
steps.rename_observations,
steps.add_batch_dim,
RenameObservationsProcessorStep(rename_map={}),
AddBatchDimensionProcessorStep(),
WallXTaskProcessor(), # Process task description
steps.normalize,
steps.to_device,
NormalizerProcessorStep(
features={**config.input_features, **config.output_features},
norm_map=config.normalization_mapping,
stats=dataset_stats,
),
DeviceProcessorStep(device=config.device),
]
output_steps = [
steps.unnormalize,
steps.to_cpu,
UnnormalizerProcessorStep(
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
),
DeviceProcessorStep(device="cpu"),
]
return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
return (
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
steps=input_steps,
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
),
PolicyProcessorPipeline[PolicyAction, PolicyAction](
steps=output_steps,
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
to_transition=policy_action_to_transition,
to_output=transition_to_policy_action,
),
)
@ProcessorStepRegistry.register(name="wall_x_task_processor")
+34 -11
View File
@@ -22,14 +22,19 @@ import torch
from lerobot.configs import PipelineFeatureType, PolicyFeature
from lerobot.processor import (
AddBatchDimensionProcessorStep,
DeviceProcessorStep,
NormalizerProcessorStep,
ObservationProcessorStep,
PolicyAction,
PolicyProcessorPipeline,
ProcessorStep,
ProcessorStepRegistry,
RenameObservationsProcessorStep,
TokenizerProcessorStep,
make_default_policy_processor_steps,
make_policy_processor_pipelines,
UnnormalizerProcessorStep,
policy_action_to_transition,
transition_to_policy_action,
)
from lerobot.types import EnvTransition, TransitionKey
from lerobot.utils.constants import (
@@ -37,6 +42,8 @@ from lerobot.utils.constants import (
OBS_IMAGES,
OBS_PREFIX,
OBS_STATE,
POLICY_POSTPROCESSOR_DEFAULT_NAME,
POLICY_PREPROCESSOR_DEFAULT_NAME,
)
from .configuration_xvla import XVLAConfig
@@ -54,11 +61,10 @@ def make_xvla_pre_post_processors(
Build the LeRobot processor pipelines for XVLA.
"""
steps = make_default_policy_processor_steps(config, dataset_stats)
features = {**config.input_features, **config.output_features}
input_steps = [
steps.rename_observations,
steps.add_batch_dim,
RenameObservationsProcessorStep(rename_map={}),
AddBatchDimensionProcessorStep(),
TokenizerProcessorStep(
tokenizer_name=config.tokenizer_name,
max_length=config.tokenizer_max_length,
@@ -68,15 +74,32 @@ def make_xvla_pre_post_processors(
XVLAImageToFloatProcessorStep(),
XVLAImageNetNormalizeProcessorStep(),
XVLAAddDomainIdProcessorStep(),
steps.to_device,
steps.normalize,
DeviceProcessorStep(device=config.device),
NormalizerProcessorStep(
features=features, norm_map=config.normalization_mapping, stats=dataset_stats
),
]
output_steps = [
steps.unnormalize,
steps.to_cpu,
UnnormalizerProcessorStep(
features=config.output_features,
norm_map=config.normalization_mapping,
stats=dataset_stats,
),
DeviceProcessorStep(device="cpu"),
]
return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
return (
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
steps=input_steps,
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
),
PolicyProcessorPipeline[PolicyAction, PolicyAction](
steps=output_steps,
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
to_transition=policy_action_to_transition,
to_output=transition_to_policy_action,
),
)
# Custom XVLA processor steps
-8
View File
@@ -42,14 +42,10 @@ from .delta_action_processor import MapDeltaActionToRobotActionStep, MapTensorTo
from .device_processor import DeviceProcessorStep
from .env_processor import IsaaclabArenaProcessorStep, LiberoProcessorStep
from .factory import (
DefaultPolicyProcessorSteps,
make_default_policy_processor_steps,
make_default_pre_post_processors,
make_default_processors,
make_default_robot_action_processor,
make_default_robot_observation_processor,
make_default_teleop_action_processor,
make_policy_processor_pipelines,
)
from .gym_action_processor import (
Numpy2TorchActionProcessorStep,
@@ -133,14 +129,10 @@ __all__ = [
"ImageCropResizeProcessorStep",
"InfoProcessorStep",
"InterventionActionProcessorStep",
"DefaultPolicyProcessorSteps",
"make_default_policy_processor_steps",
"make_default_pre_post_processors",
"make_default_processors",
"make_default_teleop_action_processor",
"make_default_robot_action_processor",
"make_default_robot_observation_processor",
"make_policy_processor_pipelines",
"AbsoluteActionsProcessorStep",
"RelativeActionsProcessorStep",
"MapDeltaActionToRobotActionStep",
+2 -114
View File
@@ -14,33 +14,15 @@
# See the License for the specific language governing permissions and
# limitations under the License.
from dataclasses import dataclass
from typing import Any
from lerobot.types import RobotAction, RobotObservation
import torch
from lerobot.configs.policies import PreTrainedConfig
from lerobot.types import PolicyAction, RobotAction, RobotObservation
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
from .batch_processor import AddBatchDimensionProcessorStep
from .converters import (
observation_to_transition,
policy_action_to_transition,
robot_action_observation_to_transition,
transition_to_observation,
transition_to_policy_action,
transition_to_robot_action,
)
from .device_processor import DeviceProcessorStep
from .normalize_processor import NormalizerProcessorStep, UnnormalizerProcessorStep
from .pipeline import (
IdentityProcessorStep,
PolicyProcessorPipeline,
ProcessorStep,
RobotProcessorPipeline,
)
from .rename_processor import RenameObservationsProcessorStep
from .pipeline import IdentityProcessorStep, RobotProcessorPipeline
def make_default_teleop_action_processor() -> RobotProcessorPipeline[
@@ -79,97 +61,3 @@ def make_default_processors():
robot_action_processor = make_default_robot_action_processor()
robot_observation_processor = make_default_robot_observation_processor()
return (teleop_action_processor, robot_action_processor, robot_observation_processor)
@dataclass
class DefaultPolicyProcessorSteps:
"""The canonical processor steps shared by most policies' pre/post pipelines.
Policies compose these in their own order (step ORDER is a Hub-serialized contract
and intentionally stays explicit per policy) and interleave their custom steps.
"""
rename_observations: RenameObservationsProcessorStep
add_batch_dim: AddBatchDimensionProcessorStep
to_device: DeviceProcessorStep
normalize: NormalizerProcessorStep
unnormalize: UnnormalizerProcessorStep
to_cpu: DeviceProcessorStep
def make_default_policy_processor_steps(
config: PreTrainedConfig,
dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
*,
normalizer_device: torch.device | str | None = None,
) -> DefaultPolicyProcessorSteps:
"""Construct the canonical policy processor steps from a policy config.
Args:
config: A `PreTrainedConfig` providing `device`, `input_features`,
`output_features` and `normalization_mapping`.
dataset_stats: Dataset statistics used for (un)normalization.
normalizer_device: Device passed to `NormalizerProcessorStep` (some policies pin
their normalization stats to the policy device; most leave it unset).
"""
return DefaultPolicyProcessorSteps(
rename_observations=RenameObservationsProcessorStep(rename_map={}),
add_batch_dim=AddBatchDimensionProcessorStep(),
to_device=DeviceProcessorStep(device=config.device),
normalize=NormalizerProcessorStep(
features={**config.input_features, **config.output_features},
norm_map=config.normalization_mapping,
stats=dataset_stats,
device=normalizer_device,
),
unnormalize=UnnormalizerProcessorStep(
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
),
to_cpu=DeviceProcessorStep(device="cpu"),
)
def make_policy_processor_pipelines(
input_steps: list[ProcessorStep],
output_steps: list[ProcessorStep],
) -> tuple[
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
PolicyProcessorPipeline[PolicyAction, PolicyAction],
]:
"""Wrap pre/post step lists into the canonical policy pipeline pair.
Uses the standard pipeline names (which determine the serialized JSON filenames on
the Hub) and the standard policy-action converters on the postprocessor.
"""
return (
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
steps=input_steps,
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
),
PolicyProcessorPipeline[PolicyAction, PolicyAction](
steps=output_steps,
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
to_transition=policy_action_to_transition,
to_output=transition_to_policy_action,
),
)
def make_default_pre_post_processors(
config: PreTrainedConfig,
dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
*,
normalizer_device: torch.device | str | None = None,
) -> tuple[
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
PolicyProcessorPipeline[PolicyAction, PolicyAction],
]:
"""The pure-scaffold policy pipeline pair: Rename -> Batch -> Device -> Normalize,
and Unnormalize -> Device(cpu). Policies with custom steps or a different step order
compose `make_default_policy_processor_steps` themselves instead.
"""
s = make_default_policy_processor_steps(config, dataset_stats, normalizer_device=normalizer_device)
return make_policy_processor_pipelines(
input_steps=[s.rename_observations, s.add_batch_dim, s.to_device, s.normalize],
output_steps=[s.unnormalize, s.to_cpu],
)
@@ -126,7 +126,7 @@ class RelativeActionsProcessorStep(ProcessorStep):
observation = transition.get(TransitionKey.OBSERVATION, {})
state = observation.get(OBS_STATE) if observation else None
# Always cache state for the paired AbsoluteActionsProcessorStep.
# Always cache state for the paired AbsoluteActionsProcessorStep
if state is not None:
self._last_state = state
@@ -146,11 +146,6 @@ class RelativeActionsProcessorStep(ProcessorStep):
"""Return the cached ``observation.state`` used as the reference point for relative/absolute action conversions."""
return self._last_state
def set_cached_state(self, state: torch.Tensor | None) -> None:
"""Override the cached anchor state, e.g. to re-pin a chunk's anchor after the
per-tick pipeline overwrote it (see ``SyncInferenceEngine``)."""
self._last_state = state
def get_config(self) -> dict[str, Any]:
return {
"enabled": self.enabled,
+11
View File
@@ -43,6 +43,7 @@ from lerobot.processor import (
make_default_processors,
rename_stats,
)
from lerobot.processor.relative_action_processor import RelativeActionsProcessorStep
from lerobot.robots import make_robot_from_config
from lerobot.teleoperators import Teleoperator, make_teleoperator_from_config
from lerobot.utils.feature_utils import combine_feature_dicts, hw_to_dataset_features
@@ -51,6 +52,7 @@ from .configs import BaseStrategyConfig, DAggerStrategyConfig, RolloutConfig
from .inference import (
InferenceEngine,
RTCInferenceConfig,
SyncInferenceConfig,
create_inference_engine,
)
from .robot_wrapper import ThreadSafeRobot
@@ -397,6 +399,15 @@ def build_rollout_context(
},
)
if isinstance(cfg.inference, SyncInferenceConfig) and any(
isinstance(step, RelativeActionsProcessorStep) and step.enabled
for step in getattr(preprocessor, "steps", ())
):
raise NotImplementedError(
"SyncInferenceEngine does not support policies with relative actions for now."
"Use --inference.type=rtc or remove relative action processor steps from the policy pipeline."
)
# --- 7. Inference strategy (needs policy + pre/post + hardware) --
logger.info(
"Creating inference engine (type=%s)...",
+14 -73
View File
@@ -24,21 +24,26 @@ import torch
from lerobot.policies.pretrained import PreTrainedPolicy
from lerobot.policies.utils import make_robot_action, prepare_observation_for_inference
from lerobot.processor import PolicyProcessorPipeline, RelativeActionsProcessorStep
from lerobot.processor import PolicyProcessorPipeline
from .base import InferenceEngine
logger = logging.getLogger(__name__)
# Relative-action support: a predicted chunk of offsets is anchored to the robot
# state at prediction time, but the sync engine reruns the pre/post pipeline every
# tick, so ``RelativeActionsProcessorStep`` would re-anchor cached actions to the
# current (moved) state and drift through the chunk. We pin the anchor per chunk:
# a probe on the policy's public ``predict_action_chunk`` flags the ticks that
# predict a fresh chunk; on the others the engine restores the anchor the relative
# step overwrote. ``select_action`` stays on the hot path, so per-tick side effects
# (e.g. LingBot-VA keyframe feedback) are preserved.
# TODO(Steven): support relative-action policies. The per-tick flow refreshes
# ``RelativeActionsProcessorStep._last_state`` every call, so cached chunk
# actions popped on later ticks get reanchored to the *current* robot state and
# absolute targets drift through the chunk. Relative-action policies are
# rejected at context-build time today; RTC postprocesses the whole chunk and
# is unaffected.
#
# Candidate fix: drive the policy via ``predict_action_chunk`` and serve a
# local FIFO of postprocessed actions. Eliminates drift by construction and
# saves per-tick pre/post work, but bypasses ``select_action`` — needs
# fallbacks for SAC (raises), ACT temporal ensembling (ensembler lives in
# ``select_action``), and Diffusion-family (obs-history queues populated as a
# side effect of ``select_action``).
class SyncInferenceEngine(InferenceEngine):
@@ -68,31 +73,6 @@ class SyncInferenceEngine(InferenceEngine):
self._task = task
self._device = torch.device(device or "cpu")
self._robot_type = robot_type
# Find an enabled RelativeActionsProcessorStep to pin its anchor per chunk
# (see module comment), mirroring the RTC engine.
self._relative_step = next(
(
s
for s in getattr(preprocessor, "steps", ())
if isinstance(s, RelativeActionsProcessorStep) and s.enabled
),
None,
)
# Set by the probe for the current tick / ever, respectively.
self._chunk_predicted = False
self._ever_predicted_chunk = False
self._original_predict_action_chunk = None # set while the probe is installed
if self._relative_step is not None:
# ``action_names`` is optional on the step; fill it lazily from the
# policy/dataset so the relative<->absolute mask is built correctly. This is
# a deliberate engine->step side effect (the step is configured by its consumer).
if self._relative_step.action_names is None:
cfg_names = getattr(policy.config, "action_feature_names", None)
self._relative_step.action_names = list(cfg_names) if cfg_names else list(ordered_action_keys)
self._install_chunk_probe()
logger.info("Relative actions enabled: chunk anchor pinned per predicted chunk")
logger.info(
"SyncInferenceEngine initialized (device=%s, action_keys=%d)",
self._device,
@@ -105,11 +85,6 @@ class SyncInferenceEngine(InferenceEngine):
def stop(self) -> None:
"""No background resources to stop."""
# Undo the probe so the policy object isn't left permanently patched
# (it may outlive this engine or be reused by another).
if self._original_predict_action_chunk is not None:
self._policy.predict_action_chunk = self._original_predict_action_chunk
self._original_predict_action_chunk = None
logger.info("SyncInferenceEngine stopped")
def reset(self) -> None:
@@ -118,27 +93,6 @@ class SyncInferenceEngine(InferenceEngine):
self._policy.reset()
self._preprocessor.reset()
self._postprocessor.reset()
# New episode: the next tick predicts a fresh chunk and re-anchors.
self._chunk_predicted = False
self._ever_predicted_chunk = False
def _install_chunk_probe(self) -> None:
"""Wrap the policy's public ``predict_action_chunk`` so we learn which ticks
predict a fresh chunk (when the anchor must advance) without introspecting any
private action queue. Chunking policies call it from ``select_action``.
Wraps whatever callable is currently bound (e.g. an already-``torch.compile``d
one, since ``build_rollout_context`` compiles before building the engine); undone
in ``stop()``."""
self._original_predict_action_chunk = self._policy.predict_action_chunk
inner = self._original_predict_action_chunk
def probe(*args, **kwargs):
self._chunk_predicted = True
self._ever_predicted_chunk = True
return inner(*args, **kwargs)
self._policy.predict_action_chunk = probe
def get_action(self, obs_frame: dict | None) -> torch.Tensor | None:
"""Run the full inference pipeline on ``obs_frame`` and return an action tensor."""
@@ -153,25 +107,12 @@ class SyncInferenceEngine(InferenceEngine):
if self._device.type == "cuda" and self._policy.config.use_amp
else nullcontext()
)
# Snapshot the chunk anchor before the preprocessor overwrites it with this
# tick's state; restore it below if this tick only served a cached action.
# ``clone`` so the snapshot survives even if the cached tensor is ever mutated
# in place (today it is only rebound, but the copy is cheap for a state vector).
anchor_before = None
if self._relative_step is not None:
cached = self._relative_step.get_cached_state()
anchor_before = cached.clone() if cached is not None else None
self._chunk_predicted = False
with torch.inference_mode(), autocast_ctx:
observation = prepare_observation_for_inference(
observation, self._device, self._task, self._robot_type
)
observation = self._preprocessor(observation)
action = self._policy.select_action(observation)
# Hold the anchor only for a chunking policy serving a cached action this
# tick; policies that never chunk or that recomputed keep refreshing.
if self._relative_step is not None and self._ever_predicted_chunk and not self._chunk_predicted:
self._relative_step.set_cached_state(anchor_before)
action = self._postprocessor(action)
action_tensor = action.squeeze(0).cpu()
@@ -0,0 +1,142 @@
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Regression tests for PEFT loading when the checkpoint is a base model (see issue #3975).
Starting a fresh LoRA/PEFT fine-tune points ``--policy.path`` at a *base model* (no
``adapter_config.json``) while also setting ``use_peft=True``. This must NOT be mistaken for
loading an existing PEFT adapter. These tests lock in the base-model vs. adapter distinction
made by ``lerobot.policies.factory._has_peft_adapter_config`` and the branch it drives in
``make_policy``. They are pure/fast (no network, no ``peft``), so they run in CI.
"""
import json
from unittest.mock import MagicMock, patch
import torch
from huggingface_hub.errors import HfHubHTTPError
from lerobot.policies.factory import _has_peft_adapter_config
def test_local_base_model_dir_has_no_adapter_config(tmp_path):
# A base-model checkpoint directory (only model weights, no adapter config).
(tmp_path / "model.safetensors").write_bytes(b"")
(tmp_path / "config.json").write_text("{}")
assert _has_peft_adapter_config(str(tmp_path)) is False
def test_local_adapter_dir_has_adapter_config(tmp_path):
(tmp_path / "adapter_config.json").write_text(json.dumps({"peft_type": "LORA"}))
(tmp_path / "adapter_model.safetensors").write_bytes(b"")
assert _has_peft_adapter_config(str(tmp_path)) is True
def test_hub_base_model_repo_has_no_adapter_config():
with patch("huggingface_hub.file_exists", return_value=False) as mock_exists:
assert _has_peft_adapter_config("lerobot/lingbot_va_base") is False
mock_exists.assert_called_once()
assert mock_exists.call_args.args[1] == "adapter_config.json"
def test_hub_adapter_repo_has_adapter_config():
with patch("huggingface_hub.file_exists", return_value=True):
assert _has_peft_adapter_config("some/adapter-repo", revision="main") is True
def test_hub_lookup_error_falls_back_to_base_model():
# Offline / private / transient Hub errors must not crash; treat as "not an adapter".
with patch(
"huggingface_hub.file_exists",
side_effect=HfHubHTTPError("boom", response=MagicMock()),
):
assert _has_peft_adapter_config("some/private-repo") is False
with patch("huggingface_hub.file_exists", side_effect=OSError("offline")):
assert _has_peft_adapter_config("some/repo") is False
def _make_dummy_policy_cfg(pretrained_path, use_peft):
cfg = MagicMock()
cfg.type = "act"
cfg.device = "cpu"
cfg.pretrained_path = pretrained_path
cfg.pretrained_revision = None
cfg.use_peft = use_peft
cfg.input_features = {}
cfg.output_features = {}
return cfg
@patch("lerobot.policies.factory.validate_visual_features_consistency")
@patch("lerobot.policies.factory.env_to_policy_features", return_value={})
@patch("lerobot.policies.factory.get_policy_class")
def test_make_policy_base_model_with_use_peft_loads_base_not_adapter(
mock_get_cls, _mock_features, _mock_validate
):
"""`use_peft=True` on a base model must load the base weights, not a PEFT adapter.
Before the #3975 fix this went down the ``PeftConfig.from_pretrained`` path and failed
looking for a non-existent ``adapter_config.json``.
"""
from lerobot.policies import factory
policy_cls = MagicMock()
loaded_policy = torch.nn.Linear(1, 1) # a real nn.Module so make_policy's assert passes
policy_cls.from_pretrained.return_value = loaded_policy
mock_get_cls.return_value = policy_cls
cfg = _make_dummy_policy_cfg(pretrained_path="lerobot/lingbot_va_base", use_peft=True)
env_cfg = MagicMock()
with patch.object(factory, "_has_peft_adapter_config", return_value=False) as mock_has_adapter:
policy = factory.make_policy(cfg=cfg, env_cfg=env_cfg)
mock_has_adapter.assert_called_once()
# Base model is loaded via the normal pretrained path...
policy_cls.from_pretrained.assert_called_once()
assert policy_cls.from_pretrained.call_args.kwargs["pretrained_name_or_path"] == (
"lerobot/lingbot_va_base"
)
# ...and PEFT adapter loading is NOT attempted (would need peft + adapter_config.json).
assert policy is loaded_policy
@patch("lerobot.policies.factory.validate_visual_features_consistency")
@patch("lerobot.policies.factory.env_to_policy_features", return_value={})
@patch("lerobot.policies.factory.get_policy_class")
def test_make_policy_existing_adapter_uses_peft_loading(mock_get_cls, _mock_features, _mock_validate):
"""A real adapter checkpoint (has ``adapter_config.json``) must go through PEFT loading."""
from lerobot.policies import factory
policy_cls = MagicMock()
mock_get_cls.return_value = policy_cls
cfg = _make_dummy_policy_cfg(pretrained_path="some/adapter-repo", use_peft=True)
env_cfg = MagicMock()
policy_cls.from_pretrained.return_value = torch.nn.Linear(1, 1)
fake_peft = MagicMock()
fake_peft_config = MagicMock()
fake_peft_config.base_model_name_or_path = "lerobot/lingbot_va_base"
fake_peft.PeftConfig.from_pretrained.return_value = fake_peft_config
fake_peft.PeftModel.from_pretrained.return_value = torch.nn.Linear(1, 1)
with (
patch.object(factory, "_has_peft_adapter_config", return_value=True),
patch.dict("sys.modules", {"peft": fake_peft}),
):
factory.make_policy(cfg=cfg, env_cfg=env_cfg)
fake_peft.PeftConfig.from_pretrained.assert_called_once_with("some/adapter-repo")
fake_peft.PeftModel.from_pretrained.assert_called_once()
-7
View File
@@ -346,10 +346,3 @@ def test_state_not_modified_by_relative_processor(dataset, action_dim):
result_state = result[TransitionKey.OBSERVATION][OBS_STATE]
torch.testing.assert_close(result_state, original_state)
def test_cached_anchor_not_in_config():
"""The cached anchor is ephemeral runtime state and must not leak into the config."""
step = RelativeActionsProcessorStep(enabled=True)
step.set_cached_state(torch.tensor([[1.0, 2.0, 3.0, 4.0]]))
assert set(step.get_config()) == {"enabled", "exclude_joints", "action_names"}
-222
View File
@@ -348,225 +348,3 @@ def test_rollout_context_fields():
field_names = {f.name for f in dataclasses.fields(RolloutContext)}
assert field_names == {"runtime", "hardware", "policy", "processors", "data"}
# ---------------------------------------------------------------------------
# Sync engine: relative-action anchoring (drift-free chunk execution)
# ---------------------------------------------------------------------------
_REL_ACTION_NAMES = ["j0.pos", "j1.pos", "j2.pos", "gripper.pos"]
_REL_ACTION_DIM = len(_REL_ACTION_NAMES)
def _relative_pre_post(exclude_joints=None):
"""Pre/post processors wrapping the real relative (caches anchor) and absolute
(relative + cached state) steps, mirroring what the sync engine feeds them."""
from lerobot.processor import (
AbsoluteActionsProcessorStep,
RelativeActionsProcessorStep,
TransitionKey,
create_transition,
)
from lerobot.utils.constants import OBS_STATE
relative_step = RelativeActionsProcessorStep(
enabled=True, exclude_joints=exclude_joints or [], action_names=list(_REL_ACTION_NAMES)
)
absolute_step = AbsoluteActionsProcessorStep(enabled=True, relative_step=relative_step)
class _Pre:
steps = [relative_step]
def __call__(self, observation):
# Run the relative step so it caches the anchor, then pass the batch through.
transition = create_transition(observation={OBS_STATE: observation[OBS_STATE]})
relative_step(transition)
return observation
def reset(self):
pass
class _Post:
def __call__(self, action):
transition = create_transition(action=action)
return absolute_step(transition)[TransitionKey.ACTION]
def reset(self):
pass
return _Pre(), _Post(), relative_step
def _fake_relative_policy(chunk_rel, n_action_steps, chunking=True):
"""Fake relative-action policy for the sync engine.
``chunking=True`` buffers a chunk and serves it one action per tick, calling the
public ``predict_action_chunk`` only on refill (pi0/fastwam/lingbot). ``False``
returns an action directly and never calls it. The engine's anchor probe keys off
that public call, so the fake routes through it rather than any private queue.
"""
from collections import deque
policy = MagicMock()
policy.config.use_amp = False
policy.config.action_feature_names = list(_REL_ACTION_NAMES)
state = {"predict_calls": 0}
queue = deque(maxlen=n_action_steps)
def predict_action_chunk(_batch=None, **_kwargs):
state["predict_calls"] += 1
return chunk_rel.unsqueeze(0) # [B=1, n, dim]
def select_action(_observation):
if not chunking:
return chunk_rel[0].unsqueeze(0)
if len(queue) == 0:
actions = policy.predict_action_chunk(_observation)
queue.extend(actions.transpose(0, 1)) # [n, 1, dim]
return queue.popleft()
policy.predict_action_chunk.side_effect = predict_action_chunk
policy.select_action.side_effect = select_action
policy.reset.side_effect = queue.clear
policy._predict_state = state
return policy
def _build_sync_engine(policy, pre, post):
from lerobot.rollout import SyncInferenceEngine
return SyncInferenceEngine(
policy=policy,
preprocessor=pre,
postprocessor=post,
dataset_features={"action": {"names": list(_REL_ACTION_NAMES)}},
ordered_action_keys=list(_REL_ACTION_NAMES),
task="test",
device="cpu",
robot_type="mock",
)
def _obs_frame(state_values):
import numpy as np
return {"observation.state": np.asarray(state_values, dtype=np.float32)}
def test_sync_relative_holds_anchor_across_chunk():
"""Every action popped within a chunk must anchor to the tick-0 state (no drift)."""
n = 4
# A distinct relative offset per chunk step so a wrong anchor would be visible.
chunk_rel = torch.stack([torch.full((_REL_ACTION_DIM,), 0.1 * (i + 1)) for i in range(n)])
pre, post, relative_step = _relative_pre_post()
policy = _fake_relative_policy(chunk_rel, n_action_steps=n)
engine = _build_sync_engine(policy, pre, post)
assert engine._relative_step is relative_step # introspection wired the step
s0 = [1.0, 2.0, 3.0, 4.0]
outputs = []
for tick in range(n):
# Feed a *different* state each tick; a drifting anchor would use it.
state = [v + tick for v in s0]
outputs.append(engine.get_action(_obs_frame(state)))
# Exactly one chunk was predicted across the n ticks.
assert policy._predict_state["predict_calls"] == 1
for tick in range(n):
expected = torch.tensor(s0) + chunk_rel[tick]
torch.testing.assert_close(outputs[tick], expected)
# Next tick empties the queue -> fresh chunk -> anchor advances to the new state.
s_next = [10.0, 20.0, 30.0, 40.0]
out = engine.get_action(_obs_frame(s_next))
assert policy._predict_state["predict_calls"] == 2
torch.testing.assert_close(out, torch.tensor(s_next) + chunk_rel[0])
# The anchor now reflects the fresh-chunk state, not the held one.
torch.testing.assert_close(relative_step.get_cached_state(), torch.tensor([s_next]))
def test_sync_relative_reset_reanchors_new_episode():
"""After ``reset()`` the first tick of the next episode anchors to the new state."""
n = 3
chunk_rel = torch.stack([torch.full((_REL_ACTION_DIM,), 0.2) for _ in range(n)])
pre, post, relative_step = _relative_pre_post()
policy = _fake_relative_policy(chunk_rel, n_action_steps=n)
engine = _build_sync_engine(policy, pre, post)
# Episode 1: one tick anchors to s0 and leaves cached actions in the queue.
engine.get_action(_obs_frame([1.0, 1.0, 1.0, 1.0]))
assert policy._predict_state["predict_calls"] == 1
engine.reset() # clears the queue and the per-episode chunk flags
# Episode 2: a fresh state must produce a fresh chunk anchored to that state,
# not carry over the previous episode's anchor.
s_new = [7.0, 8.0, 9.0, 10.0]
out = engine.get_action(_obs_frame(s_new))
assert policy._predict_state["predict_calls"] == 2
torch.testing.assert_close(out, torch.tensor(s_new) + chunk_rel[0])
torch.testing.assert_close(relative_step.get_cached_state(), torch.tensor([s_new]))
def test_sync_relative_non_chunking_policy_refreshes_every_tick():
"""A policy that never calls ``predict_action_chunk`` must not freeze the anchor."""
n = 3
chunk_rel = torch.stack([torch.full((_REL_ACTION_DIM,), 0.5) for _ in range(n)])
pre, post, _ = _relative_pre_post()
policy = _fake_relative_policy(chunk_rel, n_action_steps=n, chunking=False)
engine = _build_sync_engine(policy, pre, post)
s0 = [1.0, 1.0, 1.0, 1.0]
for tick in range(3):
state = [v + tick for v in s0]
out = engine.get_action(_obs_frame(state))
# Anchor must track the current state every tick (no chunk => no hold).
torch.testing.assert_close(out, torch.tensor(state) + chunk_rel[0])
assert policy._predict_state["predict_calls"] == 0
def test_sync_engine_no_relative_step_is_none():
"""Without an enabled relative step, the engine takes the plain select_action path."""
policy = MagicMock()
policy.config.use_amp = False
engine = _build_sync_engine(policy, MagicMock(steps=[]), MagicMock())
assert engine._relative_step is None
def test_sync_relative_exclude_joints_stay_absolute():
"""With ``exclude_joints``, excluded dims pass through absolute while the relative
dims still hold the tick-0 anchor across the chunk."""
n = 4
# Distinct offset per step *and* per dim so a wrong anchor or a wrong mask shows up.
chunk_rel = torch.stack([torch.full((_REL_ACTION_DIM,), 0.1 * (i + 1)) for i in range(n)])
pre, post, relative_step = _relative_pre_post(exclude_joints=["gripper"])
policy = _fake_relative_policy(chunk_rel, n_action_steps=n)
engine = _build_sync_engine(policy, pre, post)
# gripper (last dim) is kept absolute; j0..j2 are relative.
mask = torch.tensor([1.0, 1.0, 1.0, 0.0])
s0 = [1.0, 2.0, 3.0, 4.0]
outputs = []
for tick in range(n):
state = [v + tick for v in s0] # moving state; a drifting anchor would use it
outputs.append(engine.get_action(_obs_frame(state)))
assert policy._predict_state["predict_calls"] == 1 # one chunk held across n ticks
for tick in range(n):
# relative dims: anchor held at s0; excluded gripper dim: raw predicted value.
expected = chunk_rel[tick] + torch.tensor(s0) * mask
torch.testing.assert_close(outputs[tick], expected)
def test_sync_relative_stop_restores_policy_method():
"""``stop()`` un-patches the probe so the policy object isn't permanently modified."""
n = 3
chunk_rel = torch.stack([torch.full((_REL_ACTION_DIM,), 0.2) for _ in range(n)])
pre, post, _ = _relative_pre_post()
policy = _fake_relative_policy(chunk_rel, n_action_steps=n)
original = policy.predict_action_chunk
engine = _build_sync_engine(policy, pre, post)
assert policy.predict_action_chunk is not original # probe installed
engine.stop()
assert policy.predict_action_chunk is original # restored