mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-10 19:41:45 +00:00
Compare commits
9 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| e09aae26f9 | |||
| 01a029a05e | |||
| 4ff306e002 | |||
| 8c89d4d81d | |||
| ef32c04f44 | |||
| b2b710b268 | |||
| e40b58a8df | |||
| 3e538352ca | |||
| 8a74e0ac6d |
@@ -55,7 +55,7 @@ jobs:
|
||||
github.repository == 'huggingface/lerobot'
|
||||
permissions:
|
||||
contents: read
|
||||
uses: huggingface/doc-builder/.github/workflows/build_main_documentation.yml@2430c1ec91d04667414e2fa31ecfc36c153ea391 # main
|
||||
uses: huggingface/doc-builder/.github/workflows/build_main_documentation.yml@e60a538eea9817ab312196d0d233604b01697265 # main
|
||||
with:
|
||||
commit_sha: ${{ github.sha }}
|
||||
package: lerobot
|
||||
@@ -78,7 +78,7 @@ jobs:
|
||||
permissions:
|
||||
contents: read
|
||||
pull-requests: write
|
||||
uses: huggingface/doc-builder/.github/workflows/build_pr_documentation.yml@2430c1ec91d04667414e2fa31ecfc36c153ea391 # main
|
||||
uses: huggingface/doc-builder/.github/workflows/build_pr_documentation.yml@e60a538eea9817ab312196d0d233604b01697265 # main
|
||||
with:
|
||||
commit_sha: ${{ github.event.pull_request.head.sha }}
|
||||
pr_number: ${{ github.event.number }}
|
||||
|
||||
@@ -162,11 +162,11 @@ Preliminary LeRobot integration results (GR00T-LeRobot, `eval.n_episodes >= 50`
|
||||
|
||||
| Suite | Success rate | Checkpoint |
|
||||
| ---------------- | -----------: | ------------------------------------------------------------------------------------------------------------- |
|
||||
| LIBERO Spatial | 91% | [nvidia/gr00t17-lerobot-libero_spatial-640](https://huggingface.co/nvidia/gr00t17-lerobot-libero_spatial-640) |
|
||||
| LIBERO Object | 81% | [nvidia/gr00t17-lerobot-libero_object-640](https://huggingface.co/nvidia/gr00t17-lerobot-libero_object-640) |
|
||||
| LIBERO Goal | 97% | [nvidia/gr00t17-lerobot-libero_goal-640](https://huggingface.co/nvidia/gr00t17-lerobot-libero_goal-640) |
|
||||
| LIBERO 10 (Long) | 84% | [nvidia/gr00t17-lerobot-libero_10-640](https://huggingface.co/nvidia/gr00t17-lerobot-libero_10-640) |
|
||||
| **Average** | **88.25%** | |
|
||||
| LIBERO Spatial | 95% | [nvidia/gr00t17-lerobot-libero_spatial-640](https://huggingface.co/nvidia/gr00t17-lerobot-libero_spatial-640) |
|
||||
| LIBERO Object | 100% | [nvidia/gr00t17-lerobot-libero_object-640](https://huggingface.co/nvidia/gr00t17-lerobot-libero_object-640) |
|
||||
| LIBERO Goal | 98% | [nvidia/gr00t17-lerobot-libero_goal-640](https://huggingface.co/nvidia/gr00t17-lerobot-libero_goal-640) |
|
||||
| LIBERO 10 (Long) | 93% | [nvidia/gr00t17-lerobot-libero_10-640](https://huggingface.co/nvidia/gr00t17-lerobot-libero_10-640) |
|
||||
| **Average** | **96.5%** | |
|
||||
|
||||
```bash
|
||||
export MODEL_ID=your_trained_model_on_huggingface
|
||||
|
||||
+1
-1
@@ -25,7 +25,7 @@ discord = "https://discord.gg/s3KuuzsPFb"
|
||||
|
||||
[project]
|
||||
name = "lerobot"
|
||||
version = "0.6.0"
|
||||
version = "0.6.1"
|
||||
description = "🤗 LeRobot: State-of-the-art Machine Learning for Real-World Robotics in Pytorch"
|
||||
dynamic = ["readme"]
|
||||
license = { text = "Apache-2.0" }
|
||||
|
||||
@@ -29,8 +29,10 @@ from lerobot.configs import FeatureType, PreTrainedConfig
|
||||
from lerobot.envs import EnvConfig, env_to_policy_features
|
||||
from lerobot.processor import (
|
||||
AbsoluteActionsProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
PolicyProcessorPipeline,
|
||||
RelativeActionsProcessorStep,
|
||||
UnnormalizerProcessorStep,
|
||||
batch_to_transition,
|
||||
policy_action_to_transition,
|
||||
transition_to_batch,
|
||||
@@ -84,6 +86,52 @@ def _reconnect_relative_absolute_steps(
|
||||
step.relative_step = relative_step
|
||||
|
||||
|
||||
def _ensure_relative_actions(
|
||||
preprocessor: PolicyProcessorPipeline, postprocessor: PolicyProcessorPipeline, policy_cfg
|
||||
) -> None:
|
||||
"""Enable (or inject) the relative/absolute action steps in a loaded pipeline.
|
||||
|
||||
When loading from a pretrained checkpoint, the saved processor is authoritative. If the base
|
||||
predates the relative-action feature (e.g. FastWAM/LingBot bases) its pipeline has no
|
||||
RelativeActionsProcessorStep, so lerobot-train's override cannot enable one — those override
|
||||
keys are popped before `from_pretrained` (else it raises) and we reconstruct the steps here.
|
||||
Bases that DO ship the (disabled) steps (e.g. pi0/pi05) are simply flipped on. No-op unless
|
||||
``policy_cfg.use_relative_actions`` is set, so non-relative runs are untouched.
|
||||
"""
|
||||
if not getattr(policy_cfg, "use_relative_actions", False):
|
||||
return
|
||||
|
||||
exclude_joints = list(getattr(policy_cfg, "relative_exclude_joints", []) or [])
|
||||
action_names = getattr(policy_cfg, "action_feature_names", None)
|
||||
|
||||
pre_steps = list(preprocessor.steps)
|
||||
relative_step = next((s for s in pre_steps if isinstance(s, RelativeActionsProcessorStep)), None)
|
||||
if relative_step is None:
|
||||
relative_step = RelativeActionsProcessorStep(
|
||||
enabled=True, exclude_joints=exclude_joints, action_names=action_names
|
||||
)
|
||||
# Insert right before the normalizer (raw -> relative -> normalize); fall back to the front.
|
||||
idx = next((i for i, s in enumerate(pre_steps) if isinstance(s, NormalizerProcessorStep)), 0)
|
||||
pre_steps.insert(idx, relative_step)
|
||||
preprocessor.steps = pre_steps
|
||||
else:
|
||||
relative_step.enabled = True
|
||||
relative_step.exclude_joints = exclude_joints
|
||||
relative_step.action_names = action_names
|
||||
|
||||
post_steps = list(postprocessor.steps)
|
||||
absolute_step = next((s for s in post_steps if isinstance(s, AbsoluteActionsProcessorStep)), None)
|
||||
if absolute_step is None:
|
||||
absolute_step = AbsoluteActionsProcessorStep(enabled=True, relative_step=relative_step)
|
||||
# Insert right after the unnormalizer (unnormalize -> absolute); fall back to the front.
|
||||
idx = next((i for i, s in enumerate(post_steps) if isinstance(s, UnnormalizerProcessorStep)), -1)
|
||||
post_steps.insert(idx + 1, absolute_step)
|
||||
postprocessor.steps = post_steps
|
||||
else:
|
||||
absolute_step.enabled = True
|
||||
absolute_step.relative_step = relative_step
|
||||
|
||||
|
||||
def get_policy_class(name: str) -> type[PreTrainedPolicy]:
|
||||
"""
|
||||
Retrieves a policy class by its registered name.
|
||||
@@ -320,12 +368,20 @@ def make_pre_post_processors(
|
||||
),
|
||||
)
|
||||
|
||||
# The relative/absolute override keys only match if the saved base already contains those
|
||||
# steps (e.g. pi0/pi05). For bases that predate the feature (FastWAM/LingBot) they would
|
||||
# raise "Override keys ... do not match any step". Pop them here and let
|
||||
# _ensure_relative_actions() enable-or-inject the steps after loading (handles both cases).
|
||||
pre_overrides = dict(kwargs.get("preprocessor_overrides") or {})
|
||||
post_overrides = dict(kwargs.get("postprocessor_overrides") or {})
|
||||
pre_overrides.pop("relative_actions_processor", None)
|
||||
post_overrides.pop("absolute_actions_processor", None)
|
||||
preprocessor = PolicyProcessorPipeline.from_pretrained(
|
||||
pretrained_model_name_or_path=pretrained_path,
|
||||
config_filename=kwargs.get(
|
||||
"preprocessor_config_filename", f"{POLICY_PREPROCESSOR_DEFAULT_NAME}.json"
|
||||
),
|
||||
overrides=kwargs.get("preprocessor_overrides", {}),
|
||||
overrides=pre_overrides,
|
||||
to_transition=batch_to_transition,
|
||||
to_output=transition_to_batch,
|
||||
revision=pretrained_revision,
|
||||
@@ -335,11 +391,12 @@ def make_pre_post_processors(
|
||||
config_filename=kwargs.get(
|
||||
"postprocessor_config_filename", f"{POLICY_POSTPROCESSOR_DEFAULT_NAME}.json"
|
||||
),
|
||||
overrides=kwargs.get("postprocessor_overrides", {}),
|
||||
overrides=post_overrides,
|
||||
to_transition=policy_action_to_transition,
|
||||
to_output=transition_to_policy_action,
|
||||
revision=pretrained_revision,
|
||||
)
|
||||
_ensure_relative_actions(preprocessor, postprocessor, policy_cfg)
|
||||
_reconnect_relative_absolute_steps(preprocessor, postprocessor)
|
||||
if isinstance(policy_cfg, Evo1Config):
|
||||
from .evo1.processor_evo1 import reconcile_evo1_processors
|
||||
|
||||
@@ -188,6 +188,14 @@ class FastWAMConfig(PreTrainedConfig):
|
||||
action_video_freq_ratio: int = 4
|
||||
image_size: tuple[int, int] = (224, 448)
|
||||
context_len: int = 128
|
||||
|
||||
# Relative actions: converts absolute actions to relative (action -= state) during
|
||||
# preprocessing, and reverses it at postprocessing. Requires `proprio_dim` (OBS_STATE).
|
||||
use_relative_actions: bool = False
|
||||
# Joint names to keep absolute (not converted to relative). Empty list = all dims relative.
|
||||
relative_exclude_joints: list[str] = field(default_factory=lambda: ["gripper"])
|
||||
# Populated at runtime from dataset metadata by make_policy (used to build the exclude mask).
|
||||
action_feature_names: list[str] | None = None
|
||||
model_id: str = WAN22_MODEL_ID
|
||||
tokenizer_model_id: str = WAN_T5_TOKENIZER_ID
|
||||
text_encoder_model_id: str = WAN22_DIFFUSERS_MODEL_ID
|
||||
@@ -279,6 +287,10 @@ class FastWAMConfig(PreTrainedConfig):
|
||||
finally:
|
||||
self.pretrained_path = pretrained_path
|
||||
|
||||
@property
|
||||
def chunk_size(self) -> int:
|
||||
return self.action_horizon
|
||||
|
||||
def get_optimizer_preset(self) -> AdamWConfig:
|
||||
return AdamWConfig(lr=self.optimizer_lr, weight_decay=self.optimizer_weight_decay)
|
||||
|
||||
|
||||
@@ -183,7 +183,9 @@ class FastWAMPolicy(PreTrainedPolicy):
|
||||
sample["proprio"] = state.unsqueeze(1) if state.ndim == 2 else state
|
||||
return sample
|
||||
|
||||
def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict[str, Any]]:
|
||||
def forward(
|
||||
self, batch: dict[str, Tensor], reduction: str = "mean"
|
||||
) -> tuple[Tensor, dict[str, Any]]:
|
||||
"""Compute FastWAM training loss for a LeRobot batch.
|
||||
|
||||
Args:
|
||||
@@ -191,15 +193,19 @@ class FastWAMPolicy(PreTrainedPolicy):
|
||||
(`video`, `action`, `context`, `context_mask`) or LeRobot keys
|
||||
that can be adapted (`observation.images.*`, `observation.state`,
|
||||
`action`, `action_is_pad`).
|
||||
reduction (str): "mean" returns the scalar loss (default, backward
|
||||
compatible); "none" returns per-sample losses of shape (batch_size,)
|
||||
for sample weighting (RA-BC).
|
||||
|
||||
Returns:
|
||||
tuple[Tensor, dict[str, Any]]: The scalar loss to backprop, and a dict of
|
||||
logging metrics (e.g. `loss_video`, `loss_action`) — the `(loss, output_dict)`
|
||||
contract the LeRobot training loop expects.
|
||||
tuple[Tensor, dict[str, Any]]: The loss to backprop (scalar for "mean",
|
||||
per-sample (B,) for "none"), and a dict of logging metrics (e.g.
|
||||
`loss_video`, `loss_action`) — the `(loss, output_dict)` contract the
|
||||
LeRobot training loop expects.
|
||||
"""
|
||||
|
||||
sample = self._batch_to_training_sample(batch)
|
||||
loss, metrics = self.model.training_loss(sample)
|
||||
loss, metrics = self.model.training_loss(sample, reduction=reduction)
|
||||
return loss, dict(metrics or {})
|
||||
|
||||
@torch.no_grad()
|
||||
|
||||
@@ -21,13 +21,16 @@ import torch
|
||||
|
||||
from lerobot.configs import PipelineFeatureType, PolicyFeature
|
||||
from lerobot.processor import (
|
||||
AbsoluteActionsProcessorStep,
|
||||
ActionProcessorStep,
|
||||
AddBatchDimensionProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
ProcessorStep,
|
||||
ProcessorStepRegistry,
|
||||
RelativeActionsProcessorStep,
|
||||
RenameObservationsProcessorStep,
|
||||
UnnormalizerProcessorStep,
|
||||
policy_action_to_transition,
|
||||
@@ -105,10 +108,20 @@ 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.
|
||||
|
||||
input_steps = [
|
||||
# Shared relative-action step (OpenPI order: raw -> relative -> normalize -> model ->
|
||||
# unnormalize -> absolute). The SAME instance is passed to AbsoluteActionsProcessorStep
|
||||
# below so its cached raw state (set during preprocessing) flows to postprocessing.
|
||||
relative_step = RelativeActionsProcessorStep(
|
||||
enabled=config.use_relative_actions,
|
||||
exclude_joints=getattr(config, "relative_exclude_joints", []),
|
||||
action_names=getattr(config, "action_feature_names", None),
|
||||
)
|
||||
|
||||
input_steps: list[ProcessorStep] = [
|
||||
RenameObservationsProcessorStep(rename_map={}),
|
||||
AddBatchDimensionProcessorStep(),
|
||||
DeviceProcessorStep(device=config.device),
|
||||
relative_step,
|
||||
NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
@@ -116,12 +129,13 @@ def make_fastwam_pre_post_processors(
|
||||
device=config.device,
|
||||
),
|
||||
]
|
||||
output_steps = [
|
||||
output_steps: list[ProcessorStep] = [
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features,
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=normalization_stats,
|
||||
),
|
||||
AbsoluteActionsProcessorStep(enabled=config.use_relative_actions, relative_step=relative_step),
|
||||
]
|
||||
if config.toggle_action_dimensions:
|
||||
output_steps.append(
|
||||
|
||||
@@ -1359,7 +1359,9 @@ class FastWAM(torch.nn.Module):
|
||||
pred_action = self.action_expert.post_dit(tokens_out["action"], action_pre)
|
||||
return pred_video, pred_action
|
||||
|
||||
def _compute_training_video_loss(self, inputs, pred_video, target_video, timestep_video):
|
||||
def _compute_training_video_loss(
|
||||
self, inputs, pred_video, target_video, timestep_video, reduction: str = "mean"
|
||||
):
|
||||
include_initial_video_step = inputs["first_frame_latents"] is None
|
||||
if inputs["first_frame_latents"] is not None:
|
||||
pred_video = pred_video[:, :, 1:]
|
||||
@@ -1374,9 +1376,13 @@ class FastWAM(torch.nn.Module):
|
||||
loss_video_per_sample.device,
|
||||
dtype=loss_video_per_sample.dtype,
|
||||
)
|
||||
return (loss_video_per_sample * video_weight).mean()
|
||||
weighted = loss_video_per_sample * video_weight
|
||||
# reduction="none" returns the per-sample vector (B,) for sample weighting (RA-BC).
|
||||
return weighted if reduction == "none" else weighted.mean()
|
||||
|
||||
def _compute_training_action_loss(self, inputs, pred_action, target_action, timestep_action):
|
||||
def _compute_training_action_loss(
|
||||
self, inputs, pred_action, target_action, timestep_action, reduction: str = "mean"
|
||||
):
|
||||
action_loss_token = functional.mse_loss(
|
||||
pred_action.float(), target_action.float(), reduction="none"
|
||||
).mean(dim=2)
|
||||
@@ -1393,9 +1399,11 @@ class FastWAM(torch.nn.Module):
|
||||
action_loss_per_sample.device,
|
||||
dtype=action_loss_per_sample.dtype,
|
||||
)
|
||||
return (action_loss_per_sample * action_weight).mean()
|
||||
weighted = action_loss_per_sample * action_weight
|
||||
# reduction="none" returns the per-sample vector (B,) for sample weighting (RA-BC).
|
||||
return weighted if reduction == "none" else weighted.mean()
|
||||
|
||||
def training_loss(self, sample, tiled: bool = False):
|
||||
def training_loss(self, sample, tiled: bool = False, reduction: str = "mean"):
|
||||
inputs = self.build_inputs(sample, tiled=tiled)
|
||||
targets = self._sample_training_targets(inputs)
|
||||
pred_video, pred_action = self._run_training_mot(inputs=inputs, targets=targets)
|
||||
@@ -1404,17 +1412,20 @@ class FastWAM(torch.nn.Module):
|
||||
pred_video=pred_video,
|
||||
target_video=targets["target_video"],
|
||||
timestep_video=targets["timestep_video"],
|
||||
reduction=reduction,
|
||||
)
|
||||
loss_action = self._compute_training_action_loss(
|
||||
inputs=inputs,
|
||||
pred_action=pred_action,
|
||||
target_action=targets["target_action"],
|
||||
timestep_action=targets["timestep_action"],
|
||||
reduction=reduction,
|
||||
)
|
||||
# With reduction="none" both terms are (B,), so loss_total is the per-sample loss (B,).
|
||||
loss_total = self.loss_lambda_video * loss_video + self.loss_lambda_action * loss_action
|
||||
loss_dict = {
|
||||
"loss_video": self.loss_lambda_video * float(loss_video.detach().item()),
|
||||
"loss_action": self.loss_lambda_action * float(loss_action.detach().item()),
|
||||
"loss_video": self.loss_lambda_video * float(loss_video.detach().mean().item()),
|
||||
"loss_action": self.loss_lambda_action * float(loss_action.detach().mean().item()),
|
||||
}
|
||||
return loss_total, loss_dict
|
||||
|
||||
|
||||
@@ -28,7 +28,11 @@ from dataclasses import dataclass, field
|
||||
from lerobot.configs.policies import PreTrainedConfig
|
||||
from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
|
||||
from lerobot.optim.optimizers import AdamWConfig
|
||||
from lerobot.optim.schedulers import ConstantWithWarmupSchedulerConfig, LRSchedulerConfig
|
||||
from lerobot.optim.schedulers import (
|
||||
ConstantWithWarmupSchedulerConfig,
|
||||
CosineAnnealingWithWarmupSchedulerConfig,
|
||||
LRSchedulerConfig,
|
||||
)
|
||||
from lerobot.utils.constants import ACTION
|
||||
|
||||
|
||||
@@ -92,6 +96,15 @@ class LingBotVAConfig(PreTrainedConfig):
|
||||
# (un)normalization quantiles live in the checkpoint's ``policy_postprocessor.json``, not here.
|
||||
used_action_channel_ids: list[int] = field(default_factory=lambda: list(range(7)))
|
||||
|
||||
# Relative actions: converts absolute actions to relative (action -= state) during
|
||||
# preprocessing, and reverses it at postprocessing. Requires the dataset to provide
|
||||
# observation.state whose leading dims align 1:1 with the used action channels.
|
||||
use_relative_actions: bool = False
|
||||
# Joint names to keep absolute (not converted to relative). Empty list = all dims relative.
|
||||
relative_exclude_joints: list[str] = field(default_factory=lambda: ["gripper"])
|
||||
# Populated at runtime from dataset metadata by make_policy (used to build the exclude mask).
|
||||
action_feature_names: list[str] | None = None
|
||||
|
||||
# Opt-in: VAE-decode predicted video latents to ``self.last_predicted_frames`` for saving MP4s.
|
||||
save_predicted_video: bool = False
|
||||
|
||||
@@ -112,6 +125,11 @@ class LingBotVAConfig(PreTrainedConfig):
|
||||
optimizer_weight_decay: float = 1e-4
|
||||
optimizer_grad_clip_norm: float = 1.0
|
||||
scheduler_warmup_steps: int = 1000
|
||||
# Scheduler after warmup. "constant_with_warmup" (upstream default: warmup then flat peak LR)
|
||||
# or "cosine_annealing_with_warmup" (warmup then cosine anneal peak->0 over the remaining steps).
|
||||
# Cosine tightens the loss tail and often nudges final loss down; it does NOT reduce the
|
||||
# flow-matching estimator's step-to-step noise (that's metric variance, LR-independent).
|
||||
scheduler_type: str = "constant_with_warmup"
|
||||
|
||||
def __post_init__(self):
|
||||
super().__post_init__()
|
||||
@@ -150,7 +168,10 @@ class LingBotVAConfig(PreTrainedConfig):
|
||||
)
|
||||
|
||||
def get_scheduler_preset(self) -> LRSchedulerConfig | None:
|
||||
# Upstream uses a linear warmup followed by a constant LR (warmup_constant_lambda).
|
||||
# Default (upstream): linear warmup then constant LR (warmup_constant_lambda).
|
||||
# Optionally cosine-anneal peak->0 over the remaining steps via scheduler_type.
|
||||
if self.scheduler_type == "cosine_annealing_with_warmup":
|
||||
return CosineAnnealingWithWarmupSchedulerConfig(num_warmup_steps=self.scheduler_warmup_steps)
|
||||
return ConstantWithWarmupSchedulerConfig(num_warmup_steps=self.scheduler_warmup_steps)
|
||||
|
||||
@property
|
||||
|
||||
@@ -257,8 +257,12 @@ class LingBotVAPolicy(PreTrainedPolicy):
|
||||
"grid_id": grid_id,
|
||||
}
|
||||
|
||||
def _flow_matching_loss(self, input_dict, pred):
|
||||
"""Dual-stream flow-matching loss (port of upstream ``Trainer.compute_loss``)."""
|
||||
def _flow_matching_loss(self, input_dict, pred, reduction: str = "mean"):
|
||||
"""Dual-stream flow-matching loss (port of upstream ``Trainer.compute_loss``).
|
||||
|
||||
``reduction="mean"`` returns scalar (latent_loss, action_loss); ``"none"`` returns
|
||||
per-sample vectors of shape ``(B,)`` each (averaged over latent frames) for RA-BC.
|
||||
"""
|
||||
latent_pred, action_pred = pred
|
||||
ld, ad = input_dict["latent_dict"], input_dict["action_dict"]
|
||||
action_pred = rearrange(action_pred, "b (f n) c -> b c f n 1", f=ad["targets"].shape[-3])
|
||||
@@ -278,7 +282,8 @@ class LingBotVAPolicy(PreTrainedPolicy):
|
||||
latent_loss = (
|
||||
(latent_loss * lw[:, None, :, None, None]).permute(0, 2, 3, 4, 1).flatten(0, 1).flatten(1)
|
||||
)
|
||||
latent_loss = (latent_loss.sum(dim=1) / (torch.ones_like(latent_loss).sum(dim=1) + 1e-6)).mean()
|
||||
# per (batch*frame) mean over spatial/channel -> (B*F,)
|
||||
latent_loss = latent_loss.sum(dim=1) / (torch.ones_like(latent_loss).sum(dim=1) + 1e-6)
|
||||
|
||||
amask = ad["actions_mask"].float()
|
||||
action_loss = F.mse_loss(action_pred.float(), ad["targets"].float().detach(), reduction="none")
|
||||
@@ -286,10 +291,14 @@ class LingBotVAPolicy(PreTrainedPolicy):
|
||||
(action_loss * aw[:, None, :, None, None] * amask).permute(0, 2, 3, 4, 1).flatten(0, 1).flatten(1)
|
||||
)
|
||||
amask_f = amask.permute(0, 2, 3, 4, 1).flatten(0, 1).flatten(1)
|
||||
action_loss = (action_loss.sum(dim=1) / (amask_f.sum(dim=1) + 1e-6)).mean()
|
||||
return latent_loss, action_loss
|
||||
action_loss = action_loss.sum(dim=1) / (amask_f.sum(dim=1) + 1e-6)
|
||||
|
||||
def training_loss_from_streams(self, latents, actions, actions_mask, text_emb):
|
||||
if reduction == "none":
|
||||
# (B*F,) -> (B, F) -> (B,): per-sample losses for RA-BC weighting.
|
||||
return latent_loss.reshape(bn, fn).mean(dim=1), action_loss.reshape(bn, fn).mean(dim=1)
|
||||
return latent_loss.mean(), action_loss.mean()
|
||||
|
||||
def training_loss_from_streams(self, latents, actions, actions_mask, text_emb, reduction: str = "mean"):
|
||||
"""Core dual-stream training loss given prepared latents / actions / text embeddings.
|
||||
|
||||
``latents``: ``[B, in_channels, F, h, w]`` (normalized video latents).
|
||||
@@ -318,20 +327,24 @@ class LingBotVAPolicy(PreTrainedPolicy):
|
||||
"window_size": int(torch.randint(4, 65, (1,)).item()),
|
||||
}
|
||||
pred = self.transformer(input_dict, train_mode=True)
|
||||
latent_loss, action_loss = self._flow_matching_loss(input_dict, pred)
|
||||
latent_loss, action_loss = self._flow_matching_loss(input_dict, pred, reduction)
|
||||
# reduction="none": latent_loss/action_loss are (B,) -> loss is per-sample (B,).
|
||||
loss = latent_loss + action_loss
|
||||
return loss, {"latent_loss": latent_loss.item(), "action_loss": action_loss.item()}
|
||||
return loss, {"latent_loss": latent_loss.mean().item(), "action_loss": action_loss.mean().item()}
|
||||
|
||||
def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict | None]:
|
||||
def forward(self, batch: dict[str, Tensor], reduction: str = "mean") -> tuple[Tensor, dict | None]:
|
||||
"""Training forward: dual-stream flow-matching loss.
|
||||
|
||||
Builds the (video-latent, action, text) training streams from a LeRobot batch
|
||||
(VAE-encoding the camera frames and UMT5-encoding the task), then runs the flow-matching
|
||||
dual-stream loss. Requires the policy to be built with ``attn_mode='flex'``.
|
||||
|
||||
``reduction="mean"`` returns the scalar loss (default); ``"none"`` returns per-sample
|
||||
losses of shape ``(B,)`` for sample weighting (RA-BC).
|
||||
"""
|
||||
self._ensure_frozen_modules()
|
||||
latents, actions, actions_mask, text_emb = self._build_training_streams(batch)
|
||||
return self.training_loss_from_streams(latents, actions, actions_mask, text_emb)
|
||||
return self.training_loss_from_streams(latents, actions, actions_mask, text_emb, reduction=reduction)
|
||||
|
||||
@torch.no_grad()
|
||||
def _build_training_streams(self, batch):
|
||||
|
||||
@@ -25,12 +25,14 @@ import torch
|
||||
|
||||
from lerobot.configs.types import FeatureType, NormalizationMode
|
||||
from lerobot.processor import (
|
||||
AbsoluteActionsProcessorStep,
|
||||
AddBatchDimensionProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
ProcessorStep,
|
||||
RelativeActionsProcessorStep,
|
||||
RenameObservationsProcessorStep,
|
||||
UnnormalizerProcessorStep,
|
||||
)
|
||||
@@ -52,9 +54,19 @@ def make_lingbot_va_pre_post_processors(
|
||||
]:
|
||||
"""Build the pre/post processor pipelines for LingBot-VA."""
|
||||
|
||||
# Shared relative-action step (OpenPI order: raw -> relative -> normalize -> model ->
|
||||
# unnormalize -> absolute). The SAME instance is passed to AbsoluteActionsProcessorStep
|
||||
# below so its cached raw state (set during preprocessing) flows to postprocessing.
|
||||
relative_step = RelativeActionsProcessorStep(
|
||||
enabled=config.use_relative_actions,
|
||||
exclude_joints=getattr(config, "relative_exclude_joints", []),
|
||||
action_names=getattr(config, "action_feature_names", None),
|
||||
)
|
||||
|
||||
input_steps: list[ProcessorStep] = [
|
||||
RenameObservationsProcessorStep(rename_map={}),
|
||||
AddBatchDimensionProcessorStep(),
|
||||
relative_step,
|
||||
NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
@@ -63,13 +75,16 @@ def make_lingbot_va_pre_post_processors(
|
||||
DeviceProcessorStep(device=config.device),
|
||||
]
|
||||
|
||||
# Unnormalize actions from [-1, 1] to physical units (QUANTILES) using q01/q99 restored from the checkpoint.
|
||||
# Unnormalize actions back to physical units. Config-driven norm_map (was hardcoded QUANTILES)
|
||||
# so it stays symmetric with the preprocessor's NormalizerProcessorStep — required for
|
||||
# use_relative_actions with ACTION=IDENTITY (and unchanged for QUANTILES runs).
|
||||
output_steps: list[ProcessorStep] = [
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features,
|
||||
norm_map={FeatureType.ACTION: NormalizationMode.QUANTILES},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
),
|
||||
AbsoluteActionsProcessorStep(enabled=config.use_relative_actions, relative_step=relative_step),
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
]
|
||||
|
||||
|
||||
@@ -29,6 +29,7 @@ from huggingface_hub import HfApi, ModelCard, ModelCardData, hf_hub_download, sa
|
||||
from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE
|
||||
from huggingface_hub.errors import HfHubHTTPError
|
||||
from safetensors.torch import load_model as load_model_as_safetensor, save_model as save_model_as_safetensor
|
||||
import torch
|
||||
from torch import Tensor, nn
|
||||
|
||||
from lerobot.__version__ import __version__
|
||||
@@ -221,6 +222,14 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
|
||||
|
||||
@classmethod
|
||||
def _load_as_safetensor(cls, model: T, model_file: str, map_location: str, strict: bool) -> T:
|
||||
# safetensors' load_file maps the bare string "cuda" to cuda:0 regardless of the current
|
||||
# device (unlike torch's .to("cuda"), which honors torch.cuda.current_device()). Under
|
||||
# multi-GPU accelerate/FSDP every rank would then load its weights onto GPU 0, OOMing it
|
||||
# before sharding. Resolve "cuda" to the concrete current-device index so each rank loads
|
||||
# onto its own GPU.
|
||||
if map_location == "cuda" and torch.cuda.is_available():
|
||||
map_location = f"cuda:{torch.cuda.current_device()}"
|
||||
|
||||
# Create base kwargs
|
||||
kwargs = {"strict": strict}
|
||||
|
||||
|
||||
@@ -51,6 +51,12 @@ def to_relative_actions(actions: Tensor, state: Tensor, mask: Sequence[bool]) ->
|
||||
# DeviceProcessorStep moves the transition, so it can be on CPU while actions are on CUDA.
|
||||
if state.device != actions.device or state.dtype != actions.dtype:
|
||||
state = state.to(device=actions.device, dtype=actions.dtype)
|
||||
# When the observation is temporally stacked (e.g. LingBot-VA loads several obs steps via
|
||||
# observation_delta_indices, giving state shape (B, T_obs, state_dim)), the relative reference
|
||||
# is the CURRENT frame (delta == 0, i.e. index 0). Collapse to it so the offset broadcasts over
|
||||
# the action horizon. pi0/pi05 pass a 2D (B, state_dim) state and are unaffected.
|
||||
if state.ndim == 3:
|
||||
state = state[:, 0]
|
||||
state_offset = state[..., :dims] * mask_t
|
||||
if actions.ndim == 3:
|
||||
state_offset = state_offset.unsqueeze(-2)
|
||||
@@ -73,6 +79,10 @@ def to_absolute_actions(actions: Tensor, state: Tensor, mask: Sequence[bool]) ->
|
||||
# DeviceProcessorStep moves the transition, so it can be on CPU while actions are on CUDA.
|
||||
if state.device != actions.device or state.dtype != actions.dtype:
|
||||
state = state.to(device=actions.device, dtype=actions.dtype)
|
||||
# Mirror to_relative_actions: collapse a temporally-stacked (B, T_obs, state_dim) state to the
|
||||
# current frame (index 0) so the round-trip stays symmetric with the relative conversion.
|
||||
if state.ndim == 3:
|
||||
state = state[:, 0]
|
||||
state_offset = state[..., :dims] * mask_t
|
||||
if actions.ndim == 3:
|
||||
state_offset = state_offset.unsqueeze(-2)
|
||||
@@ -102,6 +112,12 @@ class RelativeActionsProcessorStep(ProcessorStep):
|
||||
exclude_joints: list[str] = field(default_factory=list)
|
||||
action_names: list[str] | None = None
|
||||
_last_state: torch.Tensor | None = field(default=None, init=False, repr=False)
|
||||
# When True, ``__call__`` stops refreshing ``_last_state``. Chunked inference
|
||||
# engines set this so every action popped from a policy's action queue is
|
||||
# re-anchored to the state captured when the chunk was predicted, instead of
|
||||
# drifting against the current per-tick state. Episode-scoped runtime state,
|
||||
# not part of ``get_config``.
|
||||
_hold_state: bool = field(default=False, init=False, repr=False)
|
||||
|
||||
def _build_mask(self, action_dim: int) -> list[bool]:
|
||||
if not self.exclude_joints or self.action_names is None:
|
||||
@@ -126,8 +142,9 @@ class RelativeActionsProcessorStep(ProcessorStep):
|
||||
observation = transition.get(TransitionKey.OBSERVATION, {})
|
||||
state = observation.get(OBS_STATE) if observation else None
|
||||
|
||||
# Always cache state for the paired AbsoluteActionsProcessorStep
|
||||
if state is not None:
|
||||
# Always cache state for the paired AbsoluteActionsProcessorStep, unless a
|
||||
# chunked inference engine has frozen the anchor for the current chunk.
|
||||
if state is not None and not self._hold_state:
|
||||
self._last_state = state
|
||||
|
||||
if not self.enabled:
|
||||
@@ -146,6 +163,15 @@ class RelativeActionsProcessorStep(ProcessorStep):
|
||||
"""Return the cached ``observation.state`` used as the reference point for relative/absolute action conversions."""
|
||||
return self._last_state
|
||||
|
||||
def set_hold(self, flag: bool) -> None:
|
||||
"""Freeze (``True``) or resume (``False``) refreshing of the cached anchor state.
|
||||
|
||||
While held, ``__call__`` keeps the previously cached ``_last_state`` so the
|
||||
paired :class:`AbsoluteActionsProcessorStep` re-anchors every action of a
|
||||
predicted chunk to the same state, avoiding intra-chunk drift.
|
||||
"""
|
||||
self._hold_state = flag
|
||||
|
||||
def get_config(self) -> dict[str, Any]:
|
||||
return {
|
||||
"enabled": self.enabled,
|
||||
|
||||
@@ -43,7 +43,6 @@ 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
|
||||
@@ -52,7 +51,6 @@ from .configs import BaseStrategyConfig, DAggerStrategyConfig, RolloutConfig
|
||||
from .inference import (
|
||||
InferenceEngine,
|
||||
RTCInferenceConfig,
|
||||
SyncInferenceConfig,
|
||||
create_inference_engine,
|
||||
)
|
||||
from .robot_wrapper import ThreadSafeRobot
|
||||
@@ -399,15 +397,6 @@ 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,26 +24,32 @@ 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
|
||||
from lerobot.processor import PolicyProcessorPipeline, RelativeActionsProcessorStep
|
||||
|
||||
from .base import InferenceEngine
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# 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.
|
||||
# Relative-action support (drift-free anchoring)
|
||||
# ----------------------------------------------
|
||||
# Relative-action policies predict a *chunk* of offsets anchored to the robot
|
||||
# state at chunk-prediction time. ``select_action`` serves that chunk one action
|
||||
# per tick from an internal ``_action_queue``, recomputing only when the queue is
|
||||
# empty. The per-tick flow here runs the full pre/post pipeline every call, and
|
||||
# ``RelativeActionsProcessorStep`` would otherwise refresh its cached anchor state
|
||||
# on every tick — so actions popped from the queue on later ticks would be
|
||||
# re-anchored to the *current* (already-moved) state and absolute targets would
|
||||
# drift through the chunk.
|
||||
#
|
||||
# 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``).
|
||||
# Fix: detect chunk boundaries by inspecting the policy's ``_action_queue`` length
|
||||
# *before* running the pipeline, and freeze the relative step's cached anchor
|
||||
# (``set_hold``) on ticks that pop a cached action. The whole chunk is then
|
||||
# anchored to a single state, exactly like RTC. ``select_action`` stays on the
|
||||
# hot path, so policy-specific side effects (e.g. LingBot-VA's per-tick keyframe
|
||||
# feedback) are preserved. Policies without an ``_action_queue`` (e.g. ACT
|
||||
# temporal ensembling, which recomputes every tick) fall back to refreshing the
|
||||
# anchor every tick, which is the correct behaviour there.
|
||||
|
||||
|
||||
class SyncInferenceEngine(InferenceEngine):
|
||||
@@ -73,6 +79,24 @@ class SyncInferenceEngine(InferenceEngine):
|
||||
self._task = task
|
||||
self._device = torch.device(device or "cpu")
|
||||
self._robot_type = robot_type
|
||||
|
||||
# Relative-action policies need the chunk anchor held while cached actions
|
||||
# are popped (see module docstring). Introspect the preprocessor for an
|
||||
# enabled RelativeActionsProcessorStep, mirroring the RTC engine.
|
||||
self._relative_step = next(
|
||||
(
|
||||
s
|
||||
for s in getattr(preprocessor, "steps", ())
|
||||
if isinstance(s, RelativeActionsProcessorStep) and s.enabled
|
||||
),
|
||||
None,
|
||||
)
|
||||
if self._relative_step is not None:
|
||||
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)
|
||||
logger.info("Relative actions enabled: chunk anchor will be held per chunk")
|
||||
|
||||
logger.info(
|
||||
"SyncInferenceEngine initialized (device=%s, action_keys=%d)",
|
||||
self._device,
|
||||
@@ -93,6 +117,21 @@ class SyncInferenceEngine(InferenceEngine):
|
||||
self._policy.reset()
|
||||
self._preprocessor.reset()
|
||||
self._postprocessor.reset()
|
||||
# ``policy.reset()`` empties ``_action_queue`` so the next ``get_action``
|
||||
# recomputes and refreshes the anchor; clear any leftover hold defensively.
|
||||
if self._relative_step is not None:
|
||||
self._relative_step.set_hold(False)
|
||||
|
||||
def _policy_will_recompute(self) -> bool:
|
||||
"""True if the next ``select_action`` will predict a fresh chunk (queue empty/absent).
|
||||
|
||||
Relative-action policies expose an ``_action_queue`` deque that is refilled
|
||||
only when empty. When it is non-empty the upcoming ``select_action`` will
|
||||
pop a cached action, so the anchor state must be held. Policies without the
|
||||
attribute recompute every tick, so we always refresh the anchor.
|
||||
"""
|
||||
queue = getattr(self._policy, "_action_queue", None)
|
||||
return queue is None or len(queue) == 0
|
||||
|
||||
def get_action(self, obs_frame: dict | None) -> torch.Tensor | None:
|
||||
"""Run the full inference pipeline on ``obs_frame`` and return an action tensor."""
|
||||
@@ -107,13 +146,25 @@ class SyncInferenceEngine(InferenceEngine):
|
||||
if self._device.type == "cuda" and self._policy.config.use_amp
|
||||
else nullcontext()
|
||||
)
|
||||
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)
|
||||
action = self._postprocessor(action)
|
||||
# For relative-action policies, hold the cached anchor on ticks that pop a
|
||||
# cached action so the whole chunk stays anchored to the state captured when
|
||||
# it was predicted. Decided before the pipeline runs (the queue is drained
|
||||
# inside ``select_action``); always released in ``finally`` so a hold never
|
||||
# leaks across ticks or on exception.
|
||||
hold_anchor = self._relative_step is not None and not self._policy_will_recompute()
|
||||
if self._relative_step is not None:
|
||||
self._relative_step.set_hold(hold_anchor)
|
||||
try:
|
||||
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)
|
||||
action = self._postprocessor(action)
|
||||
finally:
|
||||
if self._relative_step is not None:
|
||||
self._relative_step.set_hold(False)
|
||||
action_tensor = action.squeeze(0).cpu()
|
||||
|
||||
# Reorder to match dataset action ordering so the caller can treat
|
||||
|
||||
@@ -20,6 +20,7 @@ Requires: pip install 'lerobot[training]' (includes dataset + accelerate + wand
|
||||
|
||||
import dataclasses
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from contextlib import nullcontext
|
||||
@@ -461,6 +462,14 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
|
||||
# declares language columns; otherwise stay on PyTorch's default
|
||||
# collate so non-language training runs are unaffected.
|
||||
collate_fn = lerobot_collate_fn if dataset.meta.has_language_columns else None
|
||||
# Multi-node fork-OOM mitigation: on EFA nodes (vm.overcommit_memory=0, no swap),
|
||||
# forking dataloader workers from a multi-GB rank reserve-charges the rank's full virtual
|
||||
# footprint, so 8 ranks x num_workers forking at once trips OSError(ENOMEM) despite free
|
||||
# RAM. Honor LEROBOT_DATALOADER_MP_CONTEXT (e.g. "forkserver"/"spawn") to spawn workers
|
||||
# from a clean context instead of fork(). Only meaningful when workers are used.
|
||||
dataloader_mp_context = os.environ.get("LEROBOT_DATALOADER_MP_CONTEXT") or None
|
||||
if cfg.num_workers == 0:
|
||||
dataloader_mp_context = None
|
||||
dataloader = torch.utils.data.DataLoader(
|
||||
dataset,
|
||||
num_workers=cfg.num_workers,
|
||||
@@ -472,6 +481,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
|
||||
collate_fn=collate_fn,
|
||||
prefetch_factor=cfg.prefetch_factor if cfg.num_workers > 0 else None,
|
||||
persistent_workers=cfg.persistent_workers and cfg.num_workers > 0,
|
||||
multiprocessing_context=dataloader_mp_context,
|
||||
)
|
||||
|
||||
# Build eval dataloader if a held-out split exists
|
||||
@@ -499,6 +509,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
|
||||
collate_fn=eval_collate_fn,
|
||||
prefetch_factor=cfg.prefetch_factor if cfg.num_workers > 0 else None,
|
||||
persistent_workers=cfg.persistent_workers and cfg.num_workers > 0,
|
||||
multiprocessing_context=dataloader_mp_context,
|
||||
)
|
||||
|
||||
# Prepare everything with accelerator
|
||||
|
||||
@@ -346,3 +346,41 @@ 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)
|
||||
|
||||
|
||||
# Anchor-hold semantics (used by the synchronous rollout engine to avoid intra-chunk drift)
|
||||
|
||||
|
||||
def _state_transition(state):
|
||||
return batch_to_transition({OBS_STATE: state})
|
||||
|
||||
|
||||
def test_set_hold_freezes_cached_anchor_state():
|
||||
"""While held, the cached anchor is not overwritten, but the observation passes through."""
|
||||
step = RelativeActionsProcessorStep(enabled=True)
|
||||
|
||||
state_a = torch.tensor([[1.0, 2.0, 3.0, 4.0]])
|
||||
step(_state_transition(state_a))
|
||||
torch.testing.assert_close(step.get_cached_state(), state_a)
|
||||
|
||||
# Freeze: a new state must NOT replace the cached anchor.
|
||||
step.set_hold(True)
|
||||
state_b = torch.tensor([[9.0, 9.0, 9.0, 9.0]])
|
||||
out = step(_state_transition(state_b))
|
||||
torch.testing.assert_close(step.get_cached_state(), state_a)
|
||||
# The transition's observation state is still the current one (hold only affects caching).
|
||||
torch.testing.assert_close(out[TransitionKey.OBSERVATION][OBS_STATE], state_b)
|
||||
|
||||
# Release: caching resumes.
|
||||
step.set_hold(False)
|
||||
state_c = torch.tensor([[5.0, 6.0, 7.0, 8.0]])
|
||||
step(_state_transition(state_c))
|
||||
torch.testing.assert_close(step.get_cached_state(), state_c)
|
||||
|
||||
|
||||
def test_hold_state_not_in_config():
|
||||
"""Hold is ephemeral runtime state and must not leak into the serialized config."""
|
||||
step = RelativeActionsProcessorStep(enabled=True)
|
||||
step.set_hold(True)
|
||||
assert "_hold_state" not in step.get_config()
|
||||
assert set(step.get_config()) == {"enabled", "exclude_joints", "action_names"}
|
||||
|
||||
@@ -348,3 +348,164 @@ 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):
|
||||
"""Build fake pre/post processors wrapping real relative/absolute steps.
|
||||
|
||||
The preprocessor runs the ``RelativeActionsProcessorStep`` (caching/holding the
|
||||
anchor state) and passes the observation through; the postprocessor runs the
|
||||
paired ``AbsoluteActionsProcessorStep`` (relative + cached state) and returns the
|
||||
absolute action tensor. Shapes mirror 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):
|
||||
# observation carries a batched OBS_STATE tensor; run the relative step so
|
||||
# it caches (or holds) the anchor, then pass the batch through unchanged.
|
||||
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, with_queue=True):
|
||||
"""Fake chunk policy: refills an ``_action_queue`` with ``chunk_rel`` when empty."""
|
||||
from collections import deque
|
||||
|
||||
policy = MagicMock()
|
||||
policy.config.use_amp = False
|
||||
policy.config.action_feature_names = list(_REL_ACTION_NAMES)
|
||||
state = {"predict_calls": 0}
|
||||
|
||||
if with_queue:
|
||||
policy._action_queue = deque(maxlen=n_action_steps)
|
||||
else:
|
||||
# Ensure the attribute is truly absent so getattr(...) falls back.
|
||||
del policy._action_queue
|
||||
|
||||
def select_action(_observation):
|
||||
if with_queue:
|
||||
if len(policy._action_queue) == 0:
|
||||
state["predict_calls"] += 1
|
||||
policy._action_queue.extend(chunk_rel[i].unsqueeze(0) for i in range(n_action_steps))
|
||||
return policy._action_queue.popleft()
|
||||
# No queue: recompute every tick (like temporal ensembling).
|
||||
state["predict_calls"] += 1
|
||||
return chunk_rel[0].unsqueeze(0)
|
||||
|
||||
policy.select_action.side_effect = select_action
|
||||
policy.reset.side_effect = lambda: policy._action_queue.clear() if with_queue else None
|
||||
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 -> recompute -> anchor refreshes 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])
|
||||
assert relative_step._hold_state is False # released after every call
|
||||
|
||||
|
||||
def test_sync_relative_fallback_without_action_queue():
|
||||
"""A policy without ``_action_queue`` refreshes the anchor every tick."""
|
||||
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, with_queue=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 tracks the current state every tick.
|
||||
torch.testing.assert_close(out, torch.tensor(state) + chunk_rel[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
|
||||
|
||||
Reference in New Issue
Block a user