mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-28 20:26:05 +00:00
Compare commits
1 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| a3755f07c5 |
@@ -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.
|
||||
|
||||
@@ -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 24–32 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.
|
||||
|
||||
|
||||
@@ -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).
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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. "
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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,
|
||||
),
|
||||
)
|
||||
|
||||
@@ -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
@@ -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,
|
||||
),
|
||||
)
|
||||
|
||||
@@ -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,
|
||||
),
|
||||
)
|
||||
|
||||
@@ -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,
|
||||
),
|
||||
)
|
||||
|
||||
@@ -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,
|
||||
),
|
||||
)
|
||||
|
||||
@@ -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,
|
||||
),
|
||||
)
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)...",
|
||||
|
||||
@@ -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()
|
||||
@@ -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"}
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user