From 287c823f13d23d0c7b35eba9cb2b90880cfc5240 Mon Sep 17 00:00:00 2001 From: Caroline Pascal Date: Tue, 16 Jun 2026 17:58:59 +0200 Subject: [PATCH 01/38] fix(features copy): adding deepcopy on LeRobot dataset features to avoid shallow copy leaks (#3826) * fix(features copy): adding deepcopy on LeRobot dataset features to avoid shallow copy leaks * tests(test): adding new test --- src/lerobot/datasets/dataset_metadata.py | 3 ++- src/lerobot/datasets/dataset_tools.py | 5 ++++- tests/datasets/test_datasets.py | 17 ++++++++++++++++- 3 files changed, 22 insertions(+), 3 deletions(-) diff --git a/src/lerobot/datasets/dataset_metadata.py b/src/lerobot/datasets/dataset_metadata.py index 39a1b6d2b..b496e4f65 100644 --- a/src/lerobot/datasets/dataset_metadata.py +++ b/src/lerobot/datasets/dataset_metadata.py @@ -15,6 +15,7 @@ # limitations under the License. import contextlib from collections.abc import Callable +from copy import deepcopy from pathlib import Path import numpy as np @@ -709,7 +710,7 @@ class LeRobotDatasetMetadata: obj.root.mkdir(parents=True, exist_ok=False) - features = {**features, **DEFAULT_FEATURES} + features = {**deepcopy(features), **DEFAULT_FEATURES} _validate_feature_names(features) obj.tasks = None diff --git a/src/lerobot/datasets/dataset_tools.py b/src/lerobot/datasets/dataset_tools.py index 91dc66af2..9aca859b4 100644 --- a/src/lerobot/datasets/dataset_tools.py +++ b/src/lerobot/datasets/dataset_tools.py @@ -27,6 +27,7 @@ import logging import shutil from collections.abc import Callable from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed +from copy import deepcopy from pathlib import Path import datasets @@ -1101,7 +1102,9 @@ def _copy_episodes_metadata_and_stats( if dst_meta.video_keys and src_dataset.meta.video_keys: for key in dst_meta.video_keys: if key in src_dataset.meta.features: - dst_meta.info.features[key]["info"] = src_dataset.meta.info.features[key].get("info", {}) + dst_meta.info.features[key]["info"] = deepcopy( + src_dataset.meta.info.features[key].get("info", {}) + ) write_info(dst_meta.info, dst_meta.root) diff --git a/tests/datasets/test_datasets.py b/tests/datasets/test_datasets.py index 19c314fd6..1d2fb1d55 100644 --- a/tests/datasets/test_datasets.py +++ b/tests/datasets/test_datasets.py @@ -51,7 +51,7 @@ from lerobot.robots import make_robot_from_config from lerobot.transforms import ImageTransforms, ImageTransformsConfig from lerobot.utils.constants import ACTION, DONE, OBS_IMAGES, OBS_STATE, OBS_STR, REWARD from lerobot.utils.feature_utils import hw_to_dataset_features -from tests.fixtures.constants import DUMMY_CHW, DUMMY_HWC, DUMMY_REPO_ID +from tests.fixtures.constants import DUMMY_CHW, DUMMY_HWC, DUMMY_MOTOR_FEATURES, DUMMY_REPO_ID from tests.mocks.mock_robot import MockRobotConfig from tests.utils import require_x86_64_kernel @@ -133,6 +133,21 @@ def test_dataset_feature_with_forward_slash_raises_error(): ) +def test_create_does_not_mutate_input_features(tmp_path, empty_lerobot_dataset_factory): + # ``create`` must deep-copy features so a dataset built from another's features stays independent. + dataset = empty_lerobot_dataset_factory( + root=tmp_path / "ds1", features=DUMMY_MOTOR_FEATURES, use_videos=False + ) + dataset_copy = empty_lerobot_dataset_factory( + root=tmp_path / "ds2", features=dataset.meta.features, use_videos=False + ) + + original_shape = dataset.meta.info.features["state"]["shape"] + dataset_copy.meta.info.features["state"]["shape"] = (999,) + + assert dataset.meta.info.features["state"]["shape"] == original_shape + + def test_add_frame_missing_task(tmp_path, empty_lerobot_dataset_factory): features = {"state": {"dtype": "float32", "shape": (1,), "names": None}} dataset = empty_lerobot_dataset_factory(root=tmp_path / "test", features=features) From 2b0834bcb84ca529ede50b2de3afcbd77325e93c Mon Sep 17 00:00:00 2001 From: Caroline Pascal Date: Wed, 17 Jun 2026 11:40:17 +0200 Subject: [PATCH 02/38] fix(cameras): snapshot stop_event in read loops to avoid None deref (#3812) * Do not set stop_event to None when stopping thread * fix(cameras): snapshot stop_event in read loops to avoid None deref The background read loops accessed self.stop_event repeatedly while _stop_read_thread() can reassign it to None after join(). Reading the attribute across the loop condition (and a mid-loop re-check) was a time-of-check/time-of-use race: stop_event could flip to None between the `is None` test and the `.is_set()` call, raising AttributeError on the worker thread. Snapshot self.stop_event into a local once, guard it, and loop on the local Event. The Event object is thread-safe and lives for the thread's lifetime; _stop_read_thread() always calls .set() before nulling the attribute, so the local observes the stop and exits cleanly. This also lets us drop the redundant pre-lock stop check. Applies to OpenCVCamera, RealSenseCamera, and ZMQ camera. --------- Co-authored-by: Anes Benmerzoug --- src/lerobot/cameras/opencv/camera_opencv.py | 5 +++-- src/lerobot/cameras/realsense/camera_realsense.py | 5 +++-- src/lerobot/cameras/zmq/camera_zmq.py | 5 +++-- 3 files changed, 9 insertions(+), 6 deletions(-) diff --git a/src/lerobot/cameras/opencv/camera_opencv.py b/src/lerobot/cameras/opencv/camera_opencv.py index 3e92eaf06..b3c20e8dd 100644 --- a/src/lerobot/cameras/opencv/camera_opencv.py +++ b/src/lerobot/cameras/opencv/camera_opencv.py @@ -442,11 +442,12 @@ class OpenCVCamera(Camera): Stops on DeviceNotConnectedError, logs other errors and continues. """ - if self.stop_event is None: + stop_event = self.stop_event + if stop_event is None: raise RuntimeError(f"{self}: stop_event is not initialized before starting read loop.") failure_count = 0 - while not self.stop_event.is_set(): + while not stop_event.is_set(): try: raw_frame = self._read_from_hardware() processed_frame = self._postprocess_image(raw_frame) diff --git a/src/lerobot/cameras/realsense/camera_realsense.py b/src/lerobot/cameras/realsense/camera_realsense.py index e156e6d14..80008e9f9 100644 --- a/src/lerobot/cameras/realsense/camera_realsense.py +++ b/src/lerobot/cameras/realsense/camera_realsense.py @@ -471,11 +471,12 @@ class RealSenseCamera(Camera): Stops on DeviceNotConnectedError, logs other errors and continues. """ - if self.stop_event is None: + stop_event = self.stop_event + if stop_event is None: raise RuntimeError(f"{self}: stop_event is not initialized before starting read loop.") failure_count = 0 - while not self.stop_event.is_set(): + while not stop_event.is_set(): try: frame = self._read_from_hardware() color_frame_raw = frame.get_color_frame() diff --git a/src/lerobot/cameras/zmq/camera_zmq.py b/src/lerobot/cameras/zmq/camera_zmq.py index 1b0be5de6..f3df17814 100644 --- a/src/lerobot/cameras/zmq/camera_zmq.py +++ b/src/lerobot/cameras/zmq/camera_zmq.py @@ -246,11 +246,12 @@ class ZMQCamera(Camera): """ Internal loop run by the background thread for asynchronous reading. """ - if self.stop_event is None: + stop_event = self.stop_event + if stop_event is None: raise RuntimeError(f"{self}: stop_event is not initialized.") failure_count = 0 - while not self.stop_event.is_set(): + while not stop_event.is_set(): try: frame = self._read_from_hardware() capture_time = time.perf_counter() From da92db8fc0c935950a56b1ea61fa9b211ef3ac30 Mon Sep 17 00:00:00 2001 From: Caroline Pascal Date: Wed, 17 Jun 2026 11:50:09 +0200 Subject: [PATCH 03/38] fix(image transforms): cleaning up image_transforms implementation in LeRobotDataset (#3829) --- src/lerobot/datasets/dataset_reader.py | 12 ++++++++++++ src/lerobot/datasets/lerobot_dataset.py | 12 +++++------- 2 files changed, 17 insertions(+), 7 deletions(-) diff --git a/src/lerobot/datasets/dataset_reader.py b/src/lerobot/datasets/dataset_reader.py index 59aaa40e5..d7289ac48 100644 --- a/src/lerobot/datasets/dataset_reader.py +++ b/src/lerobot/datasets/dataset_reader.py @@ -74,6 +74,8 @@ class DatasetReader: self.episodes = episodes self._tolerance_s = tolerance_s self._video_backend = video_backend + if image_transforms is not None and not callable(image_transforms): + raise TypeError("image_transforms must be callable or None.") self._image_transforms = image_transforms self._return_uint8 = return_uint8 @@ -86,6 +88,16 @@ class DatasetReader: check_delta_timestamps(delta_timestamps, meta.fps, tolerance_s) self.delta_indices = get_delta_indices(delta_timestamps, meta.fps) + def set_image_transforms(self, image_transforms: Callable | None) -> None: + """Replace the transform applied to visual observations.""" + if image_transforms is not None and not callable(image_transforms): + raise TypeError("image_transforms must be callable or None.") + self._image_transforms = image_transforms + + def clear_image_transforms(self) -> None: + """Remove the transform applied to visual observations.""" + self._image_transforms = None + def try_load(self) -> bool: """Attempt to load from local cache. Returns True if data is sufficient.""" try: diff --git a/src/lerobot/datasets/lerobot_dataset.py b/src/lerobot/datasets/lerobot_dataset.py index d0dcf087d..d1e65fef1 100644 --- a/src/lerobot/datasets/lerobot_dataset.py +++ b/src/lerobot/datasets/lerobot_dataset.py @@ -201,8 +201,6 @@ class LeRobotDataset(torch.utils.data.Dataset): super().__init__() self.repo_id = repo_id self._requested_root = Path(root) if root else None - self.reader = None - self.set_image_transforms(image_transforms) self.delta_timestamps = delta_timestamps self.tolerance_s = tolerance_s self.revision = revision if revision else CODEBASE_VERSION @@ -249,6 +247,7 @@ class LeRobotDataset(torch.utils.data.Dataset): image_transforms=image_transforms, return_uint8=self._return_uint8, ) + self.image_transforms = image_transforms # Load actual data if force_cache_sync or not self.reader.try_load(): @@ -505,15 +504,14 @@ class LeRobotDataset(torch.utils.data.Dataset): def set_image_transforms(self, image_transforms: Callable | None) -> None: """Replace the transform applied to visual observations.""" - if image_transforms is not None and not callable(image_transforms): - raise TypeError("image_transforms must be callable or None.") + self._ensure_reader().set_image_transforms(image_transforms) self.image_transforms = image_transforms - if self.reader is not None: - self.reader._image_transforms = image_transforms def clear_image_transforms(self) -> None: """Remove the transform applied to visual observations.""" - self.set_image_transforms(None) + if self.reader is not None: + self.reader.set_image_transforms(None) + self.image_transforms = None # ── Hub methods (stay on facade) ────────────────────────────────── From 8bf6056d147a703accbb53ed5415356a4cf9a8e5 Mon Sep 17 00:00:00 2001 From: Nicolas Rabault Date: Wed, 17 Jun 2026 18:22:21 +0200 Subject: [PATCH 04/38] docs: add LeLab web interface to README (#3831) --- README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/README.md b/README.md index fa3e9e1a3..2a330d823 100644 --- a/README.md +++ b/README.md @@ -136,6 +136,7 @@ Learn how to implement your own simulation environment or benchmark and distribu - **[X](https://x.com/LeRobotHF):** Follow us on X to stay up-to-date with the latest developments. - **[Robot Learning Tutorial](https://huggingface.co/spaces/lerobot/robot-learning-tutorial):** A free, hands-on course to learn robot learning using LeRobot. - **[T-Shirt Folding Experiment](https://huggingface.co/spaces/lerobot/robot-folding):** An end-to-end demonstration of folding t-shirts with LeRobot. +- **[LeLab](https://github.com/huggingface/leLab):** A web interface for LeRobot — teleoperate, calibrate, record datasets, replay, and train your SO arm from the browser, no CLI required. ## Citation From 552b4c3563a8801748c15d0f9a65a42e8db38485 Mon Sep 17 00:00:00 2001 From: Khalil Meftah Date: Fri, 19 Jun 2026 18:30:00 +0200 Subject: [PATCH 05/38] Add third-party env plugin discovery (#3823) * feat(envs): add env plugin discovery - Add 'lerobot_env_' to third-party plugin discovery prefixes, completing the plugin system for all component types (robots, cameras, teleoperators, policies, and now environments). External packages named lerobot_env_* can self-register EnvConfig subclasses on import, enabling --env.type= resolution without lerobot code changes. * feat(envs): add generic observation passthrough - Add generic observation passthrough in preprocess_observation() for unhandled ndarray/tensor keys, replacing the pattern of adding per-env hardcoded key handlers. Extra keys are forwarded as observation. and can be shaped by env-specific ProcessorSteps via get_env_processors(). --------- Co-authored-by: Steven Palma --- src/lerobot/envs/utils.py | 20 ++++++++++++++++++++ src/lerobot/utils/import_utils.py | 10 ++++++++-- 2 files changed, 28 insertions(+), 2 deletions(-) diff --git a/src/lerobot/envs/utils.py b/src/lerobot/envs/utils.py index 6e6f352e9..8b9c4f94b 100644 --- a/src/lerobot/envs/utils.py +++ b/src/lerobot/envs/utils.py @@ -126,6 +126,26 @@ def preprocess_observation(observations: dict[str, np.ndarray]) -> dict[str, Ten if "camera_obs" in observations: return_observations[f"{OBS_STR}.camera_obs"] = observations["camera_obs"] + # Pass through any remaining ndarray/tensor keys not already handled above, + # so env plugins can expose extra observation keys via get_env_processors(). + _handled = {"pixels", "environment_state", "agent_pos", "robot_state", "policy", "camera_obs"} + for key, value in observations.items(): + if key in _handled: + continue + target = f"{OBS_STR}.{key}" + if target in return_observations: + continue + if isinstance(value, np.ndarray): + val = torch.from_numpy(value).float() + if val.dim() == 1: + val = val.unsqueeze(0) + return_observations[target] = val + elif isinstance(value, Tensor): + val = value.float() + if val.dim() == 1: + val = val.unsqueeze(0) + return_observations[target] = val + return return_observations diff --git a/src/lerobot/utils/import_utils.py b/src/lerobot/utils/import_utils.py index 5dbce2c5b..b0d894c04 100644 --- a/src/lerobot/utils/import_utils.py +++ b/src/lerobot/utils/import_utils.py @@ -216,9 +216,15 @@ def register_third_party_plugins() -> None: This function uses `importlib.metadata` to find packages installed in the environment (including editable installs) starting with 'lerobot_robot_', 'lerobot_camera_', - 'lerobot_teleoperator_', or 'lerobot_policy_' and imports them. + 'lerobot_teleoperator_', 'lerobot_policy_', or 'lerobot_env_' and imports them. """ - prefixes = ("lerobot_robot_", "lerobot_camera_", "lerobot_teleoperator_", "lerobot_policy_") + prefixes = ( + "lerobot_robot_", + "lerobot_camera_", + "lerobot_teleoperator_", + "lerobot_policy_", + "lerobot_env_", + ) imported: list[str] = [] failed: list[str] = [] From b3d74f80f042a137b5ba46cf977e85a4da0e6ad3 Mon Sep 17 00:00:00 2001 From: Khalil Meftah Date: Fri, 19 Jun 2026 18:31:12 +0200 Subject: [PATCH 06/38] Fix batch wandb logging metrics and handle scalar stats (#3821) * fix(logging): batch wandb metrics - Batch all metrics into a single wandb.log() call instead of one per key, reducing API overhead. - Add support for list-valued metrics by expanding them to indexed keys (e.g. metric_0, metric_1). * fix(stats): handle scalar stats robustly - Wrap cast_stats_to_numpy with np.atleast_1d to prevent 0-d arrays from scalar stats causing shape mismatches downstream. * fix(logging): remove unused list-valued metric expansion --------- Co-authored-by: Steven Palma --- src/lerobot/common/wandb_utils.py | 20 +++++++++++--------- src/lerobot/datasets/io_utils.py | 2 +- 2 files changed, 12 insertions(+), 10 deletions(-) diff --git a/src/lerobot/common/wandb_utils.py b/src/lerobot/common/wandb_utils.py index b782cd751..c229b5eaa 100644 --- a/src/lerobot/common/wandb_utils.py +++ b/src/lerobot/common/wandb_utils.py @@ -180,24 +180,26 @@ class WandBLogger: self._wandb_custom_step_key.add(new_custom_key) self._wandb.define_metric(new_custom_key, hidden=True) + batch_data = {} for k, v in d.items(): + # Skip the custom step key here, it's added to the batch below. + if custom_step_key is not None and k == custom_step_key: + continue + if not isinstance(v, (int | float | str)): logging.warning( f'WandB logging of key "{k}" was ignored as its type "{type(v)}" is not handled by this wrapper.' ) continue - # Do not log the custom step key itself. - if self._wandb_custom_step_key is not None and k in self._wandb_custom_step_key: - continue + batch_data[f"{mode}/{k}"] = v + if batch_data: if custom_step_key is not None: - value_custom_step = d[custom_step_key] - data = {f"{mode}/{k}": v, f"{mode}/{custom_step_key}": value_custom_step} - self._wandb.log(data) - continue - - self._wandb.log(data={f"{mode}/{k}": v}, step=step) + batch_data[f"{mode}/{custom_step_key}"] = d[custom_step_key] + self._wandb.log(batch_data) + else: + self._wandb.log(data=batch_data, step=step) def log_video(self, video_path: str, step: int, mode: str = "train"): if mode not in {"train", "eval"}: diff --git a/src/lerobot/datasets/io_utils.py b/src/lerobot/datasets/io_utils.py index b6344942c..be94f3b3a 100644 --- a/src/lerobot/datasets/io_utils.py +++ b/src/lerobot/datasets/io_utils.py @@ -154,7 +154,7 @@ def cast_stats_to_numpy(stats: dict) -> dict[str, dict[str, np.ndarray]]: Returns: dict: The statistics dictionary with values cast to numpy arrays. """ - stats = {key: np.array(value) for key, value in flatten_dict(stats).items()} + stats = {key: np.atleast_1d(np.array(value)) for key, value in flatten_dict(stats).items()} return unflatten_dict(stats) From b06ad40888ff5045d014d1072ff1b1cac83e4010 Mon Sep 17 00:00:00 2001 From: Khalil Meftah Date: Fri, 19 Jun 2026 18:32:47 +0200 Subject: [PATCH 07/38] feat(hub): add pretrained_revision to pin Hub model versions (#3820) - Add pretrained_revision field to PreTrainedConfig (policies) and RewardModelConfig (reward models), and thread it through make_policy(), make_pre_post_processors(), and make_reward_model() so that weights and processor configs can be loaded from a specific Hub commit, branch, or tag. Defaults to None (latest version, preserving current behavior). Dataset and env hub loading already supported revision pinning. Co-authored-by: Steven Palma --- src/lerobot/configs/policies.py | 2 ++ src/lerobot/configs/rewards.py | 2 ++ src/lerobot/policies/factory.py | 4 ++++ src/lerobot/rewards/factory.py | 1 + src/lerobot/scripts/lerobot_train.py | 1 + 5 files changed, 10 insertions(+) diff --git a/src/lerobot/configs/policies.py b/src/lerobot/configs/policies.py index 91701af6d..b0f003519 100644 --- a/src/lerobot/configs/policies.py +++ b/src/lerobot/configs/policies.py @@ -79,6 +79,8 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC): # type: igno # Either the repo ID of a model hosted on the Hub or a path to a directory containing weights # saved using `Policy.save_pretrained`. If not provided, the policy is initialized from scratch. pretrained_path: Path | None = None + # Optional Hub revision (commit hash, branch, or tag) to pin the pretrained model version. + pretrained_revision: str | None = None def __post_init__(self) -> None: if not self.device or not is_torch_device_available(self.device): diff --git a/src/lerobot/configs/rewards.py b/src/lerobot/configs/rewards.py index 7e99e7f71..92490bc9f 100644 --- a/src/lerobot/configs/rewards.py +++ b/src/lerobot/configs/rewards.py @@ -56,6 +56,8 @@ class RewardModelConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC): device: str | None = None pretrained_path: str | None = None + # Optional Hub revision (commit hash, branch, or tag) to pin the pretrained reward model version. + pretrained_revision: str | None = None push_to_hub: bool = False repo_id: str | None = None diff --git a/src/lerobot/policies/factory.py b/src/lerobot/policies/factory.py index a42b38ba4..b82eaeb72 100644 --- a/src/lerobot/policies/factory.py +++ b/src/lerobot/policies/factory.py @@ -252,6 +252,7 @@ class ProcessorConfigKwargs(TypedDict, total=False): def make_pre_post_processors( policy_cfg: PreTrainedConfig, pretrained_path: str | None = None, + pretrained_revision: str | None = None, **kwargs: Unpack[ProcessorConfigKwargs], ) -> tuple[ PolicyProcessorPipeline[dict[str, Any], dict[str, Any]], @@ -309,6 +310,7 @@ def make_pre_post_processors( overrides=kwargs.get("preprocessor_overrides", {}), to_transition=batch_to_transition, to_output=transition_to_batch, + revision=pretrained_revision, ) postprocessor = PolicyProcessorPipeline.from_pretrained( pretrained_model_name_or_path=pretrained_path, @@ -318,6 +320,7 @@ def make_pre_post_processors( overrides=kwargs.get("postprocessor_overrides", {}), to_transition=policy_action_to_transition, to_output=transition_to_policy_action, + revision=pretrained_revision, ) _reconnect_relative_absolute_steps(preprocessor, postprocessor) return preprocessor, postprocessor @@ -557,6 +560,7 @@ def make_policy( # Load a pretrained policy and override the config if needed (for example, if there are inference-time # hyperparameters that we want to vary). 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: # Load a pretrained PEFT model on top of the policy. The pretrained path points to the folder/repo diff --git a/src/lerobot/rewards/factory.py b/src/lerobot/rewards/factory.py index 2d73ae575..fee90c211 100644 --- a/src/lerobot/rewards/factory.py +++ b/src/lerobot/rewards/factory.py @@ -124,6 +124,7 @@ def make_reward_model(cfg: RewardModelConfig, **kwargs) -> PreTrainedRewardModel if cfg.pretrained_path: kwargs["pretrained_name_or_path"] = cfg.pretrained_path + kwargs["revision"] = cfg.pretrained_revision reward_model = reward_cls.from_pretrained(**kwargs) else: reward_model = reward_cls(**kwargs) diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index 70a5e9e9d..c94223c62 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/src/lerobot/scripts/lerobot_train.py @@ -345,6 +345,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): preprocessor, postprocessor = make_pre_post_processors( policy_cfg=cfg.policy, pretrained_path=processor_pretrained_path, + pretrained_revision=getattr(cfg.policy, "pretrained_revision", None), **processor_kwargs, ) From 2d7a42011a4f8e05a8c85d5fb908da258d4cc7b1 Mon Sep 17 00:00:00 2001 From: Khalil Meftah Date: Sun, 21 Jun 2026 11:48:45 +0200 Subject: [PATCH 08/38] fix(policies): support offline batch inference for ACT and Diffusion (#3822) - Guard ACT's KL divergence computation against None latent params to prevent crashes during eval when use_vae is set but the forward path returns no VAE outputs. - Add offline batch fallback to Diffusion's predict_action_chunk() so it works with dataloader batches (empty queues) in addition to the existing online rollout path (populated queues). This enables batched action prediction for offline evaluation. --- src/lerobot/policies/act/modeling_act.py | 2 +- .../policies/diffusion/modeling_diffusion.py | 20 +++++++++++++++---- 2 files changed, 17 insertions(+), 5 deletions(-) diff --git a/src/lerobot/policies/act/modeling_act.py b/src/lerobot/policies/act/modeling_act.py index 5651fbfb1..1432b68a5 100644 --- a/src/lerobot/policies/act/modeling_act.py +++ b/src/lerobot/policies/act/modeling_act.py @@ -148,7 +148,7 @@ class ACTPolicy(PreTrainedPolicy): l1_loss = (abs_err * valid_mask).sum() / num_valid.clamp_min(1) loss_dict = {"l1_loss": l1_loss.item()} - if self.config.use_vae: + if self.config.use_vae and log_sigma_x2_hat is not None: # Calculate Dₖₗ(latent_pdf || standard_normal). Note: After computing the KL-divergence for # each dimension independently, we sum over the latent dimension to get the total # KL-divergence per batch element, then take the mean over the batch. diff --git a/src/lerobot/policies/diffusion/modeling_diffusion.py b/src/lerobot/policies/diffusion/modeling_diffusion.py index 9fbe1f703..8758a7e29 100644 --- a/src/lerobot/policies/diffusion/modeling_diffusion.py +++ b/src/lerobot/policies/diffusion/modeling_diffusion.py @@ -101,11 +101,23 @@ class DiffusionPolicy(PreTrainedPolicy): @torch.no_grad() def predict_action_chunk(self, batch: dict[str, Tensor], noise: Tensor | None = None) -> Tensor: - """Predict a chunk of actions given environment observations.""" - # stack n latest observations from the queue - batch = {k: torch.stack(list(self._queues[k]), dim=1) for k in batch if k in self._queues} - actions = self.diffusion.generate_actions(batch, noise=noise) + """Predict a chunk of actions given environment observations. + Supports two modes: + - Online (queues populated via select_action): stacks observations from internal queues. + - Offline (empty queues, e.g. dataloader batch): uses the batch directly. + """ + queues_populated = any(len(q) > 0 for q in self._queues.values()) + if queues_populated: + batch = {k: torch.stack(list(self._queues[k]), dim=1) for k in batch if k in self._queues} + else: + batch = dict(batch) + if self.config.image_features: + for key in self.config.image_features: + if batch[key].ndim == 4: + batch[key] = batch[key].unsqueeze(1) + batch[OBS_IMAGES] = torch.stack([batch[key] for key in self.config.image_features], dim=-4) + actions = self.diffusion.generate_actions(batch, noise=noise) return actions @torch.no_grad() From 73782447f2ca420f8d71c9bd0a169ece5968d2d6 Mon Sep 17 00:00:00 2001 From: Maxime Ellerbach Date: Mon, 22 Jun 2026 13:51:21 +0200 Subject: [PATCH 09/38] feat(train): FSDP checkpoint saving (#3810) * feat(train): FSDP checkpoint saving * adding docs for FSDP * adding a test for the fsdp checkpoint path * cleanup * fixing final upload to hub * refactored initial implementation to use torch fsdp api and adding new tests --- docs/source/multi_gpu_training.mdx | 55 ++++++++++++ src/lerobot/common/train_utils.py | 88 +++++++++++++++++-- src/lerobot/optim/__init__.py | 2 + src/lerobot/optim/optimizers.py | 35 ++++++-- src/lerobot/policies/pretrained.py | 43 +++++++++- src/lerobot/scripts/lerobot_train.py | 32 ++++++- tests/optim/test_optimizers.py | 39 +++++++++ tests/policies/test_policies.py | 24 ++++++ tests/training/test_multi_gpu.py | 121 ++++++++++++++++++++++++--- tests/utils/test_train_utils.py | 15 ++++ 10 files changed, 423 insertions(+), 31 deletions(-) diff --git a/docs/source/multi_gpu_training.mdx b/docs/source/multi_gpu_training.mdx index d7369e8f8..7907340c3 100644 --- a/docs/source/multi_gpu_training.mdx +++ b/docs/source/multi_gpu_training.mdx @@ -113,6 +113,61 @@ accelerate launch --num_processes=2 $(which lerobot-train) \ --policy=act ``` +## Training Large Models with FSDP + +DDP replicates the full model on every GPU, so a model that doesn't fit on one GPU won't fit under +DDP either. For large models, use **FSDP** (Fully Sharded Data Parallel), which shards parameters, +gradients, and optimizer state across GPUs. See the [accelerate FSDP guide](https://huggingface.co/docs/accelerate/usage_guides/fsdp) for background. + +An example on how to launch LeRobot training with FSDP across 4 GPUs (1 machine): + +```bash +accelerate launch --config_file fsdp.yaml --num_processes=4 $(which lerobot-train) \ + --dataset.repo_id=${HF_USER}/my_dataset \ + --policy.type= \ + --output_dir=outputs/train/my_policy_fsdp +``` + +A minimal `fsdp.yaml` (FSDP1; shards params/grads/optimizer — ZeRO-3-equivalent): + +```yaml +compute_environment: LOCAL_MACHINE +distributed_type: FSDP +mixed_precision: bf16 +num_machines: 1 +num_processes: 4 +fsdp_config: + fsdp_version: 1 + fsdp_sharding_strategy: FULL_SHARD # params + grads + optimizer (ZeRO-3) + fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP + fsdp_transformer_layer_cls_to_wrap: # repeated block class to shard + fsdp_use_orig_params: true # required: optimizer is built pre-prepare + fsdp_state_dict_type: FULL_STATE_DICT +``` + +Set `fsdp_transformer_layer_cls_to_wrap` to your model's repeated transformer-block class so each +block is sharded as its own unit. `fsdp_use_orig_params: true` is required because LeRobot builds the +optimizer before `accelerator.prepare()`. + +### FSDP checkpoints + +LeRobot gathers the full state dict across all ranks and the main process writes it as a single +`model.safetensors`, loadable as usual with `Policy.from_pretrained(...)`. Two things to look out for: + +- **Checkpoints store fp32 weights.** Under mixed precision (`bf16`/`fp16`) FSDP keeps an fp32 master + copy, and the checkpoint saves it (~2× the bf16 size on disk) so training can resume consistently + with the fp32 optimizer state; `from_pretrained` casts back to the policy dtype on load. FSDP-specific + caveat: an fp32 checkpoint is materialized in full precision on the target device _before_ casting, + so loading it for inference on a tight GPU can OOM even when the bf16 model would fit — load on CPU + first, or cast `model.safetensors` to the deployment dtype offline. +- The sharded optimizer state is gathered into a full (world-size-independent) state dict and saved + alongside the model in the same `optimizer_state.safetensors` / `optimizer_param_groups.json` + format as single-GPU training, so **resume-from-checkpoint is supported** with `--resume=true`. + Resume reshards both the model and the optimizer state to the _current_ FSDP topology, so you can + resume an FSDP checkpoint on a different number of GPUs. Note that the data sampler is only + sample-exact when the world size and batch size match the original run (a warning is logged + otherwise); the optimizer/model state itself is unaffected. + ## Notes - The `--policy.use_amp` flag in `lerobot-train` is only used when **not** running with accelerate. When using accelerate, mixed precision is controlled by accelerate's configuration. diff --git a/src/lerobot/common/train_utils.py b/src/lerobot/common/train_utils.py index 2d23b4003..5ae593bb8 100644 --- a/src/lerobot/common/train_utils.py +++ b/src/lerobot/common/train_utils.py @@ -21,6 +21,7 @@ from torch.optim.lr_scheduler import LRScheduler from lerobot.configs.train import TrainPipelineConfig from lerobot.optim import ( load_optimizer_state, + load_optimizer_state_dict, load_scheduler_state, save_optimizer_state, save_scheduler_state, @@ -98,6 +99,8 @@ def save_checkpoint( postprocessor: PolicyProcessorPipeline | None = None, num_processes: int | None = None, batch_size: int | None = None, + model_state_dict: dict | None = None, + optim_state_dict: dict | None = None, ) -> None: """This function creates the following directory structure: @@ -127,9 +130,18 @@ def save_checkpoint( resume. Defaults to None (not recorded). batch_size (int | None, optional): Per-process batch size to record for sample-exact resume. Defaults to None (not recorded). + model_state_dict: Pre-gathered full (unsharded) model state dict. Required under FSDP, + where `policy.state_dict()` would return sharded tensors; the caller gathers it via a + cross-rank collective and passes it here so rank 0 can write it directly. It holds + FSDP's fp32 master weights and is saved as-is (the loader casts to the policy dtype on + read). When None (DDP / single-GPU), the model is saved the normal way. Defaults to None. + optim_state_dict: Pre-gathered full (unsharded) optimizer state dict. Required under FSDP + (gathered alongside `model_state_dict` via `gather_fsdp_state_dicts`); saved in the same + safetensors format as the single-GPU path. When None, `optimizer.state_dict()` is used. + Defaults to None. """ pretrained_dir = checkpoint_dir / PRETRAINED_MODEL_DIR - policy.save_pretrained(pretrained_dir) + policy.save_pretrained(pretrained_dir, state_dict=model_state_dict) cfg.save_pretrained(pretrained_dir) if cfg.peft is not None: # When using PEFT, policy.save_pretrained will only write the adapter weights + config, not the @@ -140,7 +152,13 @@ def save_checkpoint( if postprocessor is not None: postprocessor.save_pretrained(pretrained_dir) save_training_state( - checkpoint_dir, step, optimizer, scheduler, num_processes=num_processes, batch_size=batch_size + checkpoint_dir, + step, + optimizer, + scheduler, + num_processes=num_processes, + batch_size=batch_size, + optim_state_dict=optim_state_dict, ) @@ -151,6 +169,7 @@ def save_training_state( scheduler: LRScheduler | None = None, num_processes: int | None = None, batch_size: int | None = None, + optim_state_dict: dict | None = None, ) -> None: """ Saves the training step, optimizer state, scheduler state, and rng state. @@ -164,19 +183,21 @@ def save_training_state( Defaults to None. num_processes (int | None, optional): Distributed world size to record. Defaults to None. batch_size (int | None, optional): Per-process batch size to record. Defaults to None. + optim_state_dict: Pre-gathered full optimizer state dict (for FSDP). Saved instead of + `optimizer.state_dict()` when provided. Defaults to None. """ save_dir = checkpoint_dir / TRAINING_STATE_DIR save_dir.mkdir(parents=True, exist_ok=True) save_training_step(train_step, save_dir, num_processes=num_processes, batch_size=batch_size) save_rng_state(save_dir) if optimizer is not None: - save_optimizer_state(optimizer, save_dir) + save_optimizer_state(optimizer, save_dir, optim_state_dict=optim_state_dict) if scheduler is not None: save_scheduler_state(scheduler, save_dir) def load_training_state( - checkpoint_dir: Path, optimizer: Optimizer, scheduler: LRScheduler | None + checkpoint_dir: Path, optimizer: Optimizer, scheduler: LRScheduler | None, load_optimizer: bool = True ) -> tuple[int, Optimizer, LRScheduler | None]: """ Loads the training step, optimizer state, scheduler state, and rng state. @@ -186,6 +207,10 @@ def load_training_state( checkpoint_dir (Path): The checkpoint directory. Should contain a 'training_state' dir. optimizer (Optimizer): The optimizer to load the state_dict to. scheduler (LRScheduler | None): The scheduler to load the state_dict to (can be None). + load_optimizer (bool, optional): Whether to load the optimizer state from disk. Defaults to + True. Set to False under FSDP, where the sharded optimizer state must be loaded after + `accelerator.prepare()` via `load_fsdp_optimizer_state` (the optimizer is returned + untouched here). Raises: NotADirectoryError: If 'checkpoint_dir' doesn't contain a 'training_state' dir @@ -200,8 +225,61 @@ def load_training_state( load_rng_state(training_state_dir) step = load_training_step(training_state_dir) - optimizer = load_optimizer_state(optimizer, training_state_dir) + if load_optimizer: + optimizer = load_optimizer_state(optimizer, training_state_dir) if scheduler is not None: scheduler = load_scheduler_state(scheduler, training_state_dir) return step, optimizer, scheduler + + +def gather_fsdp_state_dicts(model, optimizer) -> tuple[dict, dict]: + """Gather the full (unsharded) model and optimizer state dicts under FSDP. + + `model.state_dict()` and `FSDP.optim_state_dict(...)` are cross-rank collectives, so this must be + called on *every* rank with the prepared (FSDP-wrapped) `model` and `optimizer`. With + `rank0_only=True` and `offload_to_cpu=True`, every rank runs the all-gather but only rank 0 + materializes the full dicts (the others get empty dicts) and they are kept on CPU to bound GPU + memory. The returned optimizer state dict is keyed by parameter FQNs and is world-size + independent; `load_fsdp_optimizer_state` reshards it on resume. + + Returns: + (model_state_dict, optim_state_dict): full dicts on rank 0, empty dicts on other ranks. + """ + from torch.distributed.fsdp import ( + FullOptimStateDictConfig, + FullStateDictConfig, + FullyShardedDataParallel as FSDP, # noqa F401 + StateDictType, + ) + + state_cfg = FullStateDictConfig(offload_to_cpu=True, rank0_only=True) + optim_cfg = FullOptimStateDictConfig(offload_to_cpu=True, rank0_only=True) + with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, state_cfg, optim_cfg): + model_state_dict = model.state_dict() + optim_state_dict = FSDP.optim_state_dict(model, optimizer) + return model_state_dict, optim_state_dict + + +def load_fsdp_optimizer_state(model, optimizer, checkpoint_dir: Path) -> None: + """Load the FSDP optimizer state (saved as safetensors) and reshard it into the optimizer. + + This is a cross-rank collective and must be called on every rank *after* `accelerator.prepare()` + with the prepared (FSDP-wrapped) `model` and `optimizer`. The saved state is the full, + world-size-independent optimizer state (keyed by parameter FQNs); `FSDP.optim_state_dict_to_load` + reshards it to the current FSDP topology, so resume on a different number of GPUs works. + """ + from torch.distributed.fsdp import ( + FullOptimStateDictConfig, + FullStateDictConfig, + FullyShardedDataParallel as FSDP, # noqa F401 + StateDictType, + ) + + # Every rank reads the same full state from the (shared) checkpoint dir, so rank0_only=False. + full_osd = load_optimizer_state_dict(checkpoint_dir / TRAINING_STATE_DIR) + state_cfg = FullStateDictConfig(rank0_only=False) + optim_cfg = FullOptimStateDictConfig(rank0_only=False) + with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, state_cfg, optim_cfg): + sharded_osd = FSDP.optim_state_dict_to_load(model=model, optim=optimizer, optim_state_dict=full_osd) + optimizer.load_state_dict(sharded_osd) diff --git a/src/lerobot/optim/__init__.py b/src/lerobot/optim/__init__.py index 46676027b..2d564c25f 100644 --- a/src/lerobot/optim/__init__.py +++ b/src/lerobot/optim/__init__.py @@ -20,6 +20,7 @@ from .optimizers import ( SGDConfig as SGDConfig, XVLAAdamWConfig as XVLAAdamWConfig, load_optimizer_state, + load_optimizer_state_dict, save_optimizer_state, ) from .schedulers import ( @@ -50,6 +51,7 @@ __all__ = [ "VQBeTSchedulerConfig", # State management "load_optimizer_state", + "load_optimizer_state_dict", "load_scheduler_state", "save_optimizer_state", "save_scheduler_state", diff --git a/src/lerobot/optim/optimizers.py b/src/lerobot/optim/optimizers.py index 0bdd7a37e..0a462e1aa 100644 --- a/src/lerobot/optim/optimizers.py +++ b/src/lerobot/optim/optimizers.py @@ -27,7 +27,7 @@ from lerobot.utils.constants import ( OPTIMIZER_PARAM_GROUPS, OPTIMIZER_STATE, ) -from lerobot.utils.io_utils import deserialize_json_into_object, write_json +from lerobot.utils.io_utils import deserialize_json_into_object, load_json, write_json from lerobot.utils.utils import flatten_dict, unflatten_dict # Type alias for parameters accepted by optimizer build() methods. @@ -281,28 +281,37 @@ class MultiAdamConfig(OptimizerConfig): def save_optimizer_state( - optimizer: torch.optim.Optimizer | dict[str, torch.optim.Optimizer], save_dir: Path + optimizer: torch.optim.Optimizer | dict[str, torch.optim.Optimizer], + save_dir: Path, + optim_state_dict: dict | None = None, ) -> None: """Save optimizer state to disk. Args: optimizer: Either a single optimizer or a dictionary of optimizers. save_dir: Directory to save the optimizer state. + optim_state_dict: Pre-gathered optimizer state dict (for FSDP, where the sharded state must + be gathered across ranks first). If provided, it is saved directly instead of calling + ``optimizer.state_dict()``. Only supported for a single optimizer. Defaults to None. """ if isinstance(optimizer, dict): # Handle dictionary of optimizers + if optim_state_dict is not None: + raise ValueError("optim_state_dict is not supported for a dict of optimizers") for name, opt in optimizer.items(): optimizer_dir = save_dir / name optimizer_dir.mkdir(exist_ok=True, parents=True) _save_single_optimizer_state(opt, optimizer_dir) else: # Handle single optimizer - _save_single_optimizer_state(optimizer, save_dir) + _save_single_optimizer_state(optimizer, save_dir, optim_state_dict=optim_state_dict) -def _save_single_optimizer_state(optimizer: torch.optim.Optimizer, save_dir: Path) -> None: +def _save_single_optimizer_state( + optimizer: torch.optim.Optimizer, save_dir: Path, optim_state_dict: dict | None = None +) -> None: """Save a single optimizer's state to disk.""" - state = optimizer.state_dict() + state = dict(optim_state_dict) if optim_state_dict is not None else optimizer.state_dict() param_groups = state.pop("param_groups") flat_state = flatten_dict(state) save_file(flat_state, save_dir / OPTIMIZER_STATE) @@ -356,3 +365,19 @@ def _load_single_optimizer_state(optimizer: torch.optim.Optimizer, save_dir: Pat optimizer.load_state_dict(loaded_state_dict) return optimizer + + +def load_optimizer_state_dict(save_dir: Path) -> dict: + """Read a saved optimizer state dict (safetensors + json) back into a plain dict. + + Unlike `load_optimizer_state`, this does not load into an optimizer and preserves the original + ``state`` keys verbatim (e.g. FSDP parameter FQNs, which are not integer-castable). It is used by + the FSDP resume path, where the full state must be resharded via `FSDP.optim_state_dict_to_load` + before being loaded into the (sharded) optimizer. + """ + flat_state = load_file(save_dir / OPTIMIZER_STATE) + state = unflatten_dict(flat_state) + return { + "state": state.get("state", {}), + "param_groups": load_json(save_dir / OPTIMIZER_PARAM_GROUPS), + } diff --git a/src/lerobot/policies/pretrained.py b/src/lerobot/policies/pretrained.py index a69487f3f..a7aabb3f3 100644 --- a/src/lerobot/policies/pretrained.py +++ b/src/lerobot/policies/pretrained.py @@ -23,7 +23,7 @@ from typing import TypedDict, TypeVar, Unpack import packaging import safetensors -from huggingface_hub import HfApi, ModelCard, ModelCardData, hf_hub_download +from huggingface_hub import HfApi, ModelCard, ModelCardData, hf_hub_download, save_torch_state_dict 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 @@ -129,10 +129,43 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC): if not getattr(cls, "name", None): raise TypeError(f"Class {cls.__name__} must define 'name'") - def _save_pretrained(self, save_directory: Path) -> None: + def save_pretrained( + self, + save_directory: str | Path, + *, + state_dict: dict[str, Tensor] | None = None, + repo_id: str | None = None, + push_to_hub: bool = False, + card_kwargs: dict | None = None, + **push_to_hub_kwargs, + ) -> str | None: + """Save the policy to a directory (and optionally push to the Hub). + + Overrides `HubMixin.save_pretrained` to add a `state_dict` argument (mirroring + `transformers.PreTrainedModel.save_pretrained`). Under FSDP, `self.state_dict()` would + return sharded tensors, so the caller gathers the full state dict via a cross-rank + collective and passes it here for `_save_pretrained` to write directly. + """ + save_directory = Path(save_directory) + save_directory.mkdir(parents=True, exist_ok=True) + self._save_pretrained(save_directory, state_dict=state_dict) + if push_to_hub: + if repo_id is None: + repo_id = save_directory.name + return self.push_to_hub(repo_id=repo_id, card_kwargs=card_kwargs, **push_to_hub_kwargs) + return None + + def _save_pretrained(self, save_directory: Path, state_dict: dict[str, Tensor] | None = None) -> None: self.config._save_pretrained(save_directory) model_to_save = self.module if hasattr(self, "module") else self - save_model_as_safetensor(model_to_save, str(save_directory / SAFETENSORS_SINGLE_FILE)) + if state_dict is None: + save_model_as_safetensor(model_to_save, str(save_directory / SAFETENSORS_SINGLE_FILE)) + return + # A pre-gathered (e.g. FSDP full) state dict was supplied: write it directly. + # `save_torch_state_dict` discards shared-tensor duplicates just like `save_model` does; + # pin `max_shard_size` above the total size so the output stays a single `model.safetensors` + total_bytes = sum(t.numel() * t.element_size() for t in state_dict.values()) + save_torch_state_dict(state_dict, str(save_directory), max_shard_size=max(total_bytes, 1)) @classmethod def from_pretrained( @@ -270,6 +303,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC): self, cfg: TrainPipelineConfig, peft_model=None, + state_dict: dict[str, Tensor] | None = None, ): api = HfApi() repo_id = api.create_repo( @@ -287,7 +321,8 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC): peft_model.save_pretrained(saved_path) self.config.save_pretrained(saved_path) else: - self.save_pretrained(saved_path) # Calls _save_pretrained and stores model tensors + # Calls _save_pretrained and stores model tensors + self.save_pretrained(saved_path, state_dict=state_dict) card = self.generate_model_card( cfg.dataset.repo_id, self.config.type, self.config.license, self.config.tags, cfg=cfg diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index c94223c62..45281dac9 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/src/lerobot/scripts/lerobot_train.py @@ -34,8 +34,10 @@ from torch.optim import Optimizer from tqdm import tqdm from lerobot.common.train_utils import ( + gather_fsdp_state_dicts, get_step_checkpoint_dir, get_step_identifier, + load_fsdp_optimizer_state, load_training_batch_size, load_training_num_processes, load_training_state, @@ -189,6 +191,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): require_package("accelerate", extra="training") from accelerate import Accelerator + from accelerate.utils import DistributedDataParallelKwargs, DistributedType cfg.validate() @@ -197,8 +200,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): # We set step_scheduler_with_optimizer=False to prevent accelerate from adjusting the lr_scheduler steps based on the num_processes # We set find_unused_parameters=True to handle models with conditional computation if accelerator is None: - from accelerate.utils import DistributedDataParallelKwargs - ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True) # Accelerate auto-detects the device based on the available hardware and ignores the policy.device setting. # Force the device to be CPU when the active config's device is set to CPU (works for both policy and reward model training). @@ -371,7 +372,12 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): step = 0 # number of policy updates (forward + backward + optim) if cfg.resume: - step, optimizer, lr_scheduler = load_training_state(cfg.checkpoint_path, optimizer, lr_scheduler) + # Under FSDP the optimizer state is sharded and must be loaded after `accelerator.prepare()` + # (see load_fsdp_optimizer_state below), so skip the optimizer here and load it then. + is_fsdp = accelerator.distributed_type == DistributedType.FSDP + step, optimizer, lr_scheduler = load_training_state( + cfg.checkpoint_path, optimizer, lr_scheduler, load_optimizer=not is_fsdp + ) num_learnable_params = sum(p.numel() for p in policy.parameters() if p.requires_grad) num_total_params = sum(p.numel() for p in policy.parameters()) @@ -461,6 +467,12 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): policy, optimizer, dataloader, lr_scheduler = accelerator.prepare( policy, optimizer, dataloader, lr_scheduler ) + + # FSDP optimizer state is sharded across ranks, so it can only be loaded once the optimizer and + # model are FSDP-wrapped (i.e. after `prepare`). Collective: every rank must participate. + if cfg.resume and accelerator.distributed_type == DistributedType.FSDP: + load_fsdp_optimizer_state(policy, optimizer, cfg.checkpoint_path) + dl_iter = cycle(dataloader) policy.train() @@ -559,6 +571,14 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): train_tracker.reset_averages() if cfg.save_checkpoint and is_saving_step: + # Under FSDP, gathering the full model + optimizer state dicts is a cross-rank collective, + # so all ranks must participate; rank 0 then writes the materialized dicts. For DDP / + # single-GPU the state dicts are saved the normal way inside save_checkpoint. + is_fsdp = accelerator.distributed_type == DistributedType.FSDP + if is_fsdp: + model_state_dict, optim_state_dict = gather_fsdp_state_dicts(policy, optimizer) + else: + model_state_dict, optim_state_dict = None, None if is_main_process: logging.info(f"Checkpoint policy after step {step}") checkpoint_dir = get_step_checkpoint_dir(cfg.output_dir, cfg.steps, step) @@ -573,6 +593,8 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): postprocessor=postprocessor, num_processes=accelerator.num_processes, batch_size=cfg.batch_size, + model_state_dict=model_state_dict, + optim_state_dict=optim_state_dict, ) update_last_checkpoint(checkpoint_dir) if wandb_logger: @@ -635,6 +657,8 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): if eval_env: close_envs(eval_env) + is_fsdp = accelerator.distributed_type == DistributedType.FSDP + model_state_dict = accelerator.get_state_dict(policy) if is_fsdp else None if is_main_process: logging.info("End of training") @@ -644,7 +668,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): if not cfg.is_reward_model_training and cfg.policy.use_peft: unwrapped_model.push_model_to_hub(cfg, peft_model=unwrapped_model) else: - unwrapped_model.push_model_to_hub(cfg) + unwrapped_model.push_model_to_hub(cfg, state_dict=model_state_dict) preprocessor.push_to_hub(active_cfg.repo_id) postprocessor.push_to_hub(active_cfg.repo_id) diff --git a/tests/optim/test_optimizers.py b/tests/optim/test_optimizers.py index d18565562..5b480f70d 100644 --- a/tests/optim/test_optimizers.py +++ b/tests/optim/test_optimizers.py @@ -20,6 +20,7 @@ from lerobot.optim.optimizers import ( MultiAdamConfig, SGDConfig, load_optimizer_state, + load_optimizer_state_dict, save_optimizer_state, ) from lerobot.utils.constants import ( @@ -65,6 +66,44 @@ def test_save_and_load_optimizer_state(model_params, optimizer, tmp_path): torch.testing.assert_close(optimizer.state_dict(), loaded_optimizer.state_dict()) +def test_save_and_load_fsdp_optimizer_state_dict_roundtrip(tmp_path): + """The FSDP full optimizer state dict is keyed by parameter FQNs (dotted strings), not the + integer indices of the single-GPU path. Verify it survives the safetensors save -> read + round-trip used by the FSDP save/resume path (save_optimizer_state(optim_state_dict=...) then + load_optimizer_state_dict), which the flatten/unflatten "/" separator must not corrupt.""" + full_osd = { + "state": { + "model.layers.0.weight": { + "step": torch.tensor(3.0), + "exp_avg": torch.randn(4, 4), + "exp_avg_sq": torch.randn(4, 4), + }, + "model.layers.0.bias": { + "step": torch.tensor(3.0), + "exp_avg": torch.randn(4), + "exp_avg_sq": torch.randn(4), + }, + }, + "param_groups": [ + {"lr": 1e-4, "betas": [0.9, 0.999], "eps": 1e-8, "weight_decay": 0.0, "params": [0, 1]} + ], + } + + save_optimizer_state( + torch.optim.Adam([torch.nn.Parameter(torch.randn(1))]), tmp_path, optim_state_dict=full_osd + ) + assert (tmp_path / OPTIMIZER_STATE).is_file() + assert (tmp_path / OPTIMIZER_PARAM_GROUPS).is_file() + + loaded = load_optimizer_state_dict(tmp_path) + # FQN keys must be preserved verbatim (not int-cast, not split on their dots). + assert set(loaded["state"].keys()) == set(full_osd["state"].keys()) + for fqn, sub in full_osd["state"].items(): + for k, v in sub.items(): + torch.testing.assert_close(loaded["state"][fqn][k], v) + assert loaded["param_groups"] == full_osd["param_groups"] + + @pytest.fixture def base_params_dict(): return { diff --git a/tests/policies/test_policies.py b/tests/policies/test_policies.py index e9388b3ed..285b87d4c 100644 --- a/tests/policies/test_policies.py +++ b/tests/policies/test_policies.py @@ -23,6 +23,7 @@ import torch pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])") +from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE from packaging import version from safetensors.torch import load_file @@ -300,6 +301,29 @@ def test_save_and_load_pretrained(dummy_dataset_metadata, tmp_path, policy_name: torch.testing.assert_close(list(policy.parameters()), list(loaded_policy.parameters()), rtol=0, atol=0) +def test_save_pretrained_with_state_dict(dummy_dataset_metadata, tmp_path): + """Exercise the FSDP checkpoint path: save_pretrained with a pre-gathered state_dict.""" + policy_cls = get_policy_class("act") + policy_cfg = make_policy_config("act") + features = dataset_to_policy_features(dummy_dataset_metadata.features) + policy_cfg.output_features = {key: ft for key, ft in features.items() if ft.type is FeatureType.ACTION} + policy_cfg.input_features = { + key: ft for key, ft in features.items() if key not in policy_cfg.output_features + } + policy = policy_cls(policy_cfg) + policy.to(policy_cfg.device) + + save_dir = tmp_path / "fsdp_state_dict" + policy.save_pretrained(save_dir, state_dict=policy.state_dict()) + + # A single, unsharded safetensors file (no sharded set + index). + assert (save_dir / SAFETENSORS_SINGLE_FILE).is_file() + assert not (save_dir / f"{SAFETENSORS_SINGLE_FILE}.index.json").exists() + + loaded_policy = policy_cls.from_pretrained(save_dir, config=policy_cfg) + torch.testing.assert_close(list(policy.parameters()), list(loaded_policy.parameters()), rtol=0, atol=0) + + @pytest.mark.parametrize("multikey", [True, False]) def test_multikey_construction(multikey: bool): """ diff --git a/tests/training/test_multi_gpu.py b/tests/training/test_multi_gpu.py index 638dc3131..d37f1e35d 100644 --- a/tests/training/test_multi_gpu.py +++ b/tests/training/test_multi_gpu.py @@ -58,7 +58,46 @@ def download_dataset(repo_id, episodes): print(f"Dataset {repo_id} downloaded successfully") -def run_accelerate_training(config_args, num_processes=4, temp_dir=None): +def _write_multi_gpu_config(f, num_processes): + f.write("compute_environment: LOCAL_MACHINE\n") + f.write("distributed_type: MULTI_GPU\n") + f.write("mixed_precision: 'no'\n") + f.write(f"num_processes: {num_processes}\n") + f.write("use_cpu: false\n") + f.write("gpu_ids: all\n") + f.write("downcast_bf16: 'no'\n") + f.write("machine_rank: 0\n") + f.write("main_training_function: main\n") + f.write("num_machines: 1\n") + f.write("rdzv_backend: static\n") + f.write("same_network: true\n") + + +def _write_fsdp_config(f, num_processes): + # FSDP1 with FULL_SHARD (ZeRO-3-equivalent) and FULL_STATE_DICT, matching + # docs/source/multi_gpu_training.mdx. ACT's repeated transformer blocks are the wrap units; + # fsdp_use_orig_params is required because LeRobot builds the optimizer before prepare(). + f.write("compute_environment: LOCAL_MACHINE\n") + f.write("distributed_type: FSDP\n") + f.write("mixed_precision: 'no'\n") + f.write(f"num_processes: {num_processes}\n") + f.write("use_cpu: false\n") + f.write("gpu_ids: all\n") + f.write("machine_rank: 0\n") + f.write("main_training_function: main\n") + f.write("num_machines: 1\n") + f.write("rdzv_backend: static\n") + f.write("same_network: true\n") + f.write("fsdp_config:\n") + f.write(" fsdp_version: 1\n") + f.write(" fsdp_sharding_strategy: FULL_SHARD\n") + f.write(" fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP\n") + f.write(" fsdp_transformer_layer_cls_to_wrap: ACTEncoderLayer,ACTDecoderLayer\n") + f.write(" fsdp_use_orig_params: true\n") + f.write(" fsdp_state_dict_type: FULL_STATE_DICT\n") + + +def run_accelerate_training(config_args, num_processes=4, temp_dir=None, distributed_type="MULTI_GPU"): """ Helper function to run training with accelerate launch. @@ -66,6 +105,7 @@ def run_accelerate_training(config_args, num_processes=4, temp_dir=None): config_args: List of config arguments to pass to lerobot_train.py num_processes: Number of processes (GPUs) to use temp_dir: Temporary directory for outputs + distributed_type: "MULTI_GPU" (DDP) or "FSDP" — selects the generated accelerate config. Returns: subprocess.CompletedProcess result @@ -75,18 +115,10 @@ def run_accelerate_training(config_args, num_processes=4, temp_dir=None): # Write YAML config with open(config_path, "w") as f: - f.write("compute_environment: LOCAL_MACHINE\n") - f.write("distributed_type: MULTI_GPU\n") - f.write("mixed_precision: 'no'\n") - f.write(f"num_processes: {num_processes}\n") - f.write("use_cpu: false\n") - f.write("gpu_ids: all\n") - f.write("downcast_bf16: 'no'\n") - f.write("machine_rank: 0\n") - f.write("main_training_function: main\n") - f.write("num_machines: 1\n") - f.write("rdzv_backend: static\n") - f.write("same_network: true\n") + if distributed_type == "FSDP": + _write_fsdp_config(f, num_processes) + else: + _write_multi_gpu_config(f, num_processes) cmd = [ "accelerate", @@ -211,3 +243,66 @@ class TestMultiGPUTraining: # Verify optimizer state exists optimizer_state = training_state_dir / "optimizer_state.safetensors" assert optimizer_state.exists(), f"No optimizer state in checkpoint {checkpoint_dir}" + + def test_fsdp_optimizer_save_and_resume(self): + """ + Test that FSDP saves the (gathered) optimizer state and can resume from it. + + Trains a few steps under FSDP, verifies the gathered optimizer state is written next to the + rest of the training state, then resumes from the checkpoint for more steps and checks it + completes without shape/key errors in the FSDP optimizer load path. + """ + # Pre-download dataset to avoid race conditions + download_dataset("lerobot/pusht", episodes=[0]) + + with tempfile.TemporaryDirectory() as temp_dir: + output_dir = Path(temp_dir) / "outputs" + + config_args = [ + "--dataset.repo_id=lerobot/pusht", + "--dataset.episodes=[0]", + "--policy.type=act", + "--policy.device=cuda", + "--policy.push_to_hub=false", + f"--output_dir={output_dir}", + "--batch_size=4", + "--steps=10", + "--eval_freq=-1", + "--log_freq=5", + "--save_freq=10", + "--seed=42", + "--num_workers=0", + ] + + result = run_accelerate_training( + config_args, num_processes=2, temp_dir=temp_dir, distributed_type="FSDP" + ) + assert result.returncode == 0, ( + f"FSDP training failed:\nSTDOUT:\n{result.stdout}\n\nSTDERR:\n{result.stderr}" + ) + + # The gathered optimizer state must be written under FSDP (proves the save collective ran), + # in the same safetensors format as single-GPU training. + training_state_dir = output_dir / "checkpoints" / "last" / "training_state" + optimizer_state = training_state_dir / "optimizer_state.safetensors" + optimizer_param_groups = training_state_dir / "optimizer_param_groups.json" + assert optimizer_state.exists(), f"FSDP optimizer state not saved in {training_state_dir}" + assert optimizer_param_groups.exists(), ( + f"FSDP optimizer param groups not saved in {training_state_dir}" + ) + + # Resume from the checkpoint for more steps. A successful run proves load_fsdp_optimizer + # accepts the saved state and reshards it without shape/key errors. + resume_config = output_dir / "checkpoints" / "last" / "pretrained_model" / "train_config.json" + resume_args = [ + f"--config_path={resume_config}", + "--resume=true", + "--steps=20", + ] + resume_result = run_accelerate_training( + resume_args, num_processes=2, temp_dir=temp_dir, distributed_type="FSDP" + ) + assert resume_result.returncode == 0, ( + f"FSDP resume failed:\nSTDOUT:\n{resume_result.stdout}\n\nSTDERR:\n{resume_result.stderr}" + ) + assert "End of training" in resume_result.stdout or "End of training" in resume_result.stderr diff --git a/tests/utils/test_train_utils.py b/tests/utils/test_train_utils.py index c171763c2..e3705409b 100644 --- a/tests/utils/test_train_utils.py +++ b/tests/utils/test_train_utils.py @@ -136,3 +136,18 @@ def test_save_load_training_state(tmp_path, optimizer, scheduler): assert loaded_step == 10 assert loaded_optimizer is optimizer assert loaded_scheduler is scheduler + + +def test_load_training_state_skip_optimizer(tmp_path, optimizer, scheduler): + # FSDP loads optimizer separately (after accelerator.prepare) + # load_training_state(load_optimizer=False) must restore step + scheduler but leave the + # optimizer untouched and never touch the on-disk optimizer state. + save_training_state(tmp_path, 10, optimizer, scheduler) + with patch("lerobot.common.train_utils.load_optimizer_state") as mock_load_optimizer_state: + loaded_step, loaded_optimizer, loaded_scheduler = load_training_state( + tmp_path, optimizer, scheduler, load_optimizer=False + ) + mock_load_optimizer_state.assert_not_called() + assert loaded_step == 10 + assert loaded_optimizer is optimizer + assert loaded_scheduler is scheduler From 6f0ba4be3853b773ad2e8887be2761bc5474fc89 Mon Sep 17 00:00:00 2001 From: Khalil Meftah Date: Tue, 23 Jun 2026 14:03:57 +0200 Subject: [PATCH 10/38] Record eval rollouts as LeRobot datasets (#3825) * feat(eval): record eval rollouts as raw LeRobot datasets - Record raw env observations inline during rollout(), before preprocess_observation() transforms them. Uses LeRobotDataset.create() with add_frame()/save_episode(). - Supports vectorized envs: each env in the batch records independently, with save_episode() called per env on termination. Each task gets its own dataset under output_dir/recordings/{task_group}_{task_id}/. Enabled via --eval.recording=true; disabled by default. * fix(eval): use FeatureType enum comparison instead of string value * refactor(eval): per-env datasets recording, no double reset - Extract _infer_shape_from_obs() to reduce nesting in feature conversion - Move dataset creation into rollout() using its own env.reset() observation, eliminating the extra reset in run_one() - Replace deepcopy with _shallow_copy_obs() for raw observation stashing - Support batch_size > 1: each parallel env records to its own dataset (single env skips the env_0/ nesting for simplicity) - One-time warning for env_features keys missing from observations - Pass recording_dir + env_features through the call chain instead of a pre-built recording_dataset object * refactor(eval): remove shape inference and shallow copy helpers * feat(eval): optionally push recorded eval datasets to the Hub * fix(eval): address review comments - Wrap rollout loop in try/finally so finalize() runs on crash/interrupt - Guard push_to_hub with num_episodes > 0 to avoid pushing empty datasets - Hoist loop-invariant multi_env and base_repo_id out of creation loop --- src/lerobot/configs/default.py | 9 + src/lerobot/scripts/lerobot_eval.py | 308 +++++++++++++++++++++------- 2 files changed, 248 insertions(+), 69 deletions(-) diff --git a/src/lerobot/configs/default.py b/src/lerobot/configs/default.py index b809e71d9..648e03f33 100644 --- a/src/lerobot/configs/default.py +++ b/src/lerobot/configs/default.py @@ -73,8 +73,17 @@ class EvalConfig: # `use_async_envs` specifies whether to use asynchronous environments (multiprocessing). # Defaults to True; automatically downgraded to SyncVectorEnv when batch_size=1. use_async_envs: bool = True + # Whether to record eval rollouts as a LeRobot dataset on disk. + recording: bool = False + # If set, push recorded eval datasets to the Hub under this repo id (one repo per task, + # suffixed by task and env index). Requires recording=true. + recording_repo_id: str | None = None + # Whether the pushed recording repositories should be private. + recording_private: bool = False def __post_init__(self) -> None: + if self.recording_repo_id is not None and not self.recording: + raise ValueError("eval.recording_repo_id requires eval.recording=true.") if self.batch_size == 0: self.batch_size = self._auto_batch_size() if self.batch_size > self.n_episodes: diff --git a/src/lerobot/scripts/lerobot_eval.py b/src/lerobot/scripts/lerobot_eval.py index d45483d21..1ec4ea75f 100644 --- a/src/lerobot/scripts/lerobot_eval.py +++ b/src/lerobot/scripts/lerobot_eval.py @@ -72,8 +72,9 @@ from termcolor import colored from torch import Tensor, nn from tqdm import trange -from lerobot.configs import parser +from lerobot.configs import FeatureType, parser from lerobot.configs.eval import EvalPipelineConfig +from lerobot.datasets.lerobot_dataset import LeRobotDataset from lerobot.envs import ( check_env_attributes_and_types, close_envs, @@ -84,7 +85,7 @@ from lerobot.envs import ( from lerobot.policies import PreTrainedPolicy, make_policy, make_pre_post_processors from lerobot.processor import PolicyProcessorPipeline from lerobot.types import PolicyAction -from lerobot.utils.constants import ACTION, DONE, OBS_STR, REWARD +from lerobot.utils.constants import ACTION, DONE, OBS_IMAGE, OBS_IMAGES, OBS_STR, REWARD from lerobot.utils.device_utils import get_safe_torch_device from lerobot.utils.import_utils import register_third_party_plugins from lerobot.utils.io_utils import write_video @@ -95,6 +96,65 @@ from lerobot.utils.utils import ( ) +def _env_features_to_dataset_features(env_features: dict) -> dict: + """Convert EnvConfig.features to the dict format expected by LeRobotDataset.create().""" + features = {} + for key, ft in env_features.items(): + shape = tuple(ft.shape) + if ft.type is FeatureType.VISUAL: + features[key] = {"dtype": "video", "shape": shape, "names": ["height", "width", "channel"]} + else: + features[key] = {"dtype": "float32", "shape": shape, "names": None} + features["next.reward"] = {"dtype": "float32", "shape": (1,), "names": None} + features["next.success"] = {"dtype": "bool", "shape": (1,), "names": None} + features["next.done"] = {"dtype": "bool", "shape": (1,), "names": None} + return features + + +def _build_raw_frame( + raw_obs: dict, + env_idx: int, + action: np.ndarray, + reward: float, + success: bool, + done: bool, + task: str, + env_features: dict, +) -> dict: + """Build a dataset frame from raw env observations for one env index. + + Keys in the frame match the keys in env_features so they align with the + dataset schema created by _env_features_to_dataset_features(). + """ + frame: dict[str, Any] = {} + for key in env_features: + if key == ACTION: + continue + if key.startswith("next."): + continue + if "pixels" in raw_obs and isinstance(raw_obs["pixels"], dict): + for cam_name, img in raw_obs["pixels"].items(): + candidate = f"{OBS_IMAGES}.{cam_name}" + if candidate == key: + frame[key] = img[env_idx] + if key in frame: + continue + if "pixels" in raw_obs and not isinstance(raw_obs["pixels"], dict) and key in ("pixels", OBS_IMAGE): + frame[key] = raw_obs["pixels"][env_idx] + continue + if key in raw_obs and isinstance(raw_obs[key], np.ndarray): + val = raw_obs[key][env_idx] + if val.dtype == np.float64: + val = val.astype(np.float32) + frame[key] = val + frame[ACTION] = action + frame["next.reward"] = np.atleast_1d(np.float32(reward)) + frame["next.success"] = np.atleast_1d(np.bool_(success)) + frame["next.done"] = np.atleast_1d(np.bool_(done)) + frame["task"] = task + return frame + + def rollout( env: gym.vector.VectorEnv, policy: PreTrainedPolicy, @@ -105,6 +165,10 @@ def rollout( seeds: list[int] | None = None, return_observations: bool = False, render_callback: Callable[[gym.vector.VectorEnv], None] | None = None, + recording_dir: Path | None = None, + env_features: dict | None = None, + recording_repo_id: str | None = None, + recording_private: bool = False, ) -> dict: """Run a batched policy rollout once through a batch of environments. @@ -145,6 +209,33 @@ def rollout( if render_callback is not None: render_callback(env) + recording_datasets: list[LeRobotDataset] | None = None + raw_observation = None + task_desc = "" + if recording_dir is not None and env_features is not None: + features = _env_features_to_dataset_features(env_features) + fps = env.unwrapped.metadata.get("render_fps", 30) + recording_datasets = [] + multi_env = env.num_envs > 1 + base_repo_id = recording_repo_id or "eval_recording" + for i in range(env.num_envs): + root = str(recording_dir / f"env_{i}") if multi_env else str(recording_dir) + repo_id = f"{base_repo_id}_env_{i}" if multi_env else base_repo_id + recording_datasets.append( + LeRobotDataset.create( + repo_id=repo_id, + fps=fps, + features=features, + root=root, + use_videos=True, + ) + ) + raw_observation = deepcopy(observation) + try: + task_desc = list(env.call("task_description"))[0] + except (AttributeError, NotImplementedError): + task_desc = "" + all_observations = [] all_actions = [] all_rewards = [] @@ -162,80 +253,112 @@ def rollout( leave=False, ) check_env_attributes_and_types(env) - while not np.all(done) and step < max_steps: - # Numpy array to tensor and changing dictionary keys to LeRobot policy format. - observation = preprocess_observation(observation) - if return_observations: - all_observations.append(deepcopy(observation)) + try: + while not np.all(done) and step < max_steps: + # Numpy array to tensor and changing dictionary keys to LeRobot policy format. + observation = preprocess_observation(observation) + if return_observations: + all_observations.append(deepcopy(observation)) - # Infer "task" from sub-environments (prefer natural language description). - # env.call() works with both SyncVectorEnv and AsyncVectorEnv. - try: - observation["task"] = list(env.call("task_description")) - except (AttributeError, NotImplementedError): + # Infer "task" from sub-environments (prefer natural language description). + # env.call() works with both SyncVectorEnv and AsyncVectorEnv. try: - observation["task"] = list(env.call("task")) + observation["task"] = list(env.call("task_description")) except (AttributeError, NotImplementedError): - observation["task"] = [""] * env.num_envs + try: + observation["task"] = list(env.call("task")) + except (AttributeError, NotImplementedError): + observation["task"] = [""] * env.num_envs - # Apply environment-specific preprocessing (e.g., LiberoProcessorStep for LIBERO) - observation = env_preprocessor(observation) + # Apply environment-specific preprocessing (e.g., LiberoProcessorStep for LIBERO) + observation = env_preprocessor(observation) - observation = preprocessor(observation) - with torch.inference_mode(): - action = policy.select_action(observation) - action = postprocessor(action) + observation = preprocessor(observation) + with torch.inference_mode(): + action = policy.select_action(observation) + action = postprocessor(action) - action_transition = {ACTION: action} - action_transition = env_postprocessor(action_transition) - action = action_transition[ACTION] + action_transition = {ACTION: action} + action_transition = env_postprocessor(action_transition) + action = action_transition[ACTION] - # Convert to CPU / numpy. - action_numpy: np.ndarray = action.to("cpu").numpy() - assert action_numpy.ndim == 2, "Action dimensions should be (batch, action_dim)" + # Convert to CPU / numpy. + action_numpy: np.ndarray = action.to("cpu").numpy() + assert action_numpy.ndim == 2, "Action dimensions should be (batch, action_dim)" - # Apply the next action. - observation, reward, terminated, truncated, info = env.step(action_numpy) - if render_callback is not None: - render_callback(env) + # Apply the next action. + observation, reward, terminated, truncated, info = env.step(action_numpy) + if render_callback is not None: + render_callback(env) - # VectorEnv stores is_success in `info["final_info"][env_index]["is_success"]`. "final_info" isn't - # available if none of the envs finished. - if "final_info" in info: - final_info = info["final_info"] - if not isinstance(final_info, dict): - raise RuntimeError( - "Unsupported `final_info` format: expected dict (Gymnasium >= 1.0). " - "You're likely using an older version of gymnasium (< 1.0). Please upgrade." + # VectorEnv stores is_success in `info["final_info"][env_index]["is_success"]`. "final_info" isn't + # available if none of the envs finished. + if "final_info" in info: + final_info = info["final_info"] + if not isinstance(final_info, dict): + raise RuntimeError( + "Unsupported `final_info` format: expected dict (Gymnasium >= 1.0). " + "You're likely using an older version of gymnasium (< 1.0). Please upgrade." + ) + successes = final_info["is_success"].tolist() + elif "is_success" in info: + is_success = info["is_success"] + successes = ( + is_success.tolist() + if hasattr(is_success, "tolist") + else [bool(is_success)] * env.num_envs ) - successes = final_info["is_success"].tolist() - elif "is_success" in info: - is_success = info["is_success"] - successes = ( - is_success.tolist() if hasattr(is_success, "tolist") else [bool(is_success)] * env.num_envs + else: + successes = [False] * env.num_envs + + if recording_datasets is not None and raw_observation is not None: + prev_done = done.copy() + for env_idx in range(env.num_envs): + if prev_done[env_idx]: + continue + frame = _build_raw_frame( + raw_observation, + env_idx, + action_numpy[env_idx], + reward[env_idx], + successes[env_idx], + bool(terminated[env_idx] | truncated[env_idx]), + task_desc, + recording_datasets[env_idx].features, + ) + recording_datasets[env_idx].add_frame(frame) + if terminated[env_idx] or truncated[env_idx]: + recording_datasets[env_idx].save_episode() + raw_observation = deepcopy(observation) + + # Keep track of which environments are done so far. + # Mark the episode as done if we reach the maximum step limit. + # This ensures that the rollout always terminates cleanly at `max_steps`, + # and allows logging/saving (e.g., videos) to be triggered consistently. + done = terminated | truncated | done + if step + 1 == max_steps: + done = np.ones_like(done, dtype=bool) + + all_actions.append(torch.from_numpy(action_numpy)) + all_rewards.append(torch.from_numpy(reward)) + all_dones.append(torch.from_numpy(done)) + all_successes.append(torch.tensor(successes)) + + step += 1 + running_success_rate = ( + einops.reduce(torch.stack(all_successes, dim=1), "b n -> b", "any").numpy().mean() ) - else: - successes = [False] * env.num_envs - - # Keep track of which environments are done so far. - # Mark the episode as done if we reach the maximum step limit. - # This ensures that the rollout always terminates cleanly at `max_steps`, - # and allows logging/saving (e.g., videos) to be triggered consistently. - done = terminated | truncated | done - if step + 1 == max_steps: - done = np.ones_like(done, dtype=bool) - - all_actions.append(torch.from_numpy(action_numpy)) - all_rewards.append(torch.from_numpy(reward)) - all_dones.append(torch.from_numpy(done)) - all_successes.append(torch.tensor(successes)) - - step += 1 - running_success_rate = ( - einops.reduce(torch.stack(all_successes, dim=1), "b n -> b", "any").numpy().mean() - ) - progbar.set_postfix({"running_success_rate": f"{running_success_rate.item() * 100:.1f}%"}) - progbar.update() + progbar.set_postfix({"running_success_rate": f"{running_success_rate.item() * 100:.1f}%"}) + progbar.update() + finally: + if recording_datasets is not None: + for ds in recording_datasets: + ds.finalize() + if recording_repo_id is not None: + if ds.num_episodes > 0: + ds.push_to_hub(private=recording_private) + else: + logging.warning("No episodes recorded for %s — skipping push to hub.", ds.repo_id) # Track the final observation. if return_observations: @@ -273,6 +396,10 @@ def eval_policy( videos_dir: Path | None = None, return_episode_data: bool = False, start_seed: int | None = None, + recording_dir: Path | None = None, + env_features: dict | None = None, + recording_repo_id: str | None = None, + recording_private: bool = False, ) -> dict: """ Args: @@ -361,6 +488,10 @@ def eval_policy( seeds=list(seeds) if seeds else None, return_observations=return_episode_data, render_callback=render_frame if max_episodes_rendered > 0 else None, + recording_dir=recording_dir, + env_features=env_features, + recording_repo_id=recording_repo_id, + recording_private=recording_private, ) # Figure out where in each rollout sequence the first done condition was encountered (results after @@ -563,6 +694,10 @@ def eval_main(cfg: EvalPipelineConfig): # Create environment-specific preprocessor and postprocessor (e.g., for LIBERO environments) env_preprocessor, env_postprocessor = make_env_pre_post_processors(env_cfg=cfg.env, policy_cfg=cfg.policy) + recording_dir = Path(cfg.output_dir) / "recordings" if cfg.eval.recording else None + max_episodes_rendered = 0 if cfg.eval.recording else 10 + videos_dir = None if cfg.eval.recording else Path(cfg.output_dir) / "videos" + with torch.no_grad(), torch.autocast(device_type=device.type) if cfg.policy.use_amp else nullcontext(): info = eval_policy_all( envs=envs, @@ -572,10 +707,15 @@ def eval_main(cfg: EvalPipelineConfig): preprocessor=preprocessor, postprocessor=postprocessor, n_episodes=cfg.eval.n_episodes, - max_episodes_rendered=10, - videos_dir=Path(cfg.output_dir) / "videos", + max_episodes_rendered=max_episodes_rendered, + videos_dir=videos_dir, + return_episode_data=False, start_seed=cfg.seed, max_parallel_tasks=cfg.env.max_parallel_tasks, + recording_dir=recording_dir, + env_features=cfg.env.features if cfg.eval.recording else None, + recording_repo_id=cfg.eval.recording_repo_id, + recording_private=cfg.eval.recording_private, ) print("Overall Aggregated Metrics:") print(info["overall"]) @@ -618,6 +758,10 @@ def eval_one( videos_dir: Path | None, return_episode_data: bool, start_seed: int | None, + recording_dir: Path | None = None, + env_features: dict | None = None, + recording_repo_id: str | None = None, + recording_private: bool = False, ) -> TaskMetrics: """Evaluates one task_id of one suite using the provided vec env.""" @@ -635,6 +779,10 @@ def eval_one( videos_dir=task_videos_dir, return_episode_data=return_episode_data, start_seed=start_seed, + recording_dir=recording_dir, + env_features=env_features, + recording_repo_id=recording_repo_id, + recording_private=recording_private, ) per_episode = task_result["per_episode"] @@ -661,6 +809,10 @@ def run_one( videos_dir: Path | None, return_episode_data: bool, start_seed: int | None, + recording_dir: Path | None = None, + env_features: dict | None = None, + recording_repo_id: str | None = None, + recording_private: bool = False, ): """ Run eval_one for a single (task_group, task_id, env). @@ -672,7 +824,13 @@ def run_one( task_videos_dir = videos_dir / f"{task_group}_{task_id}" task_videos_dir.mkdir(parents=True, exist_ok=True) - # Call the existing eval_one (assumed to return TaskMetrics-like dict) + task_recording_dir = None + task_repo_id = None + if recording_dir is not None and env_features is not None: + task_recording_dir = recording_dir / f"{task_group}_{task_id}" + if recording_repo_id is not None: + task_repo_id = f"{recording_repo_id}_{task_group}_{task_id}" + metrics = eval_one( env, policy=policy, @@ -685,8 +843,12 @@ def run_one( videos_dir=task_videos_dir, return_episode_data=return_episode_data, start_seed=start_seed, + recording_dir=task_recording_dir, + env_features=env_features, + recording_repo_id=task_repo_id, + recording_private=recording_private, ) - # ensure we always provide video_paths key to simplify accumulation + if max_episodes_rendered > 0: metrics.setdefault("video_paths", []) return task_group, task_id, metrics @@ -702,6 +864,10 @@ def eval_policy_all( n_episodes: int, *, max_episodes_rendered: int = 0, + recording_dir: Path | None = None, + env_features: dict | None = None, + recording_repo_id: str | None = None, + recording_private: bool = False, videos_dir: Path | None = None, return_episode_data: bool = False, start_seed: int | None = None, @@ -761,6 +927,10 @@ def eval_policy_all( videos_dir=videos_dir, return_episode_data=return_episode_data, start_seed=start_seed, + recording_dir=recording_dir, + env_features=env_features, + recording_repo_id=recording_repo_id, + recording_private=recording_private, ) if max_parallel_tasks <= 1: From 79d4976ae23bbc959232e495be5849d9dfd0d71b Mon Sep 17 00:00:00 2001 From: Jiwen Cai Date: Wed, 24 Jun 2026 17:23:25 +0800 Subject: [PATCH 11/38] fix(deps): pin cmeel-urdfdom <5 and cmeel-tinyxml2 <11 in placo-dep (#3873) placo pulls in pin (Pinocchio), whose binary wheels dlopen specific cmeel sonames (liburdfdom_sensor.so.4.0, libtinyxml2.so.10) but declare only `>=` floors on their cmeel packages. The 2026-05-21 major bumps (cmeel-urdfdom 6.0.0 -> .so.6, cmeel-tinyxml2 11.0.0 -> .so.11) ship newer sonames, so left unpinned the resolver grabs them and `import placo` fails at load with "liburdfdom_sensor.so.4.0: cannot open shared object file". #3647 capped placo and hardened the kinematics import, but the guard only defers the failure: constructing RobotKinematics still raises. Pin the cmeel packages to the 4.x / 10.x ABI the placo/pin wheels are built against (there is no cmeel-urdfdom 5.x; <5 selects 4.x). Regenerated uv.lock with uv 0.8.0 to match CI; the only resolution change is the two cmeel versions (plus a deterministic decord platform-marker cascade from 4.0.1's wider wheel set). Fixes #3755 --- pyproject.toml | 9 ++++++++- uv.lock | 48 +++++++++++++++++++++++++++++++++++------------- 2 files changed, 43 insertions(+), 14 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 0dc86d7ff..3961ded19 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -140,7 +140,14 @@ av-dep = ["av>=15.0.0,<16.0.0"] pygame-dep = ["pygame>=2.5.1,<2.7.0"] # NOTE: 0.9.16 links against liburdfdom_sensor.so.4, which is unavailable on Ubuntu 24.04 # (noble ships urdfdom 3.x). Cap below 0.9.16 until system urdfdom 4.x is broadly available. -placo-dep = ["placo>=0.9.6,<0.9.16"] +# +# NOTE: placo pulls in pin (Pinocchio), whose binary wheels dlopen specific cmeel sonames +# (liburdfdom_sensor.so.4.0, libtinyxml2.so.10) but declare only `>=` floors on their cmeel +# packages. The 2026-05-21 major bumps (cmeel-urdfdom 6.0.0 -> .so.6, cmeel-tinyxml2 11.0.0 +# -> .so.11) ship newer sonames, so left unpinned the resolver grabs them and `import placo` +# fails at load with "liburdfdom_sensor.so.4.0: cannot open shared object file" (see #3755). +# There is no cmeel-urdfdom 5.x; <5 selects the 4.x ABI the placo/pin wheels are built against. +placo-dep = ["placo>=0.9.6,<0.9.16", "cmeel-urdfdom>=4,<5", "cmeel-tinyxml2<11"] transformers-dep = ["transformers>=5.4.0,<5.6.0"] grpcio-dep = ["grpcio>=1.73.1,<2.0.0", "protobuf>=6.31.1,<8.0.0"] accelerate-dep = ["accelerate>=1.14.0,<2.0.0"] diff --git a/uv.lock b/uv.lock index 5c55f8a3d..e24b1d884 100644 --- a/uv.lock +++ b/uv.lock @@ -768,34 +768,46 @@ wheels = [ [[package]] name = "cmeel-tinyxml2" -version = "11.0.0" +version = "10.0.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "cmeel" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/9b/96/4311533fee0a364bb605b585762f04c249f47857b33548a8ea837a7eb860/cmeel_tinyxml2-11.0.0.tar.gz", hash = "sha256:85d9c7680b3369af4c6b40a0dce70bbd84aa67832755622e57eb260cd95abe40", size = 645900, upload-time = "2026-05-21T11:49:32.652Z" } +sdist = { url = "https://files.pythonhosted.org/packages/28/9f/030eca702c485f7a641f975f167fa93164911b3329f005fb0730ff5e793f/cmeel_tinyxml2-10.0.0.tar.gz", hash = "sha256:00252aefc1c94a55b89f25ad08ee79fda2da8d1d94703e051598ddb52a9088fe", size = 645297, upload-time = "2025-02-06T10:29:00.106Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/79/f0/90c1640c53b623359d75ab1c70bdf19dc0afe82722bc5df57d09f8eaf83a/cmeel_tinyxml2-11.0.0-0-py3-none-macosx_10_9_x86_64.whl", hash = "sha256:b0bd974e549b8c444626671a8e645897603ebf5225734cbe04a9dd3461477754", size = 111719, upload-time = "2026-05-21T11:49:25.999Z" }, - { url = "https://files.pythonhosted.org/packages/56/40/166447150a31bc3b794ffb493d5a634f67ffbc75dd8b4c46373701b7ef15/cmeel_tinyxml2-11.0.0-0-py3-none-macosx_11_0_arm64.whl", hash = "sha256:9a1406f408262c37ae7c4566b1d67801c4b10c4980903fb1ef0ba45fa4407072", size = 109146, upload-time = "2026-05-21T11:49:27.829Z" }, - { url = "https://files.pythonhosted.org/packages/d8/ca/3cc665afe2d76999f15454bb3b2f7c05f0088ad7de35718648291a536fd9/cmeel_tinyxml2-11.0.0-0-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:6f830007917c3e36f26b27d170ce84a619a62f46104d3cce435dff0125dd665f", size = 157109, upload-time = "2026-05-21T11:49:29.358Z" }, - { url = "https://files.pythonhosted.org/packages/87/4e/dcc0d9756d93be734d824e2a570cc9ac68909a1d7d3b6fc87c2fb32726c0/cmeel_tinyxml2-11.0.0-0-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:18674156bd41f3993dc1d5199da04fa496674358daa6588090fb9f86c71917b0", size = 148825, upload-time = "2026-05-21T11:49:31.035Z" }, + { url = "https://files.pythonhosted.org/packages/fe/5d/bc3a932eb7996a0a789979426a9bb8a3948bf57f3f17bab87dddbef62433/cmeel_tinyxml2-10.0.0-0-py3-none-macosx_10_9_x86_64.whl", hash = "sha256:924499bb1b60b9a17bd001d12a9af88ddbee4ca888638ae684ba7f0f3ce49e87", size = 111913, upload-time = "2025-02-06T10:28:45.723Z" }, + { url = "https://files.pythonhosted.org/packages/92/bf/67d11e123313c034712896e94038291fe506bb099bdb75a136392002ffd0/cmeel_tinyxml2-10.0.0-0-py3-none-macosx_11_0_arm64.whl", hash = "sha256:26a1eb30c2a00bfc172e89ed015a18b8efb2b383546252ca8859574aed684686", size = 109487, upload-time = "2025-02-06T10:28:47.546Z" }, + { url = "https://files.pythonhosted.org/packages/ca/48/d8c81ce19b4b278ed0e8f81f93ae8670209bf3a9ac20141b9c386bb40cc7/cmeel_tinyxml2-10.0.0-0-py3-none-manylinux_2_17_i686.whl", hash = "sha256:53d86e02864c712f51f9a9adfcd8b6046b2ed51d44a0c34a8438d93b72b48325", size = 160118, upload-time = "2025-02-06T10:28:49.627Z" }, + { url = "https://files.pythonhosted.org/packages/87/4e/62193e27c9581f8ba7aeaeca7805632a64f2f4a824b1db37ad02ee953e8a/cmeel_tinyxml2-10.0.0-0-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:74112e2e9473afbf6ee2d25c9942553e9f6a40465e714533db72db48bc7658e1", size = 158477, upload-time = "2025-02-06T10:28:51.667Z" }, + { url = "https://files.pythonhosted.org/packages/14/f9/d0420c39e9ade99beeec61cd3abc68880fe6e14d85e9df292af8fabe65c8/cmeel_tinyxml2-10.0.0-0-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:ecd6e99caa2a06ac0d4b333b740c20fca526d0ca426f99eb5c0a0039117afdb6", size = 147025, upload-time = "2025-02-06T10:28:53.944Z" }, + { url = "https://files.pythonhosted.org/packages/66/9e/df63147fc162ab487217fa5596778ab7a81a82d9b3ce4236fd3a1e48cecb/cmeel_tinyxml2-10.0.0-0-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:30993fffb7032a45d5d3b1e5670cb879dad667a13144cd68c8f4e0371a8a3d2e", size = 150958, upload-time = "2025-02-06T10:28:55.301Z" }, + { url = "https://files.pythonhosted.org/packages/0e/a8/b03567275fd83f5af33ddb61de942689dec72c5b21bec01e6a5b11101aa5/cmeel_tinyxml2-10.0.0-0-py3-none-musllinux_1_2_i686.whl", hash = "sha256:8c09ede51784af54211a6225884dc7ddbb02ea1681656d173060c7ad2a5b9a3c", size = 160300, upload-time = "2025-02-06T10:28:57.189Z" }, + { url = "https://files.pythonhosted.org/packages/3a/ec/2781635b66c1059ca1243ae0f5a0410e171a5d8b8a71be3e34cb172f9f2d/cmeel_tinyxml2-10.0.0-0-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:3bd511d6d0758224efdebc23d3ead6e94f0755b04141ebf7d5493377829e8332", size = 149184, upload-time = "2025-02-06T10:28:58.734Z" }, ] [[package]] name = "cmeel-urdfdom" -version = "6.0.0" +version = "4.0.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "cmeel" }, { name = "cmeel-console-bridge" }, { name = "cmeel-tinyxml2" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/21/75/4e8aff079e98582aeeb8e752805081da0c2dea405e79bafeefb555defe9f/cmeel_urdfdom-6.0.0.tar.gz", hash = "sha256:65c0fdc6021300fc55b2d0c03ab64dedc328034a74e40498e671bc894bb1dcf7", size = 303688, upload-time = "2026-05-21T12:08:56.663Z" } +sdist = { url = "https://files.pythonhosted.org/packages/31/09/be81a5e7db56f34b6ccdbe7afe855c95a18c8439e173519e0146e9276a8c/cmeel_urdfdom-4.0.1.tar.gz", hash = "sha256:2e3f41e8483889e195b574acb326a4464cf11a3c0a8724031ac28bcda2223efc", size = 291511, upload-time = "2025-02-12T12:07:09.699Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/f2/40/51ba667135f01631179eee1614557193f8453740f248302d1b8b7f9f693e/cmeel_urdfdom-6.0.0-0-py3-none-macosx_10_9_x86_64.whl", hash = "sha256:53d55cebb137a6e4dac6c16fa53f2dc2b7b9b5cda644bd1637a5bb849cd96e52", size = 381501, upload-time = "2026-05-21T12:08:48.758Z" }, - { url = "https://files.pythonhosted.org/packages/3c/d1/2b49a8c940fa75abc13df9842c14e577e6a82d5854b6d52597ce3bb04894/cmeel_urdfdom-6.0.0-0-py3-none-macosx_11_0_arm64.whl", hash = "sha256:0ef424735bd30f4afa4d1b4ddca9b297498c43005ddd775c080e55f62e9e0466", size = 377159, upload-time = "2026-05-21T12:08:50.485Z" }, - { url = "https://files.pythonhosted.org/packages/db/ac/0efde3a48220b55707bafb6d2e2dcca562f99dcd5c2c15311f7696eeacce/cmeel_urdfdom-6.0.0-0-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:0436f5230f1484c8e583284ef48c7291b230ada3dc5fb2937941f582e72409ec", size = 506000, upload-time = "2026-05-21T12:08:52.273Z" }, - { url = "https://files.pythonhosted.org/packages/32/d4/dfd617e598100e4e53ae3d228a968facff80bae53038fb18e2dccb1ab03a/cmeel_urdfdom-6.0.0-0-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:7ab1be680a8ec866d5422c617b641d1f0e38774061df28b8b426fb26edce6337", size = 530049, upload-time = "2026-05-21T12:08:54.224Z" }, + { url = "https://files.pythonhosted.org/packages/be/d0/20147dd6bb723afc44a58d89ea624df2bad1bed7b898a2df112aaca4a479/cmeel_urdfdom-4.0.1-0-py3-none-macosx_10_9_x86_64.whl", hash = "sha256:2fe56939c6b47f6ec57021aac154123da47ecdcd79a217f3a5e3c4b705a07dee", size = 300860, upload-time = "2025-02-12T12:06:58.536Z" }, + { url = "https://files.pythonhosted.org/packages/8e/98/f832bca347e2d987c6b0ebb6930caf7b2c402535324aeed466b6aa2c4513/cmeel_urdfdom-4.0.1-0-py3-none-macosx_11_0_arm64.whl", hash = "sha256:00a0aba78b68c428b27abeed1db58d73e65319ed966911a0e97b37367442e756", size = 300616, upload-time = "2025-02-12T12:07:00.556Z" }, + { url = "https://files.pythonhosted.org/packages/cf/10/bf5765b6f388037cff166a754a0958ac2fee34ca3c0975ef64d0324e4647/cmeel_urdfdom-4.0.1-0-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:a701a8f9671331f11b18ecf37a6537db546a21e6a0e5d0ff53341fea0693ed7f", size = 385951, upload-time = "2025-02-12T12:07:02.556Z" }, + { url = "https://files.pythonhosted.org/packages/c3/82/cb3f8f587d293a17bdbea15b50cdaa4a1e28e04583eb4cb4821685b89466/cmeel_urdfdom-4.0.1-0-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:12e39fc388c077d79fc9b3841d3d972a1da90b90de754d3363194c1540e18abf", size = 399619, upload-time = "2025-02-12T12:07:04.388Z" }, + { url = "https://files.pythonhosted.org/packages/24/77/322d7ac92c692d8dfaeda9de2d937087d15e2b564dc457d656e5fde3991d/cmeel_urdfdom-4.0.1-0-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:c4a83925df1d5923c4485c3eb2b80b3a61b14f119ab724fb5bd04cec494690ee", size = 373969, upload-time = "2025-02-12T12:07:06.222Z" }, + { url = "https://files.pythonhosted.org/packages/9f/63/bdc6b55cc8bd99bb9dce6be801b30feffaa1c3841ecb7f4fe4d137424518/cmeel_urdfdom-4.0.1-0-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:4c4f44270971b3d05c45a4e21b1fb2df7e05a750363ae918f59532bff0bfe0e1", size = 388237, upload-time = "2025-02-12T12:07:08.326Z" }, + { url = "https://files.pythonhosted.org/packages/1d/2d/8463fc23230612daf4da1e31d3229f47708381f3ae4d1500f0f007ac0f92/cmeel_urdfdom-4.0.1-1-py3-none-macosx_10_9_x86_64.whl", hash = "sha256:f7535158f45992eb2ba79e90d9db1bf9adc3846d9c7ed3e7a8c1c4d5343afa37", size = 301006, upload-time = "2025-02-13T11:42:08.8Z" }, + { url = "https://files.pythonhosted.org/packages/0f/d5/c8cdf500e49300d85624cbc3ef804107ddcdc9c541b1d3f726bfb58a9fc1/cmeel_urdfdom-4.0.1-1-py3-none-macosx_11_0_arm64.whl", hash = "sha256:fef2a01a00d61d41b3d35dd4958bba973e9025c26eea1d3c9880932f4dba89a5", size = 300758, upload-time = "2025-02-13T11:42:10.449Z" }, + { url = "https://files.pythonhosted.org/packages/cf/b3/2f7bac1544113a7f8e0f6d8b1fab5e75c6a3d27ffbb584b03267251b2165/cmeel_urdfdom-4.0.1-1-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:7a52eb36950ce982014d99a55717ca29985da056e3705f20746f15d3244c1f7a", size = 386043, upload-time = "2025-02-13T11:42:11.923Z" }, + { url = "https://files.pythonhosted.org/packages/86/03/8bdeb36ba6a3e8125d523ecfc010403049e463fe589f9896858d4bdcaf1e/cmeel_urdfdom-4.0.1-1-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:9f3b9c80b10d7246821ff61c2573f799e3da23d483e6f7367ddcad8a48baf58f", size = 399719, upload-time = "2025-02-13T11:42:14.325Z" }, + { url = "https://files.pythonhosted.org/packages/3f/ed/43f99e7512460294cd8acc5753ba25f8a20bdf28d62e143eaf3ec7a28bb6/cmeel_urdfdom-4.0.1-1-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:2de69f47e8312cc09157624802d5bdaad6406443f863fb4b9ec62a19b4de3c72", size = 374073, upload-time = "2025-02-13T11:42:17.907Z" }, + { url = "https://files.pythonhosted.org/packages/17/c6/2e9bde6d7c02c1cf203ea896f8ce1afd441412f09b44830f1ee4a96d77de/cmeel_urdfdom-4.0.1-1-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:7708c1402de450fbeab21f7ca264a9a4676ed4c1cdf8d84d840bc5d057aac920", size = 388337, upload-time = "2025-02-13T11:42:19.657Z" }, ] [[package]] @@ -1147,7 +1159,7 @@ name = "decord" version = "0.6.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "numpy", marker = "(platform_machine != 'arm64' and sys_platform == 'darwin') or (platform_machine == 'AMD64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'linux')" }, + { name = "numpy", marker = "(platform_machine != 'arm64' and platform_machine != 's390x' and sys_platform == 'darwin') or (platform_machine == 'AMD64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine != 's390x' and sys_platform != 'darwin' and sys_platform != 'linux')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/11/79/936af42edf90a7bd4e41a6cac89c913d4b47fa48a26b042d5129a9242ee3/decord-0.6.0-py3-none-manylinux2010_x86_64.whl", hash = "sha256:51997f20be8958e23b7c4061ba45d0efcd86bffd5fe81c695d0befee0d442976", size = 13602299, upload-time = "2021-06-14T21:30:55.486Z" }, @@ -2788,6 +2800,8 @@ accelerate-dep = [ all = [ { name = "accelerate" }, { name = "av" }, + { name = "cmeel-tinyxml2" }, + { name = "cmeel-urdfdom" }, { name = "contourpy" }, { name = "datasets" }, { name = "debugpy" }, @@ -2971,6 +2985,8 @@ hardware = [ ] hilserl = [ { name = "av" }, + { name = "cmeel-tinyxml2" }, + { name = "cmeel-urdfdom" }, { name = "datasets" }, { name = "grpcio" }, { name = "gym-hil" }, @@ -2993,6 +3009,8 @@ intelrealsense = [ { name = "pyrealsense2-macosx", marker = "sys_platform == 'darwin'" }, ] kinematics = [ + { name = "cmeel-tinyxml2" }, + { name = "cmeel-urdfdom" }, { name = "placo" }, ] lekiwi = [ @@ -3066,6 +3084,8 @@ pi = [ { name = "transformers" }, ] placo-dep = [ + { name = "cmeel-tinyxml2" }, + { name = "cmeel-urdfdom" }, { name = "placo" }, ] pusht = [ @@ -3186,6 +3206,8 @@ requires-dist = [ { name = "accelerate", marker = "extra == 'accelerate-dep'", specifier = ">=1.14.0,<2.0.0" }, { name = "av", marker = "extra == 'av-dep'", specifier = ">=15.0.0,<16.0.0" }, { name = "cmake", specifier = ">=3.29.0.1,<4.2.0" }, + { name = "cmeel-tinyxml2", marker = "extra == 'placo-dep'", specifier = "<11" }, + { name = "cmeel-urdfdom", marker = "extra == 'placo-dep'", specifier = ">=4,<5" }, { name = "contourpy", marker = "extra == 'matplotlib-dep'", specifier = ">=1.3.0,<2.0.0" }, { name = "datasets", marker = "extra == 'dataset'", specifier = ">=4.7.0,<5.0.0" }, { name = "debugpy", marker = "extra == 'dev'", specifier = ">=1.8.1,<1.9.0" }, From 536b9621b2d8b0193be71508140701a18f79c953 Mon Sep 17 00:00:00 2001 From: Alexandre Edmond <145270396+AlexandreEDMOND@users.noreply.github.com> Date: Wed, 24 Jun 2026 11:44:03 +0200 Subject: [PATCH 12/38] Fix pi0fast model id in docs (#3855) --- docs/source/pi0fast.mdx | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/source/pi0fast.mdx b/docs/source/pi0fast.mdx index f7272acc5..15dff8071 100644 --- a/docs/source/pi0fast.mdx +++ b/docs/source/pi0fast.mdx @@ -96,7 +96,7 @@ lerobot-train \ --policy.type=pi0_fast \ --output_dir=./outputs/pi0fast_training \ --job_name=pi0fast_training \ - --policy.pretrained_path=lerobot/pi0_fast_base \ + --policy.pretrained_path=lerobot/pi0fast-base \ --policy.dtype=bfloat16 \ --policy.gradient_checkpointing=true \ --policy.chunk_size=10 \ @@ -187,7 +187,7 @@ lerobot-train \ --dataset.repo_id=lerobot/libero \ --output_dir=outputs/libero_pi0fast \ --job_name=libero_pi0fast \ - --policy.path=lerobot/pi0fast_base \ + --policy.path=lerobot/pi0fast-base \ --policy.dtype=bfloat16 \ --steps=100000 \ --save_freq=20000 \ From 508d18f8a1e62c17380510323878fc72b91e54ec Mon Sep 17 00:00:00 2001 From: someone114514 <130830487+someone114514@users.noreply.github.com> Date: Wed, 24 Jun 2026 23:59:07 -0700 Subject: [PATCH 13/38] Fix ACT policy type examples in docs (#3792) --- README.md | 2 +- docs/source/multi_gpu_training.mdx | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 2a330d823..f72952102 100644 --- a/README.md +++ b/README.md @@ -97,7 +97,7 @@ Training a policy is as simple as running a script configuration: ```bash lerobot-train \ - --policy=act \ + --policy.type=act \ --dataset.repo_id=lerobot/aloha_mobile_cabinet ``` diff --git a/docs/source/multi_gpu_training.mdx b/docs/source/multi_gpu_training.mdx index 7907340c3..7c212364e 100644 --- a/docs/source/multi_gpu_training.mdx +++ b/docs/source/multi_gpu_training.mdx @@ -95,7 +95,7 @@ If you want to scale your hyperparameters when using multiple GPUs, you should d accelerate launch --num_processes=2 $(which lerobot-train) \ --optimizer.lr=2e-4 \ --dataset.repo_id=lerobot/pusht \ - --policy=act + --policy.type=act ``` **Training Steps Scaling:** @@ -110,7 +110,7 @@ accelerate launch --num_processes=2 $(which lerobot-train) \ --batch_size=8 \ --steps=50000 \ --dataset.repo_id=lerobot/pusht \ - --policy=act + --policy.type=act ``` ## Training Large Models with FSDP From b4e454c0ff3637d918e7aab41ee4b4096d79cd90 Mon Sep 17 00:00:00 2001 From: Steven Palma Date: Thu, 25 Jun 2026 10:58:39 +0200 Subject: [PATCH 14/38] feat(utils): display-independent keyboard controls for recording (Wayland / headless / macOS) (#3875) * feat(utils): headless keyboard control * refactor(utils): consolidate keyboard listener creation * fix(rollout): remove import require guard for pynput --------- Co-authored-by: Leo Toff Co-authored-by: Stefano Maestri Co-authored-by: Sahil Chande <85823961+SahilChande@users.noreply.github.com> Co-authored-by: Vinayak Agarwal <63502278+Vinayak-Agarwal-2004@users.noreply.github.com> Co-authored-by: Abdul Rahim Mirani --- docs/source/il_robots.mdx | 16 +- docs/source/lekiwi.mdx | 2 +- examples/lekiwi/evaluate.py | 3 +- examples/lekiwi/record.py | 2 +- examples/phone_to_so100/evaluate.py | 3 +- examples/phone_to_so100/record.py | 2 +- examples/so100_to_so100_EE/evaluate.py | 3 +- examples/so100_to_so100_EE/record.py | 2 +- src/lerobot/common/control_utils.py | 84 ---- src/lerobot/rollout/strategies/dagger.py | 90 +--- src/lerobot/rollout/strategies/episodic.py | 5 +- src/lerobot/rollout/strategies/highlight.py | 58 +-- src/lerobot/scripts/lerobot_record.py | 9 +- .../teleoperators/gamepad/gamepad_utils.py | 10 + .../teleoperators/keyboard/teleop_keyboard.py | 21 +- src/lerobot/utils/keyboard_input.py | 440 ++++++++++++++++++ tests/utils/test_keyboard_input.py | 228 +++++++++ 17 files changed, 758 insertions(+), 220 deletions(-) create mode 100644 src/lerobot/utils/keyboard_input.py create mode 100644 tests/utils/test_keyboard_input.py diff --git a/docs/source/il_robots.mdx b/docs/source/il_robots.mdx index 53ae5af82..6a820e0db 100644 --- a/docs/source/il_robots.mdx +++ b/docs/source/il_robots.mdx @@ -390,9 +390,17 @@ Set the flow of data recording using command-line arguments: Control the data recording flow using keyboard shortcuts: -- Press **Right Arrow (`→`)**: Early stop the current episode or reset time and move to the next. -- Press **Left Arrow (`←`)**: Cancel the current episode and re-record it. -- Press **Escape (`ESC`)**: Immediately stop the session, encode videos, and upload the dataset. +- Press **Right Arrow (`→`)** or **`n`**: Early stop the current episode or reset time and move to the next. +- Press **Left Arrow (`←`)** or **`r`**: Cancel the current episode and re-record it. +- Press **Escape (`ESC`)** or **`q`**: Immediately stop the session, encode videos, and upload the dataset. + + + +These control-flow shortcuts work on **X11, Wayland, and headless/SSH** sessions. When a global keyboard backend isn't available (Wayland, a headless machine, or macOS without Accessibility permission), `lerobot-record` automatically reads the same keys from the terminal — launch it from an interactive terminal and keep it focused. You can also use the letter equivalents **`n`** (next, same as `→`), **`r`** (re-record, same as `←`) and **`q`** (quit, same as `ESC`). No `$DISPLAY` setup is required. + +This applies to the recording control flow only. Keyboard **teleoperation** (driving the robot with the keyboard) still needs a global key backend, so it works only on an X11 session, a Windows desktop, or macOS with Accessibility/Input Monitoring granted — not on Wayland or headless sessions. + + #### Tips for gathering data @@ -406,7 +414,7 @@ If you want to dive deeper into this important topic, you can check out the [blo #### Troubleshooting: -- On Linux, if the left and right arrow keys and escape key don't have any effect during data recording, make sure you've set the `$DISPLAY` environment variable. See [pynput limitations](https://pynput.readthedocs.io/en/latest/limitations.html#linux). +- On Linux, the recording control-flow keys (arrow keys, Escape) work on X11, Wayland, and headless/SSH sessions as long as `lerobot-record` runs in an interactive terminal — no `$DISPLAY` setup is needed. If the keys have no effect, make sure you are in an interactive (TTY) terminal, not a piped/non-TTY session, and that it is focused; the letter equivalents `n` / `r` / `q` also work. Keyboard _teleoperation_ (as opposed to the recording control flow) still requires a global key backend — an X11 session, a Windows desktop, or macOS with Accessibility/Input Monitoring granted — and is unavailable on Wayland or headless machines. See [pynput limitations](https://pynput.readthedocs.io/en/latest/limitations.html#linux). ## Visualize a dataset diff --git a/docs/source/lekiwi.mdx b/docs/source/lekiwi.mdx index 7e7c1a680..739073b65 100644 --- a/docs/source/lekiwi.mdx +++ b/docs/source/lekiwi.mdx @@ -319,7 +319,7 @@ If you want to dive deeper into this important topic, you can check out the [blo #### Troubleshooting: -- On Linux, if the left and right arrow keys and escape key don't have any effect during data recording, make sure you've set the `$DISPLAY` environment variable. See [pynput limitations](https://pynput.readthedocs.io/en/latest/limitations.html#linux). +- On Linux, the recording control-flow keys (arrow keys, Escape) work on X11, Wayland, and headless/SSH sessions as long as you run the recording from an interactive terminal (keep it focused) — no `$DISPLAY` setup is needed; the letter equivalents `n` / `r` / `q` also work. Note that **keyboard teleoperation of the LeKiwi base** is different: it relies on a global key backend and therefore works only on an X11 session, a Windows desktop, or macOS with Accessibility/Input Monitoring granted — not on Wayland or headless machines. See [pynput limitations](https://pynput.readthedocs.io/en/latest/limitations.html#linux). ## Replay an episode diff --git a/examples/lekiwi/evaluate.py b/examples/lekiwi/evaluate.py index 3ddcd1f14..13bb6ac28 100644 --- a/examples/lekiwi/evaluate.py +++ b/examples/lekiwi/evaluate.py @@ -17,7 +17,7 @@ import logging import time -from lerobot.common.control_utils import init_keyboard_listener, predict_action +from lerobot.common.control_utils import predict_action from lerobot.datasets import LeRobotDataset from lerobot.policies import make_pre_post_processors from lerobot.policies.act import ACTPolicy @@ -26,6 +26,7 @@ from lerobot.processor import make_default_processors from lerobot.robots.lekiwi import LeKiwiClient, LeKiwiClientConfig from lerobot.utils.constants import ACTION, OBS_STR from lerobot.utils.feature_utils import build_dataset_frame, hw_to_dataset_features +from lerobot.utils.keyboard_input import init_keyboard_listener from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.utils import log_say from lerobot.utils.visualization_utils import init_rerun, log_rerun_data diff --git a/examples/lekiwi/record.py b/examples/lekiwi/record.py index 2c581f5ff..f62a9eb49 100644 --- a/examples/lekiwi/record.py +++ b/examples/lekiwi/record.py @@ -14,7 +14,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from lerobot.common.control_utils import init_keyboard_listener from lerobot.datasets import LeRobotDataset from lerobot.processor import make_default_processors from lerobot.robots.lekiwi import LeKiwiClient, LeKiwiClientConfig @@ -23,6 +22,7 @@ from lerobot.teleoperators.keyboard import KeyboardTeleop, KeyboardTeleopConfig from lerobot.teleoperators.so_leader import SO100Leader, SO100LeaderConfig from lerobot.utils.constants import ACTION, OBS_STR from lerobot.utils.feature_utils import hw_to_dataset_features +from lerobot.utils.keyboard_input import init_keyboard_listener from lerobot.utils.utils import log_say from lerobot.utils.visualization_utils import init_rerun diff --git a/examples/phone_to_so100/evaluate.py b/examples/phone_to_so100/evaluate.py index e859123d0..d1fb4de67 100644 --- a/examples/phone_to_so100/evaluate.py +++ b/examples/phone_to_so100/evaluate.py @@ -18,7 +18,7 @@ import logging import time from lerobot.cameras.opencv import OpenCVCameraConfig -from lerobot.common.control_utils import init_keyboard_listener, predict_action +from lerobot.common.control_utils import predict_action from lerobot.configs import FeatureType, PolicyFeature from lerobot.datasets import LeRobotDataset, aggregate_pipeline_dataset_features, create_initial_features from lerobot.model.kinematics import RobotKinematics @@ -41,6 +41,7 @@ from lerobot.robots.so_follower.robot_kinematic_processor import ( from lerobot.types import RobotAction, RobotObservation from lerobot.utils.constants import ACTION, OBS_STR from lerobot.utils.feature_utils import build_dataset_frame, combine_feature_dicts +from lerobot.utils.keyboard_input import init_keyboard_listener from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.utils import log_say from lerobot.utils.visualization_utils import init_rerun, log_rerun_data diff --git a/examples/phone_to_so100/record.py b/examples/phone_to_so100/record.py index 87b8e49fd..612e94ab9 100644 --- a/examples/phone_to_so100/record.py +++ b/examples/phone_to_so100/record.py @@ -15,7 +15,6 @@ # limitations under the License. from lerobot.cameras.opencv import OpenCVCameraConfig -from lerobot.common.control_utils import init_keyboard_listener from lerobot.datasets import LeRobotDataset, aggregate_pipeline_dataset_features, create_initial_features from lerobot.model.kinematics import RobotKinematics from lerobot.processor import ( @@ -39,6 +38,7 @@ from lerobot.teleoperators.phone.config_phone import PhoneOS from lerobot.teleoperators.phone.phone_processor import MapPhoneActionToRobotAction from lerobot.types import RobotAction, RobotObservation from lerobot.utils.feature_utils import combine_feature_dicts +from lerobot.utils.keyboard_input import init_keyboard_listener from lerobot.utils.utils import log_say from lerobot.utils.visualization_utils import init_rerun diff --git a/examples/so100_to_so100_EE/evaluate.py b/examples/so100_to_so100_EE/evaluate.py index 63def68d0..2a2022623 100644 --- a/examples/so100_to_so100_EE/evaluate.py +++ b/examples/so100_to_so100_EE/evaluate.py @@ -18,7 +18,7 @@ import logging import time from lerobot.cameras.opencv import OpenCVCameraConfig -from lerobot.common.control_utils import init_keyboard_listener, predict_action +from lerobot.common.control_utils import predict_action from lerobot.configs import FeatureType, PolicyFeature from lerobot.datasets import LeRobotDataset, aggregate_pipeline_dataset_features, create_initial_features from lerobot.model.kinematics import RobotKinematics @@ -41,6 +41,7 @@ from lerobot.robots.so_follower.robot_kinematic_processor import ( from lerobot.types import RobotAction, RobotObservation from lerobot.utils.constants import ACTION, OBS_STR from lerobot.utils.feature_utils import build_dataset_frame, combine_feature_dicts +from lerobot.utils.keyboard_input import init_keyboard_listener from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.utils import log_say from lerobot.utils.visualization_utils import init_rerun, log_rerun_data diff --git a/examples/so100_to_so100_EE/record.py b/examples/so100_to_so100_EE/record.py index a0b92da3b..3706ee4f5 100644 --- a/examples/so100_to_so100_EE/record.py +++ b/examples/so100_to_so100_EE/record.py @@ -16,7 +16,6 @@ from lerobot.cameras.opencv import OpenCVCameraConfig -from lerobot.common.control_utils import init_keyboard_listener from lerobot.datasets import LeRobotDataset, aggregate_pipeline_dataset_features, create_initial_features from lerobot.model.kinematics import RobotKinematics from lerobot.processor import ( @@ -36,6 +35,7 @@ from lerobot.scripts.lerobot_record import record_loop from lerobot.teleoperators.so_leader import SO100Leader, SO100LeaderConfig from lerobot.types import RobotAction, RobotObservation from lerobot.utils.feature_utils import combine_feature_dicts +from lerobot.utils.keyboard_input import init_keyboard_listener from lerobot.utils.utils import log_say from lerobot.utils.visualization_utils import init_rerun diff --git a/src/lerobot/common/control_utils.py b/src/lerobot/common/control_utils.py index ddaf77d26..e3130643d 100644 --- a/src/lerobot/common/control_utils.py +++ b/src/lerobot/common/control_utils.py @@ -17,12 +17,9 @@ from __future__ import annotations ######################################################################################## # Utilities ######################################################################################## -import logging import time -import traceback from contextlib import nullcontext from copy import copy -from functools import cache from typing import TYPE_CHECKING, Any import numpy as np @@ -43,34 +40,6 @@ from lerobot.robots import Robot from lerobot.types import PolicyAction -@cache -def is_headless(): - """ - Detects if the Python script is running in a headless environment (e.g., without a display). - - This function attempts to import `pynput`, a library that requires a graphical environment. - If the import fails, it assumes the environment is headless. The result is cached to avoid - re-running the check. - - Returns: - True if the environment is determined to be headless, False otherwise. - """ - try: - import pynput # noqa - - return False - except Exception: - print( - "Error trying to import pynput. Switching to headless mode. " - "As a result, the video stream from the cameras won't be shown, " - "and you won't be able to change the control flow with keyboards. " - "For more info, see traceback below.\n" - ) - traceback.print_exc() - print() - return True - - def predict_action( observation: dict[str, np.ndarray], policy: PreTrainedPolicy, @@ -122,59 +91,6 @@ def predict_action( return action -def init_keyboard_listener(): - """ - Initializes a non-blocking keyboard listener for real-time user interaction. - - This function sets up a listener for specific keys (right arrow, left arrow, escape) to control - the program flow during execution, such as stopping recording or exiting loops. It gracefully - handles headless environments where keyboard listening is not possible. - - Returns: - A tuple containing: - - The `pynput.keyboard.Listener` instance, or `None` if in a headless environment. - - A dictionary of event flags (e.g., `exit_early`) that are set by key presses. - """ - # Allow to exit early while recording an episode or resetting the environment, - # by tapping the right arrow key '->'. This might require a sudo permission - # to allow your terminal to monitor keyboard events. - events = {} - events["exit_early"] = False - events["rerecord_episode"] = False - events["stop_recording"] = False - - if is_headless(): - logging.warning( - "Headless environment detected. On-screen cameras display and keyboard inputs will not be available." - ) - listener = None - return listener, events - - # Only import pynput if not in a headless environment - from pynput import keyboard - - def on_press(key): - try: - if key == keyboard.Key.right: - print("Right arrow key pressed. Exiting loop...") - events["exit_early"] = True - elif key == keyboard.Key.left: - print("Left arrow key pressed. Exiting loop and rerecord the last episode...") - events["rerecord_episode"] = True - events["exit_early"] = True - elif key == keyboard.Key.esc: - print("Escape key pressed. Stopping data recording...") - events["stop_recording"] = True - events["exit_early"] = True - except Exception as e: - print(f"Error handling key press: {e}") - - listener = keyboard.Listener(on_press=on_press) - listener.start() - - return listener, events - - def sanity_check_dataset_name(repo_id, policy_cfg): """ Validates the dataset repository name against the presence of a policy configuration. diff --git a/src/lerobot/rollout/strategies/dagger.py b/src/lerobot/rollout/strategies/dagger.py index 8791a5502..21d1e8e98 100644 --- a/src/lerobot/rollout/strategies/dagger.py +++ b/src/lerobot/rollout/strategies/dagger.py @@ -47,8 +47,6 @@ from __future__ import annotations import contextlib import enum import logging -import os -import sys import time from concurrent.futures import Future, ThreadPoolExecutor from threading import Event, Lock @@ -58,7 +56,6 @@ import numpy as np from lerobot.common.control_utils import ( follower_smooth_move_to, - is_headless, teleop_smooth_move_to, teleop_supports_feedback, ) @@ -66,7 +63,7 @@ from lerobot.datasets import VideoEncodingManager from lerobot.datasets.utils import DEFAULT_VIDEO_FILE_SIZE_IN_MB from lerobot.utils.constants import ACTION, OBS_STR from lerobot.utils.feature_utils import build_dataset_frame -from lerobot.utils.import_utils import _pynput_available +from lerobot.utils.keyboard_input import create_key_listener from lerobot.utils.pedal import start_pedal_listener from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.utils import log_say @@ -75,19 +72,6 @@ from ..configs import DAggerKeyboardConfig, DAggerPedalConfig, DAggerStrategyCon from ..context import RolloutContext from .core import RolloutStrategy, estimate_max_episode_seconds, safe_push_to_hub, send_next_action -PYNPUT_AVAILABLE = _pynput_available -keyboard = None -if PYNPUT_AVAILABLE: - try: - if ("DISPLAY" not in os.environ) and ("linux" in sys.platform): - logging.info("No DISPLAY set. Skipping pynput import.") - PYNPUT_AVAILABLE = False - else: - from pynput import keyboard - except Exception as e: - PYNPUT_AVAILABLE = False - logging.info(f"Could not import pynput: {e}") - logger = logging.getLogger(__name__) @@ -180,64 +164,36 @@ class DAggerEvents: def _init_dagger_keyboard(events: DAggerEvents, cfg: DAggerKeyboardConfig): - """Initialise keyboard listener with DAgger 3-key controls. + """Initialise a keyboard listener for DAgger's 3 controls. - Returns the pynput Listener (or ``None`` in headless mode or when - pynput is unavailable). + Backend selection (pynput on X11 / trusted-macOS / Windows, a terminal reader on + Wayland / headless TTY) is delegated to :func:`create_key_listener`. Returns the + listener (exposing ``stop()``) or ``None`` when no keyboard backend is usable. """ - if not PYNPUT_AVAILABLE or is_headless(): - logger.warning("Headless environment or pynput unavailable — keyboard controls disabled") - return None - - # Map config key names to pynput Key objects for special keys - special_keys = { - "space": keyboard.Key.space, - "tab": keyboard.Key.tab, - "enter": keyboard.Key.enter, - } - - def _resolve_key(key) -> str | None: - """Resolve a pynput key event to a config-comparable string.""" - if key == keyboard.Key.esc: - return "esc" - for name, pynput_key in special_keys.items(): - if key == pynput_key: - return name - if hasattr(key, "char") and key.char: - return key.char - return None - - # Build mapping: resolved key string -> DAgger event name + # Map config key names to DAgger event names. key_to_event = { cfg.pause_resume: "pause_resume", cfg.correction: "correction", } - def on_press(key): - try: - resolved = _resolve_key(key) - if resolved is None: - return - if resolved == "esc": - logger.info("Stop recording...") - events.stop_recording.set() - return - if resolved in key_to_event: - events.request_transition(key_to_event[resolved]) - if resolved == cfg.upload: - events.upload_requested.set() - except Exception as e: - logger.debug("Key error: %s", e) + def dispatch(name: str) -> None: + """Apply a resolved key name to the DAgger events.""" + if name == "esc": + logger.info("Stop recording...") + events.stop_recording.set() + return + if name in key_to_event: + events.request_transition(key_to_event[name]) + if name == cfg.upload: + events.upload_requested.set() - listener = keyboard.Listener(on_press=on_press) - listener.start() - logger.info( - "DAgger keyboard listener started (pause_resume='%s', correction='%s', upload='%s', ESC=stop)", - cfg.pause_resume, - cfg.correction, - cfg.upload, + return create_key_listener( + dispatch, + controls_help=( + f"pause_resume='{cfg.pause_resume}', correction='{cfg.correction}', " + f"upload='{cfg.upload}', ESC=stop" + ), ) - return listener def _init_dagger_pedal(events: DAggerEvents, cfg: DAggerPedalConfig): @@ -328,7 +284,7 @@ class DAggerStrategy(RolloutStrategy): logger.info("Stopping DAgger recording") log_say("Stopping DAgger recording", play_sounds) - if self._listener is not None and not is_headless(): + if self._listener is not None: logger.info("Stopping keyboard listener") self._listener.stop() diff --git a/src/lerobot/rollout/strategies/episodic.py b/src/lerobot/rollout/strategies/episodic.py index e925fb2ea..e70e66787 100644 --- a/src/lerobot/rollout/strategies/episodic.py +++ b/src/lerobot/rollout/strategies/episodic.py @@ -35,14 +35,13 @@ import time from lerobot.common.control_utils import ( follower_smooth_move_to, - init_keyboard_listener, - is_headless, teleop_smooth_move_to, teleop_supports_feedback, ) from lerobot.datasets import VideoEncodingManager from lerobot.utils.constants import ACTION, OBS_STR from lerobot.utils.feature_utils import build_dataset_frame +from lerobot.utils.keyboard_input import init_keyboard_listener from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.utils import log_say from lerobot.utils.visualization_utils import log_rerun_data @@ -307,7 +306,7 @@ class EpisodicStrategy(RolloutStrategy): log_say("Stop recording", play_sounds, blocking=True) - if not is_headless() and self._listener is not None: + if self._listener is not None: self._listener.stop() if ctx.data.dataset is not None: diff --git a/src/lerobot/rollout/strategies/highlight.py b/src/lerobot/rollout/strategies/highlight.py index baff70da7..385a9e2b6 100644 --- a/src/lerobot/rollout/strategies/highlight.py +++ b/src/lerobot/rollout/strategies/highlight.py @@ -18,17 +18,14 @@ from __future__ import annotations import contextlib import logging -import os -import sys import time from concurrent.futures import Future, ThreadPoolExecutor from threading import Event as ThreadingEvent, Lock -from lerobot.common.control_utils import is_headless from lerobot.datasets import VideoEncodingManager from lerobot.utils.constants import ACTION, OBS_STR from lerobot.utils.feature_utils import build_dataset_frame -from lerobot.utils.import_utils import _pynput_available, require_package +from lerobot.utils.keyboard_input import create_key_listener from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.utils import log_say @@ -37,19 +34,6 @@ from ..context import RolloutContext from ..ring_buffer import RolloutRingBuffer from .core import RolloutStrategy, safe_push_to_hub, send_next_action -PYNPUT_AVAILABLE = _pynput_available -keyboard = None -if PYNPUT_AVAILABLE: - try: - if ("DISPLAY" not in os.environ) and ("linux" in sys.platform): - logging.info("No DISPLAY set. Skipping pynput import.") - PYNPUT_AVAILABLE = False - else: - from pynput import keyboard - except Exception as e: - PYNPUT_AVAILABLE = False - logging.info(f"Could not import pynput: {e}") - logger = logging.getLogger(__name__) @@ -72,7 +56,6 @@ class HighlightStrategy(RolloutStrategy): def __init__(self, config: HighlightStrategyConfig): super().__init__(config) - require_package("pynput", extra="pynput-dep") self._ring: RolloutRingBuffer | None = None self._listener = None self._save_requested = ThreadingEvent() @@ -234,30 +217,27 @@ class HighlightStrategy(RolloutStrategy): logger.info("Highlight strategy teardown complete") def _setup_keyboard(self, shutdown_event: ThreadingEvent) -> None: - """Set up keyboard listener for save and push keys.""" - if is_headless(): - logger.warning("Headless environment — highlight keys unavailable") - return + """Set up a keyboard listener for the save and push keys. - try: - save_key = self.config.save_key - push_key = self.config.push_key + Backend selection (pynput on X11 / trusted-macOS / Windows, a terminal reader on + Wayland / headless TTY) is delegated to :func:`create_key_listener`. + """ + save_key = self.config.save_key + push_key = self.config.push_key - def on_press(key): - with contextlib.suppress(Exception): - if hasattr(key, "char") and key.char == save_key: - self._save_requested.set() - elif hasattr(key, "char") and key.char == push_key: - self._push_requested.set() - elif key == keyboard.Key.esc: - self._save_requested.clear() - shutdown_event.set() + def dispatch(name: str) -> None: + """Apply a resolved key name to the highlight events.""" + if name == save_key: + self._save_requested.set() + elif name == push_key: + self._push_requested.set() + elif name == "esc": + self._save_requested.clear() + shutdown_event.set() - self._listener = keyboard.Listener(on_press=on_press) - self._listener.start() - logger.info("Keyboard listener started (save='%s', push='%s', ESC=stop)", save_key, push_key) - except ImportError: - logger.warning("pynput not available — keyboard listener disabled") + self._listener = create_key_listener( + dispatch, controls_help=f"save='{save_key}', push='{push_key}', ESC=stop" + ) def _background_push(self, dataset, cfg) -> None: """Queue a Hub push on the single-worker executor.""" diff --git a/src/lerobot/scripts/lerobot_record.py b/src/lerobot/scripts/lerobot_record.py index 0deb54b90..4d5518c7c 100644 --- a/src/lerobot/scripts/lerobot_record.py +++ b/src/lerobot/scripts/lerobot_record.py @@ -96,11 +96,7 @@ from lerobot.cameras.opencv import OpenCVCameraConfig # noqa: F401 from lerobot.cameras.reachy2_camera import Reachy2CameraConfig # noqa: F401 from lerobot.cameras.realsense import RealSenseCameraConfig # noqa: F401 from lerobot.cameras.zmq import ZMQCameraConfig # noqa: F401 -from lerobot.common.control_utils import ( - init_keyboard_listener, - is_headless, - sanity_check_dataset_robot_compatibility, -) +from lerobot.common.control_utils import sanity_check_dataset_robot_compatibility from lerobot.configs import parser from lerobot.configs.dataset import DatasetRecordConfig from lerobot.datasets import ( @@ -155,6 +151,7 @@ from lerobot.teleoperators.keyboard import KeyboardTeleop from lerobot.utils.constants import ACTION, OBS_STR from lerobot.utils.feature_utils import build_dataset_frame, combine_feature_dicts from lerobot.utils.import_utils import register_third_party_plugins +from lerobot.utils.keyboard_input import init_keyboard_listener from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.utils import ( init_logging, @@ -508,7 +505,7 @@ def record( if teleop and teleop.is_connected: teleop.disconnect() - if not is_headless() and listener: + if listener is not None: listener.stop() if cfg.dataset.push_to_hub: diff --git a/src/lerobot/teleoperators/gamepad/gamepad_utils.py b/src/lerobot/teleoperators/gamepad/gamepad_utils.py index c1531ca84..22dbb7cca 100644 --- a/src/lerobot/teleoperators/gamepad/gamepad_utils.py +++ b/src/lerobot/teleoperators/gamepad/gamepad_utils.py @@ -18,6 +18,7 @@ import logging from typing import TYPE_CHECKING from lerobot.utils.import_utils import _hidapi_available, _pygame_available, require_package +from lerobot.utils.keyboard_input import pynput_can_capture from ..utils import TeleopEvents @@ -123,6 +124,15 @@ class KeyboardController(InputController): def start(self): """Start the keyboard listener.""" + if not pynput_can_capture(): + logging.warning( + "Keyboard control is unavailable in this environment. pynput cannot capture keys " + "on Wayland or headless machines, or on macOS without Accessibility / Input " + "Monitoring permission. Keyboard motion will be inactive." + ) + self.running = False + return + from pynput import keyboard def on_press(key): diff --git a/src/lerobot/teleoperators/keyboard/teleop_keyboard.py b/src/lerobot/teleoperators/keyboard/teleop_keyboard.py index 801789bcb..872cc7a26 100644 --- a/src/lerobot/teleoperators/keyboard/teleop_keyboard.py +++ b/src/lerobot/teleoperators/keyboard/teleop_keyboard.py @@ -15,8 +15,6 @@ # limitations under the License. import logging -import os -import sys import time from queue import Queue from typing import Any @@ -24,6 +22,7 @@ from typing import Any from lerobot.types import RobotAction from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected from lerobot.utils.import_utils import _pynput_available, require_package +from lerobot.utils.keyboard_input import pynput_can_capture from ..teleoperator import Teleoperator from ..utils import TeleopEvents @@ -37,14 +36,10 @@ PYNPUT_AVAILABLE = _pynput_available keyboard = None if PYNPUT_AVAILABLE: try: - if ("DISPLAY" not in os.environ) and ("linux" in sys.platform): - logging.info("No DISPLAY set. Skipping pynput import.") - PYNPUT_AVAILABLE = False - else: - from pynput import keyboard + from pynput import keyboard except Exception as e: PYNPUT_AVAILABLE = False - logging.info(f"Could not import pynput: {e}") + logging.info("Could not import pynput keyboard backend: %s", e) class KeyboardTeleop(Teleoperator): @@ -88,7 +83,7 @@ class KeyboardTeleop(Teleoperator): @check_if_already_connected def connect(self) -> None: - if PYNPUT_AVAILABLE: + if PYNPUT_AVAILABLE and pynput_can_capture(): logging.info("pynput is available - enabling local keyboard listener.") self.listener = keyboard.Listener( on_press=self._on_press, @@ -96,7 +91,13 @@ class KeyboardTeleop(Teleoperator): ) self.listener.start() else: - logging.info("pynput not available - skipping local keyboard listener.") + logging.warning( + "Keyboard teleoperation is unavailable in this environment. pynput can only " + "capture key events on an X11 session (Linux), a Windows desktop, or macOS with " + "Accessibility / Input Monitoring granted - not on Wayland or headless machines. " + "This keyboard teleoperator will produce no actions; use an X11 session, a " + "gamepad, or a leader-arm teleoperator instead." + ) self.listener = None def calibrate(self) -> None: diff --git a/src/lerobot/utils/keyboard_input.py b/src/lerobot/utils/keyboard_input.py new file mode 100644 index 000000000..00c0f53ec --- /dev/null +++ b/src/lerobot/utils/keyboard_input.py @@ -0,0 +1,440 @@ +# Copyright 2026 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. + +"""Display-independent keyboard input for interactive controls. + +This module centralizes everything related to *discrete* keyboard controls +(end-episode-early, re-record, stop, and the rollout strategies' custom keys): + +* environment detection — :func:`is_headless`, :func:`is_wayland`, + :func:`pynput_can_capture` (the single predicate every call-site should use to + decide whether ``pynput`` can actually capture keys here); +* a shared key mapping — :func:`apply_recording_control`; and +* two interchangeable backends behind one ``(listener, events)`` contract: + the ``pynput`` global listener (X11 / trusted-macOS / Windows) and a + standard-library :class:`TerminalKeyListener` that reads the controlling TTY + (Wayland / headless-SSH-with-TTY / macOS without Accessibility permission). + +NOTE: *continuous* key-state teleoperation ("hold a key to keep moving") is +deliberately NOT served here. A terminal in cbreak mode delivers only key-down +bytes — there is no key-release event — so the held-key model cannot be +reproduced. Those teleoperators stay on ``pynput`` and use +:func:`pynput_can_capture` to warn instead of silently doing nothing. +""" + +from __future__ import annotations + +import atexit +import contextlib +import logging +import os +import platform +import select +import sys +import threading +import time +from collections.abc import Callable +from functools import cache +from typing import TYPE_CHECKING + +from .import_utils import _pynput_available + +logger = logging.getLogger(__name__) + +# POSIX-only terminal modules (absent on Windows, where the pynput backend is used). +if TYPE_CHECKING: + import termios + import tty + + _TERMIOS_AVAILABLE = True +else: + try: + import termios + import tty + + _TERMIOS_AVAILABLE = True + except ImportError: # POSIX-only modules; unavailable on Windows + termios = tty = None + _TERMIOS_AVAILABLE = False + +keyboard = None +if _pynput_available: + try: + from pynput import keyboard + except Exception as e: # e.g. no reachable X display on a headless Linux box + logger.info("Could not import pynput keyboard backend: %s", e) + + +@cache +def is_headless() -> bool: + """Return ``True`` when no display server is available. + + * Linux: headless when neither ``DISPLAY`` (X11) nor ``WAYLAND_DISPLAY`` is set. + * macOS / Windows: a display is always assumed to be present. A genuinely GUI-less + Mac/Windows CI host would be misclassified but it doesn't matter, because the + sys.stdin.isatty() gate returns None there regardless. + """ + if platform.system() == "Linux": + return not (os.environ.get("DISPLAY") or os.environ.get("WAYLAND_DISPLAY")) + return False + + +@cache +def is_wayland() -> bool: + """Return ``True`` when running under a Wayland session. + + ``pynput`` relies on an X11 backend. Under Wayland it still imports (XWayland + is usually present and ``$DISPLAY`` is set) but cannot capture *global* + hotkeys, so the documented arrow/Esc shortcuts silently do nothing. This case + is invisible to :func:`is_headless`, hence the dedicated check. + """ + return os.environ.get("XDG_SESSION_TYPE", "").lower() == "wayland" or bool( + os.environ.get("WAYLAND_DISPLAY") + ) + + +@cache +def pynput_can_capture() -> bool: + """Return ``True`` when a ``pynput`` global listener can actually capture keys. + + This is the single predicate every keyboard call-site should use to choose + between the ``pynput`` backend and a fallback. It is intentionally + conservative: + + * Linux: only a real X11 session (a display is present *and* it is not Wayland). + * macOS: ``True`` here — Accessibility / Input-Monitoring permission + (``IS_TRUSTED``) can only be confirmed at runtime *after* starting a + listener, so :func:`init_keyboard_listener` refines this with + :func:`pynput_listener_is_trusted`. + * Windows: ``True`` (the low-level global hook needs no special permission). + + Always ``False`` when ``pynput`` is not installed. + """ + if not _pynput_available: + return False + if platform.system() == "Linux": + return not is_headless() and not is_wayland() + return True + + +def pynput_listener_is_trusted(listener, timeout_s: float = 1.0) -> bool: + """Best-effort check that a freshly started ``pynput`` listener can capture. + + On macOS, ``pynput`` sets ``listener.IS_TRUSTED`` on its *listener thread* + once the Quartz event tap is created; the class default is ``False``. We + therefore wait for the thread to either flip it ``True`` (trusted) or for a + short timeout to elapse (untrusted — it stays ``False`` forever). This biases + toward the common trusted case (returns as soon as the flag flips) and only + pays the full ``timeout_s`` on an already-broken untrusted machine. + + On non-macOS backends the attribute is absent and capture is assumed to work. + """ + if platform.system() != "Darwin": + return True + deadline = time.perf_counter() + timeout_s + while time.perf_counter() < deadline: + if getattr(listener, "IS_TRUSTED", False): + return True + time.sleep(0.02) + return bool(getattr(listener, "IS_TRUSTED", False)) + + +def apply_recording_control(control: str, events: dict) -> None: + """Apply a recording control-flow key press to the shared ``events`` dict. + + Centralizes the mapping so the ``pynput`` and terminal backends behave + identically. ``control`` is one of ``"right"`` (end the loop early), ``"left"`` + (re-record the last episode), or ``"esc"`` (stop recording). + """ + if control == "right": + print("Right arrow key pressed. Exiting loop...") + events["exit_early"] = True + elif control == "left": + print("Left arrow key pressed. Exiting loop and rerecord the last episode...") + events["rerecord_episode"] = True + events["exit_early"] = True + elif control == "esc": + print("Escape key pressed. Stopping data recording...") + events["stop_recording"] = True + events["exit_early"] = True + + +# Terminal arrow keys arrive as a 3-byte escape sequence whose *final* byte identifies +# the direction. Two encodings exist depending on the terminal's cursor-key mode — CSI +# ("ESC [ X") and SS3 ("ESC O X", common over SSH/tmux) — but both share the same final +# byte, so this single table decodes either. Looked up by TerminalKeyListener._parse; +# an unknown final byte yields None (sequence ignored). +_ARROW_FINAL_BYTES = {"A": "up", "B": "down", "C": "right", "D": "left"} + + +class TerminalKeyListener: + """Display-independent keyboard listener that reads keys from the controlling TTY. + + Used as the Wayland / headless / macOS-untrusted equivalent of the ``pynput`` + listener for *discrete* controls. It puts the terminal into cbreak mode with + echo disabled and reads bytes on a daemon thread, decoding them into logical + key names that are passed to ``on_key``: + + * arrow keys (``ESC [ C`` / ``ESC O C`` …) -> ``"right"`` / ``"left"`` / ``"up"`` / ``"down"`` + * a bare ``ESC`` -> ``"esc"`` + * Enter / Tab / Space / Backspace -> ``"enter"`` / ``"tab"`` / ``"space"`` / ``"backspace"`` + * any other printable byte -> that character (e.g. ``"n"``, ``"s"``) + + Only key-down events are produced (terminals have no key-release), so this is + suitable for discrete commands but NOT for continuous "hold-to-move" teleop. + + The terminal is restored on :meth:`stop` and also via an ``atexit`` hook, so a + crash or Ctrl-C never leaves the shell in a no-echo cbreak state. POSIX-only + (``termios`` / ``tty`` / ``select``); those modules are imported lazily so this + file stays importable on Windows (where ``pynput`` is used instead). + """ + + def __init__(self, on_key: Callable[[str], None]): + self._on_key = on_key + self._running = False + self._thread: threading.Thread | None = None + self._fd: int | None = None + self._old_attrs = None + + def _read_char(self, timeout: float) -> str | None: + """Return one character from stdin within ``timeout`` seconds, or ``None``.""" + if self._fd is None: + return None + ready, _, _ = select.select([self._fd], [], [], timeout) + if not ready: + return None + try: + data = os.read(self._fd, 1) + except OSError: + return None + if not data: + return None + return data.decode(errors="ignore") + + def _parse(self, ch: str) -> str | None: + """Decode one (possibly multi-byte) key starting at ``ch`` into a key name.""" + if ch == "\x1b": + # Possible CSI / SS3 escape sequence (arrow keys) or a bare ESC. Use + # short follow-up reads so a lone ESC is not mistaken for a sequence. + ch2 = self._read_char(timeout=0.02) + if ch2 is None: + return "esc" + if ch2 in ("[", "O"): + ch3 = self._read_char(timeout=0.02) + return _ARROW_FINAL_BYTES.get(ch3 or "") + # Some other escape sequence (e.g. Alt+key); ignore it. + return None + if ch in ("\r", "\n"): + return "enter" + if ch == "\t": + return "tab" + if ch == " ": + return "space" + if ch in ("\x7f", "\x08"): + return "backspace" + if ch.isprintable(): + return ch + return None + + def _run(self) -> None: + while self._running: + ch = self._read_char(timeout=0.05) + if ch is None: + continue + name = self._parse(ch) + if name is None: + continue + try: + self._on_key(name) + except Exception as e: # never let a handler error kill the reader thread + logger.debug("Terminal key handler error: %s", e) + + def start(self) -> None: + """Switch the terminal to cbreak mode (echo off) and read keys on a daemon thread. + + No-op when stdin is not a TTY (piped/redirected input) or on platforms + without ``termios`` (e.g. Windows), so non-interactive runs are unaffected. + """ + if not sys.stdin.isatty(): + return + if not _TERMIOS_AVAILABLE: # POSIX-only modules (e.g. unavailable on Windows) + logger.warning("Terminal keyboard input is not supported on this platform.") + return + + self._fd = sys.stdin.fileno() + self._old_attrs = termios.tcgetattr(self._fd) + tty.setcbreak(self._fd) + # Explicitly disable ECHO so arrow-key escape sequences (e.g. ^[[C) are not + # echoed as garbage into the recording terminal. (Independent of the + # version-specific behavior of tty.setcbreak.) + new_attrs = termios.tcgetattr(self._fd) + new_attrs[3] &= ~termios.ECHO # index 3 == lflags + termios.tcsetattr(self._fd, termios.TCSADRAIN, new_attrs) + # Safety net: restore the terminal even if stop() is never reached (crash). + atexit.register(self.stop) + + self._running = True + self._thread = threading.Thread(target=self._run, daemon=True) + self._thread.start() + + def stop(self) -> None: + """Stop the reader thread and restore the original terminal attributes. + + Idempotent: safe to call multiple times (e.g. explicitly and via atexit). + """ + self._running = False + thread = self._thread + if thread is not None: + thread.join(timeout=0.5) + self._thread = None + if self._fd is not None and self._old_attrs is not None and _TERMIOS_AVAILABLE: + try: + termios.tcsetattr(self._fd, termios.TCSADRAIN, self._old_attrs) + finally: + self._old_attrs = None + with contextlib.suppress(Exception): + atexit.unregister(self.stop) + + +# Map pynput key objects to the same canonical names TerminalKeyListener emits, so a +# single dispatch works across both backends. Empty when pynput is unavailable. +if keyboard is not None: + _PYNPUT_KEY_NAMES = { + keyboard.Key.right: "right", + keyboard.Key.left: "left", + keyboard.Key.up: "up", + keyboard.Key.down: "down", + keyboard.Key.esc: "esc", + keyboard.Key.enter: "enter", + keyboard.Key.tab: "tab", + keyboard.Key.space: "space", + keyboard.Key.backspace: "backspace", + } +else: + _PYNPUT_KEY_NAMES = {} + + +def _resolve_pynput_key(key) -> str | None: + """Resolve a pynput key event to the canonical name TerminalKeyListener also emits. + + Special keys map through :data:`_PYNPUT_KEY_NAMES`; character keys fall back to their + ``.char`` (e.g. ``"n"``). Returns ``None`` for keys with no mapping and no character. + """ + name = _PYNPUT_KEY_NAMES.get(key) + if name is not None: + return name + # ``or None`` keeps the historical truthy-char semantics: an empty/None char is "no key". + return getattr(key, "char", None) or None + + +def create_key_listener(dispatch: Callable[[str], None], *, controls_help: str = ""): + """Start a keyboard listener that routes resolved key names to ``dispatch``. + + Shared backend selection used by recording and the rollout strategies: + + * the ``pynput`` global listener on X11 / trusted-macOS / Windows (on macOS the + listener's ``IS_TRUSTED`` flag is checked after start, and an untrusted listener is + stopped so the terminal backend is used instead); + * the stdlib :class:`TerminalKeyListener` on Wayland / headless sessions with a TTY; + * ``None`` when no backend is usable (non-interactive / piped runs). + + Both backends pass ``dispatch`` the same canonical key names ("right" / "left" / "up" / + "down" / "esc" / "enter" / "tab" / "space" / "backspace", or a character), so one + ``dispatch`` works regardless of backend. ``controls_help`` is an optional hint + appended to the log messages. + + Returns the listener (exposing ``.stop()``) or ``None``. + """ + suffix = f" ({controls_help})" if controls_help else "" + + if pynput_can_capture() and keyboard is not None: + + def on_press(key): + with contextlib.suppress(Exception): + name = _resolve_pynput_key(key) + if name is not None: + dispatch(name) + + listener = keyboard.Listener(on_press=on_press) + listener.start() + if pynput_listener_is_trusted(listener): + logger.info("Keyboard listener started%s.", suffix) + return listener + # macOS without Accessibility / Input-Monitoring permission: the listener never + # fires. Stop it and fall through to the terminal backend. + logger.warning( + "pynput keyboard listener is not trusted (missing macOS Accessibility / " + "Input Monitoring permission); falling back to terminal keyboard input." + ) + listener.stop() + + if sys.stdin.isatty(): + listener = TerminalKeyListener(dispatch) + listener.start() + logger.info("Using terminal keyboard input — keep this terminal focused%s.", suffix) + return listener + + logger.warning( + "Keyboard controls unavailable: no usable display (Wayland/headless) and stdin is " + "not an interactive terminal%s.", + suffix, + ) + return None + + +def init_keyboard_listener(): + """Initialize a non-blocking keyboard listener for interactive recording controls. + + Backend selection: + + * ``pynput`` global listener when :func:`pynput_can_capture` is true (real + X11, macOS, Windows). On macOS the listener's ``IS_TRUSTED`` flag is checked + after start; if the process lacks Accessibility / Input-Monitoring + permission, the listener is stopped and the terminal backend is used. + * a :class:`TerminalKeyListener` reading the controlling TTY when ``pynput`` + cannot capture (Wayland / headless-SSH / macOS-untrusted) *and* stdin is a TTY. + * otherwise no listener (non-interactive / piped runs) — recording relies on + the episode/reset timers (or Ctrl+C). + + Both backends accept the same controls: Right/Left/Esc, plus the single-byte letter + equivalents ``n`` (next), ``r`` (re-record) and ``q`` (quit). The letters are the most + reliable choice over high-latency SSH/VNC links, where arrow-key escape sequences can + be split, delayed, or intercepted by the terminal. + + Returns: + A tuple ``(listener, events)`` where ``listener`` exposes ``.stop()`` or is + ``None``, and ``events`` is the dict of flags (``exit_early``, + ``rerecord_episode``, ``stop_recording``) set by key presses. + """ + events = { + "exit_early": False, + "rerecord_episode": False, + "stop_recording": False, + } + + # Accept the single-byte letter equivalents n/r/q alongside the arrow/Esc keys: the + # letters are immune to the escape-sequence split/delay/interception that affects arrows + # over laggy SSH/VNC links. Case-insensitive so Shift+letter still works. + def on_key(name: str) -> None: + key = name.lower() + if key in ("right", "n"): + apply_recording_control("right", events) + elif key in ("left", "r"): + apply_recording_control("left", events) + elif key in ("esc", "q"): + apply_recording_control("esc", events) + # other keys (incl. up/down) are intentionally ignored + + listener = create_key_listener(on_key, controls_help="Right/Left/Esc, or n=next, r=re-record, q=quit") + return listener, events diff --git a/tests/utils/test_keyboard_input.py b/tests/utils/test_keyboard_input.py new file mode 100644 index 000000000..2f0dee889 --- /dev/null +++ b/tests/utils/test_keyboard_input.py @@ -0,0 +1,228 @@ +#!/usr/bin/env python + +# Copyright 2026 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. + +"""Unit tests for the display-independent keyboard input helpers. + +These cover the parts most likely to regress: the environment-detection decision +table (the heart of the Wayland/headless fix), the macOS trust probe, the control +mapping, the terminal escape-sequence parsing, and backend selection. They require +neither ``pynput`` nor a real terminal. +""" + +import io +import platform +import sys + +import pytest + +import lerobot.utils.keyboard_input as ki +from lerobot.utils.keyboard_input import ( + TerminalKeyListener, + apply_recording_control, + create_key_listener, + init_keyboard_listener, + is_headless, + is_wayland, + pynput_can_capture, + pynput_listener_is_trusted, +) + + +@pytest.fixture(autouse=True) +def _clear_detection_caches(): + """The detection helpers are ``@cache``-decorated; clear around each test.""" + for fn in (is_headless, is_wayland, pynput_can_capture): + fn.cache_clear() + yield + for fn in (is_headless, is_wayland, pynput_can_capture): + fn.cache_clear() + + +def _set_platform(monkeypatch, name): + monkeypatch.setattr(platform, "system", lambda: name) + + +def _set_tty(monkeypatch, is_tty): + stdin = io.StringIO("") + stdin.isatty = lambda: is_tty + monkeypatch.setattr(sys, "stdin", stdin) + + +# --- Environment detection (the core of the fix) --------------------------- +@pytest.mark.parametrize( + ("system", "env", "expected"), + [ + ("Linux", {}, True), # no display server + ("Linux", {"DISPLAY": ":0"}, False), # X11 + ("Linux", {"WAYLAND_DISPLAY": "wayland-0"}, False), # Wayland + ("Darwin", {}, False), # display always assumed present + ], +) +def test_is_headless(monkeypatch, system, env, expected): + _set_platform(monkeypatch, system) + monkeypatch.delenv("DISPLAY", raising=False) + monkeypatch.delenv("WAYLAND_DISPLAY", raising=False) + for key, value in env.items(): + monkeypatch.setenv(key, value) + assert is_headless() is expected + + +@pytest.mark.parametrize( + ("env", "expected"), + [ + ({"XDG_SESSION_TYPE": "wayland"}, True), + ({"WAYLAND_DISPLAY": "wayland-0"}, True), + ({"XDG_SESSION_TYPE": "x11"}, False), + ({}, False), + ], +) +def test_is_wayland(monkeypatch, env, expected): + monkeypatch.delenv("XDG_SESSION_TYPE", raising=False) + monkeypatch.delenv("WAYLAND_DISPLAY", raising=False) + for key, value in env.items(): + monkeypatch.setenv(key, value) + assert is_wayland() is expected + + +@pytest.mark.parametrize( + ("system", "env", "pynput_available", "expected"), + [ + ("Linux", {"DISPLAY": ":0"}, True, True), # X11 + ("Linux", {"DISPLAY": ":0", "WAYLAND_DISPLAY": "wayland-0"}, True, False), # Wayland + ("Linux", {}, True, False), # headless + ("Darwin", {}, True, True), + ("Linux", {"DISPLAY": ":0"}, False, False), # pynput not installed + ], +) +def test_pynput_can_capture(monkeypatch, system, env, pynput_available, expected): + _set_platform(monkeypatch, system) + monkeypatch.setattr(ki, "_pynput_available", pynput_available) + for var in ("DISPLAY", "WAYLAND_DISPLAY", "XDG_SESSION_TYPE"): + monkeypatch.delenv(var, raising=False) + for key, value in env.items(): + monkeypatch.setenv(key, value) + assert pynput_can_capture() is expected + + +# --- macOS trust probe ------------------------------------------------------ +class _FakeListener: + def __init__(self, is_trusted): + self.IS_TRUSTED = is_trusted + + +def test_pynput_listener_is_trusted(monkeypatch): + _set_platform(monkeypatch, "Linux") + assert pynput_listener_is_trusted(_FakeListener(False)) is True # non-macOS: always assumed ok + _set_platform(monkeypatch, "Darwin") + assert pynput_listener_is_trusted(_FakeListener(False), timeout_s=0.05) is False + + +# --- Control mapping -------------------------------------------------------- +def test_apply_recording_control(): + events = {"exit_early": False, "rerecord_episode": False, "stop_recording": False} + apply_recording_control("left", events) + assert events == {"exit_early": True, "rerecord_episode": True, "stop_recording": False} + apply_recording_control("esc", events) + assert events["stop_recording"] is True + apply_recording_control("up", events) # unknown control -> no-op (no error) + + +# --- Terminal escape-sequence parsing (the tricky bit) ---------------------- +def _drive(listener, byte_seq): + """Run the listener's read loop over a scripted list of bytes (no real terminal).""" + script = list(byte_seq) + + def fake_read(timeout): + if script: + return script.pop(0) + listener._running = False + return None + + listener._read_char = fake_read + listener._running = True + listener._run() + + +@pytest.mark.parametrize( + ("byte_seq", "expected"), + [ + (["\x1b", "[", "C"], ["right"]), # CSI arrow + (["\x1b", "O", "D"], ["left"]), # SS3 arrow (e.g. over SSH/tmux) + (["\x1b"], ["esc"]), # bare ESC + (["\x1b", "[", "A"], ["up"]), # decoded even though the record handler ignores it + (["n"], ["n"]), # letter passthrough + ], +) +def test_terminal_parsing(byte_seq, expected): + collected = [] + _drive(TerminalKeyListener(collected.append), byte_seq) + assert collected == expected + + +# --- Backend selection ------------------------------------------------------ +def test_init_selects_terminal_when_pynput_cannot_capture(monkeypatch): + monkeypatch.setattr(ki, "pynput_can_capture", lambda: False) + _set_tty(monkeypatch, is_tty=True) + monkeypatch.setattr(TerminalKeyListener, "start", lambda self: None) # avoid touching termios + listener, _ = init_keyboard_listener() + assert isinstance(listener, TerminalKeyListener) + + +def test_init_returns_none_without_tty(monkeypatch): + monkeypatch.setattr(ki, "pynput_can_capture", lambda: False) + _set_tty(monkeypatch, is_tty=False) + listener, _ = init_keyboard_listener() + assert listener is None + + +@pytest.mark.parametrize( + ("key", "flag"), + [("right", "exit_early"), ("r", "rerecord_episode"), ("q", "stop_recording")], +) +def test_init_terminal_key_routing(monkeypatch, key, flag): + """Arrows and their letter equivalents drive the same events (terminal backend).""" + monkeypatch.setattr(ki, "pynput_can_capture", lambda: False) + _set_tty(monkeypatch, is_tty=True) + monkeypatch.setattr(TerminalKeyListener, "start", lambda self: None) + listener, events = init_keyboard_listener() + listener._on_key(key) + assert events[flag] is True + + +# --- Shared factory + pynput key resolver ----------------------------------- +def test_resolve_pynput_key_char_fallback(): + """Unmapped keys fall back to ``.char`` (and yield None when there is none).""" + assert ki._resolve_pynput_key(type("K", (), {"char": "s"})()) == "s" + assert ki._resolve_pynput_key(type("K", (), {"char": None})()) is None + assert ki._resolve_pynput_key(type("K", (), {"char": ""})()) is None # empty char -> no key + + +def test_create_key_listener_routes_to_dispatch(monkeypatch): + """The terminal backend forwards canonical key names straight to ``dispatch``.""" + monkeypatch.setattr(ki, "pynput_can_capture", lambda: False) + _set_tty(monkeypatch, is_tty=True) + monkeypatch.setattr(TerminalKeyListener, "start", lambda self: None) + seen = [] + listener = create_key_listener(seen.append, controls_help="save='s'") + assert isinstance(listener, TerminalKeyListener) + listener._on_key("space") + assert seen == ["space"] + + +def test_create_key_listener_none_without_tty(monkeypatch): + monkeypatch.setattr(ki, "pynput_can_capture", lambda: False) + _set_tty(monkeypatch, is_tty=False) + assert create_key_listener(lambda name: None) is None From 324086abc34bd5d3c2095b37ec6e2dc74762be31 Mon Sep 17 00:00:00 2001 From: Eric Chan Date: Thu, 25 Jun 2026 04:58:08 -0700 Subject: [PATCH 15/38] Update follower arm description in documentation (#3780) Signed-off-by: Eric Chan --- docs/source/so101.mdx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/so101.mdx b/docs/source/so101.mdx index 1274b8282..5b4ed0985 100644 --- a/docs/source/so101.mdx +++ b/docs/source/so101.mdx @@ -122,7 +122,7 @@ The video below shows the sequence of steps for setting the motor ids. #### Follower -Connect the usb cable from your computer and the power supply to the follower arm's controller board. Then, run the following command or run the API example with the port you got from the previous step. You'll also need to give your leader arm a name with the `id` parameter. +Connect the usb cable from your computer and the power supply to the follower arm's controller board. Then, run the following command or run the API example with the port you got from the previous step. You'll also need to give your follower arm a name with the `id` parameter. From c3f180e115a5833befd16f06467dfc2e12c943e1 Mon Sep 17 00:00:00 2001 From: Khalil Meftah Date: Thu, 25 Jun 2026 14:19:35 +0200 Subject: [PATCH 16/38] refactor(policies): clean MolmoAct2 to follow EO1/TOPReward patterns (#3724) Align the MolmoAct2 implementation with lerobot codebase conventions: - Rename hf_model/ to molmoact2_hf_model/ - Slim config: move all I/O and runtime logic to modeling - Remove blanket from 8 vendored files, fix 66 lint issues - Deduplicate _hf_token() and _resolve_checkpoint_location() - Make huggingface_hub imports lazy - Remove custom MolmoAct2CosineDecayWithWarmupSchedulerConfig, use base class - Extract 13 static/classmethods from MolmoAct2Policy to free functions - Replace print() with logger in vendored action_tokenizer - Add module docstrings, class docstring, and key method docstrings - Add module-level loggers to modeling and processor - Fix docs: pip to uv install, deduplicate README symlink - Remove shebangs from all files --- docs/source/molmoact2.mdx | 6 +- src/lerobot/policies/molmoact2/README.md | 2 +- src/lerobot/policies/molmoact2/__init__.py | 2 - .../molmoact2/configuration_molmoact2.py | 279 +------- .../policies/molmoact2/modeling_molmoact2.py | 630 ++++++++++++------ .../__init__.py | 4 - .../action_tokenizer.py | 17 +- .../configuration_molmoact2.py | 5 +- .../image_processing_molmoact2.py | 30 +- .../inference.py | 13 +- .../modeling_molmoact2.py | 23 +- .../processing_molmoact2.py | 42 +- .../video_processing_molmoact2.py | 63 +- .../policies/molmoact2/processor_molmoact2.py | 106 ++- tests/policies/molmoact2/test_molmoact2.py | 83 +-- 15 files changed, 611 insertions(+), 694 deletions(-) rename src/lerobot/policies/molmoact2/{hf_model => molmoact2_hf_model}/__init__.py (94%) rename src/lerobot/policies/molmoact2/{hf_model => molmoact2_hf_model}/action_tokenizer.py (96%) rename src/lerobot/policies/molmoact2/{hf_model => molmoact2_hf_model}/configuration_molmoact2.py (99%) rename src/lerobot/policies/molmoact2/{hf_model => molmoact2_hf_model}/image_processing_molmoact2.py (98%) rename src/lerobot/policies/molmoact2/{hf_model => molmoact2_hf_model}/inference.py (99%) rename src/lerobot/policies/molmoact2/{hf_model => molmoact2_hf_model}/modeling_molmoact2.py (99%) rename src/lerobot/policies/molmoact2/{hf_model => molmoact2_hf_model}/processing_molmoact2.py (95%) rename src/lerobot/policies/molmoact2/{hf_model => molmoact2_hf_model}/video_processing_molmoact2.py (98%) diff --git a/docs/source/molmoact2.mdx b/docs/source/molmoact2.mdx index c6ae24e9e..7927c6694 100644 --- a/docs/source/molmoact2.mdx +++ b/docs/source/molmoact2.mdx @@ -17,7 +17,7 @@ the paper, see [allenai/molmoact2](https://github.com/allenai/molmoact2). Install LeRobot with the MolmoAct2 optional dependencies: ```bash -pip install -e ".[molmoact2]" +uv sync --locked --extra molmoact2 ``` To run the models in this repository, you need an NVIDIA GPU. The measurements @@ -46,8 +46,8 @@ The repo has been tested with Ubuntu 22.04. To use MolmoAct2 in a LeRobot training config, set: -```python -policy.type=molmoact2 +```bash +--policy.type=molmoact2 ``` ## Training diff --git a/src/lerobot/policies/molmoact2/README.md b/src/lerobot/policies/molmoact2/README.md index ef419516d..9756785d9 120000 --- a/src/lerobot/policies/molmoact2/README.md +++ b/src/lerobot/policies/molmoact2/README.md @@ -1 +1 @@ -../../../../docs/source/policy_molmoact2_README.md \ No newline at end of file +../../../../docs/source/molmoact2.mdx \ No newline at end of file diff --git a/src/lerobot/policies/molmoact2/__init__.py b/src/lerobot/policies/molmoact2/__init__.py index bfef53bb2..a4e7695c2 100644 --- a/src/lerobot/policies/molmoact2/__init__.py +++ b/src/lerobot/policies/molmoact2/__init__.py @@ -1,5 +1,3 @@ -#!/usr/bin/env python - # Copyright 2026 The Allen Institute for Artificial Intelligence and The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/src/lerobot/policies/molmoact2/configuration_molmoact2.py b/src/lerobot/policies/molmoact2/configuration_molmoact2.py index de2585281..53aefdee6 100644 --- a/src/lerobot/policies/molmoact2/configuration_molmoact2.py +++ b/src/lerobot/policies/molmoact2/configuration_molmoact2.py @@ -1,5 +1,3 @@ -#!/usr/bin/env python - # Copyright 2026 The Allen Institute for Artificial Intelligence and The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -16,16 +14,9 @@ from __future__ import annotations -import json -import math -import os -from contextlib import suppress from dataclasses import dataclass, field -from pathlib import Path from typing import Any -from huggingface_hub import snapshot_download - from lerobot.configs import FeatureType, NormalizationMode, PolicyFeature, PreTrainedConfig from lerobot.optim import ( AdamWConfig, @@ -37,146 +28,6 @@ from lerobot.utils.constants import ACTION, OBS_STATE from ..rtc.configuration_rtc import RTCConfig -MOLMOACT2_DEFAULT_NUM_IMAGES = 2 -MOLMOACT2_IMAGE_TOKENS_PER_IMAGE = 196 -MOLMOACT2_FIXED_PROMPT_TOKEN_BUDGET = 80 -MOLMOACT2_TASK_TOKEN_BUDGET = 32 -MOLMOACT2_SEQUENCE_LENGTH_MARGIN = 32 -MOLMOACT2_SEQUENCE_LENGTH_MULTIPLE = 64 -MOLMOACT2_DISCRETE_ACTION_WRAPPER_TOKENS = 4 -MOLMOACT2_MIN_DISCRETE_ACTION_TOKENS_PER_STEP = 6 -MOLMOACT2_DISCRETE_ACTION_TOKENS_PER_DIM = 0.95 - - -def _hf_token() -> str | None: - return os.environ.get("HF_TOKEN") or os.environ.get("HF_ACCESS_TOKEN") - - -def _resolve_checkpoint_location( - checkpoint_path: str, - *, - revision: str | None = None, - force_download: bool = False, -) -> str: - checkpoint_path = str(checkpoint_path or "").strip() - if not checkpoint_path: - raise ValueError("MolmoAct2 policy requires `checkpoint_path`.") - local_path = Path(checkpoint_path).expanduser() - if local_path.exists(): - return str(local_path) - return snapshot_download( - repo_id=checkpoint_path, - repo_type="model", - revision=revision, - force_download=force_download, - ignore_patterns=["*.py", "*.pyc", "__pycache__/*"], - token=_hf_token(), - ) - - -def _load_hf_norm_metadata_for_tag( - checkpoint_path: str, - *, - revision: str | None, - force_download: bool, - norm_tag: str | None, -) -> dict[str, Any]: - norm_tag = str(norm_tag or "").strip() - if not norm_tag: - return {} - checkpoint_location = Path( - _resolve_checkpoint_location( - checkpoint_path, - revision=revision, - force_download=force_download, - ) - ) - norm_stats_filename = "norm_stats.json" - config_path = checkpoint_location / "config.json" - if config_path.exists(): - with suppress(OSError, json.JSONDecodeError): - norm_stats_filename = str( - json.loads(config_path.read_text()).get("norm_stats_filename") or norm_stats_filename - ) - stats_path = checkpoint_location / norm_stats_filename - if not stats_path.exists(): - raise FileNotFoundError( - f"MolmoAct2 HF checkpoint is missing {norm_stats_filename!r}; cannot resolve norm_tag={norm_tag!r}." - ) - payload = json.loads(stats_path.read_text()) - metadata_by_tag = payload.get("metadata_by_tag") - if not isinstance(metadata_by_tag, dict): - raise ValueError(f"MolmoAct2 norm stats file {stats_path} has no metadata_by_tag mapping.") - metadata = metadata_by_tag.get(norm_tag) - if not isinstance(metadata, dict): - available = sorted(str(tag) for tag in metadata_by_tag) - raise ValueError(f"Unknown MolmoAct2 norm_tag={norm_tag!r}. Available tags: {available}.") - return metadata - - -@LRSchedulerConfig.register_subclass("molmoact2_cosine_decay_with_warmup") -@dataclass -class MolmoAct2CosineDecayWithWarmupSchedulerConfig(CosineDecayWithWarmupSchedulerConfig): - """MolmoAct2-local cosine scheduler with optional decay-step auto-match. - - LeRobot's generic cosine scheduler keeps an explicit integer decay length. - For MolmoAct2, leaving num_decay_steps unset means "decay across this run's - training steps"; build() is the first point where num_training_steps is known. - """ - - num_decay_steps: int | None - - def build(self, optimizer, num_training_steps: int): - return CosineDecayWithWarmupSchedulerConfig( - peak_lr=self.peak_lr, - decay_lr=self.decay_lr, - num_warmup_steps=self.num_warmup_steps, - num_decay_steps=num_training_steps if self.num_decay_steps is None else self.num_decay_steps, - ).build(optimizer, num_training_steps=num_training_steps) - - -def _round_up(value: int, multiple: int) -> int: - return int(math.ceil(value / multiple) * multiple) - - -def infer_molmoact2_max_sequence_length( - *, - num_images: int, - state_dim: int, - action_dim: int, - action_horizon: int, - include_discrete_action: bool, -) -> int: - """Infer the padded text/image sequence cap from MolmoAct2's fixed token layout.""" - if num_images < 1: - num_images = MOLMOACT2_DEFAULT_NUM_IMAGES - if state_dim < 0: - state_dim = 0 - if action_dim < 1: - action_dim = 1 - if action_horizon < 1: - action_horizon = 1 - - image_tokens = num_images * MOLMOACT2_IMAGE_TOKENS_PER_IMAGE - prompt_tokens = ( - MOLMOACT2_FIXED_PROMPT_TOKEN_BUDGET - + MOLMOACT2_TASK_TOKEN_BUDGET - + state_dim - + MOLMOACT2_SEQUENCE_LENGTH_MARGIN - ) - action_tokens = 0 - if include_discrete_action: - action_tokens_per_step = max( - MOLMOACT2_MIN_DISCRETE_ACTION_TOKENS_PER_STEP, - math.ceil(action_dim * MOLMOACT2_DISCRETE_ACTION_TOKENS_PER_DIM), - ) - action_tokens = MOLMOACT2_DISCRETE_ACTION_WRAPPER_TOKENS + action_horizon * action_tokens_per_step - - return _round_up( - image_tokens + prompt_tokens + action_tokens, - MOLMOACT2_SEQUENCE_LENGTH_MULTIPLE, - ) - @PreTrainedConfig.register_subclass("molmoact2") @dataclass @@ -255,7 +106,7 @@ class MolmoAct2Config(PreTrainedConfig): optimizer_grad_clip_norm: float = 1.0 scheduler_warmup_steps: int = 200 - scheduler_decay_steps: int | None = None + scheduler_decay_steps: int = 100_000 scheduler_decay_lr: float = 1e-6 normalization_mapping: dict[str, NormalizationMode] = field( @@ -333,41 +184,6 @@ class MolmoAct2Config(PreTrainedConfig): if self.max_sequence_length is not None and self.max_sequence_length < 1: raise ValueError(f"max_sequence_length must be >= 1 or None, got {self.max_sequence_length}.") - def inferred_max_sequence_length( - self, - *, - num_images: int | None = None, - state_dim: int | None = None, - action_dim: int | None = None, - action_horizon: int | None = None, - include_discrete_action: bool | None = None, - ) -> int: - if self.max_sequence_length is not None: - return int(self.max_sequence_length) - - if num_images is None: - num_images = len(self.image_keys) or len(self.image_features) or MOLMOACT2_DEFAULT_NUM_IMAGES - if state_dim is None: - state_feature = self.robot_state_feature - state_dim = int(state_feature.shape[0]) if state_feature is not None else 0 - if action_dim is None: - action_feature = self.action_feature - action_dim = ( - int(action_feature.shape[0]) if action_feature is not None else self.expected_max_action_dim - ) - if action_horizon is None: - action_horizon = self.chunk_size - if include_discrete_action is None: - include_discrete_action = self.action_mode in {"discrete", "both"} - - return infer_molmoact2_max_sequence_length( - num_images=int(num_images), - state_dim=int(state_dim), - action_dim=int(action_dim), - action_horizon=int(action_horizon), - include_discrete_action=bool(include_discrete_action), - ) - @property def observation_delta_indices(self) -> None: return None @@ -390,7 +206,7 @@ class MolmoAct2Config(PreTrainedConfig): ) def get_scheduler_preset(self) -> LRSchedulerConfig | None: - return MolmoAct2CosineDecayWithWarmupSchedulerConfig( + return CosineDecayWithWarmupSchedulerConfig( peak_lr=self.optimizer_lr, decay_lr=self.scheduler_decay_lr, num_warmup_steps=self.scheduler_warmup_steps, @@ -426,94 +242,3 @@ class MolmoAct2Config(PreTrainedConfig): shape=(self.expected_max_action_dim,), ) self.output_features[ACTION] = action_feature - - def apply_norm_tag_metadata(self) -> None: - if not str(self.norm_tag or "").strip(): - return - metadata = _load_hf_norm_metadata_for_tag( - self.checkpoint_path, - revision=self.checkpoint_revision, - force_download=bool(self.checkpoint_force_download), - norm_tag=self.norm_tag, - ) - if metadata.get("action_horizon") is not None: - self.chunk_size = int(metadata["action_horizon"]) - if metadata.get("n_action_steps") is not None: - self.n_action_steps = int(metadata["n_action_steps"]) - if not self.setup_type and metadata.get("setup_type") is not None: - self.setup_type = str(metadata["setup_type"]) - if not self.control_mode and metadata.get("control_mode") is not None: - self.control_mode = str(metadata["control_mode"]) - - def saved_policy_action_mode(self) -> str | None: - pretrained_path = getattr(self, "pretrained_path", None) - if pretrained_path is None: - return None - config_path = Path(pretrained_path) / "config.json" - if not config_path.exists(): - return None - try: - mode = json.loads(config_path.read_text()).get("action_mode") - except (OSError, json.JSONDecodeError): - return None - if mode in {"continuous", "discrete", "both"}: - return str(mode) - return None - - def training_action_mode(self, saved_policy_action_mode: str | None = None) -> str: - return saved_policy_action_mode or self.action_mode - - def validate_inference_action_mode(self, saved_policy_action_mode: str | None = None) -> None: - requested_mode = self.inference_action_mode - if requested_mode is None: - return - training_mode = self.training_action_mode(saved_policy_action_mode) - if requested_mode == "continuous" and training_mode == "discrete": - raise ValueError( - "MolmoAct2 checkpoint was trained with action_mode='discrete' and cannot run " - "continuous inference." - ) - if requested_mode == "discrete" and training_mode == "continuous": - raise ValueError( - "MolmoAct2 checkpoint was trained with action_mode='continuous' and cannot run " - "discrete inference. Train with action_mode='both' or action_mode='discrete' first." - ) - - def validate_checkpoint_action_mode( - self, - checkpoint_action_mode: str, - *, - has_action_expert: bool, - ) -> None: - if self.action_mode == "both" and checkpoint_action_mode != "both": - raise ValueError( - f"action_mode='both' requires checkpoint action_mode='both', got {checkpoint_action_mode!r}." - ) - if self.action_mode == "discrete" and checkpoint_action_mode not in {"discrete", "both"}: - raise ValueError( - f"action_mode='discrete' requires checkpoint action_mode in {{'discrete', 'both'}}, " - f"got {checkpoint_action_mode!r}." - ) - if self.action_mode in {"continuous", "both"} and not has_action_expert: - raise ValueError("Continuous MolmoAct2 training requires an action expert checkpoint.") - - def resolve_inference_action_mode( - self, - requested_mode: str | None, - saved_policy_action_mode: str | None = None, - ) -> str: - training_mode = self.training_action_mode(saved_policy_action_mode) - if requested_mode is None: - requested_mode = self.inference_action_mode - if requested_mode is None: - raise ValueError( - "MolmoAct2 inference requires `inference_action_mode` to be set explicitly " - "to either 'continuous' or 'discrete'." - ) - if requested_mode not in {"continuous", "discrete"}: - raise ValueError("MolmoAct2 inference_action_mode must be either 'continuous' or 'discrete'.") - if requested_mode == "continuous" and training_mode == "discrete": - raise ValueError("MolmoAct2 action_mode='discrete' checkpoint cannot run continuous inference.") - if requested_mode == "discrete" and training_mode == "continuous": - raise ValueError("MolmoAct2 action_mode='continuous' checkpoint cannot run discrete inference.") - return requested_mode diff --git a/src/lerobot/policies/molmoact2/modeling_molmoact2.py b/src/lerobot/policies/molmoact2/modeling_molmoact2.py index f86be0904..2cc85ab02 100644 --- a/src/lerobot/policies/molmoact2/modeling_molmoact2.py +++ b/src/lerobot/policies/molmoact2/modeling_molmoact2.py @@ -1,5 +1,3 @@ -#!/usr/bin/env python - # Copyright 2026 The Allen Institute for Artificial Intelligence and The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -14,9 +12,22 @@ # See the License for the specific language governing permissions and # limitations under the License. +"""MolmoAct2 policy for LeRobot. + +MolmoAct2 is a VLM-based robotics policy from Allen AI that combines a +Molmo vision-language backbone with a per-layer flow-matching action expert +for continuous action generation, plus an optional discrete action token +head. This module wraps the vendored HF model implementation +(``molmoact2_hf_model/``) into the LeRobot ``PreTrainedPolicy`` interface. + +Paper: https://allenai.org/blog/molmoact2 +Code: https://github.com/allenai/molmoact2 +""" + from __future__ import annotations import json +import logging import os import types from collections import deque @@ -35,13 +46,58 @@ from lerobot.utils.constants import ACTION from lerobot.utils.import_utils import _scipy_available, _transformers_available, require_package from ..rtc.modeling_rtc import RTCProcessor -from .configuration_molmoact2 import MolmoAct2Config, _hf_token, _resolve_checkpoint_location +from .configuration_molmoact2 import MolmoAct2Config + +logger = logging.getLogger(__name__) + + +def _hf_token() -> str | None: + return os.environ.get("HF_TOKEN") or os.environ.get("HF_ACCESS_TOKEN") + + +def _resolve_checkpoint_location( + checkpoint_path: str, + *, + revision: str | None = None, + force_download: bool = False, +) -> str: + """Resolve a checkpoint path to a local directory, downloading from Hub if needed.""" + checkpoint_path = str(checkpoint_path or "").strip() + if not checkpoint_path: + raise ValueError("MolmoAct2 policy requires `checkpoint_path`.") + from pathlib import Path + + local_path = Path(checkpoint_path).expanduser() + if local_path.exists(): + return str(local_path) + from huggingface_hub import snapshot_download + + return snapshot_download( + repo_id=checkpoint_path, + repo_type="model", + revision=revision, + force_download=force_download, + ignore_patterns=["*.py", "*.pyc", "__pycache__/*"], + token=_hf_token(), + ) + + +def _torch_dtype(dtype: str) -> torch.dtype: + """Convert a dtype name string to a torch.dtype.""" + if dtype == "float32": + return torch.float32 + if dtype == "bfloat16": + return torch.bfloat16 + if dtype == "float16": + return torch.float16 + raise ValueError(f"Unsupported dtype: {dtype}") + if TYPE_CHECKING or _transformers_available: from transformers.utils import SAFE_WEIGHTS_INDEX_NAME, SAFE_WEIGHTS_NAME - from .hf_model.configuration_molmoact2 import MolmoAct2Config as HFMolmoAct2Config - from .hf_model.modeling_molmoact2 import MolmoAct2ForConditionalGeneration + from .molmoact2_hf_model.configuration_molmoact2 import MolmoAct2Config as HFMolmoAct2Config + from .molmoact2_hf_model.modeling_molmoact2 import MolmoAct2ForConditionalGeneration else: SAFE_WEIGHTS_INDEX_NAME = "model.safetensors.index.json" SAFE_WEIGHTS_NAME = "model.safetensors" @@ -49,7 +105,7 @@ else: MolmoAct2ForConditionalGeneration = None if TYPE_CHECKING or (_transformers_available and _scipy_available): - from .hf_model.action_tokenizer import UniversalActionProcessor + from .molmoact2_hf_model.action_tokenizer import UniversalActionProcessor else: UniversalActionProcessor = None @@ -70,6 +126,156 @@ _MODEL_INPUT_KEYS = { } +def _load_hf_norm_metadata_for_tag( + checkpoint_path: str, + *, + revision: str | None, + force_download: bool, + norm_tag: str | None, +) -> dict[str, Any]: + """Read per-tag metadata from the checkpoint's ``norm_stats.json``.""" + norm_tag = str(norm_tag or "").strip() + if not norm_tag: + return {} + from contextlib import suppress + from pathlib import Path + + checkpoint_location = Path( + _resolve_checkpoint_location( + checkpoint_path, + revision=revision, + force_download=force_download, + ) + ) + norm_stats_filename = "norm_stats.json" + config_path = checkpoint_location / "config.json" + if config_path.exists(): + with suppress(OSError, json.JSONDecodeError): + norm_stats_filename = str( + json.loads(config_path.read_text()).get("norm_stats_filename") or norm_stats_filename + ) + stats_path = checkpoint_location / norm_stats_filename + if not stats_path.exists(): + raise FileNotFoundError( + f"MolmoAct2 HF checkpoint is missing {norm_stats_filename!r}; cannot resolve norm_tag={norm_tag!r}." + ) + payload = json.loads(stats_path.read_text()) + metadata_by_tag = payload.get("metadata_by_tag") + if not isinstance(metadata_by_tag, dict): + raise ValueError(f"MolmoAct2 norm stats file {stats_path} has no metadata_by_tag mapping.") + metadata = metadata_by_tag.get(norm_tag) + if not isinstance(metadata, dict): + available = sorted(str(tag) for tag in metadata_by_tag) + raise ValueError(f"Unknown MolmoAct2 norm_tag={norm_tag!r}. Available tags: {available}.") + return metadata + + +def _apply_norm_tag_metadata(config: MolmoAct2Config) -> None: + """Populate config fields from the checkpoint's norm-tag metadata.""" + if not str(config.norm_tag or "").strip(): + return + metadata = _load_hf_norm_metadata_for_tag( + config.checkpoint_path, + revision=config.checkpoint_revision, + force_download=bool(config.checkpoint_force_download), + norm_tag=config.norm_tag, + ) + if metadata.get("action_horizon") is not None: + config.chunk_size = int(metadata["action_horizon"]) + if metadata.get("n_action_steps") is not None: + config.n_action_steps = int(metadata["n_action_steps"]) + if not config.setup_type and metadata.get("setup_type") is not None: + config.setup_type = str(metadata["setup_type"]) + if not config.control_mode and metadata.get("control_mode") is not None: + config.control_mode = str(metadata["control_mode"]) + + +def _saved_policy_action_mode(config: MolmoAct2Config) -> str | None: + """Read the action mode from a LeRobot-saved checkpoint's ``config.json``.""" + from pathlib import Path + + pretrained_path = getattr(config, "pretrained_path", None) + if pretrained_path is None: + return None + config_path = Path(pretrained_path) / "config.json" + if not config_path.exists(): + return None + try: + mode = json.loads(config_path.read_text()).get("action_mode") + except (OSError, json.JSONDecodeError): + return None + if mode in {"continuous", "discrete", "both"}: + return str(mode) + return None + + +def _training_action_mode(config: MolmoAct2Config, saved_policy_action_mode: str | None = None) -> str: + return saved_policy_action_mode or config.action_mode + + +def _validate_inference_action_mode( + config: MolmoAct2Config, saved_policy_action_mode: str | None = None +) -> None: + """Check that the requested inference mode is compatible with the training mode.""" + requested_mode = config.inference_action_mode + if requested_mode is None: + return + training_mode = _training_action_mode(config, saved_policy_action_mode) + if requested_mode == "continuous" and training_mode == "discrete": + raise ValueError( + "MolmoAct2 checkpoint was trained with action_mode='discrete' and cannot run " + "continuous inference." + ) + if requested_mode == "discrete" and training_mode == "continuous": + raise ValueError( + "MolmoAct2 checkpoint was trained with action_mode='continuous' and cannot run " + "discrete inference. Train with action_mode='both' or action_mode='discrete' first." + ) + + +def _validate_checkpoint_action_mode( + config: MolmoAct2Config, + checkpoint_action_mode: str, + *, + has_action_expert: bool, +) -> None: + """Check that the checkpoint's action mode is compatible with the config.""" + if config.action_mode == "both" and checkpoint_action_mode != "both": + raise ValueError( + f"action_mode='both' requires checkpoint action_mode='both', got {checkpoint_action_mode!r}." + ) + if config.action_mode == "discrete" and checkpoint_action_mode not in {"discrete", "both"}: + raise ValueError( + f"action_mode='discrete' requires checkpoint action_mode in {{'discrete', 'both'}}, " + f"got {checkpoint_action_mode!r}." + ) + if config.action_mode in {"continuous", "both"} and not has_action_expert: + raise ValueError("Continuous MolmoAct2 training requires an action expert checkpoint.") + + +def _resolve_inference_action_mode( + config: MolmoAct2Config, + requested_mode: str | None, + saved_policy_action_mode: str | None = None, +) -> str: + """Resolve the final inference action mode, validating compatibility.""" + training_mode = _training_action_mode(config, saved_policy_action_mode) + if requested_mode is None: + requested_mode = config.inference_action_mode + if requested_mode is None: + raise ValueError( + "MolmoAct2 inference requires `inference_action_mode` to be set explicitly " + "to either 'continuous' or 'discrete'." + ) + if requested_mode not in {"continuous", "discrete"}: + raise ValueError("MolmoAct2 inference_action_mode must be either 'continuous' or 'discrete'.") + if requested_mode == "continuous" and training_mode == "discrete": + raise ValueError("MolmoAct2 action_mode='discrete' checkpoint cannot run continuous inference.") + if requested_mode == "discrete" and training_mode == "continuous": + raise ValueError("MolmoAct2 action_mode='continuous' checkpoint cannot run discrete inference.") + return requested_mode + + def _strict_load_safetensors_weights(model: torch.nn.Module, checkpoint_location: str) -> None: index_path = os.path.join(checkpoint_location, SAFE_WEIGHTS_INDEX_NAME) single_file_path = os.path.join(checkpoint_location, SAFE_WEIGHTS_NAME) @@ -103,16 +309,6 @@ def _strict_load_safetensors_weights(model: torch.nn.Module, checkpoint_location ) -def _torch_dtype(dtype: str) -> torch.dtype: - if dtype == "float32": - return torch.float32 - if dtype == "bfloat16": - return torch.bfloat16 - if dtype == "float16": - return torch.float16 - raise ValueError(f"Unsupported dtype: {dtype}") - - def _sample_beta_timesteps( *, batch_size: int, @@ -136,7 +332,180 @@ def _sample_beta_timesteps( return time_offset + scale * samples +def _mask_discrete_action_spans( + *, + input_ids: Tensor, + mask: Tensor, + start_token_id: int | None, + end_token_id: int | None, +) -> Tensor: + if start_token_id is None or end_token_id is None: + return mask + mask = mask.clone() + for batch_idx in range(input_ids.shape[0]): + row = input_ids[batch_idx] + starts = (row == int(start_token_id)).nonzero(as_tuple=False).flatten().tolist() + ends = (row == int(end_token_id)).nonzero(as_tuple=False).flatten().tolist() + end_ptr = 0 + for start in starts: + while end_ptr < len(ends) and ends[end_ptr] < start: + end_ptr += 1 + if end_ptr >= len(ends): + mask[batch_idx, start:] = False + break + end = int(ends[end_ptr]) + mask[batch_idx, start : end + 1] = False + end_ptr += 1 + return mask + + +def _drop_trivial_attention_mask(model_inputs: dict[str, Tensor]) -> dict[str, Tensor]: + attention_mask = model_inputs.get("attention_mask") + if torch.is_tensor(attention_mask) and bool(attention_mask.to(dtype=torch.bool).all().item()): + model_inputs = dict(model_inputs) + model_inputs.pop("attention_mask", None) + return model_inputs + + +def _expand_mask(mask: Tensor | None, num_flow_timesteps: int) -> Tensor | None: + if mask is None: + return None + return ( + mask.unsqueeze(1) + .expand(-1, num_flow_timesteps, *([-1] * (mask.ndim - 1))) + .reshape(mask.shape[0] * num_flow_timesteps, *mask.shape[1:]) + ) + + +def _action_dim_valid_mask(target: Tensor, action_dim_is_pad: Tensor | None) -> Tensor | None: + if action_dim_is_pad is None: + return None + mask = ~action_dim_is_pad.to(device=target.device, dtype=torch.bool) + if mask.ndim == 1: + mask = mask.unsqueeze(0) + if mask.shape[-1] != target.shape[-1]: + raise ValueError( + f"action_dim_is_pad width {mask.shape[-1]} does not match target width {target.shape[-1]}." + ) + if mask.shape[0] == 1 and target.shape[0] != 1: + mask = mask.expand(target.shape[0], -1) + if mask.shape[0] != target.shape[0]: + raise ValueError( + f"action_dim_is_pad batch {mask.shape[0]} does not match target batch {target.shape[0]}." + ) + while mask.ndim < target.ndim: + mask = mask.unsqueeze(1) + return mask + + +def _mask_action_dim_tensor(tensor: Tensor, action_dim_is_pad: Tensor | None) -> Tensor: + if action_dim_is_pad is None: + return tensor + valid_mask = _action_dim_valid_mask(tensor, action_dim_is_pad) + if valid_mask is None: + return tensor + return tensor.masked_fill(~valid_mask, 0) + + +def _apply_action_dim_padding_mask(loss: Tensor, action_dim_is_pad: Tensor | None) -> Tensor: + valid_mask = _action_dim_valid_mask(loss, action_dim_is_pad) + if valid_mask is None: + return loss + valid = valid_mask.to(dtype=loss.dtype) + denom = valid.sum(dim=-1).clamp_min(1.0) + return (loss * valid).sum(dim=-1) / denom + + +def _apply_action_chunk_padding_mask(loss: Tensor, action_horizon_is_pad: Tensor | None) -> Tensor: + if action_horizon_is_pad is None: + return loss + valid_action = ( + (~action_horizon_is_pad.to(device=loss.device, dtype=torch.bool)).unsqueeze(1).unsqueeze(-1) + ) + return loss * valid_action + + +def _combine_rollout_seeds(first_seed: int, batch_size: int) -> int: + seed = 0 + for idx in range(batch_size): + seed = (seed + (idx + 1) * (first_seed + idx)) % (2**63 - 1) + return seed + + +def _rollout_task_signature(batch: dict[str, Any]) -> tuple[Any, ...] | None: + task = batch.get("task") + if task is None: + task = batch.get("observation.language") + if task is None: + return None + if isinstance(task, str): + return (task,) + if isinstance(task, (list, tuple)): + return tuple(str(item) for item in task) + return (str(task),) + + +def _extract_discrete_token_bins( + generated_ids: list[int], + start_token_id: int, + end_token_id: int, + token_id_to_bin: dict[int, int], +) -> list[int]: + start_idx = None + end_idx = None + for idx, token_id in enumerate(generated_ids): + if token_id == start_token_id: + start_idx = idx + break + if start_idx is not None: + for idx in range(start_idx + 1, len(generated_ids)): + if generated_ids[idx] == end_token_id: + end_idx = idx + break + span_start = 0 if start_idx is None else start_idx + 1 + span_end = len(generated_ids) if end_idx is None else end_idx + return [ + int(token_id_to_bin[token_id]) + for token_id in generated_ids[span_start:span_end] + if token_id in token_id_to_bin + ] + + +def _weighted_mean(values: Tensor, weights: Tensor | None) -> Tensor: + if weights is None: + return values.mean() + weights = weights.to(device=values.device, dtype=values.dtype) + return torch.dot(values, weights) / weights.sum().clamp_min(1.0) + + +def _weighted_per_example( + values: Tensor, + weights: Tensor | None, + example_indices: Tensor, + batch_size: int, +) -> Tensor: + values = values.float() + if weights is None: + weights = torch.ones_like(values) + else: + weights = weights.to(device=values.device, dtype=values.dtype) + loss_sum = torch.zeros(batch_size, device=values.device, dtype=torch.float32) + weight_sum = torch.zeros(batch_size, device=values.device, dtype=torch.float32) + loss_sum.scatter_add_(0, example_indices, values * weights) + weight_sum.scatter_add_(0, example_indices, weights) + global_weight_sum = weight_sum.sum().clamp_min(1.0) + return loss_sum * float(batch_size) / global_weight_sum + + class MolmoAct2Policy(PreTrainedPolicy): + """MolmoAct2 policy wrapping the vendored HF model for LeRobot. + + Supports three training modes via ``config.action_mode``: + ``"continuous"`` (flow-matching only), ``"discrete"`` (autoregressive + token prediction only), or ``"both"`` (joint loss). At inference, + ``config.inference_action_mode`` selects which head generates actions. + """ + config_class = MolmoAct2Config name = "molmoact2" @@ -149,10 +518,10 @@ class MolmoAct2Policy(PreTrainedPolicy): **kwargs, ): super().__init__(config, *inputs, **kwargs) - self.config.apply_norm_tag_metadata() + _apply_norm_tag_metadata(self.config) self.config.validate_features() del inputs, kwargs, dataset_stats, dataset_meta - self._checkpoint_action_mode = self.config.saved_policy_action_mode() + self._checkpoint_action_mode = _saved_policy_action_mode(self.config) self._action_queue: deque[Tensor] = deque(maxlen=self.config.n_action_steps) self._rollout_action_generator: torch.Generator | None = None self._rollout_task_key: tuple[Any, ...] | None = None @@ -160,7 +529,7 @@ class MolmoAct2Policy(PreTrainedPolicy): self.rtc_processor: RTCProcessor | None = None self.action_tokenizer: Any | None = None self._load_hf_model() - self.config.validate_inference_action_mode(self._checkpoint_action_mode) + _validate_inference_action_mode(self.config, self._checkpoint_action_mode) if self.config.enable_lora_vlm: self._apply_lora_adapters() self.init_rtc_processor() @@ -212,7 +581,8 @@ class MolmoAct2Policy(PreTrainedPolicy): "`policy.checkpoint_force_download=true` after the updated files are pushed." ) checkpoint_action_mode = str(self.model.config.action_mode) - self.config.validate_checkpoint_action_mode( + _validate_checkpoint_action_mode( + self.config, checkpoint_action_mode, has_action_expert=bool(getattr(self.model.config, "add_action_expert", False)), ) @@ -226,6 +596,7 @@ class MolmoAct2Policy(PreTrainedPolicy): self.train(self.training) def reset(self) -> None: + """Clear the action queue and rollout generator between episodes.""" self._action_queue = deque(maxlen=self.config.n_action_steps) self._rollout_action_generator = None @@ -334,6 +705,7 @@ class MolmoAct2Policy(PreTrainedPolicy): param.requires_grad = False def get_optim_params(self) -> list[dict[str, Any]]: + """Return optimizer param groups with per-component learning rates.""" vit_params: list[Tensor] = [] connector_params: list[Tensor] = [] action_expert_params: list[Tensor] = [] @@ -419,33 +791,6 @@ class MolmoAct2Policy(PreTrainedPolicy): return int(value) raise RuntimeError("MolmoAct2 could not resolve an action generation horizon.") - @staticmethod - def _mask_discrete_action_spans( - *, - input_ids: Tensor, - mask: Tensor, - start_token_id: int | None, - end_token_id: int | None, - ) -> Tensor: - if start_token_id is None or end_token_id is None: - return mask - mask = mask.clone() - for batch_idx in range(input_ids.shape[0]): - row = input_ids[batch_idx] - starts = (row == int(start_token_id)).nonzero(as_tuple=False).flatten().tolist() - ends = (row == int(end_token_id)).nonzero(as_tuple=False).flatten().tolist() - end_ptr = 0 - for start in starts: - while end_ptr < len(ends) and ends[end_ptr] < start: - end_ptr += 1 - if end_ptr >= len(ends): - mask[batch_idx, start:] = False - break - end = int(ends[end_ptr]) - mask[batch_idx, start : end + 1] = False - end_ptr += 1 - return mask - def _encoder_attention_mask_for_action_expert( self, *, @@ -470,21 +815,13 @@ class MolmoAct2Policy(PreTrainedPolicy): eos_token_id = getattr(self.model.config, "eos_token_id", None) if eos_token_id is not None: mask &= input_ids != int(eos_token_id) - return self._mask_discrete_action_spans( + return _mask_discrete_action_spans( input_ids=input_ids, mask=mask, start_token_id=getattr(self.model.config, "action_start_token_id", None), end_token_id=getattr(self.model.config, "action_end_token_id", None), ) - @staticmethod - def _drop_trivial_attention_mask(model_inputs: dict[str, Tensor]) -> dict[str, Tensor]: - attention_mask = model_inputs.get("attention_mask") - if torch.is_tensor(attention_mask) and bool(attention_mask.to(dtype=torch.bool).all().item()): - model_inputs = dict(model_inputs) - model_inputs.pop("attention_mask", None) - return model_inputs - def _load_discrete_action_tokenizer(self) -> Any: if self.action_tokenizer is None: require_package("transformers", extra="molmoact2") @@ -498,27 +835,7 @@ class MolmoAct2Policy(PreTrainedPolicy): return self.action_tokenizer def _resolve_inference_action_mode(self, requested_mode: str | None) -> str: - return self.config.resolve_inference_action_mode(requested_mode, self._checkpoint_action_mode) - - @staticmethod - def _combine_rollout_seeds(first_seed: int, batch_size: int) -> int: - seed = 0 - for idx in range(batch_size): - seed = (seed + (idx + 1) * (first_seed + idx)) % (2**63 - 1) - return seed - - @staticmethod - def _rollout_task_signature(batch: dict[str, Any]) -> tuple[Any, ...] | None: - task = batch.get("task") - if task is None: - task = batch.get("observation.language") - if task is None: - return None - if isinstance(task, str): - return (task,) - if isinstance(task, (list, tuple)): - return tuple(str(item) for item in task) - return (str(task),) + return _resolve_inference_action_mode(self.config, requested_mode, self._checkpoint_action_mode) def _rollout_generator_for_inputs( self, @@ -532,7 +849,7 @@ class MolmoAct2Policy(PreTrainedPolicy): if self._rollout_action_generator is not None: return self._rollout_action_generator - task_signature = self._rollout_task_signature(batch) + task_signature = _rollout_task_signature(batch) if task_signature != self._rollout_task_key: self._rollout_task_key = task_signature self._rollout_index_for_task = 0 @@ -545,72 +862,10 @@ class MolmoAct2Policy(PreTrainedPolicy): device if device.type == "cuda" and torch.cuda.is_available() else torch.device("cpu") ) generator = torch.Generator(device=generator_device) - generator.manual_seed(self._combine_rollout_seeds(first_seed, batch_size)) + generator.manual_seed(_combine_rollout_seeds(first_seed, batch_size)) self._rollout_action_generator = generator return generator - @staticmethod - def _expand_mask(mask: Tensor | None, num_flow_timesteps: int) -> Tensor | None: - if mask is None: - return None - return ( - mask.unsqueeze(1) - .expand(-1, num_flow_timesteps, *([-1] * (mask.ndim - 1))) - .reshape(mask.shape[0] * num_flow_timesteps, *mask.shape[1:]) - ) - - @staticmethod - def _action_dim_valid_mask(target: Tensor, action_dim_is_pad: Tensor | None) -> Tensor | None: - if action_dim_is_pad is None: - return None - mask = ~action_dim_is_pad.to(device=target.device, dtype=torch.bool) - if mask.ndim == 1: - mask = mask.unsqueeze(0) - if mask.shape[-1] != target.shape[-1]: - raise ValueError( - f"action_dim_is_pad width {mask.shape[-1]} does not match target width {target.shape[-1]}." - ) - if mask.shape[0] == 1 and target.shape[0] != 1: - mask = mask.expand(target.shape[0], -1) - if mask.shape[0] != target.shape[0]: - raise ValueError( - f"action_dim_is_pad batch {mask.shape[0]} does not match target batch {target.shape[0]}." - ) - while mask.ndim < target.ndim: - mask = mask.unsqueeze(1) - return mask - - @classmethod - def _mask_action_dim_tensor(cls, tensor: Tensor, action_dim_is_pad: Tensor | None) -> Tensor: - if not cls._mask_enabled_static(action_dim_is_pad): - return tensor - valid_mask = cls._action_dim_valid_mask(tensor, action_dim_is_pad) - if valid_mask is None: - return tensor - return tensor.masked_fill(~valid_mask, 0) - - @staticmethod - def _mask_enabled_static(action_dim_is_pad: Tensor | None) -> bool: - return action_dim_is_pad is not None - - @classmethod - def _apply_action_dim_padding_mask(cls, loss: Tensor, action_dim_is_pad: Tensor | None) -> Tensor: - valid_mask = cls._action_dim_valid_mask(loss, action_dim_is_pad) - if valid_mask is None: - return loss - valid = valid_mask.to(dtype=loss.dtype) - denom = valid.sum(dim=-1).clamp_min(1.0) - return (loss * valid).sum(dim=-1) / denom - - @staticmethod - def _apply_action_chunk_padding_mask(loss: Tensor, action_horizon_is_pad: Tensor | None) -> Tensor: - if action_horizon_is_pad is None: - return loss - valid_action = ( - (~action_horizon_is_pad.to(device=loss.device, dtype=torch.bool)).unsqueeze(1).unsqueeze(-1) - ) - return loss * valid_action - def _prepare_flow_matching_tensors( self, *, @@ -649,7 +904,7 @@ class MolmoAct2Policy(PreTrainedPolicy): ) if self.config.mask_action_dim_padding: - actions = self._mask_action_dim_tensor(actions, action_dim_is_pad) + actions = _mask_action_dim_tensor(actions, action_dim_is_pad) expected_noise_shape = (batch_size, num_flow_timesteps, actions.shape[1], actions.shape[2]) if noise is None: @@ -661,7 +916,7 @@ class MolmoAct2Policy(PreTrainedPolicy): f"flow noise must have shape {expected_noise_shape}, got {tuple(noise.shape)}." ) if self.config.mask_action_dim_padding: - noise = self._mask_action_dim_tensor(noise, action_dim_is_pad) + noise = _mask_action_dim_tensor(noise, action_dim_is_pad) t_broadcast = timesteps.view(batch_size, num_flow_timesteps, 1, 1) actions_expanded = actions.unsqueeze(1).expand(-1, num_flow_timesteps, -1, -1) @@ -789,7 +1044,7 @@ class MolmoAct2Policy(PreTrainedPolicy): valid_action = None if action_attention_mask is not None: valid_action = action_attention_mask.to(device=device, dtype=actions.dtype).unsqueeze(-1) - valid_action = self._expand_mask(valid_action, num_flow_timesteps) + valid_action = _expand_mask(valid_action, num_flow_timesteps) rope_cache = None if len(action_expert.blocks) > 0 and action_expert.blocks[0].self_attn.rope is not None: @@ -804,14 +1059,14 @@ class MolmoAct2Policy(PreTrainedPolicy): batch_size, actions.dtype, ) - cross_mask = self._expand_mask(cross_mask, num_flow_timesteps) + cross_mask = _expand_mask(cross_mask, num_flow_timesteps) self_mask = action_expert._build_self_attention_mask( action_attention_mask, actions.shape[1], device, actions.dtype, ) - self_mask = self._expand_mask(self_mask, num_flow_timesteps) + self_mask = _expand_mask(self_mask, num_flow_timesteps) conditioning = self._action_time_conditioning(action_expert, timesteps_flat) action_hidden = action_expert.action_embed(xt_flat) @@ -871,8 +1126,8 @@ class MolmoAct2Policy(PreTrainedPolicy): if k_norm is not None: k_ctx = k_norm(k_ctx.transpose(1, 2)).transpose(1, 2) if num_flow_timesteps != 1: - k_ctx = self._expand_mask(k_ctx, num_flow_timesteps) - v_ctx = self._expand_mask(v_ctx, num_flow_timesteps) + k_ctx = _expand_mask(k_ctx, num_flow_timesteps) + v_ctx = _expand_mask(v_ctx, num_flow_timesteps) next_action_hidden = action_block( layer_action_hidden, @@ -912,9 +1167,9 @@ class MolmoAct2Policy(PreTrainedPolicy): ) loss = F.mse_loss(pred_velocity, target_velocity, reduction="none") - loss = self._apply_action_chunk_padding_mask(loss, batch.get("action_horizon_is_pad")) + loss = _apply_action_chunk_padding_mask(loss, batch.get("action_horizon_is_pad")) if self.config.mask_action_dim_padding: - loss = self._apply_action_dim_padding_mask(loss, batch.get("action_dim_is_pad")) + loss = _apply_action_dim_padding_mask(loss, batch.get("action_dim_is_pad")) loss = loss.reshape(batch_size, -1).mean(dim=1) if reduction == "mean": loss = loss.mean() @@ -933,32 +1188,6 @@ class MolmoAct2Policy(PreTrainedPolicy): example_weights[nonempty] = 2.0 / torch.sqrt(token_counts[nonempty]) return example_weights[:, None].expand_as(valid_positions)[valid_positions].to(dtype=torch.float32) - @staticmethod - def _weighted_mean(values: Tensor, weights: Tensor | None) -> Tensor: - if weights is None: - return values.mean() - weights = weights.to(device=values.device, dtype=values.dtype) - return torch.dot(values, weights) / weights.sum().clamp_min(1.0) - - @staticmethod - def _weighted_per_example( - values: Tensor, - weights: Tensor | None, - example_indices: Tensor, - batch_size: int, - ) -> Tensor: - values = values.float() - if weights is None: - weights = torch.ones_like(values) - else: - weights = weights.to(device=values.device, dtype=values.dtype) - loss_sum = torch.zeros(batch_size, device=values.device, dtype=torch.float32) - weight_sum = torch.zeros(batch_size, device=values.device, dtype=torch.float32) - loss_sum.scatter_add_(0, example_indices, values * weights) - weight_sum.scatter_add_(0, example_indices, weights) - global_weight_sum = weight_sum.sum().clamp_min(1.0) - return loss_sum * float(batch_size) / global_weight_sum - def _discrete_loss_from_backbone_outputs( self, batch: dict[str, Tensor], @@ -992,56 +1221,28 @@ class MolmoAct2Policy(PreTrainedPolicy): token_weights = self._discrete_token_weights(valid_positions) if reduction == "none": example_indices = valid_positions.nonzero(as_tuple=False)[:, 0].to(device=hidden_states.device) - ce_loss = self._weighted_per_example( + ce_loss = _weighted_per_example( token_ce_loss, token_weights, example_indices, int(labels.shape[0]), ) else: - ce_loss = self._weighted_mean(token_ce_loss, token_weights) + ce_loss = _weighted_mean(token_ce_loss, token_weights) if not self.config.softmax_auxiliary_loss: return ce_loss, None if reduction == "none": - z_loss = self.config.softmax_auxiliary_loss_scale * self._weighted_per_example( + z_loss = self.config.softmax_auxiliary_loss_scale * _weighted_per_example( log_z.pow(2), token_weights, example_indices, int(labels.shape[0]), ) else: - z_loss = self.config.softmax_auxiliary_loss_scale * self._weighted_mean( - log_z.pow(2), token_weights - ) + z_loss = self.config.softmax_auxiliary_loss_scale * _weighted_mean(log_z.pow(2), token_weights) return ce_loss, z_loss - @staticmethod - def _extract_discrete_token_bins( - generated_ids: list[int], - start_token_id: int, - end_token_id: int, - token_id_to_bin: dict[int, int], - ) -> list[int]: - start_idx = None - end_idx = None - for idx, token_id in enumerate(generated_ids): - if token_id == start_token_id: - start_idx = idx - break - if start_idx is not None: - for idx in range(start_idx + 1, len(generated_ids)): - if generated_ids[idx] == end_token_id: - end_idx = idx - break - span_start = 0 if start_idx is None else start_idx + 1 - span_end = len(generated_ids) if end_idx is None else end_idx - return [ - int(token_id_to_bin[token_id]) - for token_id in generated_ids[span_start:span_end] - if token_id in token_id_to_bin - ] - def _action_token_id_to_bin(self) -> dict[int, int]: method = getattr(self.model, "_action_token_id_to_bin", None) if callable(method): @@ -1179,7 +1380,7 @@ class MolmoAct2Policy(PreTrainedPolicy): chunks: list[Tensor] = [] for token_row in generated_token_ids: generated_ids = [int(token_id) for token_id in token_row.detach().cpu().tolist()] - discrete_token_ids = self._extract_discrete_token_bins( + discrete_token_ids = _extract_discrete_token_bins( generated_ids, int(self.model.config.action_start_token_id), int(self.model.config.action_end_token_id), @@ -1218,7 +1419,7 @@ class MolmoAct2Policy(PreTrainedPolicy): model_inputs: dict[str, Tensor], action_dim: int, ) -> Tensor: - model_inputs = self._drop_trivial_attention_mask(model_inputs) + model_inputs = _drop_trivial_attention_mask(model_inputs) max_steps = self._discrete_generation_max_steps() static_cache, attention_bias = self._make_discrete_ar_graph_decode_inputs( model_inputs, @@ -1294,7 +1495,7 @@ class MolmoAct2Policy(PreTrainedPolicy): generator=generator, ) if self.config.mask_action_dim_padding: - trajectory = self._mask_action_dim_tensor(trajectory, action_dim_is_pad) + trajectory = _mask_action_dim_tensor(trajectory, action_dim_is_pad) action_context = action_expert.prepare_context( encoder_kv_states=encoder_kv_states, @@ -1327,7 +1528,7 @@ class MolmoAct2Policy(PreTrainedPolicy): modulation=step_modulation, ) if mask_enabled: - velocity = self._mask_action_dim_tensor(velocity, action_dim_is_pad) + velocity = _mask_action_dim_tensor(velocity, action_dim_is_pad) return velocity if self._rtc_enabled(): @@ -1352,7 +1553,7 @@ class MolmoAct2Policy(PreTrainedPolicy): trajectory = trajectory + dt * velocity if mask_enabled: - trajectory = self._mask_action_dim_tensor(trajectory, action_dim_is_pad) + trajectory = _mask_action_dim_tensor(trajectory, action_dim_is_pad) if self.rtc_processor is not None and self.rtc_processor.is_debug_enabled(): self.rtc_processor.track(time=float(flow_timestep[0].item()), x_t=trajectory, v_t=velocity) @@ -1363,6 +1564,7 @@ class MolmoAct2Policy(PreTrainedPolicy): batch: dict[str, Tensor], reduction: str = "mean", ) -> tuple[Tensor, dict[str, Any]]: + """Compute training loss (flow-matching and/or discrete token loss).""" if reduction not in {"mean", "none"}: raise ValueError(f"Unsupported reduction={reduction!r}. Expected 'mean' or 'none'.") model_inputs = self._model_inputs(batch) @@ -1422,6 +1624,7 @@ class MolmoAct2Policy(PreTrainedPolicy): @torch.no_grad() def predict_action_chunk(self, batch: dict[str, Tensor], **kwargs) -> Tensor: + """Generate an action chunk via continuous flow matching or discrete AR decoding.""" if "action_mode" in kwargs: raise TypeError( "MolmoAct2 predict_action_chunk got unexpected keyword argument 'action_mode'; " @@ -1476,6 +1679,7 @@ class MolmoAct2Policy(PreTrainedPolicy): @torch.no_grad() def select_action(self, batch: dict[str, Tensor], **kwargs) -> Tensor: + """Pop one action step from the queue, regenerating the chunk when empty.""" if self._rtc_enabled(): raise AssertionError("RTC is not supported for select_action, use it with predict_action_chunk") self.eval() diff --git a/src/lerobot/policies/molmoact2/hf_model/__init__.py b/src/lerobot/policies/molmoact2/molmoact2_hf_model/__init__.py similarity index 94% rename from src/lerobot/policies/molmoact2/hf_model/__init__.py rename to src/lerobot/policies/molmoact2/molmoact2_hf_model/__init__.py index 39b15cb3a..4436c9fda 100644 --- a/src/lerobot/policies/molmoact2/hf_model/__init__.py +++ b/src/lerobot/policies/molmoact2/molmoact2_hf_model/__init__.py @@ -1,5 +1,3 @@ -#!/usr/bin/env python - # Copyright 2026 The Allen Institute for Artificial Intelligence and The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -13,5 +11,3 @@ # 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. - -# ruff: noqa diff --git a/src/lerobot/policies/molmoact2/hf_model/action_tokenizer.py b/src/lerobot/policies/molmoact2/molmoact2_hf_model/action_tokenizer.py similarity index 96% rename from src/lerobot/policies/molmoact2/hf_model/action_tokenizer.py rename to src/lerobot/policies/molmoact2/molmoact2_hf_model/action_tokenizer.py index f7dacbce6..11a228731 100644 --- a/src/lerobot/policies/molmoact2/hf_model/action_tokenizer.py +++ b/src/lerobot/policies/molmoact2/molmoact2_hf_model/action_tokenizer.py @@ -1,5 +1,3 @@ -#!/usr/bin/env python - # Copyright 2026 The Allen Institute for Artificial Intelligence and The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -14,23 +12,19 @@ # See the License for the specific language governing permissions and # limitations under the License. -# ruff: noqa - import logging -import os from pathlib import Path from typing import ClassVar import numpy as np from tokenizers import ByteLevelBPETokenizer from tokenizers.trainers import BpeTrainer -from huggingface_hub import snapshot_download from transformers import PreTrainedTokenizerFast from transformers.processing_utils import ProcessorMixin +from ..modeling_molmoact2 import _hf_token -def _hf_token() -> str | None: - return os.environ.get("HF_TOKEN") or os.environ.get("HF_ACCESS_TOKEN") +logger = logging.getLogger(__name__) def _resolve_tokenizer_location( @@ -42,6 +36,8 @@ def _resolve_tokenizer_location( local_path = Path(str(tokenizer_path)).expanduser() if local_path.exists(): return str(local_path) + from huggingface_hub import snapshot_download + return snapshot_download( repo_id=str(tokenizer_path), repo_type="model", @@ -134,9 +130,8 @@ class UniversalActionProcessor(ProcessorMixin): ), ( f"Decoded DCT coefficients have shape {decoded_dct_coeff.shape}, expected ({self.time_horizon}, {self.action_dim})" ) - except Exception as e: - print(f"Error decoding tokens: {e}") - print(f"Tokens: {token}") + except Exception: + logger.warning("Error decoding tokens: %s", token, exc_info=True) decoded_dct_coeff = np.zeros((self.time_horizon, self.action_dim)) decoded_actions.append(idct(decoded_dct_coeff / self.scale, axis=0, norm="ortho")) return np.stack(decoded_actions) diff --git a/src/lerobot/policies/molmoact2/hf_model/configuration_molmoact2.py b/src/lerobot/policies/molmoact2/molmoact2_hf_model/configuration_molmoact2.py similarity index 99% rename from src/lerobot/policies/molmoact2/hf_model/configuration_molmoact2.py rename to src/lerobot/policies/molmoact2/molmoact2_hf_model/configuration_molmoact2.py index 29da68c14..df5449bef 100644 --- a/src/lerobot/policies/molmoact2/hf_model/configuration_molmoact2.py +++ b/src/lerobot/policies/molmoact2/molmoact2_hf_model/configuration_molmoact2.py @@ -1,5 +1,3 @@ -#!/usr/bin/env python - # Copyright 2026 The Allen Institute for Artificial Intelligence and The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -14,13 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -# ruff: noqa """ MolmoAct2 configuration """ -from typing import Optional, Any +from typing import Any from transformers import PretrainedConfig from transformers.modeling_rope_utils import rope_config_validation diff --git a/src/lerobot/policies/molmoact2/hf_model/image_processing_molmoact2.py b/src/lerobot/policies/molmoact2/molmoact2_hf_model/image_processing_molmoact2.py similarity index 98% rename from src/lerobot/policies/molmoact2/hf_model/image_processing_molmoact2.py rename to src/lerobot/policies/molmoact2/molmoact2_hf_model/image_processing_molmoact2.py index a172c8477..acc709cb5 100644 --- a/src/lerobot/policies/molmoact2/hf_model/image_processing_molmoact2.py +++ b/src/lerobot/policies/molmoact2/molmoact2_hf_model/image_processing_molmoact2.py @@ -1,5 +1,3 @@ -#!/usr/bin/env python - # Copyright 2026 The Allen Institute for Artificial Intelligence and The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -14,33 +12,28 @@ # See the License for the specific language governing permissions and # limitations under the License. -# ruff: noqa """Image processor class for MolmoAct2""" -from typing import Optional, Union -import numpy as np import einops +import numpy as np import torch import torchvision.transforms - +from transformers.feature_extraction_utils import BatchFeature +from transformers.image_processing_utils import BaseImageProcessor, get_size_dict +from transformers.image_transforms import convert_to_rgb from transformers.image_utils import ( IMAGENET_STANDARD_MEAN, IMAGENET_STANDARD_STD, ImageInput, PILImageResampling, make_flat_list_of_images, - valid_images, to_numpy_array, + valid_images, ) -from transformers.image_transforms import convert_to_rgb from transformers.processing_utils import ImagesKwargs -from transformers.image_processing_utils import BaseImageProcessor, get_size_dict -from transformers.utils import logging -from transformers.feature_extraction_utils import BatchFeature from transformers.utils import TensorType, logging - logger = logging.get_logger(__name__) @@ -73,8 +66,8 @@ def resize_image( )(image) resized = torch.clip(resized, 0.0, 1.0).to(dtype) else: - assert image.dtype == torch.uint8, "SigLIP expects float images or uint8 images, but got {}".format( - image.dtype + assert image.dtype == torch.uint8, ( + f"SigLIP expects float images or uint8 images, but got {image.dtype}" ) in_min = 0.0 in_max = 255.0 @@ -96,7 +89,6 @@ def resize_image( def select_tiling(h, w, patch_size, max_num_crops): """Divide in image of size [w, h] in up to max_num_patches of size patch_size""" original_size = np.stack([h, w]) # [1, 2] - original_res = h * w tilings = [] for i in range(1, max_num_crops + 1): for j in range(1, max_num_crops + 1): @@ -406,13 +398,17 @@ class MolmoAct2ImageProcessor(BaseImageProcessor): image_std: float | list[float] | None = None, do_convert_rgb: bool = True, max_crops: int = 8, - overlap_margins: list[int] = [4, 4], + overlap_margins: list[int] | None = None, crop_mode: str = "overlap-and-resize-c2", patch_size: int = 14, - pooling_size: list[int] = [2, 2], + pooling_size: list[int] | None = None, **kwargs, ) -> None: super().__init__(**kwargs) + if overlap_margins is None: + overlap_margins = [4, 4] + if pooling_size is None: + pooling_size = [2, 2] size = size if size is not None else {"height": 378, "width": 378} size = get_size_dict(size, default_to_square=True) self.size = size diff --git a/src/lerobot/policies/molmoact2/hf_model/inference.py b/src/lerobot/policies/molmoact2/molmoact2_hf_model/inference.py similarity index 99% rename from src/lerobot/policies/molmoact2/hf_model/inference.py rename to src/lerobot/policies/molmoact2/molmoact2_hf_model/inference.py index 2c0243880..428800a8c 100644 --- a/src/lerobot/policies/molmoact2/hf_model/inference.py +++ b/src/lerobot/policies/molmoact2/molmoact2_hf_model/inference.py @@ -1,5 +1,3 @@ -#!/usr/bin/env python - # Copyright 2026 The Allen Institute for Artificial Intelligence and The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -14,16 +12,15 @@ # See the License for the specific language governing permissions and # limitations under the License. -# ruff: noqa """Inference utilities for MolmoAct2""" -from dataclasses import dataclass -from typing import Any, Optional, Tuple from collections.abc import Iterable, Sequence +from dataclasses import dataclass +from typing import Any import torch -from torch.nn import functional as F +from torch.nn import functional as F # noqa: N812 from transformers.cache_utils import Cache from transformers.configuration_utils import PretrainedConfig @@ -679,7 +676,7 @@ def _clone_static_inputs(inputs: _ActionFlowInputs) -> _ActionFlowInputs: def _copy_context_(dst: Any, src: Any) -> None: - for (dst_k, dst_v), (src_k, src_v) in zip(dst.kv_contexts, src.kv_contexts): + for (dst_k, dst_v), (src_k, src_v) in zip(dst.kv_contexts, src.kv_contexts, strict=False): dst_k.copy_(src_k) dst_v.copy_(src_v) if src.cross_mask is not None: @@ -689,7 +686,7 @@ def _copy_context_(dst: Any, src: Any) -> None: if src.valid_action is not None: dst.valid_action.copy_(src.valid_action) if src.rope_cache is not None: - for dst_tensor, src_tensor in zip(dst.rope_cache, src.rope_cache): + for dst_tensor, src_tensor in zip(dst.rope_cache, src.rope_cache, strict=False): dst_tensor.copy_(src_tensor) diff --git a/src/lerobot/policies/molmoact2/hf_model/modeling_molmoact2.py b/src/lerobot/policies/molmoact2/molmoact2_hf_model/modeling_molmoact2.py similarity index 99% rename from src/lerobot/policies/molmoact2/hf_model/modeling_molmoact2.py rename to src/lerobot/policies/molmoact2/molmoact2_hf_model/modeling_molmoact2.py index 4c36b04c8..e2edbe68d 100644 --- a/src/lerobot/policies/molmoact2/hf_model/modeling_molmoact2.py +++ b/src/lerobot/policies/molmoact2/molmoact2_hf_model/modeling_molmoact2.py @@ -1,5 +1,3 @@ -#!/usr/bin/env python - # Copyright 2026 The Allen Institute for Artificial Intelligence and The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -14,24 +12,25 @@ # See the License for the specific language governing permissions and # limitations under the License. -# ruff: noqa """Modeling code for MolmoAct2""" +# ruff: noqa: N806 + import json import math import os import re +from collections.abc import Callable, Mapping, Sequence from copy import deepcopy from dataclasses import dataclass -from typing import Any, Dict, List, Optional, Tuple, Union -from collections.abc import Callable, Mapping, Sequence +from typing import Any, Optional import numpy as np import torch import torch.utils.checkpoint from torch import nn -from torch.nn import functional as F +from torch.nn import functional as F # noqa: N812 from torch.nn.attention import SDPBackend, sdpa_kernel from transformers.activations import ACT2FN from transformers.cache_utils import Cache, DynamicCache @@ -647,7 +646,7 @@ class ActionExpert(nn.Module): f"got {len(encoder_kv_states)}." ) kv_contexts = [] - for block, (k_in, v_in) in zip(self.blocks, encoder_kv_states): + for block, (k_in, v_in) in zip(self.blocks, encoder_kv_states, strict=False): k_ctx = self._project_kv_tensor(k_in, self.context_k_proj) v_ctx = self._project_kv_tensor(v_in, self.context_v_proj) k_norm = block.cross_attn.k_norm @@ -732,7 +731,7 @@ class ActionExpert(nn.Module): timesteps: Sequence[torch.Tensor], ) -> Sequence[ActionExpertStepModulation]: cache = [] - for idx, step_t in enumerate(timesteps): + for _idx, step_t in enumerate(timesteps): conditioning = self._time_conditioning(step_t) block_modulations = [] for block in self.blocks: @@ -786,8 +785,8 @@ class ActionExpert(nn.Module): x = self.action_embed(actions) if context.valid_action is not None: x = x * context.valid_action - for idx, (block, kv_context, block_modulation) in enumerate( - zip(self.blocks, context.kv_contexts, block_modulations) + for _idx, (block, kv_context, block_modulation) in enumerate( + zip(self.blocks, context.kv_contexts, block_modulations, strict=False) ): x = block( x, @@ -2874,7 +2873,7 @@ class MolmoAct2Model(MolmoAct2PreTrainedModel): depth_mask=depth_mask, encoder_attention_mask=encoder_attention_mask, ) - for gate, source in zip(gate_head, sources) + for gate, source in zip(gate_head, sources, strict=False) ] return gates, depth_mask gate = self._depth_gate_from_source( @@ -4458,7 +4457,7 @@ class MolmoAct2ForConditionalGeneration(MolmoAct2PreTrainedModel, GenerationMixi ```python >>> from PIL import Image >>> import requests - >>> from lerobot.policies.molmoact2.hf_model.modeling_molmoact2 import MolmoAct2ForConditionalGeneration + >>> from lerobot.policies.molmoact2.molmoact2_hf_model.modeling_molmoact2 import MolmoAct2ForConditionalGeneration >>> from lerobot.policies.molmoact2.processor_molmoact2 import _load_local_molmoact2_processor >>> model = MolmoAct2ForConditionalGeneration.from_pretrained("...") diff --git a/src/lerobot/policies/molmoact2/hf_model/processing_molmoact2.py b/src/lerobot/policies/molmoact2/molmoact2_hf_model/processing_molmoact2.py similarity index 95% rename from src/lerobot/policies/molmoact2/hf_model/processing_molmoact2.py rename to src/lerobot/policies/molmoact2/molmoact2_hf_model/processing_molmoact2.py index 7b8775faa..6a73d2465 100644 --- a/src/lerobot/policies/molmoact2/hf_model/processing_molmoact2.py +++ b/src/lerobot/policies/molmoact2/molmoact2_hf_model/processing_molmoact2.py @@ -1,5 +1,3 @@ -#!/usr/bin/env python - # Copyright 2026 The Allen Institute for Artificial Intelligence and The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -14,45 +12,39 @@ # See the License for the specific language governing permissions and # limitations under the License. -# ruff: noqa """ Processor class for MolmoAct2. """ -from typing import Optional, Union -import dataclasses - import numpy as np - +from transformers import AutoTokenizer +from transformers.feature_extraction_utils import BatchFeature from transformers.image_utils import ImageInput -from transformers.video_utils import VideoInput from transformers.processing_utils import ( - Unpack, ProcessingKwargs, ProcessorMixin, + Unpack, ) -from transformers.feature_extraction_utils import BatchFeature -from transformers.tokenization_utils_base import TextInput, PreTokenizedInput +from transformers.tokenization_utils_base import PreTokenizedInput, TextInput from transformers.utils import logging +from transformers.video_utils import VideoInput -from transformers import AutoTokenizer -from .image_processing_molmoact2 import MolmoAct2ImagesKwargs, MolmoAct2ImageProcessor -from .video_processing_molmoact2 import MolmoAct2VideoProcessorKwargs, MolmoAct2VideoProcessor - +from .image_processing_molmoact2 import MolmoAct2ImageProcessor, MolmoAct2ImagesKwargs +from .video_processing_molmoact2 import MolmoAct2VideoProcessor, MolmoAct2VideoProcessorKwargs logger = logging.get_logger(__name__) # Special tokens, these should be present in any tokenizer we use since the preprocessor uses them -IMAGE_PATCH_TOKEN = f"" # Where to insert high-res tokens -IMAGE_LOW_RES_TOKEN = f"" # Where to insert low-res tokens -IM_START_TOKEN = f"" -LOW_RES_IMAGE_START_TOKEN = f"" -FRAME_START_TOKEN = f"" -IM_END_TOKEN = f"" -FRAME_END_TOKEN = f"" -IM_COL_TOKEN = f"" +IMAGE_PATCH_TOKEN = "" # nosec B105 # Where to insert high-res tokens +IMAGE_LOW_RES_TOKEN = "" # nosec B105 # Where to insert low-res tokens +IM_START_TOKEN = "" # nosec B105 +LOW_RES_IMAGE_START_TOKEN = "" # nosec B105 +FRAME_START_TOKEN = "" # nosec B105 +IM_END_TOKEN = "" # nosec B105 +FRAME_END_TOKEN = "" # nosec B105 +IM_COL_TOKEN = "" # nosec B105 IMAGE_PROMPT = "<|image|>" VIDEO_PROMPT = "<|video|>" @@ -224,7 +216,7 @@ class MolmoAct2Processor(ProcessorMixin): input_ids = input_ids[None, :] attention_mask = attention_mask[None, :] - B, S = input_ids.shape + B, S = input_ids.shape # noqa: N806 # Handle zero-length sequence if S == 0: @@ -364,7 +356,7 @@ class MolmoAct2Processor(ProcessorMixin): assert num_videos in {0, 1}, "At most one video is supported for now" video_grids_i = video_grids[index : index + num_videos] metadata_i = video_metadata[index : index + num_videos] - for video_grid, metadata in zip(video_grids_i, metadata_i): + for video_grid, metadata in zip(video_grids_i, metadata_i, strict=False): video_string = self.get_video_string( video_grid, metadata.timestamps, diff --git a/src/lerobot/policies/molmoact2/hf_model/video_processing_molmoact2.py b/src/lerobot/policies/molmoact2/molmoact2_hf_model/video_processing_molmoact2.py similarity index 98% rename from src/lerobot/policies/molmoact2/hf_model/video_processing_molmoact2.py rename to src/lerobot/policies/molmoact2/molmoact2_hf_model/video_processing_molmoact2.py index 644d5a691..bf4e44dde 100644 --- a/src/lerobot/policies/molmoact2/hf_model/video_processing_molmoact2.py +++ b/src/lerobot/policies/molmoact2/molmoact2_hf_model/video_processing_molmoact2.py @@ -1,5 +1,3 @@ -#!/usr/bin/env python - # Copyright 2026 The Allen Institute for Artificial Intelligence and The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -14,25 +12,23 @@ # See the License for the specific language governing permissions and # limitations under the License. -# ruff: noqa """Video processor class for MolmoAct2""" -from functools import partial import os import warnings +from collections.abc import Callable from contextlib import redirect_stdout +from functools import partial from io import BytesIO from urllib.parse import urlparse -from typing import Optional, Union -from collections.abc import Callable +import einops import numpy as np import requests -import einops import torch import torchvision.transforms - +from transformers.feature_extraction_utils import BatchFeature from transformers.image_utils import ( IMAGENET_STANDARD_MEAN, IMAGENET_STANDARD_STD, @@ -41,27 +37,24 @@ from transformers.image_utils import ( SizeDict, validate_kwargs, ) -from transformers.video_utils import ( - VideoInput, - is_valid_video, - make_batched_videos, - make_batched_metadata, - VideoMetadata, -) from transformers.processing_utils import Unpack, VideosKwargs -from transformers.video_processing_utils import BaseVideoProcessor -from transformers.utils import logging -from transformers.feature_extraction_utils import BatchFeature from transformers.utils import ( + TensorType, is_av_available, is_decord_available, is_torchcodec_available, is_yt_dlp_available, - TensorType, logging, to_numpy, ) - +from transformers.video_processing_utils import BaseVideoProcessor +from transformers.video_utils import ( + VideoInput, + VideoMetadata, + is_valid_video, + make_batched_metadata, + make_batched_videos, +) logger = logging.get_logger(__name__) @@ -102,8 +95,8 @@ def resize_image( )(image) resized = torch.clip(resized, 0.0, 1.0).to(dtype) else: - assert image.dtype == torch.uint8, "SigLIP expects float images or uint8 images, but got {}".format( - image.dtype + assert image.dtype == torch.uint8, ( + f"SigLIP expects float images or uint8 images, but got {image.dtype}" ) in_min = 0.0 in_max = 255.0 @@ -548,9 +541,8 @@ def get_target_fps( step_size = max(int(video_fps / target_fps), 1) num_frames_sampled_at_fps = int(total_frames / step_size) if num_frames_sampled == 0: - if "uniform" in frame_sample_mode: - if num_frames_sampled_at_fps > max_frames: - break + if "uniform" in frame_sample_mode and num_frames_sampled_at_fps > max_frames: + break selected_target_fps = target_fps num_frames_sampled = num_frames_sampled_at_fps @@ -779,13 +771,15 @@ class MolmoAct2VideoProcessor(BaseVideoProcessor): elif is_torchcodec_available(): warnings.warn( "`decord` is not installed and cannot be used to decode the video by default. " - "Falling back to `torchcodec`." + "Falling back to `torchcodec`.", + stacklevel=2, ) backend = "torchcodec" else: warnings.warn( "`decord` is not installed and cannot be used to decode the video by default. " - "Falling back to `PyAV`." + "Falling back to `PyAV`.", + stacklevel=2, ) backend = "pyav" @@ -795,7 +789,8 @@ class MolmoAct2VideoProcessor(BaseVideoProcessor): *[ self.fetch_videos(x, sample_timestamps_fn=sample_timestamps_fn) for x in video_url_or_urls - ] + ], + strict=False, ) ) else: @@ -821,7 +816,7 @@ class MolmoAct2VideoProcessor(BaseVideoProcessor): assert video_metadata[0].fps is not None, "FPS must be provided for video input" sampled_videos = [] sampled_metadata = [] - for video, metadata in zip(videos, video_metadata): + for video, metadata in zip(videos, video_metadata, strict=False): indices = sample_indices_fn(metadata=metadata) metadata.frames_indices = indices sampled_videos.append(video[indices]) @@ -985,11 +980,11 @@ class MolmoAct2VideoProcessor(BaseVideoProcessor): pixel_values_videos = np.concatenate(batch_crops, 0) video_token_pooling = np.concatenate(batch_pooled_patches_idx, 0) - data = dict( - pixel_values_videos=pixel_values_videos, - video_token_pooling=video_token_pooling, - video_grids=video_grids, - ) + data = { + "pixel_values_videos": pixel_values_videos, + "video_token_pooling": video_token_pooling, + "video_grids": video_grids, + } return BatchFeature(data, tensor_type=return_tensors) diff --git a/src/lerobot/policies/molmoact2/processor_molmoact2.py b/src/lerobot/policies/molmoact2/processor_molmoact2.py index 6c7a3ed5c..1303e94a1 100644 --- a/src/lerobot/policies/molmoact2/processor_molmoact2.py +++ b/src/lerobot/policies/molmoact2/processor_molmoact2.py @@ -1,5 +1,3 @@ -#!/usr/bin/env python - # Copyright 2026 The Allen Institute for Artificial Intelligence and The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -14,10 +12,18 @@ # See the License for the specific language governing permissions and # limitations under the License. +"""MolmoAct2 pre/post processing pipeline. + +Builds the multimodal prompt (images, discretised state, task text), +tokenises it via the vendored MolmoAct2 processor, and handles quantile +normalisation with optional per-dimension gripper masking. +""" + from __future__ import annotations import json -import os +import logging +import math import re from contextlib import suppress from copy import deepcopy @@ -27,7 +33,6 @@ from typing import TYPE_CHECKING, Any import numpy as np import torch -from huggingface_hub import snapshot_download from torch import Tensor from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature @@ -54,14 +59,71 @@ from lerobot.utils.constants import ( ) from lerobot.utils.import_utils import _scipy_available, _transformers_available, require_package -from .configuration_molmoact2 import MolmoAct2Config, infer_molmoact2_max_sequence_length +from .configuration_molmoact2 import MolmoAct2Config +from .modeling_molmoact2 import _hf_token, _resolve_checkpoint_location + +logger = logging.getLogger(__name__) + +MOLMOACT2_DEFAULT_NUM_IMAGES = 2 +MOLMOACT2_IMAGE_TOKENS_PER_IMAGE = 196 +MOLMOACT2_FIXED_PROMPT_TOKEN_BUDGET = 80 +MOLMOACT2_TASK_TOKEN_BUDGET = 32 +MOLMOACT2_SEQUENCE_LENGTH_MARGIN = 32 +MOLMOACT2_SEQUENCE_LENGTH_MULTIPLE = 64 +MOLMOACT2_DISCRETE_ACTION_WRAPPER_TOKENS = 4 +MOLMOACT2_MIN_DISCRETE_ACTION_TOKENS_PER_STEP = 6 +MOLMOACT2_DISCRETE_ACTION_TOKENS_PER_DIM = 0.95 + + +def _round_up(value: int, multiple: int) -> int: + return int(math.ceil(value / multiple) * multiple) + + +def infer_molmoact2_max_sequence_length( + *, + num_images: int, + state_dim: int, + action_dim: int, + action_horizon: int, + include_discrete_action: bool, +) -> int: + """Infer the padded text/image sequence cap from MolmoAct2's fixed token layout.""" + if num_images < 1: + num_images = MOLMOACT2_DEFAULT_NUM_IMAGES + if state_dim < 0: + state_dim = 0 + if action_dim < 1: + action_dim = 1 + if action_horizon < 1: + action_horizon = 1 + + image_tokens = num_images * MOLMOACT2_IMAGE_TOKENS_PER_IMAGE + prompt_tokens = ( + MOLMOACT2_FIXED_PROMPT_TOKEN_BUDGET + + MOLMOACT2_TASK_TOKEN_BUDGET + + state_dim + + MOLMOACT2_SEQUENCE_LENGTH_MARGIN + ) + action_tokens = 0 + if include_discrete_action: + action_tokens_per_step = max( + MOLMOACT2_MIN_DISCRETE_ACTION_TOKENS_PER_STEP, + math.ceil(action_dim * MOLMOACT2_DISCRETE_ACTION_TOKENS_PER_DIM), + ) + action_tokens = MOLMOACT2_DISCRETE_ACTION_WRAPPER_TOKENS + action_horizon * action_tokens_per_step + + return _round_up( + image_tokens + prompt_tokens + action_tokens, + MOLMOACT2_SEQUENCE_LENGTH_MULTIPLE, + ) + if TYPE_CHECKING or _transformers_available: from transformers import Qwen2Tokenizer - from .hf_model.image_processing_molmoact2 import MolmoAct2ImageProcessor - from .hf_model.processing_molmoact2 import MolmoAct2Processor - from .hf_model.video_processing_molmoact2 import MolmoAct2VideoProcessor + from .molmoact2_hf_model.image_processing_molmoact2 import MolmoAct2ImageProcessor + from .molmoact2_hf_model.processing_molmoact2 import MolmoAct2Processor + from .molmoact2_hf_model.video_processing_molmoact2 import MolmoAct2VideoProcessor else: Qwen2Tokenizer = None MolmoAct2ImageProcessor = None @@ -69,7 +131,7 @@ else: MolmoAct2VideoProcessor = None if TYPE_CHECKING or (_transformers_available and _scipy_available): - from .hf_model.action_tokenizer import UniversalActionProcessor + from .molmoact2_hf_model.action_tokenizer import UniversalActionProcessor else: UniversalActionProcessor = None @@ -97,32 +159,6 @@ _QUESTION_PREFIX_PATTERNS = tuple( ) -def _hf_token() -> str | None: - return os.environ.get("HF_TOKEN") or os.environ.get("HF_ACCESS_TOKEN") - - -def _resolve_checkpoint_location( - checkpoint_path: str, - *, - revision: str | None = None, - force_download: bool = False, -) -> str: - checkpoint_path = str(checkpoint_path or "").strip() - if not checkpoint_path: - raise ValueError("MolmoAct2 policy requires `checkpoint_path`.") - local_path = Path(checkpoint_path).expanduser() - if local_path.exists(): - return str(local_path) - return snapshot_download( - repo_id=checkpoint_path, - repo_type="model", - revision=revision, - force_download=force_download, - ignore_patterns=["*.py", "*.pyc", "__pycache__/*"], - token=_hf_token(), - ) - - def _load_hf_norm_stats_for_tag( checkpoint_path: str, *, diff --git a/tests/policies/molmoact2/test_molmoact2.py b/tests/policies/molmoact2/test_molmoact2.py index 3631bcc9b..5fba72913 100644 --- a/tests/policies/molmoact2/test_molmoact2.py +++ b/tests/policies/molmoact2/test_molmoact2.py @@ -1,5 +1,3 @@ -#!/usr/bin/env python - # Copyright 2026 The Allen Institute for Artificial Intelligence and The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -35,16 +33,16 @@ pytest.importorskip("scipy") from lerobot.configs import FeatureType, NormalizationMode, PolicyFeature from lerobot.policies import get_policy_class, make_policy_config from lerobot.policies.molmoact2 import ( - configuration_molmoact2 as molmoact2_config, modeling_molmoact2 as molmoact2_modeling, processor_molmoact2 as molmoact2_processor, ) -from lerobot.policies.molmoact2.configuration_molmoact2 import ( - MolmoAct2Config, - MolmoAct2CosineDecayWithWarmupSchedulerConfig, - infer_molmoact2_max_sequence_length, +from lerobot.policies.molmoact2.configuration_molmoact2 import MolmoAct2Config +from lerobot.policies.molmoact2.modeling_molmoact2 import ( + MolmoAct2Policy, + _apply_action_chunk_padding_mask, + _apply_action_dim_padding_mask, + _combine_rollout_seeds, ) -from lerobot.policies.molmoact2.modeling_molmoact2 import MolmoAct2Policy from lerobot.policies.molmoact2.processor_molmoact2 import ( MolmoAct2ClampNormalizedProcessorStep, MolmoAct2MaskedNormalizerProcessorStep, @@ -53,6 +51,7 @@ from lerobot.policies.molmoact2.processor_molmoact2 import ( _add_gripper_masks_to_stats, _build_discrete_state_string, _normalize_question_text, + infer_molmoact2_max_sequence_length, make_molmoact2_pre_post_processors, ) from lerobot.policies.rtc.configuration_rtc import RTCConfig @@ -71,34 +70,38 @@ def test_molmoact2_policy_registration(): assert cfg.per_episode_seed is False assert cfg.eval_seed is None assert cfg.normalize_language is True - assert cfg.get_scheduler_preset().num_decay_steps is None + assert cfg.get_scheduler_preset().num_decay_steps == 100_000 assert cfg.action_delta_indices == list(range(cfg.chunk_size)) assert get_policy_class("molmoact2") is MolmoAct2Policy def test_molmoact2_checkpoint_download_ignores_remote_python(monkeypatch): + import huggingface_hub + download_kwargs = {} def fake_snapshot_download(**kwargs): download_kwargs.update(kwargs) return "/tmp/downloaded-molmoact2" - monkeypatch.setattr(molmoact2_config, "snapshot_download", fake_snapshot_download) + monkeypatch.setattr(huggingface_hub, "snapshot_download", fake_snapshot_download) - checkpoint_location = molmoact2_config._resolve_checkpoint_location("allenai/MolmoAct2") + checkpoint_location = molmoact2_modeling._resolve_checkpoint_location("allenai/MolmoAct2") assert checkpoint_location == "/tmp/downloaded-molmoact2" assert download_kwargs["ignore_patterns"] == ["*.py", "*.pyc", "__pycache__/*"] -def test_molmoact2_scheduler_decay_steps_auto_match_training_steps(): +def test_molmoact2_scheduler_auto_scales_to_training_steps(): + from lerobot.optim import CosineDecayWithWarmupSchedulerConfig + param = torch.nn.Parameter(torch.ones(())) optimizer = torch.optim.AdamW([param], lr=0.001) - config = MolmoAct2CosineDecayWithWarmupSchedulerConfig( + config = CosineDecayWithWarmupSchedulerConfig( peak_lr=0.01, decay_lr=0.001, num_warmup_steps=10, - num_decay_steps=None, + num_decay_steps=100_000, ) scheduler = config.build(optimizer, num_training_steps=100) @@ -123,9 +126,7 @@ def test_molmoact2_rollout_generator_uses_eval_seed_per_task(): batch_size=3, device=torch.device("cpu"), ) - expected_first = torch.Generator().manual_seed( - MolmoAct2Policy._combine_rollout_seeds(first_seed=1000, batch_size=3) - ) + expected_first = torch.Generator().manual_seed(_combine_rollout_seeds(first_seed=1000, batch_size=3)) assert torch.allclose(torch.rand(4, generator=first), torch.rand(4, generator=expected_first)) policy.reset() @@ -134,9 +135,7 @@ def test_molmoact2_rollout_generator_uses_eval_seed_per_task(): batch_size=3, device=torch.device("cpu"), ) - expected_second = torch.Generator().manual_seed( - MolmoAct2Policy._combine_rollout_seeds(first_seed=1003, batch_size=3) - ) + expected_second = torch.Generator().manual_seed(_combine_rollout_seeds(first_seed=1003, batch_size=3)) assert torch.allclose(torch.rand(4, generator=second), torch.rand(4, generator=expected_second)) policy.reset() @@ -145,9 +144,7 @@ def test_molmoact2_rollout_generator_uses_eval_seed_per_task(): batch_size=3, device=torch.device("cpu"), ) - expected_new_task = torch.Generator().manual_seed( - MolmoAct2Policy._combine_rollout_seeds(first_seed=1000, batch_size=3) - ) + expected_new_task = torch.Generator().manual_seed(_combine_rollout_seeds(first_seed=1000, batch_size=3)) assert torch.allclose(torch.rand(4, generator=new_task), torch.rand(4, generator=expected_new_task)) @@ -537,36 +534,26 @@ def test_train_action_expert_only_requires_continuous_action_mode(): def test_molmoact2_sequence_length_is_inferred_from_fixed_token_budget(): - cfg = MolmoAct2Config( - action_mode="both", - chunk_size=10, - n_action_steps=10, - image_keys=["observation.images.image", "observation.images.wrist_image"], - input_features={OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(8,))}, - output_features={ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(7,))}, - ) - - assert cfg.max_sequence_length is None - assert cfg.inferred_max_sequence_length() == 640 - assert cfg.inferred_max_sequence_length(include_discrete_action=False) == 576 assert ( infer_molmoact2_max_sequence_length( - num_images=2, - state_dim=8, - action_dim=7, - action_horizon=30, - include_discrete_action=True, + num_images=2, state_dim=8, action_dim=7, action_horizon=10, include_discrete_action=True + ) + == 640 + ) + assert ( + infer_molmoact2_max_sequence_length( + num_images=2, state_dim=8, action_dim=7, action_horizon=10, include_discrete_action=False + ) + == 576 + ) + assert ( + infer_molmoact2_max_sequence_length( + num_images=2, state_dim=8, action_dim=7, action_horizon=30, include_discrete_action=True ) == 768 ) -def test_molmoact2_sequence_length_override_is_preserved(): - cfg = MolmoAct2Config(max_sequence_length=1024) - - assert cfg.inferred_max_sequence_length(num_images=2, state_dim=8, action_dim=7) == 1024 - - def test_train_action_expert_only_freezes_non_action_expert_params(): class DummyBackbone(torch.nn.Module): def __init__(self): @@ -963,7 +950,7 @@ def test_action_dim_padding_loss_reduces_like_old_trainer(): ] ) - reduced = MolmoAct2Policy._apply_action_dim_padding_mask(loss, action_dim_is_pad) + reduced = _apply_action_dim_padding_mask(loss, action_dim_is_pad) expected = torch.stack( [ @@ -979,7 +966,7 @@ def test_action_chunk_padding_keeps_old_mean_denominator(): loss = torch.ones(1, 2, 4, 3) action_horizon_is_pad = torch.tensor([[False, False, True, True]]) - masked = MolmoAct2Policy._apply_action_chunk_padding_mask(loss, action_horizon_is_pad) + masked = _apply_action_chunk_padding_mask(loss, action_horizon_is_pad) assert masked.mean().item() == 0.5 From 6a788fbdb02cabfae60f7408636945df0b1eafa0 Mon Sep 17 00:00:00 2001 From: Khalil Meftah Date: Thu, 25 Jun 2026 15:31:24 +0200 Subject: [PATCH 17/38] Add inline offline validation with train/eval split (#3824) * refactor(training): rename eval_freq to env_eval_freq - Rename eval_freq to env_eval_freq to distinguish sim environment evaluation from offline loss evaluation. * feat(training): add inline offline validation with train/eval split - Add eval_split config for balanced per-task holdout - Add eval_steps for periodic inline eval loss computation - Add max_eval_samples to cap eval cost * fix(datasets): remap absolute indices in __getitem__ for filtered datasets * fix(train): vectorize eval subset selection for max_eval_samples * fix(datasets): Move the remapping into EpisodeAwareSampler via absolute_to_relative_idx * fix(validation): add eval_split range check and eval_steps warning Validate eval_split is in [0.0, 1.0) to prevent garbage splits from out-of-range values. Raise when eval_steps > 0 but eval_split is 0.0 since no offline eval will run. * fix(train): prepare eval dataloader with accelerator for multi-GPU Prepare eval_dataloader through accelerator.prepare() so eval data is sharded across ranks instead of duplicated. Reduce eval_loss across ranks with mean reduction for consistent logging. * fix(test): rename eval_freq to env_eval_freq for multi-GPU training --- .github/workflows/benchmark_tests.yml | 6 +- Makefile | 8 +-- docs/source/hilserl.mdx | 2 +- docs/source/libero.mdx | 2 +- docs/source/libero_plus.mdx | 2 +- docs/source/metaworld.mdx | 4 +- docs/source/molmoact2.mdx | 4 +- docs/source/multi_task_dit.mdx | 2 +- docs/source/robocasa.mdx | 2 +- docs/source/vlabench.mdx | 2 +- src/lerobot/configs/default.py | 4 ++ src/lerobot/configs/train.py | 10 +++- src/lerobot/datasets/__init__.py | 3 +- src/lerobot/datasets/factory.py | 79 +++++++++++++++++++++++++ src/lerobot/datasets/lerobot_dataset.py | 12 ++++ src/lerobot/datasets/sampler.py | 7 ++- src/lerobot/scripts/lerobot_train.py | 76 +++++++++++++++++++++--- tests/training/test_multi_gpu.py | 6 +- 18 files changed, 199 insertions(+), 32 deletions(-) diff --git a/.github/workflows/benchmark_tests.yml b/.github/workflows/benchmark_tests.yml index b82c59a8b..3493e5048 100644 --- a/.github/workflows/benchmark_tests.yml +++ b/.github/workflows/benchmark_tests.yml @@ -167,9 +167,9 @@ jobs: # ── LIBERO TRAIN+EVAL SMOKE ────────────────────────────────────────────── # Train SmolVLA for 1 step (batch_size=1, dataset episode 0 only) then - # immediately runs eval inside the training loop (eval_freq=1, 1 episode). + # immediately runs eval inside the training loop (env_eval_freq=1, 1 episode). # Tests the full train→eval-within-training pipeline end-to-end. - - name: Run Libero train+eval smoke (1 step, eval_freq=1) + - name: Run Libero train+eval smoke (1 step, env_eval_freq=1) if: env.HF_USER_TOKEN != '' run: | docker run --name libero-train-smoke --gpus all \ @@ -196,7 +196,7 @@ jobs: --output_dir=/tmp/train-smoke \ --steps=1 \ --batch_size=1 \ - --eval_freq=1 \ + --env_eval_freq=1 \ --eval.n_episodes=1 \ --eval.batch_size=1 \ --eval.use_async_envs=false \ diff --git a/Makefile b/Makefile index d3987101f..ea3b6e261 100644 --- a/Makefile +++ b/Makefile @@ -58,7 +58,7 @@ test-act-ete-train: --dataset.episodes="[0]" \ --batch_size=2 \ --steps=4 \ - --eval_freq=2 \ + --env_eval_freq=2 \ --eval.n_episodes=1 \ --eval.batch_size=1 \ --save_freq=2 \ @@ -96,7 +96,7 @@ test-diffusion-ete-train: --dataset.episodes="[0]" \ --batch_size=2 \ --steps=2 \ - --eval_freq=2 \ + --env_eval_freq=2 \ --eval.n_episodes=1 \ --eval.batch_size=1 \ --save_checkpoint=true \ @@ -126,7 +126,7 @@ test-tdmpc-ete-train: --dataset.episodes="[0]" \ --batch_size=2 \ --steps=2 \ - --eval_freq=2 \ + --env_eval_freq=2 \ --eval.n_episodes=1 \ --eval.batch_size=1 \ --save_checkpoint=true \ @@ -161,7 +161,7 @@ test-smolvla-ete-train: --dataset.episodes="[0]" \ --batch_size=2 \ --steps=4 \ - --eval_freq=2 \ + --env_eval_freq=2 \ --eval.n_episodes=1 \ --eval.batch_size=1 \ --save_freq=2 \ diff --git a/docs/source/hilserl.mdx b/docs/source/hilserl.mdx index 76e985cfe..09a370f3d 100644 --- a/docs/source/hilserl.mdx +++ b/docs/source/hilserl.mdx @@ -719,7 +719,7 @@ Example configuration for training the [reward classifier](https://huggingface.c "num_workers": 4, "steps": 5000, "log_freq": 10, - "eval_freq": 1000, + "env_eval_freq": 1000, "save_freq": 1000, "save_checkpoint": true, "seed": 2, diff --git a/docs/source/libero.mdx b/docs/source/libero.mdx index 043348690..b95af1d27 100644 --- a/docs/source/libero.mdx +++ b/docs/source/libero.mdx @@ -143,7 +143,7 @@ lerobot-train \ --batch_size=4 \ --eval.batch_size=1 \ --eval.n_episodes=1 \ - --eval_freq=1000 + --env_eval_freq=1000 ``` ## Reproducing published results diff --git a/docs/source/libero_plus.mdx b/docs/source/libero_plus.mdx index 4249bf49e..b065649fa 100644 --- a/docs/source/libero_plus.mdx +++ b/docs/source/libero_plus.mdx @@ -173,7 +173,7 @@ lerobot-train \ --batch_size=4 \ --eval.batch_size=1 \ --eval.n_episodes=1 \ - --eval_freq=1000 + --env_eval_freq=1000 ``` ## Relationship to LIBERO diff --git a/docs/source/metaworld.mdx b/docs/source/metaworld.mdx index 8e629dea9..b7accdfa2 100644 --- a/docs/source/metaworld.mdx +++ b/docs/source/metaworld.mdx @@ -120,11 +120,11 @@ lerobot-train \ --batch_size=4 \ --eval.batch_size=1 \ --eval.n_episodes=1 \ - --eval_freq=1000 + --env_eval_freq=1000 ``` ## Practical tips - Use the one-hot task conditioning for multi-task training (MT10/MT50 conventions) so policies have explicit task context. - Inspect the dataset task descriptions and the `info["is_success"]` keys when writing post-processing or logging so your success metrics line up with the benchmark. -- Adjust `batch_size`, `steps`, and `eval_freq` to match your compute budget. +- Adjust `batch_size`, `steps`, and `env_eval_freq` to match your compute budget. diff --git a/docs/source/molmoact2.mdx b/docs/source/molmoact2.mdx index 7927c6694..9a377031c 100644 --- a/docs/source/molmoact2.mdx +++ b/docs/source/molmoact2.mdx @@ -103,7 +103,7 @@ accelerate launch \ --batch_size=32 \ --num_workers=4 \ --log_freq=20 \ - --eval_freq=-1 \ + --env_eval_freq=-1 \ --save_checkpoint=true \ --save_freq=2000 ``` @@ -142,7 +142,7 @@ accelerate launch \ --batch_size=32 \ --num_workers=4 \ --log_freq=20 \ - --eval_freq=-1 \ + --env_eval_freq=-1 \ --save_checkpoint=true \ --save_freq=2000 ``` diff --git a/docs/source/multi_task_dit.mdx b/docs/source/multi_task_dit.mdx index 450d8a9f2..ebe46489a 100644 --- a/docs/source/multi_task_dit.mdx +++ b/docs/source/multi_task_dit.mdx @@ -314,7 +314,7 @@ lerobot-train \ --steps=30000 \ --save_freq=1000 \ --log_freq=100 \ - --eval_freq=1000 \ + --env_eval_freq=1000 \ --policy.type=multi_task_dit \ --policy.device=cuda \ --policy.horizon=32 \ diff --git a/docs/source/robocasa.mdx b/docs/source/robocasa.mdx index f6a784e72..5a335a484 100644 --- a/docs/source/robocasa.mdx +++ b/docs/source/robocasa.mdx @@ -166,7 +166,7 @@ lerobot-train \ --output_dir=./outputs/smolvla_robocasa_CloseFridge \ --steps=100000 \ --batch_size=4 \ - --eval_freq=5000 \ + --env_eval_freq=5000 \ --eval.batch_size=1 \ --eval.n_episodes=5 \ --save_freq=10000 diff --git a/docs/source/vlabench.mdx b/docs/source/vlabench.mdx index da579d674..9d45da4ec 100644 --- a/docs/source/vlabench.mdx +++ b/docs/source/vlabench.mdx @@ -165,7 +165,7 @@ lerobot-train \ --output_dir=./outputs/smolvla_vlabench_primitive \ --steps=100000 \ --batch_size=4 \ - --eval_freq=5000 \ + --env_eval_freq=5000 \ --eval.batch_size=1 \ --eval.n_episodes=1 \ --save_freq=10000 diff --git a/src/lerobot/configs/default.py b/src/lerobot/configs/default.py index 648e03f33..9b5433005 100644 --- a/src/lerobot/configs/default.py +++ b/src/lerobot/configs/default.py @@ -39,8 +39,12 @@ class DatasetConfig: # This reduces memory and speeds up DataLoader IPC. The training pipeline handles the conversion. return_uint8: bool = False streaming: bool = False + # Fraction of episodes held out per task for offline evaluation (0.0 = disabled). + eval_split: float = 0.0 def __post_init__(self) -> None: + if not (0.0 <= self.eval_split < 1.0): + raise ValueError(f"eval_split must be in [0.0, 1.0), got {self.eval_split}") if self.episodes is not None: if any(ep < 0 for ep in self.episodes): raise ValueError( diff --git a/src/lerobot/configs/train.py b/src/lerobot/configs/train.py index bac1a946b..00fbe81b2 100644 --- a/src/lerobot/configs/train.py +++ b/src/lerobot/configs/train.py @@ -100,8 +100,13 @@ class TrainPipelineConfig(HubMixin): prefetch_factor: int = 4 persistent_workers: bool = True steps: int = 100_000 - eval_freq: int = 20_000 + # Run policy in the simulation environment every N steps to measure reward/success (0 = disabled). + env_eval_freq: int = 20_000 log_freq: int = 200 + # Compute eval loss on held-out episodes every N steps (0 = disabled). Requires eval_split > 0. + eval_steps: int = 0 + # Cap on total eval samples, split uniformly across tasks (0 = use all held-out data). + max_eval_samples: int = 0 tolerance_s: float = 1e-4 save_checkpoint: bool = True # Checkpoint is saved every `save_freq` training iterations and after the last training step. @@ -208,6 +213,9 @@ class TrainPipelineConfig(HubMixin): self.optimizer = active_cfg.get_optimizer_preset() self.scheduler = active_cfg.get_scheduler_preset() + if self.eval_steps > 0 and self.dataset.eval_split == 0.0: + raise ValueError("eval_steps > 0 requires dataset.eval_split > 0.0 to hold out eval data.") + if hasattr(active_cfg, "push_to_hub") and active_cfg.push_to_hub and not active_cfg.repo_id: raise ValueError("'repo_id' argument missing. Please specify it to push the model to the hub.") diff --git a/src/lerobot/datasets/__init__.py b/src/lerobot/datasets/__init__.py index bd12a7248..7715a115e 100644 --- a/src/lerobot/datasets/__init__.py +++ b/src/lerobot/datasets/__init__.py @@ -35,7 +35,7 @@ from .dataset_tools import ( remove_feature, split_dataset, ) -from .factory import make_dataset, resolve_delta_timestamps +from .factory import make_dataset, make_train_eval_datasets, resolve_delta_timestamps from .image_writer import safe_stop_image_writer from .io_utils import load_episodes, write_stats from .language import ( @@ -89,6 +89,7 @@ __all__ = [ "get_feature_stats", "load_episodes", "make_dataset", + "make_train_eval_datasets", "merge_datasets", "modify_features", "modify_tasks", diff --git a/src/lerobot/datasets/factory.py b/src/lerobot/datasets/factory.py index cbbe83dc8..cd29ee99e 100644 --- a/src/lerobot/datasets/factory.py +++ b/src/lerobot/datasets/factory.py @@ -14,6 +14,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import logging +import math from pprint import pformat import torch @@ -130,3 +131,81 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas dataset.meta.stats[key][stats_type] = torch.tensor(stats, dtype=torch.float32) return dataset + + +def make_train_eval_datasets( + cfg: TrainPipelineConfig, +) -> tuple[LeRobotDataset | MultiLeRobotDataset, LeRobotDataset | None]: + """Create train and optional eval datasets by splitting episodes based on eval_split. + + The last ceil(n_episodes * eval_split) episodes per task are held out for evaluation. + If eval_split == 0.0, returns (full_dataset, None). + """ + full_dataset = make_dataset(cfg) + + if cfg.dataset.eval_split == 0.0: + return full_dataset, None + + base_episodes = ( + full_dataset.episodes if full_dataset.episodes is not None else list(range(full_dataset.num_episodes)) + ) + + episode_tasks = full_dataset.meta.episodes["tasks"] + task_to_episodes: dict[str, list[int]] = {} + for ep_idx in base_episodes: + task_key = episode_tasks[ep_idx][0] if episode_tasks[ep_idx] else "" + task_to_episodes.setdefault(task_key, []).append(ep_idx) + + train_episodes, eval_episodes = [], [] + for eps in task_to_episodes.values(): + n_eval = math.ceil(len(eps) * cfg.dataset.eval_split) + train_episodes.extend(eps[: len(eps) - n_eval]) + eval_episodes.extend(eps[len(eps) - n_eval :]) + + if not train_episodes: + raise ValueError( + f"eval_split={cfg.dataset.eval_split} leaves 0 training episodes from {len(base_episodes)} total." + ) + + logging.info( + f"Train/eval split: {len(train_episodes)} train, {len(eval_episodes)} eval " + f"(eval_split={cfg.dataset.eval_split}, {len(task_to_episodes)} tasks)" + ) + + delta_timestamps = resolve_delta_timestamps(cfg.trainable_config, full_dataset.meta) + + train_image_transforms = ( + ImageTransforms(cfg.dataset.image_transforms) if cfg.dataset.image_transforms.enable else None + ) + + train_dataset = LeRobotDataset( + cfg.dataset.repo_id, + root=cfg.dataset.root, + episodes=train_episodes, + delta_timestamps=delta_timestamps, + image_transforms=train_image_transforms, + revision=cfg.dataset.revision, + video_backend=cfg.dataset.video_backend, + return_uint8=True, + tolerance_s=cfg.tolerance_s, + ) + + eval_dataset = LeRobotDataset( + cfg.dataset.repo_id, + root=cfg.dataset.root, + episodes=eval_episodes, + delta_timestamps=delta_timestamps, + image_transforms=None, + revision=cfg.dataset.revision, + video_backend=cfg.dataset.video_backend, + return_uint8=True, + tolerance_s=cfg.tolerance_s, + ) + + if cfg.dataset.use_imagenet_stats: + for ds in (train_dataset, eval_dataset): + for key in ds.meta.camera_keys: + for stats_type, stats in IMAGENET_STATS.items(): + ds.meta.stats[key][stats_type] = torch.tensor(stats, dtype=torch.float32) + + return train_dataset, eval_dataset diff --git a/src/lerobot/datasets/lerobot_dataset.py b/src/lerobot/datasets/lerobot_dataset.py index d1e65fef1..49e77b53a 100644 --- a/src/lerobot/datasets/lerobot_dataset.py +++ b/src/lerobot/datasets/lerobot_dataset.py @@ -369,6 +369,18 @@ class LeRobotDataset(torch.utils.data.Dataset): self.reader.load_and_activate() return self.reader.hf_dataset + @property + def absolute_to_relative_idx(self) -> dict[int, int] | None: + """Mapping from absolute frame indices to HF dataset row positions. + + Non-None only for episode-filtered datasets where absolute indices + (from metadata) differ from row positions in the loaded HF dataset. + """ + reader = self._ensure_reader() + if reader.hf_dataset is None: + reader.load_and_activate() + return reader._absolute_to_relative_idx + # ── Writer-delegated methods ────────────────────────────────────── def add_frame(self, frame: dict) -> None: diff --git a/src/lerobot/datasets/sampler.py b/src/lerobot/datasets/sampler.py index af85dff9b..aee6ce46d 100644 --- a/src/lerobot/datasets/sampler.py +++ b/src/lerobot/datasets/sampler.py @@ -53,6 +53,7 @@ class EpisodeAwareSampler: drop_n_last_frames: int = 0, shuffle: bool = False, seed: int = 0, + absolute_to_relative_idx: dict[int, int] | None = None, ): """ Args: @@ -107,6 +108,7 @@ class EpisodeAwareSampler: self.seed = seed self._epoch = 0 self._start_index = 0 + self._absolute_to_relative = absolute_to_relative_idx @property def indices(self) -> list[int]: @@ -132,7 +134,10 @@ class EpisodeAwareSampler: def _frame_index(self, position: int) -> int: episode = int(np.searchsorted(self._cum_lengths, position, side="right")) position_in_episode = position - (int(self._cum_lengths[episode - 1]) if episode > 0 else 0) - return int(self._starts[episode]) + position_in_episode + absolute_idx = int(self._starts[episode]) + position_in_episode + if self._absolute_to_relative is not None: + return self._absolute_to_relative[absolute_idx] + return absolute_idx def __iter__(self) -> Iterator[int]: # Advance epoch state eagerly, not on first consumption of the generator. diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index 45281dac9..82c265c93 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/src/lerobot/scripts/lerobot_train.py @@ -47,7 +47,8 @@ from lerobot.common.train_utils import ( from lerobot.common.wandb_utils import WandBLogger from lerobot.configs import parser from lerobot.configs.train import TrainPipelineConfig -from lerobot.datasets import EpisodeAwareSampler, compute_sampler_state, make_dataset +from lerobot.datasets import EpisodeAwareSampler, compute_sampler_state +from lerobot.datasets.factory import make_train_eval_datasets from lerobot.envs import close_envs, make_env, make_env_pre_post_processors from lerobot.optim.factory import make_optimizer_and_scheduler from lerobot.policies import PreTrainedPolicy, make_policy, make_pre_post_processors @@ -245,19 +246,19 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): # LeRobotDataset skips its snapshot_download when try_load() succeeds, so no rank re-downloads. if is_main_process: logging.info("Creating dataset") - dataset = make_dataset(cfg) + dataset, eval_dataset = make_train_eval_datasets(cfg) accelerator.wait_for_everyone() # Other ranks read from the shared copy populated by the main process. if not is_main_process: - dataset = make_dataset(cfg) + dataset, eval_dataset = make_train_eval_datasets(cfg) # Create environment used for evaluating checkpoints during training on simulation data. # On real-world data, no need to create an environment as evaluations are done outside train.py, # using the eval.py instead, with gym_dora environment and dora-rs. eval_env = None - if cfg.eval_freq > 0 and cfg.env is not None and is_main_process: + if cfg.env_eval_freq > 0 and cfg.env is not None and is_main_process: logging.info("Creating env") eval_env = make_env(cfg.env, n_envs=cfg.eval.batch_size, use_async_envs=cfg.eval.use_async_envs) @@ -413,6 +414,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): drop_n_last_frames=getattr(active_cfg, "drop_n_last_frames", 0), shuffle=True, seed=cfg.seed if cfg.seed is not None else 0, + absolute_to_relative_idx=dataset.absolute_to_relative_idx, ) if cfg.resume and step > 0: # The resume offset depends on the (num_processes, batch_size) that produced `step`, so @@ -462,11 +464,43 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): persistent_workers=cfg.persistent_workers and cfg.num_workers > 0, ) + # Build eval dataloader if a held-out split exists + eval_dataloader = None + if eval_dataset is not None: + eval_ds = eval_dataset + if cfg.max_eval_samples > 0 and hasattr(eval_dataset, "hf_dataset"): + task_arr = eval_dataset.hf_dataset.data.column("task_index").to_numpy() + unique_tasks = sorted(set(task_arr.tolist())) + per_task = max(1, cfg.max_eval_samples // len(unique_tasks)) + selected: list[int] = [] + for t in unique_tasks: + frames = (task_arr == t).nonzero()[0][:per_task] + selected.extend(frames.tolist()) + eval_ds = torch.utils.data.Subset(eval_dataset, selected) + + eval_collate_fn = lerobot_collate_fn if dataset.meta.has_language_columns else None + eval_dataloader = torch.utils.data.DataLoader( + eval_ds, + batch_size=cfg.batch_size, + shuffle=False, + num_workers=cfg.num_workers, + pin_memory=device.type == "cuda", + drop_last=False, + 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, + ) + # Prepare everything with accelerator accelerator.wait_for_everyone() - policy, optimizer, dataloader, lr_scheduler = accelerator.prepare( - policy, optimizer, dataloader, lr_scheduler - ) + if eval_dataloader is not None: + policy, optimizer, dataloader, lr_scheduler, eval_dataloader = accelerator.prepare( + policy, optimizer, dataloader, lr_scheduler, eval_dataloader + ) + else: + policy, optimizer, dataloader, lr_scheduler = accelerator.prepare( + policy, optimizer, dataloader, lr_scheduler + ) # FSDP optimizer state is sharded across ranks, so it can only be loaded once the optimizer and # model are FSDP-wrapped (i.e. after `prepare`). Collective: every rank must participate. @@ -547,7 +581,8 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): train_tracker.step() is_log_step = cfg.log_freq > 0 and step % cfg.log_freq == 0 is_saving_step = step % cfg.save_freq == 0 or step == cfg.steps - is_eval_step = cfg.eval_freq > 0 and step % cfg.eval_freq == 0 + is_env_eval_step = cfg.env_eval_freq > 0 and step % cfg.env_eval_freq == 0 + is_eval_step = cfg.eval_steps > 0 and eval_dataloader is not None and step % cfg.eval_steps == 0 if is_log_step: # Collective reduce must run on every rank, before the main-process gate below. @@ -570,6 +605,29 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): wandb_logger.log_dict(wandb_log_dict, step) train_tracker.reset_averages() + if is_eval_step: + policy.eval() + eval_loss_sum = 0.0 + n_eval_batches = 0 + with torch.no_grad(), accelerator.autocast(): + for eval_batch in eval_dataloader: + for cam_key in dataset.meta.camera_keys: + if cam_key in eval_batch and eval_batch[cam_key].dtype == torch.uint8: + eval_batch[cam_key] = eval_batch[cam_key].to(dtype=torch.float32) / 255.0 + eval_batch = preprocessor(eval_batch) + loss, _ = policy.forward(eval_batch) + eval_loss_sum += loss.item() + n_eval_batches += 1 + eval_loss = eval_loss_sum / max(n_eval_batches, 1) + eval_loss = torch.tensor(eval_loss, device=device) + eval_loss = accelerator.reduce(eval_loss, reduction="mean").item() + policy.train() + + if is_main_process: + logging.info(f"step {step}: eval_loss={eval_loss:.4f}") + if wandb_logger: + wandb_logger.log_dict({"eval_loss": eval_loss}, step=step, mode="eval") + if cfg.save_checkpoint and is_saving_step: # Under FSDP, gathering the full model + optimizer state dicts is a cross-rank collective, # so all ranks must participate; rank 0 then writes the materialized dicts. For DDP / @@ -602,7 +660,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): accelerator.wait_for_everyone() - if cfg.env and is_eval_step: + if cfg.env and is_env_eval_step: if is_main_process: step_id = get_step_identifier(step, cfg.steps) logging.info(f"Eval policy at step {step}") diff --git a/tests/training/test_multi_gpu.py b/tests/training/test_multi_gpu.py index d37f1e35d..1a75d21cb 100644 --- a/tests/training/test_multi_gpu.py +++ b/tests/training/test_multi_gpu.py @@ -166,7 +166,7 @@ class TestMultiGPUTraining: f"--output_dir={output_dir}", "--batch_size=4", "--steps=10", - "--eval_freq=-1", + "--env_eval_freq=-1", "--log_freq=5", "--save_freq=10", "--seed=42", @@ -209,7 +209,7 @@ class TestMultiGPUTraining: f"--output_dir={output_dir}", "--batch_size=4", "--steps=20", - "--eval_freq=-1", + "--env_eval_freq=-1", "--log_freq=5", "--save_freq=10", "--seed=42", @@ -267,7 +267,7 @@ class TestMultiGPUTraining: f"--output_dir={output_dir}", "--batch_size=4", "--steps=10", - "--eval_freq=-1", + "--env_eval_freq=-1", "--log_freq=5", "--save_freq=10", "--seed=42", From 3dd19d043e2f3fe5673b13ea0ebe4f31884c0797 Mon Sep 17 00:00:00 2001 From: Caroline Pascal Date: Sat, 27 Jun 2026 14:21:21 +0200 Subject: [PATCH 18/38] feat(depth maps): adding support for depth in LeRobot (#3644) * feat(depth): add depth quantization helpers and tests * feat(video): add ffv1 to supported codecs * feat(depth): persist depth metadata * feat(depth): extend quantization tools to better fit the encoding/decoding pipeline * feat(depth): plumb DepthEncoderConfig through LeRobotDataset and DatasetWriter * feat(depth): wire StreamingVideoEncoder + writer to depth encoder * feat(depth): wire DatasetReader to decode_depth_frames * feat(cameras/realsense): expose async depth in metric meters * feat(features): route 2D camera shapes to observation.depth. * feat(robots/so_follower): emit + populate depth keys when use_depth * feat(record): plumb DepthEncoderConfig through lerobot-record * feat(viz): render depth observations as rr.DepthImage in Viridis * feat(depth maps writer): adding support for raw depth maps recording with image writer * chore(format): format code * feat(depth shape): ensuring depth maps shape is always including the channel * feat(is_depth): simplifying is_depth nested name + legacy support * fix(stop_event): fixing stop_event race condition in camera classes * fix(plumbing): fixing missing parts in the depth maps pipeline * chore(typos): fixing typos * test(fix): fixing exisiting tests to still work with latest features * tests(depth): adding new tests for depth integration validation * feat(pix_fmt channels): use PyAv to check get pixel formats number of channels * feat(refactor): refactor DepthEncoderConfig quantization pipeline, so that the methods do not live in the config class. Add pixel format - channels validation.Move the default pixel format for depth in the config file. * fix(pre-commit): fixing mutable defautl value * fix(info): fixing info metadata update when is_depth_map was set * tests(typos): fixing typos in tests * fix(realsense): fixing typo in realsense serial number * fix(normalization): restricting 255 normalization to non depth/uint8 images only * fix(typo): fixing typo * fix(TIFF): add missing quantization and cleanup for TIFF files * feat(batched dequantization): optimizing dequantize_depth for torch based batched dequantization * feat(tools): adding depth support in LeRobotDataset edition tools * test(aggregate): extending aggregation tests to depth frames * test(cleaning): cleaning up tests * fix(from_video_info): fixing early validation issue in from_video_info * fix(typo): fixing typo * fix(is_depth): adding missing doctrings and is_depth arguments in video decoding functions Co-authored-by: Wensi (Vince) Ai <59036629+wensi-ai@users.noreply.github.com> * fix(depth units): fixing depth units output for the realsense cameras * feat(output unit): adding support for output unit specification at dataset reading/training time Co-authored-by: Wensi (Vince) Ai <59036629+wensi-ai@users.noreply.github.com> * test(depth): cleaning up depth tests * test(depth encoding): updating and cleaning video/depth encoding tests * chore(format): formatting code * docs(depth): improving depth maps docs * test(fix): fixing depth tests * test(dataset tools): adding missing tests for new dataset edition tools features * chore(format): formatting code * fix(pyav check): fixing PyAV option validation for integer codec options by normalizing numeric values before calling `is_integer()` Co-authored-by: Wensi (Vince) Ai <59036629+wensi-ai@users.noreply.github.com> * docs(mermaid): fixing mermaid diagram * fix(rebase): rebase follow up corrections * feat(dataset tools): adding missing docstrings and features for depth fill support in dataset edition tools * docs(docstring): updating docstrings * docs(dataset tools): updating docs * fix(save images): fixing image saving in dataset tools * fix(update video info): fixing update video info logic to match the recording and editing use cases * test(reencode): fixing reencoding monkeypatch * fix(review): add Claude review * chore(format): format code * fix(update video info): ditching the differentiated approahces for video info update - video info are always updated unless for preserved keys. * chore(rebase): fixing rebase merge conflicts * test(visualization): fixing visualization tests * feat(docstrings): adding explicit docstring for encoding parameters. Docstrigns will now show up as description in the CLI --help. * feat(mm as default): adding a global DEFAULT_DEPTH_UNIT variable setting mm as default depth unit * fix(RGB <-> camera): renaming camera_encoder to rgb_encoder for clarity * chore(TODO): removing deprecated TODO * doc(write_u16_plane): improving docstrings for write_u16_plane * feat(units): adding constants for depth frames units (m and mm) * fix(spam): replacing spamming warning but a debug log * feat(leagcy metadata): adding automatic metadata update for legacy 'video.is_depth_map' feature * fix(copy&reindex): fixing metadat reshaping for single channel frames * fix(ImageNet): excluding dpeth frames from ImageNet stats * fix(PyAV container seek): fixing initial PyAV container seek to be robust againsy codec choice * feat(lerobot-dataset-viz): adding support for depth in lerobot-dataset-viz * fix(compress): removing rerun compression for DepthImages * fix(signle channel squeeze): fixing single channel squeezing * chore(format): format code * fix(streaming): adding support for dequantization in streaming_dataset.py * refactor(read depth): factorizing depth reading methods for realsense camera and adding support for depth-only usage * chore(renaming): fixing missed RGBEncoderConfig renamings * docs(renaming): reflecting renamings in a clearer way in the docs * chore(annotation): excluding depth from the annotation pipeline * feat(robots): adding depth support in compatible follower robots * feat(LeSadKiwi): excluding LeKiwi from depth support (for now) * chore(fail): removing misplaced file * chore(fail): removing misplaced file * fix(remove ffv1): removing ffv1 as it does not support MP4 * docs(cheat sheet): adding depth and video encoding to the cheat sheet * fix(lossless): tuning depth encoding parameters for lossless depth storage * test(fix): fixing failing tests * depth(ZMQ): excluding ZMQ from depth support * Revert "depth(ZMQ): excluding ZMQ from depth support" This reverts commit b95cf4e4c2bb1c188263bbcdcfbd6f3aea034ecb. * fix(image transforms): excluding depth frames from images transforms * fix(typo): typo * fix(stats): fixing stats computation for depth frames * fix(TIFF vs. pytorch): adding an extra uint16 to float32 conversion for depth maps stored as raw TIFF images * fix(typos): fixing typos * test(dtype): fixing stats computation typing tests --------- Signed-off-by: Steven Palma Co-authored-by: Wensi (Vince) Ai <59036629+wensi-ai@users.noreply.github.com> Co-authored-by: Steven Palma Co-authored-by: Wensi Ai --- docs/source/cameras.mdx | 8 + docs/source/cheat-sheet.mdx | 30 ++ docs/source/earthrover_mini_plus.mdx | 2 +- docs/source/groot.mdx | 2 +- docs/source/hope_jr.mdx | 4 +- docs/source/il_robots.mdx | 2 +- docs/source/lerobot-dataset-v3.mdx | 2 +- docs/source/reachy2.mdx | 4 +- docs/source/streaming_video_encoding.mdx | 40 +- docs/source/using_dataset_tools.mdx | 57 +- docs/source/video_encoding_parameters.mdx | 97 +++- .../annotations/steerable_pipeline/frames.py | 10 +- src/lerobot/async_inference/helpers.py | 5 +- src/lerobot/cameras/opencv/camera_opencv.py | 4 +- .../cameras/realsense/camera_realsense.py | 183 ++++--- .../realsense/configuration_realsense.py | 6 + src/lerobot/cameras/zmq/camera_zmq.py | 2 + src/lerobot/configs/__init__.py | 15 +- src/lerobot/configs/dataset.py | 8 +- src/lerobot/configs/default.py | 11 +- src/lerobot/configs/video.py | 159 ++++-- src/lerobot/datasets/compute_stats.py | 22 +- src/lerobot/datasets/dataset_metadata.py | 53 +- src/lerobot/datasets/dataset_reader.py | 30 +- src/lerobot/datasets/dataset_tools.py | 185 ++++--- src/lerobot/datasets/dataset_writer.py | 56 +- src/lerobot/datasets/depth_utils.py | 268 ++++++++++ src/lerobot/datasets/factory.py | 3 + src/lerobot/datasets/feature_utils.py | 2 +- src/lerobot/datasets/image_writer.py | 68 ++- src/lerobot/datasets/io_utils.py | 47 +- src/lerobot/datasets/lerobot_dataset.py | 58 ++- src/lerobot/datasets/pyav_utils.py | 51 +- src/lerobot/datasets/streaming_dataset.py | 47 +- src/lerobot/datasets/utils.py | 5 +- src/lerobot/datasets/video_utils.py | 239 ++++++--- src/lerobot/policies/utils.py | 3 +- src/lerobot/robots/hope_jr/hope_jr_arm.py | 26 +- src/lerobot/robots/hope_jr/hope_jr_hand.py | 26 +- .../robots/koch_follower/koch_follower.py | 26 +- src/lerobot/robots/lekiwi/lekiwi.py | 6 + src/lerobot/robots/lekiwi/lekiwi_client.py | 7 + .../robots/omx_follower/omx_follower.py | 26 +- .../openarm_follower/openarm_follower.py | 26 +- .../rebot_b601_follower.py | 26 +- src/lerobot/robots/so_follower/so_follower.py | 25 +- src/lerobot/robots/unitree_g1/unitree_g1.py | 16 +- src/lerobot/rollout/context.py | 6 +- src/lerobot/scripts/lerobot_dataset_viz.py | 44 +- src/lerobot/scripts/lerobot_edit_dataset.py | 53 +- src/lerobot/scripts/lerobot_record.py | 14 +- src/lerobot/scripts/lerobot_rollout.py | 6 +- src/lerobot/utils/feature_utils.py | 35 +- src/lerobot/utils/visualization_utils.py | 5 +- tests/annotations/test_frames.py | 5 +- tests/datasets/test_aggregate.py | 82 ++- tests/datasets/test_compute_stats.py | 39 +- tests/datasets/test_dataset_metadata.py | 49 +- tests/datasets/test_dataset_tools.py | 144 +++++- tests/datasets/test_dataset_writer.py | 14 +- tests/datasets/test_datasets.py | 4 + tests/datasets/test_depth.py | 247 +++++++++ tests/datasets/test_image_writer.py | 4 +- .../datasets/test_streaming_video_encoder.py | 45 +- tests/datasets/test_video_encoding.py | 487 +++++++++++++----- tests/fixtures/constants.py | 46 +- tests/fixtures/dataset_factories.py | 38 ++ tests/scripts/test_edit_dataset_parsing.py | 45 ++ tests/utils/test_visualization_utils.py | 9 +- 69 files changed, 2740 insertions(+), 679 deletions(-) create mode 100644 src/lerobot/datasets/depth_utils.py create mode 100644 tests/datasets/test_depth.py diff --git a/docs/source/cameras.mdx b/docs/source/cameras.mdx index 2dc2859dd..02714d591 100644 --- a/docs/source/cameras.mdx +++ b/docs/source/cameras.mdx @@ -157,6 +157,14 @@ finally: +### Working with depth + +The Intel RealSense and Reachy 2 cameras can capture both color and depth in lockstep. Calling `read()` returns the **color** frame as `(H, W, 3)` `uint8`. Calling `read_depth()` returns the **depth map** as `(H, W, 1)` `uint16`, where each pixel value is the distance from the sensor expressed in **millimetres**. A pixel value of `0` typically means "no measurement available" (out-of-range, occluded, or low-confidence). + +During recording, the control loop peeks the freshest buffered frames non-blockingly via `read_latest()` (color) and `read_latest_depth()` (depth), adding the depth map as a sibling feature (e.g. `front_depth` next to `front`). + +For how depth streams are stored and encoded when recording a dataset, see the [Depth streams](./video_encoding_parameters#depth-streams) section of the video encoding guide. + ## Use your phone's camera diff --git a/docs/source/cheat-sheet.mdx b/docs/source/cheat-sheet.mdx index a6afa14c2..45952c5b3 100644 --- a/docs/source/cheat-sheet.mdx +++ b/docs/source/cheat-sheet.mdx @@ -89,6 +89,36 @@ Control the data recording flow using keyboard shortcuts: - Press **Left Arrow (`←`)**: Delete current episode and retry. - Press **Escape (`ESC`)**: Stop, encode videos, and upload. +### Recording depth + +Intel RealSense cameras (`type: intelrealsense`) record a depth stream when you set `use_depth: true`. Depth is quantized to 12-bit codes and stored as its own video. + +```bash +lerobot-record \ + ... \ + --robot.cameras="{ head: {type: intelrealsense, serial_number_or_name: \"0123456789\", width: 640, height: 480, fps: 30, use_depth: true} }" \ + --dataset.repo_id=${HF_USER}/so101_depth_test \ + --dataset.single_task="put the red brick in a bowl" \ + --dataset.depth_encoder.depth_min=0.01 \ + --dataset.depth_encoder.depth_max=10.0 \ + --dataset.depth_encoder.shift=0.0 \ + --dataset.depth_encoder.use_log=true +``` + +### Video encoding parameters + +RGB and depth streams are encoded independently via the `--dataset.rgb_encoder.*` and `--dataset.depth_encoder.*` keys. + +```bash +lerobot-record \ + ... \ + --dataset.rgb_encoder.vcodec=h264 \ + --dataset.rgb_encoder.pix_fmt=yuv420p \ + --dataset.rgb_encoder.crf=23 \ + --dataset.depth_encoder.vcodec=hevc \ + --dataset.depth_encoder.extra_options='{"x265-params": "lossless=1"}' +``` + ### Training Depending on your hardware training the policy might take a few hours. That's how you train simple `ACT` policy: diff --git a/docs/source/earthrover_mini_plus.mdx b/docs/source/earthrover_mini_plus.mdx index 508c0e3a9..f3b324093 100644 --- a/docs/source/earthrover_mini_plus.mdx +++ b/docs/source/earthrover_mini_plus.mdx @@ -194,7 +194,7 @@ lerobot-record \ --dataset.single_task="Navigate around obstacles" \ --dataset.streaming_encoding=true \ --dataset.encoder_threads=2 \ - # --dataset.camera_encoder.vcodec=auto \ + # --dataset.rgb_encoder.vcodec=auto \ --display_data=true ``` diff --git a/docs/source/groot.mdx b/docs/source/groot.mdx index a10b5e369..3ab202fb2 100644 --- a/docs/source/groot.mdx +++ b/docs/source/groot.mdx @@ -124,7 +124,7 @@ lerobot-rollout\ --dataset.single_task="Grab and handover the red cube to the other arm" \ --dataset.streaming_encoding=true \ --dataset.encoder_threads=2 \ - # --dataset.camera_encoder.vcodec=auto \ + # --dataset.rgb_encoder.vcodec=auto \ --policy.path=/groot-bimanual \ # your trained model --duration=600 ``` diff --git a/docs/source/hope_jr.mdx b/docs/source/hope_jr.mdx index 1f3b08fd7..c29a9f216 100644 --- a/docs/source/hope_jr.mdx +++ b/docs/source/hope_jr.mdx @@ -232,7 +232,7 @@ lerobot-record \ --dataset.private=true \ --dataset.streaming_encoding=true \ --dataset.encoder_threads=2 \ - # --dataset.camera_encoder.vcodec=auto \ + # --dataset.rgb_encoder.vcodec=auto \ --display_data=true ``` @@ -278,6 +278,6 @@ lerobot-record \ --dataset.num_episodes=10 \ --dataset.streaming_encoding=true \ --dataset.encoder_threads=2 \ - # --dataset.camera_encoder.vcodec=auto \ + # --dataset.rgb_encoder.vcodec=auto \ --policy.path=outputs/train/hopejr_hand/checkpoints/last/pretrained_model ``` diff --git a/docs/source/il_robots.mdx b/docs/source/il_robots.mdx index 6a820e0db..0f14bd133 100644 --- a/docs/source/il_robots.mdx +++ b/docs/source/il_robots.mdx @@ -207,7 +207,7 @@ lerobot-record \ --dataset.num_episodes=5 \ --dataset.single_task="Grab the black cube" \ --dataset.streaming_encoding=true \ - # --dataset.camera_encoder.vcodec=auto \ + # --dataset.rgb_encoder.vcodec=auto \ --dataset.encoder_threads=2 ``` diff --git a/docs/source/lerobot-dataset-v3.mdx b/docs/source/lerobot-dataset-v3.mdx index 21cb232d3..0647af0b0 100644 --- a/docs/source/lerobot-dataset-v3.mdx +++ b/docs/source/lerobot-dataset-v3.mdx @@ -44,7 +44,7 @@ lerobot-record \ --dataset.num_episodes=5 \ --dataset.single_task="Grab the black cube" \ --dataset.streaming_encoding=true \ - # --dataset.camera_encoder.vcodec=auto \ + # --dataset.rgb_encoder.vcodec=auto \ --dataset.encoder_threads=2 ``` diff --git a/docs/source/reachy2.mdx b/docs/source/reachy2.mdx index 4b08569db..7f975af43 100644 --- a/docs/source/reachy2.mdx +++ b/docs/source/reachy2.mdx @@ -161,7 +161,7 @@ lerobot-record \ --dataset.private=true \ --dataset.streaming_encoding=true \ --dataset.encoder_threads=2 \ - # --dataset.camera_encoder.vcodec=auto \ + # --dataset.rgb_encoder.vcodec=auto \ --display_data=true ``` @@ -203,7 +203,7 @@ lerobot-record \ --dataset.private=true \ --dataset.streaming_encoding=true \ --dataset.encoder_threads=2 \ - # --dataset.camera_encoder.vcodec=auto \ + # --dataset.rgb_encoder.vcodec=auto \ --display_data=true ``` diff --git a/docs/source/streaming_video_encoding.mdx b/docs/source/streaming_video_encoding.mdx index 96e049eb3..0be32b717 100644 --- a/docs/source/streaming_video_encoding.mdx +++ b/docs/source/streaming_video_encoding.mdx @@ -17,7 +17,7 @@ This makes `save_episode()` near-instant (the video is already encoded by the ti | Parameter | CLI Flag | Type | Default | Description | | ----------------------- | --------------------------------- | ------------- | ------------- | ----------------------------------------------------------------- | | `streaming_encoding` | `--dataset.streaming_encoding` | `bool` | `True` | Enable real-time encoding during capture | -| `vcodec` | `--dataset.camera_encoder.vcodec` | `str` | `"libsvtav1"` | Video codec. `"auto"` detects best HW encoder | +| `vcodec` | `--dataset.rgb_encoder.vcodec` | `str` | `"libsvtav1"` | Video codec. `"auto"` detects best HW encoder | | `encoder_threads` | `--dataset.encoder_threads` | `int \| None` | `None` (auto) | Threads per encoder instance. `None` will leave the vcoded decide | | `encoder_queue_maxsize` | `--dataset.encoder_queue_maxsize` | `int` | `30` | Max buffered frames per camera (~1s at 30fps). Consumes RAM | @@ -82,15 +82,15 @@ Use HW encoding when: ### Available HW Encoders -| Encoder | Platform | Hardware | CLI Value | -| ------------------- | ------------- | ------------------------------------------------------------------------------------------------ | --------------------------------------------------- | -| `h264_videotoolbox` | macOS | Apple Silicon / Intel | `--dataset.camera_encoder.vcodec=h264_videotoolbox` | -| `hevc_videotoolbox` | macOS | Apple Silicon / Intel | `--dataset.camera_encoder.vcodec=hevc_videotoolbox` | -| `h264_nvenc` | Linux/Windows | NVIDIA GPU | `--dataset.camera_encoder.vcodec=h264_nvenc` | -| `hevc_nvenc` | Linux/Windows | NVIDIA GPU | `--dataset.camera_encoder.vcodec=hevc_nvenc` | -| `h264_vaapi` | Linux | Intel/AMD GPU | `--dataset.camera_encoder.vcodec=h264_vaapi` | -| `h264_qsv` | Linux/Windows | Intel Quick Sync | `--dataset.camera_encoder.vcodec=h264_qsv` | -| `auto` | Any | Probes the system for available HW encoders. Falls back to `libsvtav1` if no HW encoder is found | `--dataset.camera_encoder.vcodec=auto` | +| Encoder | Platform | Hardware | CLI Value | +| ------------------- | ------------- | ------------------------------------------------------------------------------------------------ | ------------------------------------------------ | +| `h264_videotoolbox` | macOS | Apple Silicon / Intel | `--dataset.rgb_encoder.vcodec=h264_videotoolbox` | +| `hevc_videotoolbox` | macOS | Apple Silicon / Intel | `--dataset.rgb_encoder.vcodec=hevc_videotoolbox` | +| `h264_nvenc` | Linux/Windows | NVIDIA GPU | `--dataset.rgb_encoder.vcodec=h264_nvenc` | +| `hevc_nvenc` | Linux/Windows | NVIDIA GPU | `--dataset.rgb_encoder.vcodec=hevc_nvenc` | +| `h264_vaapi` | Linux | Intel/AMD GPU | `--dataset.rgb_encoder.vcodec=h264_vaapi` | +| `h264_qsv` | Linux/Windows | Intel Quick Sync | `--dataset.rgb_encoder.vcodec=h264_qsv` | +| `auto` | Any | Probes the system for available HW encoders. Falls back to `libsvtav1` if no HW encoder is found | `--dataset.rgb_encoder.vcodec=auto` | > [!NOTE] > In order to use the HW accelerated encoders you might need to upgrade your GPU drivers. @@ -100,15 +100,15 @@ Use HW encoding when: ## 5. Troubleshooting -| Symptom | Likely Cause | Fix | -| ------------------------------------------------------------------ | -------------------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| System freezes or choppy robot movement or Rerun visualization lag | CPU starved (100% load usage) | Close other apps, reduce encoding throughput, lower `encoder_threads`, use `h264`, use `display_data=False`. If the CPU continues to be at 100% then it might be insufficient for your setup, consider `--dataset.streaming_encoding=false` or HW encoding (`--dataset.camera_encoder.vcodec=auto`) | -| "Encoder queue full" warnings or dropped frames in dataset | Encoder can't keep up (Queue overflow) | If CPU is not at 100%: Increase `encoder_threads`, increase `encoder_queue_maxsize` or use HW encoding (`--dataset.camera_encoder.vcodec=auto`). | -| High RAM usage | Queue filling faster than encoding | `encoder_threads` too low or CPU insufficient. Reduce `encoder_queue_maxsize` or use HW encoding | -| Large video files | Using HW encoder or H.264 | Expected trade-off. Switch to `libsvtav1` if CPU allows | -| `save_episode()` still slow | `streaming_encoding` is `False` | Set `--dataset.streaming_encoding=true` | -| Encoder thread crash | Codec not available or invalid settings | Check `vcodec` is installed, try `--dataset.camera_encoder.vcodec=auto` | -| Recorded dataset is missing frames | CPU/GPU starvation or occasional load spikes | If ~5% of frames are missing, your system is likely overloaded — follow the recommendations above. If fewer frames are missing (~2%), they are probably due to occasional transient load spikes (often at startup) and can be considered expected. | +| Symptom | Likely Cause | Fix | +| ------------------------------------------------------------------ | -------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | +| System freezes or choppy robot movement or Rerun visualization lag | CPU starved (100% load usage) | Close other apps, reduce encoding throughput, lower `encoder_threads`, use `h264`, use `display_data=False`. If the CPU continues to be at 100% then it might be insufficient for your setup, consider `--dataset.streaming_encoding=false` or HW encoding (`--dataset.rgb_encoder.vcodec=auto`) | +| "Encoder queue full" warnings or dropped frames in dataset | Encoder can't keep up (Queue overflow) | If CPU is not at 100%: Increase `encoder_threads`, increase `encoder_queue_maxsize` or use HW encoding (`--dataset.rgb_encoder.vcodec=auto`). | +| High RAM usage | Queue filling faster than encoding | `encoder_threads` too low or CPU insufficient. Reduce `encoder_queue_maxsize` or use HW encoding | +| Large video files | Using HW encoder or H.264 | Expected trade-off. Switch to `libsvtav1` if CPU allows | +| `save_episode()` still slow | `streaming_encoding` is `False` | Set `--dataset.streaming_encoding=true` | +| Encoder thread crash | Codec not available or invalid settings | Check `vcodec` is installed, try `--dataset.rgb_encoder.vcodec=auto` | +| Recorded dataset is missing frames | CPU/GPU starvation or occasional load spikes | If ~5% of frames are missing, your system is likely overloaded — follow the recommendations above. If fewer frames are missing (~2%), they are probably due to occasional transient load spikes (often at startup) and can be considered expected. | ## 6. Recommended Configurations @@ -146,7 +146,7 @@ On very constrained systems, streaming encoding may compete too heavily with the # 2camsx 640x480x3 @30fps: Requires some tuning. # Use H.264, disable streaming, consider batching encoding -lerobot-record --dataset.camera_encoder.vcodec=h264 --dataset.streaming_encoding=false ... +lerobot-record --dataset.rgb_encoder.vcodec=h264 --dataset.streaming_encoding=false ... ``` ## 7. Closing note diff --git a/docs/source/using_dataset_tools.mdx b/docs/source/using_dataset_tools.mdx index 49247a6c1..e9299d298 100644 --- a/docs/source/using_dataset_tools.mdx +++ b/docs/source/using_dataset_tools.mdx @@ -11,8 +11,9 @@ LeRobot provides several utilities for manipulating datasets: 3. **Merge Datasets** - Combine multiple datasets into one. The datasets must have identical features, and episodes are concatenated in the order specified in `repo_ids` 4. **Add Features** - Add new features to a dataset 5. **Remove Features** - Remove features from a dataset -6. **Convert to Video** - Convert image-based datasets to video format for efficient storage -7. **Show the Info of Datasets** - Show the summary of datasets information such as number of episode etc. +6. **Convert to Video** - Convert image-based datasets to video format for efficient storage (RGB and depth cameras are encoded with separate encoders) +7. **Re-encode Videos** - Re-encode an existing video dataset's RGB and/or depth streams with new encoder settings +8. **Show the Info of Datasets** - Show the summary of datasets information such as number of episode etc. The core implementation is in `lerobot.datasets.dataset_tools`. An example script detailing how to use the tools API is available in `examples/dataset/use_dataset_tools.py`. @@ -117,10 +118,19 @@ lerobot-edit-dataset \ --repo_id lerobot/pusht_image \ --operation.type convert_image_to_video \ --operation.output_dir outputs/pusht_video \ - --operation.camera_encoder.vcodec libsvtav1 \ - --operation.camera_encoder.pix_fmt yuv420p \ - --operation.camera_encoder.g 2 \ - --operation.camera_encoder.crf 30 + --operation.rgb_encoder.vcodec libsvtav1 \ + --operation.rgb_encoder.pix_fmt yuv420p \ + --operation.rgb_encoder.g 2 \ + --operation.rgb_encoder.crf 30 + +# Convert a dataset that includes depth maps, customizing the depth encoder +lerobot-edit-dataset \ + --repo_id lerobot/pusht_image \ + --operation.type convert_image_to_video \ + --operation.output_dir outputs/pusht_video \ + --operation.depth_encoder.depth_min 0.01 \ + --operation.depth_encoder.depth_max 10.0 \ + --operation.depth_encoder.use_log true # Convert only specific episodes lerobot-edit-dataset \ @@ -147,11 +157,42 @@ lerobot-edit-dataset \ **Parameters:** - `output_dir`: Custom output directory (optional - by default uses `new_repo_id` or `{repo_id}_video`) -- `camera_encoder`: Video encoder settings — all sub-fields accessible via `--operation.camera_encoder.. See [Video Encoding Parameters](./video_encoding_parameters) for more details. +- `rgb_encoder`: Video encoder settings applied to RGB cameras — all sub-fields accessible via `--operation.rgb_encoder.`. See [Video Encoding Parameters](./video_encoding_parameters) for more details. +- `depth_encoder`: Video encoder settings applied to depth-map cameras (e.g. from an Intel RealSense). In addition to the standard encoder fields it exposes the depth quantization knobs (`depth_min`, `depth_max`, `shift`, `use_log`), accessible via `--operation.depth_encoder.`. These quantization settings are persisted to the dataset metadata so depth can be dequantized back to physical units on load. See the [Depth streams](./video_encoding_parameters#depth-streams) section for details. - `episode_indices`: List of specific episodes to convert (default: all episodes) - `num_workers`: Number of parallel workers for processing (default: 4) -**Note:** The resulting dataset will be a proper LeRobotDataset with all cameras encoded as videos in the `videos/` directory, with parquet files containing only metadata (no raw image data). All episodes, stats, and tasks are preserved. +**Note:** The resulting dataset will be a proper LeRobotDataset with all cameras encoded as videos in the `videos/` directory, with parquet files containing only metadata (no raw image data). Depth-map cameras are detected automatically and routed to the `depth_encoder`, while RGB cameras use the `rgb_encoder`. All episodes, stats, and tasks are preserved. + +#### Re-encode Videos + +Re-encode the videos of an existing video dataset with different encoder settings, without going back to raw frames. RGB videos use the `rgb_encoder` and depth videos use the `depth_encoder`. Provide only the encoder(s) you want to re-encode; the other stream type is left untouched. + +```bash +# Re-encode all RGB videos with new settings (saves to lerobot/pusht_reencoded by default) +lerobot-edit-dataset \ + --repo_id lerobot/pusht \ + --operation.type reencode_videos \ + --operation.rgb_encoder.vcodec h264 \ + --operation.rgb_encoder.pix_fmt yuv420p \ + --operation.rgb_encoder.crf 23 + +# Re-encode both RGB and depth videos in a dataset with depth maps +lerobot-edit-dataset \ + --repo_id lerobot/pusht_depth \ + --operation.type reencode_videos \ + --operation.rgb_encoder.vcodec h264 \ + --operation.depth_encoder.crf 50 +``` + +**Parameters:** + +- `rgb_encoder`: Encoder settings applied to every RGB video. Omit to skip re-encoding RGB videos. +- `depth_encoder`: Encoder settings applied to every depth video. Omit to skip re-encoding depth videos. +- `num_workers`: Number of parallel workers for processing. + +> [!NOTE] +> When re-encoding depth videos, the existing depth quantization parameters (`depth_min`, `depth_max`, `shift`, `use_log`) and the `is_depth_map` flag are **preserved** — re-encoding only changes the codec/quality of the stored stream, not how depth is dequantized on load. ### Show the information of datasets diff --git a/docs/source/video_encoding_parameters.mdx b/docs/source/video_encoding_parameters.mdx index 0b5b99b2b..132d25056 100644 --- a/docs/source/video_encoding_parameters.mdx +++ b/docs/source/video_encoding_parameters.mdx @@ -2,15 +2,15 @@ When video storage is enabled, LeRobot stores each camera stream as an **MP4** file instead of saving one image file per timestep. Video encoding compresses across time, which usually cuts dataset size and I/O compared to a pile of PNG, while keeping MP4 — a format every player and loader understands. -Encoding frames into an MP4 is a full FFmpeg pipeline: choice of encoder, pixel format, GOP/keyframes, quality vs. speed, and optional extra encoder flags. Most of these knobs are user-tunable through `camera_encoder`, a nested `VideoEncoderConfig` (`lerobot.configs.video.VideoEncoderConfig`) passed through PyAV. +Encoding frames into an MP4 is a full FFmpeg pipeline: choice of encoder, pixel format, GOP/keyframes, quality vs. speed, and optional extra encoder flags. Most of these knobs are user-tunable through `rgb_encoder`, a nested `RGBEncoderConfig` (`lerobot.configs.video.RGBEncoderConfig`) passed through PyAV. -You can set these parameters from the CLI with `--dataset.camera_encoder.` (e.g. with `lerobot-record` or `lerobot-rollout`). The same block applies to every camera video stream in that run. +You can set these parameters from the CLI with `--dataset.rgb_encoder.` (e.g. with `lerobot-record` or `lerobot-rollout`). The same block applies to every camera video stream in that run. - Video storage must be on for `camera_encoder` to have any effect — + Video storage must be on for `rgb_encoder` to have any effect — `use_videos=True` in Python APIs, or `--dataset.video=true` on the CLI (the - recording default). With video off, inputs stay as images and `camera_encoder` - is ignored. + recording default). With video off, inputs stay as images and `rgb_encoder` is + ignored. For details on **when** frames are written vs. encoded (streaming vs. post-episode), queues, and other top-level `--dataset.*` switches, see [Streaming Video Encoding](./streaming_video_encoding). For an encoding-parameter comparison and experiments, see the [video-benchmark Space](https://huggingface.co/spaces/lerobot/video-benchmark). @@ -33,9 +33,9 @@ lerobot-record \ --dataset.single_task="Grab the cube" \ --dataset.streaming_encoding=true \ --dataset.encoder_threads=2 \ - --dataset.camera_encoder.vcodec=h264 \ - --dataset.camera_encoder.preset=fast \ - --dataset.camera_encoder.extra_options={"tune": "film", "profile:v": "high", "bf": 2} \ + --dataset.rgb_encoder.vcodec=h264 \ + --dataset.rgb_encoder.preset=fast \ + --dataset.rgb_encoder.extra_options={"tune": "film", "profile:v": "high", "bf": 2} \ --display_data=true ``` @@ -50,7 +50,7 @@ Only override these parameters if you have a specific reason to, and measure the -All flags below are prefixed with `--dataset.camera_encoder.` on the CLI. +All flags below are prefixed with `--dataset.rgb_encoder.` on the CLI. | Parameter | Type | Default | Description | | --------------- | ---------------- | ------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | @@ -65,6 +65,77 @@ All flags below are prefixed with `--dataset.camera_encoder.` on the CLI. --- +## Depth streams + +Depth maps (Intel RealSense, Reachy 2) are stored as their **own video streams** alongside the RGB streams. Raw depth (`uint16` millimetres or `float32` metres) can't survive an 8-bit codec, so LeRobot **quantizes** each map to a 12-bit code (`[0, 4095]`) — logarithmically by default, to match the `1/depth` error profile of depth sensors — then packs it into a high-bit-depth pixel format (`gray12le`) and encodes it with a 12-bit codec. + +```mermaid +flowchart LR + A["Raw depth (uint16 mm / float32 m)"] --> B["Clip to depth_min, depth_max"] + B --> C["Quantize to 12-bit code 0–4095 (log or linear)"] + C --> D["Pack into gray12le"] + D --> E["Encode video (hevc Main 12)"] + E --> F[("MP4 + metadata: depth_min/max, shift, use_log")] + F -. "load time (depth_output_unit)" .-> G["Dequantize to mm or m"] + + classDef input fill:#e3f2fd,stroke:#1565c0,color:#0d47a1; + classDef encode fill:#ede7f6,stroke:#5e35b1,color:#311b92; + classDef store fill:#fff8e1,stroke:#f9a825,color:#e65100; + classDef load fill:#e8f5e9,stroke:#2e7d32,color:#1b5e20; + + class A input; + class B,C,D,E encode; + class F store; + class G load; +``` + +Configure the depth pipeline through a parallel **`depth_encoder`** block (`DepthEncoderConfig`). It shares every `RGBEncoderConfig` field (`vcodec`, `pix_fmt`, `crf`, …) and adds four quantizer knobs, set via `--dataset.depth_encoder.`: + +```bash +lerobot-record \ + ... \ + --dataset.depth_encoder.vcodec=hevc \ + --dataset.depth_encoder.depth_min=0.05 \ + --dataset.depth_encoder.depth_max=5.0 \ + --dataset.depth_encoder.use_log=true +``` + +| Parameter | Type | Default | Description | +| --------------- | ------- | ------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------- | +| `vcodec` | `str` | `"hevc"` | HEVC Main 12 (a 12-bit-capable codec, MP4-compatible). | +| `extra_options` | `dict` | `{"x265-params": "lossless=1"}` | **Depth defaults to lossless** (exact round-trip); `crf` is ignored. Pass `extra_options={}` and set `crf` for a smaller lossy stream. | +| `pix_fmt` | `str` | `"gray12le"` | Single-channel 12-bit pixel format used to carry the quantized codes. | +| `depth_min` | `float` | `0.01` | Depth in metres mapped to quantum `0`. Values below are clipped on decode. | +| `depth_max` | `float` | `10.0` | Depth in metres mapped to quantum `4095`. Values above are clipped on decode. | +| `shift` | `float` | `3.5` | Pre-log offset (metres) used in logarithmic quantization for numerical stability near zero. Must satisfy `depth_min + shift > 0`. | +| `use_log` | `bool` | `True` | If `true`, quantize in log-space (recommended for typical depth sensors). Set to `false` for uniform/linear quantization. | + +> [!TIP] +> `depth_min`, `depth_max`, and `shift` are always interpreted in **metres**, regardless of the input depth's unit. Inputs are auto-detected: integer arrays (e.g. `uint16` millimetres straight from a RealSense) are treated as millimetres, floating arrays as metres. +> Pick `depth_min` / `depth_max` to bracket the actual working range of your sensor — quanta outside that range saturate, which can crush detail at the boundaries. + +Depth features are flagged with `"is_depth_map": true` in `meta/info.json`, and their quantizer settings (`video.depth_min`, `video.depth_max`, `video.shift`, `video.use_log`) are persisted — which is what lets depth be **dequantized back to physical units** on load. + +### Output unit at load time + +`depth_encoder` is a **record-time** concern. The unit that depth maps are dequantized to on _load_ (e.g. during training) is set separately by the read-time flag `--dataset.depth_output_unit`: + +```bash +lerobot-train \ + --dataset.repo_id=/ \ + --dataset.depth_output_unit=m \ + --policy.type=act +``` + +| Parameter | Type | Default | Description | +| ------------------- | ----- | ------- | -------------------------------------------------------------------------------------------- | +| `depth_output_unit` | `str` | `"mm"` | Physical unit depth maps are dequantized to on load: `"mm"` (millimetres) or `"m"` (metres). | + +> [!TIP] +> This is purely a decode-time presentation choice — it does **not** alter the stored video or its metadata, so the same dataset can be read as `mm` or `m` without re-encoding. It has no effect on datasets without depth cameras. + +--- + ## Persistence in dataset metadata After the first episode of a video stream is encoded, the encoder configuration is **persisted into the dataset metadata** (`meta/info.json`) under each video feature, alongside the values probed from the file itself. For a video feature `observation.images.`, the layout in `info.json` is: @@ -82,7 +153,7 @@ After the first episode of a video stream is encoded, the encoder configuration "video.pix_fmt": "yuv420p", "video.fps": 30, "video.channels": 3, - "video.is_depth_map": false, + "is_depth_map": false, "video.g": 2, "video.crf": 30, "video.preset": "fast", @@ -97,12 +168,12 @@ After the first episode of a video stream is encoded, the encoder configuration Two sources contribute to the `info` block: -- **Stream-derived** (read back from the encoded MP4 with PyAV): `video.height`, `video.width`, `video.codec`, `video.pix_fmt`, `video.fps`, `video.channels`, `video.is_depth_map`, plus `audio.*` if an audio stream is present. -- **Encoder-derived** (taken from `VideoEncoderConfig`): `video.g`, `video.crf`, `video.preset`, `video.fast_decode`, `video.video_backend`, `video.extra_options`. +- **Stream-derived** (read back from the encoded MP4 with PyAV): `video.height`, `video.width`, `video.codec`, `video.pix_fmt`, `video.fps`, `video.channels`, `is_depth_map`, plus `audio.*` if an audio stream is present. +- **Encoder-derived** (taken from `RGBEncoderConfig` or `DepthEncoderConfig`): `video.g`, `video.crf`, `video.preset`, `video.fast_decode`, `video.video_backend`, `video.extra_options`. This block is populated **once**, from the **first** episode. It assumes every - episode in the dataset was encoded with the same `camera_encoder`. Changing + episode in the dataset was encoded with the same `rgb_encoder`. Changing encoder settings partway through a recording is not supported — the `info.json` will only reflect the parameters used for the first episode. diff --git a/src/lerobot/annotations/steerable_pipeline/frames.py b/src/lerobot/annotations/steerable_pipeline/frames.py index a6c904673..5a6a5879c 100644 --- a/src/lerobot/annotations/steerable_pipeline/frames.py +++ b/src/lerobot/annotations/steerable_pipeline/frames.py @@ -36,7 +36,7 @@ from typing import Any, Protocol import PIL.Image import torch -from lerobot.configs.video import VideoEncoderConfig +from lerobot.configs import RGBEncoderConfig from lerobot.datasets.video_utils import decode_video_frames, reencode_video from .reader import EpisodeRecord, snap_to_frame @@ -164,7 +164,9 @@ class VideoFrameProvider: # only for video-stored cameras. Image-stored cameras (also in # ``camera_keys``) would KeyError, so restrict the list — and the # default — to video keys. - keys = list(self._meta.video_keys) + # Depth cameras are excluded from the annotation pipeline for now. + depth_keys = set(self._meta.depth_keys) + keys = [key for key in self._meta.video_keys if key not in depth_keys] # Last-resort fallback: if metadata didn't surface any video keys but # the caller explicitly named a camera (``--vlm.camera_key=...``), # trust them — the key is by definition known to exist on the dataset. @@ -276,12 +278,12 @@ class VideoFrameProvider: from_timestamp = float(ep[f"videos/{self.camera_key}/from_timestamp"]) to_timestamp = float(ep[f"videos/{self.camera_key}/to_timestamp"]) src = self.root / self._meta.get_video_file_path(record.episode_index, self.camera_key) - encoder = VideoEncoderConfig(vcodec="h264", pix_fmt="yuv420p", g=None, crf=23, preset="ultrafast") + encoder = RGBEncoderConfig(vcodec="h264", pix_fmt="yuv420p", g=None, crf=23, preset="ultrafast") try: reencode_video( src, out_path, - camera_encoder=encoder, + video_encoder=encoder, overwrite=True, start_time_s=from_timestamp, end_time_s=to_timestamp, diff --git a/src/lerobot/async_inference/helpers.py b/src/lerobot/async_inference/helpers.py index 4931c68c5..54f0ca69f 100644 --- a/src/lerobot/async_inference/helpers.py +++ b/src/lerobot/async_inference/helpers.py @@ -105,8 +105,9 @@ def raw_observation_to_observation( def prepare_image(image: torch.Tensor) -> torch.Tensor: - """Minimal preprocessing to turn int8 images to float32 in [0, 1], and create a memory-contiguous tensor""" - image = image.type(torch.float32) / 255 + """Minimal preprocessing to turn RGB uint8 images to float32 in [0, 1], and create a memory-contiguous tensor""" + if image.dtype == torch.uint8: + image = image.type(torch.float32) / 255 image = image.contiguous() return image diff --git a/src/lerobot/cameras/opencv/camera_opencv.py b/src/lerobot/cameras/opencv/camera_opencv.py index b3c20e8dd..e50d24c01 100644 --- a/src/lerobot/cameras/opencv/camera_opencv.py +++ b/src/lerobot/cameras/opencv/camera_opencv.py @@ -436,7 +436,7 @@ class OpenCVCamera(Camera): Internal loop run by the background thread for asynchronous reading. On each iteration: - 1. Reads a color frame + 1. Reads a color frame (blocking call) 2. Stores result in latest_frame and updates timestamp (thread-safe) 3. Sets new_frame_event to notify listeners @@ -485,6 +485,8 @@ class OpenCVCamera(Camera): if self.thread is not None and self.thread.is_alive(): self.thread.join(timeout=2.0) + if self.thread.is_alive(): + logger.warning(f"{self} read thread did not terminate within timeout.") self.thread = None self.stop_event = None diff --git a/src/lerobot/cameras/realsense/camera_realsense.py b/src/lerobot/cameras/realsense/camera_realsense.py index 80008e9f9..29cb1e5e0 100644 --- a/src/lerobot/cameras/realsense/camera_realsense.py +++ b/src/lerobot/cameras/realsense/camera_realsense.py @@ -128,6 +128,7 @@ class RealSenseCamera(Camera): self.fps = config.fps self.color_mode = config.color_mode + self.use_rgb = config.use_rgb self.use_depth = config.use_depth self.warmup_s = config.warmup_s @@ -195,12 +196,15 @@ class RealSenseCamera(Camera): # NOTE(Steven/Caroline): Enforcing at least one second of warmup as RS cameras need a bit of time before the first read. If we don't wait, the first read from the warmup will raise. self.warmup_s = max(self.warmup_s, 1) + warmup_read = self.async_read if self.use_rgb else self.async_read_depth start_time = time.time() while time.time() - start_time < self.warmup_s: - self.async_read(timeout_ms=self.warmup_s * 1000) + warmup_read(timeout_ms=self.warmup_s * 1000) time.sleep(0.1) with self.frame_lock: - if self.latest_color_frame is None or self.use_depth and self.latest_depth_frame is None: + if (self.use_rgb and self.latest_color_frame is None) or ( + self.use_depth and self.latest_depth_frame is None + ): raise ConnectionError(f"{self} failed to capture frames during warmup.") logger.info(f"{self} connected.") @@ -268,13 +272,13 @@ class RealSenseCamera(Camera): ) if len(found_devices) > 1: - serial_numbers = [dev["serial_number"] for dev in found_devices] + serial_numbers = [dev["id"] for dev in found_devices] raise ValueError( f"Multiple RealSense cameras found with name '{name}'. " f"Please use a unique serial number instead. Found SNs: {serial_numbers}" ) - serial_number = str(found_devices[0]["serial_number"]) + serial_number = str(found_devices[0]["id"]) return serial_number def _configure_rs_pipeline_config(self, rs_config: Any) -> None: @@ -282,15 +286,17 @@ class RealSenseCamera(Camera): rs.config.enable_device(rs_config, self.serial_number) if self.width and self.height and self.fps: - rs_config.enable_stream( - rs.stream.color, self.capture_width, self.capture_height, rs.format.rgb8, self.fps - ) + if self.use_rgb: + rs_config.enable_stream( + rs.stream.color, self.capture_width, self.capture_height, rs.format.rgb8, self.fps + ) if self.use_depth: rs_config.enable_stream( rs.stream.depth, self.capture_width, self.capture_height, rs.format.z16, self.fps ) else: - rs_config.enable_stream(rs.stream.color) + if self.use_rgb: + rs_config.enable_stream(rs.stream.color) if self.use_depth: rs_config.enable_stream(rs.stream.depth) @@ -298,8 +304,9 @@ class RealSenseCamera(Camera): def _configure_capture_settings(self) -> None: """Sets fps, width, and height from device stream if not already configured. - Uses the color stream profile to update unset attributes. Handles rotation by - swapping width/height when needed. Original capture dimensions are always stored. + Uses the color stream profile (or the depth stream profile when the color + stream is disabled) to update unset attributes. Handles rotation by swapping + width/height when needed. Original capture dimensions are always stored. Raises: DeviceNotConnectedError: If device is not connected. @@ -308,7 +315,8 @@ class RealSenseCamera(Camera): if self.rs_profile is None: raise RuntimeError(f"{self}: rs_profile must be initialized before use.") - stream = self.rs_profile.get_stream(rs.stream.color).as_video_stream_profile() + rs_stream = rs.stream.color if self.use_rgb else rs.stream.depth + stream = self.rs_profile.get_stream(rs_stream).as_video_stream_profile() if self.fps is None: self.fps = stream.fps() @@ -323,6 +331,14 @@ class RealSenseCamera(Camera): self.width, self.height = actual_width, actual_height self.capture_width, self.capture_height = actual_width, actual_height + def _read(self, read_depth: bool = False) -> NDArray[Any]: + """Shared helper for :meth:`read`/:meth:`read_depth`: wait for a fresh color or depth frame.""" + if self.thread is None or not self.thread.is_alive(): + raise RuntimeError(f"{self} read thread is not running.") + + self.new_frame_event.clear() + return self._async_read(timeout_ms=10000, read_depth=read_depth) + @check_if_not_connected def read_depth(self, timeout_ms: int = 200) -> NDArray[Any]: """ @@ -332,8 +348,8 @@ class RealSenseCamera(Camera): from the camera hardware via the RealSense pipeline. Returns: - np.ndarray: The depth map as a NumPy array (height, width) - of type `np.uint16` (raw depth values in millimeters) and rotation. + np.ndarray: The depth map as a NumPy array (height, width, 1) + of type `np.uint16` (raw depth values in millimeters). Raises: DeviceNotConnectedError: If the camera is not connected. @@ -349,20 +365,7 @@ class RealSenseCamera(Camera): f"Failed to capture depth frame '.read_depth()'. Depth stream is not enabled for {self}." ) - if self.thread is None or not self.thread.is_alive(): - raise RuntimeError(f"{self} read thread is not running.") - - self.new_frame_event.clear() - - _ = self.async_read(timeout_ms=10000) - - with self.frame_lock: - depth_map = self.latest_depth_frame - - if depth_map is None: - raise RuntimeError("No depth frame available. Ensure camera is streaming.") - - return depth_map + return self._read(read_depth=True) def _read_from_hardware(self): if self.rs_pipeline is None: @@ -405,12 +408,10 @@ class RealSenseCamera(Camera): f"{self} read() timeout_ms parameter is deprecated and will be removed in future versions." ) - if self.thread is None or not self.thread.is_alive(): - raise RuntimeError(f"{self} read thread is not running.") + if not self.use_rgb: + raise RuntimeError(f"{self}: cannot read color — camera was configured with use_rgb=False.") - self.new_frame_event.clear() - - frame = self.async_read(timeout_ms=10000) + frame = self._read() read_duration_ms = (time.perf_counter() - start_time) * 1e3 logger.debug(f"{self} read took: {read_duration_ms:.1f}ms") @@ -465,8 +466,8 @@ class RealSenseCamera(Camera): Internal loop run by the background thread for asynchronous reading. On each iteration: - 1. Reads a color frame with 500ms timeout - 2. Stores result in latest_frame and updates timestamp (thread-safe) + 1. Reads a color/depth frame (blocking call with 10s timeout) + 2. Stores result in latest_color_frame/latest_depth_frame and updates timestamp (thread-safe) 3. Sets new_frame_event to notify listeners Stops on DeviceNotConnectedError, logs other errors and continues. @@ -479,19 +480,24 @@ class RealSenseCamera(Camera): while not stop_event.is_set(): try: frame = self._read_from_hardware() - color_frame_raw = frame.get_color_frame() - color_frame = np.asanyarray(color_frame_raw.get_data()) - processed_color_frame = self._postprocess_image(color_frame) + + if self.use_rgb: + color_frame_raw = frame.get_color_frame() + color_frame = np.asanyarray(color_frame_raw.get_data()) + processed_color_frame = self._postprocess_image(color_frame) if self.use_depth: depth_frame_raw = frame.get_depth_frame() depth_frame = np.asanyarray(depth_frame_raw.get_data()) processed_depth_frame = self._postprocess_image(depth_frame, depth_frame=True) + if processed_depth_frame.ndim == 2: # (H, W) -> (H, W, 1) + processed_depth_frame = processed_depth_frame[..., np.newaxis] capture_time = time.perf_counter() with self.frame_lock: - self.latest_color_frame = processed_color_frame + if self.use_rgb: + self.latest_color_frame = processed_color_frame if self.use_depth: self.latest_depth_frame = processed_depth_frame self.latest_timestamp = capture_time @@ -523,6 +529,8 @@ class RealSenseCamera(Camera): if self.thread is not None and self.thread.is_alive(): self.thread.join(timeout=2.0) + if self.thread.is_alive(): # pragma: no cover + logger.warning(f"{self} read thread did not terminate within timeout.") self.thread = None self.stop_event = None @@ -533,7 +541,26 @@ class RealSenseCamera(Camera): self.latest_timestamp = None self.new_frame_event.clear() - # NOTE(Steven): Missing implementation for depth for now + def _async_read(self, timeout_ms: float, read_depth: bool = False) -> NDArray[Any]: + """Shared helper for :meth:`async_read`/:meth:`async_read_depth`: return the latest buffered frame.""" + if self.thread is None or not self.thread.is_alive(): + raise RuntimeError(f"{self} read thread is not running.") + + if not self.new_frame_event.wait(timeout=timeout_ms / 1000.0): + raise TimeoutError( + f"Timed out waiting for frame from camera {self} after {timeout_ms} ms. " + f"Read thread alive: {self.thread.is_alive()}." + ) + + with self.frame_lock: + frame = self.latest_depth_frame if read_depth else self.latest_color_frame + self.new_frame_event.clear() + + if frame is None: + raise RuntimeError(f"Internal error: Event set but no frame available for {self}.") + + return frame + @check_if_not_connected def async_read(self, timeout_ms: float = 200) -> NDArray[Any]: """ @@ -558,25 +585,31 @@ class RealSenseCamera(Camera): RuntimeError: If the background thread died unexpectedly or another error occurs. """ + if not self.use_rgb: + raise RuntimeError(f"{self}: cannot read color — camera was configured with use_rgb=False.") + + return self._async_read(timeout_ms=timeout_ms) + + def _read_latest(self, max_age_ms: int, read_depth: bool = False) -> NDArray[Any]: + """Shared helper for :meth:`read_latest`/:meth:`read_latest_depth`: peek the latest buffered frame.""" if self.thread is None or not self.thread.is_alive(): raise RuntimeError(f"{self} read thread is not running.") - if not self.new_frame_event.wait(timeout=timeout_ms / 1000.0): - raise TimeoutError( - f"Timed out waiting for frame from camera {self} after {timeout_ms} ms. " - f"Read thread alive: {self.thread.is_alive()}." - ) - with self.frame_lock: - frame = self.latest_color_frame - self.new_frame_event.clear() + frame = self.latest_depth_frame if read_depth else self.latest_color_frame + timestamp = self.latest_timestamp - if frame is None: - raise RuntimeError(f"Internal error: Event set but no frame available for {self}.") + if frame is None or timestamp is None: + raise RuntimeError(f"{self} has not captured any frames yet.") + + age_ms = (time.perf_counter() - timestamp) * 1e3 + if age_ms > max_age_ms: + raise TimeoutError( + f"{self} latest frame is too old: {age_ms:.1f} ms (max allowed: {max_age_ms} ms)." + ) return frame - # NOTE(Steven): Missing implementation for depth for now @check_if_not_connected def read_latest(self, max_age_ms: int = 500) -> NDArray[Any]: """Return the most recent (color) frame captured immediately (Peeking). @@ -593,24 +626,48 @@ class RealSenseCamera(Camera): DeviceNotConnectedError: If the camera is not connected. RuntimeError: If the camera is connected but has not captured any frames yet. """ + if not self.use_rgb: + raise RuntimeError(f"{self}: cannot read color — camera was configured with use_rgb=False.") - if self.thread is None or not self.thread.is_alive(): - raise RuntimeError(f"{self} read thread is not running.") + return self._read_latest(max_age_ms=max_age_ms) - with self.frame_lock: - frame = self.latest_color_frame - timestamp = self.latest_timestamp + @check_if_not_connected + def async_read_depth(self, timeout_ms: float = 200) -> NDArray[np.uint16]: + """Read the latest depth frame asynchronously, in millimeters. - if frame is None or timestamp is None: - raise RuntimeError(f"{self} has not captured any frames yet.") + Mirrors :meth:`async_read` but returns the depth stream rather than the + color stream. Output is ``np.uint16`` of shape ``(H, W, 1)``, where each + pixel is the distance from the sensor in millimeters. - age_ms = (time.perf_counter() - timestamp) * 1e3 - if age_ms > max_age_ms: - raise TimeoutError( - f"{self} latest frame is too old: {age_ms:.1f} ms (max allowed: {max_age_ms} ms)." - ) + Raises: + DeviceNotConnectedError: If the camera is not connected. + RuntimeError: If ``use_depth`` is ``False`` for this camera, or if + the background read thread is not running. + TimeoutError: If no frame becomes available within ``timeout_ms``. + """ + if not self.use_depth: + raise RuntimeError(f"{self}: cannot read depth — camera was configured with use_depth=False.") - return frame + return self._async_read(timeout_ms=timeout_ms, read_depth=True) + + @check_if_not_connected + def read_latest_depth(self, max_age_ms: int = 500) -> NDArray[Any]: + """Return the most recent depth frame in millimeters (peeking). + + Non-blocking counterpart of :meth:`read_latest` for the depth stream. + Output is ``np.uint16`` of shape ``(H, W, 1)``, where each pixel is the + distance from the sensor in millimeters. + + Raises: + DeviceNotConnectedError: If the camera is not connected. + RuntimeError: If ``use_depth`` is ``False`` for this camera, or if + no depth frame has been captured yet. + TimeoutError: If the latest depth frame is older than ``max_age_ms``. + """ + if not self.use_depth: + raise RuntimeError(f"{self}: cannot read depth — camera was configured with use_depth=False.") + + return self._read_latest(max_age_ms=max_age_ms, read_depth=True) def disconnect(self) -> None: """ diff --git a/src/lerobot/cameras/realsense/configuration_realsense.py b/src/lerobot/cameras/realsense/configuration_realsense.py index 71b083b00..018675195 100644 --- a/src/lerobot/cameras/realsense/configuration_realsense.py +++ b/src/lerobot/cameras/realsense/configuration_realsense.py @@ -42,12 +42,14 @@ class RealSenseCameraConfig(CameraConfig): height: Requested frame height in pixels for the color stream. serial_number_or_name: Unique serial number or human-readable name to identify the camera. color_mode: Color mode for image output (RGB or BGR). Defaults to RGB. + use_rgb: Whether to enable the color stream. Defaults to True. use_depth: Whether to enable depth stream. Defaults to False. rotation: Image rotation setting (0°, 90°, 180°, or 270°). Defaults to no rotation. warmup_s: Time reading frames before returning from connect (in seconds) Note: - Either name or serial_number must be specified. + - At least one of `use_rgb` or `use_depth` must be enabled. - Depth stream configuration (if enabled) will use the same FPS as the color stream. - The actual resolution and FPS may be adjusted by the camera to the nearest supported mode. - For `fps`, `width` and `height`, either all of them need to be set, or none of them. @@ -55,6 +57,7 @@ class RealSenseCameraConfig(CameraConfig): serial_number_or_name: str color_mode: ColorMode = ColorMode.RGB + use_rgb: bool = True use_depth: bool = False rotation: Cv2Rotation = Cv2Rotation.NO_ROTATION warmup_s: int = 1 @@ -63,6 +66,9 @@ class RealSenseCameraConfig(CameraConfig): self.color_mode = ColorMode(self.color_mode) self.rotation = Cv2Rotation(self.rotation) + if not self.use_rgb and not self.use_depth: + raise ValueError("At least one of `use_rgb` or `use_depth` must be enabled.") + values = (self.fps, self.width, self.height) if any(v is not None for v in values) and any(v is None for v in values): raise ValueError( diff --git a/src/lerobot/cameras/zmq/camera_zmq.py b/src/lerobot/cameras/zmq/camera_zmq.py index f3df17814..cd32a117b 100644 --- a/src/lerobot/cameras/zmq/camera_zmq.py +++ b/src/lerobot/cameras/zmq/camera_zmq.py @@ -293,6 +293,8 @@ class ZMQCamera(Camera): if self.thread is not None and self.thread.is_alive(): self.thread.join(timeout=2.0) + if self.thread.is_alive(): + logger.warning(f"{self} read thread did not terminate within timeout.") self.thread = None self.stop_event = None diff --git a/src/lerobot/configs/__init__.py b/src/lerobot/configs/__init__.py index be4491811..fa5942129 100644 --- a/src/lerobot/configs/__init__.py +++ b/src/lerobot/configs/__init__.py @@ -33,10 +33,15 @@ from .types import ( RTCAttentionSchedule, ) from .video import ( + DEFAULT_DEPTH_UNIT, VALID_VIDEO_CODECS, VIDEO_ENCODER_INFO_KEYS, + DepthEncoderConfig, + RGBEncoderConfig, VideoEncoderConfig, - camera_encoder_defaults, + depth_encoder_defaults, + encoder_config_from_video_info, + rgb_encoder_defaults, ) __all__ = [ @@ -57,9 +62,15 @@ __all__ = [ "WandBConfig", "load_recipe", "VideoEncoderConfig", + "RGBEncoderConfig", + "DepthEncoderConfig", # Defaults - "camera_encoder_defaults", + "rgb_encoder_defaults", + "depth_encoder_defaults", + # Factories + "encoder_config_from_video_info", # Constants + "DEFAULT_DEPTH_UNIT", "VALID_VIDEO_CODECS", "VIDEO_ENCODER_INFO_KEYS", ] diff --git a/src/lerobot/configs/dataset.py b/src/lerobot/configs/dataset.py index c40c0fae2..7d30ca038 100644 --- a/src/lerobot/configs/dataset.py +++ b/src/lerobot/configs/dataset.py @@ -18,7 +18,7 @@ from dataclasses import dataclass, field from datetime import datetime from pathlib import Path -from .video import VideoEncoderConfig, camera_encoder_defaults +from .video import DepthEncoderConfig, RGBEncoderConfig, depth_encoder_defaults, rgb_encoder_defaults @dataclass @@ -58,8 +58,10 @@ class DatasetRecordConfig: # Set to 1 for immediate encoding (default behavior), or higher for batched encoding video_encoding_batch_size: int = 1 # Video encoder settings for camera MP4s (codec, quality, GOP, etc.). Tuned via CLI nested keys, - # e.g. ``--dataset.camera_encoder.vcodec=h264`` (see ``VideoEncoderConfig``). - camera_encoder: VideoEncoderConfig = field(default_factory=camera_encoder_defaults) + # e.g. ``--dataset.rgb_encoder.vcodec=h264`` (see ``RGBEncoderConfig``). + rgb_encoder: RGBEncoderConfig = field(default_factory=rgb_encoder_defaults) + # Video encoder settings for depth-map MP4s (codec, quality, GOP, etc.). Tuned via CLI nested keys. + depth_encoder: DepthEncoderConfig = field(default_factory=depth_encoder_defaults) # Enable streaming video encoding: encode frames in real-time during capture instead # of writing PNG images first. Makes save_episode() near-instant. More info in the documentation: https://huggingface.co/docs/lerobot/streaming_video_encoding streaming_encoding: bool = False diff --git a/src/lerobot/configs/default.py b/src/lerobot/configs/default.py index 9b5433005..4f24b9dac 100644 --- a/src/lerobot/configs/default.py +++ b/src/lerobot/configs/default.py @@ -19,6 +19,8 @@ from dataclasses import dataclass, field from lerobot.transforms import ImageTransformsConfig from lerobot.utils.import_utils import get_safe_default_video_backend +from .video import DEFAULT_DEPTH_UNIT, DEPTH_METER_UNIT, DEPTH_MILLIMETER_UNIT + @dataclass class DatasetConfig: @@ -35,14 +37,21 @@ class DatasetConfig: revision: str | None = None use_imagenet_stats: bool = True video_backend: str = field(default_factory=get_safe_default_video_backend) - # When True, video frames are returned as uint8 tensors (0-255) instead of float32 (0.0-1.0). + # When True, RGB video frames are returned as uint8 tensors (0-255) instead of float32 (0.0-1.0). # This reduces memory and speeds up DataLoader IPC. The training pipeline handles the conversion. return_uint8: bool = False + # Physical unit depth maps are dequantized to at load time: "mm" (millimeters) or "m" (metres). + # Has no effect on datasets without depth cameras. + depth_output_unit: str = DEFAULT_DEPTH_UNIT streaming: bool = False # Fraction of episodes held out per task for offline evaluation (0.0 = disabled). eval_split: float = 0.0 def __post_init__(self) -> None: + if self.depth_output_unit not in (DEPTH_METER_UNIT, DEPTH_MILLIMETER_UNIT): + raise ValueError( + f"depth_output_unit must be '{DEPTH_METER_UNIT}' or '{DEPTH_MILLIMETER_UNIT}', got {self.depth_output_unit!r}" + ) if not (0.0 <= self.eval_split < 1.0): raise ValueError(f"eval_split must be in [0.0, 1.0), got {self.eval_split}") if self.episodes is not None: diff --git a/src/lerobot/configs/video.py b/src/lerobot/configs/video.py index bf2471453..3ea834508 100644 --- a/src/lerobot/configs/video.py +++ b/src/lerobot/configs/video.py @@ -20,7 +20,7 @@ from __future__ import annotations import logging from dataclasses import dataclass, field -from typing import Any +from typing import Any, ClassVar, Self from lerobot.utils.import_utils import require_package @@ -40,7 +40,6 @@ VALID_VIDEO_CODECS: frozenset[str] = frozenset({"h264", "hevc", "libsvtav1", "au # Aliases for legacy video codec names. VIDEO_CODECS_ALIASES: dict[str, str] = {"av1": "libsvtav1"} - LIBSVTAV1_DEFAULT_PRESET: int = 12 # Keys persisted under ``features[*]["info"]`` as ``video.`` (from :class:`VideoEncoderConfig`). @@ -52,40 +51,45 @@ VIDEO_ENCODER_INFO_KEYS: frozenset[str] = frozenset( f"video.{name}" for name in VIDEO_ENCODER_INFO_FIELD_NAMES ) +# Default depth quantization and encoding parameters. +DEPTH_QUANT_BITS: int = 12 +DEPTH_QMAX: int = (1 << DEPTH_QUANT_BITS) - 1 # 4095 + +DEFAULT_DEPTH_MIN: float = 0.01 +DEFAULT_DEPTH_MAX: float = 10.0 +DEFAULT_DEPTH_SHIFT: float = 3.5 +DEFAULT_DEPTH_USE_LOG: bool = True +DEFAULT_DEPTH_PIX_FMT: str = "gray12le" + +DEPTH_METER_UNIT: str = "m" +DEPTH_MILLIMETER_UNIT: str = "mm" +DEFAULT_DEPTH_UNIT: str = DEPTH_MILLIMETER_UNIT + +# Depth-specific tuning fields persisted under ``features[*]["info"]`` as ``video.``. +DEPTH_ENCODER_INFO_FIELD_NAMES: frozenset[str] = frozenset({"depth_min", "depth_max", "shift", "use_log"}) + @dataclass class VideoEncoderConfig: - """Video encoder configuration. + """Video encoder configuration.""" - Attributes: - vcodec: Video encoder name. ``"auto"`` is resolved during - construction (HW encoder if available, else ``libsvtav1``). - pix_fmt: Pixel format (e.g. ``"yuv420p"``). - g: GOP size (keyframe interval). - crf: Quality level — mapped to the native quality parameter of the - codec (``crf`` for software, ``qp`` for NVENC/VAAPI, - ``q:v`` for VideoToolbox, ``global_quality`` for QSV). - preset: Speed/quality preset. Accepted type is per-codec. - fast_decode: Fast-decode tuning. For ``libsvtav1`` this is a level (0-2) - embedded in ``svtav1-params``. For ``h264`` and ``hevc`` non-zero values - set ``tune=fastdecode``. Ignored for other codecs. - video_backend: Python to be used for encoding. Only ``"pyav"`` - is currently supported. - extra_options: Free-form dictionary of additional video encoder options - (e.g. ``{"tune": "film", "profile:v": "high", "bf": 2}``). - """ - - vcodec: str = "libsvtav1" # TODO(CarolinePascal): rename to codec ? - pix_fmt: str = "yuv420p" - g: int | None = 2 - crf: int | float | None = 30 - preset: int | str | None = None - fast_decode: int = 0 + vcodec: str = "libsvtav1" # Video codec name. "auto" picks a hardware codec if available, else libsvtav1. + pix_fmt: str = "yuv420p" # Pixel format (e.g. yuv420p). + g: int | None = 2 # GOP size (keyframe interval). + crf: int | float | None = 30 # Quality level. Lower means better quality and larger files. + preset: int | str | None = None # Speed/quality preset. Accepted values are codec-specific. + fast_decode: int = 0 # Fast-decode tuning. Accepted values are codec-specific, 0 disables it. # TODO(CarolinePascal): add torchcodec support + find a way to unify the # two backends (encoding and decoding). - video_backend: str = "pyav" + video_backend: str = "pyav" # Encoding backend. Only "pyav" is currently supported. + # Extra codec options merged last, e.g. {"tune": "film"}. extra_options: dict[str, Any] = field(default_factory=dict) + # Source-data channel count this encoder is expected to handle. ``None`` + # disables the pix_fmt channel-count check; concrete subclasses set it + # (3 for RGB, 1 for depth, etc.). + _DEFAULT_CHANNELS: ClassVar[int | None] = None + def __post_init__(self) -> None: self.resolve_vcodec() # Empty-constructor ergonomics: ``VideoEncoderConfig()`` must "just work". @@ -94,9 +98,9 @@ class VideoEncoderConfig: self.validate() @classmethod - def from_video_info(cls, video_info: dict | None) -> VideoEncoderConfig: - """Reconstruct a :class:`VideoEncoderConfig` from a video feature's ``info`` block. - Missing or ``None`` values fall back to the class defaults. + def _kwargs_from_video_info(cls, video_info: dict | None) -> dict[str, Any]: + """Parse the ``video.*`` keys of a feature ``info`` block into + constructor kwargs. """ video_info = video_info or {} kwargs: dict[str, Any] = {} @@ -115,7 +119,15 @@ class VideoEncoderConfig: continue kwargs[field_name] = value - return cls(**kwargs) + return kwargs + + @classmethod + def from_video_info(cls, video_info: dict | None) -> Self: + """Reconstruct an encoder config from a video feature's ``info`` block. + + Missing or ``None`` values fall back to the class defaults. + """ + return cls(**cls._kwargs_from_video_info(video_info)) def detect_available_encoders(self, encoders: list[str] | str) -> list[str]: """Return the subset of available encoders based on the specified video backend. @@ -138,7 +150,9 @@ class VideoEncoderConfig: require_package("av", extra="dataset") from lerobot.datasets import check_video_encoder_parameters_pyav - check_video_encoder_parameters_pyav(self.vcodec, self.pix_fmt, self.get_codec_options()) + check_video_encoder_parameters_pyav( + self.vcodec, self.pix_fmt, self.get_codec_options(), channels=self._DEFAULT_CHANNELS + ) def resolve_vcodec(self) -> None: """Check ``vcodec`` and, when it is ``"auto"``, pick a concrete encoder. @@ -230,6 +244,79 @@ class VideoEncoderConfig: return opts -def camera_encoder_defaults() -> VideoEncoderConfig: - """Return a :class:`VideoEncoderConfig` with RGB-camera defaults.""" - return VideoEncoderConfig() +@dataclass +class RGBEncoderConfig(VideoEncoderConfig): + """Encoder configuration for RGB camera streams. + + Identical to :class:`VideoEncoderConfig` but declares the 3-channel + source-data layout so ``pix_fmt`` is validated against RGB inputs. + """ + + _DEFAULT_CHANNELS: ClassVar[int] = 3 + + +def rgb_encoder_defaults() -> RGBEncoderConfig: + """Return a :class:`RGBEncoderConfig` with RGB-camera defaults.""" + return RGBEncoderConfig() + + +@dataclass +class DepthEncoderConfig(VideoEncoderConfig): + """Encoder configuration for depth-map streams. + + Inherits the full :class:`VideoEncoderConfig` surface (codec, GOP, CRF, + preset, ``extra_options``…) and adds the parameters of the depth quantizer. + Defaults flip ``vcodec`` to ``"hevc"`` (Main 12 profile) and ``pix_fmt`` to + ``"gray12le"``. + """ + + vcodec: str = "hevc" # Video codec name. Defaults to HEVC Main 12 (a 12-bit-capable codec). + pix_fmt: str = "gray12le" # Pixel format. Defaults to 12-bit grayscale. + extra_options: dict[str, Any] = field(default_factory=lambda: {"x265-params": "lossless=1"}) + + depth_min: float = DEFAULT_DEPTH_MIN # Minimum depth in meters, mapped to the lowest quantum. + depth_max: float = DEFAULT_DEPTH_MAX # Maximum depth in meters, mapped to the highest quantum. + shift: float = DEFAULT_DEPTH_SHIFT # Pre-log offset in meters for numerical stability near zero. + use_log: bool = DEFAULT_DEPTH_USE_LOG # Use logarithmic quantization (True) or linear (False). + + _DEFAULT_CHANNELS: ClassVar[int] = 1 + + @classmethod + def _kwargs_from_video_info(cls, video_info: dict | None) -> dict[str, Any]: + """Layer the depth-specific tuning (``depth_min`` / ``depth_max`` / + ``shift`` / ``use_log``) on top of the base parser. Missing keys + fall back to the class defaults. + """ + kwargs = super()._kwargs_from_video_info(video_info) + video_info = video_info or {} + for name in DEPTH_ENCODER_INFO_FIELD_NAMES: + value = video_info.get(f"video.{name}") + if value is not None: + kwargs[name] = value + return kwargs + + +def depth_encoder_defaults() -> DepthEncoderConfig: + """Return a :class:`DepthEncoderConfig` with depth-camera defaults.""" + return DepthEncoderConfig() + + +def encoder_config_from_video_info(video_info: dict | None) -> VideoEncoderConfig: + """Build the appropriate encoder config from a feature's ``info`` block. + + Dispatches to :class:`DepthEncoderConfig` when the dict marks the feature + as a depth map and to :class:`RGBEncoderConfig` + otherwise. + + Args: + video_info: A feature's ``info`` dict as persisted in ``info.json``, + or ``None`` (treated as an empty dict). + + Returns: + A :class:`DepthEncoderConfig` for depth features, otherwise a + :class:`RGBEncoderConfig`. + """ + video_info = video_info or {} + is_depth = bool(video_info.get("is_depth_map") or video_info.get("video.is_depth_map")) + cls: type[VideoEncoderConfig] = DepthEncoderConfig if is_depth else RGBEncoderConfig + return cls.from_video_info(video_info) diff --git a/src/lerobot/datasets/compute_stats.py b/src/lerobot/datasets/compute_stats.py index 09765c130..88f7ea226 100644 --- a/src/lerobot/datasets/compute_stats.py +++ b/src/lerobot/datasets/compute_stats.py @@ -242,12 +242,12 @@ def sample_images(image_paths: list[str]) -> np.ndarray: images = None for i, idx in enumerate(sampled_indices): path = image_paths[idx] - # we load as uint8 to reduce memory usage + # we load RGB images as uint8 to reduce memory usage; depth keeps its native dtype img = load_image_as_numpy(path, dtype=np.uint8, channel_first=True) img = auto_downsample_height_width(img) if images is None: - images = np.empty((len(sampled_indices), *img.shape), dtype=np.uint8) + images = np.empty((len(sampled_indices), *img.shape), dtype=img.dtype) images[i] = img @@ -506,8 +506,10 @@ def compute_episode_stats( Each statistics dictionary contains min, max, mean, std, count, and quantiles. Note: - Image statistics are normalized to [0,1] range and have shape (3,1,1) for - per-channel values when dtype is 'image' or 'video'. + For 'image'/'video' features, stats are computed per channel and kept with a + leading channel axis (e.g. shape (3, 1, 1) for RGB). RGB stats are divided by + 255 to land in [0, 1]; depth maps (features flagged with ``is_depth_map``) skip + this rescaling and remain in their stored units. """ if quantile_list is None: quantile_list = DEFAULT_QUANTILES @@ -531,8 +533,12 @@ def compute_episode_stats( ) if features[key]["dtype"] in ["image", "video"]: + normalization_factor = ( + 255.0 if not (features[key].get("info") or {}).get("is_depth_map", False) else 1.0 + ) ep_stats[key] = { - k: v if k == "count" else np.squeeze(v / 255.0, axis=0) for k, v in ep_stats[key].items() + k: v if k == "count" else np.squeeze(v / normalization_factor, axis=0) + for k, v in ep_stats[key].items() } return ep_stats @@ -552,8 +558,10 @@ def _validate_stat_value(value: np.ndarray, key: str, feature_key: str) -> None: if key == "count" and value.shape != (1,): raise ValueError(f"Shape of 'count' must be (1), but is {value.shape} instead.") - if "image" in feature_key and key != "count" and value.shape != (3, 1, 1): - raise ValueError(f"Shape of quantile '{key}' must be (3,1,1), but is {value.shape} instead.") + if "image" in feature_key and key != "count" and value.shape not in ((3, 1, 1), (1, 1, 1)): + raise ValueError( + f"Shape of quantile '{key}' must be (3,1,1) or (1,1,1) but is {value.shape} instead." + ) def _assert_type_and_shape(stats_list: list[dict[str, dict]]): diff --git a/src/lerobot/datasets/dataset_metadata.py b/src/lerobot/datasets/dataset_metadata.py index b496e4f65..ea329668c 100644 --- a/src/lerobot/datasets/dataset_metadata.py +++ b/src/lerobot/datasets/dataset_metadata.py @@ -14,7 +14,8 @@ # See the License for the specific language governing permissions and # limitations under the License. import contextlib -from collections.abc import Callable +import logging +from collections.abc import Callable, Iterable from copy import deepcopy from pathlib import Path @@ -338,6 +339,25 @@ class LeRobotDatasetMetadata: """Keys to access visual modalities stored as videos.""" return [key for key, ft in self.features.items() if ft["dtype"] == "video"] + @property + def depth_keys(self) -> list[str]: + """Keys to access depth-map modalities stored as videos or images. + + A depth key is a feature whose ``info`` dict carries ``"is_depth_map": True`` + (or the legacy ``"video.is_depth_map"`` inside ``info`` or ``video_info``). + """ + + def _is_depth(ft: dict) -> bool: + info = ft.get("info") or {} + video_info = ft.get("video_info") or {} + return ( + info.get("is_depth_map", False) + or info.get("video.is_depth_map", False) + or video_info.get("video.is_depth_map", False) + ) + + return [key for key, ft in self.features.items() if _is_depth(ft)] + @property def camera_keys(self) -> list[str]: """Keys to access visual modalities (regardless of their storage method).""" @@ -581,29 +601,48 @@ class LeRobotDatasetMetadata: def update_video_info( self, video_key: str | None = None, - camera_encoder: VideoEncoderConfig | None = None, + video_encoder: VideoEncoderConfig | None = None, + preserve_keys: Iterable[str] | None = None, ) -> None: - """Populate per-feature video info in ``info.json``. + """Populate or refresh per-feature video info in ``info.json``. Warning: this function writes info from first episode videos, implicitly assuming that all videos have been encoded the same way. Also, this means it assumes the first episode exists. + Always re-probes the videos and overwrites existing info for every recomputed + key. ``preserve_keys`` lists keys whose existing values must be kept (e.g. + data-intrinsic entries like ``is_depth_map`` and depth quantization params) + instead of being recomputed. + Args: video_key: If provided, only update this video key. Otherwise update all video keys in the dataset. - camera_encoder: Encoder configuration used to produce the + video_encoder: Encoder configuration used to produce the videos. When provided, its fields are recorded as ``video.`` entries alongside the stream-derived ``video.*`` entries (see :func:`get_video_info`). + preserve_keys: Keys whose existing values are kept instead of being + recomputed. ``None`` (default) recomputes every key. """ if video_key is not None and video_key not in self.video_keys: raise ValueError(f"Video key {video_key} not found in dataset") video_keys = [video_key] if video_key is not None else self.video_keys + preserve_set = set(preserve_keys or ()) for key in video_keys: - if not self.features[key].get("info", None): - video_path = self.root / self.video_path.format(video_key=key, chunk_index=0, file_index=0) - self.info.features[key]["info"] = get_video_info(video_path, camera_encoder=camera_encoder) + existing = self.features[key].get("info") or {} + video_path = self.root / self.video_path.format(video_key=key, chunk_index=0, file_index=0) + new_info = get_video_info(video_path, video_encoder=video_encoder) + # Drop preserved keys so the existing values win on merge. + new_info = {k: v for k, v in new_info.items() if k not in preserve_set} + merged = {**existing, **new_info} + # Migrate the legacy depth marker to the canonical key. + if "video.is_depth_map" in merged: + logging.warning( + f"Migrating legacy 'video.is_depth_map' to 'is_depth_map' for feature {key!r}." + ) + merged.setdefault("is_depth_map", merged.pop("video.is_depth_map")) + self.info.features[key]["info"] = merged def update_chunk_settings( self, diff --git a/src/lerobot/datasets/dataset_reader.py b/src/lerobot/datasets/dataset_reader.py index d7289ac48..e8e07301e 100644 --- a/src/lerobot/datasets/dataset_reader.py +++ b/src/lerobot/datasets/dataset_reader.py @@ -22,7 +22,10 @@ from pathlib import Path import datasets import torch +from lerobot.configs import DEFAULT_DEPTH_UNIT, DepthEncoderConfig + from .dataset_metadata import LeRobotDatasetMetadata +from .depth_utils import dequantize_depth from .feature_utils import ( check_delta_timestamps, get_delta_indices, @@ -51,6 +54,7 @@ class DatasetReader: delta_timestamps: dict[str, list[float]] | None, image_transforms: Callable | None, return_uint8: bool = False, + depth_output_unit: str = DEFAULT_DEPTH_UNIT, ): """Initialize the reader with metadata, filtering, and transform config. @@ -68,6 +72,10 @@ class DatasetReader: relative timestamp offsets for temporal context windows. image_transforms: Optional torchvision v2 transform applied to visual features. + return_uint8: If True, return RGB video frames as raw uint8 tensors + instead of normalized float32. + depth_output_unit: Physical unit depth maps are dequantized to + (``"m"`` or ``"mm"``). Defaults to ``"mm"``. """ self._meta = meta self.root = root @@ -78,6 +86,7 @@ class DatasetReader: raise TypeError("image_transforms must be callable or None.") self._image_transforms = image_transforms self._return_uint8 = return_uint8 + self._depth_output_unit = depth_output_unit self.hf_dataset: datasets.Dataset | None = None self._absolute_to_relative_idx: dict[int, int] | None = None @@ -88,6 +97,11 @@ class DatasetReader: check_delta_timestamps(delta_timestamps, meta.fps, tolerance_s) self.delta_indices = get_delta_indices(delta_timestamps, meta.fps) + self._depth_encoder_configs: dict[str, DepthEncoderConfig] = { + vid_key: DepthEncoderConfig.from_video_info(self._meta.features[vid_key].get("info")) + for vid_key in self._meta.depth_keys + } + def set_image_transforms(self, image_transforms: Callable | None) -> None: """Replace the transform applied to visual observations.""" if image_transforms is not None and not callable(image_transforms): @@ -259,7 +273,18 @@ class DatasetReader: self._tolerance_s, self._video_backend, return_uint8=self._return_uint8, + is_depth=vid_key in self._meta.depth_keys, ) + if vid_key in self._meta.depth_keys: + depth_encoder = self._depth_encoder_configs[vid_key] + frames = dequantize_depth( + frames, + depth_min=depth_encoder.depth_min, + depth_max=depth_encoder.depth_max, + shift=depth_encoder.shift, + use_log=depth_encoder.use_log, + output_unit=self._depth_output_unit, + ) return vid_key, frames.squeeze(0) items = list(query_timestamps.items()) @@ -299,8 +324,9 @@ class DatasetReader: item = {**video_frames, **item} if self._image_transforms is not None: - image_keys = self._meta.camera_keys - for cam in image_keys: + for cam in self._meta.camera_keys: + if cam in self._meta.depth_keys: + continue item[cam] = self._image_transforms(item[cam]) # Add task as a string diff --git a/src/lerobot/datasets/dataset_tools.py b/src/lerobot/datasets/dataset_tools.py index 9aca859b4..31e075d7c 100644 --- a/src/lerobot/datasets/dataset_tools.py +++ b/src/lerobot/datasets/dataset_tools.py @@ -37,7 +37,15 @@ import pyarrow.parquet as pq import torch from tqdm import tqdm -from lerobot.configs import VideoEncoderConfig, camera_encoder_defaults +from lerobot.configs import ( + DepthEncoderConfig, + RGBEncoderConfig, + VideoEncoderConfig, + depth_encoder_defaults, + encoder_config_from_video_info, + rgb_encoder_defaults, +) +from lerobot.configs.video import DEPTH_ENCODER_INFO_FIELD_NAMES from lerobot.utils.constants import ACTION, HF_LEROBOT_HOME, OBS_IMAGE, OBS_STATE from lerobot.utils.utils import flatten_dict @@ -48,6 +56,7 @@ from .compute_stats import ( compute_relative_action_stats, ) from .dataset_metadata import LeRobotDatasetMetadata +from .image_writer import write_image from .io_utils import ( get_parquet_file_size_in_mb, load_episodes, @@ -62,12 +71,13 @@ from .utils import ( DEFAULT_DATA_FILE_SIZE_IN_MB, DEFAULT_DATA_PATH, DEFAULT_EPISODES_PATH, + DEPTH_FILE_PATTERN, + IMAGE_FILE_PATTERN, VIDEO_DIR, update_chunk_file_indices, ) from .video_utils import ( encode_video_frames, - get_video_info, reencode_video, ) @@ -601,7 +611,7 @@ def _keep_episodes_from_video_with_av( output_path: Path, episodes_to_keep: list[tuple[int, int]], fps: float, - camera_encoder: VideoEncoderConfig, + video_encoder: VideoEncoderConfig, ) -> None: """Keep only specified episodes from a video file using PyAV. @@ -615,7 +625,7 @@ def _keep_episodes_from_video_with_av( Ranges are half-open intervals: [start_frame, end_frame), where start_frame is inclusive and end_frame is exclusive. fps: Frame rate of the video. - camera_encoder: Video encoder settings used to re-encode the kept frames. + video_encoder: Video encoder settings used to re-encode the kept frames. """ from fractions import Fraction @@ -640,13 +650,13 @@ def _keep_episodes_from_video_with_av( # Convert fps to Fraction for PyAV compatibility. fps_fraction = Fraction(fps).limit_denominator(1000) - codec_options = camera_encoder.get_codec_options(as_strings=True) - v_out = out.add_stream(camera_encoder.vcodec, rate=fps_fraction, options=codec_options) + codec_options = video_encoder.get_codec_options(as_strings=True) + v_out = out.add_stream(video_encoder.vcodec, rate=fps_fraction, options=codec_options) # PyAV type stubs don't distinguish video streams from audio/subtitle streams. v_out.width = v_in.codec_context.width v_out.height = v_in.codec_context.height - v_out.pix_fmt = camera_encoder.pix_fmt + v_out.pix_fmt = video_encoder.pix_fmt # Set time_base to match the frame rate for proper timestamp handling. v_out.time_base = Fraction(1, int(fps)) @@ -733,7 +743,7 @@ def _copy_and_reindex_videos( for video_key in src_dataset.meta.video_keys: logging.info(f"Processing videos for {video_key}") - camera_encoder = VideoEncoderConfig.from_video_info( + video_encoder = encoder_config_from_video_info( src_dataset.meta.info.features.get(video_key, {}).get("info") ) @@ -817,7 +827,7 @@ def _copy_and_reindex_videos( dst_video_path, episodes_to_keep_ranges, src_dataset.meta.fps, - camera_encoder, + video_encoder, ) cumulative_ts = 0.0 @@ -874,11 +884,11 @@ def _copy_and_reindex_episodes_metadata( episode_meta.update(video_metadata[new_idx]) # Extract episode statistics from parquet metadata. - # Note (maractingi): When pandas/pyarrow serializes numpy arrays with shape (3, 1, 1) to parquet, + # When pandas/pyarrow serializes numpy arrays with shape (C, 1, 1) to parquet, # they are being deserialized as nested object arrays like: # array([array([array([0.])]), array([array([0.])]), array([array([0.])])]) # This happens particularly with image/video statistics. We need to detect and flatten - # these nested structures back to proper (3, 1, 1) arrays so aggregate_stats can process them. + # these nested structures back to proper (C, 1, 1) arrays so aggregate_stats can process them. episode_stats = {} for key in src_episode_full: if key.startswith("stats/"): @@ -894,15 +904,16 @@ def _copy_and_reindex_episodes_metadata( if feature_name in src_dataset.meta.features: feature_dtype = src_dataset.meta.features[feature_name]["dtype"] if feature_dtype in ["image", "video"] and stat_name != "count": + # Stats are channel-first (C, 1, 1) if isinstance(value, np.ndarray) and value.dtype == object: flat_values = [] for item in value: while isinstance(item, np.ndarray): item = item.flatten()[0] flat_values.append(item) - value = np.array(flat_values, dtype=np.float64).reshape(3, 1, 1) - elif isinstance(value, np.ndarray) and value.shape == (3,): - value = value.reshape(3, 1, 1) + value = np.array(flat_values, dtype=np.float64).reshape(-1, 1, 1) + elif isinstance(value, np.ndarray) and value.ndim == 1: + value = value.reshape(-1, 1, 1) episode_stats[feature_name][stat_name] = value @@ -1153,15 +1164,15 @@ def _save_episode_images_for_video( # Get all items for this episode episode_dataset = imgs_dataset.select(range(from_idx, to_idx)) + is_depth = img_key in dataset.meta.depth_keys + frame_pattern = DEPTH_FILE_PATTERN if is_depth else IMAGE_FILE_PATTERN + # Define function to save a single image def save_single_image(i_item_tuple): i, item = i_item_tuple - img = item[img_key] - # Use frame-XXXXXX.png format to match encode_video_frames expectations - img.save(str(imgs_dir / f"frame-{i:06d}.png"), quality=100) + write_image(item[img_key], imgs_dir / frame_pattern.format(frame_index=i)) return i - # Save images with proper naming convention for encode_video_frames (frame-XXXXXX.png) items = list(enumerate(episode_dataset)) with ThreadPoolExecutor(max_workers=num_workers) as executor: @@ -1193,13 +1204,14 @@ def _save_batch_episodes_images( hf_dataset = dataset.hf_dataset.with_format(None) imgs_dataset = hf_dataset.select_columns(img_key) + is_depth = img_key in dataset.meta.depth_keys + frame_pattern = DEPTH_FILE_PATTERN if is_depth else IMAGE_FILE_PATTERN + # Define function to save a single image with global frame index # Defined once outside the loop to avoid repeated closure creation def save_single_image(i_item_tuple, base_frame_idx, img_key_param): i, item = i_item_tuple - img = item[img_key_param] - # Use global frame index for naming - img.save(str(imgs_dir / f"frame-{base_frame_idx + i:06d}.png"), quality=100) + write_image(item[img_key_param], imgs_dir / frame_pattern.format(frame_index=base_frame_idx + i)) return i episode_durations = [] @@ -1290,7 +1302,7 @@ def _estimate_frame_size_via_calibration( episode_indices: list[int], temp_dir: Path, fps: int, - camera_encoder: VideoEncoderConfig, + video_encoder: VideoEncoderConfig, num_calibration_frames: int = 30, ) -> float: """Estimate MB per frame by encoding a small calibration sample. @@ -1304,7 +1316,7 @@ def _estimate_frame_size_via_calibration( episode_indices: List of episode indices being processed. temp_dir: Temporary directory for calibration files. fps: Frames per second for video encoding. - camera_encoder: Video encoder settings used for calibration encoding. + video_encoder: Video encoder settings used for calibration encoding. num_calibration_frames: Number of frames to use for calibration (default: 30). Returns: @@ -1329,10 +1341,11 @@ def _estimate_frame_size_via_calibration( hf_dataset = dataset.hf_dataset.with_format(None) sample_indices = range(from_idx, from_idx + num_frames) - # Save calibration frames + # Save calibration frames using the suffix/format the encoder expects. + is_depth = img_key in dataset.meta.depth_keys + frame_pattern = DEPTH_FILE_PATTERN if is_depth else IMAGE_FILE_PATTERN for i, idx in enumerate(sample_indices): - img = hf_dataset[idx][img_key] - img.save(str(calibration_dir / f"frame-{i:06d}.png"), quality=100) + write_image(hf_dataset[idx][img_key], calibration_dir / frame_pattern.format(frame_index=i)) # Encode calibration video calibration_video_path = calibration_dir / "calibration.mp4" @@ -1340,7 +1353,7 @@ def _estimate_frame_size_via_calibration( imgs_dir=calibration_dir, video_path=calibration_video_path, fps=fps, - camera_encoder=camera_encoder, + video_encoder=video_encoder, overwrite=True, ) @@ -1613,6 +1626,7 @@ def recompute_stats( raise ValueError(f"No parquet files found in {data_dir}") all_episode_stats = [] + # TODO: enable image and video stats re-computation numeric_keys = [k for k, v in features_to_compute.items() if v["dtype"] not in ["image", "video"]] for parquet_path in tqdm(parquet_files, desc="Computing stats from data files"): @@ -1658,7 +1672,8 @@ def convert_image_to_video_dataset( dataset: LeRobotDataset, output_dir: Path | None = None, repo_id: str | None = None, - camera_encoder: VideoEncoderConfig | None = None, + rgb_encoder: RGBEncoderConfig | None = None, + depth_encoder: DepthEncoderConfig | None = None, episode_indices: list[int] | None = None, num_workers: int = 4, max_episodes_per_batch: int | None = None, @@ -1670,21 +1685,32 @@ def convert_image_to_video_dataset( LeRobot dataset structure with videos stored in chunked MP4 files. Args: - dataset: The source LeRobot dataset with images - output_dir: Root directory where the edited dataset will be stored. If not specified, defaults to $HF_LEROBOT_HOME/repo_id. Equivalent to new_root in EditDatasetConfig. - repo_id: Edited dataset identifier. Equivalent to new_repo_id in EditDatasetConfig. - camera_encoder: Video encoder settings - (``None`` uses :func:`~lerobot.configs.camera_encoder_defaults`). - episode_indices: List of episode indices to convert (None = all episodes) - num_workers: Number of threads for parallel processing (default: 4) - max_episodes_per_batch: Maximum episodes per video batch to avoid memory issues (None = no limit) - max_frames_per_batch: Maximum frames per video batch to avoid memory issues (None = no limit) + dataset: The source LeRobot dataset with images. + output_dir: Root directory where the converted dataset will be stored. When + ``None``, defaults to ``$HF_LEROBOT_HOME/repo_id``. Equivalent to + ``new_root`` in ``EditDatasetConfig``. + repo_id: Converted dataset identifier. Equivalent to ``new_repo_id`` in + ``EditDatasetConfig``. + rgb_encoder: Video encoder settings applied to RGB cameras. When ``None``, + :func:`~lerobot.configs.video.rgb_encoder_defaults` is used. + depth_encoder: Video encoder settings applied to depth-map cameras, including + the quantization parameters persisted to the dataset metadata. When + ``None``, :func:`~lerobot.configs.video.depth_encoder_defaults` is used. + episode_indices: Episode indices to convert. When ``None``, all episodes are + converted. + num_workers: Number of threads for parallel processing. + max_episodes_per_batch: Maximum episodes per video batch, to bound memory use. + ``None`` means no limit. + max_frames_per_batch: Maximum frames per video batch, to bound memory use. + ``None`` means no limit. Returns: - New LeRobotDataset with images encoded as videos + A new :class:`LeRobotDataset` with images encoded as videos. """ - if camera_encoder is None: - camera_encoder = camera_encoder_defaults() + if rgb_encoder is None: + rgb_encoder = rgb_encoder_defaults() + if depth_encoder is None: + depth_encoder = depth_encoder_defaults() # Check that it's an image dataset if len(dataset.meta.video_keys) > 0: @@ -1709,10 +1735,7 @@ def convert_image_to_video_dataset( logging.info( f"Converting {len(episode_indices)} episodes with {len(img_keys)} cameras from {dataset.repo_id}" ) - logging.info( - f"Video codec: {camera_encoder.vcodec}, pixel format: {camera_encoder.pix_fmt}, " - f"GOP: {camera_encoder.g}, CRF: {camera_encoder.crf}" - ) + logging.info(f"RGB video encoder: {rgb_encoder}, depth video encoder: {depth_encoder}") # Create new features dict, converting image features to video features new_features = {} @@ -1774,6 +1797,8 @@ def convert_image_to_video_dataset( episode_lengths = {ep_idx: dataset.meta.episodes["length"][ep_idx] for ep_idx in episode_indices} for img_key in tqdm(img_keys, desc="Processing cameras"): + target_encoder = depth_encoder if img_key in dataset.meta.depth_keys else rgb_encoder + # Estimate size per frame by encoding a small calibration sample # This provides accurate compression ratio for the specific codec parameters size_per_frame_mb = _estimate_frame_size_via_calibration( @@ -1782,7 +1807,7 @@ def convert_image_to_video_dataset( episode_indices=episode_indices, temp_dir=temp_dir, fps=fps, - camera_encoder=camera_encoder, + video_encoder=target_encoder, ) logging.info(f"Processing camera: {img_key}") @@ -1824,7 +1849,7 @@ def convert_image_to_video_dataset( imgs_dir=imgs_dir, video_path=video_path, fps=fps, - camera_encoder=camera_encoder, + video_encoder=target_encoder, overwrite=True, ) @@ -1863,16 +1888,11 @@ def convert_image_to_video_dataset( new_meta.info.total_tasks = dataset.meta.total_tasks new_meta.info.splits = {"train": f"0:{len(episode_indices)}"} - # Update video info for all image keys (now videos) - # We need to manually set video info since update_video_info() checks video_keys first + # Update video info for all image keys (now videos). They are registered as + # video features above, so update_video_info populates their (still-empty) info. for img_key in img_keys: - if not new_meta.features[img_key].get("info", None): - video_path = new_meta.root / new_meta.video_path.format( - video_key=img_key, chunk_index=0, file_index=0 - ) - new_meta.info.features[img_key]["info"] = get_video_info( - video_path, camera_encoder=camera_encoder - ) + target_encoder = depth_encoder if img_key in dataset.meta.depth_keys else rgb_encoder + new_meta.update_video_info(video_key=img_key, video_encoder=target_encoder) write_info(new_meta.info, new_meta.root) @@ -1899,11 +1919,11 @@ def convert_image_to_video_dataset( def _reencode_video_worker(args: tuple) -> Path: """Picklable worker for :func:`reencode_dataset`'s process pool.""" - video_path, camera_encoder, encoder_threads = args + video_path, video_encoder, encoder_threads = args reencode_video( input_video_path=video_path, output_video_path=video_path, - camera_encoder=camera_encoder, + video_encoder=video_encoder, encoder_threads=encoder_threads, overwrite=True, ) @@ -1912,7 +1932,8 @@ def _reencode_video_worker(args: tuple) -> Path: def reencode_dataset( dataset: LeRobotDataset, - camera_encoder: VideoEncoderConfig, + rgb_encoder: RGBEncoderConfig | None = None, + depth_encoder: DepthEncoderConfig | None = None, encoder_threads: int | None = None, num_workers: int | None = None, ) -> LeRobotDataset: @@ -1923,8 +1944,11 @@ def reencode_dataset( Args: dataset: An existing :class:`LeRobotDataset` whose videos will be re-encoded. - camera_encoder: Target encoder configuration applied to every video - file. + rgb_encoder: Target encoder configuration applied to every RGB video + file. If ``None``, re-encoding is skipped for RGB videos. + depth_encoder: Target encoder configuration applied to every depth video + file. If ``None``, re-encoding is skipped for depth videos. + Quantization parameters will not override the ones in the current dataset. encoder_threads: Per-encoder thread count forwarded to :func:`reencode_video`. ``None`` lets the codec decide. num_workers: Number of parallel processes. ``None`` or ``0`` means @@ -1936,23 +1960,35 @@ def reencode_dataset( on disk. """ meta = dataset.meta - video_paths_list = [] + video_keys_encoders_dict = {} + video_keys_paths_dict = {} + + if rgb_encoder is None and depth_encoder is None: + raise ValueError("Either rgb_encoder or depth_encoder must be provided") # Only re-encode if the videos are not already encoded with the given video encoding parameters for video_key in meta.video_keys: current_info = meta.info.features[video_key].get("info", {}) - current_encoder = VideoEncoderConfig.from_video_info(current_info) - if current_encoder != camera_encoder: - video_paths_list.extend((meta.root / VIDEO_DIR / video_key).rglob("*.mp4")) + current_encoder = encoder_config_from_video_info(current_info) + target_encoder = depth_encoder if video_key in meta.depth_keys else rgb_encoder + if target_encoder is None: + logging.info(f"No encoder provided for {video_key} video. Skipping re-encoding.") + elif current_encoder != target_encoder: + video_keys_paths_dict[video_key] = list((meta.root / VIDEO_DIR / video_key).rglob("*.mp4")) + video_keys_encoders_dict[video_key] = target_encoder else: - logging.info(f"{video_key} videos are already encoded with {camera_encoder}. Nothing to do.") + logging.info(f"{video_key} videos are already encoded with {target_encoder}. Nothing to do.") - if len(video_paths_list) == 0: + if len(video_keys_paths_dict) == 0: logging.warning("Dataset has no videos to re-encode.") return dataset - logging.info(f"Re-encoding {len(video_paths_list)} video file(s) with {camera_encoder}") + logging.info(f"Re-encoding {sum(len(paths) for paths in video_keys_paths_dict.values())} video file(s).") - worker_args = [(vp, camera_encoder, encoder_threads) for vp in video_paths_list] + worker_args = [ + (path, encoder, encoder_threads) + for video_key, encoder in video_keys_encoders_dict.items() + for path in video_keys_paths_dict[video_key] + ] if num_workers and num_workers > 1: with ProcessPoolExecutor(max_workers=num_workers) as pool: futures = [pool.submit(_reencode_video_worker, args) for args in worker_args] @@ -1966,10 +2002,15 @@ def reencode_dataset( for args in tqdm(worker_args, desc="Re-encoding videos"): _reencode_video_worker(args) - # Refresh video info in metadata for every video key. - for vid_key in meta.video_keys: - video_path = meta.root / meta.get_video_file_path(0, vid_key) - meta.info.features[vid_key]["info"] = get_video_info(video_path, camera_encoder=camera_encoder) + # Refresh video info in metadata for every re-encoded key. Re-encoding only + # changes codec/container params, so for depth videos we preserve ``is_depth_map`` + # and the depth quantization params (``video.depth_min`` / ``video.depth_max`` / + # ...), which describe the data rather than the codec and must survive a transcode. + # RGB videos pass an empty set: still a refresh, but nothing to preserve. + depth_preserve_keys = {"is_depth_map", *(f"video.{n}" for n in DEPTH_ENCODER_INFO_FIELD_NAMES)} + for video_key, encoder in video_keys_encoders_dict.items(): + preserve_keys = depth_preserve_keys if video_key in meta.depth_keys else set() + meta.update_video_info(video_key=video_key, video_encoder=encoder, preserve_keys=preserve_keys) write_info(meta.info, meta.root) logging.info("Dataset metadata updated.") diff --git a/src/lerobot/datasets/dataset_writer.py b/src/lerobot/datasets/dataset_writer.py index 633c00c1a..1aee1497c 100644 --- a/src/lerobot/datasets/dataset_writer.py +++ b/src/lerobot/datasets/dataset_writer.py @@ -31,7 +31,13 @@ import PIL.Image import pyarrow.parquet as pq import torch -from lerobot.configs import VideoEncoderConfig, camera_encoder_defaults +from lerobot.configs import ( + DepthEncoderConfig, + RGBEncoderConfig, + VideoEncoderConfig, + depth_encoder_defaults, + rgb_encoder_defaults, +) from .compute_stats import compute_episode_stats from .dataset_metadata import LeRobotDatasetMetadata @@ -48,6 +54,7 @@ from .io_utils import ( write_info, ) from .utils import ( + DEFAULT_DEPTH_PATH, DEFAULT_EPISODES_PATH, DEFAULT_IMAGE_PATH, update_chunk_file_indices, @@ -67,17 +74,22 @@ def _encode_video_worker( episode_index: int, root: Path, fps: int, - camera_encoder: VideoEncoderConfig | None = None, + video_encoder: VideoEncoderConfig | None = None, encoder_threads: int | None = None, ) -> Path: temp_path = Path(tempfile.mkdtemp(dir=root)) / f"{video_key}_{episode_index:03d}.mp4" - fpath = DEFAULT_IMAGE_PATH.format(image_key=video_key, episode_index=episode_index, frame_index=0) + path_template = ( + DEFAULT_DEPTH_PATH + if video_encoder is not None and isinstance(video_encoder, DepthEncoderConfig) + else DEFAULT_IMAGE_PATH + ) + fpath = path_template.format(image_key=video_key, episode_index=episode_index, frame_index=0) img_dir = (root / fpath).parent encode_video_frames( img_dir, temp_path, fps, - camera_encoder=camera_encoder, + video_encoder=video_encoder, encoder_threads=encoder_threads, overwrite=True, ) @@ -96,7 +108,8 @@ class DatasetWriter: self, meta: LeRobotDatasetMetadata, root: Path, - camera_encoder: VideoEncoderConfig | None, + rgb_encoder: RGBEncoderConfig | None, + depth_encoder: DepthEncoderConfig | None, encoder_threads: int | None, batch_encoding_size: int, streaming_encoder: StreamingVideoEncoder | None = None, @@ -108,8 +121,11 @@ class DatasetWriter: meta: Dataset metadata instance (used for feature schema, chunk settings, and episode persistence). root: Local dataset root directory. - camera_encoder: Video encoder settings applied to all cameras. - ``None`` uses :func:`~lerobot.configs.camera_encoder_defaults`. + rgb_encoder: Video encoder settings applied to RGB cameras. When + ``None``, :func:`~lerobot.configs.video.rgb_encoder_defaults` is used. + depth_encoder: Video encoder settings applied to depth cameras, including + the quantization parameters. When ``None``, + :func:`~lerobot.configs.video.depth_encoder_defaults` is used. encoder_threads: Number of encoder threads (global). ``None`` lets the codec decide. batch_encoding_size: Number of episodes to accumulate before @@ -120,7 +136,8 @@ class DatasetWriter: """ self._meta = meta self._root = root - self._camera_encoder = camera_encoder or camera_encoder_defaults() + self._rgb_encoder = rgb_encoder or rgb_encoder_defaults() + self._depth_encoder = depth_encoder or depth_encoder_defaults() self._encoder_threads = encoder_threads self._batch_encoding_size = batch_encoding_size self._streaming_encoder = streaming_encoder @@ -145,7 +162,8 @@ class DatasetWriter: return ep_buffer def _get_image_file_path(self, episode_index: int, image_key: str, frame_index: int) -> Path: - fpath = DEFAULT_IMAGE_PATH.format( + path_template = DEFAULT_DEPTH_PATH if image_key in self._meta.depth_keys else DEFAULT_IMAGE_PATH + fpath = path_template.format( image_key=image_key, episode_index=episode_index, frame_index=frame_index ) return self._root / fpath @@ -195,6 +213,7 @@ class DatasetWriter: if frame_index == 0 and self._streaming_encoder is not None: self._streaming_encoder.start_episode( video_keys=list(self._meta.video_keys), + depth_video_keys=list(self._meta.depth_keys), temp_dir=self._root, ) @@ -282,10 +301,13 @@ class DatasetWriter: if use_streaming: streaming_results = self._streaming_encoder.finish_episode() for video_key in self._meta.video_keys: + normalization_factor = 255.0 if video_key not in self._meta.depth_keys else 1.0 temp_path, video_stats = streaming_results[video_key] if video_stats is not None: ep_stats[video_key] = { - k: v if k == "count" else np.squeeze(v.reshape(1, -1, 1, 1) / 255.0, axis=0) + k: v + if k == "count" + else np.squeeze(v.reshape(1, -1, 1, 1) / normalization_factor, axis=0) for k, v in video_stats.items() } ep_metadata.update(self._save_episode_video(video_key, episode_index, temp_path=temp_path)) @@ -300,7 +322,7 @@ class DatasetWriter: episode_index, self._root, self._meta.fps, - self._camera_encoder, + self._depth_encoder if video_key in self._meta.depth_keys else self._rgb_encoder, self._encoder_threads, ): video_key for video_key in self._meta.video_keys @@ -511,7 +533,12 @@ class DatasetWriter: # Update video info (only needed when first episode is encoded) if episode_index == 0: - self._meta.update_video_info(video_key, camera_encoder=self._camera_encoder) + self._meta.update_video_info( + video_key, + video_encoder=self._depth_encoder + if video_key in self._meta.depth_keys + else self._rgb_encoder, + ) write_info(self._meta.info, self._meta.root) metadata = { @@ -578,13 +605,14 @@ class DatasetWriter: self.image_writer.wait_until_done() def _encode_temporary_episode_video(self, video_key: str, episode_index: int) -> Path: - """Use ffmpeg to convert frames stored as png into mp4 videos.""" + """Use ffmpeg to convert frames stored as png/tiff into mp4 videos.""" + is_depth = video_key in self._meta.depth_keys return _encode_video_worker( video_key, episode_index, self._root, self._meta.fps, - self._camera_encoder, + self._depth_encoder if is_depth else self._rgb_encoder, self._encoder_threads, ) diff --git a/src/lerobot/datasets/depth_utils.py b/src/lerobot/datasets/depth_utils.py new file mode 100644 index 000000000..801c86a09 --- /dev/null +++ b/src/lerobot/datasets/depth_utils.py @@ -0,0 +1,268 @@ +#!/usr/bin/env python + +# Copyright 2026 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. +""" +Depth encoding/decoding helpers for :class:`DepthEncoderConfig`. +""" + +import math +from typing import Literal + +import av +import numpy as np +import torch +from numpy.typing import NDArray + +from lerobot.configs.video import ( + DEFAULT_DEPTH_MAX, + DEFAULT_DEPTH_MIN, + DEFAULT_DEPTH_PIX_FMT, + DEFAULT_DEPTH_SHIFT, + DEFAULT_DEPTH_USE_LOG, + DEPTH_METER_UNIT, + DEPTH_MILLIMETER_UNIT, + DEPTH_QMAX, +) + +from .image_writer import squeeze_single_channel +from .pyav_utils import write_u16_plane + +_MM_PER_METRE = 1000.0 +_UINT16_MAX = 65535 + + +def _validate_log_quant_params(depth_min: float, shift: float) -> None: + """Ensure ``log(depth_min + shift)`` is finite.""" + if depth_min + shift <= 0: + raise ValueError( + f"depth_min + shift must be positive for logarithmic quantization, " + f"got depth_min={depth_min} + shift={shift} = {depth_min + shift}" + ) + + +def _depth_input_to_float32_and_unit( + depth: NDArray[np.integer] | NDArray[np.floating], + input_unit: Literal["auto", DEPTH_METER_UNIT, DEPTH_MILLIMETER_UNIT], +) -> tuple[NDArray[np.float32], Literal[DEPTH_METER_UNIT, DEPTH_MILLIMETER_UNIT]]: + """Convert depth to float32 in the chosen unit, and return the resolved unit.""" + resolved_unit = ( + (DEPTH_METER_UNIT if np.issubdtype(depth.dtype, np.floating) else DEPTH_MILLIMETER_UNIT) + if input_unit == "auto" + else input_unit + ) + return depth.astype(np.float32, order="K"), resolved_unit + + +def quantize_depth( + depth: NDArray[np.uint16] | NDArray[np.float32] | torch.Tensor, + depth_min: float = DEFAULT_DEPTH_MIN, + depth_max: float = DEFAULT_DEPTH_MAX, + shift: float = DEFAULT_DEPTH_SHIFT, + use_log: bool = DEFAULT_DEPTH_USE_LOG, + pix_fmt: str = DEFAULT_DEPTH_PIX_FMT, + video_backend: str | None = "pyav", + input_unit: Literal["auto", DEPTH_METER_UNIT, DEPTH_MILLIMETER_UNIT] = "auto", +) -> NDArray[np.uint16] | av.VideoFrame: + """Quantize depth to 12-bit codes (``uint16``, values ``0…DEPTH_QMAX``). + + Depth maps are packed into 12-bit integer frames so they fit in standard + high-bit-depth pixel formats (e.g. ``yuv420p12le`` / ``gray12le``) + and can be encoded by widely supported video codecs (e.g. HEVC Main 12). + Logarithmic quantization is the default because it allocates more quanta + to near-range depth, which matches the (1/depth) error profile of typical + depth sensors. Math is ported from BEHAVIOR-1K's ``obs_utils.py``. + + **Input units**: + + - ``input_unit="auto"`` (default): infer from dtype (floating = m, non-floating = mm). + - ``input_unit="mm"``: interpret input values as millimetres. + - ``input_unit="m"``: interpret input values as metres. + + Quantization math runs in the **resolved input unit**. + + ``depth_min``, ``depth_max``, and ``shift`` are always in **metres**. + + Args: + depth: Depth map; ``torch.Tensor`` is moved to CPU for conversion. + depth_min: Depth (metres) at quantum ``0``. + depth_max: Depth (metres) at quantum :data:`DEPTH_QMAX`. + shift: Depth shift (metres); used in log mode. Must satisfy ``depth_min + shift > 0``. + use_log: If ``True`` (default), quantize in log space. + video_backend: Video backend to use for encoding. Defaults to "pyav". + input_unit: Input unit policy (``"auto"``, ``"mm"``, ``"m"``). + + Returns: + ``numpy.ndarray``, ``dtype=uint16``, same shape as ``depth``, values in + ``[0, DEPTH_QMAX]``. + + Raises: + ValueError: If ``input_unit`` is not ``"auto"``, ``"mm"``, or ``"m"``. + ValueError: If ``use_log=True`` and ``depth_min + shift <= 0``. + """ + if input_unit not in ("auto", DEPTH_METER_UNIT, DEPTH_MILLIMETER_UNIT): + raise ValueError( + f"input_unit must be 'auto', '{DEPTH_METER_UNIT}', or '{DEPTH_MILLIMETER_UNIT}', got {input_unit!r}" + ) + + if isinstance(depth, torch.Tensor): + depth = depth.detach().cpu().numpy() + + # Squeeze single-channel dim: (H, W, 1) or (1, H, W) → (H, W) + depth = squeeze_single_channel(depth) + + depth_f, resolved_unit = _depth_input_to_float32_and_unit(depth, input_unit=input_unit) + + # Convert depth_min, depth_max, and shift to the resolved input unit. + depth_min_u = ( + np.float32(depth_min) if resolved_unit == DEPTH_METER_UNIT else np.float32(depth_min * _MM_PER_METRE) + ) + depth_max_u = ( + np.float32(depth_max) if resolved_unit == DEPTH_METER_UNIT else np.float32(depth_max * _MM_PER_METRE) + ) + shift_u = np.float32(shift) if resolved_unit == DEPTH_METER_UNIT else np.float32(shift * _MM_PER_METRE) + + # Normalization and quantization is performed in the resolved input unit. + if use_log: + _validate_log_quant_params(depth_min, shift) + log_min = math.log(float(depth_min_u + shift_u)) + log_max = math.log(float(depth_max_u + shift_u)) + norm = (np.log(depth_f + shift_u) - log_min) / (log_max - log_min) + else: + norm = (depth_f - depth_min_u) / (depth_max_u - depth_min_u) + + quantized = np.rint(norm * DEPTH_QMAX).clip(0, DEPTH_QMAX).astype(np.uint16, copy=False) + + if video_backend == "pyav": + frame = av.VideoFrame.from_ndarray(quantized, format=pix_fmt) + write_u16_plane(frame.planes[0], quantized) + return frame + else: + return quantized + + +def dequantize_depth( + quantized: NDArray[np.uint16] | av.VideoFrame | torch.Tensor, + depth_min: float = DEFAULT_DEPTH_MIN, + depth_max: float = DEFAULT_DEPTH_MAX, + shift: float = DEFAULT_DEPTH_SHIFT, + use_log: bool = DEFAULT_DEPTH_USE_LOG, + pix_fmt: str = DEFAULT_DEPTH_PIX_FMT, + output_unit: Literal[DEPTH_METER_UNIT, DEPTH_MILLIMETER_UNIT] = DEPTH_MILLIMETER_UNIT, + output_tensor: bool = True, + output_channel_last: bool = False, +) -> NDArray[np.uint16] | NDArray[np.float32] | torch.Tensor: + """Inverse of :func:`quantize_depth`. + + Decoding inverts the same normalized code mapping as :func:`quantize_depth` + using ``depth_min`` / ``depth_max`` / ``shift`` (in metres), then returns + the requested output unit. Tuning arguments **must match** :func:`quantize_depth`. + + Accepted input layouts : + + - ``(H, W, 1)`` or ``(H, W)`` — single frame with channel-last. + - ``(..., 1, H, W)`` — batched frames with channel-first. + - ``(..., H, W, 1)`` — batched frames with channel-last. + Output layout is determined by ``output_channel_last``. + + Args: + quantized: 12-bit codes in ``[0, DEPTH_QMAX]``. ``np.ndarray``, + ``av.VideoFrame``, or ``torch.Tensor`` (any integer or float dtype). + depth_min, depth_max, shift, use_log: Same as :func:`quantize_depth` (metres). + pix_fmt: Pixel format used to extract the plane from an ``av.VideoFrame``. + output_unit: ``"mm"`` returns ``uint16`` millimetres (rint, clip + ``[0, 65535]``) when returning a numpy array, or ``float32`` mm when + ``output_tensor=True``. ``"m"`` returns ``float32`` metres in + ``[depth_min, depth_max]``. + output_tensor: If True, return a ``torch.Tensor`` instead of a numpy array. + + Returns: + Depth map in the requested unit and dtype. + + Raises: + ValueError: If ``output_unit`` is not ``"m"`` or ``"mm"``. + ValueError: If ``use_log=True`` and ``depth_min + shift <= 0``. + """ + if output_unit not in (DEPTH_METER_UNIT, DEPTH_MILLIMETER_UNIT): + raise ValueError( + f"output_unit must be '{DEPTH_METER_UNIT}' or '{DEPTH_MILLIMETER_UNIT}', got {output_unit!r}" + ) + if use_log: + _validate_log_quant_params(depth_min, shift) + + if isinstance(quantized, av.VideoFrame): + quantized = quantized.to_ndarray(format=pix_fmt) + + # Compute the scale and offset first. + depth_min_m = float(depth_min) + depth_max_m = float(depth_max) + shift_m = float(shift) + if use_log: + log_min = math.log(depth_min_m + shift_m) + log_max = math.log(depth_max_m + shift_m) + scale = (log_max - log_min) / DEPTH_QMAX + offset = log_min + else: + scale = (depth_max_m - depth_min_m) / DEPTH_QMAX + offset = depth_min_m + + # ── Torch path: stay on the input device, single fp32 allocation. ──────── + if isinstance(quantized, torch.Tensor): + if quantized.ndim >= 3: + # Drop the single-channel dimension so the math runs on (..., H, W). + quantized = quantized.squeeze(-3) if quantized.shape[-3] == 1 else quantized.squeeze(-1) + + # Single allocation we own; everything else is in-place. + buf = quantized.to(dtype=torch.float32, copy=True) + buf.mul_(scale).add_(offset) + if use_log: + buf.exp_().sub_(shift_m) + buf.clamp_(depth_min_m, depth_max_m) + buf.unsqueeze_(-1) if output_channel_last else buf.unsqueeze_(-3) + + if output_unit == DEPTH_METER_UNIT: + return buf if output_tensor else buf.cpu().numpy() + + # mm path: round + clamp in float32, skipping the uint16 round-trip + # when returning a tensor (torch.uint16 is poorly supported). + buf.mul_(_MM_PER_METRE).round_().clamp_(0.0, _UINT16_MAX) + if output_tensor: + return buf + return buf.cpu().numpy().astype(np.uint16, copy=False) + + # ── NumPy path: single fp32 allocation, ``out=`` for in-place math. ───── + arr = np.asarray(quantized) + if arr.ndim >= 3: + # Drop the single-channel dimension so the math runs on (..., H, W). + arr = np.squeeze(arr, axis=-3) if arr.shape[-3] == 1 else np.squeeze(arr, axis=-1) + + buf = np.empty(arr.shape, dtype=np.float32) + np.multiply(arr, scale, out=buf) + np.add(buf, offset, out=buf) + if use_log: + np.exp(buf, out=buf) + np.subtract(buf, shift_m, out=buf) + np.clip(buf, depth_min_m, depth_max_m, out=buf) + buf = np.expand_dims(buf, axis=-1) if output_channel_last else np.expand_dims(buf, axis=-3) + + if output_unit == DEPTH_METER_UNIT: + return torch.from_numpy(buf) if output_tensor else buf + + np.multiply(buf, _MM_PER_METRE, out=buf) + np.rint(buf, out=buf) + np.clip(buf, 0.0, _UINT16_MAX, out=buf) + if output_tensor: + # torch.uint16 support is very limited; return float32 millimetres. + return torch.from_numpy(buf) + return buf.astype(np.uint16, copy=False) diff --git a/src/lerobot/datasets/factory.py b/src/lerobot/datasets/factory.py index cd29ee99e..da7b4365a 100644 --- a/src/lerobot/datasets/factory.py +++ b/src/lerobot/datasets/factory.py @@ -97,6 +97,7 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas revision=cfg.dataset.revision, video_backend=cfg.dataset.video_backend, return_uint8=True, + depth_output_unit=cfg.dataset.depth_output_unit, tolerance_s=cfg.tolerance_s, ) else: @@ -127,6 +128,8 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas if cfg.dataset.use_imagenet_stats: for key in dataset.meta.camera_keys: + if key in dataset.meta.depth_keys: + continue # Exclude depth keys from ImageNet stats for stats_type, stats in IMAGENET_STATS.items(): dataset.meta.stats[key][stats_type] = torch.tensor(stats, dtype=torch.float32) diff --git a/src/lerobot/datasets/feature_utils.py b/src/lerobot/datasets/feature_utils.py index 56264408f..343b2fdcc 100644 --- a/src/lerobot/datasets/feature_utils.py +++ b/src/lerobot/datasets/feature_utils.py @@ -336,7 +336,7 @@ def validate_feature_image_or_video( Args: name (str): The name of the feature. - expected_shape (list[str]): The expected shape (C, H, W). + expected_shape (list[str]): The expected shape, e.g. (C, H, W) or (H, W, C). value: The image data to validate. Returns: diff --git a/src/lerobot/datasets/image_writer.py b/src/lerobot/datasets/image_writer.py index 8fb5804a5..41790b46a 100644 --- a/src/lerobot/datasets/image_writer.py +++ b/src/lerobot/datasets/image_writer.py @@ -41,11 +41,51 @@ def safe_stop_image_writer(func): return wrapper -def image_array_to_pil_image(image_array: np.ndarray, range_check: bool = True) -> PIL.Image.Image: - # TODO(aliberts): handle 1 channel and 4 for depth images - if image_array.ndim != 3: - raise ValueError(f"The array has {image_array.ndim} dimensions, but 3 is expected for an image.") +def squeeze_single_channel(array: np.ndarray) -> np.ndarray: + """Drop a leading or trailing singleton channel dim: ``(1, H, W)`` / ``(H, W, 1)`` -> ``(H, W)``. + Unlike ``array.squeeze()``, this only removes the channel axis, never an ``H`` or ``W`` of size 1. + """ + if array.ndim == 3: + if array.shape[0] == 1: + return array[0] + if array.shape[-1] == 1: + return array[..., 0] + return array + + +def image_array_to_pil_image(image_array: np.ndarray, range_check: bool = True) -> PIL.Image.Image: + """Convert a NumPy array to a PIL Image, preserving precision for grayscale. + + Behaviour by shape: + + - ``(H, W)`` or ``(1, H, W)`` / ``(H, W, 1)``: single-channel grayscale. + The native dtype is preserved using the matching PIL mode + (``I;16`` / ``F``). This is the path used for raw depth maps (no rescaling, clamping, or downcasting) + - ``(3, H, W)`` / ``(H, W, 3)``: RGB. Channels-first inputs are transposed + to channels-last. Float inputs in ``[0, 1]`` are scaled to ``uint8`` + (existing behaviour, gated by ``range_check``). + + Other shapes / channel counts raise ``NotImplementedError`` or + ``ValueError``. + """ + # TODO(CarolinePascal): 4 dimensions RGB-D images + if image_array.ndim not in (2, 3): + raise ValueError(f"The array has {image_array.ndim} dimensions, but 2 or 3 is expected for an image.") + + # Squeeze 3D single-channel inputs to 2D so depth maps work whether the + # caller emits (H, W), (1, H, W), or (H, W, 1). + image_array = squeeze_single_channel(image_array) + + if image_array.ndim == 2: + if image_array.dtype not in [np.uint16, np.float32]: + raise ValueError( + f"Unsupported single-channel image dtype: {image_array.dtype}. " + f"Supported dtypes: {sorted(str(d) for d in [np.uint16, np.float32])}." + ) + return PIL.Image.fromarray(np.ascontiguousarray(image_array)) + + # 3D path: must be RGB (3 channels), channels-first or channels-last. if image_array.shape[0] == 3: # Transpose from pytorch convention (C, H, W) to (H, W, C) image_array = image_array.transpose(1, 2, 0) @@ -71,13 +111,29 @@ def image_array_to_pil_image(image_array: np.ndarray, range_check: bool = True) return PIL.Image.fromarray(image_array) +def save_kwargs_for_path(fpath: Path, compress_level: int) -> dict: + """Pick the right format-specific kwargs for :meth:`PIL.Image.Image.save`. + + PNG uses ``compress_level`` (0-9, zlib). TIFF uses ``compression`` (raw) for lossless raw depth maps. + """ + suffix = Path(fpath).suffix.lower() + if suffix == ".png": + return {"compress_level": compress_level} + if suffix in (".tif", ".tiff"): + return {"compression": "raw"} + else: + raise ValueError(f"Unsupported image file extension: {suffix}") + + def write_image(image: np.ndarray | PIL.Image.Image, fpath: Path, compress_level: int = 1): """ Saves a NumPy array or PIL Image to a file. This function handles both NumPy arrays and PIL Image objects, converting the former to a PIL Image before saving. It includes error handling for - the save operation. + the save operation. The output format is inferred from the *fpath* + extension: ``.png`` → PNG with ``compress_level``, ``.tiff`` / ``.tif`` + → lossless raw depth maps (TIFF). Args: image (np.ndarray | PIL.Image.Image): The image data to save. @@ -101,7 +157,7 @@ def write_image(image: np.ndarray | PIL.Image.Image, fpath: Path, compress_level img = image else: raise TypeError(f"Unsupported image type: {type(image)}") - img.save(fpath, compress_level=compress_level) + img.save(fpath, **save_kwargs_for_path(fpath, compress_level)) except Exception as e: logger.error("Error writing image %s: %s", fpath, e) diff --git a/src/lerobot/datasets/io_utils.py b/src/lerobot/datasets/io_utils.py index be94f3b3a..868a114f5 100644 --- a/src/lerobot/datasets/io_utils.py +++ b/src/lerobot/datasets/io_utils.py @@ -226,28 +226,50 @@ def load_image_as_numpy( Args: fpath (str | Path): Path to the image file. dtype (np.dtype): The desired data type of the output array. If floating, - pixels are scaled to [0, 1]. + pixels are scaled to [0, 1]. Only used for RGB images. channel_first (bool): If True, converts the image to (C, H, W) format. Otherwise, it remains in (H, W, C) format. Returns: np.ndarray: The image as a numpy array. """ - img = PILImage.open(fpath).convert("RGB") - img_array = np.array(img, dtype=dtype) + is_depth = fpath.endswith(".tiff") or fpath.endswith(".tif") + if is_depth: + # Preserve the native depth dtype (uint16 -> "I;16", float32 -> "F"). + img = PILImage.open(fpath) + img_array = np.array(img) + else: + img = PILImage.open(fpath).convert("RGB") + img_array = np.array(img, dtype=dtype) + if np.issubdtype(dtype, np.floating): + img_array /= 255.0 if channel_first: # (H, W, C) -> (C, H, W) - img_array = np.transpose(img_array, (2, 0, 1)) - if np.issubdtype(dtype, np.floating): - img_array /= 255.0 + img_array = img_array[np.newaxis, ...] if img_array.ndim == 2 else np.transpose(img_array, (2, 0, 1)) return img_array +# PIL modes for 16-bit unsigned depth maps. +UINT16_PIL_MODES = {"I;16", "I;16B", "I;16L"} + + +def pil_to_chw_tensor(img: PILImage.Image) -> torch.Tensor: + """Convert a PIL image to a channel-first tensor. + + ``uint16`` depth maps become ``float32 (1, H, W)`` in native units (``ToTensor`` + would overflow them to ``int16``); all other modes use the standard ``ToTensor`` path. + """ + if img.mode in UINT16_PIL_MODES: + return torch.from_numpy(np.array(img, dtype=np.float32))[None, ...] + return transforms.ToTensor()(img) + + def hf_transform_to_torch(items_dict: dict[str, list[Any]]) -> dict[str, list[torch.Tensor | str]]: """Convert a batch from a Hugging Face dataset to torch tensors. This transform function converts items from Hugging Face dataset format (pyarrow) - to torch tensors. Importantly, images are converted from PIL objects (H, W, C, uint8) - to a torch image representation (C, H, W, float32) in the range [0, 1]. Other + to torch tensors. RGB images are converted from PIL objects (H, W, C, uint8) + to a torch image representation (C, H, W, float32) in the range [0, 1]. Depth + maps are returned as float32 (1, H, W) in their native units. Other types are converted to torch.tensor. Args: @@ -262,8 +284,7 @@ def hf_transform_to_torch(items_dict: dict[str, list[Any]]) -> dict[str, list[to continue first_item = items_dict[key][0] if isinstance(first_item, PILImage.Image): - to_tensor = transforms.ToTensor() - items_dict[key] = [to_tensor(img) for img in items_dict[key]] + items_dict[key] = [pil_to_chw_tensor(img) for img in items_dict[key]] elif first_item is None or isinstance(first_item, dict): pass else: @@ -329,7 +350,11 @@ def item_to_torch(item: dict) -> dict: """ skip_keys = {"task", *LANGUAGE_COLUMNS} for key, val in item.items(): - if isinstance(val, (np.ndarray | list)) and key not in skip_keys: + if key in skip_keys: + continue + if isinstance(val, PILImage.Image): + item[key] = pil_to_chw_tensor(val) + elif isinstance(val, (np.ndarray | list)): # Convert numpy arrays and lists to torch tensors item[key] = torch.tensor(val) return item diff --git a/src/lerobot/datasets/lerobot_dataset.py b/src/lerobot/datasets/lerobot_dataset.py index 49e77b53a..f600f1804 100644 --- a/src/lerobot/datasets/lerobot_dataset.py +++ b/src/lerobot/datasets/lerobot_dataset.py @@ -24,7 +24,7 @@ import torch.utils from huggingface_hub import HfApi, snapshot_download from huggingface_hub.errors import RevisionNotFoundError -from lerobot.configs import VideoEncoderConfig +from lerobot.configs import DEFAULT_DEPTH_UNIT, DepthEncoderConfig, RGBEncoderConfig from lerobot.utils.constants import HF_LEROBOT_HUB_CACHE from .dataset_metadata import CODEBASE_VERSION, LeRobotDatasetMetadata @@ -58,8 +58,10 @@ class LeRobotDataset(torch.utils.data.Dataset): download_videos: bool = True, video_backend: str | None = None, return_uint8: bool = False, + depth_output_unit: str = DEFAULT_DEPTH_UNIT, batch_encoding_size: int = 1, - camera_encoder: VideoEncoderConfig | None = None, + rgb_encoder: RGBEncoderConfig | None = None, + depth_encoder: DepthEncoderConfig | None = None, encoder_threads: int | None = None, streaming_encoding: bool = False, encoder_queue_maxsize: int = 30, @@ -183,8 +185,11 @@ class LeRobotDataset(torch.utils.data.Dataset): You can also use the 'pyav' decoder used by Torchvision, which used to be the default option, or 'video_reader' which is another decoder of Torchvision. batch_encoding_size (int, optional): Number of episodes to accumulate before batch encoding videos. Set to 1 for immediate encoding (default), or higher for batched encoding. Defaults to 1. - camera_encoder (VideoEncoderConfig | None, optional): Video encoder settings for cameras - (codec, quality, etc.). When ``None``, :func:`~lerobot.configs.video.camera_encoder_defaults` + rgb_encoder (RGBEncoderConfig | None, optional): Video encoder settings for cameras + (codec, quality, etc.). When ``None``, :func:`~lerobot.configs.video.rgb_encoder_defaults` + is used by the writer. + depth_encoder (DepthEncoderConfig | None, optional): Video encoder settings for depth cameras + (codec, quality, etc.). When ``None``, :func:`~lerobot.configs.video.depth_encoder_defaults` is used by the writer. encoder_threads (int | None, optional): Number of encoder threads (global). ``None`` lets the codec decide. @@ -206,6 +211,7 @@ class LeRobotDataset(torch.utils.data.Dataset): self.revision = revision if revision else CODEBASE_VERSION self._video_backend = video_backend if video_backend else get_safe_default_video_backend() self._return_uint8 = return_uint8 + self._depth_output_unit = depth_output_unit self._batch_encoding_size = batch_encoding_size self._encoder_threads = encoder_threads @@ -246,6 +252,7 @@ class LeRobotDataset(torch.utils.data.Dataset): delta_timestamps=delta_timestamps, image_transforms=image_transforms, return_uint8=self._return_uint8, + depth_output_unit=self._depth_output_unit, ) self.image_transforms = image_transforms @@ -271,14 +278,16 @@ class LeRobotDataset(torch.utils.data.Dataset): if streaming_encoding and len(self.meta.video_keys) > 0: streaming_enc = self._build_streaming_encoder( self.meta.fps, - camera_encoder, + rgb_encoder, + depth_encoder, encoder_queue_maxsize, encoder_threads, ) self.writer = DatasetWriter( meta=self.meta, root=self.root, - camera_encoder=camera_encoder, + rgb_encoder=rgb_encoder, + depth_encoder=depth_encoder, encoder_threads=encoder_threads, batch_encoding_size=batch_encoding_size, streaming_encoder=streaming_enc, @@ -314,19 +323,22 @@ class LeRobotDataset(torch.utils.data.Dataset): delta_timestamps=self.delta_timestamps, image_transforms=self.image_transforms, return_uint8=self._return_uint8, + depth_output_unit=self._depth_output_unit, ) return self.reader @staticmethod def _build_streaming_encoder( fps: int, - camera_encoder: VideoEncoderConfig | None, + rgb_encoder: RGBEncoderConfig | None, + depth_encoder: DepthEncoderConfig | None, encoder_queue_maxsize: int, encoder_threads: int | None, ) -> StreamingVideoEncoder: return StreamingVideoEncoder( fps=fps, - camera_encoder=camera_encoder, + rgb_encoder=rgb_encoder, + depth_encoder=depth_encoder, queue_maxsize=encoder_queue_maxsize, encoder_threads=encoder_threads, ) @@ -655,7 +667,8 @@ class LeRobotDataset(torch.utils.data.Dataset): image_writer_threads: int = 0, video_backend: str | None = None, batch_encoding_size: int = 1, - camera_encoder: VideoEncoderConfig | None = None, + rgb_encoder: RGBEncoderConfig | None = None, + depth_encoder: DepthEncoderConfig | None = None, metadata_buffer_size: int = 10, streaming_encoding: bool = False, encoder_queue_maxsize: int = 30, @@ -686,8 +699,10 @@ class LeRobotDataset(torch.utils.data.Dataset): video_backend: Video decoding backend (used when reading back). batch_encoding_size: Number of episodes to accumulate before batch-encoding videos. ``1`` means encode immediately. - camera_encoder: Video encoder settings for cameras (codec, quality, etc.). - When ``None``, :func:`~lerobot.configs.video.camera_encoder_defaults` is used. + rgb_encoder: Video encoder settings for cameras (codec, quality, etc.). + When ``None``, :func:`~lerobot.configs.video.rgb_encoder_defaults` is used. + depth_encoder: Video encoder settings for depth cameras (codec, quality, etc.). + When ``None``, :func:`~lerobot.configs.video.depth_encoder_defaults` is used. encoder_threads: Number of encoder threads (global). ``None`` lets the codec decide. metadata_buffer_size: Number of episode metadata records to buffer @@ -722,6 +737,7 @@ class LeRobotDataset(torch.utils.data.Dataset): obj.episodes = None obj._video_backend = video_backend if video_backend is not None else get_safe_default_video_backend() obj._return_uint8 = False + obj._depth_output_unit = DEFAULT_DEPTH_UNIT obj._batch_encoding_size = batch_encoding_size obj._encoder_threads = encoder_threads @@ -731,12 +747,13 @@ class LeRobotDataset(torch.utils.data.Dataset): streaming_enc = None if streaming_encoding and len(obj.meta.video_keys) > 0: streaming_enc = cls._build_streaming_encoder( - fps, camera_encoder, encoder_queue_maxsize, encoder_threads + fps, rgb_encoder, depth_encoder, encoder_queue_maxsize, encoder_threads ) obj.writer = DatasetWriter( meta=obj.meta, root=obj.root, - camera_encoder=camera_encoder, + rgb_encoder=rgb_encoder, + depth_encoder=depth_encoder, encoder_threads=encoder_threads, batch_encoding_size=batch_encoding_size, streaming_encoder=streaming_enc, @@ -759,7 +776,8 @@ class LeRobotDataset(torch.utils.data.Dataset): force_cache_sync: bool = False, video_backend: str | None = None, batch_encoding_size: int = 1, - camera_encoder: VideoEncoderConfig | None = None, + rgb_encoder: RGBEncoderConfig | None = None, + depth_encoder: DepthEncoderConfig | None = None, encoder_threads: int | None = None, image_writer_processes: int = 0, image_writer_threads: int = 0, @@ -787,8 +805,10 @@ class LeRobotDataset(torch.utils.data.Dataset): video_backend: Video decoding backend for reading back data. batch_encoding_size: Number of episodes to accumulate before batch-encoding videos. - camera_encoder: Video encoder settings for cameras (codec, quality, etc.). - When ``None``, :func:`~lerobot.configs.video.camera_encoder_defaults` is used. + rgb_encoder: Video encoder settings for cameras (codec, quality, etc.). + When ``None``, :func:`~lerobot.configs.video.rgb_encoder_defaults` is used. + depth_encoder: Video encoder settings for depth cameras (codec, quality, etc.). + When ``None``, :func:`~lerobot.configs.video.depth_encoder_defaults` is used. encoder_threads: Number of encoder threads (global). ``None`` lets the codec decide. image_writer_processes: Subprocesses for async image writing. @@ -816,6 +836,7 @@ class LeRobotDataset(torch.utils.data.Dataset): obj.episodes = None obj._video_backend = video_backend if video_backend else get_safe_default_video_backend() obj._return_uint8 = False + obj._depth_output_unit = DEFAULT_DEPTH_UNIT obj._batch_encoding_size = batch_encoding_size if obj._requested_root is not None: @@ -835,12 +856,13 @@ class LeRobotDataset(torch.utils.data.Dataset): streaming_enc = None if streaming_encoding and len(obj.meta.video_keys) > 0: streaming_enc = cls._build_streaming_encoder( - obj.meta.fps, camera_encoder, encoder_queue_maxsize, encoder_threads + obj.meta.fps, rgb_encoder, depth_encoder, encoder_queue_maxsize, encoder_threads ) obj.writer = DatasetWriter( meta=obj.meta, root=obj.root, - camera_encoder=camera_encoder, + rgb_encoder=rgb_encoder, + depth_encoder=depth_encoder, encoder_threads=encoder_threads, batch_encoding_size=batch_encoding_size, streaming_encoder=streaming_enc, diff --git a/src/lerobot/datasets/pyav_utils.py b/src/lerobot/datasets/pyav_utils.py index d291f8b40..7b7d1e5de 100644 --- a/src/lerobot/datasets/pyav_utils.py +++ b/src/lerobot/datasets/pyav_utils.py @@ -24,6 +24,7 @@ import logging from typing import Any import av +import numpy as np logger = logging.getLogger(__name__) @@ -31,6 +32,34 @@ FFMPEG_NUMERIC_OPTION_TYPES = ("INT", "INT64", "UINT64", "FLOAT", "DOUBLE") FFMPEG_INTEGER_OPTION_TYPES = ("INT", "INT64", "UINT64") +def write_u16_plane(plane: av.video.plane.VideoPlane, src: np.ndarray, fill_value: int | None = None) -> None: + """Copy a 2D ``uint16`` image into the plane's memory buffer, row by row. + + For speed, each row is padded to a wider size than ``width``, so the true row width in + memory is ``plane.line_size`` (bytes), not ``width``. Copying as one straight stream + would skew the image, so we write only the first ``width`` columns of each row and + leave the padding untouched. + + Args: + plane: Destination 16-bit plane. + src: Source image, shape ``(height, width)``, dtype ``uint16``. + fill_value: If given, every pixel (padding included) is set to this first, so the + padding holds clean data instead of garbage. + """ + height, width = src.shape + stride_u16 = plane.line_size // np.dtype(np.uint16).itemsize + dst = np.frombuffer(plane, dtype=np.uint16).reshape(height, stride_u16) + if fill_value is not None: + dst.fill(fill_value) + dst[:, :width] = src + + +@functools.cache +def get_pix_fmt_channels(pix_fmt: str) -> int: + """Return the number of components (channels) for *pix_fmt*.""" + return len(av.VideoFormat(pix_fmt).components) + + @functools.cache def get_codec(vcodec: str) -> av.codec.Codec | None: """PyAV write-mode ``Codec`` for *vcodec*, or ``None`` if unavailable.""" @@ -92,7 +121,7 @@ def _check_option_value(vcodec: str, label: str, value: Any, opt: av.option.Opti f"{label}={value!r} is not numeric; codec {vcodec!r} expects a number for this option." ) from e elif isinstance(value, (float, int)): - num_val = value + num_val = float(value) else: raise ValueError( f"{label}={value!r} is not numeric; codec {vcodec!r} expects a number for this option." @@ -142,6 +171,16 @@ def _check_pixel_format(vcodec: str, pix_fmt: str) -> None: ) +def _check_pix_fmt_channels(pix_fmt: str, channels: int) -> None: + """Ensure *pix_fmt* can carry at least *channels* components.""" + pix_fmt_channels = get_pix_fmt_channels(pix_fmt) + if pix_fmt_channels < channels: + raise ValueError( + f"pix_fmt={pix_fmt!r} carries only {pix_fmt_channels} component(s) " + f"but the source data has {channels} channel(s)." + ) + + def _check_codec_options(vcodec: str, codec_options: dict[str, Any]) -> None: """Validate merged encoder options (typed) against the codec's published AVOptions.""" supported_options = _get_codec_options_by_name(vcodec) @@ -156,12 +195,18 @@ def _check_codec_options(vcodec: str, codec_options: dict[str, Any]) -> None: _check_option_value(vcodec, key, value, supported_options[key]) -def check_video_encoder_parameters_pyav(vcodec: str, pix_fmt: str, codec_options: dict[str, Any]) -> None: +def check_video_encoder_parameters_pyav( + vcodec: str, + pix_fmt: str, + codec_options: dict[str, Any], + channels: int | None = None, +) -> None: """Verify *config* is compatible with the bundled FFmpeg build. Checks pixel format, abstract tuning-field compatibility, and each merged encoder option from :meth:`~lerobot.configs.video.VideoEncoderConfig.get_codec_options` against PyAV (including numeric ``extra_options`` present in that dict). + When given, additionally verify that *pix_fmt* carries as many components as the source data channels. No-op when ``config.vcodec`` isn't in the local FFmpeg build. Raises: @@ -171,4 +216,6 @@ def check_video_encoder_parameters_pyav(vcodec: str, pix_fmt: str, codec_options if not options: raise ValueError(f"Codec {vcodec!r} is not available in the bundled FFmpeg build") _check_pixel_format(vcodec, pix_fmt) + if channels is not None: + _check_pix_fmt_channels(pix_fmt, channels) _check_codec_options(vcodec, codec_options) diff --git a/src/lerobot/datasets/streaming_dataset.py b/src/lerobot/datasets/streaming_dataset.py index 3c1e4a73c..4c4ae59bf 100644 --- a/src/lerobot/datasets/streaming_dataset.py +++ b/src/lerobot/datasets/streaming_dataset.py @@ -22,9 +22,11 @@ import numpy as np import torch from datasets import load_dataset +from lerobot.configs import DEFAULT_DEPTH_UNIT, DepthEncoderConfig from lerobot.utils.constants import HF_LEROBOT_HOME, LOOKAHEAD_BACKTRACKTABLE, LOOKBACK_BACKTRACKTABLE from .dataset_metadata import CODEBASE_VERSION, LeRobotDatasetMetadata +from .depth_utils import dequantize_depth from .feature_utils import get_delta_indices from .io_utils import item_to_torch from .utils import ( @@ -35,6 +37,7 @@ from .utils import ( ) from .video_utils import ( VideoDecoderCache, + decode_video_frames, decode_video_frames_torchcodec, ) @@ -252,6 +255,7 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): rng: np.random.Generator | None = None, shuffle: bool = True, return_uint8: bool = False, + depth_output_unit: str = DEFAULT_DEPTH_UNIT, ): """Initialize a StreamingLeRobotDataset. @@ -272,6 +276,8 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): seed (int, optional): Reproducibility random seed. rng (np.random.Generator | None, optional): Random number generator. shuffle (bool, optional): Whether to shuffle the dataset across exhaustions. Defaults to True. + depth_output_unit (str, optional): Physical unit depth maps are dequantized to ("m" or "mm"). + Defaults to "mm". """ super().__init__() self.repo_id = repo_id @@ -290,6 +296,7 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): self.streaming = streaming self.buffer_size = buffer_size self._return_uint8 = return_uint8 + self._depth_output_unit = depth_output_unit # We cache the video decoders to avoid re-initializing them at each frame (avoiding a ~10x slowdown) self.video_decoder_cache = None @@ -306,6 +313,11 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): # Check version check_version_compatibility(self.repo_id, self.meta._version, CODEBASE_VERSION) + self._depth_encoder_configs: dict[str, DepthEncoderConfig] = { + vid_key: DepthEncoderConfig.from_video_info(self.meta.features[vid_key].get("info")) + for vid_key in self.meta.depth_keys + } + self.delta_timestamps = None self.delta_indices = None @@ -554,13 +566,34 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): for video_key, query_ts in query_timestamps.items(): root = self.meta.url_root if self.streaming and not self.streaming_from_local else self.root video_path = f"{root}/{self.meta.get_video_file_path(ep_idx, video_key)}" - frames = decode_video_frames_torchcodec( - video_path, - query_ts, - self.tolerance_s, - decoder_cache=self.video_decoder_cache, - return_uint8=self._return_uint8, - ) + if video_key in self.meta.depth_keys: + # Depth maps are 12-bit quantized and only decodable via pyav; dequantize back + # to physical units to match the non-streaming reader. + frames = decode_video_frames( + video_path, + query_ts, + self.tolerance_s, + backend="pyav", + return_uint8=False, + is_depth=True, + ) + depth_encoder = self._depth_encoder_configs[video_key] + frames = dequantize_depth( + frames, + depth_min=depth_encoder.depth_min, + depth_max=depth_encoder.depth_max, + shift=depth_encoder.shift, + use_log=depth_encoder.use_log, + output_unit=self._depth_output_unit, + ) + else: + frames = decode_video_frames_torchcodec( + video_path, + query_ts, + self.tolerance_s, + decoder_cache=self.video_decoder_cache, + return_uint8=self._return_uint8, + ) item[video_key] = frames.squeeze(0) if len(query_ts) == 1 else frames diff --git a/src/lerobot/datasets/utils.py b/src/lerobot/datasets/utils.py index de91978ea..d30761515 100644 --- a/src/lerobot/datasets/utils.py +++ b/src/lerobot/datasets/utils.py @@ -87,11 +87,14 @@ DATA_DIR = "data" VIDEO_DIR = "videos" CHUNK_FILE_PATTERN = "chunk-{chunk_index:03d}/file-{file_index:03d}" +IMAGE_FILE_PATTERN = "frame-{frame_index:06d}.png" +DEPTH_FILE_PATTERN = "frame-{frame_index:06d}.tiff" DEFAULT_TASKS_PATH = "meta/tasks.parquet" DEFAULT_EPISODES_PATH = EPISODES_DIR + "/" + CHUNK_FILE_PATTERN + ".parquet" DEFAULT_DATA_PATH = DATA_DIR + "/" + CHUNK_FILE_PATTERN + ".parquet" DEFAULT_VIDEO_PATH = VIDEO_DIR + "/{video_key}/" + CHUNK_FILE_PATTERN + ".mp4" -DEFAULT_IMAGE_PATH = "images/{image_key}/episode-{episode_index:06d}/frame-{frame_index:06d}.png" +DEFAULT_IMAGE_PATH = "images/{image_key}/episode-{episode_index:06d}/" + IMAGE_FILE_PATTERN +DEFAULT_DEPTH_PATH = "images/{image_key}/episode-{episode_index:06d}/" + DEPTH_FILE_PATTERN LEGACY_EPISODES_PATH = "meta/episodes.jsonl" LEGACY_EPISODES_STATS_PATH = "meta/episodes_stats.jsonl" diff --git a/src/lerobot/datasets/video_utils.py b/src/lerobot/datasets/video_utils.py index ca90fba45..ef3005dd8 100644 --- a/src/lerobot/datasets/video_utils.py +++ b/src/lerobot/datasets/video_utils.py @@ -39,11 +39,17 @@ from datasets.features.features import register_feature from PIL import Image from lerobot.configs import ( + DepthEncoderConfig, + RGBEncoderConfig, VideoEncoderConfig, - camera_encoder_defaults, + depth_encoder_defaults, + rgb_encoder_defaults, ) from lerobot.utils.import_utils import get_safe_default_video_backend +from .depth_utils import quantize_depth +from .pyav_utils import get_pix_fmt_channels + logger = logging.getLogger(__name__) @@ -53,6 +59,7 @@ def decode_video_frames( tolerance_s: float, backend: str | None = None, return_uint8: bool = False, + is_depth: bool = False, ) -> torch.Tensor: """ Decodes video frames using the specified backend. @@ -64,23 +71,35 @@ def decode_video_frames( backend (str, optional): Backend to use for decoding. Defaults to "torchcodec" when available in the platform; otherwise, defaults to "pyav". The legacy value "video_reader" is accepted for one release as an alias for "pyav" and will be removed in a future version. - return_uint8 (bool): If True, return raw uint8 frames without float32 normalization. + return_uint8 (bool): For RGB videos, if True return raw uint8 frames without float32 normalization. This reduces memory for DataLoader IPC; normalization can be done on GPU afterward. + is_depth (bool): Set to True if the video is a depth map (1 channel, uint12). Returns: - torch.Tensor: Decoded frames (float32 in [0,1] by default, or uint8 if return_uint8=True). + torch.Tensor: Decoded frames (RGB: float32 in [0,1] by default, or uint8 if return_uint8=True, Depth: uint12). Currently supports torchcodec on cpu and pyav. """ + if backend != "pyav" and is_depth: + logger.debug("Decoding depth maps is only supported with the 'pyav' backend, falling back to pyav.") + # We do not actually return uint8 here, but we avoid the 255 normalization step. + return decode_video_frames_pyav( + video_path, timestamps, tolerance_s, return_uint8=False, is_depth=True + ) + if backend is None: backend = get_safe_default_video_backend() if backend == "torchcodec": return decode_video_frames_torchcodec(video_path, timestamps, tolerance_s, return_uint8=return_uint8) elif backend == "pyav": - return decode_video_frames_pyav(video_path, timestamps, tolerance_s, return_uint8=return_uint8) + return decode_video_frames_pyav( + video_path, timestamps, tolerance_s, return_uint8=return_uint8, is_depth=is_depth + ) elif backend == "video_reader": logger.warning("backend='video_reader' is deprecated and now aliases to 'pyav'.") - return decode_video_frames_pyav(video_path, timestamps, tolerance_s, return_uint8=return_uint8) + return decode_video_frames_pyav( + video_path, timestamps, tolerance_s, return_uint8=return_uint8, is_depth=is_depth + ) else: raise ValueError(f"Unsupported video backend: {backend}") @@ -91,6 +110,7 @@ def decode_video_frames_pyav( tolerance_s: float, log_loaded_timestamps: bool = False, return_uint8: bool = False, + is_depth: bool = False, ) -> torch.Tensor: """Loads frames associated to the requested timestamps of a video using PyAV. @@ -109,8 +129,9 @@ def decode_video_frames_pyav( tolerance_s: Allowed deviation in seconds between a queried timestamp and the closest decoded frame. log_loaded_timestamps: When True, log every decoded frame's timestamp at INFO level. - return_uint8: When True, return raw uint8 frames (C, H, W). Otherwise, return float32 in - [0, 1] range. + return_uint8: For RGB videos, if True return raw uint8 frames (C, H, W). + Otherwise, return float32 in [0, 1] range. + is_depth: Set to True if the video is a depth map (1 channel, uint12). Returns: torch.Tensor of shape (len(timestamps), C, H, W). @@ -132,7 +153,13 @@ def decode_video_frames_pyav( # https://pyav.basswood-io.com/docs/stable/api/container.html#av.container.InputContainer.seek with av.open(video_path) as container: stream = container.streams.video[0] - container.seek(int(first_ts * av.time_base), backward=True) + # Seek to the nearest keyframe at or before `first_ts` with a 1 frame margin + container.seek( + round(first_ts / stream.time_base) - 1, + backward=True, + any_frame=False, + stream=stream, + ) for frame in container.decode(stream): if frame.pts is None: @@ -140,9 +167,13 @@ def decode_video_frames_pyav( current_ts = float(frame.pts * stream.time_base) if log_loaded_timestamps: logger.info(f"frame loaded at timestamp={current_ts:.4f}") - # Convert to CHW uint8 to match torchcodec's output layout. - arr = frame.to_ndarray(format="rgb24") # H, W, 3 - loaded_frames.append(torch.from_numpy(arr).permute(2, 0, 1).contiguous()) + if is_depth: + arr = frame.to_ndarray(format="gray12le") # (H, W) uint12 + loaded_frames.append(torch.from_numpy(arr).unsqueeze(0).contiguous()) + else: + arr = frame.to_ndarray(format="rgb24") # (H, W, 3) + # Convert to CHW uint8 to match torchcodec's output layout. + loaded_frames.append(torch.from_numpy(arr).permute(2, 0, 1).contiguous()) loaded_ts.append(current_ts) if current_ts >= last_ts: break @@ -185,7 +216,7 @@ def decode_video_frames_pyav( f"number of queried timestamps ({len(timestamps)})" ) - if return_uint8: + if return_uint8 or is_depth: return closest_frames # convert to the pytorch format which is float32 in [0,1] range (and channel first) @@ -406,17 +437,38 @@ def encode_video_frames( imgs_dir: Path | str, video_path: Path | str, fps: int, - camera_encoder: VideoEncoderConfig | None = None, + video_encoder: VideoEncoderConfig | None = None, encoder_threads: int | None = None, *, log_level: int | None = av.logging.WARNING, overwrite: bool = False, ) -> None: - """More info on ffmpeg arguments tuning on `benchmark/video/README.md`""" - if camera_encoder is None: - camera_encoder = camera_encoder_defaults() - vcodec = camera_encoder.vcodec - pix_fmt = camera_encoder.pix_fmt + """Encode a directory of image frames into an MP4 video. + + When ``video_encoder`` is a :class:`~lerobot.configs.video.DepthEncoderConfig`, + frames are read from ``.tiff`` files and quantized to 12-bit depth codes using the + encoder's ``depth_min`` / ``depth_max`` / ``shift`` / ``use_log``; otherwise ``.png`` + RGB frames are encoded directly. + + Args: + imgs_dir: Directory containing the frames to encode, named ``frame-000000`` + onwards (``.png`` for RGB, ``.tiff`` for depth). + video_path: Output path for the encoded ``.mp4`` file. + fps: Frame rate of the output video. + video_encoder: Encoder settings (codec, pixel format, quality, ...). When + ``None``, :func:`rgb_encoder_defaults` is used. Pass a + :class:`~lerobot.configs.video.DepthEncoderConfig` to encode depth frames. + encoder_threads: Per-encoder thread count forwarded to the codec. ``None`` + lets the codec decide. + log_level: libav log level to set while encoding, or ``None`` to leave the + current logging configuration unchanged. + overwrite: When ``False`` and ``video_path`` already exists, skip encoding and + log a warning. When ``True``, re-encode and replace the existing file. + """ + if video_encoder is None: + video_encoder = rgb_encoder_defaults() + vcodec = video_encoder.vcodec + pix_fmt = video_encoder.pix_fmt video_path = Path(video_path) imgs_dir = Path(imgs_dir) @@ -428,17 +480,19 @@ def encode_video_frames( video_path.parent.mkdir(parents=True, exist_ok=True) # Get input frames - template = "frame-" + ("[0-9]" * 6) + ".png" + is_depth = isinstance(video_encoder, DepthEncoderConfig) + suffix = ".png" if not is_depth else ".tiff" + template = "frame-" + ("[0-9]" * 6) + suffix input_list = sorted( glob.glob(str(imgs_dir / template)), key=lambda x: int(x.split("-")[-1].split(".")[0]) ) if len(input_list) == 0: - raise FileNotFoundError(f"No images found in {imgs_dir}.") + raise FileNotFoundError(f"No images with suffix {suffix} found in {imgs_dir}.") with Image.open(input_list[0]) as dummy_image: width, height = dummy_image.size - video_options = camera_encoder.get_codec_options(encoder_threads, as_strings=True) + video_options = video_encoder.get_codec_options(encoder_threads, as_strings=True) # Set logging level if log_level is not None: @@ -455,8 +509,19 @@ def encode_video_frames( # Loop through input frames and encode them for input_data in input_list: with Image.open(input_data) as input_image: - input_image = input_image.convert("RGB") - input_frame = av.VideoFrame.from_image(input_image) + if is_depth: + input_frame = quantize_depth( + np.array(input_image), + depth_min=video_encoder.depth_min, + depth_max=video_encoder.depth_max, + shift=video_encoder.shift, + use_log=video_encoder.use_log, + pix_fmt=video_encoder.pix_fmt, + video_backend="pyav", + ) + else: + input_image = input_image.convert("RGB") + input_frame = av.VideoFrame.from_image(input_image) packet = output_stream.encode(input_frame) if packet: output.mux(packet) @@ -477,7 +542,7 @@ def encode_video_frames( def reencode_video( input_video_path: Path | str, output_video_path: Path | str, - camera_encoder: VideoEncoderConfig | None = None, + video_encoder: VideoEncoderConfig | None = None, encoder_threads: int | None = None, log_level: int | None = av.logging.WARNING, overwrite: bool = False, @@ -489,7 +554,7 @@ def reencode_video( Args: input_video_path: Existing video file to read. output_video_path: Path for the re-encoded file. - camera_encoder: Encoder configuration. Defaults to :func:`camera_encoder_defaults`. + video_encoder: Encoder configuration. Defaults to :func:`rgb_encoder_defaults`. encoder_threads: Optional thread count forwarded to :meth:`VideoEncoderConfig.get_codec_options`. log_level: libav log level while encoding, or ``None`` to leave logging unchanged. Defaults to WARNING. overwrite: When ``False`` and ``output_video_path`` already exists, skip and log a warning. @@ -497,7 +562,7 @@ def reencode_video( end_time_s: When set, trim the output to end at this timestamp (seconds, exclusive). """ - camera_encoder = camera_encoder or camera_encoder_defaults() + video_encoder = video_encoder or rgb_encoder_defaults() if (start_time_s is not None and start_time_s < 0) or (end_time_s is not None and end_time_s < 0): raise ValueError(f"Trim times must be non-negative, got start={start_time_s}, end={end_time_s}.") @@ -512,9 +577,9 @@ def reencode_video( output_video_path.parent.mkdir(parents=True, exist_ok=True) - video_options = camera_encoder.get_codec_options(encoder_threads, as_strings=True) - vcodec = camera_encoder.vcodec - pix_fmt = camera_encoder.pix_fmt + video_options = video_encoder.get_codec_options(encoder_threads, as_strings=True) + vcodec = video_encoder.vcodec + pix_fmt = video_encoder.pix_fmt with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp_named_file: tmp_output_video_path = tmp_named_file.name @@ -696,22 +761,21 @@ class _CameraEncoderThread(threading.Thread): self, video_path: Path, fps: int, - vcodec: str, - pix_fmt: str, - codec_options: dict[str, str], + video_encoder: VideoEncoderConfig, frame_queue: queue.Queue, result_queue: queue.Queue, stop_event: threading.Event, + encoder_threads: int | None = None, ): super().__init__(daemon=True) self.video_path = video_path self.fps = fps - self.vcodec = vcodec - self.pix_fmt = pix_fmt - self.codec_options = codec_options + self.video_encoder = video_encoder + self.is_depth = isinstance(video_encoder, DepthEncoderConfig) self.frame_queue = frame_queue self.result_queue = result_queue self.stop_event = stop_event + self.encoder_threads = encoder_threads def run(self) -> None: from .compute_stats import RunningQuantileStats, auto_downsample_height_width @@ -736,12 +800,12 @@ class _CameraEncoderThread(threading.Thread): # Sentinel: flush and close break - # Ensure HWC uint8 numpy array + # Ensure HWC (RGB or depth) uint8 (RGB only) numpy array if isinstance(frame_data, np.ndarray): - if frame_data.ndim == 3 and frame_data.shape[0] == 3: + if frame_data.ndim == 3 and frame_data.shape[0] in (1, 3): # CHW -> HWC frame_data = frame_data.transpose(1, 2, 0) - if frame_data.dtype != np.uint8: + if not self.is_depth and frame_data.dtype != np.uint8: frame_data = (frame_data * 255).astype(np.uint8) # Open container on first frame (to get width/height) @@ -749,15 +813,29 @@ class _CameraEncoderThread(threading.Thread): height, width = frame_data.shape[:2] Path(self.video_path).parent.mkdir(parents=True, exist_ok=True) container = av.open(str(self.video_path), "w") - output_stream = container.add_stream(self.vcodec, self.fps, options=self.codec_options) - output_stream.pix_fmt = self.pix_fmt + output_stream = container.add_stream( + self.video_encoder.vcodec, + self.fps, + options=self.video_encoder.get_codec_options(self.encoder_threads, as_strings=True), + ) + output_stream.pix_fmt = self.video_encoder.pix_fmt output_stream.width = width output_stream.height = height output_stream.time_base = Fraction(1, self.fps) # Encode frame with explicit timestamps - pil_img = Image.fromarray(frame_data) - video_frame = av.VideoFrame.from_image(pil_img) + if not self.is_depth: + pil_img = Image.fromarray(frame_data) + video_frame = av.VideoFrame.from_image(pil_img) + else: + video_frame = quantize_depth( + frame_data, + depth_min=self.video_encoder.depth_min, + depth_max=self.video_encoder.depth_max, + shift=self.video_encoder.shift, + use_log=self.video_encoder.use_log, + video_backend=self.video_encoder.video_backend, + ) video_frame.pts = frame_count video_frame.time_base = Fraction(1, self.fps) packet = output_stream.encode(video_frame) @@ -815,22 +893,27 @@ class StreamingVideoEncoder: def __init__( self, fps: int, - camera_encoder: VideoEncoderConfig | None = None, + rgb_encoder: RGBEncoderConfig | None = None, + depth_encoder: DepthEncoderConfig | None = None, queue_maxsize: int = 30, encoder_threads: int | None = None, ): """ Args: fps: Frames per second for the output videos. - camera_encoder: Video encoder settings applied to all cameras. - When ``None``, :func:`camera_encoder_defaults` is used. - encoder_threads: Number of encoder threads (global setting). - ``None`` lets the codec decide. + rgb_encoder: Video encoder settings applied to all RGB cameras. + When ``None``, :func:`rgb_encoder_defaults` is used. + depth_encoder: Video encoder settings applied to all depth cameras, + including the depth quantization parameters. When ``None``, + :func:`depth_encoder_defaults` is used. queue_maxsize: Max frames to buffer per camera before back-pressure drops frames. + encoder_threads: Number of encoder threads (global setting). + ``None`` lets the codec decide. """ self.fps = fps - self._camera_encoder = camera_encoder or camera_encoder_defaults() + self._rgb_encoder = rgb_encoder or rgb_encoder_defaults() + self._depth_encoder = depth_encoder or depth_encoder_defaults() self._encoder_threads = encoder_threads self.queue_maxsize = queue_maxsize @@ -843,18 +926,25 @@ class StreamingVideoEncoder: self._episode_active = False self._closed = False - def start_episode(self, video_keys: list[str], temp_dir: Path) -> None: + def start_episode( + self, video_keys: list[str], temp_dir: Path, depth_video_keys: list[str] | None = None + ) -> None: """Start encoder threads for a new episode. Args: video_keys: List of video feature keys (e.g. ["observation.images.laptop"]) temp_dir: Base directory for temporary MP4 files + depth_video_keys: List of video or image feature keys that carry depth maps (e.g. + ["observation.images.laptop_depth"]). Defaults to ``[]`` (no depth keys). """ if self._episode_active: self.cancel_episode() self._dropped_frames.clear() + if depth_video_keys is None: + depth_video_keys = [] + for video_key in video_keys: frame_queue: queue.Queue = queue.Queue(maxsize=self.queue_maxsize) result_queue: queue.Queue = queue.Queue(maxsize=1) @@ -863,17 +953,15 @@ class StreamingVideoEncoder: temp_video_dir = Path(tempfile.mkdtemp(dir=temp_dir)) video_path = temp_video_dir / f"{video_key.replace('/', '_')}_streaming.mp4" - vcodec = self._camera_encoder.vcodec - codec_options = self._camera_encoder.get_codec_options(self._encoder_threads, as_strings=True) + encoder = self._depth_encoder if video_key in depth_video_keys else self._rgb_encoder encoder_thread = _CameraEncoderThread( video_path=video_path, fps=self.fps, - vcodec=vcodec, - pix_fmt=self._camera_encoder.pix_fmt, - codec_options=codec_options, + video_encoder=encoder, frame_queue=frame_queue, result_queue=result_queue, stop_event=stop_event, + encoder_threads=self._encoder_threads, ) encoder_thread.start() @@ -1080,15 +1168,23 @@ def get_audio_info(video_path: Path | str) -> dict: def get_video_info( video_path: Path | str, - camera_encoder: VideoEncoderConfig | None = None, + video_encoder: VideoEncoderConfig | None = None, ) -> dict: """Build the ``video.*`` / ``audio.*`` info dict persisted in ``info.json``. Args: video_path: Path to the encoded video file to probe. - camera_encoder: If provided, record the exact encoder settings used to encode this + video_encoder: If provided, record the exact encoder settings used to encode this video. Stream-derived values take precedence — encoder fields are only written for keys - not already populated from the video file itself. + not already populated from the video file itself. When a + :class:`~lerobot.configs.video.DepthEncoderConfig` is passed, the depth + quantization parameters (``depth_min`` / ``depth_max`` / ``shift`` / + ``use_log``) are recorded so frames can be dequantized on read. + + Returns: + The ``video.*`` / ``audio.*`` info dict, including ``is_depth_map`` which is + ``True`` only when ``video_encoder`` is a + :class:`~lerobot.configs.video.DepthEncoderConfig`. """ logging.getLogger("libav").setLevel(av.logging.WARNING) @@ -1106,13 +1202,10 @@ def get_video_info( video_info["video.width"] = video_stream.width video_info["video.codec"] = video_stream.codec.canonical_name video_info["video.pix_fmt"] = video_stream.pix_fmt - video_info["video.is_depth_map"] = False # Calculate fps from r_frame_rate video_info["video.fps"] = int(video_stream.base_rate) - - pixel_channels = get_video_pixel_channels(video_stream.pix_fmt) - video_info["video.channels"] = pixel_channels + video_info["video.channels"] = get_pix_fmt_channels(video_stream.pix_fmt) # Reset logging level av.logging.restore_default_callback() @@ -1121,27 +1214,18 @@ def get_video_info( video_info.update(**get_audio_info(video_path)) # Add additional encoder configuration if provided - if camera_encoder is not None: - for field_name, field_value in asdict(camera_encoder).items(): + if video_encoder is not None: + for field_name, field_value in asdict(video_encoder).items(): # vcodec is already populated from the video stream if field_name == "vcodec": continue video_info.setdefault(f"video.{field_name}", field_value) + video_info["is_depth_map"] = isinstance(video_encoder, DepthEncoderConfig) + return video_info -def get_video_pixel_channels(pix_fmt: str) -> int: - if "gray" in pix_fmt or "depth" in pix_fmt or "monochrome" in pix_fmt: - return 1 - elif "rgba" in pix_fmt or "yuva" in pix_fmt: - return 4 - elif "rgb" in pix_fmt or "yuv" in pix_fmt: - return 3 - else: - raise ValueError("Unknown format") - - def get_video_duration_in_s(video_path: Path | str) -> float: """ Get the duration of a video file in seconds using PyAV. @@ -1202,10 +1286,13 @@ class VideoEncodingManager: img_dir = self.dataset.root / "images" if img_dir.exists(): png_files = list(img_dir.rglob("*.png")) - if len(png_files) == 0: + tiff_files = list(img_dir.rglob("*.tiff")) + if len(png_files) == 0 and len(tiff_files) == 0: shutil.rmtree(img_dir) logger.debug("Cleaned up empty images directory") else: - logger.debug(f"Images directory is not empty, containing {len(png_files)} PNG files") + logger.debug( + f"Images directory is not empty, containing {len(png_files)} PNG and {len(tiff_files)} TIFF files" + ) return False # Don't suppress the original exception diff --git a/src/lerobot/policies/utils.py b/src/lerobot/policies/utils.py index c37127813..f465fcff8 100644 --- a/src/lerobot/policies/utils.py +++ b/src/lerobot/policies/utils.py @@ -126,7 +126,8 @@ def prepare_observation_for_inference( for name in observation: observation[name] = torch.from_numpy(observation[name]) if "image" in name: - observation[name] = observation[name].type(torch.float32) / 255 + if observation[name].dtype == torch.uint8: + observation[name] = observation[name].type(torch.float32) / 255 observation[name] = observation[name].permute(2, 0, 1).contiguous() observation[name] = observation[name].unsqueeze(0) observation[name] = observation[name].to(device) diff --git a/src/lerobot/robots/hope_jr/hope_jr_arm.py b/src/lerobot/robots/hope_jr/hope_jr_arm.py index 4918bcae3..b606a4fe7 100644 --- a/src/lerobot/robots/hope_jr/hope_jr_arm.py +++ b/src/lerobot/robots/hope_jr/hope_jr_arm.py @@ -66,9 +66,14 @@ class HopeJrArm(Robot): @property def _cameras_ft(self) -> dict[str, tuple]: - return { - cam: (self.config.cameras[cam].height, self.config.cameras[cam].width, 3) for cam in self.cameras - } + features: dict[str, tuple] = {} + for cam in self.cameras: + cfg = self.config.cameras[cam] + if getattr(cfg, "use_rgb", True): + features[cam] = (cfg.height, cfg.width, 3) + if getattr(cfg, "use_depth", False): + features[f"{cam}_depth"] = (cfg.height, cfg.width, 1) + return features @cached_property def observation_features(self) -> dict[str, type | tuple]: @@ -139,10 +144,17 @@ class HopeJrArm(Robot): # Capture images from cameras for cam_key, cam in self.cameras.items(): - start = time.perf_counter() - obs_dict[cam_key] = cam.read_latest() - dt_ms = (time.perf_counter() - start) * 1e3 - logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms") + if getattr(cam, "use_rgb", True): + start = time.perf_counter() + obs_dict[cam_key] = cam.read_latest() + dt_ms = (time.perf_counter() - start) * 1e3 + logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms") + + if getattr(cam, "use_depth", False): + start = time.perf_counter() + obs_dict[f"{cam_key}_depth"] = cam.read_latest_depth() + dt_ms = (time.perf_counter() - start) * 1e3 + logger.debug(f"{self} read {cam_key} depth: {dt_ms:.1f}ms") return obs_dict diff --git a/src/lerobot/robots/hope_jr/hope_jr_hand.py b/src/lerobot/robots/hope_jr/hope_jr_hand.py index 566628724..ce70e7e13 100644 --- a/src/lerobot/robots/hope_jr/hope_jr_hand.py +++ b/src/lerobot/robots/hope_jr/hope_jr_hand.py @@ -102,9 +102,14 @@ class HopeJrHand(Robot): @property def _cameras_ft(self) -> dict[str, tuple]: - return { - cam: (self.config.cameras[cam].height, self.config.cameras[cam].width, 3) for cam in self.cameras - } + features: dict[str, tuple] = {} + for cam in self.cameras: + cfg = self.config.cameras[cam] + if getattr(cfg, "use_rgb", True): + features[cam] = (cfg.height, cfg.width, 3) + if getattr(cfg, "use_depth", False): + features[f"{cam}_depth"] = (cfg.height, cfg.width, 1) + return features @cached_property def observation_features(self) -> dict[str, type | tuple]: @@ -170,10 +175,17 @@ class HopeJrHand(Robot): # Capture images from cameras for cam_key, cam in self.cameras.items(): - start = time.perf_counter() - obs_dict[cam_key] = cam.read_latest() - dt_ms = (time.perf_counter() - start) * 1e3 - logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms") + if getattr(cam, "use_rgb", True): + start = time.perf_counter() + obs_dict[cam_key] = cam.read_latest() + dt_ms = (time.perf_counter() - start) * 1e3 + logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms") + + if getattr(cam, "use_depth", False): + start = time.perf_counter() + obs_dict[f"{cam_key}_depth"] = cam.read_latest_depth() + dt_ms = (time.perf_counter() - start) * 1e3 + logger.debug(f"{self} read {cam_key} depth: {dt_ms:.1f}ms") return obs_dict diff --git a/src/lerobot/robots/koch_follower/koch_follower.py b/src/lerobot/robots/koch_follower/koch_follower.py index 3f40ac738..de6f9c4a3 100644 --- a/src/lerobot/robots/koch_follower/koch_follower.py +++ b/src/lerobot/robots/koch_follower/koch_follower.py @@ -68,9 +68,14 @@ class KochFollower(Robot): @property def _cameras_ft(self) -> dict[str, tuple]: - return { - cam: (self.config.cameras[cam].height, self.config.cameras[cam].width, 3) for cam in self.cameras - } + features: dict[str, tuple] = {} + for cam in self.cameras: + cfg = self.config.cameras[cam] + if getattr(cfg, "use_rgb", True): + features[cam] = (cfg.height, cfg.width, 3) + if getattr(cfg, "use_depth", False): + features[f"{cam}_depth"] = (cfg.height, cfg.width, 1) + return features @cached_property def observation_features(self) -> dict[str, type | tuple]: @@ -192,10 +197,17 @@ class KochFollower(Robot): # Capture images from cameras for cam_key, cam in self.cameras.items(): - start = time.perf_counter() - obs_dict[cam_key] = cam.read_latest() - dt_ms = (time.perf_counter() - start) * 1e3 - logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms") + if getattr(cam, "use_rgb", True): + start = time.perf_counter() + obs_dict[cam_key] = cam.read_latest() + dt_ms = (time.perf_counter() - start) * 1e3 + logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms") + + if getattr(cam, "use_depth", False): + start = time.perf_counter() + obs_dict[f"{cam_key}_depth"] = cam.read_latest_depth() + dt_ms = (time.perf_counter() - start) * 1e3 + logger.debug(f"{self} read {cam_key} depth: {dt_ms:.1f}ms") return obs_dict diff --git a/src/lerobot/robots/lekiwi/lekiwi.py b/src/lerobot/robots/lekiwi/lekiwi.py index b73ebeab9..3712a64d3 100644 --- a/src/lerobot/robots/lekiwi/lekiwi.py +++ b/src/lerobot/robots/lekiwi/lekiwi.py @@ -72,6 +72,12 @@ class LeKiwi(Robot): ) self.arm_motors = [motor for motor in self.bus.motors if motor.startswith("arm")] self.base_motors = [motor for motor in self.bus.motors if motor.startswith("base")] + depth_cameras = [name for name, cfg in config.cameras.items() if getattr(cfg, "use_depth", False)] + if depth_cameras: + raise NotImplementedError( + f"Depth cameras are not supported on LeKiwi (got depth-enabled cameras: {depth_cameras}). " + "The host/client transport only carries color frames." + ) self.cameras = make_cameras_from_configs(config.cameras) @property diff --git a/src/lerobot/robots/lekiwi/lekiwi_client.py b/src/lerobot/robots/lekiwi/lekiwi_client.py index fd43e84fe..1bc3dadc4 100644 --- a/src/lerobot/robots/lekiwi/lekiwi_client.py +++ b/src/lerobot/robots/lekiwi/lekiwi_client.py @@ -44,6 +44,13 @@ class LeKiwiClient(Robot): self.id = config.id self.robot_type = config.type + depth_cameras = [name for name, cfg in config.cameras.items() if getattr(cfg, "use_depth", False)] + if depth_cameras: + raise NotImplementedError( + f"Depth cameras are not supported on LeKiwi (got depth-enabled cameras: {depth_cameras}). " + "The host/client transport only carries color frames." + ) + self.remote_ip = config.remote_ip self.port_zmq_cmd = config.port_zmq_cmd self.port_zmq_observations = config.port_zmq_observations diff --git a/src/lerobot/robots/omx_follower/omx_follower.py b/src/lerobot/robots/omx_follower/omx_follower.py index c30eec97a..b2cfb52e9 100644 --- a/src/lerobot/robots/omx_follower/omx_follower.py +++ b/src/lerobot/robots/omx_follower/omx_follower.py @@ -68,9 +68,14 @@ class OmxFollower(Robot): @property def _cameras_ft(self) -> dict[str, tuple]: - return { - cam: (self.config.cameras[cam].height, self.config.cameras[cam].width, 3) for cam in self.cameras - } + features: dict[str, tuple] = {} + for cam in self.cameras: + cfg = self.config.cameras[cam] + if getattr(cfg, "use_rgb", True): + features[cam] = (cfg.height, cfg.width, 3) + if getattr(cfg, "use_depth", False): + features[f"{cam}_depth"] = (cfg.height, cfg.width, 1) + return features @cached_property def observation_features(self) -> dict[str, type | tuple]: @@ -175,10 +180,17 @@ class OmxFollower(Robot): # Capture images from cameras for cam_key, cam in self.cameras.items(): - start = time.perf_counter() - obs_dict[cam_key] = cam.read_latest() - dt_ms = (time.perf_counter() - start) * 1e3 - logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms") + if getattr(cam, "use_rgb", True): + start = time.perf_counter() + obs_dict[cam_key] = cam.read_latest() + dt_ms = (time.perf_counter() - start) * 1e3 + logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms") + + if getattr(cam, "use_depth", False): + start = time.perf_counter() + obs_dict[f"{cam_key}_depth"] = cam.read_latest_depth() + dt_ms = (time.perf_counter() - start) * 1e3 + logger.debug(f"{self} read {cam_key} depth: {dt_ms:.1f}ms") return obs_dict diff --git a/src/lerobot/robots/openarm_follower/openarm_follower.py b/src/lerobot/robots/openarm_follower/openarm_follower.py index 020f24052..e2c7c8cf5 100644 --- a/src/lerobot/robots/openarm_follower/openarm_follower.py +++ b/src/lerobot/robots/openarm_follower/openarm_follower.py @@ -101,9 +101,14 @@ class OpenArmFollower(Robot): @property def _cameras_ft(self) -> dict[str, tuple]: """Camera features for observation space.""" - return { - cam: (self.config.cameras[cam].height, self.config.cameras[cam].width, 3) for cam in self.cameras - } + features: dict[str, tuple] = {} + for cam in self.cameras: + cfg = self.config.cameras[cam] + if getattr(cfg, "use_rgb", True): + features[cam] = (cfg.height, cfg.width, 3) + if getattr(cfg, "use_depth", False): + features[f"{cam}_depth"] = (cfg.height, cfg.width, 1) + return features @cached_property def observation_features(self) -> dict[str, type | tuple]: @@ -242,10 +247,17 @@ class OpenArmFollower(Robot): # Capture images from cameras for cam_key, cam in self.cameras.items(): - start = time.perf_counter() - obs_dict[cam_key] = cam.read_latest() - dt_ms = (time.perf_counter() - start) * 1e3 - logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms") + if getattr(cam, "use_rgb", True): + start = time.perf_counter() + obs_dict[cam_key] = cam.read_latest() + dt_ms = (time.perf_counter() - start) * 1e3 + logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms") + + if getattr(cam, "use_depth", False): + start = time.perf_counter() + obs_dict[f"{cam_key}_depth"] = cam.read_latest_depth() + dt_ms = (time.perf_counter() - start) * 1e3 + logger.debug(f"{self} read {cam_key} depth: {dt_ms:.1f}ms") dt_ms = (time.perf_counter() - start) * 1e3 logger.debug(f"{self} get_observation took: {dt_ms:.1f}ms") diff --git a/src/lerobot/robots/rebot_b601_follower/rebot_b601_follower.py b/src/lerobot/robots/rebot_b601_follower/rebot_b601_follower.py index ec00f4aa9..bf989702b 100644 --- a/src/lerobot/robots/rebot_b601_follower/rebot_b601_follower.py +++ b/src/lerobot/robots/rebot_b601_follower/rebot_b601_follower.py @@ -80,9 +80,14 @@ class RebotB601Follower(Robot): @property def _cameras_ft(self) -> dict[str, tuple]: - return { - cam: (self.config.cameras[cam].height, self.config.cameras[cam].width, 3) for cam in self.cameras - } + features: dict[str, tuple] = {} + for cam in self.cameras: + cfg = self.config.cameras[cam] + if getattr(cfg, "use_rgb", True): + features[cam] = (cfg.height, cfg.width, 3) + if getattr(cfg, "use_depth", False): + features[f"{cam}_depth"] = (cfg.height, cfg.width, 1) + return features @cached_property def observation_features(self) -> dict[str, type | tuple]: @@ -213,10 +218,17 @@ class RebotB601Follower(Robot): logger.debug(f"{self} read state: {dt_ms:.1f}ms") for cam_key, cam in self.cameras.items(): - start = time.perf_counter() - obs_dict[cam_key] = cam.read_latest() - dt_ms = (time.perf_counter() - start) * 1e3 - logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms") + if getattr(cam, "use_rgb", True): + start = time.perf_counter() + obs_dict[cam_key] = cam.read_latest() + dt_ms = (time.perf_counter() - start) * 1e3 + logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms") + + if getattr(cam, "use_depth", False): + start = time.perf_counter() + obs_dict[f"{cam_key}_depth"] = cam.read_latest_depth() + dt_ms = (time.perf_counter() - start) * 1e3 + logger.debug(f"{self} read {cam_key} depth: {dt_ms:.1f}ms") return obs_dict diff --git a/src/lerobot/robots/so_follower/so_follower.py b/src/lerobot/robots/so_follower/so_follower.py index 0651f566c..c6e67fafe 100644 --- a/src/lerobot/robots/so_follower/so_follower.py +++ b/src/lerobot/robots/so_follower/so_follower.py @@ -68,9 +68,13 @@ class SOFollower(Robot): @property def _cameras_ft(self) -> dict[str, tuple]: - return { - cam: (self.config.cameras[cam].height, self.config.cameras[cam].width, 3) for cam in self.cameras - } + features: dict[str, tuple] = {} + for cam in self.cameras: + if getattr(self.cameras[cam], "use_rgb", True): + features[cam] = (self.cameras[cam].height, self.cameras[cam].width, 3) + if getattr(self.cameras[cam], "use_depth", False): + features[f"{cam}_depth"] = (self.cameras[cam].height, self.cameras[cam].width, 1) + return features @cached_property def observation_features(self) -> dict[str, type | tuple]: @@ -185,10 +189,17 @@ class SOFollower(Robot): # Capture images from cameras for cam_key, cam in self.cameras.items(): - start = time.perf_counter() - obs_dict[cam_key] = cam.read_latest() - dt_ms = (time.perf_counter() - start) * 1e3 - logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms") + if getattr(cam, "use_rgb", True): + start = time.perf_counter() + obs_dict[cam_key] = cam.read_latest() + dt_ms = (time.perf_counter() - start) * 1e3 + logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms") + + if getattr(cam, "use_depth", False): + start = time.perf_counter() + obs_dict[f"{cam_key}_depth"] = cam.read_latest_depth() + dt_ms = (time.perf_counter() - start) * 1e3 + logger.debug(f"{self} read {cam_key} depth: {dt_ms:.1f}ms") return obs_dict diff --git a/src/lerobot/robots/unitree_g1/unitree_g1.py b/src/lerobot/robots/unitree_g1/unitree_g1.py index 25ec32716..5b8be0941 100644 --- a/src/lerobot/robots/unitree_g1/unitree_g1.py +++ b/src/lerobot/robots/unitree_g1/unitree_g1.py @@ -222,9 +222,14 @@ class UnitreeG1(Robot): @property def _cameras_ft(self) -> dict[str, tuple]: - return { - cam: (self.config.cameras[cam].height, self.config.cameras[cam].width, 3) for cam in self.cameras - } + features: dict[str, tuple] = {} + for cam in self.cameras: + cfg = self.config.cameras[cam] + if getattr(cfg, "use_rgb", True): + features[cam] = (cfg.height, cfg.width, 3) + if getattr(cfg, "use_depth", False): + features[f"{cam}_depth"] = (cfg.height, cfg.width, 1) + return features @cached_property def observation_features(self) -> dict[str, type | tuple]: @@ -458,7 +463,10 @@ class UnitreeG1(Robot): # Cameras - read images from ZMQ cameras for cam_name, cam in self._cameras.items(): - obs[cam_name] = cam.read_latest() + if getattr(cam, "use_rgb", True): + obs[cam_name] = cam.read_latest() + if getattr(cam, "use_depth", False): + obs[f"{cam_name}_depth"] = cam.read_latest_depth() return obs diff --git a/src/lerobot/rollout/context.py b/src/lerobot/rollout/context.py index bf5fa0fd4..62d844932 100644 --- a/src/lerobot/rollout/context.py +++ b/src/lerobot/rollout/context.py @@ -332,7 +332,8 @@ def build_rollout_context( cfg.dataset.repo_id, root=cfg.dataset.root, batch_encoding_size=cfg.dataset.video_encoding_batch_size, - camera_encoder=cfg.dataset.camera_encoder, + rgb_encoder=cfg.dataset.rgb_encoder, + depth_encoder=cfg.dataset.depth_encoder, streaming_encoding=cfg.dataset.streaming_encoding, encoder_queue_maxsize=cfg.dataset.encoder_queue_maxsize, encoder_threads=cfg.dataset.encoder_threads, @@ -367,7 +368,8 @@ def build_rollout_context( image_writer_threads=cfg.dataset.num_image_writer_threads_per_camera * len(robot.cameras if hasattr(robot, "cameras") else []), batch_encoding_size=cfg.dataset.video_encoding_batch_size, - camera_encoder=cfg.dataset.camera_encoder, + rgb_encoder=cfg.dataset.rgb_encoder, + depth_encoder=cfg.dataset.depth_encoder, streaming_encoding=cfg.dataset.streaming_encoding, encoder_queue_maxsize=cfg.dataset.encoder_queue_maxsize, encoder_threads=cfg.dataset.encoder_threads, diff --git a/src/lerobot/scripts/lerobot_dataset_viz.py b/src/lerobot/scripts/lerobot_dataset_viz.py index d07a2767d..21ae1ac9d 100644 --- a/src/lerobot/scripts/lerobot_dataset_viz.py +++ b/src/lerobot/scripts/lerobot_dataset_viz.py @@ -77,15 +77,28 @@ from lerobot.utils.constants import ACTION, DONE, OBS_STATE, REWARD from lerobot.utils.utils import init_logging +def check_chw_float32(frame: torch.Tensor) -> None: + """ + Check if a frame is a channel-first, float32 tensor. + """ + assert frame.dtype == torch.float32 + assert frame.ndim == 3 + c, h, w = frame.shape + assert c < h and c < w, f"expect channel first images, but instead {frame.shape}" + + def to_hwc_uint8_numpy(chw_float32_torch: torch.Tensor) -> np.ndarray: - assert chw_float32_torch.dtype == torch.float32 - assert chw_float32_torch.ndim == 3 - c, h, w = chw_float32_torch.shape - assert c < h and c < w, f"expect channel first images, but instead {chw_float32_torch.shape}" + check_chw_float32(chw_float32_torch) hwc_uint8_numpy = (chw_float32_torch * 255).type(torch.uint8).permute(1, 2, 0).numpy() return hwc_uint8_numpy +def to_hwc_uint16_numpy(chw_float32_torch: torch.Tensor) -> np.ndarray: + check_chw_float32(chw_float32_torch) + hwc_uint16_numpy = chw_float32_torch.round().type(torch.uint16).permute(1, 2, 0).numpy() + return hwc_uint16_numpy + + def visualize_dataset( dataset: LeRobotDataset, episode_index: int, @@ -138,6 +151,14 @@ def visualize_dataset( logging.info("Logging to Rerun") + # Use the dataset's q01/q99 depth statistics for robust depth range bounds + depth_ranges = {} + for key in dataset.meta.depth_keys: + stats = dataset.meta.stats[key] + lo = stats["q01"] if "q01" in stats else stats["min"] + hi = stats["q99"] if "q99" in stats else stats["max"] + depth_ranges[key] = (float(np.asarray(lo).item()), float(np.asarray(hi).item())) + first_index = None for batch in tqdm.tqdm(dataloader, total=len(dataloader)): if first_index is None: @@ -149,9 +170,18 @@ def visualize_dataset( # display each camera image for key in dataset.meta.camera_keys: - img = to_hwc_uint8_numpy(batch[key][i]) - img_entity = rr.Image(img).compress() if display_compressed_images else rr.Image(img) - rr.log(key, entity=img_entity) + if key in dataset.meta.depth_keys: + depth = to_hwc_uint16_numpy(batch[key][i]) + depth_entity = rr.DepthImage( + depth, + colormap=rr.components.Colormap.Viridis, + depth_range=depth_ranges[key], + ) + rr.log(key, entity=depth_entity) + else: + img = to_hwc_uint8_numpy(batch[key][i]) + img_entity = rr.Image(img).compress() if display_compressed_images else rr.Image(img) + rr.log(key, entity=img_entity) # display each dimension of action space (e.g. actuators command) if ACTION in batch: diff --git a/src/lerobot/scripts/lerobot_edit_dataset.py b/src/lerobot/scripts/lerobot_edit_dataset.py index eaadf47de..42dce438f 100644 --- a/src/lerobot/scripts/lerobot_edit_dataset.py +++ b/src/lerobot/scripts/lerobot_edit_dataset.py @@ -133,6 +133,15 @@ Convert image dataset to video format and save locally: --new_root /path/to/output/pusht_video \ --operation.type convert_image_to_video +Convert image dataset (with depth maps) to video format, customizing the depth encoder: + lerobot-edit-dataset \ + --repo_id lerobot/pusht_image \ + --new_root /path/to/output/pusht_video \ + --operation.type convert_image_to_video \ + --operation.depth_encoder.depth_min 0.01 \ + --operation.depth_encoder.depth_max 10.0 \ + --operation.depth_encoder.use_log true + Convert image dataset to video format and save with new repo_id: lerobot-edit-dataset \ --repo_id lerobot/pusht_image \ @@ -190,17 +199,17 @@ Re-encode all videos in a dataset (saves to lerobot/pusht_reencoded by default): lerobot-edit-dataset \ --repo_id lerobot/pusht \ --operation.type reencode_videos \ - --operation.camera_encoder.vcodec h264 \ - --operation.camera_encoder.pix_fmt yuv420p \ - --operation.camera_encoder.crf 23 + --operation.rgb_encoder.vcodec h264 \ + --operation.rgb_encoder.pix_fmt yuv420p \ + --operation.rgb_encoder.crf 23 Re-encode videos into a new dataset using 4 parallel processes: lerobot-edit-dataset \ --repo_id lerobot/pusht \ --new_repo_id lerobot/pusht_h264 \ --operation.type reencode_videos \ - --operation.camera_encoder.vcodec h264 \ - --operation.camera_encoder.crf 23 \ + --operation.rgb_encoder.vcodec h264 \ + --operation.rgb_encoder.crf 23 \ --operation.num_workers 4 Re-encode videos in-place (overwrites original dataset): @@ -208,9 +217,16 @@ Re-encode videos in-place (overwrites original dataset): --repo_id lerobot/pusht \ --new_repo_id lerobot/pusht \ --operation.type reencode_videos \ - --operation.camera_encoder.vcodec h264 \ + --operation.rgb_encoder.vcodec h264 \ --operation.overwrite true +Re-encode both RGB and depth videos in a dataset (depth quantization params are preserved): + lerobot-edit-dataset \ + --repo_id lerobot/pusht_depth \ + --operation.type reencode_videos \ + --operation.rgb_encoder.vcodec h264 \ + --operation.depth_encoder.extra_options '{"x265-params": "lossless=1"}' + Using JSON config file: lerobot-edit-dataset \ --config_path path/to/edit_config.json @@ -225,7 +241,13 @@ from pathlib import Path import draccus -from lerobot.configs import VideoEncoderConfig, camera_encoder_defaults, parser +from lerobot.configs import ( + DepthEncoderConfig, + RGBEncoderConfig, + depth_encoder_defaults, + parser, + rgb_encoder_defaults, +) from lerobot.datasets import ( LeRobotDataset, convert_image_to_video_dataset, @@ -287,7 +309,8 @@ class ModifyTasksConfig(OperationConfig): @dataclass class ConvertImageToVideoConfig(OperationConfig): output_dir: str | None = None - camera_encoder: VideoEncoderConfig = field(default_factory=camera_encoder_defaults) + rgb_encoder: RGBEncoderConfig = field(default_factory=rgb_encoder_defaults) + depth_encoder: DepthEncoderConfig = field(default_factory=depth_encoder_defaults) episode_indices: list[int] | None = None num_workers: int = 4 max_episodes_per_batch: int | None = None @@ -308,7 +331,8 @@ class RecomputeStatsConfig(OperationConfig): @OperationConfig.register_subclass("reencode_videos") @dataclass class ReencodeVideosConfig(OperationConfig): - camera_encoder: VideoEncoderConfig = field(default_factory=camera_encoder_defaults) + rgb_encoder: RGBEncoderConfig = field(default_factory=rgb_encoder_defaults) + depth_encoder: DepthEncoderConfig = field(default_factory=depth_encoder_defaults) num_workers: int = 0 encoder_threads: int | None = None overwrite: bool = False @@ -601,7 +625,8 @@ def handle_convert_image_to_video(cfg: EditDatasetConfig) -> None: dataset=dataset, output_dir=output_dir, repo_id=output_repo_id, - camera_encoder=getattr(cfg.operation, "camera_encoder", None) or camera_encoder_defaults(), + rgb_encoder=getattr(cfg.operation, "rgb_encoder", None) or rgb_encoder_defaults(), + depth_encoder=getattr(cfg.operation, "depth_encoder", None) or depth_encoder_defaults(), episode_indices=getattr(cfg.operation, "episode_indices", None), num_workers=getattr(cfg.operation, "num_workers", 4), max_episodes_per_batch=getattr(cfg.operation, "max_episodes_per_batch", None), @@ -719,10 +744,14 @@ def handle_reencode_videos(cfg: EditDatasetConfig) -> None: shutil.copytree(input_root, output_root) dataset = LeRobotDataset(output_repo_id, root=output_root) - logging.info(f"Re-encoding videos in {output_repo_id} with {cfg.operation.camera_encoder}") + logging.info( + f"Re-encoding videos in {output_repo_id} with RGB encoder {cfg.operation.rgb_encoder} " + f"and depth encoder {cfg.operation.depth_encoder}" + ) reencode_dataset( dataset, - camera_encoder=cfg.operation.camera_encoder, + rgb_encoder=cfg.operation.rgb_encoder, + depth_encoder=cfg.operation.depth_encoder, encoder_threads=cfg.operation.encoder_threads, num_workers=cfg.operation.num_workers, ) diff --git a/src/lerobot/scripts/lerobot_record.py b/src/lerobot/scripts/lerobot_record.py index 4d5518c7c..b759d86e0 100644 --- a/src/lerobot/scripts/lerobot_record.py +++ b/src/lerobot/scripts/lerobot_record.py @@ -79,9 +79,9 @@ lerobot-record \\ --dataset.single_task="Grab the cube" \\ --dataset.streaming_encoding=true \\ --dataset.encoder_threads=2 \\ - --dataset.camera_encoder.vcodec=h264 \\ - --dataset.camera_encoder.preset=fast \\ - --dataset.camera_encoder.extra_options={"tune": "film", "profile:v": "high", "bf": 2} \\ + --dataset.rgb_encoder.vcodec=h264 \\ + --dataset.rgb_encoder.preset=fast \\ + --dataset.rgb_encoder.extra_options={"tune": "film", "profile:v": "high", "bf": 2} \\ --display_data=true ``` """ @@ -400,7 +400,8 @@ def record( cfg.dataset.repo_id, root=cfg.dataset.root, batch_encoding_size=cfg.dataset.video_encoding_batch_size, - camera_encoder=cfg.dataset.camera_encoder, + rgb_encoder=cfg.dataset.rgb_encoder, + depth_encoder=cfg.dataset.depth_encoder, encoder_threads=cfg.dataset.encoder_threads, streaming_encoding=cfg.dataset.streaming_encoding, encoder_queue_maxsize=cfg.dataset.encoder_queue_maxsize, @@ -429,7 +430,8 @@ def record( image_writer_processes=cfg.dataset.num_image_writer_processes, image_writer_threads=cfg.dataset.num_image_writer_threads_per_camera * len(robot.cameras), batch_encoding_size=cfg.dataset.video_encoding_batch_size, - camera_encoder=cfg.dataset.camera_encoder, + rgb_encoder=cfg.dataset.rgb_encoder, + depth_encoder=cfg.dataset.depth_encoder, encoder_threads=cfg.dataset.encoder_threads, streaming_encoding=cfg.dataset.streaming_encoding, encoder_queue_maxsize=cfg.dataset.encoder_queue_maxsize, @@ -443,7 +445,7 @@ def record( if not cfg.dataset.streaming_encoding: logging.info( - "Streaming encoding is disabled. If you have capable hardware, consider enabling it for way faster episode saving. --dataset.streaming_encoding=true --dataset.encoder_threads=2 # --dataset.camera_encoder.vcodec=auto. More info in the documentation: https://huggingface.co/docs/lerobot/streaming_video_encoding" + "Streaming encoding is disabled. If you have capable hardware, consider enabling it for way faster episode saving. --dataset.streaming_encoding=true --dataset.encoder_threads=2 # --dataset.rgb_encoder.vcodec=auto. More info in the documentation: https://huggingface.co/docs/lerobot/streaming_video_encoding" ) with VideoEncodingManager(dataset): diff --git a/src/lerobot/scripts/lerobot_rollout.py b/src/lerobot/scripts/lerobot_rollout.py index 8515c4cc9..daee87bbe 100644 --- a/src/lerobot/scripts/lerobot_rollout.py +++ b/src/lerobot/scripts/lerobot_rollout.py @@ -142,9 +142,9 @@ Usage examples --robot.port=/dev/ttyACM0 \\ --task="pick up cube" --duration=60 \\ --display_data=true \\ - --dataset.camera_encoder.vcodec=h264 \\ - --dataset.camera_encoder.preset=fast \\ - --dataset.camera_encoder.extra_options={"tune": "film", "profile:v": "high", "bf": 2} + --dataset.rgb_encoder.vcodec=h264 \\ + --dataset.rgb_encoder.preset=fast \\ + --dataset.rgb_encoder.extra_options={"tune": "film", "profile:v": "high", "bf": 2} """ import logging diff --git a/src/lerobot/utils/feature_utils.py b/src/lerobot/utils/feature_utils.py index 2a4886234..38516d6ab 100644 --- a/src/lerobot/utils/feature_utils.py +++ b/src/lerobot/utils/feature_utils.py @@ -51,7 +51,9 @@ def hw_to_dataset_features( This function takes a dictionary describing hardware outputs (like joint states or camera image shapes) and formats it into the standard LeRobot feature - specification. + specification. Single-channel cameras (shape ``(H, W, 1)``) are flagged as depth + maps via ``info["is_depth_map"] = True``; three-channel cameras ``(H, W, 3)`` are + treated as RGB. Args: hw_features (dict): Dictionary mapping feature names to their type (float for @@ -61,7 +63,7 @@ def hw_to_dataset_features( use_video (bool): If True, image features are marked as "video", otherwise "image". Returns: - dict: A LeRobot features dictionary. + dict: A LeRobot features dictionary. Depth cameras carry ``info["is_depth_map"] = True``. """ features = {} joint_fts = { @@ -69,6 +71,7 @@ def hw_to_dataset_features( for key, ftype in hw_features.items() if ftype is float or (isinstance(ftype, PolicyFeature) and ftype.type != FeatureType.VISUAL) } + # TODO(CarolinePascal): we should not rely on the shape to determine if a feature is a camera ! cam_fts = {key: shape for key, shape in hw_features.items() if isinstance(shape, tuple)} if joint_fts and prefix == ACTION: @@ -86,11 +89,19 @@ def hw_to_dataset_features( } for key, shape in cam_fts.items(): - features[f"{prefix}.images.{key}"] = { - "dtype": "video" if use_video else "image", - "shape": shape, - "names": ["height", "width", "channels"], - } + dtype = "video" if use_video else "image" + if len(shape) == 3 and shape[2] in (1, 3): + features[f"{prefix}.images.{key}"] = { + "dtype": dtype, + "shape": shape, + "names": ["height", "width", "channels"], + "info": {"is_depth_map": shape[2] == 1}, + } + else: + raise ValueError( + f"Camera feature '{key}' has shape {shape}. " + f"Expected a 3-tuple (H, W, C), e.g. (480, 640, 3) for RGB or (480, 640, 1) for depth." + ) _validate_feature_names(features) return features @@ -149,11 +160,11 @@ def dataset_to_policy_features(features: dict[str, dict]) -> dict[str, PolicyFea type = FeatureType.VISUAL if len(shape) != 3: raise ValueError(f"Number of dimensions of {key} != 3 (shape={shape})") - - names = ft["names"] - # Backward compatibility for "channel" which is an error introduced in LeRobotDataset v2.0 for ported datasets. - if names[2] in ["channel", "channels"]: # (h, w, c) -> (c, h, w) - shape = (shape[2], shape[0], shape[1]) + else: + names = ft["names"] + # Backward compatibility for "channel" which is an error introduced in LeRobotDataset v2.0 for ported datasets. + if names[2] in ["channel", "channels"]: # (h, w, c) -> (c, h, w) + shape = (shape[2], shape[0], shape[1]) elif key == OBS_ENV_STATE: type = FeatureType.ENV elif key.startswith(OBS_STR): diff --git a/src/lerobot/utils/visualization_utils.py b/src/lerobot/utils/visualization_utils.py index d9d5bf6b5..e039f7b33 100644 --- a/src/lerobot/utils/visualization_utils.py +++ b/src/lerobot/utils/visualization_utils.py @@ -107,7 +107,10 @@ def log_rerun_data( for i, vi in enumerate(arr): rr.log(f"{key}_{i}", rr.Scalars(float(vi))) else: - img_entity = rr.Image(arr).compress() if compress_images else rr.Image(arr) + if arr.shape[-1] == 1: + img_entity = rr.DepthImage(arr, colormap=rr.components.Colormap.Viridis) + else: + img_entity = rr.Image(arr).compress() if compress_images else rr.Image(arr) rr.log(key, entity=img_entity, static=True) if action: diff --git a/tests/annotations/test_frames.py b/tests/annotations/test_frames.py index 5c9c58f7b..1a626533f 100644 --- a/tests/annotations/test_frames.py +++ b/tests/annotations/test_frames.py @@ -47,6 +47,7 @@ class _FakeMeta: def __init__(self, video_keys: list[str], image_keys: list[str], video_path: Path | None = None) -> None: self.video_keys = video_keys self.camera_keys = [*video_keys, *image_keys] + self.depth_keys = [] self._video_path = video_path self.episodes = {0: {f"videos/{key}/from_timestamp": 0.0 for key in video_keys}} @@ -208,14 +209,14 @@ def test_episode_clip_path_trims_via_reencode_video(tmp_path: Path, monkeypatch) def fake_reencode( input_video_path, output_video_path, - camera_encoder=None, + video_encoder=None, overwrite=False, start_time_s=None, end_time_s=None, ): captured.update( src=Path(input_video_path), - encoder=camera_encoder, + encoder=video_encoder, start_time_s=start_time_s, end_time_s=end_time_s, ) diff --git a/tests/datasets/test_aggregate.py b/tests/datasets/test_aggregate.py index e9930575f..2fafd2777 100644 --- a/tests/datasets/test_aggregate.py +++ b/tests/datasets/test_aggregate.py @@ -29,7 +29,10 @@ from lerobot.configs import VIDEO_ENCODER_INFO_KEYS from lerobot.datasets.aggregate import aggregate_datasets from lerobot.datasets.feature_utils import features_equal_for_merge from lerobot.datasets.lerobot_dataset import LeRobotDataset -from tests.fixtures.constants import DUMMY_REPO_ID +from tests.fixtures.constants import ( + DUMMY_CAMERA_FEATURES_WITH_DEPTH, + DUMMY_REPO_ID, +) def assert_data_shards_one_row_group_per_episode(root): @@ -211,6 +214,26 @@ def assert_dataset_iteration_works(aggr_ds): pass +def assert_depth_keys_preserved(aggr_ds, ds_0, ds_1): + """Test that depth keys are correctly preserved after aggregation. + + Ensures that the ``is_depth_map`` marker on visual features survives + aggregation, so that downstream consumers (e.g. the dataset reader's + depth decoding path) keep working on the merged dataset. + """ + expected_depth_keys = set(ds_0.meta.depth_keys) + assert expected_depth_keys == set(ds_1.meta.depth_keys), ( + "Source datasets disagree on depth_keys; test setup is inconsistent" + ) + actual_depth_keys = set(aggr_ds.meta.depth_keys) + assert actual_depth_keys == expected_depth_keys, ( + f"Expected depth_keys {expected_depth_keys}, got {actual_depth_keys}" + ) + for key in expected_depth_keys: + info = aggr_ds.meta.info.features[key].get("info") or {} + assert info.get("is_depth_map") is True, f"Depth marker lost on feature {key!r} after aggregation" + + def assert_video_timestamps_within_bounds(aggr_ds): """Test that all video timestamps are within valid bounds for their respective video files. @@ -260,7 +283,11 @@ def assert_video_timestamps_within_bounds(aggr_ds): def test_aggregate_datasets(tmp_path, lerobot_dataset_factory): - """Test basic aggregation functionality with standard parameters.""" + """Test basic aggregation functionality with standard parameters. + + Source datasets include both RGB and depth video features so the same + aggregation flow is exercised on the ``is_depth_map`` branch. + """ ds_0_num_frames = 400 ds_1_num_frames = 800 ds_0_num_episodes = 10 @@ -272,14 +299,21 @@ def test_aggregate_datasets(tmp_path, lerobot_dataset_factory): repo_id=f"{DUMMY_REPO_ID}_0", total_episodes=ds_0_num_episodes, total_frames=ds_0_num_frames, + camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH, ) ds_1 = lerobot_dataset_factory( root=tmp_path / "test_1", repo_id=f"{DUMMY_REPO_ID}_1", total_episodes=ds_1_num_episodes, total_frames=ds_1_num_frames, + camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH, ) + # Confirm depth was actually wired into the source datasets so the + # rest of the assertions exercise the depth aggregation path. + assert len(ds_0.meta.depth_keys) > 0, "ds_0 should expose at least one depth key" + assert len(ds_1.meta.depth_keys) > 0, "ds_1 should expose at least one depth key" + aggregate_datasets( repo_ids=[ds_0.repo_id, ds_1.repo_id], roots=[ds_0.root, ds_1.root], @@ -306,6 +340,7 @@ def test_aggregate_datasets(tmp_path, lerobot_dataset_factory): assert_episode_indices_updated_correctly(aggr_ds, ds_0, ds_1) assert_video_frames_integrity(aggr_ds, ds_0, ds_1) assert_video_timestamps_within_bounds(aggr_ds) + assert_depth_keys_preserved(aggr_ds, ds_0, ds_1) assert_dataset_iteration_works(aggr_ds) @@ -423,7 +458,11 @@ def test_aggregate_incomplete_video_encoder_info_warns_and_nuls_encoders( def test_aggregate_with_low_threshold(tmp_path, lerobot_dataset_factory): - """Test aggregation with small file size limits to force file rotation/sharding.""" + """Test aggregation with small file size limits to force file rotation/sharding. + + Depth video features are included to verify that file rotation/concat + correctly handles depth-marked features alongside regular RGB ones. + """ ds_0_num_episodes = ds_1_num_episodes = 10 ds_0_num_frames = ds_1_num_frames = 400 @@ -432,14 +471,19 @@ def test_aggregate_with_low_threshold(tmp_path, lerobot_dataset_factory): repo_id=f"{DUMMY_REPO_ID}_small_0", total_episodes=ds_0_num_episodes, total_frames=ds_0_num_frames, + camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH, ) ds_1 = lerobot_dataset_factory( root=tmp_path / "small_1", repo_id=f"{DUMMY_REPO_ID}_small_1", total_episodes=ds_1_num_episodes, total_frames=ds_1_num_frames, + camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH, ) + assert len(ds_0.meta.depth_keys) > 0, "ds_0 should expose at least one depth key" + assert len(ds_1.meta.depth_keys) > 0, "ds_1 should expose at least one depth key" + # Use the new configurable parameters to force file rotation aggregate_datasets( repo_ids=[ds_0.repo_id, ds_1.repo_id], @@ -470,6 +514,7 @@ def test_aggregate_with_low_threshold(tmp_path, lerobot_dataset_factory): assert_episode_indices_updated_correctly(aggr_ds, ds_0, ds_1) assert_video_frames_integrity(aggr_ds, ds_0, ds_1) assert_video_timestamps_within_bounds(aggr_ds) + assert_depth_keys_preserved(aggr_ds, ds_0, ds_1) assert_dataset_iteration_works(aggr_ds) # Check that multiple files were actually created due to small size limits @@ -489,7 +534,8 @@ def test_video_timestamps_regression(tmp_path, lerobot_dataset_factory): """Regression test for video timestamp bug when merging datasets. This test specifically checks that video timestamps are correctly calculated - and accumulated when merging multiple datasets. + and accumulated when merging multiple datasets. Depth video features are + included so depth timestamps are also covered by the regression. """ datasets = [] for i in range(3): @@ -498,9 +544,13 @@ def test_video_timestamps_regression(tmp_path, lerobot_dataset_factory): repo_id=f"{DUMMY_REPO_ID}_regression_{i}", total_episodes=2, total_frames=100, + camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH, ) datasets.append(ds) + for i, ds in enumerate(datasets): + assert len(ds.meta.depth_keys) > 0, f"Dataset {i} should expose at least one depth key" + aggregate_datasets( repo_ids=[ds.repo_id for ds in datasets], roots=[ds.root for ds in datasets], @@ -517,12 +567,21 @@ def test_video_timestamps_regression(tmp_path, lerobot_dataset_factory): aggr_ds = LeRobotDataset(f"{DUMMY_REPO_ID}_regression_aggr", root=tmp_path / "regression_aggr") assert_video_timestamps_within_bounds(aggr_ds) + # Depth keys must survive the merge for the regression to cover the + # ``is_depth_map`` decoding branch. + assert set(aggr_ds.meta.depth_keys) == set(datasets[0].meta.depth_keys) + depth_keys = set(aggr_ds.meta.depth_keys) for i in range(len(aggr_ds)): item = aggr_ds[i] for key in aggr_ds.meta.video_keys: assert key in item, f"Video key {key} missing from item {i}" - assert item[key].shape[0] == 3, f"Expected 3 channels for video key {key}" + # Depth frames are single-channel (1, H, W) after dequantization; + # standard RGB frames keep the 3-channel layout. + expected_channels = 1 if key in depth_keys else 3 + assert item[key].shape[0] == expected_channels, ( + f"Expected {expected_channels} channels for video key {key}, got {item[key].shape}" + ) def assert_image_schema_preserved(aggr_ds): @@ -639,25 +698,31 @@ def test_aggregate_image_datasets(tmp_path, lerobot_dataset_factory): ds_0_num_episodes = 2 ds_1_num_episodes = 3 - # Create two image-based datasets (use_videos=False) + # Create two image-based datasets (use_videos=False) with a mix of RGB + # and depth-marked cameras so the depth path is exercised in image mode. ds_0 = lerobot_dataset_factory( root=tmp_path / "image_0", repo_id=f"{DUMMY_REPO_ID}_image_0", total_episodes=ds_0_num_episodes, total_frames=ds_0_num_frames, - use_videos=False, # Image-based dataset + use_videos=False, + camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH, ) ds_1 = lerobot_dataset_factory( root=tmp_path / "image_1", repo_id=f"{DUMMY_REPO_ID}_image_1", total_episodes=ds_1_num_episodes, total_frames=ds_1_num_frames, - use_videos=False, # Image-based dataset + use_videos=False, + camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH, ) # Verify source datasets have image keys assert len(ds_0.meta.image_keys) > 0, "ds_0 should have image keys" assert len(ds_1.meta.image_keys) > 0, "ds_1 should have image keys" + # And that the depth marker actually made it onto an image feature. + assert len(ds_0.meta.depth_keys) > 0, "ds_0 should expose at least one depth key" + assert len(ds_1.meta.depth_keys) > 0, "ds_1 should expose at least one depth key" # Aggregate the datasets aggregate_datasets( @@ -692,6 +757,7 @@ def test_aggregate_image_datasets(tmp_path, lerobot_dataset_factory): # Image-specific assertions assert_image_schema_preserved(aggr_ds) assert_image_frames_integrity(aggr_ds, ds_0, ds_1) + assert_depth_keys_preserved(aggr_ds, ds_0, ds_1) # Verify images can be accessed and have correct shape sample_item = aggr_ds[0] diff --git a/tests/datasets/test_compute_stats.py b/tests/datasets/test_compute_stats.py index 0f5abfb95..9f399b85c 100644 --- a/tests/datasets/test_compute_stats.py +++ b/tests/datasets/test_compute_stats.py @@ -35,7 +35,11 @@ from lerobot.utils.constants import OBS_IMAGE, OBS_STATE def mock_load_image_as_numpy(path, dtype, channel_first): - return np.ones((3, 32, 32), dtype=dtype) if channel_first else np.ones((32, 32, 3), dtype=dtype) + is_depth = "depth" in str(path) + channels = 1 if is_depth else 3 + out_dtype = np.uint16 if is_depth else dtype + arr = np.arange(channels * 32 * 32, dtype=out_dtype).reshape(channels, 32, 32) + return arr if channel_first else arr.transpose(1, 2, 0) @pytest.fixture @@ -168,22 +172,33 @@ def test_get_feature_stats_single_value(): def test_compute_episode_stats(): + depth_key = "observation.images.depth" episode_data = { OBS_IMAGE: [f"image_{i}.jpg" for i in range(100)], + depth_key: [f"depth_{i}.tiff" for i in range(100)], OBS_STATE: np.random.rand(100, 10), } features = { OBS_IMAGE: {"dtype": "image"}, + depth_key: {"dtype": "image", "info": {"is_depth_map": True}}, OBS_STATE: {"dtype": "numeric"}, } with patch("lerobot.datasets.compute_stats.load_image_as_numpy", side_effect=mock_load_image_as_numpy): stats = compute_episode_stats(episode_data, features) - assert OBS_IMAGE in stats and OBS_STATE in stats + assert OBS_IMAGE in stats and depth_key in stats and OBS_STATE in stats assert stats[OBS_IMAGE]["count"].item() == 100 + assert stats[depth_key]["count"].item() == 100 assert stats[OBS_STATE]["count"].item() == 100 assert stats[OBS_IMAGE]["mean"].shape == (3, 1, 1) + assert stats[depth_key]["mean"].shape == (1, 1, 1) + # Depth keeps raw values: max far exceeds 255, proving no /255 and no uint8 downcast. + assert stats[depth_key]["min"].item() == 0.0 + assert stats[depth_key]["max"].item() == 1023.0 + # RGB is normalized to [0, 1]. + np.testing.assert_allclose(stats[OBS_IMAGE]["min"], 0.0) + np.testing.assert_allclose(stats[OBS_IMAGE]["max"], 1.0) def test_assert_type_and_shape_valid(): @@ -618,25 +633,31 @@ def test_compute_episode_stats_with_custom_quantiles(): def test_compute_episode_stats_with_image_data(): """Test quantile computation with image features.""" image_paths = [f"image_{i}.jpg" for i in range(50)] + depth_paths = [f"depth_{i}.tiff" for i in range(50)] episode_data = { "observation.image": image_paths, + "observation.images.depth": depth_paths, "action": np.random.normal(0, 1, (50, 5)), } features = { "observation.image": {"dtype": "image"}, + "observation.images.depth": {"dtype": "image", "info": {"is_depth_map": True}}, "action": {"dtype": "float32", "shape": (5,)}, } with patch("lerobot.datasets.compute_stats.load_image_as_numpy", side_effect=mock_load_image_as_numpy): stats = compute_episode_stats(episode_data, features) - # Image quantiles should be normalized and have correct shape - assert "q01" in stats["observation.image"] - assert "q50" in stats["observation.image"] - assert "q99" in stats["observation.image"] - assert stats["observation.image"]["q01"].shape == (3, 1, 1) - assert stats["observation.image"]["q50"].shape == (3, 1, 1) - assert stats["observation.image"]["q99"].shape == (3, 1, 1) + # RGB image quantiles should be normalized and per-channel. + for q in ("q01", "q50", "q99"): + assert stats["observation.image"][q].shape == (3, 1, 1) + + # Depth quantiles are single-channel and kept in raw (un-normalized) units. + for q in ("q01", "q50", "q99"): + assert stats["observation.images.depth"][q].shape == (1, 1, 1) + # Depth max stays in raw units (not /255, not uint8-capped); RGB is normalized. + assert stats["observation.images.depth"]["max"].item() == 1023.0 + np.testing.assert_allclose(stats["observation.image"]["max"], 1.0) # Action quantiles should have correct shape assert stats["action"]["q01"].shape == (5,) diff --git a/tests/datasets/test_dataset_metadata.py b/tests/datasets/test_dataset_metadata.py index 171d8af8b..a1630f17d 100644 --- a/tests/datasets/test_dataset_metadata.py +++ b/tests/datasets/test_dataset_metadata.py @@ -59,11 +59,13 @@ def _make_dummy_stats(features: dict) -> dict: stats = {} for key, ft in features.items(): if ft["dtype"] in ("image", "video"): + channels = ft["shape"][-1] + stat_shape = (channels, 1, 1) stats[key] = { - "max": np.ones((3, 1, 1), dtype=np.float32), - "mean": np.full((3, 1, 1), 0.5, dtype=np.float32), - "min": np.zeros((3, 1, 1), dtype=np.float32), - "std": np.full((3, 1, 1), 0.25, dtype=np.float32), + "max": np.ones(stat_shape, dtype=np.float32), + "mean": np.full(stat_shape, 0.5, dtype=np.float32), + "min": np.zeros(stat_shape, dtype=np.float32), + "std": np.full(stat_shape, 0.25, dtype=np.float32), "count": np.array([5]), } elif ft["dtype"] in ("float32", "float64", "int64"): @@ -142,6 +144,45 @@ def test_create_without_videos_has_no_video_path(tmp_path): assert meta.video_keys == [] +@pytest.mark.parametrize( + ("marker_field", "marker_key"), + [ + ("info", "is_depth_map"), + ("info", "video.is_depth_map"), + ("video_info", "video.is_depth_map"), + ], + ids=["info.is_depth_map", "info.video.is_depth_map_legacy", "video_info.video.is_depth_map_legacy"], +) +def test_depth_keys_property_filters_by_marker(tmp_path, marker_field, marker_key): + """``depth_keys`` recognises the canonical and the two legacy marker variants.""" + depth_feature = { + "dtype": "video", + "shape": (64, 96, 1), + "names": ["height", "width", "channels"], + marker_field: {marker_key: True}, + } + features = { + **VIDEO_FEATURES, + "observation.images.laptop_depth": depth_feature, + } + meta = LeRobotDatasetMetadata.create( + repo_id="test/depth_keys", + fps=DEFAULT_FPS, + features=features, + root=tmp_path / f"depth_keys_{marker_field}_{marker_key.replace('.', '_')}", + ) + + assert set(meta.video_keys) == {"observation.images.laptop", "observation.images.laptop_depth"} + assert meta.depth_keys == ["observation.images.laptop_depth"] + + +def test_depth_keys_empty_when_no_marker(tmp_path): + meta = LeRobotDatasetMetadata.create( + repo_id="test/no_depth", fps=DEFAULT_FPS, features=VIDEO_FEATURES, root=tmp_path / "no_depth" + ) + assert meta.depth_keys == [] + + def test_create_raises_on_existing_directory(tmp_path): """create() raises if root directory already exists.""" root = tmp_path / "existing" diff --git a/tests/datasets/test_dataset_tools.py b/tests/datasets/test_dataset_tools.py index d36312920..c19e7f41f 100644 --- a/tests/datasets/test_dataset_tools.py +++ b/tests/datasets/test_dataset_tools.py @@ -24,7 +24,7 @@ import torch pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])") -from lerobot.configs import VideoEncoderConfig +from lerobot.configs import DepthEncoderConfig, RGBEncoderConfig from lerobot.datasets.dataset_tools import ( add_features, convert_image_to_video_dataset, @@ -37,7 +37,9 @@ from lerobot.datasets.dataset_tools import ( split_dataset, ) from lerobot.datasets.io_utils import load_info -from tests.datasets.test_video_encoding import _add_frames, require_h264, require_libsvtav1 +from tests.datasets.test_video_encoding import require_h264, require_hevc, require_libsvtav1 +from tests.fixtures.constants import DUMMY_DEPTH_FEATURES, DUMMY_DEPTH_KEY +from tests.fixtures.dataset_factories import add_frames @pytest.fixture @@ -1251,7 +1253,7 @@ def test_convert_image_to_video_dataset(tmp_path): dataset=source_dataset, output_dir=output_dir, repo_id="lerobot/pusht_video", - camera_encoder=VideoEncoderConfig( + rgb_encoder=RGBEncoderConfig( vcodec="libsvtav1", pix_fmt="yuv420p", g=2, @@ -1332,9 +1334,131 @@ def test_convert_image_to_video_dataset_subset_episodes(tmp_path): shutil.rmtree(output_dir) +@require_libsvtav1 +@require_hevc +def test_convert_image_to_video_dataset_depth(tmp_path, empty_lerobot_dataset_factory): + """Depth image features convert to depth videos using the depth encoder. + + Mirrors :func:`test_convert_image_to_video_dataset` but with a small local + image dataset that mixes an RGB camera with a depth camera, so the + ``depth_keys`` → ``depth_encoder`` routing and ``is_depth_map`` preservation + are exercised end-to-end. + """ + features = { + "action": {"dtype": "float32", "shape": (2,), "names": ["a", "b"]}, + "observation.images.cam": { + "dtype": "image", + "shape": (64, 96, 3), + "names": ["height", "width", "channels"], + }, + "observation.images.depth": { + "dtype": "image", + "shape": (64, 96, 1), + "names": ["height", "width", "channels"], + "info": {"is_depth_map": True}, + }, + } + source_dataset = empty_lerobot_dataset_factory( + root=tmp_path / "img_ds", + features=features, + use_videos=False, + ) + + add_frames(source_dataset, num_frames=4) + source_dataset.save_episode() + source_dataset.finalize() + + # Source is an image dataset with the depth marker on the depth camera. + assert len(source_dataset.meta.video_keys) == 0 + assert "observation.images.depth" in source_dataset.meta.depth_keys + + output_dir = tmp_path / "video_ds" + with ( + patch("lerobot.datasets.dataset_metadata.get_safe_version") as mock_get_safe_version, + patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download, + ): + mock_get_safe_version.return_value = "v3.0" + mock_snapshot_download.return_value = str(output_dir) + + # Use non-default quantization params so the persisted metadata must + # come from the depth encoder (not RGB encoder defaults). + depth_encoder = DepthEncoderConfig( + vcodec="hevc", + pix_fmt="gray12le", + g=2, + crf=30, + depth_min=0.05, + depth_max=8.0, + shift=2.0, + use_log=False, + ) + video_dataset = convert_image_to_video_dataset( + dataset=source_dataset, + output_dir=output_dir, + repo_id="dummy/depth_video", + rgb_encoder=RGBEncoderConfig(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30), + depth_encoder=depth_encoder, + num_workers=1, + ) + + # Both cameras are now videos, and the depth marker survived the conversion. + assert "observation.images.cam" in video_dataset.meta.video_keys + assert "observation.images.depth" in video_dataset.meta.video_keys + assert "observation.images.depth" in video_dataset.meta.depth_keys + assert "observation.images.cam" not in video_dataset.meta.depth_keys + + depth_path = video_dataset.root / video_dataset.meta.get_video_file_path(0, "observation.images.depth") + assert depth_path.exists(), f"Depth video file should exist: {depth_path}" + + # The persisted depth-video metadata must carry the depth quantization params + # from the depth encoder (so frames dequantize correctly on read), and the RGB + # camera must not be marked as a depth map. + persisted_info = load_info(video_dataset.root) + depth_info = persisted_info.features["observation.images.depth"]["info"] + assert depth_info["is_depth_map"] is True + assert DepthEncoderConfig.from_video_info(depth_info) == depth_encoder + + cam_info = persisted_info.features["observation.images.cam"]["info"] + assert cam_info.get("is_depth_map") is False + assert "video.codec" in cam_info + + # ─── reencode_dataset ───────────────────────────────────────────────── +@require_hevc +def test_reencode_dataset_depth_uses_depth_encoder(tmp_path, empty_lerobot_dataset_factory): + """Depth videos are re-encoded with the depth encoder and keep their depth metadata. + + Depth-focused companion to :func:`test_reencode_dataset_multi_key_multiprocessing`. + """ + initial_cfg = DepthEncoderConfig(vcodec="hevc", pix_fmt="gray12le", g=2, crf=30) + dataset = empty_lerobot_dataset_factory( + root=tmp_path / "ds", + features=DUMMY_DEPTH_FEATURES, + use_videos=True, + depth_encoder=initial_cfg, + ) + + add_frames(dataset, num_frames=4) + dataset.save_episode() + dataset.finalize() + + assert DUMMY_DEPTH_KEY in dataset.meta.depth_keys + + target_cfg = DepthEncoderConfig(vcodec="hevc", pix_fmt="gray12le", g=6, crf=23) + result = reencode_dataset(dataset, depth_encoder=target_cfg, num_workers=0) + + assert result is dataset + + persisted_info = load_info(dataset.root) + depth_info = persisted_info.features[DUMMY_DEPTH_KEY].get("info", {}) + # Re-encode applied the new codec parameters to the depth video ... + assert DepthEncoderConfig.from_video_info(depth_info) == target_cfg + # ... while preserving the depth marker. + assert depth_info["is_depth_map"] is True + + @require_libsvtav1 @require_h264 def test_reencode_dataset_multi_key_multiprocessing( @@ -1342,29 +1466,29 @@ def test_reencode_dataset_multi_key_multiprocessing( ): """Re-encode a two-camera dataset with num_workers=2 and verify metadata refresh.""" features = features_factory(use_videos=True) - initial_cfg = VideoEncoderConfig(vcodec="libsvtav1", g=2, crf=30, preset=12) + initial_cfg = RGBEncoderConfig(vcodec="libsvtav1", g=2, crf=30, preset=12) dataset = empty_lerobot_dataset_factory( root=tmp_path / "ds", features=features, use_videos=True, - camera_encoder=initial_cfg, + rgb_encoder=initial_cfg, ) - _add_frames(dataset, num_frames=4) + add_frames(dataset, num_frames=4) dataset.save_episode() - _add_frames(dataset, num_frames=4) + add_frames(dataset, num_frames=4) dataset.save_episode() dataset.finalize() assert len(dataset.meta.video_keys) == 2 - target_cfg = VideoEncoderConfig(vcodec="h264", g=6, crf=23, pix_fmt="yuv420p") + target_cfg = RGBEncoderConfig(vcodec="h264", g=6, crf=23, pix_fmt="yuv420p") - result = reencode_dataset(dataset, camera_encoder=target_cfg, num_workers=2) + result = reencode_dataset(dataset, rgb_encoder=target_cfg, num_workers=2) assert result is dataset persisted_info = load_info(dataset.root) for vk in dataset.meta.video_keys: - persisted_encoder = VideoEncoderConfig.from_video_info(persisted_info.features[vk].get("info", {})) + persisted_encoder = RGBEncoderConfig.from_video_info(persisted_info.features[vk].get("info", {})) assert persisted_encoder == target_cfg diff --git a/tests/datasets/test_dataset_writer.py b/tests/datasets/test_dataset_writer.py index 8670aeebc..17785ad74 100644 --- a/tests/datasets/test_dataset_writer.py +++ b/tests/datasets/test_dataset_writer.py @@ -53,8 +53,8 @@ def _make_frame(features: dict, task: str = "Dummy task") -> dict: # ── Existing encode_video_worker tests ─────────────────────────────── -def test_encode_video_worker_forwards_camera_encoder(tmp_path): - """_encode_video_worker forwards camera_encoder to encode_video_frames.""" +def test_encode_video_worker_forwards_video_encoder(tmp_path): + """_encode_video_worker forwards video_encoder to encode_video_frames.""" video_key = "observation.images.laptop" fpath = DEFAULT_IMAGE_PATH.format(image_key=video_key, episode_index=0, frame_index=0) img_dir = tmp_path / Path(fpath).parent @@ -74,16 +74,16 @@ def test_encode_video_worker_forwards_camera_encoder(tmp_path): 0, tmp_path, fps=30, - camera_encoder=VideoEncoderConfig(vcodec="h264", preset=None), + video_encoder=VideoEncoderConfig(vcodec="h264", preset=None), encoder_threads=4, ) - assert captured_kwargs["camera_encoder"].vcodec == "h264" + assert captured_kwargs["video_encoder"].vcodec == "h264" assert captured_kwargs["encoder_threads"] == 4 -def test_encode_video_worker_default_camera_encoder(tmp_path): - """_encode_video_worker passes None camera_encoder which encode_video_frames defaults.""" +def test_encode_video_worker_default_video_encoder(tmp_path): + """_encode_video_worker passes None video_encoder which encode_video_frames defaults.""" video_key = "observation.images.laptop" fpath = DEFAULT_IMAGE_PATH.format(image_key=video_key, episode_index=0, frame_index=0) img_dir = tmp_path / Path(fpath).parent @@ -100,7 +100,7 @@ def test_encode_video_worker_default_camera_encoder(tmp_path): with patch("lerobot.datasets.dataset_writer.encode_video_frames", side_effect=mock_encode): _encode_video_worker(video_key, 0, tmp_path, fps=30) - assert captured_kwargs["camera_encoder"] is None + assert captured_kwargs["video_encoder"] is None assert captured_kwargs["encoder_threads"] is None diff --git a/tests/datasets/test_datasets.py b/tests/datasets/test_datasets.py index 1d2fb1d55..225479814 100644 --- a/tests/datasets/test_datasets.py +++ b/tests/datasets/test_datasets.py @@ -1534,6 +1534,10 @@ def test_valid_video_codecs_constant(): assert "auto" in VALID_VIDEO_CODECS assert "h264_videotoolbox" in VALID_VIDEO_CODECS assert "h264_nvenc" in VALID_VIDEO_CODECS + assert "h264_vaapi" in VALID_VIDEO_CODECS + assert "h264_qsv" in VALID_VIDEO_CODECS + assert "hevc_videotoolbox" in VALID_VIDEO_CODECS + assert "hevc_nvenc" in VALID_VIDEO_CODECS assert len(VALID_VIDEO_CODECS) == 10 diff --git a/tests/datasets/test_depth.py b/tests/datasets/test_depth.py new file mode 100644 index 000000000..a075fa6b5 --- /dev/null +++ b/tests/datasets/test_depth.py @@ -0,0 +1,247 @@ +"""Tests for the depth-integration feature. + +Covers: +- ``depth_utils`` quantize/dequantize round-trips and backend agreement. +- Image-writer support for single-channel depth. +- Hardware-feature → depth flag routing. +- Feature-to-file-format routing through the dataset writer. + +Depth metadata detection on ``LeRobotDatasetMetadata.depth_keys`` lives in +``test_dataset_metadata.py``. Depth video encoding/decoding lives in +``test_video_encoding.py``. +""" + +from pathlib import Path + +import pytest + +pytest.importorskip("av", reason="av is required (install lerobot[dataset])") + +import av +import numpy as np +import PIL.Image +import torch + +from lerobot.configs import DepthEncoderConfig +from lerobot.configs.video import ( + DEFAULT_DEPTH_MAX, + DEFAULT_DEPTH_MIN, + DEPTH_METER_UNIT, + DEPTH_MILLIMETER_UNIT, + DEPTH_QMAX, +) +from lerobot.datasets.depth_utils import dequantize_depth, quantize_depth +from lerobot.datasets.image_writer import image_array_to_pil_image, write_image +from tests.fixtures.constants import ( + DEFAULT_FPS, + DUMMY_CAMERA_FEATURES, + DUMMY_CAMERA_FEATURES_WITH_DEPTH, + DUMMY_CHW, + DUMMY_DEPTH_CAMERA_FEATURES, + DUMMY_REPO_ID, +) +from tests.fixtures.dataset_factories import add_frames + +_, H, W = DUMMY_CHW + + +def _depth_metres_ramp() -> np.ndarray: + """Linearly-spaced float32 depth in metres covering the default range.""" + return np.linspace(DEFAULT_DEPTH_MIN, DEFAULT_DEPTH_MAX, H * W, dtype=np.float32).reshape(H, W) + + +# ── 1. Quantize / dequantize round-trips ────────────────────────────── + + +class TestQuantizeDequantize: + """Numerical contract of ``quantize_depth`` / ``dequantize_depth``.""" + + @pytest.mark.parametrize("use_log", [False, True]) + @pytest.mark.parametrize("output_unit", [DEPTH_METER_UNIT, DEPTH_MILLIMETER_UNIT]) + @pytest.mark.parametrize("output_channel_last", [False, True]) + def test_roundtrip(self, use_log, output_unit, output_channel_last): + """quantize → dequantize recovers depth; layout and unit are honored.""" + depth = _depth_metres_ramp() + quantized = quantize_depth(depth, use_log=use_log, video_backend=None) + recovered = dequantize_depth( + quantized, + use_log=use_log, + output_unit=output_unit, + output_tensor=False, + output_channel_last=output_channel_last, + ) + + expected_shape = (H, W, 1) if output_channel_last else (1, H, W) + assert recovered.shape == expected_shape + + recovered_m = recovered.astype(np.float32) + if output_unit == DEPTH_MILLIMETER_UNIT: + recovered_m = recovered_m / 1000.0 + recovered_2d = recovered_m[..., 0] if output_channel_last else recovered_m[0] + + if use_log: + # Log mode: tighter near-range error than far-range (the whole point). + near = depth < 1.0 + far = depth > 8.0 + err_near = np.abs(recovered_2d[near] - depth[near]) + err_far = np.abs(recovered_2d[far] - depth[far]) + assert err_near.mean() < err_far.mean() + else: + # Linear mode: bounded by quant step + 1 mm of unit-conversion rounding. + tol = (DEFAULT_DEPTH_MAX - DEFAULT_DEPTH_MIN) / DEPTH_QMAX + 1e-3 + np.testing.assert_allclose(recovered_2d, depth, atol=tol) + + @pytest.mark.parametrize("use_log", [False, True]) + @pytest.mark.parametrize("output_unit", [DEPTH_METER_UNIT, DEPTH_MILLIMETER_UNIT]) + def test_numpy_torch_agree(self, use_log, output_unit): + """Batched torch path produces the same values as the numpy path.""" + batch_size = 3 + per_frame = np.linspace(0, DEPTH_QMAX, H * W, dtype=np.uint16).reshape(H, W) + batch_np = np.broadcast_to(per_frame[None, None, ...], (batch_size, 1, H, W)).copy() + batch_t = torch.from_numpy(batch_np.astype(np.int32)) # torch.uint16 support is patchy. + + ref = dequantize_depth(batch_np, use_log=use_log, output_unit=output_unit, output_tensor=False) + out = dequantize_depth(batch_t, use_log=use_log, output_unit=output_unit, output_tensor=True) + + assert isinstance(out, torch.Tensor) + assert out.shape == (batch_size, 1, H, W) + # ``m``: float32 noise (~10 µm in log mode, after ``exp``) — still 200× below the ~2 mm quant step. + # ``mm`` + tensor stays in float32 (no uint16 round-trip), so allow 1 mm slop. + atol = 1e-5 if output_unit == DEPTH_METER_UNIT else 1.0 + np.testing.assert_allclose(out.cpu().numpy().astype(np.float64), ref.astype(np.float64), atol=atol) + + @pytest.mark.parametrize( + "input_shape,output_shape", + [ + ((H, W), (1, H, W)), + ((1, H, W), (1, H, W)), + ((H, W, 1), (1, H, W)), + ((3, 1, H, W), (3, 1, H, W)), + ((3, H, W, 1), (3, 1, H, W)), + ], + ) + def test_input_layouts_accepted(self, input_shape, output_shape): + """All documented input layouts decode to the channel-first default.""" + quantized = np.full(input_shape, DEPTH_QMAX // 2, dtype=np.uint16) + out = dequantize_depth(quantized, output_unit=DEPTH_METER_UNIT, output_tensor=False) + assert out.shape == output_shape + + def test_pyav_frame_roundtrip(self): + """quantize → av.VideoFrame → dequantize works.""" + depth = _depth_metres_ramp() + frame = quantize_depth(depth, use_log=False, video_backend="pyav") + assert isinstance(frame, av.VideoFrame) + + recovered = dequantize_depth(frame, use_log=False, output_unit=DEPTH_METER_UNIT, output_tensor=False) + assert recovered.shape == (1, H, W) + tol = (DEFAULT_DEPTH_MAX - DEFAULT_DEPTH_MIN) / DEPTH_QMAX + 1e-3 + np.testing.assert_allclose(recovered[0], depth, atol=tol) + + def test_invalid_log_params_raises(self): + with pytest.raises(ValueError, match=r"depth_min \+ shift must be positive"): + quantize_depth(_depth_metres_ramp(), depth_min=1.0, shift=-2.0, use_log=True, video_backend=None) + + +# ── 2. Image writer depth support ───────────────────────────────────── + + +class TestImageWriterDepth: + """``image_array_to_pil_image`` and ``write_image`` for depth maps.""" + + @pytest.mark.parametrize("dtype,expected_mode", [(np.uint16, "I;16"), (np.float32, "F")]) + @pytest.mark.parametrize("shape", [(H, W), (H, W, 1), (1, H, W)]) + def test_pil_depth_modes_and_squeeze(self, dtype, expected_mode, shape): + """Single-channel depth converts to PIL with the right mode and (W, H) size.""" + arr = np.zeros(shape, dtype=dtype) + img = image_array_to_pil_image(arr) + assert img.mode == expected_mode + assert img.size == (W, H) + + def test_write_image_tiff_roundtrip(self, tmp_path): + """uint16 depth round-trips through .tiff.""" + arr = np.arange(H * W, dtype=np.uint16).reshape(H, W) + fpath = tmp_path / "depth.tiff" + write_image(arr, fpath) + with PIL.Image.open(fpath) as loaded: + recovered = np.array(loaded) + np.testing.assert_array_equal(recovered, arr) + + +# ── 3. Hardware-feature → depth flag ────────────────────────────────── + + +class TestHwToDatasetFeaturesDepth: + """``hw_to_dataset_features`` flags single-channel cameras as depth.""" + + @pytest.mark.parametrize("channels,is_depth", [(1, True), (3, False)]) + def test_depth_marker_by_channels(self, channels, is_depth): + from lerobot.utils.feature_utils import hw_to_dataset_features + + features = hw_to_dataset_features({"cam": (480, 640, channels)}, prefix="observation") + assert features["observation.images.cam"]["info"]["is_depth_map"] is is_depth + + def test_invalid_channel_count_raises(self): + from lerobot.utils.feature_utils import hw_to_dataset_features + + with pytest.raises(ValueError, match="Expected a 3-tuple"): + hw_to_dataset_features({"cam": (480, 640, 2)}, prefix="observation") + + +# ── 4. Feature-to-file-format routing ──────────────────────────────── + + +# Keys derived from DUMMY_CAMERA_FEATURES_WITH_DEPTH; pick one RGB and the depth camera. +RGB_KEY = next(iter(DUMMY_CAMERA_FEATURES)) +DEPTH_KEY = next(iter(DUMMY_DEPTH_CAMERA_FEATURES)) + + +class TestFeatureFileRouting: + """Depth vs RGB features route to the correct file format.""" + + NUM_FRAMES = 5 + + def test_image_mode_depth_tiff_rgb_png(self, tmp_path, features_factory): + """Without video encoding: depth → .tiff, RGB → .png.""" + from lerobot.datasets.lerobot_dataset import LeRobotDataset + + features = features_factory(camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH, use_videos=False) + dataset = LeRobotDataset.create( + repo_id=DUMMY_REPO_ID, + fps=DEFAULT_FPS, + features=features, + root=tmp_path / "ds", + use_videos=False, + ) + + add_frames(dataset, num_frames=self.NUM_FRAMES) + + buf = dataset.writer.episode_buffer + assert all(Path(p).suffix == ".tiff" for p in buf[DEPTH_KEY]) + assert all(Path(p).suffix == ".png" for p in buf[RGB_KEY]) + + dataset.save_episode() + dataset.finalize() + + def test_video_mode_depth_uses_depth_encoder(self, tmp_path, features_factory): + """With streaming video encoding: depth → DepthEncoderConfig, RGB does not.""" + from lerobot.datasets.lerobot_dataset import LeRobotDataset + + features = features_factory(camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH, use_videos=True) + dataset = LeRobotDataset.create( + repo_id=DUMMY_REPO_ID, + fps=DEFAULT_FPS, + features=features, + root=tmp_path / "ds", + use_videos=True, + streaming_encoding=True, + ) + + add_frames(dataset, num_frames=self.NUM_FRAMES) + + encoder = dataset.writer._streaming_encoder + assert encoder is not None + assert isinstance(encoder._threads[DEPTH_KEY].video_encoder, DepthEncoderConfig) + assert not isinstance(encoder._threads[RGB_KEY].video_encoder, DepthEncoderConfig) + + dataset.save_episode() + dataset.finalize() diff --git a/tests/datasets/test_image_writer.py b/tests/datasets/test_image_writer.py index 916b8f017..1cf2cf75c 100644 --- a/tests/datasets/test_image_writer.py +++ b/tests/datasets/test_image_writer.py @@ -94,7 +94,7 @@ def test_image_array_to_pil_image_pytorch_format(img_array_factory): def test_image_array_to_pil_image_single_channel(img_array_factory): img_array = img_array_factory(channels=1) - with pytest.raises(NotImplementedError): + with pytest.raises(ValueError, match="Unsupported single-channel image dtype"): image_array_to_pil_image(img_array) @@ -344,7 +344,7 @@ def test_with_different_image_formats(tmp_path, img_array_factory): writer = AsyncImageWriter() try: image_array = img_array_factory() - formats = ["png", "jpeg", "bmp"] + formats = ["png", "tiff", "tif"] for fmt in formats: fpath = tmp_path / f"test_image.{fmt}" write_image(image_array, fpath) diff --git a/tests/datasets/test_streaming_video_encoder.py b/tests/datasets/test_streaming_video_encoder.py index b69f24254..1ffad6854 100644 --- a/tests/datasets/test_streaming_video_encoder.py +++ b/tests/datasets/test_streaming_video_encoder.py @@ -26,7 +26,7 @@ pytest.importorskip("av", reason="av is required (install lerobot[dataset])") import av # noqa: E402 -from lerobot.configs import VideoEncoderConfig +from lerobot.configs import RGBEncoderConfig from lerobot.datasets.pyav_utils import get_codec from lerobot.datasets.video_utils import ( StreamingVideoEncoder, @@ -57,13 +57,11 @@ class TestCameraEncoderThread: result_queue: queue.Queue = queue.Queue(maxsize=1) stop_event = threading.Event() - enc_cfg = VideoEncoderConfig(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13) + enc_cfg = RGBEncoderConfig(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13) encoder_thread = _CameraEncoderThread( video_path=video_path, fps=fps, - vcodec=enc_cfg.vcodec, - pix_fmt=enc_cfg.pix_fmt, - codec_options=enc_cfg.get_codec_options(as_strings=True), + video_encoder=enc_cfg, frame_queue=frame_queue, result_queue=result_queue, stop_event=stop_event, @@ -108,13 +106,11 @@ class TestCameraEncoderThread: result_queue: queue.Queue = queue.Queue(maxsize=1) stop_event = threading.Event() - enc_cfg = VideoEncoderConfig(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13) + enc_cfg = RGBEncoderConfig(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13) encoder_thread = _CameraEncoderThread( video_path=video_path, fps=fps, - vcodec=enc_cfg.vcodec, - pix_fmt=enc_cfg.pix_fmt, - codec_options=enc_cfg.get_codec_options(as_strings=True), + video_encoder=enc_cfg, frame_queue=frame_queue, result_queue=result_queue, stop_event=stop_event, @@ -142,13 +138,11 @@ class TestCameraEncoderThread: result_queue: queue.Queue = queue.Queue(maxsize=1) stop_event = threading.Event() - enc_cfg = VideoEncoderConfig(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13) + enc_cfg = RGBEncoderConfig(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13) encoder_thread = _CameraEncoderThread( video_path=video_path, fps=fps, - vcodec=enc_cfg.vcodec, - pix_fmt=enc_cfg.pix_fmt, - codec_options=enc_cfg.get_codec_options(as_strings=True), + video_encoder=enc_cfg, frame_queue=frame_queue, result_queue=result_queue, stop_event=stop_event, @@ -171,15 +165,15 @@ class TestCameraEncoderThread: class TestStreamingVideoEncoder: def _make_encoder_config(self, **kwargs): - """Helper to build a VideoEncoderConfig.""" - return VideoEncoderConfig(**kwargs) + """Helper to build an RGBEncoderConfig.""" + return RGBEncoderConfig(**kwargs) def test_single_camera_episode(self, tmp_path): """Test encoding a single camera episode.""" video_keys = [f"{OBS_IMAGES}.laptop"] encoder = StreamingVideoEncoder( fps=30, - camera_encoder=self._make_encoder_config( + rgb_encoder=self._make_encoder_config( vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13 ), ) @@ -211,7 +205,7 @@ class TestStreamingVideoEncoder: video_keys = [f"{OBS_IMAGES}.laptop", f"{OBS_IMAGES}.phone"] encoder = StreamingVideoEncoder( fps=30, - camera_encoder=self._make_encoder_config(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30), + rgb_encoder=self._make_encoder_config(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30), ) encoder.start_episode(video_keys, tmp_path) @@ -237,7 +231,7 @@ class TestStreamingVideoEncoder: video_keys = [f"{OBS_IMAGES}.cam"] encoder = StreamingVideoEncoder( fps=30, - camera_encoder=self._make_encoder_config(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30), + rgb_encoder=self._make_encoder_config(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30), ) for ep in range(3): @@ -263,7 +257,7 @@ class TestStreamingVideoEncoder: video_keys = [f"{OBS_IMAGES}.cam"] encoder = StreamingVideoEncoder( fps=30, - camera_encoder=self._make_encoder_config(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30), + rgb_encoder=self._make_encoder_config(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30), ) encoder.start_episode(video_keys, tmp_path) @@ -309,7 +303,7 @@ class TestStreamingVideoEncoder: video_keys = [f"{OBS_IMAGES}.cam"] encoder = StreamingVideoEncoder( fps=30, - camera_encoder=self._make_encoder_config( + rgb_encoder=self._make_encoder_config( vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13 ), ) @@ -346,7 +340,7 @@ class TestStreamingVideoEncoder: video_keys = [f"{OBS_IMAGES}.cam1", f"{OBS_IMAGES}.cam2"] encoder = StreamingVideoEncoder( fps=30, - camera_encoder=self._make_encoder_config(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30), + rgb_encoder=self._make_encoder_config(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30), ) encoder.start_episode(video_keys, tmp_path) @@ -375,7 +369,7 @@ class TestStreamingVideoEncoder: def test_encoder_threads_passed_to_thread(self, tmp_path): """Test that encoder_threads is stored and passed through to encoder threads.""" video_keys = [f"{OBS_IMAGES}.cam"] - cfg = VideoEncoderConfig( + cfg = RGBEncoderConfig( vcodec="libsvtav1", pix_fmt="yuv420p", g=2, @@ -383,7 +377,7 @@ class TestStreamingVideoEncoder: ) encoder = StreamingVideoEncoder( fps=30, - camera_encoder=cfg, + rgb_encoder=cfg, encoder_threads=2, ) assert encoder._encoder_threads == 2 @@ -391,7 +385,8 @@ class TestStreamingVideoEncoder: # Verify codec options include thread tuning for libsvtav1 (lp=…) thread = encoder._threads[f"{OBS_IMAGES}.cam"] - assert "svtav1-params" in thread.codec_options or "threads" in thread.codec_options + codec_opts = thread.video_encoder.get_codec_options(encoder_threads=thread.encoder_threads) + assert "svtav1-params" in codec_opts or "threads" in codec_opts # Feed some frames and finish to ensure it works end-to-end num_frames = 10 @@ -422,7 +417,7 @@ class TestStreamingVideoEncoder: video_keys = [f"{OBS_IMAGES}.cam"] encoder = StreamingVideoEncoder( fps=30, - camera_encoder=self._make_encoder_config( + rgb_encoder=self._make_encoder_config( vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13 ), queue_maxsize=1, diff --git a/tests/datasets/test_video_encoding.py b/tests/datasets/test_video_encoding.py index 2a35f3210..80819d665 100644 --- a/tests/datasets/test_video_encoding.py +++ b/tests/datasets/test_video_encoding.py @@ -26,7 +26,7 @@ pytest.importorskip("av", reason="av is required (install lerobot[dataset])") import av # noqa: E402 -from lerobot.configs import VALID_VIDEO_CODECS, VideoEncoderConfig +from lerobot.configs import VALID_VIDEO_CODECS, DepthEncoderConfig, RGBEncoderConfig, VideoEncoderConfig from lerobot.datasets.image_writer import write_image from lerobot.datasets.lerobot_dataset import LeRobotDataset from lerobot.datasets.pyav_utils import get_codec @@ -37,7 +37,15 @@ from lerobot.datasets.video_utils import ( get_video_info, reencode_video, ) -from tests.fixtures.constants import DUMMY_VIDEO_INFO +from tests.fixtures.constants import ( + DUMMY_DEPTH_FEATURES, + DUMMY_DEPTH_KEY, + DUMMY_DEPTH_VIDEO_INFO_FULL, + DUMMY_VIDEO_FEATURES, + DUMMY_VIDEO_INFO, + DUMMY_VIDEO_KEY, +) +from tests.fixtures.dataset_factories import add_frames # Per-codec skip markers — validation tests only fire when the codec is available @@ -48,19 +56,74 @@ def _require_encoder(vcodec: str) -> pytest.MarkDecorator: require_libsvtav1 = _require_encoder("libsvtav1") require_h264 = _require_encoder("h264") +require_hevc = _require_encoder("hevc") require_videotoolbox = _require_encoder("h264_videotoolbox") require_nvenc = _require_encoder("h264_nvenc") require_vaapi = _require_encoder("h264_vaapi") require_qsv = _require_encoder("h264_qsv") -# ─── VideoEncoderConfig / codec options ────────────────────────────── +TEST_ARTIFACTS_DIR = Path(__file__).parent.parent / "artifacts" / "encoded_videos" + + +def _write_color_frames(imgs_dir: Path, num_frames: int = 4, height: int = 64, width: int = 96) -> None: + imgs_dir.mkdir(parents=True, exist_ok=True) + for i in range(num_frames): + arr = np.random.randint(0, 256, (height, width, 3), dtype=np.uint8) + write_image(arr, imgs_dir / f"frame-{i:06d}.png") + + +def _write_depth_frames(imgs_dir: Path, num_frames: int = 4, height: int = 64, width: int = 96) -> None: + """Write synthetic uint16 depth TIFFs (millimetres) for depth encoder tests. + + Uses a smooth linear ramp + per-frame offset (not white noise) so HEVC Main 12 + on ``gray12le`` compresses well. Values span ~100 mm to 10 m, covering most + of the default ``[DEPTH_MIN, DEPTH_MAX]`` metres range after + ``quantize_depth(input_unit="auto"="mm")``. + """ + imgs_dir.mkdir(parents=True, exist_ok=True) + base = np.linspace(100.0, 10_000.0, height * width, dtype=np.float32).reshape(height, width) + for i in range(num_frames): + arr = (base + 50.0 * i).clip(0, 65535).astype(np.uint16) + write_image(arr, imgs_dir / f"frame-{i:06d}.tiff") + + +def _encode_video( + path: Path, + num_frames: int = 4, + fps: int = 30, + cfg: VideoEncoderConfig | None = None, + depth: bool = False, +) -> Path: + """Write synthetic frames to a temp dir and encode them to ``path``. + + ``depth=False`` writes uint8 RGB PNG noise and encodes with ``cfg`` + (defaulting to the library default). ``depth=True`` writes synthetic uint16 + depth TIFFs and encodes with ``cfg`` or a default :class:`DepthEncoderConfig` + (HEVC Main 12 / ``gray12le``). + """ + imgs_dir = path.parent / f"imgs_{path.stem}" + if depth: + _write_depth_frames(imgs_dir, num_frames=num_frames) + cfg = cfg or DepthEncoderConfig() + else: + _write_color_frames(imgs_dir, num_frames=num_frames) + encode_video_frames(imgs_dir, path, fps=fps, video_encoder=cfg, overwrite=True) + return path + + +def _read_feature_info(dataset: LeRobotDataset, key: str = DUMMY_VIDEO_KEY) -> dict: + info = json.loads((dataset.root / INFO_PATH).read_text()) + return info["features"][key]["info"] + + +# ─── RGBEncoderConfig / codec options ────────────────────────────── class TestCodecOptions: @require_libsvtav1 def test_libsvtav1_defaults(self): - cfg = VideoEncoderConfig() + cfg = RGBEncoderConfig() opts = cfg.get_codec_options() assert opts["g"] == 2 assert opts["crf"] == 30 @@ -68,12 +131,12 @@ class TestCodecOptions: @require_libsvtav1 def test_libsvtav1_custom_preset(self): - cfg = VideoEncoderConfig(preset=8) + cfg = RGBEncoderConfig(preset=8) assert cfg.get_codec_options()["preset"] == 8 @require_h264 def test_h264_options(self): - cfg = VideoEncoderConfig(vcodec="h264", g=10, crf=23, preset=None) + cfg = RGBEncoderConfig(vcodec="h264", g=10, crf=23, preset=None) opts = cfg.get_codec_options() assert opts["g"] == 10 assert opts["crf"] == 23 @@ -81,120 +144,120 @@ class TestCodecOptions: @require_videotoolbox def test_videotoolbox_options(self): - cfg = VideoEncoderConfig(vcodec="h264_videotoolbox", g=2, crf=30, preset=None) + cfg = RGBEncoderConfig(vcodec="h264_videotoolbox", g=2, crf=30, preset=None) opts = cfg.get_codec_options() assert opts["g"] == 2 assert opts["q:v"] == 40 assert "crf" not in opts - @_require_encoder("h264_nvenc") + @require_nvenc def test_nvenc_options(self): - cfg = VideoEncoderConfig(vcodec="h264_nvenc", g=2, crf=25, preset=None) + cfg = RGBEncoderConfig(vcodec="h264_nvenc", g=2, crf=25, preset=None) opts = cfg.get_codec_options() assert opts["rc"] == 0 assert opts["qp"] == 25 assert "crf" not in opts assert opts["g"] == 2 - @_require_encoder("h264_vaapi") + @require_vaapi def test_vaapi_options(self): - cfg = VideoEncoderConfig(vcodec="h264_vaapi", crf=28, preset=None) + cfg = RGBEncoderConfig(vcodec="h264_vaapi", crf=28, preset=None) assert cfg.get_codec_options()["qp"] == 28 - @_require_encoder("h264_qsv") + @require_qsv def test_qsv_options(self): - cfg = VideoEncoderConfig(vcodec="h264_qsv", crf=25, preset=None) + cfg = RGBEncoderConfig(vcodec="h264_qsv", crf=25, preset=None) assert cfg.get_codec_options()["global_quality"] == 25 @require_h264 def test_no_g_no_crf(self): - cfg = VideoEncoderConfig(vcodec="h264", g=None, crf=None, preset=None) + cfg = RGBEncoderConfig(vcodec="h264", g=None, crf=None, preset=None) opts = cfg.get_codec_options() assert "g" not in opts assert "crf" not in opts @require_libsvtav1 def test_encoder_threads_libsvtav1(self): - cfg = VideoEncoderConfig(fast_decode=0) + cfg = RGBEncoderConfig(fast_decode=0) opts = cfg.get_codec_options(encoder_threads=4) assert "lp=4" in opts.get("svtav1-params", "") @require_h264 def test_encoder_threads_h264(self): - cfg = VideoEncoderConfig(vcodec="h264", preset=None) + cfg = RGBEncoderConfig(vcodec="h264", preset=None) assert cfg.get_codec_options(encoder_threads=2)["threads"] == 2 @require_libsvtav1 def test_fast_decode_libsvtav1(self): - cfg = VideoEncoderConfig(fast_decode=1) + cfg = RGBEncoderConfig(fast_decode=1) opts = cfg.get_codec_options() assert "fast-decode=1" in opts.get("svtav1-params", "") @require_libsvtav1 def test_libsvtav1_fast_decode_clamped_to_svt_range(self): """Out-of-range fast_decode is clamped to [0, 2] in svtav1-params (SVT-AV1 FastDecode).""" - cfg = VideoEncoderConfig(fast_decode=100) + cfg = RGBEncoderConfig(fast_decode=100) assert "fast-decode=2" in cfg.get_codec_options().get("svtav1-params", "") - cfg_neg = VideoEncoderConfig(fast_decode=-5) + cfg_neg = RGBEncoderConfig(fast_decode=-5) assert "fast-decode=0" in cfg_neg.get_codec_options().get("svtav1-params", "") @require_h264 def test_fast_decode_h264(self): - cfg = VideoEncoderConfig(vcodec="h264", fast_decode=1, preset=None) + cfg = RGBEncoderConfig(vcodec="h264", fast_decode=1, preset=None) assert cfg.get_codec_options()["tune"] == "fastdecode" @require_libsvtav1 def test_pix_fmt_unsupported_raises(self): """Passing an unsupported pix_fmt is a hard error.""" with pytest.raises(ValueError, match="pix_fmt"): - VideoEncoderConfig(pix_fmt="yuv444p") # libsvtav1 only supports yuv420p variants + RGBEncoderConfig(pix_fmt="yuv444p") # libsvtav1 only supports yuv420p variants @require_libsvtav1 @require_h264 def test_preset_default_behaviour(self): """Empty constructor picks preset=12 (libsvtav1 path); other codecs stay None.""" - assert VideoEncoderConfig().preset == 12 - assert VideoEncoderConfig(vcodec="libsvtav1").preset == 12 - assert VideoEncoderConfig(vcodec="h264").preset is None - assert VideoEncoderConfig(vcodec="h264", preset=None).preset is None + assert RGBEncoderConfig().preset == 12 + assert RGBEncoderConfig(vcodec="libsvtav1").preset == 12 + assert RGBEncoderConfig(vcodec="h264").preset is None + assert RGBEncoderConfig(vcodec="h264", preset=None).preset is None @require_h264 def test_preset_string_on_h264(self): """h264 accepts string presets and forwards them to FFmpeg.""" - cfg = VideoEncoderConfig(vcodec="h264", preset="slow") + cfg = RGBEncoderConfig(vcodec="h264", preset="slow") assert cfg.get_codec_options()["preset"] == "slow" @require_videotoolbox def test_preset_on_videotoolbox_not_set(self): """videotoolbox has no preset option at all.""" - cfg = VideoEncoderConfig(vcodec="h264_videotoolbox", preset="slow") + cfg = RGBEncoderConfig(vcodec="h264_videotoolbox", preset="slow") assert "preset" not in cfg.get_codec_options() @require_libsvtav1 def test_libsvtav1_preset_out_of_range_raises(self): """libsvtav1 preset must sit in [-2, 13] as exposed by PyAV.""" with pytest.raises(ValueError, match="out of range"): - VideoEncoderConfig(vcodec="libsvtav1", preset=100) + RGBEncoderConfig(vcodec="libsvtav1", preset=100) with pytest.raises(ValueError, match="out of range"): - VideoEncoderConfig(vcodec="libsvtav1", preset=-3) + RGBEncoderConfig(vcodec="libsvtav1", preset=-3) @require_libsvtav1 def test_libsvtav1_crf_out_of_range_raises(self): """libsvtav1 crf must sit in [0, 63].""" with pytest.raises(ValueError, match="crf.*out of range"): - VideoEncoderConfig(vcodec="libsvtav1", crf=64) + RGBEncoderConfig(vcodec="libsvtav1", crf=64) @require_libsvtav1 def test_libsvtav1_crf_rejects_python_float(self): """libsvtav1 exposes ``crf`` as an INT AVOption; Python float must not pass validation.""" with pytest.raises(ValueError, match="float values are not allowed"): - VideoEncoderConfig(vcodec="libsvtav1", crf=2.5) + RGBEncoderConfig(vcodec="libsvtav1", crf=2.5) @require_libsvtav1 def test_libsvtav1_extra_crf_rejects_fractional_string(self): """INT options reject fractional values even when supplied only via ``extra_options``.""" with pytest.raises(ValueError, match="float values are not allowed"): - VideoEncoderConfig( + RGBEncoderConfig( vcodec="libsvtav1", crf=None, extra_options={"crf": "2.5"}, @@ -203,7 +266,7 @@ class TestCodecOptions: @require_libsvtav1 def test_libsvtav1_extra_crf_rejects_float(self): with pytest.raises(ValueError, match="float values are not allowed"): - VideoEncoderConfig( + RGBEncoderConfig( vcodec="libsvtav1", crf=None, extra_options={"crf": 2.5}, @@ -212,13 +275,13 @@ class TestCodecOptions: @require_h264 def test_h264_crf_accepts_float_and_int(self): """x264 exposes crf as a FLOAT option, so both int and float are accepted.""" - assert VideoEncoderConfig(vcodec="h264", crf=23).get_codec_options()["crf"] == 23 - assert VideoEncoderConfig(vcodec="h264", crf=23.5).get_codec_options()["crf"] == 23.5 + assert RGBEncoderConfig(vcodec="h264", crf=23).get_codec_options()["crf"] == 23 + assert RGBEncoderConfig(vcodec="h264", crf=23.5).get_codec_options()["crf"] == 23.5 @require_libsvtav1 def test_validate_is_rerunnable(self): """After mutating a field, validate() re-checks and surfaces new issues.""" - cfg = VideoEncoderConfig(vcodec="libsvtav1") + cfg = RGBEncoderConfig(vcodec="libsvtav1") cfg.preset = 100 # now out of range with pytest.raises(ValueError, match="out of range"): cfg.validate() @@ -227,58 +290,58 @@ class TestCodecOptions: class TestExtraOptions: @require_libsvtav1 def test_default_is_empty_dict(self): - cfg = VideoEncoderConfig() + cfg = RGBEncoderConfig() assert cfg.extra_options == {} @require_libsvtav1 def test_unknown_key_passes_through(self): """Keys not published as AVOptions are forwarded to FFmpeg.""" - cfg = VideoEncoderConfig(extra_options={"totally_made_up_option": "value"}) + cfg = RGBEncoderConfig(extra_options={"totally_made_up_option": "value"}) assert cfg.extra_options == {"totally_made_up_option": "value"} @require_libsvtav1 def test_numeric_value_in_range_ok(self): """libsvtav1 exposes ``qp`` as INT in [0, 63].""" - cfg = VideoEncoderConfig(extra_options={"qp": 30}) + cfg = RGBEncoderConfig(extra_options={"qp": 30}) assert cfg.extra_options == {"qp": 30} @require_libsvtav1 def test_numeric_out_of_range_raises(self): with pytest.raises(ValueError, match=r"qp=.*out of range"): - VideoEncoderConfig(extra_options={"qp": 999}) + RGBEncoderConfig(extra_options={"qp": 999}) @require_libsvtav1 def test_numeric_string_accepted_in_range(self): """Numeric strings are accepted for numeric options (mirrors FFmpeg).""" - cfg = VideoEncoderConfig(extra_options={"qp": "18"}) + cfg = RGBEncoderConfig(extra_options={"qp": "18"}) assert cfg.extra_options == {"qp": "18"} @require_libsvtav1 def test_numeric_string_out_of_range_raises(self): with pytest.raises(ValueError, match=r"qp=.*out of range"): - VideoEncoderConfig(extra_options={"qp": "999"}) + RGBEncoderConfig(extra_options={"qp": "999"}) @require_libsvtav1 def test_non_numeric_string_on_numeric_option_raises(self): with pytest.raises(ValueError, match=r"qp=.*not numeric"): - VideoEncoderConfig(extra_options={"qp": "medium"}) + RGBEncoderConfig(extra_options={"qp": "medium"}) @require_libsvtav1 def test_bool_on_numeric_option_raises(self): """``bool`` is explicitly rejected for numeric options.""" with pytest.raises(ValueError, match=r"qp=.*not numeric"): - VideoEncoderConfig(extra_options={"qp": True}) + RGBEncoderConfig(extra_options={"qp": True}) @require_h264 def test_string_option_passes_through_unchecked(self): """String-typed AVOptions are NOT enum-checked (too many accept freeform).""" - cfg = VideoEncoderConfig(vcodec="h264", preset=None, extra_options={"tune": "some-future-tune"}) + cfg = RGBEncoderConfig(vcodec="h264", preset=None, extra_options={"tune": "some-future-tune"}) assert cfg.extra_options == {"tune": "some-future-tune"} @require_libsvtav1 def test_merged_into_codec_options_and_stringified(self): """Typed merge by default; ``as_strings=True`` matches FFmpeg option dict.""" - cfg = VideoEncoderConfig(extra_options={"qp": 20}) + cfg = RGBEncoderConfig(extra_options={"qp": 20}) opts = cfg.get_codec_options() assert opts["qp"] == 20 assert isinstance(opts["qp"], int) @@ -287,25 +350,25 @@ class TestExtraOptions: @require_libsvtav1 def test_structured_fields_win_on_collision(self): """A colliding extra_options key is discarded; the structured field wins.""" - cfg = VideoEncoderConfig(crf=30, extra_options={"crf": 18}) + cfg = RGBEncoderConfig(crf=30, extra_options={"crf": 18}) assert cfg.get_codec_options()["crf"] == 30 class TestEncoderDetection: @require_h264 def test_explicit_codec_kept_when_available(self): - cfg = VideoEncoderConfig(vcodec="h264") + cfg = RGBEncoderConfig(vcodec="h264") assert cfg.vcodec == "h264" @require_videotoolbox def test_auto_picks_videotoolbox_when_available(self): """``h264_videotoolbox`` sits at the top of ``HW_VIDEO_CODECS`` so it wins when present.""" - cfg = VideoEncoderConfig(vcodec="auto") + cfg = RGBEncoderConfig(vcodec="auto") assert cfg.vcodec == "h264_videotoolbox" def test_invalid_codec_raises(self): with pytest.raises(ValueError, match="Invalid vcodec"): - VideoEncoderConfig(vcodec="not_a_real_codec") + RGBEncoderConfig(vcodec="not_a_real_codec") def test_hw_encoder_names_listed_as_valid(self): assert "auto" in VALID_VIDEO_CODECS @@ -313,59 +376,6 @@ class TestEncoderDetection: assert "h264_nvenc" in VALID_VIDEO_CODECS -TEST_ARTIFACTS_DIR = Path(__file__).parent.parent / "artifacts" / "encoded_videos" - -# Default video feature set used by persistence tests. -VIDEO_FEATURES = { - "observation.images.cam": { - "dtype": "video", - "shape": (64, 96, 3), - "names": ["height", "width", "channels"], - }, - "action": {"dtype": "float32", "shape": (2,), "names": ["a", "b"]}, -} -VIDEO_KEY = "observation.images.cam" - - -def _write_frames(imgs_dir: Path, num_frames: int = 4, height: int = 64, width: int = 96) -> None: - imgs_dir.mkdir(parents=True, exist_ok=True) - for i in range(num_frames): - arr = np.random.randint(0, 256, (height, width, 3), dtype=np.uint8) - write_image(arr, imgs_dir / f"frame-{i:06d}.png") - - -def _encode_video( - path: Path, num_frames: int = 4, fps: int = 30, cfg: VideoEncoderConfig | None = None -) -> Path: - imgs_dir = path.parent / f"imgs_{path.stem}" - _write_frames(imgs_dir, num_frames=num_frames) - encode_video_frames(imgs_dir, path, fps=fps, camera_encoder=cfg, overwrite=True) - return path - - -def _read_feature_info(dataset: LeRobotDataset) -> dict: - info = json.loads((dataset.root / INFO_PATH).read_text()) - return info["features"][VIDEO_KEY]["info"] - - -def _add_frames(dataset: LeRobotDataset, num_frames: int, video_keys: list[str] | None = None) -> None: - from lerobot.utils.constants import DEFAULT_FEATURES - - if video_keys is None: - video_keys = dataset.meta.video_keys - for _ in range(num_frames): - frame: dict = {"task": "test"} - for key, ft in dataset.meta.features.items(): - if key in DEFAULT_FEATURES: - continue - shape = ft["shape"] - if key in video_keys: - frame[key] = np.random.randint(0, 256, shape, dtype=np.uint8) - else: - frame[key] = np.zeros(shape, dtype=np.float32) - dataset.add_frame(frame) - - class TestGetVideoInfo: def test_returns_all_stream_fields(self): info = get_video_info(TEST_ARTIFACTS_DIR / "clip_4frames.mp4") @@ -375,7 +385,7 @@ class TestGetVideoInfo: assert info["video.pix_fmt"] == "yuv420p" assert info["video.fps"] == 30 assert info["video.channels"] == 3 - assert info["video.is_depth_map"] is False + assert info["is_depth_map"] is False assert info["has_audio"] is False assert "video.g" not in info assert "video.crf" not in info @@ -383,9 +393,9 @@ class TestGetVideoInfo: @require_libsvtav1 def test_merges_encoder_config_as_video_prefixed_entries(self): - cfg = VideoEncoderConfig(vcodec="libsvtav1", g=2, crf=30, preset=12) + cfg = RGBEncoderConfig(vcodec="libsvtav1", g=2, crf=30, preset=12) - info = get_video_info(TEST_ARTIFACTS_DIR / "clip_4frames.mp4", camera_encoder=cfg) + info = get_video_info(TEST_ARTIFACTS_DIR / "clip_4frames.mp4", video_encoder=cfg) assert info["video.g"] == 2 assert info["video.crf"] == 30 @@ -396,13 +406,18 @@ class TestGetVideoInfo: @require_libsvtav1 def test_stream_derived_keys_take_precedence_over_config(self): - cfg = VideoEncoderConfig(vcodec="libsvtav1", pix_fmt="yuv420p") + cfg = RGBEncoderConfig(vcodec="libsvtav1", pix_fmt="yuv420p") - info = get_video_info(TEST_ARTIFACTS_DIR / "clip_4frames.mp4", camera_encoder=cfg) + info = get_video_info(TEST_ARTIFACTS_DIR / "clip_4frames.mp4", video_encoder=cfg) assert info["video.codec"] # populated from stream, not from config's vcodec assert info["video.pix_fmt"] == "yuv420p" + def test_depth_encoder_config_sets_is_depth_map_true(self): + """A ``DepthEncoderConfig`` causes ``get_video_info`` to mark the stream as depth.""" + info = get_video_info(TEST_ARTIFACTS_DIR / "clip_4frames.mp4", video_encoder=DepthEncoderConfig()) + assert info["is_depth_map"] is True + class TestEncodeVideoFrames: @require_libsvtav1 @@ -434,7 +449,7 @@ class TestEncodeVideoFrames: def test_overwrite_false_skips_existing_file(self, tmp_path): imgs_dir = tmp_path / "imgs" - _write_frames(imgs_dir) + _write_color_frames(imgs_dir) video_path = tmp_path / "out.mp4" sentinel = b"pre-existing content" video_path.write_bytes(sentinel) @@ -446,7 +461,7 @@ class TestEncodeVideoFrames: @require_libsvtav1 def test_overwrite_true_replaces_existing_file(self, tmp_path): imgs_dir = tmp_path / "imgs" - _write_frames(imgs_dir) + _write_color_frames(imgs_dir) video_path = tmp_path / "out.mp4" video_path.write_bytes(b"stale content") @@ -458,10 +473,10 @@ class TestEncodeVideoFrames: @require_libsvtav1 def test_custom_encoder_config_fields_stored_in_info(self, tmp_path): """All stream-derived and encoder config fields are present after encoding.""" - cfg = VideoEncoderConfig(vcodec="libsvtav1", g=4, crf=25, preset=10) + cfg = RGBEncoderConfig(vcodec="libsvtav1", g=4, crf=25, preset=10) video_path = _encode_video(tmp_path / "out.mp4", num_frames=4, fps=30, cfg=cfg) - info = get_video_info(video_path, camera_encoder=cfg) + info = get_video_info(video_path, video_encoder=cfg) # Stream-derived assert info["video.height"] == 64 @@ -470,7 +485,7 @@ class TestEncodeVideoFrames: assert info["video.codec"] == "av1" assert info["video.pix_fmt"] == "yuv420p" assert info["video.fps"] == 30 - assert info["video.is_depth_map"] is False + assert info["is_depth_map"] is False assert info["has_audio"] is False # Encoder config assert info["video.g"] == 4 @@ -487,15 +502,15 @@ class TestReencodeVideo: def test_reencode_video(self, tmp_path): src = TEST_ARTIFACTS_DIR / "clip_4frames.mp4" out = tmp_path / "reencoded.mp4" - cfg = VideoEncoderConfig(vcodec="h264", g=6, crf=23, pix_fmt="yuv444p") - reencode_video(src, out, camera_encoder=cfg, overwrite=True) + cfg = RGBEncoderConfig(vcodec="h264", g=6, crf=23, pix_fmt="yuv444p") + reencode_video(src, out, video_encoder=cfg, overwrite=True) assert out.exists() with av.open(str(out)) as container: n_frames = sum(1 for _ in container.decode(video=0)) assert n_frames == 4 - info = get_video_info(out, camera_encoder=cfg) + info = get_video_info(out, video_encoder=cfg) assert info["video.codec"] == "h264" assert info["video.pix_fmt"] == "yuv444p" assert info["video.height"] == 64 @@ -508,8 +523,8 @@ class TestReencodeVideo: def test_reencode_video_trim_window(self, tmp_path): src = TEST_ARTIFACTS_DIR / "clip_6frames.mp4" out = tmp_path / "trim_window.mp4" - cfg = VideoEncoderConfig(vcodec="h264") - reencode_video(src, out, camera_encoder=cfg, start_time_s=0.05, end_time_s=0.12, overwrite=True) + cfg = RGBEncoderConfig(vcodec="h264") + reencode_video(src, out, video_encoder=cfg, start_time_s=0.05, end_time_s=0.12, overwrite=True) with av.open(str(out)) as container: frames = list(container.decode(video=0)) @@ -578,12 +593,12 @@ class TestEncoderConfigPersistence: @require_libsvtav1 def test_first_episode_save_persists_encoder_config(self, tmp_path, empty_lerobot_dataset_factory): - cfg = VideoEncoderConfig(vcodec="libsvtav1", g=2, crf=30, preset=12) + cfg = RGBEncoderConfig(vcodec="libsvtav1", g=2, crf=30, preset=12) dataset = empty_lerobot_dataset_factory( - root=tmp_path / "ds", features=VIDEO_FEATURES, use_videos=True, camera_encoder=cfg + root=tmp_path / "ds", features=DUMMY_VIDEO_FEATURES, use_videos=True, rgb_encoder=cfg ) - _add_frames(dataset, num_frames=4) + add_frames(dataset, num_frames=4) dataset.save_episode() dataset.finalize() @@ -601,16 +616,16 @@ class TestEncoderConfigPersistence: @require_libsvtav1 def test_second_episode_does_not_overwrite_encoder_fields(self, tmp_path, empty_lerobot_dataset_factory): - cfg = VideoEncoderConfig(vcodec="libsvtav1", g=2, crf=30, preset=12) + cfg = RGBEncoderConfig(vcodec="libsvtav1", g=2, crf=30, preset=12) dataset = empty_lerobot_dataset_factory( - root=tmp_path / "ds", features=VIDEO_FEATURES, use_videos=True, camera_encoder=cfg + root=tmp_path / "ds", features=DUMMY_VIDEO_FEATURES, use_videos=True, rgb_encoder=cfg ) - _add_frames(dataset, num_frames=4) + add_frames(dataset, num_frames=4) dataset.save_episode() first_info = dict(_read_feature_info(dataset)) - _add_frames(dataset, num_frames=4) + add_frames(dataset, num_frames=4) dataset.save_episode() dataset.finalize() @@ -618,13 +633,13 @@ class TestEncoderConfigPersistence: class TestFromVideoInfo: - """``VideoEncoderConfig.from_video_info`` reconstructs an encoder config + """``RGBEncoderConfig.from_video_info`` reconstructs an encoder config from the ``video.*`` keys persisted in a dataset's ``info.json``. """ @require_libsvtav1 def test_reconstructs_from_dummy_video_info(self): - cfg = VideoEncoderConfig.from_video_info(DUMMY_VIDEO_INFO) + cfg = RGBEncoderConfig.from_video_info(DUMMY_VIDEO_INFO) # Canonical stream codec ``"av1"`` is aliased to the encoder name. assert cfg.vcodec == "libsvtav1" @@ -636,4 +651,220 @@ class TestFromVideoInfo: assert cfg.video_backend == DUMMY_VIDEO_INFO["video.video_backend"] # ``{}`` placeholder (typical after a merge with disagreeing sources) # must not leak into the reconstructed config. - assert cfg.extra_options == VideoEncoderConfig().extra_options + assert cfg.extra_options == RGBEncoderConfig().extra_options + + +# ─── Depth-specific encoding tests ──────────────────────────────────── + + +class TestEncodeDepthVideoFrames: + """Depth mirror of :class:`TestEncodeVideoFrames`. + + Exercises ``encode_video_frames`` end-to-end through + :class:`DepthEncoderConfig` (HEVC Main 12 / ``gray12le``) on synthetic + uint16 depth TIFFs. + """ + + @require_hevc + def test_produces_readable_file(self, tmp_path): + video_path = _encode_video(tmp_path / "out.mp4", depth=True) + + assert video_path.exists() + info = get_video_info(video_path, video_encoder=DepthEncoderConfig()) + assert info["video.height"] == 64 + assert info["video.width"] == 96 + assert info["video.codec"] == "hevc" + assert info["video.pix_fmt"] == "gray12le" + assert info["video.channels"] == 1 + assert info["is_depth_map"] is True + + @require_hevc + def test_frame_count_and_duration_match_input(self, tmp_path): + num_frames = 10 + fps = 30 + video_path = _encode_video(tmp_path / "out.mp4", num_frames=num_frames, fps=fps, depth=True) + + with av.open(str(video_path)) as container: + stream = container.streams.video[0] + actual_frames = sum(1 for _ in container.decode(stream)) + duration = ( + float(stream.duration * stream.time_base) + if stream.duration is not None + else float(container.duration / av.time_base) + ) + + assert actual_frames == num_frames + assert abs(duration - num_frames / fps) < 0.1 + + def test_overwrite_false_skips_existing_file(self, tmp_path): + """Codec-agnostic: file-system semantics must hold even without an HEVC encoder.""" + imgs_dir = tmp_path / "imgs" + _write_depth_frames(imgs_dir) + video_path = tmp_path / "out.mp4" + sentinel = b"pre-existing depth content" + video_path.write_bytes(sentinel) + + encode_video_frames(imgs_dir, video_path, fps=30, video_encoder=DepthEncoderConfig(), overwrite=False) + + assert video_path.read_bytes() == sentinel + + @require_hevc + def test_overwrite_true_replaces_existing_file(self, tmp_path): + imgs_dir = tmp_path / "imgs" + _write_depth_frames(imgs_dir) + video_path = tmp_path / "out.mp4" + video_path.write_bytes(b"stale content") + + encode_video_frames(imgs_dir, video_path, fps=30, video_encoder=DepthEncoderConfig(), overwrite=True) + + info = get_video_info(video_path, video_encoder=DepthEncoderConfig()) + assert info["video.height"] == 64 + assert info["video.pix_fmt"] == "gray12le" + assert info["is_depth_map"] is True + + @require_hevc + def test_custom_encoder_config_fields_stored_in_info(self, tmp_path): + """All stream-derived and depth-encoder config fields are present after encoding.""" + cfg = DepthEncoderConfig( + vcodec="hevc", + pix_fmt="gray12le", + g=4, + crf=25, + extra_options={}, + depth_min=0.05, + depth_max=8.0, + shift=2.5, + use_log=False, + ) + video_path = _encode_video(tmp_path / "out.mp4", num_frames=4, fps=30, cfg=cfg, depth=True) + + info = get_video_info(video_path, video_encoder=cfg) + + # Stream-derived + assert info["video.height"] == 64 + assert info["video.width"] == 96 + assert info["video.channels"] == 1 + assert info["video.codec"] == "hevc" + assert info["video.pix_fmt"] == "gray12le" + assert info["video.fps"] == 30 + assert info["is_depth_map"] is True + assert info["has_audio"] is False + # Base encoder config + assert info["video.g"] == 4 + assert info["video.crf"] == 25 + assert info["video.fast_decode"] == 0 + assert info["video.video_backend"] == "pyav" + assert info["video.extra_options"] == {} + # Depth-specific tuning + assert info["video.depth_min"] == 0.05 + assert info["video.depth_max"] == 8.0 + assert info["video.shift"] == 2.5 + assert info["video.use_log"] is False + + +class TestDepthEncoderConfigPersistence: + """Depth mirror of :class:`TestEncoderConfigPersistence`. + + ``DepthEncoderConfig`` must be stored as ``video.`` entries + (including the depth-specific ``depth_min`` / ``depth_max`` / ``shift`` / + ``use_log``) under ``info["features"][]["info"]`` when the + first episode is saved. + """ + + @require_hevc + def test_first_episode_save_persists_depth_encoder_config(self, tmp_path, empty_lerobot_dataset_factory): + cfg = DepthEncoderConfig( + vcodec="hevc", + pix_fmt="gray12le", + g=2, + crf=30, + extra_options={}, + depth_min=0.05, + depth_max=8.0, + shift=2.5, + use_log=False, + ) + dataset = empty_lerobot_dataset_factory( + root=tmp_path / "ds", features=DUMMY_DEPTH_FEATURES, use_videos=True, depth_encoder=cfg + ) + + add_frames(dataset, num_frames=4) + dataset.save_episode() + dataset.finalize() + + info = _read_feature_info(dataset, key=DUMMY_DEPTH_KEY) + + # Stream-derived + assert info["video.height"] == 64 + assert info["video.width"] == 96 + assert info["video.fps"] == 30 + assert info["video.codec"] == "hevc" + assert info["video.pix_fmt"] == "gray12le" + assert info["is_depth_map"] is True + # Base encoder config + assert info["video.g"] == 2 + assert info["video.crf"] == 30 + assert info["video.fast_decode"] == 0 + assert info["video.video_backend"] == "pyav" + assert info["video.extra_options"] == {} + # Depth-specific tuning + assert info["video.depth_min"] == 0.05 + assert info["video.depth_max"] == 8.0 + assert info["video.shift"] == 2.5 + assert info["video.use_log"] is False + + @require_hevc + def test_second_episode_does_not_overwrite_depth_encoder_fields( + self, tmp_path, empty_lerobot_dataset_factory + ): + cfg = DepthEncoderConfig( + vcodec="hevc", + pix_fmt="gray12le", + g=2, + crf=30, + depth_min=0.05, + depth_max=8.0, + shift=2.5, + use_log=False, + ) + dataset = empty_lerobot_dataset_factory( + root=tmp_path / "ds", features=DUMMY_DEPTH_FEATURES, use_videos=True, depth_encoder=cfg + ) + + add_frames(dataset, num_frames=4) + dataset.save_episode() + first_info = dict(_read_feature_info(dataset, key=DUMMY_DEPTH_KEY)) + + add_frames(dataset, num_frames=4) + dataset.save_episode() + dataset.finalize() + + assert _read_feature_info(dataset, key=DUMMY_DEPTH_KEY) == first_info + + +class TestDepthFromVideoInfo: + """``DepthEncoderConfig.from_video_info`` reconstructs a depth encoder + config from the ``video.*`` keys persisted in a dataset's ``info.json``. + + Depth mirror of :class:`TestFromVideoInfo`. + """ + + @require_hevc + def test_reconstructs_from_dummy_depth_video_info(self): + cfg = DepthEncoderConfig.from_video_info(DUMMY_DEPTH_VIDEO_INFO_FULL) + + # No alias for ``"hevc"``; the canonical stream codec is reused as-is. + assert cfg.vcodec == "hevc" + assert cfg.pix_fmt == DUMMY_DEPTH_VIDEO_INFO_FULL["video.pix_fmt"] + assert cfg.g == DUMMY_DEPTH_VIDEO_INFO_FULL["video.g"] + assert cfg.crf == DUMMY_DEPTH_VIDEO_INFO_FULL["video.crf"] + assert cfg.fast_decode == DUMMY_DEPTH_VIDEO_INFO_FULL["video.fast_decode"] + assert cfg.video_backend == DUMMY_DEPTH_VIDEO_INFO_FULL["video.video_backend"] + # ``{}`` placeholder (typical after a merge with disagreeing sources) + # must not leak into the reconstructed config. + assert cfg.extra_options == DepthEncoderConfig().extra_options + # Depth-specific tuning round-trips through ``info.json``. + assert cfg.depth_min == DUMMY_DEPTH_VIDEO_INFO_FULL["video.depth_min"] + assert cfg.depth_max == DUMMY_DEPTH_VIDEO_INFO_FULL["video.depth_max"] + assert cfg.shift == DUMMY_DEPTH_VIDEO_INFO_FULL["video.shift"] + assert cfg.use_log == DUMMY_DEPTH_VIDEO_INFO_FULL["video.use_log"] diff --git a/tests/fixtures/constants.py b/tests/fixtures/constants.py index 4d578b503..d6f4f8ae5 100644 --- a/tests/fixtures/constants.py +++ b/tests/fixtures/constants.py @@ -39,12 +39,56 @@ DUMMY_VIDEO_INFO = { "video.crf": 30, "video.preset": 12, "video.fast_decode": 0, - "video.is_depth_map": False, + "is_depth_map": False, "has_audio": False, } DUMMY_CAMERA_FEATURES = { "laptop": {"shape": (64, 96, 3), "names": ["height", "width", "channels"], "info": DUMMY_VIDEO_INFO}, "phone": {"shape": (64, 96, 3), "names": ["height", "width", "channels"], "info": DUMMY_VIDEO_INFO}, } +DUMMY_DEPTH_VIDEO_INFO = { + **DUMMY_VIDEO_INFO, + "is_depth_map": True, +} +DUMMY_DEPTH_VIDEO_INFO_FULL = { + **{k: v for k, v in DUMMY_VIDEO_INFO.items() if k != "video.preset"}, + "video.codec": "hevc", + "video.pix_fmt": "gray12le", + "is_depth_map": True, + "video.depth_min": 0.05, + "video.depth_max": 8.0, + "video.shift": 2.5, + "video.use_log": True, +} +DUMMY_DEPTH_CAMERA_FEATURES = { + "laptop_depth": { + "shape": (64, 96, 1), + "names": ["height", "width", "channels"], + "info": DUMMY_DEPTH_VIDEO_INFO, + }, +} +DUMMY_CAMERA_FEATURES_WITH_DEPTH = {**DUMMY_CAMERA_FEATURES, **DUMMY_DEPTH_CAMERA_FEATURES} DUMMY_CHW = (3, 96, 128) DUMMY_HWC = (96, 128, 3) + +# Default video feature set used by video-encoding persistence tests. +DUMMY_VIDEO_FEATURES = { + "observation.images.cam": { + "dtype": "video", + "shape": (64, 96, 3), + "names": ["height", "width", "channels"], + }, + "action": {"dtype": "float32", "shape": (2,), "names": ["a", "b"]}, +} +DUMMY_VIDEO_KEY = "observation.images.cam" + +DUMMY_DEPTH_FEATURES = { + "observation.images.depth": { + "dtype": "video", + "shape": (64, 96, 1), + "names": ["height", "width", "channels"], + "info": {"is_depth_map": True}, + }, + "action": {"dtype": "float32", "shape": (2,), "names": ["a", "b"]}, +} +DUMMY_DEPTH_KEY = "observation.images.depth" diff --git a/tests/fixtures/dataset_factories.py b/tests/fixtures/dataset_factories.py index 2f4d41ff8..100922f9c 100644 --- a/tests/fixtures/dataset_factories.py +++ b/tests/fixtures/dataset_factories.py @@ -49,6 +49,39 @@ from tests.fixtures.constants import ( ) +def add_frames(dataset: LeRobotDataset, num_frames: int) -> None: + """Append ``num_frames`` synthetic frames to ``dataset``. + + Generates per-feature payloads from ``dataset.meta``: uint16 depth ramps for + keys in ``dataset.meta.depth_keys``, uint8 random noise for video/image keys, + and float32 zeros for everything else. ``DEFAULT_FEATURES`` (timestamp, + frame_index, ...) are auto-populated by ``add_frame`` and skipped here. + """ + video_keys = dataset.meta.video_keys + depth_keys = dataset.meta.depth_keys + # Smooth gradient base reused per (H, W) to keep depth frames cheap to + # encode (HEVC Main 12 hates white noise). + _depth_base_cache: dict[tuple[int, int], np.ndarray] = {} + for i in range(num_frames): + frame: dict = {"task": "test"} + for key, ft in dataset.meta.features.items(): + if key in DEFAULT_FEATURES: + continue + shape = ft["shape"] + if key in depth_keys: + h, w, _ = shape + base = _depth_base_cache.setdefault( + (h, w), + np.linspace(100.0, 10_000.0, h * w, dtype=np.float32).reshape(h, w, 1), + ) + frame[key] = (base + 50.0 * i).clip(0, 65535).astype(np.uint16) + elif key in video_keys: + frame[key] = np.random.randint(0, 256, shape, dtype=np.uint8) + else: + frame[key] = np.zeros(shape, dtype=np.float32) + dataset.add_frame(frame) + + class LeRobotDatasetFactory(Protocol): def __call__(self, *args, **kwargs) -> LeRobotDataset: ... @@ -485,10 +518,14 @@ def lerobot_dataset_factory( hf_dataset: datasets.Dataset | None = None, data_files_size_in_mb: float = DEFAULT_DATA_FILE_SIZE_IN_MB, chunks_size: int = DEFAULT_CHUNK_SIZE, + camera_features: dict | None = None, **kwargs, ) -> LeRobotDataset: # Instantiate objects if info is None: + info_kwargs = {} + if camera_features is not None: + info_kwargs["camera_features"] = camera_features info = info_factory( total_episodes=total_episodes, total_frames=total_frames, @@ -496,6 +533,7 @@ def lerobot_dataset_factory( use_videos=use_videos, data_files_size_in_mb=data_files_size_in_mb, chunks_size=chunks_size, + **info_kwargs, ) if stats is None: stats = stats_factory(features=info.features) diff --git a/tests/scripts/test_edit_dataset_parsing.py b/tests/scripts/test_edit_dataset_parsing.py index c90cffb38..22a3c1be2 100644 --- a/tests/scripts/test_edit_dataset_parsing.py +++ b/tests/scripts/test_edit_dataset_parsing.py @@ -27,6 +27,7 @@ from lerobot.scripts.lerobot_edit_dataset import ( MergeConfig, ModifyTasksConfig, OperationConfig, + ReencodeVideosConfig, RemoveFeatureConfig, SplitConfig, _validate_config, @@ -103,3 +104,47 @@ class TestOperationTypeParsing: ) resolved_name = OperationConfig.get_choice_name(type(cfg.operation)) assert resolved_name == type_name + + +class TestDepthEncoderParsing: + """Test that the depth encoder is exposed and parsed for video operations.""" + + def test_reencode_has_default_depth_encoder(self): + cfg = parse_cfg(["--repo_id", "test/repo", "--operation.type", "reencode_videos"]) + assert isinstance(cfg.operation, ReencodeVideosConfig) + # A depth encoder is configured by default so depth videos are re-encoded too. + assert cfg.operation.depth_encoder is not None + assert hasattr(cfg.operation.depth_encoder, "depth_min") + + def test_reencode_parses_depth_encoder_overrides(self): + cfg = parse_cfg( + [ + "--repo_id", + "test/repo", + "--operation.type", + "reencode_videos", + "--operation.depth_encoder.extra_options", + '{"x265-params": "lossless=1"}', + "--operation.depth_encoder.depth_max", + "12.0", + "--operation.depth_encoder.use_log", + "false", + ] + ) + assert cfg.operation.depth_encoder.extra_options == {"x265-params": "lossless=1"} + assert cfg.operation.depth_encoder.depth_max == 12.0 + assert cfg.operation.depth_encoder.use_log is False + + def test_convert_image_to_video_parses_depth_encoder_overrides(self): + cfg = parse_cfg( + [ + "--repo_id", + "test/repo", + "--operation.type", + "convert_image_to_video", + "--operation.depth_encoder.depth_min", + "0.05", + ] + ) + assert isinstance(cfg.operation, ConvertImageToVideoConfig) + assert cfg.operation.depth_encoder.depth_min == 0.05 diff --git a/tests/utils/test_visualization_utils.py b/tests/utils/test_visualization_utils.py index 63ff76c77..5bd1552db 100644 --- a/tests/utils/test_visualization_utils.py +++ b/tests/utils/test_visualization_utils.py @@ -43,6 +43,11 @@ def mock_rerun(monkeypatch): def __init__(self, arr): self.arr = arr + class DummyDepthImage: + def __init__(self, arr, colormap=None): + self.arr = arr + self.colormap = colormap + def dummy_log(key, obj=None, **kwargs): # Accept either positional `obj` or keyword `entity` and record remaining kwargs. if obj is None and "entity" in kwargs: @@ -55,6 +60,8 @@ def mock_rerun(monkeypatch): __spec__=SimpleNamespace(name="rerun", submodule_search_locations=None), Scalars=DummyScalar, Image=DummyImage, + DepthImage=DummyDepthImage, + components=SimpleNamespace(Colormap=SimpleNamespace(Viridis="viridis")), log=dummy_log, init=lambda *a, **k: None, spawn=lambda *a, **k: None, @@ -225,7 +232,7 @@ def test_log_rerun_data_kwargs_only(mock_rerun): assert temp.value == pytest.approx(10.0) img = _obj_for(calls, "observation.gray") - assert type(img).__name__ == "DummyImage" + assert type(img).__name__ == "DummyDepthImage" # single-channel -> DepthImage assert img.arr.shape == (8, 8, 1) # remains HWC assert _kwargs_for(calls, "observation.gray").get("static", False) is True From a5821a01a20d257ea92099049c446f2f05604460 Mon Sep 17 00:00:00 2001 From: Caroline Pascal Date: Mon, 29 Jun 2026 17:28:06 +0200 Subject: [PATCH 19/38] feat(dependencies): bump rerun-sdk to `<0.34.0` (#3763) * Update upper bound to latest rerun-sdk * chore(updae): update rerun logging to use the latest features * chore(format): formatting code * feat(features names and color): improving features names and display colors when replaying an episode * feat(blueprints): switching to blueprints for backwards (and forward) compatibiltiy * feat(blueprints): switching to blueprints for backwards (and forward) compatibiltiy * feat(grid): Leveraging rerun's automatic grid arangement for improved layout * test(update): update tests * chore(colors): removing unreliable colors * chore(simplification): removing no longer needed reshape * chore(imports): cleaning up imports * fix(claude): claude reviews * chore(dependecies): update rerun ceil version * chore(scripts): recover comments * chore(utils): add guard for blueprint * fix(test): style check * fix(deps): typo bound --------- Signed-off-by: Steven Palma Co-authored-by: ntjohnson1 <24689722+ntjohnson1@users.noreply.github.com> Co-authored-by: Steven Palma Co-authored-by: Steven Palma --- pyproject.toml | 2 +- src/lerobot/scripts/lerobot_dataset_viz.py | 65 ++++++++-- src/lerobot/utils/visualization_utils.py | 69 +++++++++-- tests/utils/test_visualization_utils.py | 137 ++++++++++++++++----- uv.lock | 14 +-- 5 files changed, 221 insertions(+), 66 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 3961ded19..28a8948b7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -124,7 +124,7 @@ hardware = [ "lerobot[deepdiff-dep]", ] viz = [ - "rerun-sdk>=0.24.0,<0.27.0", + "rerun-sdk>=0.24.0,<0.34.0", ] # ── User-facing composite extras (map to CLI scripts) ───── # lerobot-record, lerobot-replay, lerobot-calibrate, lerobot-teleoperate, etc. diff --git a/src/lerobot/scripts/lerobot_dataset_viz.py b/src/lerobot/scripts/lerobot_dataset_viz.py index 21ae1ac9d..22a7208d4 100644 --- a/src/lerobot/scripts/lerobot_dataset_viz.py +++ b/src/lerobot/scripts/lerobot_dataset_viz.py @@ -77,6 +77,21 @@ from lerobot.utils.constants import ACTION, DONE, OBS_STATE, REWARD from lerobot.utils.utils import init_logging +def get_feature_names(dataset: LeRobotDataset, key: str) -> list[str]: + """Return per-dimension names for a feature from the dataset metadata. + + Only flat-list ``names`` metadata is used. Dict-style ``names`` and missing names fall back to ``{key}_{i}`` indices. + """ + feature = dataset.features[key] + dim = feature["shape"][-1] + + names = feature.get("names") + if isinstance(names, list) and len(names) == dim: + return [str(name) for name in names] + + return [f"{key}_{d}" for d in range(dim)] + + def check_chw_float32(frame: torch.Tensor) -> None: """ Check if a frame is a channel-first, float32 tensor. @@ -93,6 +108,31 @@ def to_hwc_uint8_numpy(chw_float32_torch: torch.Tensor) -> np.ndarray: return hwc_uint8_numpy +def build_blueprint_from_dataset(dataset: LeRobotDataset): + """Build a Rerun blueprint laying out camera images and time series for the given dataset. + + Camera images and scalar signals (action, state, reward, done, success) are arranged in a grid. + The per-dimension series names for ``action`` and ``state`` are applied directly + via blueprint overrides. + """ + import rerun as rr + import rerun.blueprint as rrb + + views = [rrb.Spatial2DView(origin=key, name=key) for key in dataset.meta.camera_keys] + + # Style multi-dimensional signals (action, state) with per-dimension names. + for origin, key in ((ACTION, ACTION), ("state", OBS_STATE)): + if key in dataset.features: + names = get_feature_names(dataset, key) + styling = rr.SeriesLines(names=names) + views.append(rrb.TimeSeriesView(origin=origin, name=origin, overrides={origin: styling})) + for key in (DONE, REWARD, "next.success"): + if key in dataset.features: + views.append(rrb.TimeSeriesView(origin=key, name=key)) + + return rrb.Blueprint(rrb.Grid(*views)) + + def to_hwc_uint16_numpy(chw_float32_torch: torch.Tensor) -> np.ndarray: check_chw_float32(chw_float32_torch) hwc_uint16_numpy = chw_float32_torch.round().type(torch.uint16).permute(1, 2, 0).numpy() @@ -137,7 +177,8 @@ def visualize_dataset( import rerun as rr spawn_local_viewer = mode == "local" and not save - rr.init(f"{repo_id}/episode_{episode_index}", spawn=spawn_local_viewer) + blueprint = build_blueprint_from_dataset(dataset) + rr.init(f"{repo_id}/episode_{episode_index}", spawn=spawn_local_viewer, default_blueprint=blueprint) # Manually call python garbage collector after `rr.init` to avoid hanging in a blocking flush # when iterating on a dataloader with `num_workers` > 0 @@ -163,12 +204,13 @@ def visualize_dataset( for batch in tqdm.tqdm(dataloader, total=len(dataloader)): if first_index is None: first_index = batch["index"][0].item() + # iterate over the batch for i in range(len(batch["index"])): rr.set_time("frame_index", sequence=batch["index"][i].item() - first_index) rr.set_time("timestamp", timestamp=batch["timestamp"][i].item()) - # display each camera image + # display each camera image (or depth map) for key in dataset.meta.camera_keys: if key in dataset.meta.depth_keys: depth = to_hwc_uint16_numpy(batch[key][i]) @@ -183,15 +225,13 @@ def visualize_dataset( img_entity = rr.Image(img).compress() if display_compressed_images else rr.Image(img) rr.log(key, entity=img_entity) - # display each dimension of action space (e.g. actuators command) + # display the action space (e.g. actuators command) if ACTION in batch: - for dim_idx, val in enumerate(batch[ACTION][i]): - rr.log(f"{ACTION}/{dim_idx}", rr.Scalars(val.item())) + rr.log(ACTION, rr.Scalars(batch[ACTION][i].numpy())) - # display each dimension of observed state space (e.g. agent position in joint space) + # display the observed state space (e.g. agent position in joint space) if OBS_STATE in batch: - for dim_idx, val in enumerate(batch[OBS_STATE][i]): - rr.log(f"state/{dim_idx}", rr.Scalars(val.item())) + rr.log("state", rr.Scalars(batch[OBS_STATE][i].numpy())) if DONE in batch: rr.log(DONE, rr.Scalars(batch[DONE][i].item())) @@ -202,9 +242,8 @@ def visualize_dataset( if "next.success" in batch: rr.log("next.success", rr.Scalars(batch["next.success"][i].item())) + # save .rrd locally if mode == "local" and save: - # save .rrd locally - output_dir = Path(output_dir) output_dir.mkdir(parents=True, exist_ok=True) repo_id_str = repo_id.replace("/", "_") rrd_path = output_dir / f"{repo_id_str}_episode_{episode_index}.rrd" @@ -212,7 +251,7 @@ def visualize_dataset( return rrd_path elif mode == "distant": - # stop the process from exiting since it is serving the websocket connection + # Keep the process alive while it serves the gRPC/web connection. try: while True: time.sleep(1) @@ -327,12 +366,14 @@ def main(): ) logging.warning("Setting grpc_port to ws_port value.") kwargs["grpc_port"] = kwargs.pop("ws_port") + else: + kwargs.pop("ws_port") # Always remove ws_port from kwargs init_logging() logging.info("Loading dataset") dataset = LeRobotDataset(repo_id, episodes=[args.episode_index], root=root, tolerance_s=tolerance_s) - visualize_dataset(dataset, **vars(args)) + visualize_dataset(dataset, **kwargs) if __name__ == "__main__": diff --git a/src/lerobot/utils/visualization_utils.py b/src/lerobot/utils/visualization_utils.py index e039f7b33..a0f07f0c7 100644 --- a/src/lerobot/utils/visualization_utils.py +++ b/src/lerobot/utils/visualization_utils.py @@ -38,6 +38,8 @@ def init_rerun( require_package("rerun-sdk", extra="viz", import_name="rerun") import rerun as rr + log_rerun_data.blueprint = None # Reset blueprint cache for new session + batch_size = os.getenv("RERUN_FLUSH_NUM_BYTES", "8000") os.environ["RERUN_FLUSH_NUM_BYTES"] = batch_size rr.init(session_name) @@ -63,6 +65,41 @@ def _is_scalar(x): ) +def _build_blueprint(observation_paths: set[str], action_paths: set[str], image_paths: set[str]): + """Build a Rerun blueprint laying out camera images, observation and action scalars in separate views. + + Camera images, observation and action scalars are arranged in a grid. + """ + + # Safe + zero-overhead: `log_rerun_data` already ran the `require_package` guard and imported rerun. + import rerun.blueprint as rrb + + views = [rrb.Spatial2DView(origin=path, name=path) for path in sorted(image_paths)] + + if observation_paths: + views.append(rrb.TimeSeriesView(name="observation", contents=sorted(observation_paths))) + if action_paths: + views.append(rrb.TimeSeriesView(name="action", contents=sorted(action_paths))) + + return rrb.Blueprint(rrb.Grid(*views)) + + +def _ensure_blueprint(observation_paths: set[str], action_paths: set[str], image_paths: set[str]) -> None: + """Build and send the blueprint once, from the first observation and action data.""" + if getattr(log_rerun_data, "blueprint", None) is not None: + return + + if not (observation_paths or action_paths or image_paths): + return + + # Safe + zero-overhead: `log_rerun_data` already ran the `require_package` guard and imported rerun. + import rerun as rr + + blueprint = _build_blueprint(observation_paths, action_paths, image_paths) + log_rerun_data.blueprint = blueprint + rr.send_blueprint(blueprint) + + def log_rerun_data( observation: RobotObservation | None = None, action: RobotAction | None = None, @@ -76,11 +113,15 @@ def log_rerun_data( - Scalars values (floats, ints) are logged as `rr.Scalars`. - 3D NumPy arrays that resemble images (e.g., with 1, 3, or 4 channels first) are transposed from CHW to HWC format, (optionally) compressed to JPEG and logged as `rr.Image` or `rr.EncodedImage`. - - 1D NumPy arrays are logged as a series of individual scalars, with each element indexed. - - Other multi-dimensional arrays are flattened and logged as individual scalars. + - 1D NumPy arrays are logged as a single `rr.Scalars` batch under one entity path, so that every + dimension shares the same view instead of being split across one view per element. + - Multi-dimensional **action** arrays are flattened and logged as a single `rr.Scalars` batch. Keys are automatically namespaced with "observation." or "action." if not already present. + On the first call, a blueprint is built and sent so observation and action scalars get separate + time-series views and each image gets its own spatial view. + Args: observation: An optional dictionary containing observation data to log. action: An optional dictionary containing action data to log. @@ -90,6 +131,10 @@ def log_rerun_data( require_package("rerun-sdk", extra="viz", import_name="rerun") import rerun as rr + observation_paths: set[str] = set() + action_paths: set[str] = set() + image_paths: set[str] = set() + if observation: for k, v in observation.items(): if v is None: @@ -98,20 +143,22 @@ def log_rerun_data( if _is_scalar(v): rr.log(key, rr.Scalars(float(v))) + observation_paths.add(key) elif isinstance(v, np.ndarray): arr = v # Convert CHW -> HWC when needed if arr.ndim == 3 and arr.shape[0] in (1, 3, 4) and arr.shape[-1] not in (1, 3, 4): arr = np.transpose(arr, (1, 2, 0)) if arr.ndim == 1: - for i, vi in enumerate(arr): - rr.log(f"{key}_{i}", rr.Scalars(float(vi))) + rr.log(key, rr.Scalars(arr.astype(float))) + observation_paths.add(key) else: if arr.shape[-1] == 1: img_entity = rr.DepthImage(arr, colormap=rr.components.Colormap.Viridis) else: img_entity = rr.Image(arr).compress() if compress_images else rr.Image(arr) rr.log(key, entity=img_entity, static=True) + image_paths.add(key) if action: for k, v in action.items(): @@ -121,12 +168,10 @@ def log_rerun_data( if _is_scalar(v): rr.log(key, rr.Scalars(float(v))) + action_paths.add(key) elif isinstance(v, np.ndarray): - if v.ndim == 1: - for i, vi in enumerate(v): - rr.log(f"{key}_{i}", rr.Scalars(float(vi))) - else: - # Fall back to flattening higher-dimensional arrays - flat = v.flatten() - for i, vi in enumerate(flat): - rr.log(f"{key}_{i}", rr.Scalars(float(vi))) + # Flatten any (incl. higher-dimensional) array into a single batched Scalars + rr.log(key, rr.Scalars(v.reshape(-1).astype(float))) + action_paths.add(key) + + _ensure_blueprint(observation_paths, action_paths, image_paths) diff --git a/tests/utils/test_visualization_utils.py b/tests/utils/test_visualization_utils.py index 5bd1552db..f62a697cd 100644 --- a/tests/utils/test_visualization_utils.py +++ b/tests/utils/test_visualization_utils.py @@ -30,19 +30,25 @@ from lerobot.utils.constants import OBS_STATE @pytest.fixture def mock_rerun(monkeypatch): """ - Provide a mock `rerun` module so tests don't depend on the real library. - Also reload the module-under-test so it binds to this mock `rr`. + Provide a mock `rerun` module (and `rerun.blueprint` submodule) so tests don't + depend on the real library. Also reload the module-under-test so it binds to + this mock `rr`. """ calls = [] + blueprints = [] class DummyScalar: def __init__(self, value): - self.value = float(value) + # Scalars may be built from a single float or from a 1D array batch. + self.value = value class DummyImage: def __init__(self, arr): self.arr = arr + def compress(self, *a, **k): + return self + class DummyDepthImage: def __init__(self, arr, colormap=None): self.arr = arr @@ -54,6 +60,21 @@ def mock_rerun(monkeypatch): obj = kwargs.pop("entity") calls.append((key, obj, kwargs)) + def dummy_send_blueprint(blueprint, *a, **k): + blueprints.append(blueprint) + + # Mock the `rerun.blueprint` submodule used to build the layout. + dummy_rrb = SimpleNamespace( + Spatial2DView=lambda origin=None, name=None: SimpleNamespace( + kind="Spatial2DView", origin=origin, name=name + ), + TimeSeriesView=lambda name=None, contents=None: SimpleNamespace( + kind="TimeSeriesView", name=name, contents=contents + ), + Grid=lambda *views: SimpleNamespace(kind="Grid", views=list(views)), + Blueprint=lambda root: SimpleNamespace(kind="Blueprint", root=root), + ) + dummy_rr = SimpleNamespace( __name__="rerun", __package__="rerun", @@ -63,20 +84,23 @@ def mock_rerun(monkeypatch): DepthImage=DummyDepthImage, components=SimpleNamespace(Colormap=SimpleNamespace(Viridis="viridis")), log=dummy_log, + send_blueprint=dummy_send_blueprint, init=lambda *a, **k: None, spawn=lambda *a, **k: None, + blueprint=dummy_rrb, ) - # Inject fake module into sys.modules + # Inject fake modules into sys.modules (both `rerun` and `rerun.blueprint`). monkeypatch.setitem(sys.modules, "rerun", dummy_rr) + monkeypatch.setitem(sys.modules, "rerun.blueprint", dummy_rrb) # Now import and reload the module under test, to bind to our rerun mock import lerobot.utils.visualization_utils as vu importlib.reload(vu) - # Expose both the reloaded module and the call recorder - yield vu, calls + # Expose the reloaded module, the call recorder and the captured blueprints + yield vu, calls, blueprints def _keys(calls): @@ -99,8 +123,13 @@ def _kwargs_for(calls, key): raise KeyError(f"Key {key} not found in calls: {calls}") +def _views_by_kind(blueprint, kind): + """Return the views of a given kind from the (single) blueprint's grid.""" + return [v for v in blueprint.root.views if v.kind == kind] + + def test_log_rerun_data_envtransition_scalars_and_image(mock_rerun): - vu, calls = mock_rerun + vu, calls, blueprints = mock_rerun # Build EnvTransition dict obs = { @@ -110,7 +139,7 @@ def test_log_rerun_data_envtransition_scalars_and_image(mock_rerun): } act = { "action.throttle": 0.7, - # 1D array should log individual Scalars with suffix _i + # 1D array should be logged as a single Scalars batch under one entity path "action.vector": np.array([1.0, 2.0], dtype=np.float32), } transition = { @@ -127,31 +156,28 @@ def test_log_rerun_data_envtransition_scalars_and_image(mock_rerun): # - observation.state.temperature -> Scalars # - observation.camera -> Image (HWC) with static=True # - action.throttle -> Scalars - # - action.vector_0, action.vector_1 -> Scalars + # - action.vector -> single Scalars batch (no per-element suffix) expected_keys = { f"{OBS_STATE}.temperature", "observation.camera", "action.throttle", - "action.vector_0", - "action.vector_1", + "action.vector", } assert set(_keys(calls)) == expected_keys # Check scalar types and values temp_obj = _obj_for(calls, f"{OBS_STATE}.temperature") assert type(temp_obj).__name__ == "DummyScalar" - assert temp_obj.value == pytest.approx(25.0) + assert float(temp_obj.value) == pytest.approx(25.0) throttle_obj = _obj_for(calls, "action.throttle") assert type(throttle_obj).__name__ == "DummyScalar" - assert throttle_obj.value == pytest.approx(0.7) + assert float(throttle_obj.value) == pytest.approx(0.7) - v0 = _obj_for(calls, "action.vector_0") - v1 = _obj_for(calls, "action.vector_1") - assert type(v0).__name__ == "DummyScalar" - assert type(v1).__name__ == "DummyScalar" - assert v0.value == pytest.approx(1.0) - assert v1.value == pytest.approx(2.0) + # 1D vector logged as a single batched Scalars under one entity path + vec = _obj_for(calls, "action.vector") + assert type(vec).__name__ == "DummyScalar" + np.testing.assert_allclose(np.asarray(vec.value), [1.0, 2.0]) # Check image handling: CHW -> HWC img_obj = _obj_for(calls, "observation.camera") @@ -159,9 +185,24 @@ def test_log_rerun_data_envtransition_scalars_and_image(mock_rerun): assert img_obj.arr.shape == (10, 20, 3) # transposed assert _kwargs_for(calls, "observation.camera").get("static", False) is True # static=True for images + # A blueprint should have been built and sent exactly once, and cached on the function. + assert len(blueprints) == 1 + assert vu.log_rerun_data.blueprint is blueprints[0] + + bp = blueprints[0] + # One spatial view per image path + spatial_views = _views_by_kind(bp, "Spatial2DView") + assert {v.origin for v in spatial_views} == {"observation.camera"} + + # One time-series view each for observation and action scalars + ts_views = {v.name: v for v in _views_by_kind(bp, "TimeSeriesView")} + assert set(ts_views) == {"observation", "action"} + assert ts_views["observation"].contents == [f"{OBS_STATE}.temperature"] + assert ts_views["action"].contents == ["action.throttle", "action.vector"] + def test_log_rerun_data_plain_list_ordering_and_prefixes(mock_rerun): - vu, calls = mock_rerun + vu, calls, blueprints = mock_rerun # First dict without prefixes treated as observation # Second dict without prefixes treated as action @@ -180,14 +221,12 @@ def test_log_rerun_data_plain_list_ordering_and_prefixes(mock_rerun): # First dict was treated as observation, second as action vu.log_rerun_data(observation=obs_plain, action=act_plain) - # Expected keys with auto-prefixes + # Expected keys with auto-prefixes. The 1D vector is a single batched Scalars. expected = { "observation.temp", "observation.img", "action.throttle", - "action.vec_0", - "action.vec_1", - "action.vec_2", + "action.vec", } logged = set(_keys(calls)) assert logged == expected @@ -195,11 +234,11 @@ def test_log_rerun_data_plain_list_ordering_and_prefixes(mock_rerun): # Scalars t = _obj_for(calls, "observation.temp") assert type(t).__name__ == "DummyScalar" - assert t.value == pytest.approx(1.5) + assert float(t.value) == pytest.approx(1.5) throttle = _obj_for(calls, "action.throttle") assert type(throttle).__name__ == "DummyScalar" - assert throttle.value == pytest.approx(0.3) + assert float(throttle.value) == pytest.approx(0.3) # Image stays HWC img = _obj_for(calls, "observation.img") @@ -207,15 +246,23 @@ def test_log_rerun_data_plain_list_ordering_and_prefixes(mock_rerun): assert img.arr.shape == (5, 6, 3) assert _kwargs_for(calls, "observation.img").get("static", False) is True - # Vectors - for i, val in enumerate([9, 8, 7]): - o = _obj_for(calls, f"action.vec_{i}") - assert type(o).__name__ == "DummyScalar" - assert o.value == pytest.approx(val) + # Vector logged as a single batched Scalars under one entity path + vec = _obj_for(calls, "action.vec") + assert type(vec).__name__ == "DummyScalar" + np.testing.assert_allclose(np.asarray(vec.value), [9, 8, 7]) + + # Blueprint sent once with the expected view layout + assert len(blueprints) == 1 + bp = blueprints[0] + spatial_views = _views_by_kind(bp, "Spatial2DView") + assert {v.origin for v in spatial_views} == {"observation.img"} + ts_views = {v.name: v for v in _views_by_kind(bp, "TimeSeriesView")} + assert ts_views["observation"].contents == ["observation.temp"] + assert ts_views["action"].contents == ["action.throttle", "action.vec"] def test_log_rerun_data_kwargs_only(mock_rerun): - vu, calls = mock_rerun + vu, calls, blueprints = mock_rerun vu.log_rerun_data( observation={"observation.temp": 10.0, "observation.gray": np.zeros((8, 8, 1), dtype=np.uint8)}, @@ -229,7 +276,7 @@ def test_log_rerun_data_kwargs_only(mock_rerun): temp = _obj_for(calls, "observation.temp") assert type(temp).__name__ == "DummyScalar" - assert temp.value == pytest.approx(10.0) + assert float(temp.value) == pytest.approx(10.0) img = _obj_for(calls, "observation.gray") assert type(img).__name__ == "DummyDepthImage" # single-channel -> DepthImage @@ -238,4 +285,26 @@ def test_log_rerun_data_kwargs_only(mock_rerun): a = _obj_for(calls, "action.a") assert type(a).__name__ == "DummyScalar" - assert a.value == pytest.approx(1.0) + assert float(a.value) == pytest.approx(1.0) + + # Blueprint sent once, with a spatial view for the image and time-series views for scalars + assert len(blueprints) == 1 + bp = blueprints[0] + assert {v.origin for v in _views_by_kind(bp, "Spatial2DView")} == {"observation.gray"} + ts_views = {v.name: v for v in _views_by_kind(bp, "TimeSeriesView")} + assert ts_views["observation"].contents == ["observation.temp"] + assert ts_views["action"].contents == ["action.a"] + + +def test_log_rerun_data_blueprint_sent_only_once(mock_rerun): + """The blueprint is built from the first call and not resent on subsequent calls.""" + vu, calls, blueprints = mock_rerun + + vu.log_rerun_data(observation={"temp": 1.0}, action={"a": 2.0}) + assert len(blueprints) == 1 + first_blueprint = vu.log_rerun_data.blueprint + + vu.log_rerun_data(observation={"temp": 3.0}, action={"a": 4.0}) + # Still only one blueprint, and the cached one is unchanged. + assert len(blueprints) == 1 + assert vu.log_rerun_data.blueprint is first_blueprint diff --git a/uv.lock b/uv.lock index e24b1d884..5a76fcbf8 100644 --- a/uv.lock +++ b/uv.lock @@ -3395,7 +3395,7 @@ requires-dist = [ { name = "qwen-vl-utils", marker = "extra == 'qwen-vl-utils-dep'", specifier = ">=0.0.11,<0.1.0" }, { name = "reachy2-sdk", marker = "extra == 'reachy2'", specifier = ">=1.0.15,<1.1.0" }, { name = "requests", specifier = ">=2.32.0,<3.0.0" }, - { name = "rerun-sdk", marker = "extra == 'viz'", specifier = ">=0.24.0,<0.27.0" }, + { name = "rerun-sdk", marker = "extra == 'viz'", specifier = ">=0.24.0,<0.34.0" }, { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.14.1" }, { name = "safetensors", specifier = ">=0.4.3,<1.0.0" }, { name = "scikit-image", marker = "extra == 'video-benchmark'", specifier = ">=0.23.2,<0.26.0" }, @@ -5803,21 +5803,21 @@ wheels = [ [[package]] name = "rerun-sdk" -version = "0.26.2" +version = "0.33.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "attrs" }, { name = "numpy" }, { name = "pillow" }, + { name = "psutil" }, { name = "pyarrow" }, { name = "typing-extensions" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/4b/4a/767c20e1529d74d9be5b5e55c6c26b63a6918ef3c1709fc422d08a460114/rerun_sdk-0.26.2-cp39-abi3-macosx_10_12_x86_64.whl", hash = "sha256:3d4151c9a3484e112b53d1df90c8fa07397dc7b8bfbb420f09e011eff20f1ef2", size = 93349439, upload-time = "2025-10-27T11:34:10.745Z" }, - { url = "https://files.pythonhosted.org/packages/2b/3d/d8dd0af9c287a85d51ec99d69406cc4b94a9feb1d6f192d3bbcaac9f0b81/rerun_sdk-0.26.2-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:03977d2aba4966d9a70b682eca196123fda11408fecd733441ede9916c6341e2", size = 86323042, upload-time = "2025-10-27T11:34:17.995Z" }, - { url = "https://files.pythonhosted.org/packages/13/29/53d8d98799ab32418fd4ba6834d6a5749c31f56160d3c87f52a7219887e9/rerun_sdk-0.26.2-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:b6128c3c4f014cae5be18e4d37657c5932d1bcdb2ce5e9d4b488a6eed47f7437", size = 92677274, upload-time = "2025-10-27T11:34:22.601Z" }, - { url = "https://files.pythonhosted.org/packages/f5/86/0b9c8f56398b4fc85f8e99279907c258413a297e5603f8f2537fe5806e51/rerun_sdk-0.26.2-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:a6f97b60aaa7d4e8c6124a3f6b97ce9dbd09520050955f0e0bdacb72b0eb106a", size = 98768129, upload-time = "2025-10-27T11:34:27.36Z" }, - { url = "https://files.pythonhosted.org/packages/be/e7/99fc91c0f99f69d7d43e1db0a6f6cb8273ffc02111539bfc1fee43749bad/rerun_sdk-0.26.2-cp39-abi3-win_amd64.whl", hash = "sha256:a493ad6c8357022cba2ca6f8954a81d0faf984b0b22154eb1d976bfc7649df63", size = 84267089, upload-time = "2025-10-27T11:34:32.023Z" }, + { url = "https://files.pythonhosted.org/packages/16/07/380198590b194f1c17052d672865aa4a56e606eae47665f66edfb391999d/rerun_sdk-0.33.1-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:c0115b710289022587bd2e9ecb715f98d1e87dd7d8ac48e053324131d4addc89", size = 125706253, upload-time = "2026-06-22T09:04:18.974Z" }, + { url = "https://files.pythonhosted.org/packages/f8/eb/6741bbf6868175ab126aff58d372066241c6cd2fc1c4f82ed64069728e73/rerun_sdk-0.33.1-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:14451d31bc7bd0f7c6bcc9c1213ed679ab81b65ecd8b36eae99f738219897dc5", size = 135278374, upload-time = "2026-06-22T09:04:26.768Z" }, + { url = "https://files.pythonhosted.org/packages/04/28/fd3b832652900fe3415739e7411c8af8af4c44e9e1a9d55d79e37f7f9094/rerun_sdk-0.33.1-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:0f87da7a270614074aca37846350f9e257f65081345474748f578c7da64fdeba", size = 139565470, upload-time = "2026-06-22T09:04:34.941Z" }, + { url = "https://files.pythonhosted.org/packages/19/41/339920f5a6734054c07bcae543365a7ef3368ceee3eb67906e2e38bd1d67/rerun_sdk-0.33.1-cp310-abi3-win_amd64.whl", hash = "sha256:b2c2af67f3c2a85b282669d97b52596593fcfdd19bf57c423f18827c837a6e49", size = 120411717, upload-time = "2026-06-22T09:04:42.643Z" }, ] [[package]] From 5ac3b49a5fd25d9e570ea1de4e6a81c77d603bd3 Mon Sep 17 00:00:00 2001 From: Nicolas Rabault Date: Mon, 29 Jun 2026 17:59:33 +0200 Subject: [PATCH 20/38] feat(train): run training remotely on HF Jobs via --job.target (#3856) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(train): add JobConfig group, save_checkpoint_to_hub flag, Hub checkpoint helper Introduce a JobConfig draccus group on TrainPipelineConfig (--job.target/image/ timeout/detach/tags) whose is_remote property gates remote dispatch, plus a save_checkpoint_to_hub flag and validation. Add push_checkpoint_to_hub(), which uploads a saved checkpoint directory to the model repo under checkpoints// and creates the repo idempotently (private propagates from policy.private). * feat(train): run training remotely on HF Jobs via --job.target When --job.target names a GPU flavor, train() dispatches to lerobot.jobs.submit_to_hf instead of training locally: it authenticates, ensures the dataset is on the Hub (pushing a local-only one privately), serializes a pod-compatible train_config.json (strips client-only fields, points at the model repo), submits via HfApi.run_job with HF_TOKEN/WANDB_API_KEY secrets, then streams logs and finishes when the model is pushed. Wires push_checkpoint_to_hub into the training loop behind save_checkpoint_to_hub, and tags jobs/datasets/model with 'lerobot' + --job.tags. * docs(train): document remote training on HF Jobs * test(train): skip remote-dispatch tests without the dataset extra The module imports lerobot.scripts.lerobot_train, which eagerly pulls in lerobot.datasets (dataset extra). The base fast-test CI tier runs without that extra, so collection failed there. Guard with pytest.importorskip, matching the existing tests/scripts dataset-extra tests. * refactor(jobs): hoist huggingface_hub imports to module level in hf.py huggingface_hub is a core dependency, so the per-function dynamic imports had no lazy-loading rationale. Move them to a single module-level import and update test monkeypatch targets to lerobot.jobs.hf.* accordingly. * refactor(jobs): build remote config dict via cfg.to_dict() TrainPipelineConfig.to_dict() already returns the canonical draccus encoding, so the StringIO + draccus.dump + json.loads round-trip was redundant. Use it directly and drop the now-unused io/draccus imports. * refactor(train): use module-level HfApi import in push_checkpoint_to_hub huggingface_hub is a core dependency; the in-function import was unnecessary. Move HfApi to a module-level import and point the test monkeypatches at lerobot.common.train_utils.HfApi. * refactor(configs): export JobConfig from the configs package Re-export JobConfig in lerobot/configs/__init__.py so external callers import it as `from lerobot.configs import JobConfig`, matching the other config classes. Adapt the train script and test imports. * refactor(jobs): check dataset presence with api.repo_exists Replace the dataset_info try/except RepositoryNotFoundError dance with a direct api.repo_exists(repo_id, repo_type="dataset") call, dropping the httpx/RepositoryNotFoundError test scaffolding. * chore(jobs): annotate ensure_dataset_available api param as HfApi Add the missing HfApi type hint via a TYPE_CHECKING import. * refactor(jobs): use HF_LEROBOT_HOME constant for the local cache root Resolve the local dataset cache via lerobot.utils.constants.HF_LEROBOT_HOME instead of re-reading the env var by hand, dropping the os/Path imports. Tests now patch the imported constant and assert on a stable message substring (the previous "neither" match only passed by accident, matching the test name embedded in the pytest tmp_path). * chore(jobs): guard LeRobotDataset import with require_package Surface a clear "install lerobot[dataset]" error if the datasets extra is missing, instead of a raw ImportError, before pushing a local dataset. * docs(configs): clarify the is_remote_target/is_remote split Add a comment explaining why JobConfig keeps both the staticmethod (tests a raw target string from argv before a config exists) and the property (accessor for an existing config instance). * docs(train): note how to pin a pushed model version for inference Document --policy.pretrained_revision alongside --policy.path so a specific Hub-pushed checkpoint (once --save_checkpoint_to_hub has committed several) can be selected for inference. * test(jobs): skip dataset import guard in base-deps test The fast test env installs base deps only, so require_package('datasets') raised ImportError before the mocked lerobot.datasets import was reached. Monkeypatch the guard to a no-op so the unit test exercises the upload logic. * fix(jobs): address claude review findings on remote training Resolve the claude[bot] review on #3856: - Reject reward-model training under --job.target with a clear error instead of crashing on a None policy inside build_remote_config_file. - Support --policy.path remote runs: validate() no longer requires repo_id for remote runs (it is auto-generated in submit_to_hf), and repo_id/push_to_hub are now set after validate() resolves the policy. - Narrow the bare `except Exception` in _tail_logs/_poll_until_done to (OSError, httpx.HTTPError) so programming errors surface instead of being silently retried or counted as job failures. - Install the SIGINT detach handler only on the main thread. - Generate model repo timestamps in UTC. * docs(jobs): document the model-pushed marker contract and orphaned repos Follow-up to the claude[bot] review on #3856 (non-blocking observations): - Cross-reference the "Model pushed to " log line between its producer (PreTrainedPolicy.push_model_to_hub) and the remote-run consumer in submit_to_hf, noting the contract is an early-finish optimization that falls back to status polling if it drifts. - Note in the HF Jobs guide that a failed remote run leaves its model repo on the Hub (it is not auto-deleted) and how to remove it. * feat(train): tag each pushed checkpoint with its step Address review feedback on #3856: pushing a checkpoint to the Hub now also creates a tag named after the checkpoint step, so a checkpoint can be recovered with --policy.pretrained_revision= instead of having to look up its commit sha. * fix(jobs): hoist ensure_dataset_available to a module-level import Addresses Caroline's review comment on PR #3856: the local import of ensure_dataset_available inside submit_to_hf was vestigial. dataset.py does not import hf.py, so there is no circular-import risk and no extra load cost (its heavy deps stay lazy), so make it a top-level import. * refactor(configs): untangle config_path/resume resolution in validate() Split the re-parse HACK block in TrainPipelineConfig.validate() into focused helpers (_resolve_pretrained_from_cli, _resolve_resume_checkpoint) that handle the policy path, reward-model path, and resume config_path as separate, readable units. Behavior-preserving. * feat(train): resume training from a Hub checkpoint Allow --config_path to be a Hub repo id when resuming, not only a local path. The latest checkpoint under checkpoints// is downloaded into a fresh local run dir and resumed from there (optimizer, scheduler, RNG and data order restored as for a local resume). TrainPipelineConfig.from_pretrained falls back to the latest checkpoint's train_config.json when a repo has no root config (an interrupted run that only pushed checkpoints). The download is skipped when dispatching remotely so the executor (local machine or HF Jobs pod) performs it. - add find_latest_hub_checkpoint (utils/hub) and resolve_resume_checkpoint (common/train_utils), the symmetric download counterpart to push_checkpoint_to_hub - unit tests for both helpers and the from_pretrained fallback * feat(jobs): resume a run on HF Jobs from a checkpoint When --resume is set with a remote --job.target, submit_to_hf resumes from the checkpoint repo instead of staging a fresh config. A Hub config_path is resumed in place (its checkpoint config already targets that repo); a local config_path has its checkpoint uploaded to a new private repo first and the run is forced to push back to it. The pod command carries --job.target=local so the checkpoint's saved job.target can't make the pod re-dispatch itself, and the user's CLI overrides are forwarded so a remote resume matches the same local command. ensure_dataset_available is hoisted before the resume/fresh branch since it applies to both. * docs(train): document resuming from a Hub checkpoint, locally and on jobs Show that --config_path accepts a Hub repo id for --resume, and that adding --job.target resumes on HF Jobs (uploading a local checkpoint/dataset first). * fix(jobs): default remote job timeout to 2d instead of the platform default HF Jobs applies its own short 30-minute timeout when none is sent, which silently kills long training runs. Pass an explicit, generous 2d cap by default; users can still override --job.timeout to fail fast or extend it. * fix(jobs): drop --dataset.root on resume + restore keyboard-control docs Address the latest Claude review on #3856: - _build_resume_job no longer forwards --dataset.root to the pod (a host-local path it can't read); the fresh-run path already nulls it in build_remote_config_file, so this makes resume consistent. Add a unit test for _pod_forwarded_args covering the drop in both flag forms. - Restore the display-independent keyboard-control docs (n/r/q letter equivalents + X11/Wayland/headless Tip) in il_robots.mdx that this branch was stale on relative to main (#3875). * fix(jobs): handle str-typed job stage from huggingface_hub inspect_job's status.stage is an enum (with .value) in some huggingface_hub versions and a plain str in others. The poller assumed the enum shape, raising "'str' object has no attribute 'value'" on resume for users on the str-returning version. Read it via getattr(..., "value", ...) so both shapes work, and parametrize the poll test over enum and str stages so the str case is actually exercised (the old mock only ever simulated the enum). * refactor(jobs): use relative import for ensure_dataset_available * refactor(train): hoist submit_to_hf import to module top The `from lerobot.jobs import submit_to_hf` was a function-local import in train(); it pulls no heavy/optional deps and has no circular-import risk, so move it to the top-level import block. * refactor(train): hoist _remote_target_in_argv imports to module top Move `import sys` and `from lerobot.configs import JobConfig` out of the function body and into the top-level import block. * refactor(utils): use relative import for sibling constants in hub.py `from lerobot.utils.constants import CHECKPOINTS_DIR` was the odd one out in utils/ — sibling modules there are imported relatively (.constants, .errors, .utils, ...). Match that convention. * refactor(jobs): hoist LeRobotDataset import, guard dataset extra at package init Move the `from lerobot.datasets import LeRobotDataset` import to the top of dataset.py and relocate the `require_package("datasets", extra="dataset")` guard to the jobs package __init__, per review feedback. * test(jobs): skip test_hf if datasets extra is missing lerobot.configs.train pulls in datasets at import time, so the module fails to collect without lerobot[dataset]. Guard with importorskip, matching the convention in tests/training/test_multi_gpu.py. * test(jobs): skip test_dataset if datasets extra is missing tests/jobs/test_dataset.py imports lerobot.jobs.dataset, which triggers the require_package("datasets") guard in lerobot/jobs/__init__.py at import time. Without lerobot[dataset] the module fails to collect in the base CI tier. Guard with importorskip, same as test_hf.py. --- AGENT_GUIDE.md | 2 +- docs/source/cheat-sheet.mdx | 8 + docs/source/hardware_guide.mdx | 1 + docs/source/il_robots.mdx | 57 ++- src/lerobot/common/train_utils.py | 60 +++ src/lerobot/configs/__init__.py | 3 +- src/lerobot/configs/default.py | 32 ++ src/lerobot/configs/train.py | 143 ++++-- src/lerobot/jobs/__init__.py | 23 + src/lerobot/jobs/dataset.py | 53 +++ src/lerobot/jobs/hf.py | 425 +++++++++++++++++ src/lerobot/policies/pretrained.py | 3 + src/lerobot/scripts/lerobot_train.py | 31 +- src/lerobot/utils/hub.py | 24 + tests/configs/test_resume_from_hub.py | 68 +++ tests/jobs/__init__.py | 0 tests/jobs/conftest.py | 17 + tests/jobs/test_dataset.py | 66 +++ tests/jobs/test_hf.py | 493 ++++++++++++++++++++ tests/jobs/test_job_config.py | 64 +++ tests/scripts/test_train_remote_dispatch.py | 67 +++ tests/utils/test_hub.py | 54 +++ tests/utils/test_train_utils.py | 74 ++- 23 files changed, 1720 insertions(+), 48 deletions(-) create mode 100644 src/lerobot/jobs/__init__.py create mode 100644 src/lerobot/jobs/dataset.py create mode 100644 src/lerobot/jobs/hf.py create mode 100644 tests/configs/test_resume_from_hub.py create mode 100644 tests/jobs/__init__.py create mode 100644 tests/jobs/conftest.py create mode 100644 tests/jobs/test_dataset.py create mode 100644 tests/jobs/test_hf.py create mode 100644 tests/jobs/test_job_config.py create mode 100644 tests/scripts/test_train_remote_dispatch.py create mode 100644 tests/utils/test_hub.py diff --git a/AGENT_GUIDE.md b/AGENT_GUIDE.md index 57a33fdba..03b270dce 100644 --- a/AGENT_GUIDE.md +++ b/AGENT_GUIDE.md @@ -138,7 +138,7 @@ lerobot-replay --robot.type=so101_follower --robot.port= --robot. --dataset.repo_id=${HF_USER}/my_task --dataset.episode=0 ``` -**4.9 Train** (default: ACT — fastest, lowest memory). Apple silicon: `--policy.device=mps`. See §6/§7 for policy and duration. +**4.9 Train** (default: ACT — fastest, lowest memory). Apple silicon: `--policy.device=mps`. No local GPU? Add `--job.target=` (e.g. `a10g-small`, list them with `hf jobs hardware`) to run on Hugging Face Jobs instead. See §6/§7 for policy and duration. ```bash lerobot-train \ diff --git a/docs/source/cheat-sheet.mdx b/docs/source/cheat-sheet.mdx index 45952c5b3..0531c95bf 100644 --- a/docs/source/cheat-sheet.mdx +++ b/docs/source/cheat-sheet.mdx @@ -150,6 +150,14 @@ lerobot-train \ --steps=20000 ``` +No local GPU? Add `--job.target=` (e.g. `a10g-small`) to either command and `lerobot-train` runs it on [Hugging Face Jobs](https://huggingface.co/docs/hub/jobs) instead — it uploads a local-only dataset for you and pushes the trained model. List flavors with `hf jobs hardware`. + +To resume, point `--config_path` at a checkpoint and add `--resume=true`. It accepts a local path or a Hub repo id (the latest checkpoint is fetched), and works locally or on a job by adding `--job.target=`: + +```bash +lerobot-train --config_path=${HF_USER}/policy_test --resume=true --job.target=a10g-small +``` + ### Inference Inference means running the trained policy/model on a robot. For that we use `lerobot-rollout`. You will need to provide a path to your policy. It can be a local path or a path to Hugging Face for example "lerobot/folding_latest". Your cameras configuration needs to match what was used when collecting the dataset. Duration is in seconds if unspecified, it will run forever. diff --git a/docs/source/hardware_guide.mdx b/docs/source/hardware_guide.mdx index 0998344ec..5f236d3e8 100644 --- a/docs/source/hardware_guide.mdx +++ b/docs/source/hardware_guide.mdx @@ -96,3 +96,4 @@ Notes: - The leading `nvidia-smi` is a quick sanity check that CUDA is visible inside the container — useful to fail fast if the flavor or driver mismatched. - The default Job timeout is 30 minutes; pass `--timeout 4h` (or longer) for real training. - `--flavor` maps onto the table above: `t4-small`/`t4-medium` (T4, ACT only), `l4x1`/`l4x4` (L4 24 GB), `a10g-small/large/largex2/largex4` (A10G 24 GB scaled out), `a100-large` (A100). For the current full catalogue + pricing see [https://huggingface.co/docs/hub/jobs](https://huggingface.co/docs/hub/jobs). +- Prefer not to write the `hf jobs run` wrapper yourself? `lerobot-train` can submit the job for you: just add `--job.target=` to a normal training command and it handles dataset upload, log streaming, and the final model push. See the [imitation-learning training guide](./il_robots). diff --git a/docs/source/il_robots.mdx b/docs/source/il_robots.mdx index 0f14bd133..178db13bb 100644 --- a/docs/source/il_robots.mdx +++ b/docs/source/il_robots.mdx @@ -514,6 +514,12 @@ lerobot-train \ --resume=true ``` +`--config_path` also accepts a **Hub repo id**: if a run pushed its checkpoints to the Hub (with `--save_checkpoint_to_hub=true`), you can resume straight from the repo — its latest checkpoint is downloaded and training continues, restoring the optimizer, scheduler, step counter and data order: + +```bash +lerobot-train --config_path=${HF_USER}/my_policy --resume=true +``` + If you do not want to push your model to the hub after training use `--policy.push_to_hub=false`. Additionally you can provide extra `tags` or specify a `license` for your model or make the model repo `private` by adding this: `--policy.private=true --policy.tags=\[ppo,rl\] --policy.license=mit` @@ -526,7 +532,9 @@ If your local computer doesn't have a powerful GPU you could utilize Google Cola Hugging Face jobs let's you easily select hardware and run the training in the cloud. So if you don't have a powerful GPU or you need more VRAM or just want to train a model much faster use HF Jobs! It's pay as you go and you simply pay for each second of use, you can see the pricing and additional information [here](https://huggingface.co/docs/hub/jobs). -To run the training use this command: +> **Tip:** if you just want to launch a standard training run, you can skip building the command below and use the integrated **Train on HF Jobs via `--job.target`** flow described further down — `lerobot-train` then submits the job, uploads a local-only dataset for you, and streams the logs. + +To run the training manually use this command: @@ -599,6 +607,51 @@ Once the training is started you can go to [Jobs](https://huggingface.co/setting After training the model will be pushed to hub and you can use it as any other model with LeRobot. +#### Train on HF Jobs via `--job.target` (integrated CLI) + +`lerobot-train` runs locally by default. To run on a HuggingFace GPU without constructing the Docker command yourself, pass `--job.target` with a hardware flavor name: + +```bash +lerobot-train \ + --dataset.repo_id=${HF_USER}/so101_test \ + --policy.type=act \ + --policy.repo_id=${HF_USER}/my_policy \ + --job.target=a10g-small +``` + +List available flavors and pricing with `hf jobs hardware`. The run streams its logs to your terminal; press Ctrl-C to detach (the job keeps running in the cloud). Re-attach or cancel with: + +```bash +hf jobs logs +hf jobs cancel +``` + +If your dataset exists only locally (not yet on the Hub), it is automatically pushed to a **private** Hub repo so the job can download it by `repo_id` (nothing is made public). The trained model is pushed to the model repo at the end of the run. To also push every intermediate checkpoint to the Hub as it is saved (so you can monitor progress mid-run), add `--save_checkpoint_to_hub=true` — this requires a runtime image that includes this feature. + +Every job (and any dataset pushed by the run) is tagged `lerobot` so it's easy to find on the Hub. Add your own with `--job.tags '["my-tag"]'`. + +By default the job is capped at `2d` (48h) of wall-clock. Override it with an HF Jobs duration string, e.g. `--job.timeout=4h` to fail faster or `--job.timeout=7d` for a longer run. + +> **Note:** the model repo is created up front (it holds the staged training config the job runs from). If a run fails before the model is pushed, that repo is left on the Hub so you can inspect it — it is not deleted automatically, so repeated failures can leave empty repos behind. Remove one with `hf repo delete `. + +**Prerequisites:** run `hf auth login` before submitting. For Weights & Biases integration, run `wandb login` or set `WANDB_API_KEY` on your machine — the key is forwarded to the job automatically. + +**Resuming on a job.** Adding `--job.target` to a resume command runs the resume in the cloud — the same command works locally or remotely. The checkpoint repo is the source of truth, and new checkpoints continue the lineage in the same repo: + +```bash +# resume a Hub run on a job (its checkpoints are already on the Hub) +lerobot-train --config_path=${HF_USER}/my_policy --resume=true --job.target=a10g-small + +# resume a LOCAL run on a job — the checkpoint is uploaded to a private Hub repo first, +# then the job resumes from it (a local-only dataset is uploaded the same way) +lerobot-train \ + --config_path=outputs/train/act_so101_test/checkpoints/last/pretrained_model/train_config.json \ + --resume=true \ + --job.target=a10g-small +``` + +Job settings come from the current command, so override `--job.target`, `--job.timeout`, etc. as needed; for the resumed run to itself be resumable later, keep `--save_checkpoint_to_hub=true`. + #### Upload policy checkpoints Once training is done, upload the latest checkpoint with: @@ -620,6 +673,8 @@ hf upload ${HF_USER}/act_so101_test${CKPT} \ Use `lerobot-rollout` to deploy a trained policy on your robot. You can choose different strategies depending on your needs: +The examples below load the model from `--policy.path`. To pin a specific pushed version — useful once `--save_checkpoint_to_hub=true` has committed several checkpoints — add `--policy.pretrained_revision` with a commit hash, branch, or tag. Each pushed checkpoint is tagged with its step (e.g. `--policy.pretrained_revision=010000`), so you can recover a checkpoint by step without looking up its commit sha. + ```bash diff --git a/src/lerobot/common/train_utils.py b/src/lerobot/common/train_utils.py index 5ae593bb8..b26196f14 100644 --- a/src/lerobot/common/train_utils.py +++ b/src/lerobot/common/train_utils.py @@ -15,6 +15,7 @@ # limitations under the License. from pathlib import Path +from huggingface_hub import HfApi, snapshot_download from torch.optim import Optimizer from torch.optim.lr_scheduler import LRScheduler @@ -35,6 +36,7 @@ from lerobot.utils.constants import ( TRAINING_STATE_DIR, TRAINING_STEP, ) +from lerobot.utils.hub import find_latest_hub_checkpoint from lerobot.utils.io_utils import load_json, write_json from lerobot.utils.random_utils import load_rng_state, save_rng_state @@ -283,3 +285,61 @@ def load_fsdp_optimizer_state(model, optimizer, checkpoint_dir: Path) -> None: with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, state_cfg, optim_cfg): sharded_osd = FSDP.optim_state_dict_to_load(model=model, optim=optimizer, optim_state_dict=full_osd) optimizer.load_state_dict(sharded_osd) + + +def push_checkpoint_to_hub( + checkpoint_dir: Path, + repo_id: str, + *, + private: bool | None = None, +) -> None: + """Upload a saved checkpoint directory to the Hub under checkpoints//. + + Called once per save step when save_checkpoint_to_hub is enabled, so a + timed-out or crashed run still leaves recoverable checkpoints on the Hub. + The model repo is created idempotently, and the commit is tagged with the + checkpoint step so a checkpoint can be recovered with + --policy.pretrained_revision= instead of a commit sha. + """ + api = HfApi() + api.create_repo(repo_id=repo_id, repo_type="model", private=private, exist_ok=True) + commit = api.upload_folder( + folder_path=str(checkpoint_dir), + repo_id=repo_id, + repo_type="model", + path_in_repo=f"checkpoints/{checkpoint_dir.name}", + commit_message=f"checkpoint {checkpoint_dir.name}", + ) + api.create_tag( + repo_id=repo_id, + tag=checkpoint_dir.name, + revision=commit.oid, + repo_type="model", + exist_ok=True, + ) + + +def resolve_resume_checkpoint(repo_id: str, output_dir: Path) -> Path: + """Download the latest checkpoint of a Hub training repo into a local run dir. + + The symmetric counterpart to `push_checkpoint_to_hub`: given a model repo holding + `checkpoints//{pretrained_model,training_state}` subtrees, download the highest-numbered step + into `output_dir/checkpoints//`, recreate the local `last` symlink, and return that local + checkpoint dir. Used to resume training from the Hub on a machine (or HF Jobs pod) that does not + have the original local run dir. + """ + latest = find_latest_hub_checkpoint(repo_id) + if latest is None: + raise FileNotFoundError( + f"No checkpoint found in '{repo_id}' under '{CHECKPOINTS_DIR}/'. " + "Was the run trained with --save_checkpoint_to_hub?" + ) + snapshot_download( + repo_id=repo_id, + repo_type="model", + allow_patterns=f"{latest}/*", + local_dir=str(output_dir), + ) + checkpoint_dir = output_dir / latest + update_last_checkpoint(checkpoint_dir) + return checkpoint_dir diff --git a/src/lerobot/configs/__init__.py b/src/lerobot/configs/__init__.py index fa5942129..168b367db 100644 --- a/src/lerobot/configs/__init__.py +++ b/src/lerobot/configs/__init__.py @@ -22,7 +22,7 @@ Import them directly: ``from lerobot.configs.train import TrainPipelineConfig`` """ from .dataset import DatasetRecordConfig -from .default import DatasetConfig, EvalConfig, PeftConfig, WandBConfig +from .default import DatasetConfig, EvalConfig, JobConfig, PeftConfig, WandBConfig from .policies import PreTrainedConfig from .recipe import MessageTurn, TrainingRecipe, load_recipe from .types import ( @@ -55,6 +55,7 @@ __all__ = [ "DatasetRecordConfig", "DatasetConfig", "EvalConfig", + "JobConfig", "MessageTurn", "PeftConfig", "PreTrainedConfig", diff --git a/src/lerobot/configs/default.py b/src/lerobot/configs/default.py index 4f24b9dac..38991a665 100644 --- a/src/lerobot/configs/default.py +++ b/src/lerobot/configs/default.py @@ -145,3 +145,35 @@ class PeftConfig: # If None, the PEFT library defaults to alpha=8, which may dampen high-rank adapters. # Common values are r (alpha == rank) or 2*r. lora_alpha: int | None = None + + +@dataclass +class JobConfig: + # Where training runs. None (omitted) or "local" runs on this machine. + # Any other value is an HF Jobs flavor and submits the run to HF Jobs. + # List available flavors + pricing with `hf jobs hardware` command. + target: str | None = None + # Runtime image for the remote job (ignored for local runs). + image: str = "huggingface/lerobot-gpu:latest" + # Max wall-clock for the remote job as an HF Jobs duration string (e.g. "2h"). + # Defaults to "2d": We pass an explicit, generous cap instead. Set a smaller + # value to fail fast, or a larger one for long runs. + timeout: str | None = "2d" + # Submit and exit instead of streaming the job logs in the foreground. + detach: bool = False + # Extra tags attached to the HF job and to any dataset this run pushes to the + # Hub. A "lerobot" tag is always added; e.g. --job.tags '["lelab"]' adds more. + tags: list[str] = field(default_factory=list) + + # Two entry points to the same predicate: the staticmethod tests a raw target string + # straight from argv (before any JobConfig exists, to decide dispatch early), while the + # property is the ergonomic accessor for code that already holds a config instance. + @staticmethod + def is_remote_target(target: str | None) -> bool: + """True when `target` names an HF Jobs flavor rather than a local run.""" + return target not in (None, "local") + + @property + def is_remote(self) -> bool: + """True when training should run on HF Jobs rather than this machine.""" + return self.is_remote_target(self.target) diff --git a/src/lerobot/configs/train.py b/src/lerobot/configs/train.py index 00fbe81b2..e3d354691 100644 --- a/src/lerobot/configs/train.py +++ b/src/lerobot/configs/train.py @@ -26,11 +26,12 @@ from huggingface_hub.errors import HfHubHTTPError from lerobot import envs from lerobot.optim import LRSchedulerConfig, OptimizerConfig -from lerobot.utils.hub import HubMixin +from lerobot.utils.constants import PRETRAINED_MODEL_DIR +from lerobot.utils.hub import HubMixin, find_latest_hub_checkpoint from lerobot.utils.sample_weighting import SampleWeightingConfig from . import parser -from .default import DatasetConfig, EvalConfig, PeftConfig, WandBConfig +from .default import DatasetConfig, EvalConfig, JobConfig, PeftConfig, WandBConfig from .policies import PreTrainedConfig from .rewards import RewardModelConfig @@ -83,10 +84,11 @@ class TrainPipelineConfig(HubMixin): # with the same value for `dir` its contents will be overwritten unless you set `resume` to true. output_dir: Path | None = None job_name: str | None = None - # Set `resume` to true to resume a previous run. In order for this to work, you will need to make sure - # `dir` is the directory of an existing run with at least one checkpoint in it. - # Note that when resuming a run, the default behavior is to use the configuration from the checkpoint, - # regardless of what's provided with the training command at the time of resumption. + # Set `resume` to true to resume a previous run. Pass `--config_path` pointing at either a local + # checkpoint's train_config.json or a Hub repo id holding `checkpoints//` subtrees (the + # latest checkpoint is downloaded and resumed from). Note that when resuming, the default behavior + # is to use the configuration from the checkpoint, regardless of what's provided with the training + # command at the time of resumption (CLI `--*` flags still override). resume: bool = False # `seed` is used for training (eg: model initialization, dataset shuffling) # AND for the evaluation environments. @@ -118,6 +120,13 @@ class TrainPipelineConfig(HubMixin): wandb: WandBConfig = field(default_factory=WandBConfig) peft: PeftConfig | None = None + # Where to run training (local default, or an HF Jobs flavor). See JobConfig. + job: JobConfig = field(default_factory=JobConfig) + # Push each saved checkpoint to the Hub (policy.repo_id) as it is written, not + # just the final model (useful to monitor progress mid-run). Optional; the + # final model is pushed regardless. Works the same locally and remotely. + save_checkpoint_to_hub: bool = False + # Sample weighting configuration (e.g., for RA-BC training) sample_weighting: SampleWeightingConfig | None = None @@ -137,10 +146,17 @@ class TrainPipelineConfig(HubMixin): return self.reward_model # type: ignore[return-value] return self.policy # type: ignore[return-value] - def validate(self) -> None: - # HACK: We parse again the cli args here to get the pretrained paths if there was some. - policy_path = parser.get_path_arg("policy") + def _resolve_pretrained_from_cli(self) -> None: + """Resolve the pretrained source passed on the CLI into a loaded config. + + The pretrained paths (`--policy.path`, `--reward_model.path`) and + `--config_path` are only recoverable by re-reading the CLI args: draccus + has already consumed them by the time `validate()` runs, so they are not + reflected on `self`. Exactly one source applies, in priority order: + reward-model path, policy path, then resume. + """ reward_model_path = parser.get_path_arg("reward_model") + policy_path = parser.get_path_arg("policy") if reward_model_path: cli_overrides = parser.get_cli_overrides("reward_model") @@ -149,31 +165,54 @@ class TrainPipelineConfig(HubMixin): ) self.reward_model.pretrained_path = str(Path(reward_model_path)) elif policy_path: - yaml_overrides = parser.get_yaml_overrides("policy") - cli_overrides = parser.get_cli_overrides("policy") or [] - self.policy = PreTrainedConfig.from_pretrained( - policy_path, cli_overrides=yaml_overrides + cli_overrides - ) + overrides = parser.get_yaml_overrides("policy") + (parser.get_cli_overrides("policy") or []) + self.policy = PreTrainedConfig.from_pretrained(policy_path, cli_overrides=overrides) self.policy.pretrained_path = Path(policy_path) elif self.resume: - config_path = parser.parse_arg("config_path") - if not config_path: - raise ValueError( - f"A config_path is expected when resuming a run. Please specify path to {TRAIN_CONFIG_NAME}" - ) + self._resolve_resume_checkpoint() - if not Path(config_path).resolve().exists(): - raise NotADirectoryError( - f"{config_path=} is expected to be a local path. " - "Resuming from the hub is not supported for now." - ) + def _resolve_resume_checkpoint(self) -> None: + """Point the trainable config at the checkpoint named by `--config_path`. + `config_path` is either a local path (to a checkpoint's train_config.json or its + pretrained_model/ dir) or a Hub repo id. For a Hub repo, the latest checkpoint is downloaded + into a fresh local run dir and resumed from there. The download is skipped when dispatching to + an HF Job (`job.is_remote`): the pod performs it when it runs the resume locally, and + `submit_to_hf` resolves the source repo for the remote command. + """ + config_path = parser.parse_arg("config_path") + if not config_path: + raise ValueError( + f"A config_path is expected when resuming a run. Please specify path to {TRAIN_CONFIG_NAME}" + ) + + if Path(config_path).resolve().exists(): policy_dir = Path(config_path).parent - if self.policy is not None: - self.policy.pretrained_path = policy_dir - if self.reward_model is not None: - self.reward_model.pretrained_path = str(policy_dir) self.checkpoint_path = policy_dir.parent + elif self.job.is_remote: + return + else: + from lerobot.common.train_utils import resolve_resume_checkpoint + + # `self.output_dir` was loaded from the checkpoint's config and points at the original + # run's (now-absent) local dir. Resume into a fresh local dir instead, unless the user + # passed --output_dir explicitly. + cli_output_dir = parser.parse_arg("output_dir") + if cli_output_dir: + self.output_dir = Path(cli_output_dir) + else: + now = dt.datetime.now() + self.output_dir = Path("outputs/train") / f"{now:%Y-%m-%d}/{now:%H-%M-%S}_resume" + self.checkpoint_path = resolve_resume_checkpoint(config_path, self.output_dir) + policy_dir = self.checkpoint_path / PRETRAINED_MODEL_DIR + + if self.policy is not None: + self.policy.pretrained_path = policy_dir + if self.reward_model is not None: + self.reward_model.pretrained_path = str(policy_dir) + + def validate(self) -> None: + self._resolve_pretrained_from_cli() if self.policy is None and self.reward_model is None: raise ValueError( @@ -216,9 +255,19 @@ class TrainPipelineConfig(HubMixin): if self.eval_steps > 0 and self.dataset.eval_split == 0.0: raise ValueError("eval_steps > 0 requires dataset.eval_split > 0.0 to hold out eval data.") - if hasattr(active_cfg, "push_to_hub") and active_cfg.push_to_hub and not active_cfg.repo_id: + # Remote runs auto-generate the repo_id in submit_to_hf (the policy may only be + # resolved here, from --policy.path), so don't demand it up front for them. + if ( + hasattr(active_cfg, "push_to_hub") + and active_cfg.push_to_hub + and not active_cfg.repo_id + and not self.job.is_remote + ): raise ValueError("'repo_id' argument missing. Please specify it to push the model to the hub.") + if self.save_checkpoint_to_hub and not (self.policy is not None and self.policy.repo_id): + raise ValueError("save_checkpoint_to_hub requires --policy.repo_id.") + @classmethod def __get_path_fields__(cls) -> list[str]: """Keys for draccus pretrained-path loading.""" @@ -255,22 +304,30 @@ class TrainPipelineConfig(HubMixin): elif Path(model_id).is_file(): config_file = model_id else: + dl_kwargs = { + "repo_id": model_id, + "revision": revision, + "cache_dir": cache_dir, + "force_download": force_download, + "proxies": proxies, + "resume_download": resume_download, + "token": token, + "local_files_only": local_files_only, + } try: - config_file = hf_hub_download( - repo_id=model_id, - filename=TRAIN_CONFIG_NAME, - revision=revision, - cache_dir=cache_dir, - force_download=force_download, - proxies=proxies, - resume_download=resume_download, - token=token, - local_files_only=local_files_only, - ) + config_file = hf_hub_download(filename=TRAIN_CONFIG_NAME, **dl_kwargs) except HfHubHTTPError as e: - raise FileNotFoundError( - f"{TRAIN_CONFIG_NAME} not found on the HuggingFace Hub in {model_id}" - ) from e + # No root train_config.json: this is a repo of periodic checkpoints from an + # interrupted run. Fall back to the latest checkpoint's config so the run can be + # resumed straight from the repo with `--config_path=`. + latest = find_latest_hub_checkpoint(model_id, token=token, revision=revision) + if latest is None: + raise FileNotFoundError( + f"{TRAIN_CONFIG_NAME} not found on the HuggingFace Hub in {model_id}" + ) from e + config_file = hf_hub_download( + filename=f"{latest}/{PRETRAINED_MODEL_DIR}/{TRAIN_CONFIG_NAME}", **dl_kwargs + ) cli_args = kwargs.pop("cli_args", []) # Legacy RA-BC migration only applies to framework-saved checkpoints (always JSON). diff --git a/src/lerobot/jobs/__init__.py b/src/lerobot/jobs/__init__.py new file mode 100644 index 000000000..674b98b85 --- /dev/null +++ b/src/lerobot/jobs/__init__.py @@ -0,0 +1,23 @@ +# Copyright 2025 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. + +from lerobot.utils.import_utils import require_package + +# LeRobotDataset (imported at module top in dataset.py) pulls in heavy dataset deps; +# guard the optional dependency here so importing this package fails loudly if it's missing. +require_package("datasets", extra="dataset") + +from .hf import submit_to_hf + +__all__ = ["submit_to_hf"] diff --git a/src/lerobot/jobs/dataset.py b/src/lerobot/jobs/dataset.py new file mode 100644 index 000000000..497f8445e --- /dev/null +++ b/src/lerobot/jobs/dataset.py @@ -0,0 +1,53 @@ +# Copyright 2025 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. +"""Make a training dataset reachable from an HF Job pod. + +The pod can't see the host's ~/.cache/huggingface/lerobot, so the dataset has to +live on the Hub: the pod downloads it by repo_id at train time (the forwarded +HF_TOKEN covers private datasets). A dataset already on the Hub is used as-is; a +local-only dataset is pushed to a PRIVATE repo first (never public). +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from lerobot.datasets import LeRobotDataset +from lerobot.utils.constants import HF_LEROBOT_HOME + +if TYPE_CHECKING: + from huggingface_hub import HfApi + + +def ensure_dataset_available(repo_id: str, *, api: HfApi, tags: list[str] | None = None) -> None: + """Ensure repo_id resolves on the Hub, pushing a local-only dataset privately first. + + `tags` are attached to the dataset only when we push it (an already-on-Hub + dataset is left untouched). Raises RuntimeError if the dataset is neither on + the Hub nor in the local cache. + """ + if api.repo_exists(repo_id, repo_type="dataset"): + return + + local_present = (HF_LEROBOT_HOME / repo_id / "meta" / "info.json").is_file() + if not local_present: + raise RuntimeError( + f"Dataset '{repo_id}' is not in the local cache ({HF_LEROBOT_HOME}) and could not be " + f"reached on the Hub — it may not exist, or be private and inaccessible with your " + f"token. Record or download it first, or run `hf auth login`." + ) + + print(f"[dataset] '{repo_id}' is local-only; pushing to a PRIVATE Hub repo...") + LeRobotDataset(repo_id).push_to_hub(private=True, tags=tags) + print(f"[dataset] '{repo_id}' uploaded (private). The job will download it by repo_id.") diff --git a/src/lerobot/jobs/hf.py b/src/lerobot/jobs/hf.py new file mode 100644 index 000000000..645666412 --- /dev/null +++ b/src/lerobot/jobs/hf.py @@ -0,0 +1,425 @@ +# Copyright 2025 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. +"""Run a lerobot training on HF Jobs (HuggingFace GPUs). + +Ported and simplified from lelab's runners/hf_cloud.py: no UI log queue, no +registry — just submit and stream to stdout. +""" + +from __future__ import annotations + +import copy +import datetime as dt +import json +import netrc +import os +import re +import signal +import sys +import tempfile +import threading +from pathlib import Path +from typing import TYPE_CHECKING + +import httpx +from huggingface_hub import ( + HfApi, + create_repo, + fetch_job_logs, + get_token, + inspect_job, + run_job, + upload_file, +) + +from lerobot.common.train_utils import push_checkpoint_to_hub +from lerobot.configs import parser + +from .dataset import ensure_dataset_available + +if TYPE_CHECKING: + from lerobot.configs.train import TrainPipelineConfig + +_SLUG_RE = re.compile(r"[^a-zA-Z0-9._-]+") + +_TERMINAL_STAGES = {"COMPLETED", "CANCELED", "ERROR", "DELETED"} + +# huggingface_hub 1.x runs on httpx: transient HTTP/transport failures surface as +# httpx.HTTPError and socket-level errors as OSError. Catching only these keeps real +# bugs (TypeError, AttributeError, ...) from being silently retried or counted as +# job failures. +_TRANSIENT_NET_ERRORS = (OSError, httpx.HTTPError) + +# Always attached to remote jobs and pushed datasets so LeRobot-originated work +# is identifiable on the Hub; callers (e.g. LeLab) add their own via --job.tags. +LEROBOT_TAG = "lerobot" + + +def resolve_job_tags(extra: list[str] | None) -> list[str]: + """Return the tag list for a run: the lerobot tag plus any extras, deduped, order-stable.""" + tags = [LEROBOT_TAG, *(extra or [])] + seen: set[str] = set() + return [t for t in tags if not (t in seen or seen.add(t))] + + +def resolve_wandb_api_key() -> str | None: + """Host's wandb key for forwarding to the job: $WANDB_API_KEY, else ~/.netrc.""" + key = os.environ.get("WANDB_API_KEY") + if key: + return key + try: + rc = netrc.netrc() + except (FileNotFoundError, netrc.NetrcParseError, OSError): + return None + auth = rc.authenticators("api.wandb.ai") + if auth is None: + return None + _login, _account, password = auth + return password or None + + +def build_repo_id(username: str, job_name: str, now: dt.datetime) -> str: + """Generate the model repo id for a remote run: /_.""" + slug = _SLUG_RE.sub("-", job_name).strip("-") or "train" + stamp = now.strftime("%Y-%m-%d_%H-%M-%S") + return f"{username}/{slug}_{stamp}" + + +def build_remote_config_file(cfg, repo_id: str, dest: Path, tags: list[str] | None = None) -> Path: + """Write a train_config.json for the pod, with remote overrides applied. + + The pod runs `lerobot-train --config_path=` and downloads the dataset + by repo_id into its own cache. Client-only fields are stripped so the config + is accepted by the trainer image: `job` (pure client orchestration) is always + removed, and `save_checkpoint_to_hub` is removed unless explicitly enabled — + older lerobot images reject unknown keys, so the default keeps the config + compatible with the released `lerobot-gpu` image. `tags` are merged into + policy.tags so the trained model the pod pushes carries them too. + """ + remote = copy.deepcopy(cfg) + remote.policy.push_to_hub = True + remote.policy.repo_id = repo_id + # Don't pin the client's resolved device (e.g. "mps"); let the pod auto-detect its GPU. + remote.policy.device = None + # Drop any host-local dataset root; the pod resolves the dataset by repo_id. + remote.dataset.root = None + if tags: + existing = list(remote.policy.tags or []) + remote.policy.tags = existing + [t for t in tags if t not in existing] + + # Encode to the canonical, pod-parseable dict, then drop the keys the released + # trainer image doesn't know about. + data = remote.to_dict() + data.pop("job", None) + if not remote.save_checkpoint_to_hub: + data.pop("save_checkpoint_to_hub", None) + + dest.parent.mkdir(parents=True, exist_ok=True) + dest.write_text(json.dumps(data, indent=4)) + return dest + + +def _stage_config_on_hub(cfg, repo_id: str, token: str, tags: list[str] | None = None) -> str: + """Upload train_config.json to the model repo and return the repo_id for --config_path.""" + create_repo(repo_id, repo_type="model", private=True, exist_ok=True, token=token) + with tempfile.TemporaryDirectory() as tmp: + config_path = build_remote_config_file(cfg, repo_id, Path(tmp) / "train_config.json", tags=tags) + upload_file( + path_or_fileobj=config_path, + path_in_repo="train_config.json", + repo_id=repo_id, + repo_type="model", + token=token, + ) + return repo_id + + +def _tail_logs( + job_id: str, + done: threading.Event, + success_marker: str | None = None, + success_event: threading.Event | None = None, +) -> None: + """Stream job logs to stdout, reconnecting on dropped streams until done is set. + + Each reconnect re-fetches the full buffered log, so we track how many lines + were already printed and skip them — otherwise a fast-failing job's traceback + gets reprinted on every reconnect. + + When `success_marker` appears in a line, set `success_event` and `done` so the + caller can finish as soon as the trained model lands on the Hub, rather than + waiting out the platform's post-run finalization (which can add ~30s). + """ + printed = 0 + while not done.is_set(): + try: + seen = 0 + for line in fetch_job_logs(job_id=job_id, follow=True): + seen += 1 + if seen <= printed: + continue # already shown on a previous connection + printed = seen + # fetch_job_logs yields SSE data without trailing newlines, so add one + # per entry — otherwise all log lines concatenate onto a single line. + print(line.rstrip("\n"), flush=True) + if success_marker and success_event is not None and success_marker in line: + success_event.set() + done.set() + return + if done.is_set(): + return + # Stream closed cleanly. Wait a moment so the status poller can mark + # the job terminal before we reconnect (avoids re-tailing the buffer). + if done.wait(3): + return + except _TRANSIENT_NET_ERRORS: + if done.wait(2): + return + + +def _poll_until_done( + job_id: str, + done: threading.Event, + poll_interval: float = 5.0, + status_holder: dict | None = None, + max_failures: int = 6, +) -> str | None: + """Poll inspect_job until a terminal stage or until `done` is set. + + Returns the terminal stage string, or None if `done` was set first (detach) + or after `max_failures` consecutive inspect_job errors. When a terminal stage + is reached and `status_holder` is given, records `status_holder["message"]` + (the platform's status message, e.g. "Job timeout"). + """ + failures = 0 + while not done.is_set(): + try: + info = inspect_job(job_id=job_id) + failures = 0 + # `stage` is an enum in some huggingface_hub versions and a plain str in others. + stage = getattr(info.status.stage, "value", info.status.stage) + if stage in _TERMINAL_STAGES: + if status_holder is not None: + status_holder["message"] = getattr(info.status, "message", None) + done.set() + return stage + except _TRANSIENT_NET_ERRORS: + failures += 1 + if failures >= max_failures: + done.set() + return None + done.wait(poll_interval) + return None + + +def _pod_forwarded_args( + argv: list[str], drop_names: tuple[str, ...] = (), drop_prefixes: tuple[str, ...] = () +) -> list[str]: + """User CLI overrides to replay on the pod, minus flags the submitter sets itself. + + Handles both `--name=value` and `--name value` forms. Forwarding the user's overrides (e.g. + `--steps`, `--save_checkpoint_to_hub`) makes a remote resume behave like the same local command. + """ + out: list[str] = [] + skip_next = False + for i, tok in enumerate(argv): + if skip_next: + skip_next = False + continue + name = tok.split("=", 1)[0] + if name in drop_names or any(name.startswith(p) for p in drop_prefixes): + if "=" not in tok and i + 1 < len(argv) and not argv[i + 1].startswith("--"): + skip_next = True # also drop the space-separated value + continue + out.append(tok) + return out + + +def _build_resume_job(cfg: TrainPipelineConfig, username: str) -> tuple[str, list[str]]: + """Resolve the model repo and pod command to resume a run on a job. + + A Hub `config_path` is resumed from directly: its checkpoint config already targets that repo, + so new checkpoints continue the lineage there. A local `config_path` has its checkpoint uploaded + to a new PRIVATE repo first, and the resumed run is forced to push back to it. The pod command + always carries `--job.target=local` so the checkpoint's saved `job.target` can't make the pod + re-dispatch itself. + """ + config_path = parser.parse_arg("config_path") + forwarded = _pod_forwarded_args( + sys.argv[1:], + drop_names=("--config_path", "--policy.repo_id", "--policy.push_to_hub", "--dataset.root"), + drop_prefixes=("--job.",), + ) + + if Path(config_path).exists(): + # Local checkpoint: stage it on the Hub so the pod can resume from it, and push back there. + # Resolve so a `last` symlink uploads under its real step name (digit), which the pod's + # latest-checkpoint lookup keys on. + checkpoint_dir = Path(cfg.checkpoint_path).resolve() + source_repo = build_repo_id(username, cfg.job_name or "train", dt.datetime.now(dt.UTC)) + push_checkpoint_to_hub(checkpoint_dir, source_repo, private=True) + extra = [f"--policy.repo_id={source_repo}", "--policy.push_to_hub=true"] + else: + source_repo = config_path + extra = [] + + command = [ + "lerobot-train", + *forwarded, + f"--config_path={source_repo}", + "--job.target=local", + *extra, + ] + return source_repo, command + + +def submit_to_hf(cfg: TrainPipelineConfig) -> None: + """Submit a training job to HF Jobs infrastructure. + + Validates cfg, resolves credentials, ensures the dataset is on the Hub, then either stages a + sanitized config (fresh run) or resumes from a checkpoint repo, submits the job, and tails logs + until completion or detaches immediately. Ctrl-C detaches without cancelling the remote job. + """ + token = get_token() + if not token: + raise RuntimeError("Not logged in to Hugging Face. Run `hf auth login` first.") + + api = HfApi(token=token) + user_info = api.whoami(token=token) + username = user_info["name"] + + now = dt.datetime.now(dt.UTC) + fresh_repo_id: str | None = None + if not cfg.resume: + # Resolve the model repo and mark it for push BEFORE validate(): validate() requires repo_id + # to be set whenever push_to_hub is True. (A resume reuses the checkpoint's repo instead.) + if cfg.policy is not None: + base_name = cfg.job_name or cfg.policy.type + fresh_repo_id = cfg.policy.repo_id or build_repo_id(username, base_name, now) + cfg.policy.repo_id = fresh_repo_id + cfg.policy.push_to_hub = True + else: + # Path-based policy is resolved inside validate(); fall back to a generic slug. + fresh_repo_id = build_repo_id(username, cfg.job_name or "train", now) + + cfg.validate() + + if cfg.is_reward_model_training: + raise ValueError( + "Remote training via --job.target only supports policy training, not reward models. " + "Run reward-model training locally." + ) + + secrets: dict[str, str] = {"HF_TOKEN": token} + if cfg.wandb.enable: + wandb_key = resolve_wandb_api_key() + if wandb_key is None: + raise ValueError( + "wandb is enabled but no WANDB_API_KEY found. " + "Set it via `export WANDB_API_KEY=...` or add it to ~/.netrc." + ) + secrets["WANDB_API_KEY"] = wandb_key + + tags = resolve_job_tags(cfg.job.tags) + # The dataset must be reachable from the pod for both fresh and resumed runs; a local-only + # dataset is pushed PRIVATE here. Hoisted before the resume/fresh branch since it applies to both. + ensure_dataset_available(cfg.dataset.repo_id, api=api, tags=tags) + + if cfg.resume: + repo_id, command = _build_resume_job(cfg, username) + else: + config_repo_id = _stage_config_on_hub(cfg, fresh_repo_id, token, tags=tags) + repo_id = fresh_repo_id + command = ["lerobot-train", f"--config_path={config_repo_id}"] + + print(f"Submitting job to HF Jobs (flavor={cfg.job.target}, image={cfg.job.image}) ...") + job_info = run_job( + image=cfg.job.image, + command=command, + flavor=cfg.job.target, + secrets=secrets, + timeout=cfg.job.timeout, + # HF Jobs labels are key/value; expose each tag as a queryable label. + labels=dict.fromkeys(tags, "true"), + ) + job_id = job_info.id + job_url = getattr(job_info, "url", None) + print(f"Job submitted: {job_id}") + if job_url: + print(f" Job page: {job_url}") + print(f" Model repo: https://huggingface.co/{repo_id}") + print(f" Monitor: hf jobs logs {job_id}") + print(f" Cancel: hf jobs cancel {job_id}") + + if cfg.job.detach: + return + + done = threading.Event() + detached = threading.Event() + pushed_ok = threading.Event() + stage_holder: dict[str, str | None] = {} + + def _poll() -> None: + stage_holder["stage"] = _poll_until_done(job_id, done, status_holder=stage_holder) + + poll_thread = threading.Thread(target=_poll, daemon=True) + poll_thread.start() + # Finish as soon as the model is pushed, rather than waiting out the platform's + # post-run finalization before the job stage flips to COMPLETED. This matches the + # exact log line emitted by PreTrainedPolicy.push_model_to_hub — the two must stay + # in sync. If it ever stops matching we just fall back to stage-based completion + # (~30s slower), so the contract is an optimization, not a correctness requirement. + success_marker = f"Model pushed to https://huggingface.co/{repo_id}" + log_thread = threading.Thread( + target=_tail_logs, args=(job_id, done, success_marker, pushed_ok), daemon=True + ) + log_thread.start() + + def _detach(sig, frame): + detached.set() + done.set() + print("\nDetached. Job is still running.") + print(f" Monitor: hf jobs logs {job_id}") + print(f" Cancel: hf jobs cancel {job_id}") + + # signal.signal only works on the main thread; when called from a worker thread + # (e.g. an orchestration framework) skip the Ctrl-C-detaches-instead-of-cancels + # handler rather than crashing with ValueError. + install_sigint = threading.current_thread() is threading.main_thread() + original_sigint = signal.getsignal(signal.SIGINT) if install_sigint else None + if install_sigint: + signal.signal(signal.SIGINT, _detach) + try: + # Timeout-based join so SIGINT is delivered to the main thread promptly. + while poll_thread.is_alive(): + poll_thread.join(timeout=0.5) + log_thread.join(timeout=5) + finally: + if install_sigint: + signal.signal(signal.SIGINT, original_sigint) + + if detached.is_set(): + return + + if pushed_ok.is_set(): + print(f"\nTraining complete — model pushed to https://huggingface.co/{repo_id}") + return + + stage = stage_holder.get("stage") + if stage != "COMPLETED": + message = stage_holder.get("message") + detail = f" ({message})" if message else "" + raise RuntimeError( + f"Job {job_id} ended with stage={stage}{detail}. Check logs: hf jobs logs {job_id}" + ) diff --git a/src/lerobot/policies/pretrained.py b/src/lerobot/policies/pretrained.py index a7aabb3f3..aea5f1b08 100644 --- a/src/lerobot/policies/pretrained.py +++ b/src/lerobot/policies/pretrained.py @@ -340,6 +340,9 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC): ignore_patterns=["*.tmp", "*.log"], ) + # Contract: lerobot.jobs.hf.submit_to_hf watches for this exact + # "Model pushed to " line to end a remote run early. Keep the wording + # and URL format in sync (it falls back to status polling if they drift). logging.info(f"Model pushed to {commit_info.repo_url.url}") def generate_model_card( diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index 82c265c93..f2a152df9 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/src/lerobot/scripts/lerobot_train.py @@ -20,6 +20,7 @@ Requires: pip install 'lerobot[training]' (includes dataset + accelerate + wand import dataclasses import logging +import sys import time from contextlib import nullcontext from pprint import pformat @@ -41,15 +42,17 @@ from lerobot.common.train_utils import ( load_training_batch_size, load_training_num_processes, load_training_state, + push_checkpoint_to_hub, save_checkpoint, update_last_checkpoint, ) from lerobot.common.wandb_utils import WandBLogger -from lerobot.configs import parser +from lerobot.configs import JobConfig, parser from lerobot.configs.train import TrainPipelineConfig from lerobot.datasets import EpisodeAwareSampler, compute_sampler_state from lerobot.datasets.factory import make_train_eval_datasets from lerobot.envs import close_envs, make_env, make_env_pre_post_processors +from lerobot.jobs import submit_to_hf from lerobot.optim.factory import make_optimizer_and_scheduler from lerobot.policies import PreTrainedPolicy, make_policy, make_pre_post_processors from lerobot.rewards import make_reward_pre_post_processors @@ -188,6 +191,9 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): cfg: A `TrainPipelineConfig` object containing all training configurations. accelerator: Optional Accelerator instance. If None, one will be created automatically. """ + if cfg.job.is_remote: + return submit_to_hf(cfg) + from lerobot.utils.import_utils import require_package require_package("accelerate", extra="training") @@ -655,6 +661,12 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): optim_state_dict=optim_state_dict, ) update_last_checkpoint(checkpoint_dir) + if cfg.save_checkpoint_to_hub: + push_checkpoint_to_hub( + checkpoint_dir, + cfg.policy.repo_id, + private=cfg.policy.private, + ) if wandb_logger: wandb_logger.log_policy(checkpoint_dir) @@ -735,8 +747,25 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): accelerator.end_training() +def _remote_target_in_argv() -> bool: + """True when the CLI requests a remote HF Jobs run (--job.target=).""" + target = None + args = sys.argv[1:] + for i, tok in enumerate(args): + if tok == "--job.target" and i + 1 < len(args): + target = args[i + 1] + elif tok.startswith("--job.target="): + target = tok.split("=", 1)[1] + return JobConfig.is_remote_target(target) + + def main(): register_third_party_plugins() + if _remote_target_in_argv(): + # The policy device is resolved on the remote pod, not here, so silence the + # client-side "Device '...' is not available" warning PreTrainedConfig emits + # while parsing the config (it fires before train() can dispatch remotely). + logging.getLogger("lerobot.configs.policies").setLevel(logging.ERROR) train() diff --git a/src/lerobot/utils/hub.py b/src/lerobot/utils/hub.py index 566701b31..57eb819f8 100644 --- a/src/lerobot/utils/hub.py +++ b/src/lerobot/utils/hub.py @@ -20,9 +20,33 @@ from typing import Any, TypeVar from huggingface_hub import HfApi from huggingface_hub.utils import validate_hf_hub_args +from .constants import CHECKPOINTS_DIR + T = TypeVar("T", bound="HubMixin") +def find_latest_hub_checkpoint( + repo_id: str, + *, + token: str | bool | None = None, + revision: str | None = None, +) -> str | None: + """Repo-relative path of the most recent checkpoint in a training repo. + + Training runs push checkpoints to ``checkpoints//`` (see + ``push_checkpoint_to_hub``). This lists those step dirs and returns + ``checkpoints/``, or ``None`` if the repo has no checkpoints. + """ + files = HfApi().list_repo_files(repo_id=repo_id, repo_type="model", revision=revision, token=token) + prefix = f"{CHECKPOINTS_DIR}/" + steps = { + name for f in files if f.startswith(prefix) and (name := f[len(prefix) :].split("/", 1)[0]).isdigit() + } + if not steps: + return None + return f"{CHECKPOINTS_DIR}/{max(steps, key=int)}" + + class HubMixin: """ A Mixin containing the functionality to push an object to the hub. diff --git a/tests/configs/test_resume_from_hub.py b/tests/configs/test_resume_from_hub.py new file mode 100644 index 000000000..2cc9fd7ae --- /dev/null +++ b/tests/configs/test_resume_from_hub.py @@ -0,0 +1,68 @@ +# Copyright 2025 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. + +import pytest + +import lerobot.configs.train as tc +from lerobot.configs.train import TrainPipelineConfig + + +class _FakeHTTPError(tc.HfHubHTTPError): + """HfHubHTTPError that can be raised without a real HTTP response object.""" + + def __init__(self): + pass + + +def test_from_pretrained_falls_back_to_latest_checkpoint_config(tmp_path, monkeypatch): + """A Hub repo with no root train_config.json (an interrupted run that only pushed + checkpoints/) resolves via the latest checkpoint's config.""" + # A real train_config.json written by save_pretrained, to be returned by the fallback. + parsed = tc.draccus.parse(TrainPipelineConfig, args=["--dataset.repo_id", "u/d"]) + cfg_file = tmp_path / "train_config.json" + parsed._save_pretrained(tmp_path) + assert cfg_file.is_file() + + calls = [] + + def fake_hf_hub_download(filename=None, **kwargs): + calls.append(filename) + if filename == "train_config.json": + raise _FakeHTTPError() # no root config + if filename == "checkpoints/000010/pretrained_model/train_config.json": + return str(cfg_file) + raise AssertionError(f"unexpected filename {filename}") + + monkeypatch.setattr(tc, "hf_hub_download", fake_hf_hub_download) + monkeypatch.setattr( + tc, "find_latest_hub_checkpoint", lambda repo_id, token=None, revision=None: "checkpoints/000010" + ) + + loaded = TrainPipelineConfig.from_pretrained("user/interrupted-run") + assert loaded.dataset.repo_id == "u/d" + # Tried the root config first, then fell back to the latest checkpoint's config. + assert calls == ["train_config.json", "checkpoints/000010/pretrained_model/train_config.json"] + + +def test_from_pretrained_raises_when_no_root_config_and_no_checkpoints(monkeypatch): + """No root config AND no checkpoints → a clear FileNotFoundError, not the raw HTTP error.""" + + def fake_hf_hub_download(filename=None, **kwargs): + raise _FakeHTTPError() + + monkeypatch.setattr(tc, "hf_hub_download", fake_hf_hub_download) + monkeypatch.setattr(tc, "find_latest_hub_checkpoint", lambda repo_id, token=None, revision=None: None) + + with pytest.raises(FileNotFoundError, match="train_config.json not found"): + TrainPipelineConfig.from_pretrained("user/empty-repo") diff --git a/tests/jobs/__init__.py b/tests/jobs/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/jobs/conftest.py b/tests/jobs/conftest.py new file mode 100644 index 000000000..419d2f83f --- /dev/null +++ b/tests/jobs/conftest.py @@ -0,0 +1,17 @@ +# Copyright 2025 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. + +# Importing concrete policy configs registers their draccus `--policy.type` +# choices (e.g. "act") so tests can parse them. +from lerobot.policies.act.configuration_act import ACTConfig # noqa: F401 diff --git a/tests/jobs/test_dataset.py b/tests/jobs/test_dataset.py new file mode 100644 index 000000000..1f8b9b836 --- /dev/null +++ b/tests/jobs/test_dataset.py @@ -0,0 +1,66 @@ +# Copyright 2025 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. + +from unittest.mock import MagicMock + +import pytest + +pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])") + +from lerobot.jobs.dataset import ensure_dataset_available + + +def _api_with_dataset(exists: bool): + api = MagicMock() + api.repo_exists.return_value = exists + return api + + +def _make_local_cache(tmp_path, repo_id: str) -> None: + """Create the minimal local-cache layout that ensure_dataset_available checks.""" + info = tmp_path / repo_id / "meta" / "info.json" + info.parent.mkdir(parents=True) + info.write_text("{}") + + +# Branch 1: dataset already on Hub → no push, no error (pod downloads by repo_id). +def test_dataset_already_on_hub_is_noop(): + api = _api_with_dataset(True) + assert ensure_dataset_available("user/ds", api=api) is None + api.repo_exists.assert_called_once_with("user/ds", repo_type="dataset") + + +# Branch 2: not on Hub but present locally → always push privately. +def test_dataset_local_only_uploads_privately(tmp_path, monkeypatch): + monkeypatch.setattr("lerobot.jobs.dataset.HF_LEROBOT_HOME", tmp_path) + _make_local_cache(tmp_path, "user/ds") + + api = _api_with_dataset(False) + mock_ds_cls = MagicMock() + monkeypatch.setattr("lerobot.jobs.dataset.LeRobotDataset", mock_ds_cls) + + assert ensure_dataset_available("user/ds", api=api, tags=["lerobot", "lelab"]) is None + + mock_ds_cls.assert_called_once_with("user/ds") + mock_ds_cls.return_value.push_to_hub.assert_called_once_with(private=True, tags=["lerobot", "lelab"]) + + +# Branch 3: not on Hub, NOT in local cache → RuntimeError. +def test_dataset_neither_on_hub_nor_local_raises(tmp_path, monkeypatch): + monkeypatch.setattr("lerobot.jobs.dataset.HF_LEROBOT_HOME", tmp_path) + # tmp_path is empty — no local cache. + + api = _api_with_dataset(False) + with pytest.raises(RuntimeError, match="not in the local cache"): + ensure_dataset_available("user/ds", api=api) diff --git a/tests/jobs/test_hf.py b/tests/jobs/test_hf.py new file mode 100644 index 000000000..3b275cb95 --- /dev/null +++ b/tests/jobs/test_hf.py @@ -0,0 +1,493 @@ +# Copyright 2025 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. + +import datetime as dt +import json +import threading +from types import SimpleNamespace + +import draccus +import httpx +import pytest + +pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])") + +from lerobot.configs.train import TrainPipelineConfig +from lerobot.jobs.hf import ( + _pod_forwarded_args, + _poll_until_done, + build_remote_config_file, + build_repo_id, + resolve_job_tags, + resolve_wandb_api_key, + submit_to_hf, +) + + +def test_resolve_job_tags_always_includes_lerobot_and_dedups(): + assert resolve_job_tags(None) == ["lerobot"] + assert resolve_job_tags([]) == ["lerobot"] + assert resolve_job_tags(["lelab"]) == ["lerobot", "lelab"] + # lerobot isn't duplicated if passed explicitly; order is stable. + assert resolve_job_tags(["lelab", "lerobot", "lelab"]) == ["lerobot", "lelab"] + + +def _fake_inspect(stage_value, *, as_enum=True): + # huggingface_hub returns `stage` as an enum (with `.value`) in some versions and a plain str in others. + stage = SimpleNamespace(value=stage_value) if as_enum else stage_value + return lambda job_id: SimpleNamespace(status=SimpleNamespace(stage=stage)) + + +@pytest.mark.parametrize("as_enum", [True, False], ids=["enum_stage", "str_stage"]) +def test_poll_until_done_returns_terminal_stage(monkeypatch, as_enum): + monkeypatch.setattr("lerobot.jobs.hf.inspect_job", _fake_inspect("COMPLETED", as_enum=as_enum)) + done = threading.Event() + assert _poll_until_done("j", done, poll_interval=0.01) == "COMPLETED" + assert done.is_set() + + +def test_poll_until_done_exits_when_done_already_set(monkeypatch): + # Non-terminal forever; with done pre-set the loop must not block and returns None. + monkeypatch.setattr("lerobot.jobs.hf.inspect_job", _fake_inspect("RUNNING")) + done = threading.Event() + done.set() + assert _poll_until_done("j", done, poll_interval=0.01) is None + + +def test_poll_until_done_gives_up_after_repeated_network_failures(monkeypatch): + monkeypatch.setattr( + "lerobot.jobs.hf.inspect_job", lambda job_id: (_ for _ in ()).throw(httpx.ConnectError("boom")) + ) + done = threading.Event() + result = _poll_until_done("j", done, poll_interval=0.001, max_failures=3) + assert result is None + assert done.is_set() + + +def test_poll_until_done_propagates_programming_errors(monkeypatch): + """A bug (e.g. TypeError) must surface, not be silently retried as a transient failure.""" + monkeypatch.setattr("lerobot.jobs.hf.inspect_job", lambda job_id: (_ for _ in ()).throw(TypeError("bug"))) + done = threading.Event() + with pytest.raises(TypeError): + _poll_until_done("j", done, poll_interval=0.001, max_failures=3) + + +def test_resolve_wandb_key_from_env(monkeypatch): + monkeypatch.setenv("WANDB_API_KEY", "abc123") + assert resolve_wandb_api_key() == "abc123" + + +def test_resolve_wandb_key_missing(monkeypatch, tmp_path): + monkeypatch.delenv("WANDB_API_KEY", raising=False) + monkeypatch.setenv("HOME", str(tmp_path)) # no ~/.netrc here + monkeypatch.setattr("netrc.netrc", lambda *a, **k: (_ for _ in ()).throw(FileNotFoundError())) + assert resolve_wandb_api_key() is None + + +def test_resolve_wandb_key_from_netrc(monkeypatch): + # No env var → fall back to the wandb credentials in ~/.netrc. + monkeypatch.delenv("WANDB_API_KEY", raising=False) + + class _FakeNetrc: + def authenticators(self, host): + assert host == "api.wandb.ai" + return ("login", "account", "netrc-secret") + + monkeypatch.setattr("netrc.netrc", lambda *a, **k: _FakeNetrc()) + assert resolve_wandb_api_key() == "netrc-secret" + + +def test_resolve_wandb_key_netrc_without_wandb_entry(monkeypatch): + # ~/.netrc exists but has no api.wandb.ai entry → None. + monkeypatch.delenv("WANDB_API_KEY", raising=False) + + class _FakeNetrc: + def authenticators(self, host): + return None + + monkeypatch.setattr("netrc.netrc", lambda *a, **k: _FakeNetrc()) + assert resolve_wandb_api_key() is None + + +def test_build_repo_id_sanitizes_and_timestamps(): + now = dt.datetime(2026, 6, 19, 10, 22, 3) + assert build_repo_id("alice", "act", now) == "alice/act_2026-06-19_10-22-03" + # Runs of illegal characters collapse to a single dash; edges are trimmed. + assert build_repo_id("alice", "my cool/run!!", now) == "alice/my-cool-run_2026-06-19_10-22-03" + # A name with nothing usable falls back to "train". + assert build_repo_id("alice", "///", now) == "alice/train_2026-06-19_10-22-03" + + +def test_pod_forwarded_args_drops_host_only_flags(): + """User overrides are replayed on the pod, minus flags that only make sense on the submitter. + + `--dataset.root` is a host-local path the pod can't read, so it must be dropped in both the + `--name=value` and `--name value` forms; unrelated overrides are forwarded untouched. + """ + argv = [ + "--config_path=u/d", + "--dataset.root=/local/data", + "--dataset.root", + "/other/local/data", + "--policy.repo_id=u/keep", + "--steps=10", + "--job.target=a10g-small", + ] + forwarded = _pod_forwarded_args( + argv, + drop_names=("--config_path", "--policy.repo_id", "--policy.push_to_hub", "--dataset.root"), + drop_prefixes=("--job.",), + ) + assert forwarded == ["--steps=10"] + + +def _minimal_cfg(): + return draccus.parse( + TrainPipelineConfig, + args=["--dataset.repo_id", "u/d", "--policy.type", "act", "--job.target", "a10g-small"], + ) + + +def test_validate_skips_repo_id_check_for_remote(): + """Remote runs auto-assign repo_id in submit_to_hf, so validate() must not demand it up front.""" + cfg = _minimal_cfg() # remote target, push_to_hub default True, no explicit repo_id + assert cfg.policy.repo_id is None + cfg.validate() # must not raise + + +def test_validate_requires_repo_id_for_local_push(): + """Local runs that push to the Hub still need an explicit repo_id.""" + cfg = draccus.parse( + TrainPipelineConfig, + args=["--dataset.repo_id", "u/d", "--policy.type", "act"], + ) + with pytest.raises(ValueError, match="repo_id"): + cfg.validate() + + +def test_build_remote_config_applies_overrides(tmp_path): + cfg = _minimal_cfg() + dest = tmp_path / "train_config.json" + out = build_remote_config_file(cfg, "u/run", dest) + assert out == dest + data = json.loads(dest.read_text()) + # `job` is client-only orchestration and must be stripped for the pod. + assert "job" not in data + # save_checkpoint_to_hub defaults off → omitted so older images accept the config. + assert "save_checkpoint_to_hub" not in data + assert data["policy"]["push_to_hub"] is True + assert data["policy"]["repo_id"] == "u/run" + assert data["policy"]["device"] is None # pod auto-detects its GPU + assert data["dataset"]["root"] is None # pod resolves the dataset by repo_id + # the caller's cfg must be left untouched (function works on a deep copy) + assert cfg.job.target == "a10g-small" + assert cfg.save_checkpoint_to_hub is False + + +def test_build_remote_config_includes_checkpoint_flag_when_enabled(tmp_path): + cfg = draccus.parse( + TrainPipelineConfig, + args=[ + "--dataset.repo_id", + "u/d", + "--policy.type", + "act", + "--job.target", + "a10g-small", + "--save_checkpoint_to_hub", + "true", + ], + ) + dest = tmp_path / "train_config.json" + build_remote_config_file(cfg, "u/run", dest) + data = json.loads(dest.read_text()) + # explicitly enabled → kept in the config (requires a matching trainer image). + assert data["save_checkpoint_to_hub"] is True + assert "job" not in data + + +def test_build_remote_config_merges_tags_into_policy(tmp_path): + cfg = _minimal_cfg() + dest = tmp_path / "train_config.json" + build_remote_config_file(cfg, "u/run", dest, tags=["lerobot", "lelab"]) + data = json.loads(dest.read_text()) + # tags propagate to the model the pod pushes. + assert data["policy"]["tags"] == ["lerobot", "lelab"] + + +def test_build_remote_config_merges_tags_without_duplicating(tmp_path): + cfg = _minimal_cfg() + cfg.policy.tags = ["existing", "lerobot"] + dest = tmp_path / "train_config.json" + build_remote_config_file(cfg, "u/run", dest, tags=["lerobot", "lelab"]) + data = json.loads(dest.read_text()) + # pre-existing policy tags are kept; only genuinely-new tags are appended (no dup "lerobot"). + assert data["policy"]["tags"] == ["existing", "lerobot", "lelab"] + + +def test_submit_requires_login(monkeypatch): + monkeypatch.setattr("lerobot.jobs.hf.get_token", lambda: None) + cfg = draccus.parse( + TrainPipelineConfig, + args=["--dataset.repo_id", "u/d", "--policy.type", "act", "--job.target", "a10g-small"], + ) + with pytest.raises(RuntimeError, match="hf auth login"): + submit_to_hf(cfg) + + +def test_submit_passes_validation_and_submits(monkeypatch): + """A type-based policy with no explicit repo_id is auto-assigned one and submitted.""" + from unittest.mock import MagicMock + + # Patch get_token + monkeypatch.setattr("lerobot.jobs.hf.get_token", lambda: "tok") + + # Patch HfApi so whoami returns alice + class FakeHfApi: + def __init__(self, token=None): + pass + + def whoami(self, token=None): + return {"name": "alice"} + + monkeypatch.setattr("lerobot.jobs.hf.HfApi", FakeHfApi) + + # ensure_dataset_available returns None; patch it out so no Hub access happens + # (hf.py imports it at module level, so patch it on lerobot.jobs.hf). + monkeypatch.setattr("lerobot.jobs.hf.ensure_dataset_available", lambda *a, **kw: None) + + # Patch _stage_config_on_hub to skip network + monkeypatch.setattr( + "lerobot.jobs.hf._stage_config_on_hub", + lambda cfg, repo_id, token, tags=None: repo_id, + ) + + # Patch run_job to return a fake job + fake_job = MagicMock() + fake_job.id = "job-123" + run_job_calls = [] + + def fake_run_job(**kwargs): + run_job_calls.append(kwargs) + return fake_job + + monkeypatch.setattr("lerobot.jobs.hf.run_job", fake_run_job) + + cfg = draccus.parse( + TrainPipelineConfig, + args=[ + "--dataset.repo_id", + "u/d", + "--policy.type", + "act", + "--job.target", + "a10g-small", + "--job.detach", + "true", + ], + ) + + # Must NOT raise (pre-fix this raised ValueError about missing repo_id) + submit_to_hf(cfg) + + assert len(run_job_calls) == 1, "run_job should have been called exactly once" + assert cfg.policy.repo_id is not None + assert cfg.policy.repo_id.startswith("alice/") + call = run_job_calls[0] + # The pod runs `lerobot-train --config_path=` on the requested flavor/image. + assert call["command"][0] == "lerobot-train" + assert call["command"][1].startswith("--config_path=") + assert call["flavor"] == "a10g-small" + assert call["image"] == "huggingface/lerobot-gpu:latest" + # The Hub token is forwarded so the pod can pull the (possibly private) dataset. + assert call["secrets"]["HF_TOKEN"] == "tok" + # Every job carries the lerobot tag as a queryable label. + assert call["labels"].get("lerobot") == "true" + + +def test_submit_rejects_reward_model_training(monkeypatch): + """Remote training only supports policies; reward-model runs fail fast with a clear error.""" + monkeypatch.setattr("lerobot.jobs.hf.get_token", lambda: "tok") + + class FakeHfApi: + def __init__(self, token=None): + pass + + def whoami(self, token=None): + return {"name": "alice"} + + monkeypatch.setattr("lerobot.jobs.hf.HfApi", FakeHfApi) + + cfg = _minimal_cfg() + cfg.reward_model = SimpleNamespace(type="reward") # marks this as reward-model training + monkeypatch.setattr(cfg, "validate", lambda: None) # skip pretrained-path resolution + + with pytest.raises(ValueError, match="reward model"): + submit_to_hf(cfg) + + +@pytest.mark.timeout(15) +def test_submit_returns_when_job_completes(monkeypatch): + """Non-detach path must RETURN (not hang) once the job reaches a terminal stage.""" + from types import SimpleNamespace + + monkeypatch.setattr("lerobot.jobs.hf.get_token", lambda: "tok") + + class FakeHfApi: + def __init__(self, token=None): + pass + + def whoami(self, token=None): + return {"name": "alice"} + + monkeypatch.setattr("lerobot.jobs.hf.HfApi", FakeHfApi) + monkeypatch.setattr("lerobot.jobs.hf.ensure_dataset_available", lambda *a, **kw: None) + monkeypatch.setattr( + "lerobot.jobs.hf._stage_config_on_hub", lambda cfg, repo_id, token, tags=None: repo_id + ) + monkeypatch.setattr("lerobot.jobs.hf.run_job", lambda **kw: SimpleNamespace(id="job-1", url="http://x")) + # Job is already COMPLETED on the first poll. + monkeypatch.setattr( + "lerobot.jobs.hf.inspect_job", + lambda job_id: SimpleNamespace( + status=SimpleNamespace(stage=SimpleNamespace(value="COMPLETED"), message=None) + ), + ) + # Log stream ends immediately. + monkeypatch.setattr("lerobot.jobs.hf.fetch_job_logs", lambda job_id, follow=True: iter(())) + + cfg = draccus.parse( + TrainPipelineConfig, + args=["--dataset.repo_id", "u/d", "--policy.type", "act", "--job.target", "a10g-small"], + ) + # Runs in the pytest main thread (signal handler install requires it); the + # @timeout marker fails the test instead of hanging if it regresses. + submit_to_hf(cfg) + + +@pytest.mark.timeout(15) +def test_submit_returns_on_model_pushed_marker(monkeypatch): + """Finish when the model-pushed log appears, even if the job stage never flips.""" + from types import SimpleNamespace + + monkeypatch.setattr("lerobot.jobs.hf.get_token", lambda: "tok") + + class FakeHfApi: + def __init__(self, token=None): + pass + + def whoami(self, token=None): + return {"name": "alice"} + + monkeypatch.setattr("lerobot.jobs.hf.HfApi", FakeHfApi) + monkeypatch.setattr("lerobot.jobs.hf.ensure_dataset_available", lambda *a, **kw: None) + monkeypatch.setattr( + "lerobot.jobs.hf._stage_config_on_hub", lambda cfg, repo_id, token, tags=None: repo_id + ) + monkeypatch.setattr("lerobot.jobs.hf.run_job", lambda **kw: SimpleNamespace(id="job-1", url="http://x")) + # Job stays RUNNING forever — only the log marker can end the command. + monkeypatch.setattr( + "lerobot.jobs.hf.inspect_job", + lambda job_id: SimpleNamespace( + status=SimpleNamespace(stage=SimpleNamespace(value="RUNNING"), message=None) + ), + ) + pushed_line = "INFO Model pushed to https://huggingface.co/alice/myrun" + monkeypatch.setattr("lerobot.jobs.hf.fetch_job_logs", lambda job_id, follow=True: iter([pushed_line])) + + cfg = draccus.parse( + TrainPipelineConfig, + args=[ + "--dataset.repo_id", + "u/d", + "--policy.type", + "act", + "--policy.repo_id", + "alice/myrun", + "--job.target", + "a10g-small", + ], + ) + # Must return via the model-pushed marker despite the perpetual RUNNING stage. + submit_to_hf(cfg) + + +def test_submit_raises_when_wandb_enabled_without_key(monkeypatch): + """wandb.enable with no key reachable anywhere fails fast, before submitting.""" + + monkeypatch.setattr("lerobot.jobs.hf.get_token", lambda: "tok") + + class FakeHfApi: + def __init__(self, token=None): + pass + + def whoami(self, token=None): + return {"name": "alice"} + + monkeypatch.setattr("lerobot.jobs.hf.HfApi", FakeHfApi) + monkeypatch.setattr("lerobot.jobs.hf.resolve_wandb_api_key", lambda: None) + + cfg = draccus.parse( + TrainPipelineConfig, + args=[ + "--dataset.repo_id", + "u/d", + "--policy.type", + "act", + "--job.target", + "a10g-small", + "--wandb.enable", + "true", + ], + ) + with pytest.raises(ValueError, match="WANDB_API_KEY"): + submit_to_hf(cfg) + + +@pytest.mark.timeout(15) +def test_submit_raises_when_job_ends_in_error(monkeypatch): + """A terminal non-COMPLETED stage with no model-pushed marker must raise with the status.""" + from types import SimpleNamespace + + monkeypatch.setattr("lerobot.jobs.hf.get_token", lambda: "tok") + + class FakeHfApi: + def __init__(self, token=None): + pass + + def whoami(self, token=None): + return {"name": "alice"} + + monkeypatch.setattr("lerobot.jobs.hf.HfApi", FakeHfApi) + monkeypatch.setattr("lerobot.jobs.hf.ensure_dataset_available", lambda *a, **kw: None) + monkeypatch.setattr( + "lerobot.jobs.hf._stage_config_on_hub", lambda cfg, repo_id, token, tags=None: repo_id + ) + monkeypatch.setattr("lerobot.jobs.hf.run_job", lambda **kw: SimpleNamespace(id="job-1", url="http://x")) + # Job fails: a terminal ERROR stage carrying the platform's status message. + monkeypatch.setattr( + "lerobot.jobs.hf.inspect_job", + lambda job_id: SimpleNamespace( + status=SimpleNamespace(stage=SimpleNamespace(value="ERROR"), message="Job timeout") + ), + ) + # Logs end without the model-pushed marker. + monkeypatch.setattr("lerobot.jobs.hf.fetch_job_logs", lambda job_id, follow=True: iter(())) + + cfg = draccus.parse( + TrainPipelineConfig, + args=["--dataset.repo_id", "u/d", "--policy.type", "act", "--job.target", "a10g-small"], + ) + with pytest.raises(RuntimeError, match=r"stage=ERROR \(Job timeout\)"): + submit_to_hf(cfg) diff --git a/tests/jobs/test_job_config.py b/tests/jobs/test_job_config.py new file mode 100644 index 000000000..20760fb18 --- /dev/null +++ b/tests/jobs/test_job_config.py @@ -0,0 +1,64 @@ +# Copyright 2025 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. + +import draccus +import pytest + +from lerobot.configs import JobConfig +from lerobot.configs.train import TrainPipelineConfig + + +def test_jobconfig_defaults_are_local(): + cfg = JobConfig() + assert cfg.target is None + assert cfg.is_remote is False + assert cfg.image == "huggingface/lerobot-gpu:latest" + assert cfg.timeout == "2d" + assert cfg.detach is False + + +def test_jobconfig_local_string_is_not_remote(): + assert JobConfig(target="local").is_remote is False + + +def test_jobconfig_flavor_is_remote(): + assert JobConfig(target="a10g-small").is_remote is True + + +def test_train_config_parses_job_target(): + parsed = draccus.parse( + TrainPipelineConfig, + args=["--dataset.repo_id", "u/d", "--policy.type", "act", "--job.target", "a10g-small"], + ) + assert parsed.job.target == "a10g-small" + assert parsed.job.is_remote is True + assert parsed.save_checkpoint_to_hub is False + + +def test_save_checkpoint_to_hub_requires_repo_id(): + cfg = draccus.parse( + TrainPipelineConfig, + args=[ + "--dataset.repo_id", + "u/d", + "--policy.type", + "act", + "--policy.push_to_hub", + "false", + "--save_checkpoint_to_hub", + "true", + ], + ) + with pytest.raises(ValueError, match="requires --policy.repo_id"): + cfg.validate() diff --git a/tests/scripts/test_train_remote_dispatch.py b/tests/scripts/test_train_remote_dispatch.py new file mode 100644 index 000000000..50431da9e --- /dev/null +++ b/tests/scripts/test_train_remote_dispatch.py @@ -0,0 +1,67 @@ +# Copyright 2025 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. + +import sys + +import draccus +import pytest + +# Importing lerobot_train eagerly pulls in lerobot.datasets, which needs the +# `dataset` extra. The base CI tier runs without it, so skip the whole module there. +pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])") + +from lerobot.configs.train import TrainPipelineConfig # noqa: E402 +from lerobot.policies.act.configuration_act import ( + ACTConfig, # noqa: E402, F401 (registers --policy.type act) +) +from lerobot.scripts.lerobot_train import _remote_target_in_argv, train # noqa: E402 + + +def _set_argv(monkeypatch, *args): + monkeypatch.setattr(sys, "argv", ["lerobot-train", *args]) + + +def test_remote_target_detected_space_separated(monkeypatch): + _set_argv(monkeypatch, "--policy.type", "act", "--job.target", "a10g-small") + assert _remote_target_in_argv() is True + + +def test_remote_target_detected_equals(monkeypatch): + _set_argv(monkeypatch, "--job.target=t4-small") + assert _remote_target_in_argv() is True + + +def test_local_string_is_not_remote(monkeypatch): + _set_argv(monkeypatch, "--job.target", "local") + assert _remote_target_in_argv() is False + + +def test_no_target_is_not_remote(monkeypatch): + _set_argv(monkeypatch, "--policy.type", "act") + assert _remote_target_in_argv() is False + + +def test_train_dispatches_to_submit_when_remote(monkeypatch): + """A remote --job.target short-circuits train() to the HF Jobs submitter.""" + import lerobot.scripts.lerobot_train as train_module + + captured = [] + monkeypatch.setattr(train_module, "submit_to_hf", lambda cfg: captured.append(cfg) or "submitted") + cfg = draccus.parse( + TrainPipelineConfig, + args=["--dataset.repo_id", "u/d", "--policy.type", "act", "--job.target", "a10g-small"], + ) + # Returns the submitter's result and never enters the local training path. + assert train(cfg) == "submitted" + assert captured == [cfg] diff --git a/tests/utils/test_hub.py b/tests/utils/test_hub.py new file mode 100644 index 000000000..a55631aeb --- /dev/null +++ b/tests/utils/test_hub.py @@ -0,0 +1,54 @@ +# Copyright 2025 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. + +from unittest.mock import MagicMock + +from lerobot.utils.hub import find_latest_hub_checkpoint + + +def _patch_list_files(monkeypatch, files): + api = MagicMock() + api.list_repo_files.return_value = files + # HfApi is imported into lerobot.utils.hub at module load, so patch it there. + monkeypatch.setattr("lerobot.utils.hub.HfApi", lambda *a, **k: api) + return api + + +def test_find_latest_hub_checkpoint_picks_highest_step(monkeypatch): + _patch_list_files( + monkeypatch, + [ + "README.md", + "checkpoints/000500/pretrained_model/model.safetensors", + "checkpoints/000500/training_state/training_step.json", + "checkpoints/020000/pretrained_model/model.safetensors", + "checkpoints/001000/training_state/training_step.json", + ], + ) + # Numeric max, not lexicographic — "020000" beats "001000"/"000500". + assert find_latest_hub_checkpoint("u/run") == "checkpoints/020000" + + +def test_find_latest_hub_checkpoint_ignores_non_step_entries(monkeypatch): + _patch_list_files( + monkeypatch, + ["checkpoints/last/pretrained_model/model.safetensors", "config.json"], + ) + # "last" (a symlink target name) is not a numeric step → no resolvable checkpoint. + assert find_latest_hub_checkpoint("u/run") is None + + +def test_find_latest_hub_checkpoint_none_when_no_checkpoints(monkeypatch): + _patch_list_files(monkeypatch, ["config.json", "model.safetensors"]) + assert find_latest_hub_checkpoint("u/run") is None diff --git a/tests/utils/test_train_utils.py b/tests/utils/test_train_utils.py index e3705409b..ccd769bd0 100644 --- a/tests/utils/test_train_utils.py +++ b/tests/utils/test_train_utils.py @@ -15,7 +15,9 @@ # limitations under the License. from pathlib import Path -from unittest.mock import Mock, patch +from unittest.mock import MagicMock, Mock, patch + +import pytest from lerobot.common.train_utils import ( get_step_checkpoint_dir, @@ -24,6 +26,7 @@ from lerobot.common.train_utils import ( load_training_num_processes, load_training_state, load_training_step, + push_checkpoint_to_hub, save_checkpoint, save_training_state, save_training_step, @@ -151,3 +154,72 @@ def test_load_training_state_skip_optimizer(tmp_path, optimizer, scheduler): assert loaded_step == 10 assert loaded_optimizer is optimizer assert loaded_scheduler is scheduler + + +def test_push_checkpoint_to_hub_creates_repo_and_uploads(tmp_path, monkeypatch): + ckpt = tmp_path / "010000" + (ckpt / "pretrained_model").mkdir(parents=True) + api = MagicMock() + monkeypatch.setattr("lerobot.common.train_utils.HfApi", lambda *a, **k: api) + push_checkpoint_to_hub(ckpt, "user/run", private=True) + api.create_repo.assert_called_once() + assert api.create_repo.call_args.kwargs["private"] is True + assert api.create_repo.call_args.kwargs["repo_type"] == "model" + api.upload_folder.assert_called_once() + kwargs = api.upload_folder.call_args.kwargs + assert kwargs["repo_id"] == "user/run" + assert kwargs["repo_type"] == "model" + assert kwargs["path_in_repo"] == "checkpoints/010000" + assert kwargs["folder_path"] == str(ckpt) + assert kwargs["commit_message"] == "checkpoint 010000" + # A tag named after the checkpoint step is created so the checkpoint can be + # recovered with --policy.pretrained_revision instead of a commit sha. + api.create_tag.assert_called_once() + tag_kwargs = api.create_tag.call_args.kwargs + assert tag_kwargs["tag"] == "010000" + assert tag_kwargs["revision"] == api.upload_folder.return_value.oid + assert tag_kwargs["repo_type"] == "model" + assert tag_kwargs["exist_ok"] is True + + +def test_push_checkpoint_to_hub_defaults_to_hub_default_visibility(tmp_path, monkeypatch): + ckpt = tmp_path / "010000" + (ckpt / "pretrained_model").mkdir(parents=True) + api = MagicMock() + monkeypatch.setattr("lerobot.common.train_utils.HfApi", lambda *a, **k: api) + push_checkpoint_to_hub(ckpt, "user/run") + api.create_repo.assert_called_once() + assert api.create_repo.call_args.kwargs["private"] is None + + +def test_resolve_resume_checkpoint_downloads_latest_and_links(tmp_path, monkeypatch): + from lerobot.common import train_utils + + out = tmp_path / "run" + + def fake_snapshot_download(repo_id, repo_type, allow_patterns, local_dir): + # Mimic the Hub layout the real download materializes locally. + assert allow_patterns == "checkpoints/020000/*" + (Path(local_dir) / "checkpoints" / "020000" / "pretrained_model").mkdir(parents=True) + return local_dir + + monkeypatch.setattr("lerobot.common.train_utils.snapshot_download", fake_snapshot_download) + monkeypatch.setattr( + "lerobot.common.train_utils.find_latest_hub_checkpoint", lambda repo_id: "checkpoints/020000" + ) + + checkpoint_dir = train_utils.resolve_resume_checkpoint("u/run", out) + + assert checkpoint_dir == out / CHECKPOINTS_DIR / "020000" + last = out / CHECKPOINTS_DIR / LAST_CHECKPOINT_LINK + assert last.is_symlink() + # `last` points at the downloaded step dir. + assert (last.parent / last.readlink()).resolve() == checkpoint_dir.resolve() + + +def test_resolve_resume_checkpoint_raises_without_checkpoints(tmp_path, monkeypatch): + from lerobot.common import train_utils + + monkeypatch.setattr("lerobot.common.train_utils.find_latest_hub_checkpoint", lambda repo_id: None) + with pytest.raises(FileNotFoundError, match="No checkpoint"): + train_utils.resolve_resume_checkpoint("u/run", tmp_path / "run") From 18eee1b477e37f1f57883ae768b5ef231f63fe49 Mon Sep 17 00:00:00 2001 From: Maxime Ellerbach Date: Mon, 29 Jun 2026 18:50:04 +0200 Subject: [PATCH 21/38] refactor(vla-jepa): removing gpu roundtrip (#3750) * refactor(vla-jepa): removing gpu roundtrip for the preprocessing part * major refactor of the forward pass and model input conversion * linting * adressing suggestions from reviews * removing redundant state dtype conversion * avoiding recreating the same tensor each foward pass * api simplification of `_encode_qwen` * avoiding useless video assembly during inference * guard against video=None for the wm loss --- .../policies/vla_jepa/modeling_vla_jepa.py | 464 ++++++------------ .../policies/vla_jepa/qwen_interface.py | 46 +- tests/policies/vla_jepa/conftest.py | 20 +- tests/policies/vla_jepa/test_vla_jepa.py | 48 +- 4 files changed, 228 insertions(+), 350 deletions(-) diff --git a/src/lerobot/policies/vla_jepa/modeling_vla_jepa.py b/src/lerobot/policies/vla_jepa/modeling_vla_jepa.py index 45d83e652..9c689d3c7 100644 --- a/src/lerobot/policies/vla_jepa/modeling_vla_jepa.py +++ b/src/lerobot/policies/vla_jepa/modeling_vla_jepa.py @@ -17,12 +17,10 @@ from __future__ import annotations import logging from collections import deque from pathlib import Path -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any -import numpy as np import torch import torch.nn.functional as F # noqa: N812 -from PIL import Image from torch import Tensor, nn from lerobot.policies.pretrained import PreTrainedPolicy, T @@ -55,12 +53,13 @@ class VLAJEPAModel(nn.Module): - DiT-B: flow-matching action head for future action prediction - V-JEPA: world model for video frame prediction - Input: List[dict] native format (same as original starVLA) - - "image": List[PIL.Image] (multi-view images) - - "video": np.ndarray [V, T, H, W, 3] - - "lang": str (task instruction) - - "action": np.ndarray [T, action_dim] (optional, training only) - - "state": np.ndarray [1, state_dim] (optional) + Inputs are batched tensors kept on the model device + - images: List[List[Tensor [C, H, W]]] (float [0,1]) — per sample, per view (Qwen messages) + - instructions: List[str] + - videos: Tensor [B, V, T, C, H, W] (float [0,1], world model only) + - actions: Tensor [B, T, action_dim] (optional, training only) + - state: Tensor [B, 1, state_dim] (optional) + - action_is_pad: Tensor [B, T] (optional) """ def __init__(self, config: VLAJEPAConfig) -> None: @@ -75,6 +74,11 @@ class VLAJEPAModel(nn.Module): self.action_tokens, self.action_token_ids, self.embodied_action_token_id = ( self.qwen.expand_tokenizer() ) + self.register_buffer( + "_action_token_ids_t", + torch.tensor(self.action_token_ids, dtype=torch.long), + persistent=False, + ) # Action head (flow-matching DiT) self.action_model = VLAJEPAActionHead(config, cross_attention_dim=self.qwen.model.config.hidden_size) @@ -161,166 +165,123 @@ class VLAJEPAModel(nn.Module): # ---- Native VLA-JEPA forward (follows original VLA_JEPA.py) ---- - def forward(self, examples: list[dict]) -> dict[str, Tensor]: - """ - Native forward pass following original starVLA VLA_JEPA.forward. - - Args: - examples: List of per-sample dicts with keys: - "image" : List[PIL.Image] — multi-view images - "video" : np.ndarray [V, T, H, W, 3] - "lang" : str — task instruction - "action" : np.ndarray [T, action_dim] (optional) - "state" : np.ndarray [1, state_dim] (optional) - - Returns: - dict with "action_loss" and "wm_loss" keys (scalar Tensors). - """ - # Unpack native format (same pattern as original VLA_JEPA.py) - batch_images = [ex["image"] for ex in examples] # List[List[PIL.Image]] - batch_videos = [ex["video"] for ex in examples] # List[np.ndarray] - instructions = [ex["lang"] for ex in examples] # List[str] - has_action = "action" in examples[0] and examples[0]["action"] is not None - actions = [ex["action"] for ex in examples] if has_action else None - has_state = "state" in examples[0] and examples[0]["state"] is not None - state = [ex["state"] for ex in examples] if has_state else None - action_is_pad = ( - [ex["action_is_pad"] for ex in examples] - if has_action and "action_is_pad" in examples[0] and examples[0]["action_is_pad"] is not None - else None - ) - - # Stack videos: [B, V, T, H, W, 3] -> [B, V, T, 3, H, W] - batch_videos = np.stack(batch_videos) - batch_videos = batch_videos.transpose(0, 1, 2, 5, 3, 4) # [B, V, T, 3, H, W] - - # Adjust number of views for the world model: - # - fewer views than expected: duplicate the first view to fill up - # - more views than expected: keep only the first num_views_world_model views - num_views_world_model = self.config.jepa_tubelet_size - if batch_videos.shape[1] < num_views_world_model: - num_missing_views = num_views_world_model - batch_videos.shape[1] - first_view = np.repeat(batch_videos[:, :1], num_missing_views, axis=1) - batch_videos = np.concatenate([batch_videos, first_view], axis=1) - elif batch_videos.shape[1] > num_views_world_model: - batch_videos = batch_videos[:, :num_views_world_model] - - # ---- Step 1: QwenVL encode (same as original) ---- + def _encode_qwen( + self, images: list[list[Tensor]], instructions: list[str], *, need_action_tokens: bool + ) -> tuple[Tensor, Tensor, Tensor | None]: + """Run Qwen and gather the embodied-action (and optionally action) token hidden states.""" qwen_inputs = self.qwen.build_inputs( - images=batch_images, + images=images, instructions=instructions, action_prompt=self.replace_prompt, embodied_prompt=self.embodied_replace_prompt, ) - - # Locate embodied-action tokens (always needed for action head) - embodied_mask = qwen_inputs["input_ids"] == self.embodied_action_token_id - embodied_indices = embodied_mask.nonzero(as_tuple=True) - - # Locate action tokens (only needed for world model predictor) - if self.config.enable_world_model: - action_mask = torch.isin( - qwen_inputs["input_ids"], - torch.tensor(self.action_token_ids, device=qwen_inputs["input_ids"].device), - ) - action_indices = action_mask.nonzero(as_tuple=True) + input_ids = qwen_inputs["input_ids"] + embodied_idx = (input_ids == self.embodied_action_token_id).nonzero(as_tuple=True) + action_idx = None + if need_action_tokens: + action_mask = torch.isin(input_ids, self._action_token_ids_t) + action_idx = action_mask.nonzero(as_tuple=True) device_type = next(self.parameters()).device.type - with torch.autocast(device_type=device_type, dtype=torch.bfloat16): last_hidden = self._qwen_last_decoder_hidden(qwen_inputs) # [B, seq_len, H] b, _, h = last_hidden.shape + embodied_action_tokens = last_hidden[embodied_idx[0], embodied_idx[1], :].view(b, -1, h) + action_tokens = ( + last_hidden[action_idx[0], action_idx[1], :].view(b, -1, h) + if action_idx is not None + else None + ) + return embodied_action_tokens, action_tokens - if self.config.enable_world_model: - action_tokens = last_hidden[action_indices[0], action_indices[1], :].view(b, -1, h) + def _world_model_loss(self, videos: Tensor, action_tokens: Tensor) -> Tensor: + """JEPA encode + predictor L1 loss. `videos` is [B, V, T, C, H, W] float in [0, 1].""" + # Match the world model's expected view count: pad with the first view, or trim extras. + num_views = self.config.jepa_tubelet_size + if videos.shape[1] < num_views: + missing = num_views - videos.shape[1] + videos = torch.cat([videos, videos[:, :1].repeat(1, missing, 1, 1, 1, 1)], dim=1) + elif videos.shape[1] > num_views: + videos = videos[:, :num_views] - embodied_action_tokens = last_hidden[embodied_indices[0], embodied_indices[1], :].view(b, -1, h) + b, v, t_frames, c, h_img, w_img = videos.shape + flat = videos.reshape(b * v, t_frames, c, h_img, w_img) + # Fast (torchvision) video processor on-device, do_rescale=False (frames already in [0, 1]). + video_pixels = self.video_processor( + videos=list(flat), + return_tensors="pt", + device=self.video_encoder.device, + do_rescale=False, + )["pixel_values_videos"] # [B*V, T, C, H, W] - # ---- Step 2+3: JEPA Encoder + Predictor ---- - device_wm = last_hidden.device - if not self.config.enable_world_model: - wm_loss = torch.tensor(0.0, device=device_wm) + with torch.no_grad(): + video_embeddings = self.video_encoder.get_vision_features(pixel_values_videos=video_pixels) + # Merge views: [B*V, ...] -> [B, ..., V*embed_dim] + video_embeddings = torch.cat(torch.chunk(video_embeddings, chunks=v, dim=0), dim=2) + + tubelet_size = self.video_encoder.config.tubelet_size + # num_video_frames raw frames → t_enc_total temporal positions after tubelet compression + t_enc_total = self.config.num_video_frames // tubelet_size + if t_enc_total < 2: + return torch.zeros((), device=video_embeddings.device) + + # Shift-by-one JEPA split: input_states = positions 0..T-2, gt_states = positions 1..T-1 + t_enc_ctx = t_enc_total - 1 + tokens_per_frame = video_embeddings.shape[1] // t_enc_total + input_states = video_embeddings[:, : tokens_per_frame * t_enc_ctx, :] + gt_states = video_embeddings[:, tokens_per_frame:, :] + + expected_actions = t_enc_ctx * self.config.num_action_tokens_per_timestep + if action_tokens.shape[1] < expected_actions: + pad = action_tokens[:, -1:].repeat(1, expected_actions - action_tokens.shape[1], 1) + action_tokens = torch.cat([action_tokens, pad], dim=1) + + predicted_states = self.video_predictor( + input_states.float(), action_tokens[:, :expected_actions].float() + ) + return F.l1_loss(predicted_states, gt_states.float(), reduction="mean") + + def _action_loss( + self, + embodied_action_tokens: Tensor, + actions: Tensor, + state: Tensor | None, + action_is_pad: Tensor | None, + ) -> Tensor: + """Flow-matching action-head loss, repeated over `repeated_diffusion_steps`.""" + device_type = next(self.parameters()).device.type + with torch.autocast(device_type=device_type, dtype=torch.float32): + r = self.config.repeated_diffusion_steps + horizon = self.config.chunk_size + actions_target = actions[:, -horizon:, :].to(torch.float32).repeat(r, 1, 1) + embodied = embodied_action_tokens.repeat(r, 1, 1) + state_rep = state.to(embodied_action_tokens.dtype).repeat(r, 1, 1) if state is not None else None + pad_rep = action_is_pad[:, -horizon:].repeat(r, 1) if action_is_pad is not None else None + return self.action_model(embodied, actions_target, state_rep, pad_rep) + + def forward( + self, + images: list[list[Tensor]], + instructions: list[str], + videos: Tensor | None = None, + actions: Tensor | None = None, + state: Tensor | None = None, + action_is_pad: Tensor | None = None, + ) -> dict[str, Tensor]: + """Native forward: Qwen encode → optional world-model loss → optional action-head loss.""" + embodied_action_tokens, action_tokens = self._encode_qwen( + images, instructions, need_action_tokens=self.config.enable_world_model + ) + + if self.config.enable_world_model and videos is not None: + wm_loss = self._world_model_loss(videos, action_tokens) else: - b, v, t_frames, c, h_img, w_img = batch_videos.shape - batch_videos_flat = batch_videos.reshape(b * v, t_frames, c, h_img, w_img) + wm_loss = torch.zeros((), device=embodied_action_tokens.device) - video_pixels = self.video_processor(videos=list(batch_videos_flat), return_tensors="pt")[ - "pixel_values_videos" - ].to(self.video_encoder.device) # [B*V, T, C, H, W] - - with torch.no_grad(): - video_embeddings = self.video_encoder.get_vision_features(pixel_values_videos=video_pixels) - # Merge views: [B*V, ...] -> [B, ..., V*embed_dim] - video_embeddings = torch.cat(torch.chunk(video_embeddings, chunks=v, dim=0), dim=2) - - tubelet_size = self.video_encoder.config.tubelet_size - device_wm = video_embeddings.device - # num_video_frames raw frames → t_enc_total temporal positions after tubelet compression - t_enc_total = self.config.num_video_frames // tubelet_size - - if t_enc_total < 2: - wm_loss = torch.tensor(0.0, device=device_wm) - else: - # Shift-by-one JEPA split (matches original VLA_JEPA.py lines 231-232): - # input_states: positions 0..T-2, gt_states: positions 1..T-1 - t_enc_ctx = t_enc_total - 1 - tokens_per_frame = video_embeddings.shape[1] // t_enc_total - - input_states = video_embeddings[:, : tokens_per_frame * t_enc_ctx, :] - gt_states = video_embeddings[:, tokens_per_frame:, :] - - expected_actions = t_enc_ctx * self.config.num_action_tokens_per_timestep - if action_tokens.shape[1] < expected_actions: - pad = action_tokens[:, -1:].repeat(1, expected_actions - action_tokens.shape[1], 1) - action_tokens = torch.cat([action_tokens, pad], dim=1) - - predicted_states = self.video_predictor( - input_states.float(), - action_tokens[:, :expected_actions].float(), - ) - - wm_loss = F.l1_loss(predicted_states, gt_states.float(), reduction="mean") - - if not has_action: + if actions is None: return {"wm_loss": wm_loss} - # ---- Step 4: Action Head ---- - with torch.autocast(device_type=device_type, dtype=torch.float32): - actions_tensor = torch.tensor( - np.array(actions), device=last_hidden.device, dtype=torch.float32 - ) # [B, T_full, action_dim] - action_horizon = self.config.chunk_size - actions_target = actions_tensor[:, -action_horizon:, :] - - state_tensor = None - if state is not None: - state_tensor = torch.tensor( - np.array(state), device=last_hidden.device, dtype=last_hidden.dtype - ) # [B, 1, state_dim] - - repeated_diffusion_steps = self.config.repeated_diffusion_steps - actions_target = actions_target.repeat(repeated_diffusion_steps, 1, 1) - embodied_action_tokens = embodied_action_tokens.repeat(repeated_diffusion_steps, 1, 1) - if state_tensor is not None: - state_tensor = state_tensor.repeat(repeated_diffusion_steps, 1, 1) - - action_is_pad_rep = None - if action_is_pad is not None: - pad_tensor = torch.stack( - [ - p.to(actions_target.device) - if isinstance(p, Tensor) - else torch.tensor(p, device=actions_target.device) - for p in action_is_pad - ] - ) # [B, T_full] - pad_tensor = pad_tensor[:, -action_horizon:] # [B, action_horizon] - action_is_pad_rep = pad_tensor.repeat(repeated_diffusion_steps, 1) # [B*R, action_horizon] - - action_loss = self.action_model( - embodied_action_tokens, actions_target, state_tensor, action_is_pad_rep - ) - + action_loss = self._action_loss(embodied_action_tokens, actions, state, action_is_pad) return {"action_loss": action_loss, "wm_loss": wm_loss * self.config.world_model_loss_weight} # ---- Native predict_action (follows original VLA_JEPA.predict_action) ---- @@ -328,58 +289,23 @@ class VLAJEPAModel(nn.Module): @torch.no_grad() def predict_action( self, - batch_images: list[list[Image.Image]], + images: list[list[Tensor]], instructions: list[str], - state: np.ndarray | None = None, - ) -> np.ndarray: - """ - Native action prediction following original VLA_JEPA.predict_action. - - Args: - batch_images: List of samples; each is List[PIL.Image] (multi-view). - instructions: Task instructions, one per sample. - state: Optional [B, state_dim] numpy array. - - Returns: - np.ndarray [B, action_horizon, action_dim] — predicted actions. - """ + state: Tensor | None = None, + ) -> Tensor: + """Predict an action chunk. `images` is per-sample, per-view float [0,1] [C, H, W] tensors.""" if self.config.resize_images_to is not None: height, width = self.config.resize_images_to - resampling = getattr(Image, "Resampling", Image).BOX - batch_images = [ - [image.resize((width, height), resample=resampling) for image in sample_images] - for sample_images in batch_images + images = [ + [F.interpolate(img[None], size=(height, width), mode="area")[0] for img in views] + for views in images ] - qwen_inputs = self.qwen.build_inputs( - images=batch_images, - instructions=instructions, - action_prompt=self.replace_prompt, - embodied_prompt=self.embodied_replace_prompt, + embodied_action_tokens, _ = self._encode_qwen(images, instructions, need_action_tokens=False) + return self.action_model.predict_action( + embodied_action_tokens.float(), state.float() if state is not None else None ) - embodied_mask = qwen_inputs["input_ids"] == self.embodied_action_token_id - embodied_indices = embodied_mask.nonzero(as_tuple=True) - - device_type = next(self.parameters()).device.type - - with torch.autocast(device_type=device_type, dtype=torch.bfloat16): - last_hidden = self._qwen_last_decoder_hidden(qwen_inputs) # [B, seq_len, H] - b, _, h = last_hidden.shape - embodied_action_tokens = last_hidden[embodied_indices[0], embodied_indices[1], :].view(b, -1, h) - - state_tensor = None - if state is not None: - state_tensor = torch.from_numpy(np.array(state)).to( - device=last_hidden.device, dtype=last_hidden.dtype - ) - - pred_actions = self.action_model.predict_action( - embodied_action_tokens.float(), state_tensor.float() if state_tensor is not None else None - ) # [B, action_horizon, action_dim] - - return pred_actions.detach().cpu().numpy() - # ============================================================================ # LeRobot Adapter Layer - converts between LeRobot batch format and native VLA-JEPA format @@ -390,9 +316,9 @@ class VLAJEPAPolicy(PreTrainedPolicy): """ LeRobot adapter for VLA-JEPA. - Converts LeRobot's standard batch format (dict[str, Tensor]) to the native - VLA-JEPA format (List[dict]), calls the native model, and converts outputs - back to LeRobot format. + Converts LeRobot's standard batch format (dict[str, Tensor]) to the batched tensors + the native model expects (keeping everything on-device), calls the native model, and + converts outputs back to LeRobot format. """ config_class = VLAJEPAConfig @@ -419,9 +345,8 @@ class VLAJEPAPolicy(PreTrainedPolicy): # ---- Format Conversion: LeRobot → Native ---- - def _prepare_model_inputs(self, batch: dict[str, Tensor]) -> list[dict]: - """ - Convert LeRobot batch format to native VLA-JEPA examples format. + def _prepare_model_inputs(self, batch: dict[str, Tensor], training=True) -> dict[str, Any]: + """Convert a LeRobot batch to the model's batched, on-device inputs. LeRobot format: batch = { @@ -431,65 +356,25 @@ class VLAJEPAPolicy(PreTrainedPolicy): "task": str | List[str], (optional instruction) } - Native format (List[dict]): - { - "image": List[PIL.Image], # multi-view images per sample - "video": np.ndarray [V, T, H, W, 3], - "lang": str, # task instruction - "action": np.ndarray [T, action_dim], # optional - "state": np.ndarray [1, state_dim], # optional - } + Returns the kwargs for `VLAJEPAModel.forward` / `.predict_action` (everything stays + on the batch device; no per-sample shredding): `images` (per-sample, per-view list for + Qwen messages), `instructions`, and the batched `videos` / `actions` / `state` / + `action_is_pad` when present. """ - # Determine batch size from the first image feature image_keys = list(self.config.image_features.keys()) if not image_keys: raise ValueError("VLAJEPA requires at least one image feature.") - first_key = image_keys[0] - first_tensor = batch[first_key] - batch_size = first_tensor.shape[0] + batch_size = batch[image_keys[0]].shape[0] - # ---- Collect images per sample ---- - # images_per_sample[b][v] = PIL.Image for view v - images_per_sample: list[list[Image.Image]] = [[] for _ in range(batch_size)] + # Current-frame image per view ([B, C, H, W]); regroup per sample for Qwen messages. + frames = [] for key in image_keys: - tensor = batch[key] # [B, C, H, W] or [B, T, C, H, W] - if tensor.ndim == 5: - # observation_delta_indices = [0, 1, ..., num_video_frames-1] - # index 0 is the current observation (delta=0) - tensor = tensor[:, 0] - for b in range(batch_size): - images_per_sample[b].append(self.model.qwen.tensor_to_pil(tensor[b])) + t = batch[key] + if t.ndim == 5: # [B, T, C, H, W] -> current observation (delta=0) + t = t[:, 0] + frames.append(self.model.qwen.to_pixel_values(t)) + images = [[frame[b] for frame in frames] for b in range(batch_size)] - # ---- Collect videos per sample ---- - # Build video arrays: for each sample, stack views as [V, T, H, W, 3] - # Check whether any image feature has a time dimension - video_source = None - for k in image_keys: - if k in batch: - video_source = batch[k] # Use first available for shape inspection - break - - if video_source is None: - raise ValueError("No image data found in batch for video construction.") - - videos_per_sample = [] - for b in range(batch_size): - sample_views = [] - for k in image_keys: - t = batch[k][b] # [C, H, W] or [T, C, H, W] - if t.ndim == 3: - t = t.unsqueeze(0) # [1, C, H, W] - # Convert to [T, H, W, 3] numpy - t_np = t.permute(0, 2, 3, 1).detach().cpu().float().numpy() - # Clamp to [0, 255] - if t_np.max() <= 1.0: - t_np = t_np * 255.0 - t_np = np.rint(t_np.clip(0, 255)).astype(np.uint8) - sample_views.append(t_np) - # Stack views: [V, T, H, W, 3] - videos_per_sample.append(np.stack(sample_views, axis=0)) - - # ---- Collect instructions ---- tasks = batch.get("task") if tasks is None: instructions = ["Execute the robot action."] * batch_size @@ -498,52 +383,32 @@ class VLAJEPAPolicy(PreTrainedPolicy): else: instructions = list(tasks) - # ---- Collect actions (training only) ---- - actions_list = None - action_is_pad_list = None - actions_tensor = batch.get(ACTION) - if actions_tensor is not None: - if actions_tensor.ndim == 2: - actions_tensor = actions_tensor.unsqueeze(1) - actions_list = [actions_tensor[b].detach().cpu().float().numpy() for b in range(batch_size)] - action_is_pad_tensor = batch.get("action_is_pad") - if action_is_pad_tensor is not None: - action_is_pad_list = [action_is_pad_tensor[b].detach().cpu() for b in range(batch_size)] + inputs: dict[str, Any] = {"images": images, "instructions": instructions} - # ---- Collect state ---- - state_list = None - state_tensor = batch.get(OBS_STATE) - if state_tensor is not None: - if state_tensor.ndim > 2: - state_tensor = state_tensor[:, -1, :] - if state_tensor.ndim == 2: - state_tensor = state_tensor.unsqueeze(1) # [B, 1, state_dim] - state_list = [state_tensor[b].detach().cpu().float().numpy() for b in range(batch_size)] + # Videos [B, V, T, C, H, W] - only assembled during training when the world model consumes them. + if self.model.config.enable_world_model and training: + views = [batch[k].unsqueeze(1) if batch[k].ndim == 4 else batch[k] for k in image_keys] + inputs["videos"] = self.model.qwen.to_pixel_values(torch.stack(views, dim=1)) - # ---- Assemble native examples ---- - examples = [] - for b in range(batch_size): - example = { - "image": images_per_sample[b], - "video": videos_per_sample[b], - "lang": instructions[b], - } - if actions_list is not None: - example["action"] = actions_list[b] - if action_is_pad_list is not None: - example["action_is_pad"] = action_is_pad_list[b] - if state_list is not None: - example["state"] = state_list[b] - examples.append(example) + actions = batch.get(ACTION) + if actions is not None: + inputs["actions"] = (actions.unsqueeze(1) if actions.ndim == 2 else actions).float() + if (pad := batch.get("action_is_pad")) is not None: + inputs["action_is_pad"] = pad - return examples + state = batch.get(OBS_STATE) + if state is not None: + if state.ndim > 2: + state = state[:, -1, :] + inputs["state"] = (state.unsqueeze(1) if state.ndim == 2 else state).float() # [B, 1, dim] + + return inputs # ---- LeRobot Policy Interface ---- def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict]: """LeRobot train forward: convert → native forward → aggregate losses.""" - examples = self._prepare_model_inputs(batch) - native_output = self.model.forward(examples) + native_output = self.model.forward(**self._prepare_model_inputs(batch, training=True)) ref = next(iter(native_output.values())) zero = torch.zeros((), device=ref.device, dtype=ref.dtype) @@ -561,16 +426,9 @@ class VLAJEPAPolicy(PreTrainedPolicy): self.eval() self._queues = populate_queues(self._queues, batch, exclude_keys=[ACTION]) - examples = self._prepare_model_inputs(batch) - batch_images = [ex["image"] for ex in examples] - instructions = [ex["lang"] for ex in examples] - - state_np = None - if "state" in examples[0] and examples[0]["state"] is not None: - state_np = np.stack([ex["state"] for ex in examples]) - - actions_np = self.model.predict_action(batch_images, instructions, state_np) - return torch.from_numpy(actions_np).to(device=self.config.device, dtype=torch.float32) + inputs = self._prepare_model_inputs(batch, training=False) + actions = self.model.predict_action(inputs["images"], inputs["instructions"], inputs.get("state")) + return actions.to(device=self.config.device, dtype=torch.float32) @torch.no_grad() def select_action(self, batch: dict[str, Tensor], noise: Tensor | None = None) -> Tensor: diff --git a/src/lerobot/policies/vla_jepa/qwen_interface.py b/src/lerobot/policies/vla_jepa/qwen_interface.py index 24f530efc..bcad1f558 100644 --- a/src/lerobot/policies/vla_jepa/qwen_interface.py +++ b/src/lerobot/policies/vla_jepa/qwen_interface.py @@ -17,9 +17,7 @@ from __future__ import annotations from collections.abc import Sequence from typing import TYPE_CHECKING -import numpy as np import torch -from PIL import Image from lerobot.utils.import_utils import _transformers_available @@ -78,7 +76,7 @@ class Qwen3VLInterface(torch.nn.Module): def build_inputs( self, - images: Sequence[Sequence[Image.Image]], + images: Sequence[Sequence[torch.Tensor]], instructions: Sequence[str], action_prompt: str, embodied_prompt: str, @@ -94,24 +92,42 @@ class Qwen3VLInterface(torch.nn.Module): content.append({"type": "text", "text": prompt}) messages.append([{"role": "user", "content": content}]) + # The Qwen image processor is a torchvision-backed fast processor: passing the + # images as GPU tensors (with `device`) keeps the whole vision pipeline on-device + # and avoids a GPU->CPU->GPU roundtrip. The image tensors are forwarded through + # apply_chat_template untouched into Qwen3VLProcessor.__call__. + # do_rescale=False: images already arrive as float in [0, 1] (the dataset decoder + # yields float32/255 and VISUAL normalization is IDENTITY), so we skip the + # processor's /255 rescale instead of round-tripping through uint8. batch_inputs = self.processor.apply_chat_template( messages, tokenize=True, add_generation_prompt=True, return_dict=True, - processor_kwargs={"padding": True, "return_tensors": "pt"}, + processor_kwargs={ + "padding": True, + "return_tensors": "pt", + "device": self.model.device, + "do_rescale": False, + }, ) return batch_inputs.to(self.model.device) @staticmethod - def tensor_to_pil(image_tensor: torch.Tensor) -> Image.Image: - image = image_tensor.detach().cpu() - if image.ndim == 3 and image.shape[0] in (1, 3): - image = image.permute(1, 2, 0) - image = image.float() - if image.max() <= 1.0: - image = image * 255.0 - image = image.clamp(0, 255).round().to(torch.uint8).numpy() - if image.shape[-1] == 1: - image = np.repeat(image, 3, axis=-1) - return Image.fromarray(image) + def to_pixel_values(image_tensor: torch.Tensor) -> torch.Tensor: + """Prepare an image/video tensor for the fast processors (used with do_rescale=False). + + The dataset decoder yields float32 in [0, 1] (channels-first) and VISUAL + normalization is IDENTITY, so the tensor already arrives in [0, 1]; we pass it + through as float and let the processors normalize (no rescale, no uint8 + quantization). A single channel is expanded to 3 to match the RGB processors. + + Works for any channels-first layout (channel dim is -3): [C, H, W], [B, C, H, W], + [T, C, H, W], [B, V, T, C, H, W], ... + """ + image = image_tensor.detach().float() + if image.shape[-3] == 1: + repeats = [1] * image.ndim + repeats[-3] = 3 + image = image.repeat(*repeats) + return image diff --git a/tests/policies/vla_jepa/conftest.py b/tests/policies/vla_jepa/conftest.py index 5301b5bc7..dd40ca9ea 100644 --- a/tests/policies/vla_jepa/conftest.py +++ b/tests/policies/vla_jepa/conftest.py @@ -8,7 +8,6 @@ from types import SimpleNamespace import numpy as np import pytest import torch -from PIL import Image from torch import Tensor, nn from lerobot.configs.types import FeatureType, PolicyFeature @@ -191,7 +190,7 @@ class _FakeQwenInterface(nn.Module): def build_inputs( self, - images: list[list[Image.Image]], + images: list[list[Tensor]], instructions: list[str], action_prompt: str, embodied_prompt: str, @@ -214,12 +213,13 @@ class _FakeQwenInterface(nn.Module): } @staticmethod - def tensor_to_pil(image_tensor: Tensor) -> Image.Image: - image = image_tensor.detach().cpu() - if image.ndim == 3 and image.shape[0] in (1, 3): - image = image.permute(1, 2, 0) - image = (image.float().clamp(0, 1) * 255).to(torch.uint8).numpy() - return Image.fromarray(image) + def to_pixel_values(image_tensor: Tensor) -> Tensor: + image = image_tensor.detach().float() + if image.shape[-3] == 1: + repeats = [1] * image.ndim + repeats[-3] = 3 + image = image.repeat(*repeats) + return image class _FakeVideoEncoder(nn.Module): @@ -242,12 +242,14 @@ class _FakeVideoEncoder(nn.Module): class _FakeVideoProcessor: - def __call__(self, videos, return_tensors: str) -> dict[str, Tensor]: + def __call__(self, videos, return_tensors: str, device=None, **kwargs) -> dict[str, Tensor]: assert return_tensors == "pt" if isinstance(videos, list): pixel_values = torch.stack([torch.as_tensor(v) for v in videos]) else: pixel_values = torch.as_tensor(videos).unsqueeze(0) + if device is not None: + pixel_values = pixel_values.to(device) return {"pixel_values_videos": pixel_values} diff --git a/tests/policies/vla_jepa/test_vla_jepa.py b/tests/policies/vla_jepa/test_vla_jepa.py index 70194dd59..a3e24a660 100644 --- a/tests/policies/vla_jepa/test_vla_jepa.py +++ b/tests/policies/vla_jepa/test_vla_jepa.py @@ -211,40 +211,42 @@ def test_reset_clears_action_queue(patch_vla_jepa_external_models: None) -> None def test_prepare_model_inputs_training_format(patch_vla_jepa_external_models: None) -> None: - from PIL import Image - policy = VLAJEPAPolicy(make_config()) - examples = policy._prepare_model_inputs(make_train_batch()) + inputs = policy._prepare_model_inputs(make_train_batch()) - assert len(examples) == BATCH_SIZE - for ex in examples: - assert set(ex) >= {"image", "video", "lang", "action", "state"} - assert len(ex["image"]) == 1 and isinstance(ex["image"][0], Image.Image) - assert ex["video"].ndim == 5 and ex["video"].dtype == np.uint8 # [V,T,H,W,C] - assert ex["action"].shape == (ACTION_HORIZON, ACTION_DIM) - assert ex["state"].shape == (1, STATE_DIM) + assert set(inputs) >= {"images", "instructions", "videos", "actions", "state"} + # images: per-sample, per-view [C, H, W] float tensors (kept as a list for Qwen messages) + assert len(inputs["images"]) == BATCH_SIZE and len(inputs["images"][0]) == 1 + img = inputs["images"][0][0] + assert isinstance(img, torch.Tensor) and img.dtype == torch.float32 and img.ndim == 3 + assert len(inputs["instructions"]) == BATCH_SIZE + # videos: batched [B, V, T, C, H, W] float + assert inputs["videos"].ndim == 6 and inputs["videos"].shape[0] == BATCH_SIZE + assert inputs["videos"].dtype == torch.float32 + assert inputs["actions"].shape == (BATCH_SIZE, ACTION_HORIZON, ACTION_DIM) + assert inputs["state"].shape == (BATCH_SIZE, 1, STATE_DIM) def test_prepare_model_inputs_inference_omits_action(patch_vla_jepa_external_models: None) -> None: policy = VLAJEPAPolicy(make_config()) - for ex in policy._prepare_model_inputs(make_inference_batch()): - assert "action" not in ex - assert "image" in ex and "video" in ex and "lang" in ex + inputs = policy._prepare_model_inputs(make_inference_batch()) + assert "actions" not in inputs and "action_is_pad" not in inputs + assert {"images", "instructions", "state"} <= set(inputs) def test_prepare_model_inputs_missing_task_uses_default(patch_vla_jepa_external_models: None) -> None: policy = VLAJEPAPolicy(make_config()) batch = make_inference_batch() del batch["task"] - examples = policy._prepare_model_inputs(batch) - assert all(isinstance(ex["lang"], str) and len(ex["lang"]) > 0 for ex in examples) + instructions = policy._prepare_model_inputs(batch)["instructions"] + assert all(isinstance(s, str) and len(s) > 0 for s in instructions) def test_prepare_model_inputs_string_task_broadcast(patch_vla_jepa_external_models: None) -> None: policy = VLAJEPAPolicy(make_config()) batch = make_inference_batch() batch["task"] = "open the drawer" - assert all(ex["lang"] == "open the drawer" for ex in policy._prepare_model_inputs(batch)) + assert policy._prepare_model_inputs(batch)["instructions"] == ["open the drawer"] * BATCH_SIZE def test_prepare_model_inputs_no_state_omitted(patch_vla_jepa_external_models: None) -> None: @@ -253,7 +255,7 @@ def test_prepare_model_inputs_no_state_omitted(patch_vla_jepa_external_models: N policy = VLAJEPAPolicy(make_config()) batch = make_inference_batch() del batch[OBS_STATE] - assert all("state" not in ex for ex in policy._prepare_model_inputs(batch)) + assert "state" not in policy._prepare_model_inputs(batch) # --------------------------------------------------------------------------- @@ -446,14 +448,14 @@ def test_postprocessor_applied_after_predict_action_chunk( """ from lerobot.policies.vla_jepa.processor_vla_jepa import make_vla_jepa_pre_post_processors - raw_actions = np.zeros((BATCH_SIZE, ACTION_HORIZON, ACTION_DIM), dtype=np.float32) + raw_actions = torch.zeros((BATCH_SIZE, ACTION_HORIZON, ACTION_DIM), dtype=torch.float32) cfg = make_config() cfg.clip_normalized_actions = False cfg.binarize_gripper_action = False policy = VLAJEPAPolicy(cfg) policy.eval() - monkeypatch.setattr(policy.model, "predict_action", lambda *a, **kw: raw_actions.copy()) + monkeypatch.setattr(policy.model, "predict_action", lambda *a, **kw: raw_actions.clone()) dataset_stats = _make_dataset_stats() _, postprocessor = make_vla_jepa_pre_post_processors(cfg, dataset_stats) @@ -564,9 +566,9 @@ def test_single_view_is_duplicated_for_world_model(patch_vla_jepa_external_model original_processor = policy.model.video_processor class _CapturingProcessor: - def __call__(self, videos: list, return_tensors: str) -> dict: + def __call__(self, videos: list, return_tensors: str, **kwargs) -> dict: captured_videos.extend(videos) - return original_processor(videos=videos, return_tensors=return_tensors) + return original_processor(videos=videos, return_tensors=return_tensors, **kwargs) policy.model.video_processor = _CapturingProcessor() policy.forward(_make_multiview_train_batch(num_views=1)) @@ -587,9 +589,9 @@ def test_excess_views_trimmed_for_world_model(patch_vla_jepa_external_models: No original_processor = policy.model.video_processor class _CapturingProcessor: - def __call__(self, videos: list, return_tensors: str) -> dict: + def __call__(self, videos: list, return_tensors: str, **kwargs) -> dict: captured_videos.extend(videos) - return original_processor(videos=videos, return_tensors=return_tensors) + return original_processor(videos=videos, return_tensors=return_tensors, **kwargs) policy.model.video_processor = _CapturingProcessor() policy.forward(_make_multiview_train_batch(num_views=3)) From 2f2b5679510a35aa83fdd8e9f986e134666618bc Mon Sep 17 00:00:00 2001 From: Khalil Meftah Date: Mon, 29 Jun 2026 18:52:59 +0200 Subject: [PATCH 22/38] Enable MolmoAct2 rollout on SO-100/101 with calibration correction (#3879) * fix(rollout): improve visual feature mismatch error with --rename_map hint * feat(policies): add joint frame transform and hardware deployment docs for MolmoAct2 Add MolmoAct2StateFrameTransformStep and MolmoAct2ActionFrameTransformStep processor steps for cross-calibration compatibility on SO-100/101. Add joint_signs and joint_offsets config fields. Add hardware deployment section to molmoact2.mdx with camera naming convention, joint frame correction, and safety guidance. * chore(docs): address PR comment * fix: address reviewer comments --- docs/source/molmoact2.mdx | 62 ++++++++++++ .../molmoact2/configuration_molmoact2.py | 13 +++ .../policies/molmoact2/processor_molmoact2.py | 95 +++++++++++++++++++ src/lerobot/rollout/context.py | 4 +- tests/policies/molmoact2/test_molmoact2.py | 35 +++++++ 5 files changed, 208 insertions(+), 1 deletion(-) diff --git a/docs/source/molmoact2.mdx b/docs/source/molmoact2.mdx index 9a377031c..9eb449ca9 100644 --- a/docs/source/molmoact2.mdx +++ b/docs/source/molmoact2.mdx @@ -386,6 +386,68 @@ These results demonstrate MolmoAct2's strong performance across diverse robotic manipulation tasks. To reproduce them, follow the instructions in the LIBERO evaluation section. +## Hardware Deployment (lerobot-rollout) + +LeRobot-format checkpoints are available on the Hub for direct use with +`lerobot-rollout`. Each checkpoint uses specific camera names that must +match your robot's camera configuration. + +### Camera naming convention + +Each checkpoint expects specific `observation.images.*` keys. +If your robot cameras have different names, use `--rename_map` to map them: + +| Checkpoint | Camera keys | Description | +| ----------------------------- | ---------------------- | ------------------------ | +| MolmoAct2-LIBERO-LeRobot | `image`, `wrist_image` | LIBERO sim cameras | +| MolmoAct2-BimanualYAM-LeRobot | `top`, `left`, `right` | YAM 3-camera setup | +| MolmoAct2-DROID-LeRobot | `cam0`, `cam1` | External + wrist | +| MolmoAct2-SO100_101-LeRobot | `cam0`, `cam1` | Primary + secondary view | + +Example with an SO-100 robot using top and side cameras: + +```bash +lerobot-rollout \ + --policy.path=lerobot/MolmoAct2-SO100_101-LeRobot \ + --rename_map='{"observation.images.top": "observation.images.cam0", "observation.images.side": "observation.images.cam1"}' \ + --robot.type=so100_follower \ + --robot.port=/dev/ttyACM0 \ + --robot.cameras='{ + top: {type: opencv, index_or_path: 0, width: 640, height: 480, fps: 30}, + side: {type: opencv, index_or_path: 2, width: 640, height: 480, fps: 30} + }' \ + --task="pick up the red cube" --duration=30 +``` + +To use a wrist camera instead, just change the rename mapping: + +```bash +--rename_map='{"observation.images.top": "observation.images.cam0", "observation.images.wrist": "observation.images.cam1"}' +``` + +### Joint frame transform (SO-100/101 zero-shot) + + +The MolmoAct2-SO100_101 checkpoint was trained on data that uses a different +joint calibration convention than LeRobot >= 0.5.0. Without a frame +correction, the arm may move in the wrong direction. + +This affects both **zero-shot deployment** and **fine-tuning** from the +original checkpoint. The pretrained weights expect the old convention, so +all joint data (observations and actions) must be transformed to match. + +The converted LeRobot checkpoint (`lerobot/MolmoAct2-SO100_101-LeRobot`) +already includes this correction in its processor pipeline. If you convert +or fine-tune the checkpoint yourself, set the following in the policy config (`configuration_molmoact2.py`): + +- `joint_signs`: `[1, -1, 1, 1, 1, 1]` (flips shoulder_lift direction) +- `joint_offsets`: `[0, 90, 90, 0, 0, 0]` (shifts shoulder_lift and elbow_flex by 90°) + +See the [backward compatibility guide](./backwardcomp) for details on the +calibration change. + + + ## Differences From the Original Implementation This LeRobot port is intended to match MolmoAct2 behavior while using LeRobot's diff --git a/src/lerobot/policies/molmoact2/configuration_molmoact2.py b/src/lerobot/policies/molmoact2/configuration_molmoact2.py index 53aefdee6..bf9437ba9 100644 --- a/src/lerobot/policies/molmoact2/configuration_molmoact2.py +++ b/src/lerobot/policies/molmoact2/configuration_molmoact2.py @@ -79,6 +79,15 @@ class MolmoAct2Config(PreTrainedConfig): eval_seed: int | None = None rtc_config: RTCConfig | None = None + # Joint frame transform for cross-calibration compatibility. + # Some MolmoAct2 checkpoints were trained on data using a different joint + # convention than the current LeRobot calibration. Set both to apply a + # sign/offset correction at runtime (state before model, action after). + # See: https://huggingface.co/docs/lerobot/backwardcomp + # Default is None (no transform). Both must be set together. + joint_signs: list[float] | None = None + joint_offsets: list[float] | None = None + # Default is full finetuning with gradients from the action expert flowing into the VLM. enable_lora_vlm: bool = False lora_rank: int = 64 @@ -123,6 +132,10 @@ class MolmoAct2Config(PreTrainedConfig): def __post_init__(self) -> None: super().__post_init__() + if (self.joint_signs is None) != (self.joint_offsets is None): + raise ValueError("joint_signs and joint_offsets must both be set or both be None.") + if self.joint_signs is not None and len(self.joint_signs) != len(self.joint_offsets): + raise ValueError("joint_signs and joint_offsets must have the same length.") if self.action_mode not in {"continuous", "discrete", "both"}: raise ValueError( f"Unsupported action_mode={self.action_mode!r}. " diff --git a/src/lerobot/policies/molmoact2/processor_molmoact2.py b/src/lerobot/policies/molmoact2/processor_molmoact2.py index 1303e94a1..d2db817ef 100644 --- a/src/lerobot/policies/molmoact2/processor_molmoact2.py +++ b/src/lerobot/policies/molmoact2/processor_molmoact2.py @@ -1005,6 +1005,93 @@ class MolmoAct2PackInputsProcessorStep(ProcessorStep): return features +@ProcessorStepRegistry.register(name="molmoact2_state_frame_transform") +@dataclass +class MolmoAct2StateFrameTransformStep(ProcessorStep): + """Convert robot state from arm frame to model frame before normalization. + + Required for zero-shot deployment of MolmoAct2-SO100_101 on SO-100/101 + arms calibrated with LeRobot >= 0.5.0 (v3.0 convention). The checkpoint + was trained on data using a different joint convention (sign flip on + shoulder_lift, 90 deg offset on shoulder_lift and elbow_flex). + + No-op when joint_signs and joint_offsets are None (default), so this + step has no effect on fine-tuned models or other embodiments. + + state_model = signs * arm_state + offsets + + See: https://huggingface.co/docs/lerobot/backwardcomp + """ + + joint_signs: list[float] | None = None + joint_offsets: list[float] | None = None + + def __call__(self, transition: EnvTransition) -> EnvTransition: + if self.joint_signs is None or self.joint_offsets is None: + return transition + observation = transition.get(TransitionKey.OBSERVATION) + if not isinstance(observation, dict) or OBS_STATE not in observation: + return transition + transition = transition.copy() + observation = observation.copy() + state = torch.as_tensor(observation[OBS_STATE], dtype=torch.float32).clone() + n = len(self.joint_signs) + signs = torch.tensor(self.joint_signs, dtype=torch.float32, device=state.device) + offsets = torch.tensor(self.joint_offsets, dtype=torch.float32, device=state.device) + state[..., :n] = signs * state[..., :n] + offsets + observation[OBS_STATE] = state + transition[TransitionKey.OBSERVATION] = observation + return transition + + def transform_features( + self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]] + ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]: + return features + + def get_config(self) -> dict[str, Any]: + return {"joint_signs": self.joint_signs, "joint_offsets": self.joint_offsets} + + +@ProcessorStepRegistry.register(name="molmoact2_action_frame_transform") +@dataclass +class MolmoAct2ActionFrameTransformStep(ProcessorStep): + """Convert model action from model frame back to arm frame after unnormalization. + + Inverse of MolmoAct2StateFrameTransformStep. Required for zero-shot + MolmoAct2-SO100_101 on SO-100/101 arms. No-op when both fields are None. + + action_arm = signs * (model_action - offsets) + + See: https://huggingface.co/docs/lerobot/backwardcomp + """ + + joint_signs: list[float] | None = None + joint_offsets: list[float] | None = None + + def __call__(self, transition: EnvTransition) -> EnvTransition: + if self.joint_signs is None or self.joint_offsets is None: + return transition + action = transition.get(TransitionKey.ACTION) + if action is None: + return transition + transition = transition.copy() + action = torch.as_tensor(action, dtype=torch.float32).clone() + n = len(self.joint_signs) + signs = torch.tensor(self.joint_signs, dtype=torch.float32, device=action.device) + offsets = torch.tensor(self.joint_offsets, dtype=torch.float32, device=action.device) + action[..., :n] = signs * (action[..., :n] - offsets) + transition[TransitionKey.ACTION] = action + return transition + + def transform_features( + self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]] + ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]: + return features + + def get_config(self) -> dict[str, Any]: + return {"joint_signs": self.joint_signs, "joint_offsets": self.joint_offsets} + + @ProcessorStepRegistry.register(name="molmoact2_clamp_action") @dataclass class MolmoAct2ClampActionProcessorStep(ProcessorStep): @@ -1067,6 +1154,10 @@ def make_molmoact2_pre_post_processors( input_steps: list[ProcessorStep] = [ RenameObservationsProcessorStep(rename_map={}), AddBatchDimensionProcessorStep(), + MolmoAct2StateFrameTransformStep( + joint_signs=config.joint_signs, + joint_offsets=config.joint_offsets, + ), MolmoAct2MaskedNormalizerProcessorStep( features={**config.input_features, **config.output_features}, norm_map=config.normalization_mapping, @@ -1102,6 +1193,10 @@ def make_molmoact2_pre_post_processors( norm_map=config.normalization_mapping, stats=masked_dataset_stats, ), + MolmoAct2ActionFrameTransformStep( + joint_signs=config.joint_signs, + joint_offsets=config.joint_offsets, + ), DeviceProcessorStep(device="cpu"), ] diff --git a/src/lerobot/rollout/context.py b/src/lerobot/rollout/context.py index 62d844932..20a7d715a 100644 --- a/src/lerobot/rollout/context.py +++ b/src/lerobot/rollout/context.py @@ -320,7 +320,9 @@ def build_rollout_context( raise ValueError( f"Visual feature mismatch between policy and robot hardware.\n" f"Policy expects: {expected_visuals}\n" - f"Robot provides: {provided_visuals}" + f"Robot provides: {provided_visuals}\n" + f"Use --rename_map to map camera names, e.g. " + f"""--rename_map='{{"observation.images.top": "observation.images.cam0"}}'""" ) # --- 5. Dataset ------------- diff --git a/tests/policies/molmoact2/test_molmoact2.py b/tests/policies/molmoact2/test_molmoact2.py index 5fba72913..095b73180 100644 --- a/tests/policies/molmoact2/test_molmoact2.py +++ b/tests/policies/molmoact2/test_molmoact2.py @@ -44,10 +44,12 @@ from lerobot.policies.molmoact2.modeling_molmoact2 import ( _combine_rollout_seeds, ) from lerobot.policies.molmoact2.processor_molmoact2 import ( + MolmoAct2ActionFrameTransformStep, MolmoAct2ClampNormalizedProcessorStep, MolmoAct2MaskedNormalizerProcessorStep, MolmoAct2MaskedUnnormalizerProcessorStep, MolmoAct2PackInputsProcessorStep, + MolmoAct2StateFrameTransformStep, _add_gripper_masks_to_stats, _build_discrete_state_string, _normalize_question_text, @@ -926,6 +928,39 @@ def test_question_normalization_matches_release_prompt_style(): ) +def test_joint_frame_transform_round_trip(): + signs = [1.0, -1.0, 1.0, 1.0, 1.0, 1.0] + offsets = [0.0, 90.0, 90.0, 0.0, 0.0, 0.0] + original_state = torch.tensor([[10.0, -90.0, -120.0, 30.0, 0.0, -45.0]]) + + state_step = MolmoAct2StateFrameTransformStep(joint_signs=signs, joint_offsets=offsets) + action_step = MolmoAct2ActionFrameTransformStep(joint_signs=signs, joint_offsets=offsets) + + transition = { + TransitionKey.OBSERVATION: {OBS_STATE: original_state.clone()}, + } + transformed = state_step(transition) + model_state = transformed[TransitionKey.OBSERVATION][OBS_STATE] + + action_transition = {TransitionKey.ACTION: model_state.clone()} + recovered = action_step(action_transition) + recovered_state = recovered[TransitionKey.ACTION] + + assert torch.allclose(recovered_state, original_state) + + +def test_joint_frame_transform_noop_when_none(): + state_step = MolmoAct2StateFrameTransformStep(joint_signs=None, joint_offsets=None) + action_step = MolmoAct2ActionFrameTransformStep(joint_signs=None, joint_offsets=None) + state = torch.tensor([[10.0, -90.0, -120.0]]) + + state_transition = {TransitionKey.OBSERVATION: {OBS_STATE: state}} + assert state_step(state_transition) is state_transition + + action_transition = {TransitionKey.ACTION: state} + assert action_step(action_transition) is action_transition + + def test_action_padding_marks_only_real_dimensions(): step = object.__new__(MolmoAct2PackInputsProcessorStep) step.max_action_dim = 32 From 0da98afd639dc3e4977f10591ee940ab833e1361 Mon Sep 17 00:00:00 2001 From: Khalil Meftah Date: Tue, 30 Jun 2026 17:17:50 +0200 Subject: [PATCH 23/38] Feat(robot): add MIT control mode to ReBot (#3778) * fix(config): update joint limits for RebotB601Follower and RebotArm102Leader * feat(config): add MIT control mode ReBot - Add configurable arm control mode (mit default, pos_vel fallback) with tunable mit_kp / mit_kd - Add optional gripper control mode (force_pos default, mit optional) with gripper_mit_kp / gripper_mit_kd - Update tests for MIT arm routing, gripper mode routing, and revised joint limits * fix(robots): restore joint clipping and wrist_yaw fallback in ReBot B601 send_action * feat(robot): increase gripper velocity and torque for rebot arm --- .../bi_rebot_b601_follower.py | 12 +++++ .../config_rebot_b601_follower.py | 29 ++++++++--- .../rebot_b601_follower.py | 48 ++++++++++++++----- .../config_rebot_102_leader.py | 2 +- tests/robots/test_rebot_b601_follower.py | 23 +++++++-- 5 files changed, 92 insertions(+), 22 deletions(-) diff --git a/src/lerobot/robots/bi_rebot_b601_follower/bi_rebot_b601_follower.py b/src/lerobot/robots/bi_rebot_b601_follower/bi_rebot_b601_follower.py index c320cee8b..7780ac0fe 100644 --- a/src/lerobot/robots/bi_rebot_b601_follower/bi_rebot_b601_follower.py +++ b/src/lerobot/robots/bi_rebot_b601_follower/bi_rebot_b601_follower.py @@ -65,7 +65,13 @@ class BiRebotB601Follower(BimanualMixin, Robot): cameras=left_arm_cameras, motor_can_ids=config.left_arm_config.motor_can_ids, pos_vel_velocity=config.left_arm_config.pos_vel_velocity, + control_mode=config.left_arm_config.control_mode, + mit_kp=config.left_arm_config.mit_kp, + mit_kd=config.left_arm_config.mit_kd, + gripper_control_mode=config.left_arm_config.gripper_control_mode, gripper_torque_ratio=config.left_arm_config.gripper_torque_ratio, + gripper_mit_kp=config.left_arm_config.gripper_mit_kp, + gripper_mit_kd=config.left_arm_config.gripper_mit_kd, joint_limits=config.left_arm_config.joint_limits, ) @@ -80,7 +86,13 @@ class BiRebotB601Follower(BimanualMixin, Robot): cameras=config.right_arm_config.cameras, motor_can_ids=config.right_arm_config.motor_can_ids, pos_vel_velocity=config.right_arm_config.pos_vel_velocity, + control_mode=config.right_arm_config.control_mode, + mit_kp=config.right_arm_config.mit_kp, + mit_kd=config.right_arm_config.mit_kd, + gripper_control_mode=config.right_arm_config.gripper_control_mode, gripper_torque_ratio=config.right_arm_config.gripper_torque_ratio, + gripper_mit_kp=config.right_arm_config.gripper_mit_kp, + gripper_mit_kd=config.right_arm_config.gripper_mit_kd, joint_limits=config.right_arm_config.joint_limits, ) diff --git a/src/lerobot/robots/rebot_b601_follower/config_rebot_b601_follower.py b/src/lerobot/robots/rebot_b601_follower/config_rebot_b601_follower.py index 096548afb..88af07e59 100644 --- a/src/lerobot/robots/rebot_b601_follower/config_rebot_b601_follower.py +++ b/src/lerobot/robots/rebot_b601_follower/config_rebot_b601_follower.py @@ -65,18 +65,33 @@ class RebotB601FollowerConfig: } ) - # Target velocity for joints running in POS_VEL mode, in degrees/s. A scalar is - # applied to every joint; a list provides one value per joint (in motor order). - pos_vel_velocity: float | list[float] = field(default_factory=lambda: [150.0] * 7) + # Max speed (deg/s) per joint for POS_VEL arms and FORCE_POS gripper (motor order). + pos_vel_velocity: float | list[float] = field( + default_factory=lambda: [150.0, 150.0, 150.0, 150.0, 150.0, 150.0, 900.0] + ) - # Torque/current ratio for the gripper's FORCE_POS mode, in range [0, 1]. - gripper_torque_ratio: float = 0.1 + # Arm control: "mit" or "pos_vel". + control_mode: str = "mit" + + # MIT kp/kd per arm joint (motor order). Unused when control_mode="pos_vel". + mit_kp: float | list[float] = field(default_factory=lambda: [45.0, 45.0, 45.0, 8.0, 9.0, 8.0, 8.0]) + mit_kd: float | list[float] = field(default_factory=lambda: [12.0, 12.0, 12.0, 1.0, 1.0, 1.0, 1.0]) + + # Gripper control: "force_pos" or "mit". + gripper_control_mode: str = "force_pos" + + # FORCE_POS only: max grip force, in [0, 1]. + gripper_torque_ratio: float = 0.07 + + # MIT only. + gripper_mit_kp: float = 8.0 + gripper_mit_kd: float = 0.3 # Soft joint limits (degrees). These are clipped against on every action. joint_limits: dict[str, tuple[float, float]] = field( default_factory=lambda: { - "shoulder_pan": (-145.0, 145.0), - "shoulder_lift": (-170.0, 1.0), + "shoulder_pan": (-150.0, 150.0), + "shoulder_lift": (-200.0, 1.0), "elbow_flex": (-200.0, 1.0), "wrist_flex": (-80.0, 90.0), "wrist_yaw": (-90.0, 90.0), diff --git a/src/lerobot/robots/rebot_b601_follower/rebot_b601_follower.py b/src/lerobot/robots/rebot_b601_follower/rebot_b601_follower.py index bf989702b..bff0b86c1 100644 --- a/src/lerobot/robots/rebot_b601_follower/rebot_b601_follower.py +++ b/src/lerobot/robots/rebot_b601_follower/rebot_b601_follower.py @@ -174,11 +174,25 @@ class RebotB601Follower(Robot): print(f"Calibration saved to {self.calibration_fpath}") def configure(self) -> None: + if self.config.control_mode not in ("pos_vel", "mit"): + raise ValueError( + f"Unsupported control_mode '{self.config.control_mode}'. Use 'pos_vel' or 'mit'." + ) + if self.config.gripper_control_mode not in ("force_pos", "mit"): + raise ValueError( + f"Unsupported gripper_control_mode '{self.config.gripper_control_mode}'. " + "Use 'force_pos' or 'mit'." + ) + use_mit = self.config.control_mode == "mit" + gripper_use_mit = self.config.gripper_control_mode == "mit" self.bus.enable_all() for motor_name, motor in self.motors.items(): - target_mode = ( - MotorBridgeMode.FORCE_POS if motor_name == GRIPPER_MOTOR else MotorBridgeMode.POS_VEL - ) + if motor_name == GRIPPER_MOTOR: + target_mode = MotorBridgeMode.MIT if gripper_use_mit else MotorBridgeMode.FORCE_POS + elif use_mit: + target_mode = MotorBridgeMode.MIT + else: + target_mode = MotorBridgeMode.POS_VEL for attempt in range(_ENSURE_MODE_RETRIES + 1): try: motor.ensure_mode(target_mode) @@ -264,22 +278,34 @@ class RebotB601Follower(Robot): goal_present_pos = {key: (g, present_pos.get(key, g)) for key, g in goal_pos.items()} goal_pos = ensure_safe_goal_position(goal_present_pos, self.config.max_relative_target) + use_mit = self.config.control_mode == "mit" for motor_name, position_deg in goal_pos.items(): motor = self.motors.get(motor_name) if motor is None: continue idx = self.motor_names.index(motor_name) - vel_deg_s = ( - self.config.pos_vel_velocity[idx] - if isinstance(self.config.pos_vel_velocity, list) - else self.config.pos_vel_velocity - ) pos_rad = math.radians(position_deg) - vel_rad = math.radians(vel_deg_s) if motor_name == GRIPPER_MOTOR: - motor.send_force_pos(pos_rad, vel_rad, self.config.gripper_torque_ratio) + if self.config.gripper_control_mode == "mit": + motor.send_mit(pos_rad, 0.0, self.config.gripper_mit_kp, self.config.gripper_mit_kd, 0.0) + else: + vel_deg_s = ( + self.config.pos_vel_velocity[idx] + if isinstance(self.config.pos_vel_velocity, list) + else self.config.pos_vel_velocity + ) + motor.send_force_pos(pos_rad, math.radians(vel_deg_s), self.config.gripper_torque_ratio) + elif use_mit: + kp = self.config.mit_kp[idx] if isinstance(self.config.mit_kp, list) else self.config.mit_kp + kd = self.config.mit_kd[idx] if isinstance(self.config.mit_kd, list) else self.config.mit_kd + motor.send_mit(pos_rad, 0.0, kp, kd, 0.0) else: - motor.send_pos_vel(pos_rad, vel_rad) + vel_deg_s = ( + self.config.pos_vel_velocity[idx] + if isinstance(self.config.pos_vel_velocity, list) + else self.config.pos_vel_velocity + ) + motor.send_pos_vel(pos_rad, math.radians(vel_deg_s)) return {f"{motor}.pos": val for motor, val in goal_pos.items()} diff --git a/src/lerobot/teleoperators/rebot_102_leader/config_rebot_102_leader.py b/src/lerobot/teleoperators/rebot_102_leader/config_rebot_102_leader.py index d1beea2ed..81a1c2c39 100644 --- a/src/lerobot/teleoperators/rebot_102_leader/config_rebot_102_leader.py +++ b/src/lerobot/teleoperators/rebot_102_leader/config_rebot_102_leader.py @@ -65,7 +65,7 @@ class RebotArm102LeaderConfig: joint_ranges: dict[str, list[int]] = field( default_factory=lambda: { "shoulder_pan": [-150, 150], - "shoulder_lift": [-170, 1], + "shoulder_lift": [-200, 1], "elbow_flex": [-200, 1], "wrist_flex": [-80, 90], "wrist_yaw": [-90, 90], diff --git a/tests/robots/test_rebot_b601_follower.py b/tests/robots/test_rebot_b601_follower.py index 553675be0..0861827be 100644 --- a/tests/robots/test_rebot_b601_follower.py +++ b/tests/robots/test_rebot_b601_follower.py @@ -91,10 +91,11 @@ def test_get_observation_converts_to_degrees(follower): def test_send_action_clips_to_joint_limits(follower): - # shoulder_pan limit is (-145, 145); request beyond the upper bound. + # shoulder_pan limit is (-150, 150); request beyond the upper bound. returned = follower.send_action({"shoulder_pan.pos": 999.0}) - assert returned["shoulder_pan.pos"] == 145.0 - follower.motors["shoulder_pan"].send_pos_vel.assert_called_once() + assert returned["shoulder_pan.pos"] == 150.0 + # Default control_mode is "mit", so arm joints are driven via send_mit. + follower.motors["shoulder_pan"].send_mit.assert_called_once() def test_send_action_routes_gripper_to_force_pos(follower): @@ -103,6 +104,22 @@ def test_send_action_routes_gripper_to_force_pos(follower): follower.motors["gripper"].send_pos_vel.assert_not_called() +def test_gripper_mit_mode_routes_to_send_mit(): + bus_mock = _make_bus_mock() + with ( + patch(f"{_MODULE}.require_package", lambda *a, **kw: None), + patch(f"{_MODULE}.MotorBridgeController") as controller_cls, + patch(f"{_MODULE}.MotorBridgeMode", MagicMock()), + ): + controller_cls.from_dm_serial.return_value = bus_mock + cfg = RebotB601FollowerRobotConfig(port="/dev/null", gripper_control_mode="mit") + robot = RebotB601Follower(cfg) + robot.connect(calibrate=False) + robot.send_action({"gripper.pos": -10.0}) + robot.motors["gripper"].send_mit.assert_called_once() + robot.motors["gripper"].send_force_pos.assert_not_called() + + def test_bimanual_prefixes_features(): with patch(f"{_MODULE}.require_package", lambda *a, **kw: None): cfg = BiRebotB601FollowerConfig( From 8414188db0b178b947985a7a9a91314708837315 Mon Sep 17 00:00:00 2001 From: Caroline Pascal Date: Tue, 30 Jun 2026 20:21:06 +0200 Subject: [PATCH 24/38] fix(datasets dependency): removing datasets dependency in pretrained.py (#3897) --- src/lerobot/policies/pretrained.py | 56 ++++++++++++++-------------- src/lerobot/scripts/lerobot_train.py | 4 +- 2 files changed, 29 insertions(+), 31 deletions(-) diff --git a/src/lerobot/policies/pretrained.py b/src/lerobot/policies/pretrained.py index aea5f1b08..702569b8c 100644 --- a/src/lerobot/policies/pretrained.py +++ b/src/lerobot/policies/pretrained.py @@ -11,6 +11,8 @@ # 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. +from __future__ import annotations + import abc import builtins import dataclasses @@ -19,7 +21,7 @@ import os from importlib.resources import files from pathlib import Path from tempfile import TemporaryDirectory -from typing import TypedDict, TypeVar, Unpack +from typing import TYPE_CHECKING, TypedDict, TypeVar, Unpack import packaging import safetensors @@ -38,10 +40,13 @@ from .utils import log_model_loading_keys T = TypeVar("T", bound="PreTrainedPolicy") +if TYPE_CHECKING: + from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata + def _build_card_context( cfg: TrainPipelineConfig | None, - dataset_repo_id: str | None, + dataset_meta: LeRobotDatasetMetadata | None, input_features: dict | None, output_features: dict | None, ) -> dict: @@ -72,30 +77,16 @@ def _build_card_context( "lerobot_version": __version__, } - if dataset_repo_id: - dataset_cfg = getattr(cfg, "dataset", None) - try: - from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata - - meta = LeRobotDatasetMetadata( - dataset_repo_id, - root=getattr(dataset_cfg, "root", None), - revision=getattr(dataset_cfg, "revision", None), - ) - context["dataset"] = { - "repo_id": dataset_repo_id, - "episodes": meta.total_episodes, - "frames": meta.total_frames, - "fps": meta.fps, - "tasks": [str(task) for task in meta.tasks.index], - } - context["robot_type"] = meta.robot_type - context["cameras"] = [key.split(".")[-1] for key in meta.camera_keys] - except Exception as e: # noqa: BLE001 — dataset details are optional, never fail the push - logging.warning( - f"Could not load dataset metadata for '{dataset_repo_id}'; those sections will be " - f"omitted from the model card. ({e})" - ) + if dataset_meta is not None: + context["dataset"] = { + "repo_id": dataset_meta.repo_id, + "episodes": dataset_meta.total_episodes, + "frames": dataset_meta.total_frames, + "fps": dataset_meta.fps, + "tasks": [str(task) for task in dataset_meta.tasks.index], + } + context["robot_type"] = dataset_meta.robot_type + context["cameras"] = [key.split(".")[-1] for key in dataset_meta.camera_keys] return context @@ -304,6 +295,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC): cfg: TrainPipelineConfig, peft_model=None, state_dict: dict[str, Tensor] | None = None, + dataset_meta: LeRobotDatasetMetadata | None = None, ): api = HfApi() repo_id = api.create_repo( @@ -325,7 +317,12 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC): self.save_pretrained(saved_path, state_dict=state_dict) card = self.generate_model_card( - cfg.dataset.repo_id, self.config.type, self.config.license, self.config.tags, cfg=cfg + cfg.dataset.repo_id, + self.config.type, + self.config.license, + self.config.tags, + cfg=cfg, + dataset_meta=dataset_meta, ) card.save(str(saved_path / "README.md")) @@ -352,6 +349,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC): license: str | None, tags: list[str] | None, cfg: TrainPipelineConfig | None = None, + dataset_meta: LeRobotDatasetMetadata | None = None, ) -> ModelCard: base_model_mapping = { "smolvla": "lerobot/smolvla_base", @@ -372,7 +370,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC): ) context = _build_card_context( - cfg, dataset_repo_id, self.config.input_features, self.config.output_features + cfg, dataset_meta, self.config.input_features, self.config.output_features ) # Used by the template to pre-fill commands and the "Fine-tuned from" line. context["policy_repo_id"] = getattr(self.config, "repo_id", None) @@ -389,7 +387,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC): self, peft_config=None, peft_cli_overrides: dict | None = None, - ) -> "PreTrainedPolicy": + ) -> PreTrainedPolicy: """ Wrap this policy with PEFT adapters for parameter-efficient fine-tuning. diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index f2a152df9..44c94a1eb 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/src/lerobot/scripts/lerobot_train.py @@ -736,9 +736,9 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): unwrapped_model = accelerator.unwrap_model(policy) # PEFT only applies when training a policy — reward models use the plain path. if not cfg.is_reward_model_training and cfg.policy.use_peft: - unwrapped_model.push_model_to_hub(cfg, peft_model=unwrapped_model) + unwrapped_model.push_model_to_hub(cfg, peft_model=unwrapped_model, dataset_meta=dataset.meta) else: - unwrapped_model.push_model_to_hub(cfg, state_dict=model_state_dict) + unwrapped_model.push_model_to_hub(cfg, state_dict=model_state_dict, dataset_meta=dataset.meta) preprocessor.push_to_hub(active_cfg.repo_id) postprocessor.push_to_hub(active_cfg.repo_id) From 141c3532068e71a5959b7c9fd7427adec1b37c29 Mon Sep 17 00:00:00 2001 From: Maxime Ellerbach Date: Wed, 1 Jul 2026 14:35:57 +0200 Subject: [PATCH 25/38] feat(policies): Add FastWAM Policy (#3834) * Add FastWAM policy * Add FastWAM policy review updates * big refactor to use models from diffusers and transformers * changing reproducable results * preparing for training adding some temporary debug code aswell to visualize model output * re-parenting of some layers to enable proper zero-3 FSDP * linting * small fix for the preprocessor and padded images * removing some preprocessors * removing temporary debug code * cleaning up * updating uv lock after rebasing * adding lazy imports * linting * fixing stale assertion * make tokenizer/text-encoder model ids configurable + some nits * moving and renaming files to have a cleaner file tree * removed asserts from the model, added guard instead and completely removed useless asserts * cleaning up imports * removing is_main_process and custom logging logic * removing unused / stale attention path, removing some of the stale forwards within wan/models --------- Co-authored-by: ZibinDong Co-authored-by: Steven Palma --- docs/source/_toctree.yml | 2 + docs/source/fastwam.mdx | 167 ++ docs/source/policy_fastwam_README.md | 56 + pyproject.toml | 8 +- src/lerobot/policies/__init__.py | 2 + src/lerobot/policies/factory.py | 15 + src/lerobot/policies/fastwam/README.md | 1 + src/lerobot/policies/fastwam/__init__.py | 23 + .../policies/fastwam/configuration_fastwam.py | 399 ++++ .../policies/fastwam/modeling_fastwam.py | 440 ++++ .../policies/fastwam/processor_fastwam.py | 142 ++ src/lerobot/policies/fastwam/wan/README.md | 34 + src/lerobot/policies/fastwam/wan/__init__.py | 33 + src/lerobot/policies/fastwam/wan/adapters.py | 108 + .../policies/fastwam/wan/components.py | 175 ++ src/lerobot/policies/fastwam/wan/model.py | 341 +++ src/lerobot/policies/fastwam/wan/modular.py | 1912 +++++++++++++++++ src/lerobot/policies/fastwam/wan/video_dit.py | 800 +++++++ tests/policies/fastwam/test_fastwam_policy.py | 391 ++++ uv.lock | 11 +- 20 files changed, 5057 insertions(+), 3 deletions(-) create mode 100644 docs/source/fastwam.mdx create mode 100644 docs/source/policy_fastwam_README.md create mode 120000 src/lerobot/policies/fastwam/README.md create mode 100644 src/lerobot/policies/fastwam/__init__.py create mode 100644 src/lerobot/policies/fastwam/configuration_fastwam.py create mode 100644 src/lerobot/policies/fastwam/modeling_fastwam.py create mode 100644 src/lerobot/policies/fastwam/processor_fastwam.py create mode 100644 src/lerobot/policies/fastwam/wan/README.md create mode 100644 src/lerobot/policies/fastwam/wan/__init__.py create mode 100644 src/lerobot/policies/fastwam/wan/adapters.py create mode 100644 src/lerobot/policies/fastwam/wan/components.py create mode 100644 src/lerobot/policies/fastwam/wan/model.py create mode 100644 src/lerobot/policies/fastwam/wan/modular.py create mode 100644 src/lerobot/policies/fastwam/wan/video_dit.py create mode 100644 tests/policies/fastwam/test_fastwam_policy.py diff --git a/docs/source/_toctree.yml b/docs/source/_toctree.yml index 5d847a94d..dcd14e131 100644 --- a/docs/source/_toctree.yml +++ b/docs/source/_toctree.yml @@ -69,6 +69,8 @@ title: VLA-JEPA - local: eo1 title: EO-1 + - local: fastwam + title: FastWAM - local: groot title: NVIDIA GR00T N1.5 - local: xvla diff --git a/docs/source/fastwam.mdx b/docs/source/fastwam.mdx new file mode 100644 index 000000000..18b4775f8 --- /dev/null +++ b/docs/source/fastwam.mdx @@ -0,0 +1,167 @@ +# FastWAM + +FastWAM is a World Action Model policy for robot control. The LeRobot integration exposes FastWAM through the standard policy API so it can be configured with `policy.type=fastwam`, trained with `lerobot-train`, and loaded through the LeRobot pretrained policy interface. + +## Model Overview + +FastWAM keeps video modeling during training, but uses direct action prediction at inference time instead of iteratively generating future observations. This LeRobot policy wraps the FastWAM action model, adapts LeRobot batches to FastWAM training samples, and provides the standard processor pipeline for normalization and action postprocessing. + +The implementation initializes the visual world-model components from `Wan-AI/Wan2.2-TI2V-5B` by default and predicts action chunks with shape `[batch, action_horizon, action_dim]`. + +### What the LeRobot Integration Covers + +- Standard `policy.type=fastwam` configuration through LeRobot +- Image, state, action, and language-task batch adaptation +- Action chunk inference through `select_action` and `predict_action_chunk` +- Checkpoint save/load through the LeRobot policy APIs +- Configurable LIBERO gripper action postprocessing + +## Installation Requirements + +Install LeRobot from source, then install FastWAM dependencies: + +```bash +pip install -e ".[fastwam]" +``` + +This installs the FastWAM policy extra from `pyproject.toml`: `transformers`, +`diffusers`, `ftfy`, and `regex`, plus LeRobot's base dependencies. + +For LIBERO evaluation, install the benchmark dependencies too: + +```bash +pip install -e ".[fastwam,libero]" +``` + +This installs both extras. In addition to the FastWAM dependencies above, the +`libero` extra installs LeRobot dataset dependencies, `hf-libero` on Linux, and +`scipy`. + +FastWAM uses the Wan2.2 TI2V backbone. The default model id is: + +```python +policy.model_id=Wan-AI/Wan2.2-TI2V-5B +``` + +## Data Requirements + +FastWAM expects a LeRobot dataset with: + +- one or more visual observations whose widths concatenate to `policy.image_size[1]` +- `observation.state` when `policy.proprio_dim` is not `None` +- `action` +- a language task instruction through the dataset task field, or precomputed `context` and `context_mask` tensors + +The default visual setup is one image feature named `observation.images.image` with shape `(3, 224, 448)`. If the dataset uses two cameras, configure `policy.input_features` so their heights match `224` and their widths sum to `448`. + +## Usage + +Create a new FastWAM policy with: + +```bash +lerobot-train \ + --dataset.repo_id=your-org/your-dataset \ + --policy.type=fastwam \ + --policy.action_dim=7 \ + --policy.proprio_dim=8 \ + --policy.action_horizon=32 \ + --policy.n_action_steps=10 \ + --policy.image_size='[224,448]' \ + --output_dir=./outputs/fastwam_training \ + --job_name=fastwam_training \ + --steps=300000 \ + --batch_size=8 \ + --policy.device=cuda +``` + +Evaluate an existing LeRobot-format checkpoint on LIBERO-10 with: + +```bash +lerobot-eval \ + --policy.path=ZibinDong/fastwam_libero_uncond_2cam224 \ + --policy.device=cuda \ + --policy.torch_dtype=float32 \ + --policy.n_action_steps=10 \ + --env.type=libero \ + --env.task=libero_10 \ + --env.observation_height=224 \ + --env.observation_width=224 \ + --eval.batch_size=1 \ + --eval.n_episodes=50 \ + --seed=0 \ + --env.episode_length=600 +``` + +For `libero_goal`, `libero_spatial`, and `libero_object`, use +`--env.episode_length=300`. + +For real-robot rollout, use the same checkpoint path: + +```bash +lerobot-rollout \ + --robot.type=so101_follower \ + --robot.port=/dev/ttyACM0 \ + --policy.path=your-org/fastwam-real-robot +``` + +## Configuration Notes + +### Image Features + +`policy.image_size` is the size of the concatenated FastWAM image tensor as `(height, width)`. Each configured image feature must have shape `(3, height, camera_width)`, and all camera widths must sum to the configured width. + +### Action Chunking + +`policy.action_horizon` controls the number of future actions supervised during training and predicted during inference. `policy.n_action_steps` controls how many actions are consumed before the policy predicts a fresh chunk. `policy.n_action_steps` must be less than or equal to `policy.action_horizon`. + +### Wan Components + +FastWAM loads the Wan VAE, video DiT, text encoder, and tokenizer from the configured Wan model directory or Hugging Face Hub model id. LeRobot-format FastWAM checkpoints saved by `save_pretrained` also copy the local Wan component files needed by `from_pretrained`. + +### Attention Backend + +FastWAM's DiT uses PyTorch's `scaled_dot_product_attention` (SDPA) for all attention. It does **not** use FlashAttention: its Mixture-of-Transformers (MoT) routing needs arbitrary boolean `[query, key]` attention masks, which the FlashAttention varlen API cannot express. Installing the `flash-attn` package therefore has no effect on the FastWAM path. (Note that SDPA itself may still select PyTorch's own flash / memory-efficient / math kernel internally — this is unrelated to the `flash-attn` package.) + +### LIBERO Action Toggle + +FastWAM LIBERO checkpoints use `policy.toggle_action_dimensions=[-1]` by +default to match the gripper action convention used by the original FastWAM +evaluation pipeline: + +```bash +--policy.toggle_action_dimensions='[-1]' +``` + +## Results + +Evaluated on LIBERO with [`ZibinDong/fastwam_libero_uncond_2cam224`](https://huggingface.co/ZibinDong/fastwam_libero_uncond_2cam224): + +| Suite | Success rate | n_episodes | +| -------------- | -----------: | ---------: | +| libero_spatial | 97.6% | 500 | +| libero_object | 99.0% | 500 | +| libero_goal | 95.0% | 500 | +| libero_10 | 94.0% | 500 | +| **average** | **96.4%** | 2000 | + +Reproduce: `lerobot-eval --policy.path=ZibinDong/fastwam_libero_uncond_2cam224 --policy.device=cuda --policy.torch_dtype=float32 --policy.n_action_steps=10 --env.type=libero --env.task=libero_spatial --env.observation_height=256 --env.observation_width=256 --eval.batch_size=1 --eval.n_episodes=50 --seed=0 --env.episode_length=300` (1x H20 140 GB). + +## References + +- [Fast-WAM paper](https://arxiv.org/abs/2603.16666) +- [Fast-WAM project page](https://yuantianyuan01.github.io/FastWAM/) +- [Fast-WAM code](https://github.com/yuantianyuan01/FastWAM) +- [Released upstream checkpoints](https://huggingface.co/yuanty/fastwam) +- [Wan2.2 TI2V 5B](https://huggingface.co/Wan-AI/Wan2.2-TI2V-5B) + +## Citation + +```bibtex +@article{yuan2026fastwam, + title = {Fast-WAM: Do World Action Models Need Test-time Future Imagination?}, + author = {Tianyuan Yuan and Zibin Dong and Yicheng Liu and Hang Zhao}, + journal = {arXiv preprint arXiv:2603.16666}, + year = {2026}, + url = {https://arxiv.org/abs/2603.16666} +} +``` diff --git a/docs/source/policy_fastwam_README.md b/docs/source/policy_fastwam_README.md new file mode 100644 index 000000000..6af0eaa79 --- /dev/null +++ b/docs/source/policy_fastwam_README.md @@ -0,0 +1,56 @@ +## Research Paper + +Paper: https://arxiv.org/abs/2603.16666 + +## Repository + +Code: https://github.com/yuantianyuan01/FastWAM + +Project page: https://yuantianyuan01.github.io/FastWAM/ + +## Citation + +```bibtex +@article{yuan2026fastwam, + title = {Fast-WAM: Do World Action Models Need Test-time Future Imagination?}, + author = {Tianyuan Yuan and Zibin Dong and Yicheng Liu and Hang Zhao}, + journal = {arXiv preprint arXiv:2603.16666}, + year = {2026}, + url = {https://arxiv.org/abs/2603.16666} +} +``` + +## Additional Resources + +Base video model: https://huggingface.co/Wan-AI/Wan2.2-TI2V-5B + +Released upstream checkpoints: https://huggingface.co/yuanty/fastwam + +## Results + +Evaluated on LIBERO with [`ZibinDong/fastwam_libero_uncond_2cam224`](https://huggingface.co/ZibinDong/fastwam_libero_uncond_2cam224): + +| Suite | Success rate | n_episodes | +| -------------- | -----------: | ---------: | +| libero_spatial | 97.6% | 500 | +| libero_object | 99.0% | 500 | +| libero_goal | 95.0% | 500 | +| libero_10 | 94.0% | 500 | +| **average** | **96.4%** | 2000 | + +Reproduce: `lerobot-eval --policy.path=ZibinDong/fastwam_libero_uncond_2cam224 --policy.device=cuda --policy.torch_dtype=float32 --policy.n_action_steps=10 --env.type=libero --env.task=libero_spatial --env.observation_height=256 --env.observation_width=256 --eval.batch_size=1 --eval.n_episodes=50 --seed=0 --env.episode_length=300`. + +For LIBERO-10, use `--env.task=libero_10 --env.episode_length=600`: + +```bash +lerobot-eval \ + --policy.path=ZibinDong/fastwam_libero_uncond_2cam224 \ + --policy.device=cuda \ + --policy.torch_dtype=float32 \ + --policy.n_action_steps=10 \ + --env.type=libero \ + --env.task=libero_10 --env.observation_height=256 --env.observation_width=256 \ + --eval.batch_size=1 \ + --eval.n_episodes=50 \ + --seed=0 --env.episode_length=600 +``` diff --git a/pyproject.toml b/pyproject.toml index 28a8948b7..08ba7fc45 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -229,6 +229,10 @@ robometer = ["lerobot[transformers-dep]", "lerobot[qwen-vl-utils-dep]", "lerobot topreward = ["lerobot[transformers-dep]"] xvla = ["lerobot[transformers-dep]"] eo1 = ["lerobot[transformers-dep]", "lerobot[qwen-vl-utils-dep]"] +fastwam = [ + "lerobot[transformers-dep]", + "lerobot[diffusers-dep]", +] hilserl = ["lerobot[transformers-dep]", "lerobot[dataset]", "gym-hil>=0.1.14,<0.2.0", "lerobot[grpcio-dep]", "lerobot[placo-dep]"] vla_jepa = ["lerobot[transformers-dep]", "lerobot[diffusers-dep]", "lerobot[qwen-vl-utils-dep]"] @@ -308,6 +312,7 @@ all = [ "lerobot[pi]", "lerobot[molmoact2]", "lerobot[smolvla]", + "lerobot[fastwam]", # "lerobot[groot]", TODO(Steven): Gr00t requires specific installation instructions for flash-attn "lerobot[xvla]", "lerobot[hilserl]", @@ -444,7 +449,8 @@ default.extend-ignore-identifiers-re = [ "is_compileable", "ROBOTIS", "OT_VALUE", - "VanderBilt" + "VanderBilt", + "seperated_timestep", ] # TODO: Uncomment when ready to use diff --git a/src/lerobot/policies/__init__.py b/src/lerobot/policies/__init__.py index 68d23c9ca..4daa6abc5 100644 --- a/src/lerobot/policies/__init__.py +++ b/src/lerobot/policies/__init__.py @@ -18,6 +18,7 @@ from .act.configuration_act import ACTConfig as ACTConfig from .diffusion.configuration_diffusion import DiffusionConfig as DiffusionConfig from .eo1.configuration_eo1 import EO1Config as EO1Config from .factory import get_policy_class, make_policy, make_policy_config, make_pre_post_processors +from .fastwam.configuration_fastwam import FastWAMConfig as FastWAMConfig from .gaussian_actor.configuration_gaussian_actor import GaussianActorConfig as GaussianActorConfig from .groot.configuration_groot import GrootConfig as GrootConfig from .molmoact2.configuration_molmoact2 import MolmoAct2Config as MolmoAct2Config @@ -42,6 +43,7 @@ __all__ = [ "ACTConfig", "DiffusionConfig", "EO1Config", + "FastWAMConfig", "GaussianActorConfig", "GrootConfig", "MolmoAct2Config", diff --git a/src/lerobot/policies/factory.py b/src/lerobot/policies/factory.py index b82eaeb72..f5acb0170 100644 --- a/src/lerobot/policies/factory.py +++ b/src/lerobot/policies/factory.py @@ -47,6 +47,7 @@ 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 .fastwam.configuration_fastwam import FastWAMConfig from .gaussian_actor.configuration_gaussian_actor import GaussianActorConfig from .groot.configuration_groot import GrootConfig from .molmoact2.configuration_molmoact2 import MolmoAct2Config @@ -162,6 +163,10 @@ def get_policy_class(name: str) -> type[PreTrainedPolicy]: from .vla_jepa.modeling_vla_jepa import VLAJEPAPolicy return VLAJEPAPolicy + elif name == "fastwam": + from .fastwam.modeling_fastwam import FastWAMPolicy + + return FastWAMPolicy else: try: return _get_policy_cls_from_policy_name(name=name) @@ -218,6 +223,8 @@ def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig: return MolmoAct2Config(**kwargs) elif policy_type == "vla_jepa": return VLAJEPAConfig(**kwargs) + elif policy_type == "fastwam": + return FastWAMConfig(**kwargs) else: try: config_cls = PreTrainedConfig.get_choice_class(policy_type) @@ -451,6 +458,14 @@ def make_pre_post_processors( 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( diff --git a/src/lerobot/policies/fastwam/README.md b/src/lerobot/policies/fastwam/README.md new file mode 120000 index 000000000..d78b9ef36 --- /dev/null +++ b/src/lerobot/policies/fastwam/README.md @@ -0,0 +1 @@ +../../../../docs/source/policy_fastwam_README.md \ No newline at end of file diff --git a/src/lerobot/policies/fastwam/__init__.py b/src/lerobot/policies/fastwam/__init__.py new file mode 100644 index 000000000..8488e7b78 --- /dev/null +++ b/src/lerobot/policies/fastwam/__init__.py @@ -0,0 +1,23 @@ +# 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. + +from .configuration_fastwam import FastWAMConfig +from .modeling_fastwam import FastWAMPolicy +from .processor_fastwam import make_fastwam_pre_post_processors + +__all__ = [ + "FastWAMConfig", + "FastWAMPolicy", + "make_fastwam_pre_post_processors", +] diff --git a/src/lerobot/policies/fastwam/configuration_fastwam.py b/src/lerobot/policies/fastwam/configuration_fastwam.py new file mode 100644 index 000000000..a3ef4f602 --- /dev/null +++ b/src/lerobot/policies/fastwam/configuration_fastwam.py @@ -0,0 +1,399 @@ +# 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. + +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +from lerobot.configs import ( + FeatureType, + NormalizationMode, + PolicyFeature, + PreTrainedConfig, +) +from lerobot.optim import AdamWConfig +from lerobot.utils.constants import ACTION, OBS_STATE + +WAN22_MODEL_ID = "Wan-AI/Wan2.2-TI2V-5B" +WAN22_DIFFUSERS_MODEL_ID = "Wan-AI/Wan2.2-TI2V-5B-Diffusers" +FASTWAM_BASE_MODEL_ID = "lerobot/fastwam_base" +WAN_T5_TOKENIZER_ID = "google/umt5-xxl" + + +_FASTWAM_VIDEO_BASE_COMPAT_KEYS = ( + "patch_size", + "in_dim", + "hidden_dim", + "ffn_dim", + "freq_dim", + "text_dim", + "out_dim", + "num_heads", + "attn_head_dim", + "num_layers", +) + +_FASTWAM_ACTION_BASE_COMPAT_KEYS = ( + "hidden_dim", + "ffn_dim", + "num_heads", + "attn_head_dim", + "num_layers", + "text_dim", + "freq_dim", +) + + +def default_video_dit_config(action_dim: int) -> dict[str, Any]: + return { + "patch_size": [1, 2, 2], + "in_dim": 48, + "hidden_dim": 3072, + "ffn_dim": 14336, + "freq_dim": 256, + "text_dim": 4096, + "out_dim": 48, + "num_heads": 24, + "attn_head_dim": 128, + "num_layers": 30, + "eps": 1.0e-6, + "seperated_timestep": True, + "use_gradient_checkpointing": False, + "video_attention_mask_mode": "first_frame_causal", + "action_conditioned": False, + "action_dim": action_dim, + "action_group_causal_mask_mode": "group_diagonal", + "fp32_attention": True, + } + + +def default_action_dit_config(action_dim: int) -> dict[str, Any]: + return { + "action_dim": action_dim, + "hidden_dim": 1024, + "ffn_dim": 4096, + "num_heads": 24, + "attn_head_dim": 128, + "num_layers": 30, + "text_dim": 4096, + "freq_dim": 256, + "eps": 1.0e-6, + "use_gradient_checkpointing": False, + "fp32_attention": True, + } + + +def _coerce_enum(enum_cls: type, value: Any) -> Any: + if isinstance(value, enum_cls): + return value + try: + return enum_cls(value) + except (TypeError, ValueError) as exc: + member = getattr(enum_cls, str(value), None) + if member is None: + raise ValueError(f"Cannot coerce {value!r} into {enum_cls.__name__}.") from exc + return member + + +def _coerce_policy_features(features: dict[str, Any] | None) -> dict[str, PolicyFeature] | None: + if features is None: + return None + coerced = {} + for name, feature in features.items(): + if isinstance(feature, PolicyFeature): + coerced[name] = feature + continue + coerced[name] = PolicyFeature( + type=_coerce_enum(FeatureType, feature["type"]), + shape=tuple(feature["shape"]), + ) + return coerced + + +def _is_local_model_id(value: str) -> bool: + path = Path(value).expanduser() + return path.is_absolute() or value.startswith(("./", "../", "~")) or path.exists() + + +def _validate_wan_model_id(value: str, field_name: str) -> str: + if value == WAN22_MODEL_ID or _is_local_model_id(value): + return value + raise ValueError(f"`{field_name}` must be `{WAN22_MODEL_ID}` or an explicit local path, got `{value}`.") + + +def is_fastwam_base_compatible_config(config: FastWAMConfig) -> bool: + """Return whether `fastwam_base` partial weights can initialize this config.""" + + default_video_config = default_video_dit_config(config.action_dim) + default_action_config = default_action_dit_config(config.action_dim) + return all( + config.video_dit_config.get(key) == default_video_config.get(key) + for key in _FASTWAM_VIDEO_BASE_COMPAT_KEYS + ) and all( + config.action_dit_config.get(key) == default_action_config.get(key) + for key in _FASTWAM_ACTION_BASE_COMPAT_KEYS + ) + + +@PreTrainedConfig.register_subclass("fastwam") +@dataclass +class FastWAMConfig(PreTrainedConfig): + """Configuration for the FastWAM LeRobot policy. + + Args: + action_dim (int): Number of scalar action channels per timestep. + proprio_dim (int | None): Number of proprioception channels used as an + extra text-context token. `None` disables proprio conditioning. + action_horizon (int): Number of actions predicted by one policy call. + num_video_frames (int): Raw video sampling window (in dataset frames). The + model actually operates on `model_video_frames` frames after subsampling + by `action_video_freq_ratio`. + action_video_freq_ratio (int): Actions are sampled at this multiple of the + video frame rate. Video frames are taken every `action_video_freq_ratio`-th + raw frame, so the model sees `(num_video_frames - 1) // ratio + 1` frames + spanning the same time window as `action_horizon` actions (ratio actions + per video frame). + image_size (tuple[int, int]): Concatenated image size as `(height, width)`. + context_len (int): Maximum text embedding token length. + video_dit_config (dict[str, Any] | None): Wan video expert config. + action_dit_config (dict[str, Any] | None): Action expert config. + use_gradient_checkpointing (bool): Enable activation checkpointing in both DiT + experts (trades compute for memory; propagated into the DiT configs). + freeze_video_expert (bool): Freeze the ~5B Wan video expert + (`model.video_expert`) so only the action expert + proprio encoder train. + Cuts the AdamW optimizer footprint substantially; the video expert keeps its + pretrained weights. (If enabled, also set `loss.lambda_video=0` to skip the + now-gradient-free video loss compute.) + """ + + n_obs_steps: int = 1 + action_dim: int = 7 + proprio_dim: int | None = 8 + action_horizon: int = 32 + n_action_steps: int = 32 + num_video_frames: int = 33 + action_video_freq_ratio: int = 4 + image_size: tuple[int, int] = (224, 448) + context_len: int = 128 + model_id: str = WAN22_MODEL_ID + tokenizer_model_id: str = WAN_T5_TOKENIZER_ID + text_encoder_model_id: str = WAN22_DIFFUSERS_MODEL_ID + base_model_id: str | None = FASTWAM_BASE_MODEL_ID + tokenizer_max_len: int = 128 + load_text_encoder: bool = True + mot_checkpoint_mixed_attn: bool = False + torch_dtype: str = "bfloat16" + prompt_template: str = ( + "A video recorded from a robot's point of view executing the following instruction: {task}" + ) + num_inference_steps: int = 10 + inference_seed: int | None = 42 + rand_device: str = "cpu" + text_cfg_scale: float = 1.0 + negative_prompt: str = "" + sigma_shift: float | None = None + tiled: bool = False + fp32_attention: bool = True + use_gradient_checkpointing: bool = False + freeze_video_expert: bool = False + toggle_action_dimensions: list[int] = field(default_factory=list) + video_scheduler: dict[str, float | int] = field( + default_factory=lambda: {"train_shift": 5.0, "infer_shift": 5.0, "num_train_timesteps": 1000} + ) + action_scheduler: dict[str, float | int] = field( + default_factory=lambda: {"train_shift": 5.0, "infer_shift": 5.0, "num_train_timesteps": 1000} + ) + loss: dict[str, float] = field(default_factory=lambda: {"lambda_video": 1.0, "lambda_action": 1.0}) + video_dit_config: dict[str, Any] | None = None + action_dit_config: dict[str, Any] | None = None + normalization_mapping: dict[str, NormalizationMode] = field( + default_factory=lambda: { + "VISUAL": NormalizationMode.IDENTITY, + "STATE": NormalizationMode.MEAN_STD, + "ACTION": NormalizationMode.MEAN_STD, + } + ) + input_features: dict[str, PolicyFeature] | None = None + output_features: dict[str, PolicyFeature] | None = None + optimizer_lr: float = 1.0e-4 + optimizer_weight_decay: float = 1.0e-2 + + def __post_init__(self) -> None: + super().__post_init__() + self.image_size = tuple(self.image_size) + self.model_id = _validate_wan_model_id(self.model_id, "model_id") + self.input_features = _coerce_policy_features(self.input_features) + self.output_features = _coerce_policy_features(self.output_features) + self.toggle_action_dimensions = [int(dim) for dim in self.toggle_action_dimensions] + self.video_dit_config = self.video_dit_config or default_video_dit_config(self.action_dim) + self.action_dit_config = self.action_dit_config or default_action_dit_config(self.action_dim) + self.video_dit_config["fp32_attention"] = bool(self.fp32_attention) + self.action_dit_config["fp32_attention"] = bool(self.fp32_attention) + self.video_dit_config["use_gradient_checkpointing"] = bool(self.use_gradient_checkpointing) + self.action_dit_config["use_gradient_checkpointing"] = bool(self.use_gradient_checkpointing) + if self.input_features is None: + height, width = self.image_size + self.input_features = { + "observation.images.image": PolicyFeature( + type=FeatureType.VISUAL, + shape=(3, height, width), + ) + } + if self.proprio_dim is not None: + self.input_features[OBS_STATE] = PolicyFeature( + type=FeatureType.STATE, + shape=(self.proprio_dim,), + ) + if self.output_features is None: + self.output_features = {ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(self.action_dim,))} + self.validate_features() + if self.pretrained_path or self.use_peft or not self.base_model_id: + return + if not is_fastwam_base_compatible_config(self): + return + self.pretrained_path = Path(self.base_model_id) + self._auto_pretrained_path = True + + def _save_pretrained(self, save_directory: Path) -> None: + if not getattr(self, "_auto_pretrained_path", False): + super()._save_pretrained(save_directory) + return + + pretrained_path = self.pretrained_path + self.pretrained_path = None + try: + super()._save_pretrained(save_directory) + finally: + self.pretrained_path = pretrained_path + + def get_optimizer_preset(self) -> AdamWConfig: + return AdamWConfig(lr=self.optimizer_lr, weight_decay=self.optimizer_weight_decay) + + def get_scheduler_preset(self) -> None: + return None + + def set_dataset_feature_metadata(self, dataset_features: dict[str, Any]) -> None: + """Rebuild visual input features from the dataset's real camera keys. + + FastWAM's `__post_init__` installs a synthetic single-image default + (`observation.images.image` at full `image_size` width). For datasets + with one or more separately-named cameras (e.g. `observation.images.top`, + `observation.images.wrist`), this hook — invoked by `make_policy` once the + dataset metadata is known — replaces that default with the actual camera + keys, each declared at the policy's native per-camera resolution + (`image_size[0]` x `image_size[1] // num_cameras`). The accompanying + resize step in `make_fastwam_pre_post_processors` resizes raw frames to + match, so heterogeneous source resolutions (e.g. 480x640) are supported. + """ + image_keys = sorted( + key + for key, feature in dataset_features.items() + if key.startswith("observation.images.") and feature.get("dtype") in ("video", "image") + ) + if not image_keys: + return + height, total_width = self.image_size + per_cam_width = total_width // len(image_keys) + new_inputs: dict[str, PolicyFeature] = { + key: PolicyFeature(type=FeatureType.VISUAL, shape=(3, height, per_cam_width)) + for key in image_keys + } + if self.proprio_dim is not None and OBS_STATE in dataset_features: + new_inputs[OBS_STATE] = PolicyFeature(type=FeatureType.STATE, shape=(self.proprio_dim,)) + self.input_features = new_inputs + self.validate_features() + + def validate_features(self) -> None: + if self.action_dim <= 0: + raise ValueError(f"`action_dim` must be positive, got {self.action_dim}.") + if self.action_horizon <= 0: + raise ValueError(f"`action_horizon` must be positive, got {self.action_horizon}.") + if self.n_action_steps > self.action_horizon: + raise ValueError("`n_action_steps` cannot exceed `action_horizon`.") + if self.action_video_freq_ratio <= 0: + raise ValueError( + f"`action_video_freq_ratio` must be positive, got {self.action_video_freq_ratio}." + ) + # Video frames are subsampled by action_video_freq_ratio; the resulting model frame + # count must satisfy T % 4 == 1 for the VAE temporal tokenization (mirrors the + # original FastWAM dataset asserts). + if (self.num_video_frames - 1) % self.action_video_freq_ratio != 0: + raise ValueError( + f"`num_video_frames - 1` ({self.num_video_frames - 1}) must be divisible by " + f"`action_video_freq_ratio` ({self.action_video_freq_ratio})." + ) + if ((self.num_video_frames - 1) // self.action_video_freq_ratio) % 4 != 0: + raise ValueError( + f"Subsampled video transitions ({(self.num_video_frames - 1) // self.action_video_freq_ratio}) " + "must be divisible by 4 for VAE tokenization (i.e. model_video_frames % 4 == 1)." + ) + if self.action_horizon % ((self.num_video_frames - 1) // self.action_video_freq_ratio) != 0: + raise ValueError( + f"`action_horizon` ({self.action_horizon}) must be divisible by the number of " + f"video transitions ({(self.num_video_frames - 1) // self.action_video_freq_ratio})." + ) + if not self.image_features: + raise ValueError("FastWAM requires at least one image feature.") + if self.action_feature is None: + raise ValueError("FastWAM requires `action` in output_features.") + action_shape = tuple(self.action_feature.shape) + if action_shape != (self.action_dim,): + raise ValueError( + f"FastWAM action feature shape must be ({self.action_dim},), got {action_shape}." + ) + if self.proprio_dim is not None: + state_feature = self.robot_state_feature + if state_feature is None: + raise ValueError("FastWAM requires `observation.state` when `proprio_dim` is set.") + state_shape = tuple(state_feature.shape) + if state_shape != (self.proprio_dim,): + raise ValueError( + f"FastWAM state feature shape must be ({self.proprio_dim},), got {state_shape}." + ) + height, width = self.image_size + image_width_sum = 0 + for name, feature in self.image_features.items(): + shape = tuple(feature.shape) + if len(shape) != 3 or shape[0] != 3: + raise ValueError(f"FastWAM image feature `{name}` must have shape (3, H, W), got {shape}.") + if shape[1] != height: + raise ValueError(f"FastWAM image feature `{name}` height must be {height}, got {shape[1]}.") + image_width_sum += shape[2] + if image_width_sum != width: + raise ValueError(f"FastWAM image feature widths must sum to {width}, got {image_width_sum}.") + + @property + def model_video_frames(self) -> int: + """Number of video frames the model actually operates on, after subsampling the + raw `num_video_frames` window by `action_video_freq_ratio` (e.g. 33 -> 9).""" + return (self.num_video_frames - 1) // self.action_video_freq_ratio + 1 + + @property + def observation_delta_indices(self) -> list[int]: + # Load the video frames the model is supervised on: the future window subsampled by + # action_video_freq_ratio (e.g. [0, 4, 8, ..., 32] -> 9 frames). Each video frame is + # thus `action_video_freq_ratio` actions apart, while actions load at the full rate + # (`action_delta_indices` = range(action_horizon)). Returning None would load only the + # current frame, making the video target a static repeat (degenerate supervision). + return list(range(0, self.num_video_frames, self.action_video_freq_ratio)) + + @property + def action_delta_indices(self) -> list[int]: + return list(range(self.action_horizon)) + + @property + def reward_delta_indices(self) -> None: + return None diff --git a/src/lerobot/policies/fastwam/modeling_fastwam.py b/src/lerobot/policies/fastwam/modeling_fastwam.py new file mode 100644 index 000000000..10671e717 --- /dev/null +++ b/src/lerobot/policies/fastwam/modeling_fastwam.py @@ -0,0 +1,440 @@ +# 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. + +from __future__ import annotations + +import logging +from collections import deque +from typing import Any + +import torch +from torch import Tensor + +from lerobot.policies.pretrained import PreTrainedPolicy +from lerobot.utils.constants import OBS_STATE +from lerobot.utils.import_utils import require_package + +from .configuration_fastwam import FastWAMConfig +from .wan import ( + ActionDiT, + FastWAM, + MoT, + WanVideoDiT, + build_wan_tokenizer, + load_pretrained_wan_text_encoder, + load_pretrained_wan_vae, +) + + +class FastWAMPolicy(PreTrainedPolicy): + """LeRobot policy wrapper for FastWAM. + + Attention backend: FastWAM's DiT uses ``torch.nn.functional.scaled_dot_product_attention`` + (SDPA) for all attention. It does not use FlashAttention, because MoT routing requires + arbitrary boolean ``[query, key]`` masks that the FlashAttention varlen API cannot express; + installing ``flash-attn`` has no effect on the FastWAM path. (SDPA may still dispatch to + PyTorch's own flash/mem-efficient/math kernel internally, unrelated to the ``flash-attn`` package.) + + Args: + config (FastWAMConfig): FastWAM policy configuration. + dataset_stats (dict[str, dict[str, Tensor]] | None): Optional LeRobot + dataset statistics passed by the training/evaluation stack. + """ + + config_class = FastWAMConfig + name = "fastwam" + + def __init__( + self, + config: FastWAMConfig, + dataset_stats: dict[str, dict[str, Tensor]] | None = None, + **kwargs: Any, + ): + # FastWAM's Wan2.2 backbone needs transformers (UMT5 text encoder/tokenizer) and + # diffusers (Wan VAE), both behind the `fastwam` extra. Fail fast with an actionable + # message in base installs rather than deep in Wan component construction. + require_package("transformers", extra="fastwam") + require_package("diffusers", extra="fastwam") + # `make_policy`/`from_pretrained` forward extra kwargs (e.g. `dataset_meta`); the + # dataset feature metadata is already applied to `config` by make_policy upstream, + # so we accept and ignore them, matching the other LeRobot policies. + super().__init__(config, dataset_stats) + config.validate_features() + self.config = config + self.dataset_stats = dataset_stats + self.model = self._build_core_model(config) + if config.freeze_video_expert and getattr(self.model, "video_expert", None) is not None: + # Freeze the ~5B Wan video expert; get_optim_params filters on requires_grad, + # so its params drop out of the optimizer (and DDP skips them). + self.model.video_expert.requires_grad_(False) + # The transformer blocks are re-parented onto the MoTLayers (single FSDP owner), so + # `video_expert.requires_grad_` no longer reaches them — freeze them via the layers. + mot = getattr(self.model, "mot", None) + if mot is not None and getattr(mot, "layers", None) is not None: + for layer in mot.layers: + if "video" in layer.blocks: + layer.blocks["video"].requires_grad_(False) + self.reset() + + @classmethod + def _load_as_safetensor(cls, model, model_file: str, map_location: str, strict: bool): + """Shape-aware load that supports cross-embodiment fine-tuning. + + `safetensors.load_model(strict=False)` ignores missing/unexpected keys but + still raises on a shape mismatch for a shared key. When fine-tuning from a + checkpoint trained on a different embodiment (e.g. the LIBERO 7-DoF / 8-dim + checkpoint adapted to a 6-DoF / 6-dim arm), the action encoder/head and + proprio encoder legitimately differ in shape. With `strict=False` we drop + only those shape-mismatched tensors — leaving them at their freshly + initialized values — and load every compatible tensor. With `strict=True` + the standard exact-match loader is used. + """ + from safetensors import safe_open + + model_state_dict = model.state_dict() + mismatched = [] + with safe_open(model_file, framework="pt") as f: + checkpoint_keys = list(f.keys()) + for key in checkpoint_keys: + if key in model_state_dict and tuple(model_state_dict[key].shape) != tuple( + f.get_slice(key).get_shape() + ): + mismatched.append(key) + + if not mismatched: + return super()._load_as_safetensor(model, model_file, map_location, strict) + if strict: + raise RuntimeError( + f"FastWAM: {len(mismatched)} checkpoint tensors have a shape mismatch under " + f"strict=True: {mismatched}" + ) + + from safetensors.torch import load_file + + logging.warning( + "FastWAM cross-embodiment load: reinitializing %d shape-mismatched tensor(s), keeping " + "every compatible weight: %s", + len(mismatched), + mismatched, + ) + state_dict = load_file(model_file, device="cpu") + for key in mismatched: + state_dict.pop(key, None) + model.load_state_dict(state_dict, strict=False) + if map_location and map_location != "cpu": + model.to(map_location) + return model + + def get_optim_params(self) -> list[Tensor]: + # Return the trainable tensors directly (a single param group). The optimizer + # builder wraps these in a param group; returning a bare {"params": [...]} dict + # instead would make `list(...)` yield the key string "params". + params = ( + list(self.model.dit.parameters()) if hasattr(self.model, "dit") else list(self.model.parameters()) + ) + proprio_encoder = getattr(self.model, "proprio_encoder", None) + if proprio_encoder is not None: + params.extend(list(proprio_encoder.parameters())) + return [p for p in params if p.requires_grad] + + def reset(self) -> None: + self._action_queue: deque[Tensor] = deque([], maxlen=self.config.n_action_steps) + + def _batch_to_training_sample(self, batch: dict[str, Tensor]) -> dict[str, Tensor]: + """Adapt a standard LeRobot batch to the FastWAM-native sample that + `FastWAM.build_inputs` consumes (`video`, `action`, `context`/`context_mask`, + per-frame `proprio`). + + The LeRobot training loop passes raw `observation.images.*`, a single-step + `observation.state` `[B, D]`, `action`, and a language `task` string. We do + only the translation `build_inputs` can't: stack the camera frames into a + video, encode the prompt with the (frozen) text encoder (mirroring inference, + so language-conditioned datasets need no precomputed context), and give proprio + the per-frame axis `build_inputs` indexes. All shape/presence validation is + left to `build_inputs`, the single authority on the contract. + """ + sample = dict(batch) + if "video" not in sample: + sample["video"] = _stack_video_from_images(batch, self.config) + if "context" not in sample or "context_mask" not in sample: + prompt = _prompt_from_batch(batch=batch, config=self.config) + if prompt is None: + raise KeyError( + "FastWAM training requires a `task`/`prompt` to encode text context, " + "or precomputed `context`/`context_mask` in the batch." + ) + sample["context"], sample["context_mask"] = self.model.encode_prompt(prompt) + if self.config.proprio_dim is not None and "proprio" not in sample: + state = sample.get(OBS_STATE) + if state is not None: + # LeRobot gives a single-step state [B, D]; build_inputs expects + # per-frame [B, T, D] and uses frame 0, so add a T=1 axis. + 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]]: + """Compute FastWAM training loss for a LeRobot batch. + + Args: + batch (dict[str, Tensor]): Batch containing FastWAM-ready keys + (`video`, `action`, `context`, `context_mask`) or LeRobot keys + that can be adapted (`observation.images.*`, `observation.state`, + `action`, `action_is_pad`). + + 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. + """ + + sample = self._batch_to_training_sample(batch) + loss, metrics = self.model.training_loss(sample) + return loss, dict(metrics or {}) + + @torch.no_grad() + def predict_action_chunk(self, batch: dict[str, Tensor], **_: Any) -> Tensor: + """Predict a chunk of actions from the current FastWAM observation. + + Args: + batch (dict[str, Tensor]): Inference batch with `input_image` or + image observation keys, plus `context/context_mask` or `prompt`. + + Returns: + Tensor: Action chunk with shape `[B, action_horizon, action_dim]`. + """ + + self.eval() + infer_kwargs = _batch_to_infer_kwargs(batch=batch, config=self.config) + batch_size = _infer_kwargs_batch_size(infer_kwargs) + if batch_size == 1: + action = _action_from_model_output(self.model.infer_action(**infer_kwargs)) + else: + action = torch.cat( + [ + _action_from_model_output( + self.model.infer_action( + **_slice_infer_kwargs(infer_kwargs, index=i, batch_size=batch_size) + ) + ) + for i in range(batch_size) + ], + dim=0, + ) + return action.to(device=batch_device(batch), dtype=torch.float32) + + @torch.no_grad() + def select_action(self, batch: dict[str, Tensor], **kwargs: Any) -> Tensor: + self.eval() + if len(self._action_queue) == 0: + actions = self.predict_action_chunk(batch, **kwargs)[:, : self.config.n_action_steps] + self._action_queue.extend(actions.transpose(0, 1)) + return self._action_queue.popleft() + + def _build_core_model(self, config: FastWAMConfig) -> FastWAM: + """Build the FastWAM core for training / inference. + + Only the trainable parts (the MoT DiT and the proprio encoder) are + materialized empty here and then filled from the policy's + `model.safetensors` by the base `from_pretrained`. The *frozen* Wan2.2 VAE + and UMT5 text encoder are loaded with their real weights from the + `Wan-AI/Wan2.2-TI2V-5B-Diffusers` repo (cached in the HF cache, shared + across checkpoints) and are intentionally excluded from `model.safetensors` + — see `FastWAM.__init__`. The tokenizer comes from `google/umt5-xxl`. + """ + dtype = _dtype_from_name(config.torch_dtype) + device = config.device + video_expert = WanVideoDiT(**config.video_dit_config).to(device=device, dtype=dtype) + action_expert = ActionDiT(**config.action_dit_config).to(device=device, dtype=dtype) + mot = MoT( + mixtures={"video": video_expert, "action": action_expert}, + mot_checkpoint_mixed_attn=config.mot_checkpoint_mixed_attn, + ) + text_encoder = ( + load_pretrained_wan_text_encoder( + model_id=config.text_encoder_model_id, torch_dtype=dtype, device=device + ) + if config.load_text_encoder + else None + ) + return FastWAM( + video_expert=video_expert, + action_expert=action_expert, + mot=mot, + vae=load_pretrained_wan_vae(torch_dtype=dtype, device=device), + text_encoder=text_encoder, + tokenizer=build_wan_tokenizer( + model_id=config.tokenizer_model_id, tokenizer_max_len=config.tokenizer_max_len + ), + text_dim=int(config.video_dit_config["text_dim"]), + proprio_dim=config.proprio_dim, + device=device, + torch_dtype=dtype, + video_train_shift=float(config.video_scheduler["train_shift"]), + video_infer_shift=float(config.video_scheduler["infer_shift"]), + video_num_train_timesteps=int(config.video_scheduler["num_train_timesteps"]), + action_train_shift=float(config.action_scheduler["train_shift"]), + action_infer_shift=float(config.action_scheduler["infer_shift"]), + action_num_train_timesteps=int(config.action_scheduler["num_train_timesteps"]), + loss_lambda_video=float(config.loss["lambda_video"]), + loss_lambda_action=float(config.loss["lambda_action"]), + ) + + +def _scalar(value: Any) -> Any: + """Unwrap a 0-/1-element tensor (e.g. from DataLoader collation) to a Python scalar.""" + return value.item() if isinstance(value, Tensor) else value + + +def _batch_to_infer_kwargs(batch: dict[str, Tensor], config: FastWAMConfig) -> dict[str, Any]: + return { + "prompt": _prompt_from_batch(batch=batch, config=config), + "input_image": _input_image_from_batch(batch, config), + "action_horizon": config.action_horizon, + "proprio": batch.get("proprio", batch.get(OBS_STATE)), + "context": batch.get("context"), + "context_mask": batch.get("context_mask"), + "negative_prompt": batch.get("negative_prompt", config.negative_prompt), + "text_cfg_scale": float(_scalar(batch.get("text_cfg_scale", config.text_cfg_scale))), + "num_inference_steps": int(_scalar(batch.get("num_inference_steps", config.num_inference_steps))), + "sigma_shift": batch.get("sigma_shift", config.sigma_shift), + "seed": batch.get("seed", config.inference_seed), + "rand_device": batch.get("rand_device", config.rand_device), + "tiled": bool(batch.get("tiled", config.tiled)), + } + + +def _prompt_from_batch(batch: dict[str, Tensor], config: FastWAMConfig) -> Any: + prompt = batch.get("prompt") + if prompt is not None: + return prompt + + task = batch.get("task") + if task is None: + return None + if isinstance(task, str): + return config.prompt_template.format(task=task) + if isinstance(task, (list, tuple)): + return [config.prompt_template.format(task=str(item)) for item in task] + return config.prompt_template.format(task=str(task)) + + +def _action_from_model_output(output: Any) -> Tensor: + action = output["action"] if isinstance(output, dict) else output + if action.ndim == 2: + action = action.unsqueeze(0) + return action + + +def _infer_kwargs_batch_size(infer_kwargs: dict[str, Any]) -> int: + image = infer_kwargs["input_image"] + if not isinstance(image, Tensor): + raise TypeError(f"`input_image` must be a tensor, got {type(image).__name__}.") + if image.ndim == 3: + return 1 + if image.ndim == 4: + return int(image.shape[0]) + raise ValueError(f"`input_image` must be [B,C,H,W] or [C,H,W], got {tuple(image.shape)}.") + + +def _slice_infer_kwargs(infer_kwargs: dict[str, Any], *, index: int, batch_size: int) -> dict[str, Any]: + return { + key: _slice_infer_value(value, index=index, batch_size=batch_size) + for key, value in infer_kwargs.items() + } + + +def _slice_infer_value(value: Any, *, index: int, batch_size: int) -> Any: + if isinstance(value, Tensor) and value.ndim > 0 and value.shape[0] == batch_size: + return value[index : index + 1] + if isinstance(value, (list, tuple)) and len(value) == batch_size: + return value[index] + return value + + +def _dtype_from_name(name: str) -> torch.dtype: + dtype_map = {"float32": torch.float32, "float16": torch.float16, "bfloat16": torch.bfloat16} + if name not in dtype_map: + raise ValueError(f"Unsupported torch dtype `{name}`.") + return dtype_map[name] + + +def batch_device(batch: dict[str, Any]) -> torch.device: + for value in batch.values(): + if isinstance(value, Tensor): + return value.device + return torch.device("cpu") + + +def _resize_frames(frames: Tensor, size: tuple[int, int]) -> Tensor: + """Resize a frame tensor to `size` (H, W), tolerating a leading temporal/batch stack. + + `interpolate` only accepts a single leading batch dim (`[N, C, H, W]`), but FastWAM camera + tensors arrive as `[B, C, H, W]` (live eval) or `[B, T, C, H, W]` (temporal stack), so flatten + any leading dims into the batch, resize, then restore. A no-op when already at `size`. + """ + if tuple(frames.shape[-2:]) == size: + return frames + lead = frames.shape[:-3] + flat = frames.reshape(-1, *frames.shape[-3:]) + flat = torch.nn.functional.interpolate( + flat, size=size, mode="bilinear", align_corners=False, antialias=True + ) + return flat.reshape(*lead, *flat.shape[-3:]) + + +def _stack_video_from_images(batch: dict[str, Tensor], config: FastWAMConfig) -> Tensor: + # Exclude the `*_is_pad` companion tensors that delta-timestamp loading adds alongside + # each camera (shape [B, T]); they share the `observation.images.` prefix but are not frames. + image_keys = sorted(k for k in batch if k.startswith("observation.images.") and not k.endswith("_is_pad")) + if not image_keys: + raise KeyError("FastWAM batch must contain `video` or `observation.images.*` keys.") + per_cam = (int(config.image_size[0]), int(config.image_size[1]) // len(image_keys)) + images = [_resize_frames(batch[key], per_cam) for key in image_keys] + # Cameras concatenate along width (last dim) in both the single-frame and temporal case. + image = torch.cat(images, dim=-1) if len(images) > 1 else images[0] + if image.ndim == 4: + # [B, C, H, W]: a single frame (e.g. the live eval observation) -> repeat across time. + image = image.unsqueeze(2).repeat(1, 1, config.model_video_frames, 1, 1) + elif image.ndim == 5: + # [B, T, C, H, W]: temporal stack from delta-timestamp loading -> [B, C, T, H, W]. + image = image.permute(0, 2, 1, 3, 4) + else: + raise ValueError(f"Expected image batch [B,C,H,W] or temporal [B,T,C,H,W], got {tuple(image.shape)}.") + return image + + +def _input_image_from_batch(batch: dict[str, Tensor], config: FastWAMConfig) -> Tensor: + if "input_image" in batch: + return _prepare_infer_image(batch["input_image"], config) + video = batch.get("video") + if video is None: + video = _stack_video_from_images(batch, config) + if video.ndim == 5: + return _prepare_infer_image(video[:, :, 0], config) + if video.ndim == 4: + return _prepare_infer_image(video, config) + raise ValueError(f"Cannot build input image from tensor with shape {tuple(video.shape)}.") + + +def _prepare_infer_image(image: Tensor, config: FastWAMConfig) -> Tensor: + if image.ndim == 3: + image = image.unsqueeze(0) + if image.ndim != 4: + raise ValueError(f"Expected image tensor [B,C,H,W] or [C,H,W], got {tuple(image.shape)}.") + + # Resize to the full configured resolution (no-op when the video path already produced it, but + # also covers a directly-supplied `input_image`). The model owns its input resolution — see + # `_stack_video_from_images` — so we resize rather than assert on a mismatch. + target_h, target_w = int(config.image_size[0]), int(config.image_size[1]) + return _resize_frames(image, (target_h, target_w)) diff --git a/src/lerobot/policies/fastwam/processor_fastwam.py b/src/lerobot/policies/fastwam/processor_fastwam.py new file mode 100644 index 000000000..31f3b9277 --- /dev/null +++ b/src/lerobot/policies/fastwam/processor_fastwam.py @@ -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. + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import torch + +from lerobot.configs import PipelineFeatureType, PolicyFeature +from lerobot.processor import ( + ActionProcessorStep, + AddBatchDimensionProcessorStep, + DeviceProcessorStep, + NormalizerProcessorStep, + PolicyAction, + PolicyProcessorPipeline, + ProcessorStepRegistry, + 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 + + +@dataclass +@ProcessorStepRegistry.register(name="fastwam_action_toggle_processor") +class FastWAMActionToggleProcessorStep(ActionProcessorStep): + """Apply FastWAM LIBERO toggle semantics to configured action dimensions.""" + + toggle_dimensions: list[int] + + def action(self, action: PolicyAction) -> PolicyAction: + if not self.toggle_dimensions: + return action + processed_action = action.clone() + action_dim = int(processed_action.shape[-1]) + for dim in self.toggle_dimensions: + resolved_dim = dim if dim >= 0 else action_dim + dim + if resolved_dim < 0 or resolved_dim >= action_dim: + raise ValueError( + f"FastWAM action toggle dimension {dim} is out of bounds for action dim {action_dim}." + ) + value = processed_action[..., resolved_dim] + value = value * 2.0 - 1.0 + processed_action[..., resolved_dim] = torch.sign(-value) + return processed_action + + def get_config(self) -> dict[str, Any]: + return {"toggle_dimensions": self.toggle_dimensions} + + def transform_features( + self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]] + ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]: + return features + + +def make_fastwam_pre_post_processors( + config: FastWAMConfig, + dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None, +) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]: + """Create LeRobot pre- and post-processing pipelines for FastWAM. + + Args: + config (FastWAMConfig): Policy configuration controlling device and + normalization feature metadata. + dataset_stats (dict[str, dict[str, torch.Tensor]] | None): Optional + LeRobot dataset statistics used by normalization processors. + + Returns: + tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]: Input and + output processor pipelines discoverable by LeRobot. + """ + + # NOTE: no visual normalization here. VISUAL is IDENTITY (see configuration_fastwam.normalization_mapping) + # — images pass through in [0, 1] and the model maps them to the Wan VAE's [-1, 1] at the encode + # boundary. This is deliberate: `lerobot_train.py` overrides the normalizer stats with + # `dataset.meta.stats` when fine-tuning, and a real dataset's per-channel image std is the tiny + # frame-to-frame brightness variance, which would blow images far outside [-1,1] and saturate them. + # STATE/ACTION still normalize with dataset stats below. + normalization_stats: dict[str, dict[str, Any]] = dict(dataset_stats or {}) + + # NOTE: no resize step here. The model is the single authority on input resolution: it resizes + # each camera to the per-camera target (image_size split across cameras) in + # `_stack_video_from_images` / `_prepare_infer_image`, on every path (train forward, rollout and + # eval select_action). A preprocessor resize step would be both redundant (the model re-resizes + # 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 = [ + 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 = [ + 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(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, + ), + ) diff --git a/src/lerobot/policies/fastwam/wan/README.md b/src/lerobot/policies/fastwam/wan/README.md new file mode 100644 index 000000000..7b7d61033 --- /dev/null +++ b/src/lerobot/policies/fastwam/wan/README.md @@ -0,0 +1,34 @@ +# FastWAM `wan` package + +This package holds FastWAM's model implementation. It mixes a small **vendored +subset of the official Wan2.2 source tree** with FastWAM's own code, kept flat in +a single directory. + +## Vendored from Wan2.2 + +- Upstream repository: https://github.com/Wan-Video/Wan2.2 +- Upstream commit: `42bf4cfaa384bc21833865abc2f9e6c0e67233dc` +- License: Apache-2.0, matching the license in `LICENSE.txt` from the upstream repository + +Copied files: + +- `model.py` (was `wan/modules/model.py`), trimmed: the flash-attention path + (the vendored `attention.py` and the block/model `forward`s) was removed. + FastWAM's DiT uses SDPA instead (see `video_dit.py`). +- `get_sampling_sigmas` in `video_dit.py` (was `wan/utils/fm_solvers.py`), inlined + next to its only caller. + +This subset only backs FastWAM's **custom MoT video DiT**. The Wan2.2 VAE, +UMT5 text encoder, and tokenizer are no longer vendored - they come from +`diffusers.AutoencoderKLWan`, `transformers.UMT5EncoderModel`, and +`transformers.AutoTokenizer` (see `components.py` and `adapters.py`). + +## FastWAM's own code + +- `video_dit.py` builds on `model` (`sinusoidal_embedding_1d`, `rope_params`, + `rope_apply`, …) and computes attention with SDPA (`fastwam_masked_attention`). Its + `WanContinuousFlowMatchScheduler` uses `get_sampling_sigmas` for Wan-compatible + inference timesteps. +- `components.py` / `adapters.py` load the VAE, text encoder, tokenizer, and the + custom DiT weights. +- `modular.py` defines the FastWAM model (`ActionDiT`, `MoT`, `FastWAM`, …). diff --git a/src/lerobot/policies/fastwam/wan/__init__.py b/src/lerobot/policies/fastwam/wan/__init__.py new file mode 100644 index 000000000..11c0aaf6f --- /dev/null +++ b/src/lerobot/policies/fastwam/wan/__init__.py @@ -0,0 +1,33 @@ +# 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. + +from .adapters import WanVideoVAE38 +from .components import ( + build_wan_tokenizer, + load_pretrained_wan_text_encoder, + load_pretrained_wan_vae, +) +from .modular import ActionDiT, FastWAM, MoT +from .video_dit import WanVideoDiT + +__all__ = [ + "ActionDiT", + "FastWAM", + "MoT", + "WanVideoDiT", + "WanVideoVAE38", + "build_wan_tokenizer", + "load_pretrained_wan_text_encoder", + "load_pretrained_wan_vae", +] diff --git a/src/lerobot/policies/fastwam/wan/adapters.py b/src/lerobot/policies/fastwam/wan/adapters.py new file mode 100644 index 000000000..a2e8a1068 --- /dev/null +++ b/src/lerobot/policies/fastwam/wan/adapters.py @@ -0,0 +1,108 @@ +# 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. + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +if TYPE_CHECKING: + from diffusers import AutoencoderKLWan + + +class WanVideoVAE38(torch.nn.Module): + """FastWAM VAE contract over `diffusers.AutoencoderKLWan` (Wan2.2-TI2V-5B). + + 16x spatial / 4x temporal compression, 48 latent channels. diffusers' + `AutoencoderKLWan` returns *raw* latents (it does not apply `latents_mean`/ + `latents_std`), so `encode`/`decode` here apply the same standardization the + Wan reference uses — `(latents - mean) / std` — done in fp32 for stability. + `encode` uses the deterministic posterior mode, matching the original VAE + which returned the latent mean `mu`. + """ + + upsampling_factor = 16 + temporal_downsample_factor = 4 + z_dim = 48 + + def __init__( + self, + dtype: torch.dtype = torch.float32, + device: str | torch.device = "cuda", + *, + pretrained: AutoencoderKLWan, + ) -> None: + super().__init__() + # The Wan2.2 VAE is a fixed pretrained model — it is never trained from scratch, + # so a real `AutoencoderKLWan` (with weights) must always be supplied (loaded from + # the diffusers repo by `load_pretrained_wan_vae`). No random/offline build path. + self.vae = pretrained.to(device=device, dtype=dtype) + + # Read the standardization stats from the VAE's own config (diffusers populates + # these from vae/config.json) — single source of truth, no local copy. diffusers' + # encode/decode return *raw* latents, so we apply (latent - mean) / std ourselves. + # Non-persistent: kept out of state_dict. + self.register_buffer( + "latents_mean", + torch.tensor(self.vae.config.latents_mean).view(1, self.z_dim, 1, 1, 1), + persistent=False, + ) + self.register_buffer( + "latents_std", + torch.tensor(self.vae.config.latents_std).view(1, self.z_dim, 1, 1, 1), + persistent=False, + ) + + def _device_dtype(self) -> tuple[torch.device, torch.dtype]: + param = next(self.vae.parameters()) + return param.device, param.dtype + + def encode( + self, + videos: list[torch.Tensor] | torch.Tensor, + device: str | torch.device | None = None, + tiled: bool = False, + tile_size: tuple[int, int] = (34, 34), + tile_stride: tuple[int, int] = (18, 16), + ) -> torch.Tensor: + del device, tile_size, tile_stride + if tiled: + raise NotImplementedError("Tiled Wan2.2 VAE encoding is not supported by the FastWAM adapter.") + if isinstance(videos, (list, tuple)): + videos = torch.stack(list(videos)) + dev, dtype = self._device_dtype() + mu = self.vae.encode(videos.to(device=dev, dtype=dtype)).latent_dist.mode().float() + mean = self.latents_mean.float().to(mu.device) + std = self.latents_std.float().to(mu.device) + return (mu - mean) / std + + def decode( + self, + hidden_states: list[torch.Tensor] | torch.Tensor, + device: str | torch.device | None = None, + tiled: bool = False, + tile_size: tuple[int, int] = (34, 34), + tile_stride: tuple[int, int] = (18, 16), + ) -> torch.Tensor: + del device, tile_size, tile_stride + if tiled: + raise NotImplementedError("Tiled Wan2.2 VAE decoding is not supported by the FastWAM adapter.") + if isinstance(hidden_states, (list, tuple)): + hidden_states = torch.stack(list(hidden_states)) + dev, dtype = self._device_dtype() + z = hidden_states.float() + z = z * self.latents_std.float().to(z.device) + self.latents_mean.float().to(z.device) + out = self.vae.decode(z.to(device=dev, dtype=dtype)).sample + return out.float().clamp_(-1.0, 1.0) diff --git a/src/lerobot/policies/fastwam/wan/components.py b/src/lerobot/policies/fastwam/wan/components.py new file mode 100644 index 000000000..59a9b610e --- /dev/null +++ b/src/lerobot/policies/fastwam/wan/components.py @@ -0,0 +1,175 @@ +# 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. + +from __future__ import annotations + +import logging +from collections.abc import Sequence +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import torch +from huggingface_hub import snapshot_download +from safetensors.torch import load_file + +from lerobot.utils.import_utils import _diffusers_available, _transformers_available, require_package + +if TYPE_CHECKING or _transformers_available: + from transformers import AutoTokenizer, UMT5EncoderModel +else: + AutoTokenizer = None + UMT5EncoderModel = None + +if TYPE_CHECKING or _diffusers_available: + from diffusers import AutoencoderKLWan +else: + AutoencoderKLWan = None + +from .adapters import WanVideoVAE38 +from .video_dit import WanVideoDiT + +logger = logging.getLogger(__name__) + +# The custom MoT video DiT still ships in the original (non-diffusers) Wan2.2 +# repo as sharded `diffusion_pytorch_model*.safetensors`; the VAE and UMT5 text +# encoder come from the diffusers conversion. Tokenizer is the stock UMT5 one. +WAN_DIT_PATTERN = "diffusion_pytorch_model*.safetensors" +WAN_T5_TOKENIZER = "google/umt5-xxl" +WAN22_DIFFUSERS_MODEL_ID = "Wan-AI/Wan2.2-TI2V-5B-Diffusers" + + +class WanTextEncoder(torch.nn.Module): + """FastWAM text-encoder contract over `transformers.UMT5EncoderModel`. + + Exposes `.dim` (hidden size) and `forward(ids, mask) -> [B, L, dim]`, matching + the call in `FastWAM.encode_prompt`. + """ + + def __init__( + self, + dtype: torch.dtype = torch.bfloat16, + device: str | torch.device = "cuda", + *, + pretrained: torch.nn.Module, + ) -> None: + super().__init__() + # UMT5-XXL is a fixed pretrained encoder — never trained from scratch, so a real + # `UMT5EncoderModel` (with weights) must always be supplied (loaded from the + # diffusers repo by `load_pretrained_wan_text_encoder`). No random/offline build. + self.model = pretrained.to(device=device, dtype=dtype) + self.dim = int(self.model.config.d_model) + + def forward(self, ids: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: + return self.model(input_ids=ids, attention_mask=mask.long()).last_hidden_state + + +class WanTokenizer: + """UMT5 tokenizer wrapper returning `(input_ids, attention_mask)` like the + FastWAM call site expects.""" + + def __init__(self, name: str = WAN_T5_TOKENIZER, seq_len: int = 512) -> None: + require_package("transformers", extra="fastwam") + self.tokenizer = AutoTokenizer.from_pretrained(name) + self.seq_len = int(seq_len) + + def __call__( + self, + sequence: str | Sequence[str], + return_mask: bool = False, + add_special_tokens: bool = True, + **_: Any, + ): + if isinstance(sequence, str): + sequence = [sequence] + out = self.tokenizer( + list(sequence), + padding="max_length", + truncation=True, + max_length=self.seq_len, + add_special_tokens=add_special_tokens, + return_tensors="pt", + ) + if return_mask: + return out.input_ids, out.attention_mask + return out.input_ids + + +def build_wan_tokenizer(*, model_id: str = WAN_T5_TOKENIZER, tokenizer_max_len: int) -> WanTokenizer: + return WanTokenizer(name=model_id, seq_len=int(tokenizer_max_len)) + + +def load_pretrained_wan_vae(*, torch_dtype: torch.dtype, device: str) -> WanVideoVAE38: + """Load real Wan2.2 VAE weights from the diffusers repo (offline base creation).""" + require_package("diffusers", extra="fastwam") + vae = AutoencoderKLWan.from_pretrained(WAN22_DIFFUSERS_MODEL_ID, subfolder="vae", torch_dtype=torch_dtype) + return WanVideoVAE38(dtype=torch_dtype, device=device, pretrained=vae) + + +def load_pretrained_wan_text_encoder( + *, + model_id: str = WAN22_DIFFUSERS_MODEL_ID, + subfolder: str | None = "text_encoder", + torch_dtype: torch.dtype, + device: str, +) -> WanTextEncoder: + """Load UMT5-XXL encoder weights (defaults to the Wan2.2 diffusers repo). + + Must stay compatible with the tokenizer (see `build_wan_tokenizer`): the encoder's + embedding table is indexed by the tokenizer's vocabulary. + """ + require_package("transformers", extra="fastwam") + encoder = UMT5EncoderModel.from_pretrained(model_id, subfolder=subfolder, torch_dtype=torch_dtype) + return WanTextEncoder(dtype=torch_dtype, device=device, pretrained=encoder) + + +def resolve_wan_dit_paths( + model_id_or_path: str | Path, + *, + cache_dir: str | Path | None = None, + local_files_only: bool = False, + revision: str | None = None, +) -> list[Path]: + """Resolve the custom MoT DiT shards from the original Wan2.2 repo or a local dir.""" + path = Path(model_id_or_path).expanduser() + if path.is_dir(): + return sorted(path.glob(WAN_DIT_PATTERN)) + + snapshot_path = snapshot_download( + repo_id=str(model_id_or_path), + revision=revision, + cache_dir=cache_dir, + local_files_only=local_files_only, + allow_patterns=[WAN_DIT_PATTERN], + ) + return sorted(Path(snapshot_path).glob(WAN_DIT_PATTERN)) + + +def load_wan_video_dit( + paths: list[str | Path], + *, + dit_config: dict[str, Any], + torch_dtype: torch.dtype, + device: str, +) -> WanVideoDiT: + model = WanVideoDiT(**dit_config) + state_dict = _read_wan_dit_safetensors(paths) + model.load_state_dict(state_dict, strict=False) + return model.to(device=device, dtype=torch_dtype) + + +def _read_wan_dit_safetensors(paths: list[str | Path]) -> dict[str, torch.Tensor]: + state_dict = {} + for path in paths: + state_dict.update(load_file(str(path), device="cpu")) + return state_dict diff --git a/src/lerobot/policies/fastwam/wan/model.py b/src/lerobot/policies/fastwam/wan/model.py new file mode 100644 index 000000000..329d1c48a --- /dev/null +++ b/src/lerobot/policies/fastwam/wan/model.py @@ -0,0 +1,341 @@ +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. +import math + +import torch +import torch.nn as nn + + +def sinusoidal_embedding_1d(dim, position): + # preprocess + if dim % 2 != 0: + raise ValueError(f"dim must be even, got {dim}.") + half = dim // 2 + position = position.type(torch.float64) + + # calculation + sinusoid = torch.outer(position, torch.pow(10000, -torch.arange(half).to(position).div(half))) + x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1) + return x + + +@torch.amp.autocast("cuda", enabled=False) +def rope_params(max_seq_len, dim, theta=10000): + if dim % 2 != 0: + raise ValueError(f"dim must be even, got {dim}.") + freqs = torch.outer( + torch.arange(max_seq_len), 1.0 / torch.pow(theta, torch.arange(0, dim, 2).to(torch.float64).div(dim)) + ) + freqs = torch.polar(torch.ones_like(freqs), freqs) + return freqs + + +@torch.amp.autocast("cuda", enabled=False) +def rope_apply(x, grid_sizes, freqs): + n, c = x.size(2), x.size(3) // 2 + + # split freqs + freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1) + + # loop over samples + output = [] + for i, (f, h, w) in enumerate(grid_sizes.tolist()): + seq_len = f * h * w + + # precompute multipliers + x_i = torch.view_as_complex(x[i, :seq_len].to(torch.float64).reshape(seq_len, n, -1, 2)) + freqs_i = torch.cat( + [ + freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1), + freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1), + freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1), + ], + dim=-1, + ).reshape(seq_len, 1, -1) + + # apply rotary embedding + x_i = torch.view_as_real(x_i * freqs_i).flatten(2) + x_i = torch.cat([x_i, x[i, seq_len:]]) + + # append to collection + output.append(x_i) + return torch.stack(output).float() + + +class WanRMSNorm(nn.Module): + def __init__(self, dim, eps=1e-5): + super().__init__() + self.dim = dim + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x): + r""" + Args: + x(Tensor): Shape [B, L, C] + """ + return self._norm(x.float()).type_as(x) * self.weight + + def _norm(self, x): + return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps) + + +class WanLayerNorm(nn.LayerNorm): + def __init__(self, dim, eps=1e-6, elementwise_affine=False): + super().__init__(dim, elementwise_affine=elementwise_affine, eps=eps) + + def forward(self, x): + r""" + Args: + x(Tensor): Shape [B, L, C] + """ + return super().forward(x.float()).type_as(x) + + +class WanSelfAttention(nn.Module): + def __init__(self, dim, num_heads, qk_norm=True, eps=1e-6): + if dim % num_heads != 0: + raise ValueError(f"dim ({dim}) must be divisible by num_heads ({num_heads}).") + super().__init__() + self.num_heads = num_heads + self.head_dim = dim // num_heads + + # layers + self.q = nn.Linear(dim, dim) + self.k = nn.Linear(dim, dim) + self.v = nn.Linear(dim, dim) + self.o = nn.Linear(dim, dim) + self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + + # NOTE: FastWAM never runs the upstream Wan attention forward. FastWAMAttentionBlock + # reuses only the q/k/v/o/norm submodules defined above and computes attention via + # `fastwam_masked_attention` (SDPA). The original flash-attention forward was removed, + # which also collapsed the former WanCrossAttention subclass into this class (it only + # differed by its forward): self- and cross-attention now share the same projection module. + + +class WanAttentionBlock(nn.Module): + def __init__(self, dim, ffn_dim, num_heads, qk_norm=True, cross_attn_norm=False, eps=1e-6): + super().__init__() + self.dim = dim + self.ffn_dim = ffn_dim + self.num_heads = num_heads + self.qk_norm = qk_norm + self.cross_attn_norm = cross_attn_norm + self.eps = eps + + # layers + self.norm1 = WanLayerNorm(dim, eps) + self.self_attn = WanSelfAttention(dim, num_heads, qk_norm, eps) + self.norm3 = WanLayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity() + self.cross_attn = WanSelfAttention(dim, num_heads, qk_norm, eps) + self.norm2 = WanLayerNorm(dim, eps) + self.ffn = nn.Sequential( + nn.Linear(dim, ffn_dim), nn.GELU(approximate="tanh"), nn.Linear(ffn_dim, dim) + ) + + # modulation + self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + + # NOTE: The upstream Wan block forward (self-attention + cross-attention + FFN via + # flash-attention) was removed. FastWAM subclasses this block as FastWAMAttentionBlock + # and overrides forward to use SDPA with explicit boolean masks; only __init__ (the + # norm/attention/ffn submodules) is reused here. + + +class Head(nn.Module): + def __init__(self, dim, out_dim, patch_size, eps=1e-6): + super().__init__() + self.dim = dim + self.out_dim = out_dim + self.patch_size = patch_size + self.eps = eps + + # layers + out_dim = math.prod(patch_size) * out_dim + self.norm = WanLayerNorm(dim, eps) + self.head = nn.Linear(dim, out_dim) + + # modulation + self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5) + + def forward(self, x, e): + r""" + Args: + x(Tensor): Shape [B, L1, C] + e(Tensor): Shape [B, L1, C] + """ + with torch.amp.autocast("cuda", dtype=torch.float32): + e = (self.modulation.unsqueeze(0) + e.unsqueeze(2)).chunk(2, dim=2) + x = self.head(self.norm(x) * (1 + e[1].squeeze(2)) + e[0].squeeze(2)) + return x + + +class WanModel(nn.Module): + r""" + Wan diffusion backbone supporting both text-to-video and image-to-video. + """ + + def __init__( + self, + model_type="t2v", + patch_size=(1, 2, 2), + text_len=512, + in_dim=16, + dim=2048, + ffn_dim=8192, + freq_dim=256, + text_dim=4096, + out_dim=16, + num_heads=16, + num_layers=32, + qk_norm=True, + cross_attn_norm=True, + eps=1e-6, + ): + r""" + Initialize the diffusion model backbone. + + Args: + model_type (`str`, *optional*, defaults to 't2v'): + Model variant - 't2v' (text-to-video) or 'i2v' (image-to-video) + patch_size (`tuple`, *optional*, defaults to (1, 2, 2)): + 3D patch dimensions for video embedding (t_patch, h_patch, w_patch) + text_len (`int`, *optional*, defaults to 512): + Fixed length for text embeddings + in_dim (`int`, *optional*, defaults to 16): + Input video channels (C_in) + dim (`int`, *optional*, defaults to 2048): + Hidden dimension of the transformer + ffn_dim (`int`, *optional*, defaults to 8192): + Intermediate dimension in feed-forward network + freq_dim (`int`, *optional*, defaults to 256): + Dimension for sinusoidal time embeddings + text_dim (`int`, *optional*, defaults to 4096): + Input dimension for text embeddings + out_dim (`int`, *optional*, defaults to 16): + Output video channels (C_out) + num_heads (`int`, *optional*, defaults to 16): + Number of attention heads + num_layers (`int`, *optional*, defaults to 32): + Number of transformer blocks + qk_norm (`bool`, *optional*, defaults to True): + Enable query/key normalization + cross_attn_norm (`bool`, *optional*, defaults to False): + Enable cross-attention normalization + eps (`float`, *optional*, defaults to 1e-6): + Epsilon value for normalization layers + """ + + super().__init__() + + if model_type not in ["t2v", "i2v", "ti2v", "s2v"]: + raise ValueError(f"model_type must be one of ['t2v', 'i2v', 'ti2v', 's2v'], got {model_type!r}.") + self.model_type = model_type + + self.patch_size = patch_size + self.text_len = text_len + self.in_dim = in_dim + self.dim = dim + self.ffn_dim = ffn_dim + self.freq_dim = freq_dim + self.text_dim = text_dim + self.out_dim = out_dim + self.num_heads = num_heads + self.num_layers = num_layers + self.qk_norm = qk_norm + self.cross_attn_norm = cross_attn_norm + self.eps = eps + + # embeddings + self.patch_embedding = nn.Conv3d(in_dim, dim, kernel_size=patch_size, stride=patch_size) + self.text_embedding = nn.Sequential( + nn.Linear(text_dim, dim), nn.GELU(approximate="tanh"), nn.Linear(dim, dim) + ) + + self.time_embedding = nn.Sequential(nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim)) + self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(dim, dim * 6)) + + # blocks + self.blocks = nn.ModuleList( + [ + WanAttentionBlock(dim, ffn_dim, num_heads, qk_norm, cross_attn_norm, eps) + for _ in range(num_layers) + ] + ) + + # head + self.head = Head(dim, out_dim, patch_size, eps) + + # buffers (don't use register_buffer otherwise dtype will be changed in to()) + if (dim % num_heads) != 0 or (dim // num_heads) % 2 != 0: + raise ValueError( + f"dim ({dim}) must be divisible by num_heads ({num_heads}) with an even head dim." + ) + d = dim // num_heads + self.freqs = torch.cat( + [ + rope_params(1024, d - 4 * (d // 6)), + rope_params(1024, 2 * (d // 6)), + rope_params(1024, 2 * (d // 6)), + ], + dim=1, + ) + + # initialize weights + self.init_weights() + + # NOTE: The upstream Wan diffusion forward (flash-attention based) was removed. + # FastWAM's WanVideoDiT subclasses this model, rebuilds `self.blocks` with + # FastWAMAttentionBlock, and provides its own SDPA-based forward. Only the + # constructor (embeddings, blocks, head, rope buffers) and the helpers below + # (unpatchify / init_weights) are reused. WanModel is never run directly. + + def unpatchify(self, x, grid_sizes): + r""" + Reconstruct video tensors from patch embeddings. + + Args: + x (List[Tensor]): + List of patchified features, each with shape [L, C_out * prod(patch_size)] + grid_sizes (Tensor): + Original spatial-temporal grid dimensions before patching, + shape [B, 3] (3 dimensions correspond to F_patches, H_patches, W_patches) + + Returns: + List[Tensor]: + Reconstructed video tensors with shape [C_out, F, H / 8, W / 8] + """ + + c = self.out_dim + out = [] + for u, v in zip(x, grid_sizes.tolist(), strict=False): + u = u[: math.prod(v)].view(*v, *self.patch_size, c) + u = torch.einsum("fhwpqrc->cfphqwr", u) + u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size, strict=False)]) + out.append(u) + return out + + def init_weights(self): + r""" + Initialize model parameters using Xavier initialization. + """ + + # basic init + for m in self.modules(): + if isinstance(m, nn.Linear): + nn.init.xavier_uniform_(m.weight) + if m.bias is not None: + nn.init.zeros_(m.bias) + + # init embeddings + nn.init.xavier_uniform_(self.patch_embedding.weight.flatten(1)) + for m in self.text_embedding.modules(): + if isinstance(m, nn.Linear): + nn.init.normal_(m.weight, std=0.02) + for m in self.time_embedding.modules(): + if isinstance(m, nn.Linear): + nn.init.normal_(m.weight, std=0.02) + + # init output layer + nn.init.zeros_(self.head.head.weight) diff --git a/src/lerobot/policies/fastwam/wan/modular.py b/src/lerobot/policies/fastwam/wan/modular.py new file mode 100644 index 000000000..fac96776b --- /dev/null +++ b/src/lerobot/policies/fastwam/wan/modular.py @@ -0,0 +1,1912 @@ +# 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. + +from __future__ import annotations + +import logging +import re +from collections.abc import Sequence +from typing import Any + +import torch +import torch.nn as nn +import torch.nn.functional as functional +from PIL import Image + +from .components import ( + WAN22_DIFFUSERS_MODEL_ID, + WAN_T5_TOKENIZER, + build_wan_tokenizer, + load_pretrained_wan_text_encoder, + load_pretrained_wan_vae, + load_wan_video_dit, + resolve_wan_dit_paths, +) +from .video_dit import ( + FastWAMAttentionBlock, + WanContinuousFlowMatchScheduler, + fastwam_masked_attention, + gradient_checkpoint_forward, + modulate, + precompute_freqs_cis, + sinusoidal_embedding_1d, +) + +logger = logging.getLogger(__name__) + + +def _apply_block_norm(block, name: str, x: torch.Tensor) -> torch.Tensor: + apply_norm = getattr(block, f"apply_{name}", None) + if apply_norm is not None: + return apply_norm(x) + return getattr(block, name)(x) + + +class ActionHead(nn.Module): + def __init__(self, hidden_dim: int, out_dim: int, eps: float): + super().__init__() + self.norm = nn.LayerNorm(hidden_dim, eps=eps, elementwise_affine=False) + self.proj = nn.Linear(hidden_dim, out_dim) + self.modulation = nn.Parameter(torch.randn(1, 2, hidden_dim) / hidden_dim**0.5) + + def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor: + shift, scale = (self.modulation.to(dtype=t.dtype, device=t.device) + t.unsqueeze(1)).chunk(2, dim=1) + shift = shift.squeeze(1) + scale = scale.squeeze(1) + return self.proj(self.norm(x) * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)) + + +class ActionDiT(nn.Module): + def __init__( + self, + hidden_dim: int, + action_dim: int, + ffn_dim: int, + text_dim: int, + freq_dim: int, + eps: float, + num_heads: int, + attn_head_dim: int, + num_layers: int, + use_gradient_checkpointing: bool = False, + fp32_attention: bool = True, + ): + super().__init__() + self.hidden_dim = hidden_dim + self.action_dim = action_dim + self.ffn_dim = ffn_dim + self.text_dim = text_dim + self.freq_dim = freq_dim + self.num_heads = num_heads + self.attn_head_dim = attn_head_dim + + if num_heads <= 0: + raise ValueError(f"`num_heads` must be > 0, got {num_heads}") + if attn_head_dim <= 0: + raise ValueError(f"`attn_head_dim` must be > 0, got {attn_head_dim}") + if attn_head_dim % 2 != 0: + raise ValueError(f"`attn_head_dim` must be even for RoPE, got {attn_head_dim}") + + self.action_encoder = nn.Linear(action_dim, hidden_dim) + self.text_embedding = nn.Sequential( + nn.Linear(text_dim, hidden_dim), + nn.GELU(approximate="tanh"), + nn.Linear(hidden_dim, hidden_dim), + ) + self.time_embedding = nn.Sequential( + nn.Linear(freq_dim, hidden_dim), + nn.SiLU(), + nn.Linear(hidden_dim, hidden_dim), + ) + self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(hidden_dim, hidden_dim * 6)) + self.blocks = nn.ModuleList( + [ + FastWAMAttentionBlock( + hidden_dim=hidden_dim, + attn_head_dim=attn_head_dim, + num_heads=num_heads, + ffn_dim=ffn_dim, + eps=eps, + fp32_attention=fp32_attention, + ) + for _ in range(num_layers) + ] + ) + self.head = nn.Linear(hidden_dim, action_dim) + self.freqs = precompute_freqs_cis(attn_head_dim, end=1024) + self.fp32_attention = bool(fp32_attention) + + self.use_gradient_checkpointing = use_gradient_checkpointing + + def pre_dit( + self, + action_tokens: torch.Tensor, + timestep: torch.Tensor, + context: torch.Tensor, + context_mask: torch.Tensor | None = None, + ) -> dict[str, Any]: + if action_tokens.ndim != 3: + raise ValueError( + f"`action_tokens` must be 3D [B, T, action_dim], got shape {tuple(action_tokens.shape)}" + ) + if action_tokens.shape[2] != self.action_dim: + raise ValueError( + f"`action_tokens` last dim must be {self.action_dim}, got {action_tokens.shape[2]}" + ) + if timestep.ndim != 1: + raise ValueError(f"`timestep` must be 1D [B] or [1], got shape {tuple(timestep.shape)}") + if context.ndim != 3: + raise ValueError(f"`context` must be 3D [B, L, D], got shape {tuple(context.shape)}") + + batch_size = action_tokens.shape[0] + if context.shape[0] != batch_size: + raise ValueError( + f"Batch mismatch between action tokens and text context: {batch_size} vs {context.shape[0]}" + ) + if timestep.shape[0] not in (1, batch_size): + raise ValueError( + f"`timestep` length must be 1 or batch_size({batch_size}), got {timestep.shape[0]}" + ) + if timestep.shape[0] == 1 and batch_size > 1: + if self.training: + raise ValueError("During training, action timestep length must match batch_size.") + timestep = timestep.expand(batch_size) + + if context_mask is None: + context_mask = torch.ones((batch_size, context.shape[1]), dtype=torch.bool, device=context.device) + else: + if context_mask.ndim != 2: + raise ValueError(f"`context_mask` must be 2D [B, L], got shape {tuple(context_mask.shape)}") + if context_mask.shape[0] != batch_size or context_mask.shape[1] != context.shape[1]: + raise ValueError( + f"`context_mask` shape must match `context` shape [B, L], got {tuple(context_mask.shape)} vs {tuple(context.shape)}" + ) + + seq_len = action_tokens.shape[1] + if seq_len > self.freqs.shape[0]: + raise ValueError(f"Action token length {seq_len} exceeds RoPE cache {self.freqs.shape[0]}.") + + model_dtype = self.action_encoder.weight.dtype + action_tokens = action_tokens.to(dtype=model_dtype) + context = context.to(dtype=model_dtype) + t_emb = sinusoidal_embedding_1d(self.freq_dim, timestep).to(dtype=model_dtype) + t = self.time_embedding(t_emb) + t_mod = self.time_projection(t).unflatten(1, (6, self.hidden_dim)) + + tokens = self.action_encoder(action_tokens) + context_emb = self.text_embedding(context) + context_attn_mask = context_mask.unsqueeze(1).expand(-1, seq_len, -1) + freqs = self.freqs[:seq_len].view(seq_len, 1, -1).to(tokens.device) + + return { + "tokens": tokens, + "freqs": freqs, + "t": t, + "t_mod": t_mod, + "context": context_emb, + "context_mask": context_attn_mask, + "meta": { + "batch_size": batch_size, + "seq_len": seq_len, + }, + } + + def post_dit(self, tokens: torch.Tensor, pre_state: dict[str, Any]) -> torch.Tensor: + return self.head(tokens) + + def forward( + self, + action_tokens: torch.Tensor, + timestep: torch.Tensor, + context: torch.Tensor, + context_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + pre_state = self.pre_dit( + action_tokens=action_tokens, + timestep=timestep, + context=context, + context_mask=context_mask, + ) + x = pre_state["tokens"] + context = pre_state["context"] + t_mod = pre_state["t_mod"] + freqs = pre_state["freqs"] + context_mask = pre_state["context_mask"] + + for block in self.blocks: + if self.use_gradient_checkpointing: + x = gradient_checkpoint_forward( + block, + self.use_gradient_checkpointing, + x, + context, + t_mod, + freqs, + context_mask=context_mask, + ) + else: + x = block(x, context, t_mod, freqs, context_mask=context_mask) + + return self.post_dit(x, pre_state) + + +class MoTLayer(nn.Module): + """A single MoT layer: owns one transformer block per expert and runs the cross-expert + mixed-attention step for that layer. + + This exists as a module — rather than the per-layer work being inlined in ``MoT``'s loop — + so FSDP can wrap each layer as its own unit. FSDP all-gathers a wrapped module's sharded + parameters via a hook on that module's ``forward``/``__call__``. ``MoT`` drives block + submodules directly (the joint mixed attention concatenates Q/K/V across experts, so no + single block's ``forward`` is ever called), so ``MoTLayer.forward`` is the only call + boundary FSDP can hook. All three per-layer paths therefore dispatch through + ``forward(mode=...)`` so each enters via ``__call__``. + """ + + def __init__( + self, + blocks: dict[str, nn.Module], + experts: dict[str, nn.Module], + num_heads: int, + attn_head_dim: int, + fp32_attention: bool, + mot_checkpoint_mixed_attn: bool, + ): + super().__init__() + # Registered owner of this layer's blocks (one per expert) — the FSDP wrap unit. + self.blocks = nn.ModuleDict(blocks) + self.expert_order = list(blocks.keys()) + # Unregistered back-references to the experts: used only to read the live + # `use_gradient_checkpointing` flag, kept out of parameters()/state_dict(). + object.__setattr__(self, "_experts", dict(experts)) + self.num_heads = num_heads + self.attn_head_dim = attn_head_dim + self.fp32_attention = bool(fp32_attention) + self.mot_checkpoint_mixed_attn = bool(mot_checkpoint_mixed_attn) + + @staticmethod + def _split_modulation(block, t_mod: torch.Tensor): + has_seq = len(t_mod.shape) == 4 + chunk_dim = 2 if has_seq else 1 + + base_mod = block.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (base_mod + t_mod).chunk( + 6, dim=chunk_dim + ) + if has_seq: + # means t_mod has separate modulation for each token, otherwise same modulation for all tokens in the block + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + shift_msa.squeeze(2), + scale_msa.squeeze(2), + gate_msa.squeeze(2), + shift_mlp.squeeze(2), + scale_mlp.squeeze(2), + gate_mlp.squeeze(2), + ) + return shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp + + def _mixed_attention( + self, + q_cat: torch.Tensor, + k_cat: torch.Tensor, + v_cat: torch.Tensor, + attention_mask: torch.Tensor, + ) -> torch.Tensor: + attn_mask = attention_mask.to(device=q_cat.device) + + def _forward(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + return fastwam_masked_attention( + q=q, + k=k, + v=v, + num_heads=self.num_heads, + ctx_mask=attn_mask, + fp32_attention=self.fp32_attention, + ) + + if self.mot_checkpoint_mixed_attn and self.training: + return torch.utils.checkpoint.checkpoint( + _forward, + q_cat, + k_cat, + v_cat, + use_reentrant=False, + ) + return _forward(q_cat, k_cat, v_cat) + + @staticmethod + def _apply_expert_post_block( + block, + residual_x: torch.Tensor, + mixed_attn_out: torch.Tensor, + gate_msa: torch.Tensor, + shift_mlp: torch.Tensor, + scale_mlp: torch.Tensor, + gate_mlp: torch.Tensor, + context_payload: dict | None, + ) -> torch.Tensor: + if hasattr(block, "project_self_attention_output"): + projected_attn = block.project_self_attention_output(mixed_attn_out) + else: + projected_attn = block.self_attn.o(mixed_attn_out.to(dtype=block.self_attn.o.weight.dtype)) + x = residual_x + gate_msa * projected_attn + + if context_payload is not None: + context = context_payload.get("context") + if context is not None: + context_mask = context_payload.get("mask") + if context_mask is not None and context_mask.dim() == 3: + context_mask = context_mask.unsqueeze(1) + x = x + block.apply_cross_attention( + _apply_block_norm(block, "norm3", x), + context, + context_mask=context_mask, + ) + + mlp_input = modulate(_apply_block_norm(block, "norm2", x), shift_mlp, scale_mlp) + x = x + gate_mlp * block.ffn(mlp_input) + return x + + def _build_expert_attention_io( + self, + name: str, + x: torch.Tensor, + freqs: torch.Tensor | dict[str, torch.Tensor], + t_mod: torch.Tensor, + ): + """Build this expert's attention tensors and post-block states for the layer. + + Returns (q, k, v, residual_x, gate_msa, shift_mlp, scale_mlp, gate_mlp, use_gc). + """ + block = self.blocks[name] + expert = self._experts[name] + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self._split_modulation(block, t_mod) + attn_input = modulate(_apply_block_norm(block, "norm1", x), shift_msa, scale_msa) + + q, k, v = block.project_self_attention(attn_input, freqs) + + use_gradient_checkpointing = bool(getattr(expert, "use_gradient_checkpointing", False)) + return ( + q, + k, + v, + x, + gate_msa, + shift_mlp, + scale_mlp, + gate_mlp, + use_gradient_checkpointing, + ) + + def _apply_post_with_optional_checkpoint( + self, + block, + residual_x: torch.Tensor, + gate_msa: torch.Tensor, + shift_mlp: torch.Tensor, + scale_mlp: torch.Tensor, + gate_mlp: torch.Tensor, + use_gradient_checkpointing: bool, + mixed_slice: torch.Tensor, + context_payload: dict | None, + ) -> torch.Tensor: + def _post_fn( + _mixed_slice: torch.Tensor, + _x: torch.Tensor, + _gate_msa: torch.Tensor, + _shift_mlp: torch.Tensor, + _scale_mlp: torch.Tensor, + _gate_mlp: torch.Tensor, + _block=block, + _context_payload=context_payload, + ) -> torch.Tensor: + return self._apply_expert_post_block( + block=_block, + residual_x=_x, + mixed_attn_out=_mixed_slice, + gate_msa=_gate_msa, + shift_mlp=_shift_mlp, + scale_mlp=_scale_mlp, + gate_mlp=_gate_mlp, + context_payload=_context_payload, + ) + + if use_gradient_checkpointing and self.training: + return torch.utils.checkpoint.checkpoint( + _post_fn, + mixed_slice, + residual_x, + gate_msa, + shift_mlp, + scale_mlp, + gate_mlp, + use_reentrant=False, + ) + return _post_fn( + mixed_slice, + residual_x, + gate_msa, + shift_mlp, + scale_mlp, + gate_mlp, + ) + + def forward(self, mode: str, **kwargs): + if mode == "joint": + return self._forward_joint(**kwargs) + if mode == "video_prefill": + return self._forward_video_prefill(**kwargs) + if mode == "action_cached": + return self._forward_action_cached(**kwargs) + raise ValueError(f"Unknown MoTLayer forward mode: {mode!r}") + + def _forward_joint( + self, + tokens_all: dict[str, torch.Tensor], + attention_mask: torch.Tensor, + freqs_all: dict[str, torch.Tensor], + context_all: dict[str, dict | None], + t_mod_all: dict[str, torch.Tensor], + ) -> dict[str, torch.Tensor]: + q_chunks = [] + k_chunks = [] + v_chunks = [] + cached = {} + seq_lens = [] + + for name in self.expert_order: + ( + q, + k, + v, + residual_x, + gate_msa, + shift_mlp, + scale_mlp, + gate_mlp, + use_gradient_checkpointing, + ) = self._build_expert_attention_io(name, tokens_all[name], freqs_all[name], t_mod_all[name]) + + q_chunks.append(q) + k_chunks.append(k) + v_chunks.append(v) + seq_lens.append(tokens_all[name].shape[1]) + cached[name] = { + "residual_x": residual_x, + "gate_msa": gate_msa, + "shift_mlp": shift_mlp, + "scale_mlp": scale_mlp, + "gate_mlp": gate_mlp, + "use_gradient_checkpointing": use_gradient_checkpointing, + } + + q_cat = torch.cat(q_chunks, dim=1) + k_cat = torch.cat(k_chunks, dim=1) + v_cat = torch.cat(v_chunks, dim=1) + + total_seq = q_cat.shape[1] + if attention_mask.shape[0] != total_seq: + raise ValueError( + f"Attention mask seq length mismatch: mask={attention_mask.shape[0]} vs tokens={total_seq}" + ) + + mixed = self._mixed_attention(q_cat=q_cat, k_cat=k_cat, v_cat=v_cat, attention_mask=attention_mask) + + out = {} + start = 0 + for name, seq_len in zip(self.expert_order, seq_lens, strict=True): + end = start + seq_len + mixed_slice = mixed[:, start:end, :] + cached_expert = cached[name] + out[name] = self._apply_post_with_optional_checkpoint( + block=self.blocks[name], + residual_x=cached_expert["residual_x"], + gate_msa=cached_expert["gate_msa"], + shift_mlp=cached_expert["shift_mlp"], + scale_mlp=cached_expert["scale_mlp"], + gate_mlp=cached_expert["gate_mlp"], + use_gradient_checkpointing=cached_expert["use_gradient_checkpointing"], + mixed_slice=mixed_slice, + context_payload=context_all.get(name), + ) + start = end + return out + + def _forward_video_prefill( + self, + x: torch.Tensor, + freqs: torch.Tensor, + t_mod: torch.Tensor, + context_payload: dict | None, + video_attention_mask: torch.Tensor, + ): + ( + q, + k, + v, + residual_x, + gate_msa, + shift_mlp, + scale_mlp, + gate_mlp, + use_gradient_checkpointing, + ) = self._build_expert_attention_io("video", x, freqs, t_mod) + # Video prefill uses only video self-attention mask. + mixed = self._mixed_attention(q_cat=q, k_cat=k, v_cat=v, attention_mask=video_attention_mask) + x_out = self._apply_post_with_optional_checkpoint( + block=self.blocks["video"], + residual_x=residual_x, + gate_msa=gate_msa, + shift_mlp=shift_mlp, + scale_mlp=scale_mlp, + gate_mlp=gate_mlp, + use_gradient_checkpointing=use_gradient_checkpointing, + mixed_slice=mixed, + context_payload=context_payload, + ) + return x_out, k, v + + def _forward_action_cached( + self, + x: torch.Tensor, + freqs: torch.Tensor, + t_mod: torch.Tensor, + context_payload: dict | None, + k_video: torch.Tensor, + v_video: torch.Tensor, + action_attention_mask: torch.Tensor, + ) -> torch.Tensor: + ( + q_action, + k_action, + v_action, + residual_x, + gate_msa, + shift_mlp, + scale_mlp, + gate_mlp, + use_gradient_checkpointing, + ) = self._build_expert_attention_io("action", x, freqs, t_mod) + # Mixed attention: action queries attend to cached video K/V plus current action K/V. + k_cat = torch.cat([k_video, k_action], dim=1) + v_cat = torch.cat([v_video, v_action], dim=1) + mixed = self._mixed_attention( + q_cat=q_action, k_cat=k_cat, v_cat=v_cat, attention_mask=action_attention_mask + ) + return self._apply_post_with_optional_checkpoint( + block=self.blocks["action"], + residual_x=residual_x, + gate_msa=gate_msa, + shift_mlp=shift_mlp, + scale_mlp=scale_mlp, + gate_mlp=gate_mlp, + use_gradient_checkpointing=use_gradient_checkpointing, + mixed_slice=mixed, + context_payload=context_payload, + ) + + +class MoT(nn.Module): + def __init__( + self, + mixtures: dict[str, nn.Module], + mot_checkpoint_mixed_attn: bool = True, + ): + super().__init__() + if not mixtures: + raise ValueError("`mixtures` cannot be empty.") + if "video" not in mixtures or "action" not in mixtures: + raise ValueError("`mixtures` must include both 'video' and 'action' experts.") + + self.mixtures = nn.ModuleDict(mixtures) + self.expert_order = list(self.mixtures.keys()) + self.mot_checkpoint_mixed_attn = mot_checkpoint_mixed_attn + if mot_checkpoint_mixed_attn: + logger.info( + "Using gradient checkpointing for mixture attention. This will save memory but use more computation." + ) + + first_expert = self.mixtures[self.expert_order[0]] + self.num_layers = len(first_expert.blocks) + self.num_heads = first_expert.num_heads + self.attn_head_dim = first_expert.attn_head_dim + self.fp32_attention = bool(getattr(first_expert, "fp32_attention", True)) + + for name in self.expert_order[1:]: + expert = self.mixtures[name] + if len(expert.blocks) != self.num_layers: + raise ValueError( + f"All experts must have same number of layers; got {self.num_layers} and {len(expert.blocks)}" + ) + if expert.num_heads != self.num_heads: + raise ValueError( + f"All experts must have same num_heads; got {self.num_heads} and {expert.num_heads}" + ) + if expert.attn_head_dim != self.attn_head_dim: + raise ValueError( + "All experts must have same attn_head_dim; " + f"got {self.attn_head_dim} and {expert.attn_head_dim}" + ) + if bool(getattr(expert, "fp32_attention", True)) != self.fp32_attention: + raise ValueError("All experts must use the same `fp32_attention` setting.") + + logger.info(f"Initialized MoT with experts: {self.expert_order}, num_layers={self.num_layers}") + for name in self.expert_order: + expert = self.mixtures[name] + logger.info( + f" Expert '{name}': num_params={sum(p.numel() for p in expert.parameters()) / 1e9:.2f} B" + ) + + # One MoTLayer per layer, each owning that layer's block from every expert. This is the + # FSDP wrap unit: only MoTLayer.forward is ever called (MoT drives block submodules + # directly for the cross-expert mixed attention), so it is the boundary at which FSDP can + # all-gather a layer's params. The blocks are RE-PARENTED into the layers — removed from + # each expert's module registry — so they have a single owner; leaving them registered + # under both the expert and the layer would make FSDP try to manage the same params twice. + self.layers = nn.ModuleList( + [ + MoTLayer( + blocks={name: self.mixtures[name].blocks[layer_idx] for name in self.expert_order}, + experts={name: self.mixtures[name] for name in self.expert_order}, + num_heads=self.num_heads, + attn_head_dim=self.attn_head_dim, + fp32_attention=self.fp32_attention, + mot_checkpoint_mixed_attn=self.mot_checkpoint_mixed_attn, + ) + for layer_idx in range(self.num_layers) + ] + ) + for name in self.expert_order: + expert = self.mixtures[name] + kept_blocks = list(expert.blocks) + del expert._modules["blocks"] + # Keep an UNREGISTERED reference so the (unused) standalone `expert.forward` and any + # `len(expert.blocks)` still work, without re-adding the params to the expert's + # parameters()/state_dict() (which would double-register them with the MoTLayer owner). + object.__setattr__(expert, "blocks", kept_blocks) + + def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs): + # Backward-compat for checkpoints saved before the MoTLayer refactor. Then the per-layer + # blocks were keyed under the experts (`{prefix}mixtures..blocks..`, e.g. + # the released `ZibinDong/fastwam_libero_uncond_2cam224`); now they are owned by the layers + # (`{prefix}layers..blocks..`). Remap legacy keys in place so the recursion + # into `self.layers` finds them and the (now block-less) `self.mixtures` does not flag them. + legacy = re.compile(re.escape(prefix) + r"mixtures\.([^.]+)\.blocks\.(\d+)\.(.+)$") + moved = {} + for key in list(state_dict.keys()): + m = legacy.match(key) + if m is not None: + name, layer_idx, rest = m.group(1), m.group(2), m.group(3) + moved[f"{prefix}layers.{layer_idx}.blocks.{name}.{rest}"] = state_dict.pop(key) + state_dict.update(moved) + super()._load_from_state_dict(state_dict, prefix, *args, **kwargs) + + def prefill_video_cache( + self, + video_tokens: torch.Tensor, + video_freqs: torch.Tensor, + video_t_mod: torch.Tensor, + video_context_payload: dict | None, + video_attention_mask: torch.Tensor, + ) -> list[dict[str, torch.Tensor]]: + """Prefill video branch once and cache per-layer K/V for action denoising. + + Returns a list of length ``num_layers``, each entry ``{"k": ..., "v": ...}``. + """ + if "video" not in self.mixtures: + raise ValueError("MoT requires `video` expert for `prefill_video_cache`.") + if video_attention_mask.ndim != 2: + raise ValueError( + f"`video_attention_mask` must be 2D [S,S], got shape {tuple(video_attention_mask.shape)}" + ) + if video_attention_mask.shape[0] != video_attention_mask.shape[1]: + raise ValueError( + f"`video_attention_mask` must be square, got shape {tuple(video_attention_mask.shape)}" + ) + if video_attention_mask.shape[0] != video_tokens.shape[1]: + raise ValueError( + "`video_attention_mask` seq length mismatch: " + f"mask={video_attention_mask.shape[0]} vs tokens={video_tokens.shape[1]}" + ) + + x = video_tokens + kv_cache: list[dict[str, torch.Tensor]] = [] + for layer in self.layers: + x, k, v = layer( + mode="video_prefill", + x=x, + freqs=video_freqs, + t_mod=video_t_mod, + context_payload=video_context_payload, + video_attention_mask=video_attention_mask, + ) + kv_cache.append({"k": k, "v": v}) + return kv_cache + + def forward_action_with_video_cache( + self, + action_tokens: torch.Tensor, + action_freqs: torch.Tensor, + action_t_mod: torch.Tensor, + action_context_payload: dict | None, + video_kv_cache: list[dict[str, torch.Tensor]], + attention_mask: torch.Tensor, + video_seq_len: int, + ) -> torch.Tensor: + """Run action branch with cached video K/V instead of recomputing video tokens.""" + if "action" not in self.mixtures: + raise ValueError("MoT requires `action` expert for `forward_action_with_video_cache`.") + if len(video_kv_cache) != self.num_layers: + raise ValueError( + f"`video_kv_cache` must contain {self.num_layers} layers, got {len(video_kv_cache)}." + ) + if attention_mask.ndim != 2: + raise ValueError(f"`attention_mask` must be 2D [S,S], got shape {tuple(attention_mask.shape)}") + if attention_mask.shape[0] != attention_mask.shape[1]: + raise ValueError(f"`attention_mask` must be square, got shape {tuple(attention_mask.shape)}") + + action_seq_len = int(action_tokens.shape[1]) + total_seq_len = int(video_seq_len) + action_seq_len + if attention_mask.shape[0] != total_seq_len: + raise ValueError( + "`attention_mask` seq length mismatch: " + f"mask={attention_mask.shape[0]} vs expected_total={total_seq_len}" + ) + # Use the action query rows from the joint [video+action] mask. + action_attention_mask = attention_mask[video_seq_len:total_seq_len, :total_seq_len] + + x = action_tokens + for layer_idx, layer in enumerate(self.layers): + layer_cache = video_kv_cache[layer_idx] + if "k" not in layer_cache or "v" not in layer_cache: + raise ValueError(f"`video_kv_cache[{layer_idx}]` must contain `k` and `v`.") + k_video = layer_cache["k"] + v_video = layer_cache["v"] + if k_video.shape[1] != video_seq_len or v_video.shape[1] != video_seq_len: + raise ValueError(f"`video_kv_cache[{layer_idx}]` seq len mismatch, expected {video_seq_len}.") + x = layer( + mode="action_cached", + x=x, + freqs=action_freqs, + t_mod=action_t_mod, + context_payload=action_context_payload, + k_video=k_video, + v_video=v_video, + action_attention_mask=action_attention_mask, + ) + return x + + def forward( + self, + embeds_all: dict[str, torch.Tensor], + attention_mask: torch.Tensor, + freqs_all: dict[str, torch.Tensor], + context_all: dict[str, dict | None], + t_mod_all: dict[str, torch.Tensor], + ): + missing = [k for k in self.expert_order if k not in embeds_all] + if missing: + raise ValueError(f"Missing expert tokens for {missing}") + missing = [k for k in self.expert_order if k not in freqs_all] + if missing: + raise ValueError(f"Missing expert freqs for {missing}") + missing = [k for k in self.expert_order if k not in t_mod_all] + if missing: + raise ValueError(f"Missing expert t_mod for {missing}") + + if attention_mask.ndim != 2: + raise ValueError(f"`attention_mask` must be 2D [S, S], got shape {tuple(attention_mask.shape)}") + if attention_mask.shape[0] != attention_mask.shape[1]: + raise ValueError(f"`attention_mask` must be square, got shape {tuple(attention_mask.shape)}") + + # Each layer is a MoTLayer module; entering via __call__ lets FSDP all-gather that + # layer's params (the whole point of the per-layer split). + tokens_all = dict(embeds_all) + for layer in self.layers: + tokens_all = layer( + mode="joint", + tokens_all=tokens_all, + attention_mask=attention_mask, + freqs_all=freqs_all, + context_all=context_all, + t_mod_all=t_mod_all, + ) + return tokens_all + + +class FastWAM(torch.nn.Module): + """MoT world model with video/action experts.""" + + def __init__( + self, + video_expert, + action_expert: ActionDiT, + mot: MoT, + vae, + text_encoder=None, + tokenizer=None, + text_dim: int | None = None, + proprio_dim: int | None = None, + device: str = "cpu", + torch_dtype: torch.dtype = torch.float32, + video_train_shift: float = 5.0, + video_infer_shift: float = 5.0, + video_num_train_timesteps: int = 1000, + action_train_shift: float = 5.0, + action_infer_shift: float = 5.0, + action_num_train_timesteps: int = 1000, + loss_lambda_video: float = 1.0, + loss_lambda_action: float = 1.0, + ): + super().__init__() + self.mot = mot + # `video_expert` / `action_expert` are the very same module objects as + # `mot.mixtures["video"]` / `["action"]`, and `dit` is an alias of `mot`. Registering + # them as submodules too would give every expert tensor three names in `state_dict()` + # (`video_expert.*`, `mot.mixtures.video.*`, `dit.mixtures.video.*`) — a 3x-bloated + # gathered FSDP checkpoint and a doubled module tree for FSDP to traverse. Hold them as + # plain (unregistered) attributes instead — bypassing `nn.Module.__setattr__`, like the + # frozen vae/text_encoder below — so `mot` is the single registered owner and each tensor + # has one canonical name (`mot.mixtures.*` / `mot.layers.*`, matching the base checkpoint). + # Forward / freeze / optimizer code still reaches them by attribute, and device/dtype moves + # still apply via `mot`. (optimizer + freeze logic use `model.dit`.) + object.__setattr__(self, "video_expert", video_expert) + object.__setattr__(self, "action_expert", action_expert) + object.__setattr__(self, "dit", self.mot) + + # Frozen Wan2.2 components: bypass `nn.Module.__setattr__` so they are NOT + # registered as submodules. They are therefore excluded from `state_dict()` + # (lean checkpoints), `parameters()`, and DDP gradient sync, and are loaded + # with their real weights from the diffusers/transformers repos at construction. + # Device/dtype moves still reach them via the `_apply` override below. + object.__setattr__(self, "vae", vae) + object.__setattr__(self, "text_encoder", text_encoder) + self.tokenizer = tokenizer + vae.requires_grad_(False) + if text_encoder is not None: + text_encoder.requires_grad_(False) + if text_dim is None: + if self.text_encoder is None: + raise ValueError("`text_dim` is required when `text_encoder` is not loaded.") + text_dim = int(self.text_encoder.dim) + self.text_dim = int(text_dim) + self.proprio_dim = None if proprio_dim is None else int(proprio_dim) + if self.proprio_dim is not None: + self.proprio_encoder = nn.Linear(self.proprio_dim, self.text_dim).to(torch_dtype) + else: + self.proprio_encoder = None + + self.train_video_scheduler = WanContinuousFlowMatchScheduler( + num_train_timesteps=video_num_train_timesteps, + shift=video_train_shift, + ) + self.infer_video_scheduler = WanContinuousFlowMatchScheduler( + num_train_timesteps=video_num_train_timesteps, + shift=video_infer_shift, + ) + self.train_action_scheduler = WanContinuousFlowMatchScheduler( + num_train_timesteps=action_num_train_timesteps, + shift=action_train_shift, + ) + self.infer_action_scheduler = WanContinuousFlowMatchScheduler( + num_train_timesteps=action_num_train_timesteps, + shift=action_infer_shift, + ) + # Optional aliases for consistency with Wan22Core naming. + self.train_scheduler = self.train_video_scheduler + self.infer_scheduler = self.infer_video_scheduler + + self.device = torch.device(device) + self.torch_dtype = torch_dtype + self.loss_lambda_video = float(loss_lambda_video) + self.loss_lambda_action = float(loss_lambda_action) + + self.to(self.device) + + @classmethod + def from_wan22_pretrained( + cls, + device: str = "cuda", + torch_dtype: torch.dtype = torch.bfloat16, + model_id: str = "Wan-AI/Wan2.2-TI2V-5B", + tokenizer_model_id: str = WAN_T5_TOKENIZER, + text_encoder_model_id: str = WAN22_DIFFUSERS_MODEL_ID, + tokenizer_max_len: int = 512, + load_text_encoder: bool = True, + proprio_dim: int | None = None, + video_dit_config: dict[str, Any] | None = None, + action_dit_config: dict[str, Any] | None = None, + mot_checkpoint_mixed_attn: bool = True, + video_train_shift: float = 5.0, + video_infer_shift: float = 5.0, + video_num_train_timesteps: int = 1000, + action_train_shift: float = 5.0, + action_infer_shift: float = 5.0, + action_num_train_timesteps: int = 1000, + loss_lambda_video: float = 1.0, + loss_lambda_action: float = 1.0, + ): + if video_dit_config is None: + raise ValueError("`video_dit_config` is required for FastWAM.from_wan22_pretrained().") + if "text_dim" not in video_dit_config: + raise ValueError("`video_dit_config['text_dim']` is required for FastWAM.") + + # Custom MoT video DiT from the original Wan2.2 repo; frozen VAE / UMT5 from + # the diffusers conversion. This is the offline base-creation path; the + # weights it loads are then bundled into the FastWAM `model.safetensors`. + video_expert = load_wan_video_dit( + resolve_wan_dit_paths(model_id), + dit_config=video_dit_config, + torch_dtype=torch_dtype, + device=device, + ) + action_expert = ActionDiT(**action_dit_config).to(device=device, dtype=torch_dtype) + if int(action_expert.num_heads) != int(video_expert.num_heads): + raise ValueError("ActionDiT `num_heads` must match video expert for MoT mixed attention.") + if int(action_expert.attn_head_dim) != int(video_expert.attn_head_dim): + raise ValueError("ActionDiT `attn_head_dim` must match video expert for MoT mixed attention.") + if int(len(action_expert.blocks)) != int(len(video_expert.blocks)): + raise ValueError("ActionDiT `num_layers` must match video expert.") + + mot = MoT( + mixtures={"video": video_expert, "action": action_expert}, + mot_checkpoint_mixed_attn=mot_checkpoint_mixed_attn, + ) + + vae = load_pretrained_wan_vae(torch_dtype=torch_dtype, device=device) + text_encoder = ( + load_pretrained_wan_text_encoder( + model_id=text_encoder_model_id, torch_dtype=torch_dtype, device=device + ) + if load_text_encoder + else None + ) + tokenizer = build_wan_tokenizer(model_id=tokenizer_model_id, tokenizer_max_len=tokenizer_max_len) + + return cls( + video_expert=video_expert, + action_expert=action_expert, + mot=mot, + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + text_dim=int(video_dit_config["text_dim"]), + proprio_dim=proprio_dim, + device=device, + torch_dtype=torch_dtype, + video_train_shift=video_train_shift, + video_infer_shift=video_infer_shift, + video_num_train_timesteps=video_num_train_timesteps, + action_train_shift=action_train_shift, + action_infer_shift=action_infer_shift, + action_num_train_timesteps=action_num_train_timesteps, + loss_lambda_video=loss_lambda_video, + loss_lambda_action=loss_lambda_action, + ) + + def _apply(self, fn, *args, **kwargs): + # `.to()` / `.cuda()` / `.cpu()` and accelerate/DDP device moves all funnel + # through `_apply`, and the parent policy reaches us via `child._apply(fn)` + # (not `child.to()`). Propagate `fn` to the *unregistered* frozen VAE / text + # encoder here so they follow the rest of the model onto the right device, + # while staying out of `state_dict()` / `parameters()`. + super()._apply(fn, *args, **kwargs) + self.vae._apply(fn) + if self.text_encoder is not None: + self.text_encoder._apply(fn) + return self + + @staticmethod + def _check_resize_height_width(height, width, num_frames): + if height % 16 != 0: + height = (height + 15) // 16 * 16 + if width % 16 != 0: + width = (width + 15) // 16 * 16 + if num_frames % 4 != 1: + num_frames = (num_frames + 3) // 4 * 4 + 1 + return height, width, num_frames + + @torch.no_grad() + def encode_prompt(self, prompt: str | Sequence[str]): + if self.text_encoder is None or self.tokenizer is None: + raise ValueError( + "Prompt encoding requires loaded text encoder/tokenizer. " + "Set `load_text_encoder=true` or provide precomputed `context/context_mask`." + ) + ids, mask = self.tokenizer(prompt, return_mask=True, add_special_tokens=True) + ids = ids.to(self.device) + mask = mask.to(self.device, dtype=torch.bool) + prompt_emb = self.text_encoder(ids, mask) + seq_lens = mask.gt(0).sum(dim=1).long() + for i, v in enumerate(seq_lens): + prompt_emb[i, v:] = 0 + # Match FastWAM/Wan2.2 context semantics: padding embeddings are zeroed, + # while cross-attention still sees a fixed-length context. + mask = torch.ones_like(mask) + return prompt_emb.to(device=self.device), mask + + def _append_proprio_to_context( + self, + context: torch.Tensor, + context_mask: torch.Tensor, + proprio: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if self.proprio_encoder is None or proprio is None: + return context, context_mask + if proprio.ndim != 2: + raise ValueError(f"`proprio` must be 2D [B, D], got shape {tuple(proprio.shape)}") + if self.proprio_dim is None or proprio.shape[1] != self.proprio_dim: + raise ValueError(f"`proprio` last dim must be {self.proprio_dim}, got {proprio.shape[1]}") + proprio_token = self.proprio_encoder( + proprio.to(device=self.device, dtype=context.dtype).unsqueeze(1) + ).to(dtype=context.dtype) # [B, 1, D] + proprio_mask = torch.ones((context_mask.shape[0], 1), dtype=torch.bool, device=context_mask.device) + return ( + torch.cat([context, proprio_token], dim=1), + torch.cat([context_mask, proprio_mask], dim=1), + ) + + @torch.no_grad() + def _encode_video_latents(self, video_tensor, tiled=False, tile_size=(30, 52), tile_stride=(15, 26)): + # The Wan VAE expects pixels in [-1, 1]; model inputs arrive in [0, 1] (VISUAL is IDENTITY in + # the preprocessor — see configuration_fastwam.normalization_mapping). Map here, at the single + # video-encode boundary, so it is applied exactly once on every path. + video_tensor = video_tensor * 2.0 - 1.0 + z = self.vae.encode( + video_tensor, + device=self.device, + tiled=tiled, + tile_size=tile_size, + tile_stride=tile_stride, + ) + return z + + @torch.no_grad() + def _encode_input_image_latents_tensor( + self, input_image: torch.Tensor, tiled=False, tile_size=(30, 52), tile_stride=(15, 26) + ): + if input_image.ndim == 3: + input_image = input_image.unsqueeze(0) + if input_image.ndim != 4 or input_image.shape[0] != 1 or input_image.shape[1] != 3: + raise ValueError( + f"`input_image` must have shape [1,3,H,W] or [3,H,W], got {tuple(input_image.shape)}" + ) + # [0, 1] -> [-1, 1] for the Wan VAE (mirrors `_encode_video_latents`); single image-encode boundary. + input_image = input_image * 2.0 - 1.0 + image = input_image.to(device=self.device)[0].unsqueeze(1) + z = self.vae.encode( + [image], device=self.device, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride + ) + if isinstance(z, list): + z = z[0].unsqueeze(0) + return z + + def _decode_latents(self, latents, tiled=False, tile_size=(30, 52), tile_stride=(15, 26)): + video_tensor = self.vae.decode( + latents, device=self.device, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride + ) + video_tensor = video_tensor.squeeze(0).detach().float().clamp(-1, 1) + video_tensor = ((video_tensor + 1.0) * 127.5).to(torch.uint8).cpu() + frames = [] + for t in range(video_tensor.shape[1]): + frame = video_tensor[:, t].permute(1, 2, 0).numpy() + frames.append(Image.fromarray(frame)) + return frames + + def build_inputs(self, sample, tiled: bool = False): + video = sample["video"] + if "context" not in sample or "context_mask" not in sample: + raise ValueError("FastWAM training requires `sample['context']` and `sample['context_mask']`.") + context = sample["context"] + context_mask = sample["context_mask"] + proprio = sample.get("proprio", None) + if video.ndim != 5: + raise ValueError(f"`sample['video']` must be 5D [B, 3, T, H, W], got shape {tuple(video.shape)}") + if video.shape[1] != 3: + raise ValueError(f"`sample['video']` channel dimension must be 3, got shape {tuple(video.shape)}") + + batch_size, _, num_frames, height, width = video.shape + if height % 16 != 0 or width % 16 != 0: + raise ValueError(f"Video spatial dims must be multiples of 16, got H={height}, W={width}") + if num_frames % 4 != 1: + raise ValueError(f"Video T must satisfy T % 4 == 1, got T={num_frames}") + if num_frames <= 1: + raise ValueError(f"Video T must be > 1 for action-conditioned training, got T={num_frames}") + + if "action" not in sample: + raise ValueError("`sample['action']` is required for FastWAM training.") + + action = sample["action"] + if action.ndim != 3: + raise ValueError(f"`sample['action']` must be 3D [B, T, a_dim], got shape {tuple(action.shape)}") + action_horizon = int(action.shape[1]) + if action_horizon % (num_frames - 1) != 0: + raise ValueError( + f"`sample['action']` temporal dimension must be divisible by video transitions ({num_frames - 1}), got {action_horizon}" + ) + + action_is_pad = sample.get("action_is_pad", None) + if action_is_pad is not None: + if action_is_pad.ndim != 2: + raise ValueError( + f"`sample['action_is_pad']` must be 2D [B, T], got shape {tuple(action_is_pad.shape)}" + ) + if action_is_pad.shape[0] != batch_size or action_is_pad.shape[1] != action_horizon: + raise ValueError( + "`sample['action_is_pad']` shape mismatch: " + f"got {tuple(action_is_pad.shape)} vs expected ({batch_size}, {action_horizon})" + ) + + image_is_pad = sample.get("image_is_pad", None) + if image_is_pad is not None: + if image_is_pad.ndim != 2: + raise ValueError( + f"`sample['image_is_pad']` must be 2D [B, T], got shape {tuple(image_is_pad.shape)}" + ) + if image_is_pad.shape[0] != batch_size or image_is_pad.shape[1] != num_frames: + raise ValueError( + "`sample['image_is_pad']` shape mismatch: " + f"got {tuple(image_is_pad.shape)} vs expected ({batch_size}, {num_frames})" + ) + + input_video = video.to(device=self.device, dtype=self.torch_dtype, non_blocking=True) + input_latents = self._encode_video_latents(input_video, tiled=tiled) + + first_frame_latents = None + fuse_flag = False + if getattr(self.video_expert, "fuse_vae_embedding_in_latents", False): + first_frame_latents = input_latents[:, :, 0:1] + fuse_flag = True + + if context.ndim != 3 or context_mask.ndim != 2: + raise ValueError( + f"`context/context_mask` must be [B,L,D]/[B,L], got {tuple(context.shape)} and {tuple(context_mask.shape)}" + ) + context = context.to(device=self.device, dtype=self.torch_dtype, non_blocking=True) + context_mask = context_mask.to(device=self.device, dtype=torch.bool, non_blocking=True) + if self.proprio_encoder is not None: + if proprio is None: + raise ValueError("`sample['proprio']` is required when `proprio_dim` is enabled.") + if proprio.ndim != 3: + raise ValueError( + f"`sample['proprio']` must be 3D [B, T, d], got shape {tuple(proprio.shape)}" + ) + if proprio.shape[2] != self.proprio_dim: + raise ValueError( + f"`sample['proprio']` last dim must be {self.proprio_dim}, got {proprio.shape[2]}" + ) + proprio = proprio[:, 0, :] # [B, D] + context, context_mask = self._append_proprio_to_context( + context=context, + context_mask=context_mask, + proprio=proprio.to(device=self.device, dtype=self.torch_dtype), + ) + action = action.to(device=self.device, dtype=self.torch_dtype, non_blocking=True) + + if action_is_pad is not None: + action_is_pad = action_is_pad.to(device=self.device, dtype=torch.bool, non_blocking=True) + if image_is_pad is not None: + image_is_pad = image_is_pad.to(device=self.device, dtype=torch.bool, non_blocking=True) + + return { + "context": context, + "context_mask": context_mask, + "input_latents": input_latents, + "first_frame_latents": first_frame_latents, + "fuse_vae_embedding_in_latents": fuse_flag, + "action": action, + "action_is_pad": action_is_pad, + "image_is_pad": image_is_pad, + } + + @torch.no_grad() + def _build_mot_attention_mask( + self, + video_seq_len: int, + action_seq_len: int, + video_tokens_per_frame: int, + device: torch.device, + ) -> torch.Tensor: + total_seq_len = video_seq_len + action_seq_len + mask = torch.zeros((total_seq_len, total_seq_len), dtype=torch.bool, device=device) + + # video -> video + mask[:video_seq_len, :video_seq_len] = self.video_expert.build_video_to_video_mask( + video_seq_len=video_seq_len, + video_tokens_per_frame=video_tokens_per_frame, + device=device, + ) + # action -> action + mask[video_seq_len:, video_seq_len:] = True + # action -> first-frame video only + first_frame_tokens = min(video_tokens_per_frame, video_seq_len) + mask[video_seq_len:, :first_frame_tokens] = True + return mask + + def _compute_video_loss_per_sample( + self, + pred_video: torch.Tensor, + target_video: torch.Tensor, + image_is_pad: torch.Tensor | None, + include_initial_video_step: bool, + ) -> torch.Tensor: + video_loss_token = functional.mse_loss( + pred_video.float(), target_video.float(), reduction="none" + ).mean(dim=(1, 3, 4)) + if image_is_pad is None: + return video_loss_token.mean(dim=1) + + temporal_factor = int(self.vae.temporal_downsample_factor) + if temporal_factor <= 0: + raise ValueError(f"`vae.temporal_downsample_factor` must be positive, got {temporal_factor}.") + if image_is_pad.shape[1] < 1: + raise ValueError("`image_is_pad` must contain at least one frame.") + if (image_is_pad.shape[1] - 1) % temporal_factor != 0: + raise ValueError( + "Cannot align `image_is_pad` with video latent steps: " + f"num_frames={image_is_pad.shape[1]}, temporal_downsample_factor={temporal_factor}." + ) + + tail_is_pad = image_is_pad[:, 1:] + latent_tail_is_pad = tail_is_pad.view(image_is_pad.shape[0], -1, temporal_factor).all(dim=2) + if include_initial_video_step: + video_is_pad = torch.cat([image_is_pad[:, :1], latent_tail_is_pad], dim=1) + else: + video_is_pad = latent_tail_is_pad + + if video_is_pad.shape[1] != video_loss_token.shape[1]: + raise ValueError( + "Video-loss mask shape mismatch: " + f"mask steps={video_is_pad.shape[1]}, loss steps={video_loss_token.shape[1]}." + ) + + valid = (~video_is_pad).to(device=video_loss_token.device, dtype=video_loss_token.dtype) + valid_sum = valid.sum(dim=1).clamp(min=1.0) + return (video_loss_token * valid).sum(dim=1) / valid_sum + + def _sample_training_targets(self, inputs: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + input_latents = inputs["input_latents"] + batch_size = input_latents.shape[0] + action = inputs["action"] + noise_video = torch.randn_like(input_latents) + timestep_video = self.train_video_scheduler.sample_training_t( + batch_size=batch_size, + device=self.device, + dtype=input_latents.dtype, + ) + latents = self.train_video_scheduler.add_noise(input_latents, noise_video, timestep_video) + target_video = self.train_video_scheduler.training_target(input_latents, noise_video, timestep_video) + + if inputs["first_frame_latents"] is not None: + latents[:, :, 0:1] = inputs["first_frame_latents"] + noise_action = torch.randn_like(action) + timestep_action = self.train_action_scheduler.sample_training_t( + batch_size=batch_size, + device=self.device, + dtype=action.dtype, + ) + noisy_action = self.train_action_scheduler.add_noise(action, noise_action, timestep_action) + target_action = self.train_action_scheduler.training_target(action, noise_action, timestep_action) + return { + "latents": latents, + "target_video": target_video, + "noisy_action": noisy_action, + "target_action": target_action, + "timestep_video": timestep_video, + "timestep_action": timestep_action, + } + + def _run_training_mot(self, inputs: dict[str, torch.Tensor], targets: dict[str, torch.Tensor]): + video_pre = self.video_expert.pre_dit( + x=targets["latents"], + timestep=targets["timestep_video"], + context=inputs["context"], + context_mask=inputs["context_mask"], + action=inputs["action"], + fuse_vae_embedding_in_latents=inputs["fuse_vae_embedding_in_latents"], + ) + action_pre = self.action_expert.pre_dit( + action_tokens=targets["noisy_action"], + timestep=targets["timestep_action"], + context=inputs["context"], + context_mask=inputs["context_mask"], + ) + video_tokens = video_pre["tokens"] + action_tokens = action_pre["tokens"] + attention_mask = self._build_mot_attention_mask( + video_seq_len=video_tokens.shape[1], + action_seq_len=action_tokens.shape[1], + video_tokens_per_frame=int(video_pre["meta"]["tokens_per_frame"]), + device=video_tokens.device, + ) + tokens_out = self.mot( + embeds_all={ + "video": video_tokens, + "action": action_tokens, + }, + attention_mask=attention_mask, + freqs_all={ + "video": video_pre["freqs"], + "action": action_pre["freqs"], + }, + context_all={ + "video": { + "context": video_pre["context"], + "mask": video_pre["context_mask"], + }, + "action": { + "context": action_pre["context"], + "mask": action_pre["context_mask"], + }, + }, + t_mod_all={ + "video": video_pre["t_mod"], + "action": action_pre["t_mod"], + }, + ) + pred_video = self.video_expert.post_dit(tokens_out["video"], video_pre) + 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): + include_initial_video_step = inputs["first_frame_latents"] is None + if inputs["first_frame_latents"] is not None: + pred_video = pred_video[:, :, 1:] + target_video = target_video[:, :, 1:] + loss_video_per_sample = self._compute_video_loss_per_sample( + pred_video=pred_video, + target_video=target_video, + image_is_pad=inputs["image_is_pad"], + include_initial_video_step=include_initial_video_step, + ) + video_weight = self.train_video_scheduler.training_weight(timestep_video).to( + loss_video_per_sample.device, + dtype=loss_video_per_sample.dtype, + ) + return (loss_video_per_sample * video_weight).mean() + + def _compute_training_action_loss(self, inputs, pred_action, target_action, timestep_action): + action_loss_token = functional.mse_loss( + pred_action.float(), target_action.float(), reduction="none" + ).mean(dim=2) + if inputs["action_is_pad"] is not None: + valid = (~inputs["action_is_pad"]).to( + device=action_loss_token.device, + dtype=action_loss_token.dtype, + ) + valid_sum = valid.sum(dim=1).clamp(min=1.0) + action_loss_per_sample = (action_loss_token * valid).sum(dim=1) / valid_sum + else: + action_loss_per_sample = action_loss_token.mean(dim=1) + action_weight = self.train_action_scheduler.training_weight(timestep_action).to( + action_loss_per_sample.device, + dtype=action_loss_per_sample.dtype, + ) + return (action_loss_per_sample * action_weight).mean() + + def training_loss(self, sample, tiled: bool = False): + 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) + loss_video = self._compute_training_video_loss( + inputs=inputs, + pred_video=pred_video, + target_video=targets["target_video"], + timestep_video=targets["timestep_video"], + ) + loss_action = self._compute_training_action_loss( + inputs=inputs, + pred_action=pred_action, + target_action=targets["target_action"], + timestep_action=targets["timestep_action"], + ) + 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()), + } + return loss_total, loss_dict + + @torch.no_grad() + def _predict_joint_noise( + self, + latents_video: torch.Tensor, + latents_action: torch.Tensor, + timestep_video: torch.Tensor, + timestep_action: torch.Tensor, + context: torch.Tensor, + context_mask: torch.Tensor, + fuse_vae_embedding_in_latents: bool, + gt_action: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + video_pre = self.video_expert.pre_dit( + x=latents_video, + timestep=timestep_video, + context=context, + context_mask=context_mask, + action=gt_action, + fuse_vae_embedding_in_latents=fuse_vae_embedding_in_latents, + ) + action_pre = self.action_expert.pre_dit( + action_tokens=latents_action, + timestep=timestep_action, + context=context, + context_mask=context_mask, + ) + + attention_mask = self._build_mot_attention_mask( + video_seq_len=video_pre["tokens"].shape[1], + action_seq_len=action_pre["tokens"].shape[1], + video_tokens_per_frame=int(video_pre["meta"]["tokens_per_frame"]), + device=video_pre["tokens"].device, + ) + + tokens_out = self.mot( + embeds_all={ + "video": video_pre["tokens"], + "action": action_pre["tokens"], + }, + attention_mask=attention_mask, + freqs_all={ + "video": video_pre["freqs"], + "action": action_pre["freqs"], + }, + context_all={ + "video": { + "context": video_pre["context"], + "mask": video_pre["context_mask"], + }, + "action": { + "context": action_pre["context"], + "mask": action_pre["context_mask"], + }, + }, + t_mod_all={ + "video": video_pre["t_mod"], + "action": action_pre["t_mod"], + }, + ) + + pred_video = self.video_expert.post_dit(tokens_out["video"], video_pre) + pred_action = self.action_expert.post_dit(tokens_out["action"], action_pre) + return pred_video, pred_action + + @torch.no_grad() + def _predict_action_noise( + self, + first_frame_latents: torch.Tensor, + latents_action: torch.Tensor, + timestep_action: torch.Tensor, + context: torch.Tensor, + context_mask: torch.Tensor, + fuse_vae_embedding_in_latents: bool, + ) -> torch.Tensor: + timestep_video = torch.zeros_like( + timestep_action, dtype=first_frame_latents.dtype, device=self.device + ) + video_pre = self.video_expert.pre_dit( + x=first_frame_latents, + timestep=timestep_video, + context=context, + context_mask=context_mask, + action=None, + fuse_vae_embedding_in_latents=fuse_vae_embedding_in_latents, + ) + action_pre = self.action_expert.pre_dit( + action_tokens=latents_action, + timestep=timestep_action, + context=context, + context_mask=context_mask, + ) + + attention_mask = self._build_mot_attention_mask( + video_seq_len=video_pre["tokens"].shape[1], + action_seq_len=action_pre["tokens"].shape[1], + video_tokens_per_frame=int(video_pre["meta"]["tokens_per_frame"]), + device=video_pre["tokens"].device, + ) + tokens_out = self.mot( + embeds_all={ + "video": video_pre["tokens"], + "action": action_pre["tokens"], + }, + attention_mask=attention_mask, + freqs_all={ + "video": video_pre["freqs"], + "action": action_pre["freqs"], + }, + context_all={ + "video": { + "context": video_pre["context"], + "mask": video_pre["context_mask"], + }, + "action": { + "context": action_pre["context"], + "mask": action_pre["context_mask"], + }, + }, + t_mod_all={ + "video": video_pre["t_mod"], + "action": action_pre["t_mod"], + }, + ) + pred_action = self.action_expert.post_dit(tokens_out["action"], action_pre) + return pred_action + + @torch.no_grad() + def _predict_action_noise_with_cache( + self, + latents_action: torch.Tensor, + timestep_action: torch.Tensor, + context: torch.Tensor, + context_mask: torch.Tensor, + video_kv_cache: list[dict[str, torch.Tensor]], + attention_mask: torch.Tensor, + video_seq_len: int, + ) -> torch.Tensor: + action_pre = self.action_expert.pre_dit( + action_tokens=latents_action, + timestep=timestep_action, + context=context, + context_mask=context_mask, + ) + action_tokens = self.mot.forward_action_with_video_cache( + action_tokens=action_pre["tokens"], + action_freqs=action_pre["freqs"], + action_t_mod=action_pre["t_mod"], + action_context_payload={ + "context": action_pre["context"], + "mask": action_pre["context_mask"], + }, + video_kv_cache=video_kv_cache, + attention_mask=attention_mask, + video_seq_len=video_seq_len, + ) + return self.action_expert.post_dit(action_tokens, action_pre) + + def _normalize_infer_input_image( + self, + input_image: torch.Tensor, + num_video_frames: int | None = None, + ) -> tuple[torch.Tensor, int, int]: + if input_image.ndim == 3: + input_image = input_image.unsqueeze(0) + if input_image.ndim != 4 or input_image.shape[0] != 1 or input_image.shape[1] != 3: + raise ValueError( + f"`input_image` must have shape [1,3,H,W] or [3,H,W], got {tuple(input_image.shape)}" + ) + _, _, height, width = input_image.shape + if height % 16 != 0 or width % 16 != 0: + raise ValueError( + f"`input_image` must be resized before infer, expected multiples of 16 but got HxW=({height},{width})" + ) + if num_video_frames is not None: + checked_h, checked_w, checked_t = self._check_resize_height_width(height, width, num_video_frames) + if (checked_h, checked_w) != (height, width): + raise ValueError( + f"`input_image` must be resized before infer, expected multiples of 16 but got HxW=({height},{width})" + ) + if checked_t != num_video_frames: + raise ValueError(f"`num_video_frames` must satisfy T % 4 == 1, got {num_video_frames}") + return input_image, height, width + + def _normalize_infer_proprio(self, proprio: torch.Tensor | None) -> torch.Tensor | None: + if proprio is None: + return None + if self.proprio_dim is None: + raise ValueError( + "`proprio` was provided but `proprio_dim=None` so `proprio_encoder` is disabled." + ) + if proprio.ndim == 1: + proprio = proprio.unsqueeze(0) + elif proprio.ndim == 2 and proprio.shape[0] == 1: + pass + else: + raise ValueError(f"`proprio` must be [D] or [1,D], got shape {tuple(proprio.shape)}") + if proprio.shape[1] != self.proprio_dim: + raise ValueError(f"`proprio` last dim must be {self.proprio_dim}, got {proprio.shape[1]}") + return proprio.to(device=self.device, dtype=self.torch_dtype) + + def _prepare_infer_context(self, prompt, context, context_mask, proprio): + use_prompt = prompt is not None + use_context = context is not None or context_mask is not None + if use_prompt and use_context: + raise ValueError("`prompt` and `context/context_mask` are mutually exclusive.") + if not use_prompt and not use_context: + raise ValueError("Either `prompt` or both `context/context_mask` must be provided.") + if use_prompt: + context, context_mask = self.encode_prompt(prompt) + else: + context, context_mask = self._normalize_context_tensors(context, context_mask) + if proprio is not None: + context, context_mask = self._append_proprio_to_context( + context=context, + context_mask=context_mask, + proprio=proprio, + ) + return context, context_mask + + def _normalize_context_tensors(self, context, context_mask): + if context is None or context_mask is None: + raise ValueError("`context` and `context_mask` must be both provided together.") + if context.ndim == 2: + context = context.unsqueeze(0) + if context_mask.ndim == 1: + context_mask = context_mask.unsqueeze(0) + if context.ndim != 3 or context_mask.ndim != 2: + raise ValueError( + f"`context/context_mask` must be [B,L,D]/[B,L], got {tuple(context.shape)} and {tuple(context_mask.shape)}" + ) + context = context.to(device=self.device, dtype=self.torch_dtype, non_blocking=True) + context_mask = context_mask.to(device=self.device, dtype=torch.bool, non_blocking=True) + return context, context_mask + + def _make_action_latents(self, action_horizon: int, seed: int | None, rand_device: str): + generator = None if seed is None else torch.Generator(device=rand_device).manual_seed(seed) + return torch.randn( + (1, action_horizon, self.action_expert.action_dim), + generator=generator, + device=rand_device, + dtype=torch.float32, + ).to(device=self.device, dtype=self.torch_dtype) + + def _make_video_latents(self, num_video_frames: int, height: int, width: int, seed, rand_device): + latent_t = (num_video_frames - 1) // self.vae.temporal_downsample_factor + 1 + latent_h = height // self.vae.upsampling_factor + latent_w = width // self.vae.upsampling_factor + generator = None if seed is None else torch.Generator(device=rand_device).manual_seed(seed) + return torch.randn( + (1, self.vae.z_dim, latent_t, latent_h, latent_w), + generator=generator, + device=rand_device, + dtype=torch.float32, + ).to(device=self.device, dtype=self.torch_dtype) + + @torch.no_grad() + def infer_joint( + self, + prompt: str | None, + input_image: torch.Tensor, + num_video_frames: int, + action_horizon: int, + action: torch.Tensor + | None = None, # NOTE: this is gt action for conditioning videos, not for action expert + proprio: torch.Tensor | None = None, + context: torch.Tensor | None = None, + context_mask: torch.Tensor | None = None, + negative_prompt: str | None = None, + text_cfg_scale: float = 1.0, + num_inference_steps: int = 20, + sigma_shift: float | None = None, + seed: int | None = None, + rand_device: str = "cpu", + tiled: bool = False, + test_action_with_infer_action: bool = True, + ) -> dict[str, Any]: + self.eval() + if test_action_with_infer_action: + if seed is None: + raise ValueError("`test_action_with_infer_action=True` requires non-null `seed`.") + action_only_out = self.infer_action( + prompt=prompt, + input_image=input_image.clone(), + action_horizon=action_horizon, + context=context.clone() if context is not None else None, + context_mask=context_mask.clone() if context_mask is not None else None, + num_inference_steps=num_inference_steps, + sigma_shift=sigma_shift, + seed=seed, + rand_device=rand_device, + tiled=tiled, + proprio=proprio.clone() if proprio is not None else None, + )["action"] + + input_image, height, width = self._normalize_infer_input_image(input_image, num_video_frames) + if action is not None: + if action.ndim == 2: + action = action.unsqueeze(0) + if action.ndim != 3 or action.shape[0] != 1 or action.shape[1] != action_horizon: + # NOTE: This enforces action condition to have the same shape as action horizon to predict, which may be unnecessary + raise ValueError( + f"`action` must have shape [1, T, a_dim] or [T, a_dim], got {tuple(action.shape)} with action_horizon={action_horizon}" + ) + action = action.to(device=self.device, dtype=self.torch_dtype) + proprio = self._normalize_infer_proprio(proprio) + latents_video = self._make_video_latents(num_video_frames, height, width, seed, rand_device) + latents_action = self._make_action_latents(action_horizon, seed, rand_device) + + input_image = input_image.to(device=self.device, dtype=self.torch_dtype) + first_frame_latents = self._encode_input_image_latents_tensor(input_image=input_image, tiled=tiled) + latents_video[:, :, 0:1] = first_frame_latents.clone() + fuse_flag = bool(getattr(self.video_expert, "fuse_vae_embedding_in_latents", False)) + context, context_mask = self._prepare_infer_context(prompt, context, context_mask, proprio) + + infer_timesteps_video, infer_deltas_video = self.infer_video_scheduler.build_inference_schedule( + num_inference_steps=num_inference_steps, + device=self.device, + dtype=latents_video.dtype, + shift_override=sigma_shift, + ) + infer_timesteps_action, infer_deltas_action = self.infer_action_scheduler.build_inference_schedule( + num_inference_steps=num_inference_steps, + device=self.device, + dtype=latents_action.dtype, + shift_override=sigma_shift, + ) + for step_t_video, step_delta_video, step_t_action, step_delta_action in zip( + infer_timesteps_video, + infer_deltas_video, + infer_timesteps_action, + infer_deltas_action, + strict=True, + ): + timestep_video = step_t_video.unsqueeze(0).to(dtype=latents_video.dtype, device=self.device) + timestep_action = step_t_action.unsqueeze(0).to(dtype=latents_action.dtype, device=self.device) + + pred_video, pred_action = self._predict_joint_noise( + latents_video=latents_video, + latents_action=latents_action, + timestep_video=timestep_video, + timestep_action=timestep_action, + context=context, + context_mask=context_mask, + fuse_vae_embedding_in_latents=fuse_flag, + gt_action=action, + ) + + latents_video = self.infer_video_scheduler.step(pred_video, step_delta_video, latents_video) + latents_action = self.infer_action_scheduler.step(pred_action, step_delta_action, latents_action) + latents_video[:, :, 0:1] = first_frame_latents.clone() + + action_out = latents_action[0].detach().to(device="cpu", dtype=torch.float32) + if test_action_with_infer_action and not torch.allclose( + action_out, action_only_out, atol=1e-2, rtol=1e-2 + ): + max_abs_diff = (action_out - action_only_out).abs().max().item() + logger.warning( + f"Action from infer_joint and infer_action differ with max abs diff {max_abs_diff:.6f}. " + ) + + return { + "video": self._decode_latents(latents_video, tiled=tiled), + "action": action_out, + } + + @torch.no_grad() + def infer_action( + self, + prompt: str | None, + input_image: torch.Tensor, + action_horizon: int, + proprio: torch.Tensor | None = None, + context: torch.Tensor | None = None, + context_mask: torch.Tensor | None = None, + negative_prompt: str | None = None, + text_cfg_scale: float = 1.0, + num_inference_steps: int = 20, + sigma_shift: float | None = None, + seed: int | None = None, + rand_device: str = "cpu", + tiled: bool = False, + ) -> dict[str, Any]: + self.eval() + if str(getattr(self.video_expert, "video_attention_mask_mode", "")) != "first_frame_causal": + raise ValueError("`infer_action` requires `video_attention_mask_mode='first_frame_causal'`.") + + input_image, _, _ = self._normalize_infer_input_image(input_image) + proprio = self._normalize_infer_proprio(proprio) + latents_action = self._make_action_latents(action_horizon, seed, rand_device) + + input_image = input_image.to(device=self.device, dtype=self.torch_dtype) + first_frame_latents = self._encode_input_image_latents_tensor(input_image=input_image, tiled=tiled) + fuse_flag = bool(getattr(self.video_expert, "fuse_vae_embedding_in_latents", False)) + + context, context_mask = self._prepare_infer_context(prompt, context, context_mask, proprio) + + timestep_video = torch.zeros( + (first_frame_latents.shape[0],), + dtype=first_frame_latents.dtype, + device=self.device, + ) + video_pre = self.video_expert.pre_dit( + x=first_frame_latents, + timestep=timestep_video, + context=context, + context_mask=context_mask, + action=None, + fuse_vae_embedding_in_latents=fuse_flag, + ) + video_seq_len = int(video_pre["tokens"].shape[1]) + attention_mask = self._build_mot_attention_mask( + video_seq_len=video_seq_len, + action_seq_len=latents_action.shape[1], + video_tokens_per_frame=int(video_pre["meta"]["tokens_per_frame"]), + device=video_pre["tokens"].device, + ) + video_kv_cache = self.mot.prefill_video_cache( + video_tokens=video_pre["tokens"], + video_freqs=video_pre["freqs"], + video_t_mod=video_pre["t_mod"], + video_context_payload={ + "context": video_pre["context"], + "mask": video_pre["context_mask"], + }, + video_attention_mask=attention_mask[:video_seq_len, :video_seq_len], + ) + + infer_timesteps_action, infer_deltas_action = self.infer_action_scheduler.build_inference_schedule( + num_inference_steps=num_inference_steps, + device=self.device, + dtype=latents_action.dtype, + shift_override=sigma_shift, + ) + for step_t_action, step_delta_action in zip(infer_timesteps_action, infer_deltas_action, strict=True): + timestep_action = step_t_action.unsqueeze(0).to(dtype=latents_action.dtype, device=self.device) + + pred_action = self._predict_action_noise_with_cache( + latents_action=latents_action, + timestep_action=timestep_action, + context=context, + context_mask=context_mask, + video_kv_cache=video_kv_cache, + attention_mask=attention_mask, + video_seq_len=video_seq_len, + ) + + latents_action = self.infer_action_scheduler.step(pred_action, step_delta_action, latents_action) + + return { + "action": latents_action[0].detach().to(device="cpu", dtype=torch.float32), + } + + @torch.no_grad() + def infer( + self, + prompt: str | None, + input_image: torch.Tensor, + num_frames: int, + action: torch.Tensor | None = None, + action_horizon: int | None = None, + proprio: torch.Tensor | None = None, + context: torch.Tensor | None = None, + context_mask: torch.Tensor | None = None, + negative_prompt: str | None = None, + text_cfg_scale: float = 5.0, + action_cfg_scale: float = 1.0, + num_inference_steps: int = 20, + sigma_shift: float | None = None, + seed: int | None = None, + rand_device: str = "cpu", + tiled: bool = False, + ): + return self.infer_joint( + prompt=prompt, + input_image=input_image, + num_video_frames=num_frames, + action_horizon=action_horizon, + action=action, + proprio=proprio, + context=context, + context_mask=context_mask, + negative_prompt=negative_prompt, + text_cfg_scale=text_cfg_scale, + num_inference_steps=num_inference_steps, + sigma_shift=sigma_shift, + seed=seed, + rand_device=rand_device, + tiled=tiled, + ) + + def forward(self, *args, **kwargs): + return self.training_loss(*args, **kwargs) diff --git a/src/lerobot/policies/fastwam/wan/video_dit.py b/src/lerobot/policies/fastwam/wan/video_dit.py new file mode 100644 index 000000000..a98f06cde --- /dev/null +++ b/src/lerobot/policies/fastwam/wan/video_dit.py @@ -0,0 +1,800 @@ +# 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. + +import logging +from typing import Any + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as functional +from einops import rearrange + +from .model import ( + WanAttentionBlock, + WanLayerNorm, + WanModel, + WanRMSNorm, + rope_apply, + rope_params, + sinusoidal_embedding_1d, +) + +logger = logging.getLogger(__name__) + + +def get_sampling_sigmas(sampling_steps, shift): + # Vendored from Wan2.2 (formerly wan/utils/fm_solvers.py); computes the + # noise-level (sigma) schedule for Wan-compatible flow-matching inference. + sigma = np.linspace(1, 0, sampling_steps + 1)[:sampling_steps] + sigma = shift * sigma / (1 + (shift - 1) * sigma) + return sigma + + +def create_custom_forward(module): + def custom_forward(*inputs, **kwargs): + return module(*inputs, **kwargs) + + return custom_forward + + +def gradient_checkpoint_forward( + model, + use_gradient_checkpointing, + *args, + **kwargs, +): + if use_gradient_checkpointing: + model_output = torch.utils.checkpoint.checkpoint( + create_custom_forward(model), + *args, + **kwargs, + use_reentrant=False, + ) + else: + model_output = model(*args, **kwargs) + return model_output + + +def fastwam_masked_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + num_heads: int, + ctx_mask: torch.Tensor | None = None, + fp32_attention: bool = True, +) -> torch.Tensor: + """FastWAM masked attention wrapper for MoT masks and CPU test coverage. + + The official Wan attention implementation is still used as the source of + the projection/norm modules. This wrapper only replaces the final attention + kernel because FastWAM needs explicit boolean masks for video/action MoT + routing, while the upstream FlashAttention path accepts sequence lengths + but not arbitrary [query, key] masks. + """ + + q = rearrange(q, "b s (n d) -> b n s d", n=num_heads) + k = rearrange(k, "b s (n d) -> b n s d", n=num_heads) + v = rearrange(v, "b s (n d) -> b n s d", n=num_heads) + if fp32_attention: + q = q.float() + k = k.float() + v = v.float() + else: + q = q.to(dtype=v.dtype) + k = k.to(dtype=v.dtype) + x = functional.scaled_dot_product_attention(q, k, v, attn_mask=ctx_mask) + return rearrange(x, "b n s d -> b s (n d)", n=num_heads) + + +def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor): + return x * (1 + scale) + shift + + +class WanContinuousFlowMatchScheduler: + """Continuous-time Flow-Matching scheduler with shift-based Wan sampling.""" + + def __init__(self, num_train_timesteps: int = 1000, shift: float = 5.0, eps: float = 1e-10): + if num_train_timesteps <= 0: + raise ValueError(f"`num_train_timesteps` must be positive, got {num_train_timesteps}") + if shift <= 0: + raise ValueError(f"`shift` must be positive, got {shift}") + self.num_train_timesteps = int(num_train_timesteps) + self.shift = float(shift) + self.eps = float(eps) + self._y_min, self._weight_norm_const = self._precompute_training_weight_stats() + + @staticmethod + def _phi(u: torch.Tensor, shift: float) -> torch.Tensor: + return shift * u / (1.0 + (shift - 1.0) * u) + + def _precompute_training_weight_stats(self) -> tuple[float, float]: + steps = self.num_train_timesteps + u_grid = torch.linspace(1.0, 0.0, steps + 1, dtype=torch.float64)[:-1] + t_grid = self._phi(u_grid, self.shift) * float(steps) + y_grid = torch.exp(-2.0 * ((t_grid - (steps / 2.0)) / steps) ** 2) + y_min = float(y_grid.min().item()) + y_shifted_grid = y_grid - y_min + norm_const = float(y_shifted_grid.mean().item()) + return y_min, norm_const + + def sample_training_t(self, batch_size: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor: + if batch_size <= 0: + raise ValueError(f"`batch_size` must be positive, got {batch_size}") + u = torch.rand((batch_size,), device=device, dtype=torch.float32) + sigma = self._phi(u, self.shift) + timestep = sigma * float(self.num_train_timesteps) + return timestep.to(dtype=dtype) + + def training_weight(self, timestep: torch.Tensor) -> torch.Tensor: + t = timestep.to(dtype=torch.float32) + steps = float(self.num_train_timesteps) + y = torch.exp(-2.0 * ((t - (steps / 2.0)) / steps) ** 2) + y_shifted = y - self._y_min + weight = y_shifted / (self._weight_norm_const + self.eps) + if weight.numel() == 1: + return weight.reshape(()) + return weight + + def add_noise( + self, original_samples: torch.Tensor, noise: torch.Tensor, timestep: torch.Tensor + ) -> torch.Tensor: + sigma = (timestep / float(self.num_train_timesteps)).to( + original_samples.device, dtype=original_samples.dtype + ) + if sigma.ndim == 0: + return (1 - sigma) * original_samples + sigma * noise + sigma = sigma.view(-1, *([1] * (original_samples.ndim - 1))) + return (1 - sigma) * original_samples + sigma * noise + + @staticmethod + def training_target(sample: torch.Tensor, noise: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor: + del timestep + return noise - sample + + def build_inference_schedule( + self, + num_inference_steps: int, + device: torch.device, + dtype: torch.dtype, + shift_override: float | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if num_inference_steps <= 0: + raise ValueError(f"`num_inference_steps` must be positive, got {num_inference_steps}") + shift = self.shift if shift_override is None else float(shift_override) + if shift <= 0: + raise ValueError(f"`shift` must be positive, got {shift}") + + sigma_steps = torch.as_tensor( + get_sampling_sigmas(num_inference_steps, shift), + device=device, + dtype=torch.float32, + ) + timesteps = sigma_steps * float(self.num_train_timesteps) + sigma_next = torch.cat([sigma_steps[1:], sigma_steps.new_zeros(1)]) + deltas = sigma_next - sigma_steps + return timesteps.to(dtype=dtype), deltas.to(dtype=dtype) + + @staticmethod + def step(model_output: torch.Tensor, delta: torch.Tensor, sample: torch.Tensor) -> torch.Tensor: + delta = delta.to(sample.device, dtype=sample.dtype) + if delta.ndim == 0: + return sample + model_output * delta + delta = delta.view(-1, *([1] * (sample.ndim - 1))) + return sample + model_output * delta + + +def precompute_freqs_cis(dim: int, end: int = 1024, theta: float = 10000.0): + return rope_params(end, dim, theta) + + +def apply_dense_rope(x: torch.Tensor, freqs: torch.Tensor, num_heads: int) -> torch.Tensor: + x = rearrange(x, "b s (n d) -> b s n d", n=num_heads) + x_out = torch.view_as_complex(x.to(torch.float32).reshape(x.shape[0], x.shape[1], x.shape[2], -1, 2)) + freqs = freqs.to(torch.complex64) if freqs.device.type == "npu" else freqs + x_out = torch.view_as_real(x_out * freqs).flatten(2) + return x_out.to(x.dtype) + + +def _linear_input(linear: nn.Linear, x: torch.Tensor) -> torch.Tensor: + return x.to(dtype=linear.weight.dtype) + + +def _wan_layer_norm(norm: nn.Module, x: torch.Tensor) -> torch.Tensor: + if isinstance(norm, WanLayerNorm) and norm.weight is not None: + weight = norm.weight.float() + bias = norm.bias.float() if norm.bias is not None else None + return functional.layer_norm(x.float(), norm.normalized_shape, weight, bias, norm.eps).to( + dtype=x.dtype + ) + return norm(x) + + +def create_group_causal_attn_mask( + num_temporal_groups: int, num_query_per_group: int, num_key_per_group: int, mode: str = "causal" +) -> torch.Tensor: + if mode not in ["causal", "group_diagonal"]: + raise ValueError(f"`mode` must be 'causal' or 'group_diagonal', got {mode}.") + if num_temporal_groups <= 0: + raise ValueError(f"`num_temporal_groups` must be positive, got {num_temporal_groups}.") + if num_query_per_group <= 0: + raise ValueError(f"`num_query_per_group` must be positive, got {num_query_per_group}.") + if num_key_per_group <= 0: + raise ValueError(f"`num_key_per_group` must be positive, got {num_key_per_group}.") + + total_num_query_tokens = num_temporal_groups * num_query_per_group + total_num_key_tokens = num_temporal_groups * num_key_per_group + query_time_indices = torch.arange(num_temporal_groups).repeat_interleave(num_query_per_group).unsqueeze(1) + key_time_indices = torch.arange(num_temporal_groups).repeat_interleave(num_key_per_group).unsqueeze(0) + + if mode == "causal": + attn_mask = query_time_indices >= key_time_indices + else: + attn_mask = query_time_indices == key_time_indices + + if attn_mask.shape != (total_num_query_tokens, total_num_key_tokens): + raise RuntimeError("Attention mask shape mismatch.") + return attn_mask + + +class FastWAMAttentionBlock(WanAttentionBlock): + """Wan attention block with FastWAM's arbitrary boolean mask support.""" + + def __init__( + self, + hidden_dim: int, + attn_head_dim: int, + num_heads: int, + ffn_dim: int, + eps: float = 1e-6, + fp32_attention: bool = True, + ): + attention_dim = attn_head_dim * num_heads + if hidden_dim == attention_dim: + super().__init__( + dim=hidden_dim, + ffn_dim=ffn_dim, + num_heads=num_heads, + qk_norm=True, + cross_attn_norm=True, + eps=eps, + ) + else: + nn.Module.__init__(self) + self.dim = hidden_dim + self.ffn_dim = ffn_dim + self.num_heads = num_heads + self.qk_norm = True + self.cross_attn_norm = True + self.eps = eps + self.norm1 = WanLayerNorm(hidden_dim, eps) + self.self_attn = _FastWAMProjectedAttention(hidden_dim, attention_dim, num_heads, eps) + self.norm3 = WanLayerNorm(hidden_dim, eps, elementwise_affine=True) + self.cross_attn = _FastWAMProjectedAttention(hidden_dim, attention_dim, num_heads, eps) + self.norm2 = WanLayerNorm(hidden_dim, eps) + self.ffn = nn.Sequential( + nn.Linear(hidden_dim, ffn_dim), + nn.GELU(approximate="tanh"), + nn.Linear(ffn_dim, hidden_dim), + ) + self.modulation = nn.Parameter(torch.randn(1, 6, hidden_dim) / hidden_dim**0.5) + self.attn_head_dim = attn_head_dim + self.fp32_attention = bool(fp32_attention) + + @staticmethod + def split_modulation(block, t_mod: torch.Tensor): + has_seq = len(t_mod.shape) == 4 + chunk_dim = 2 if has_seq else 1 + + base_mod = block.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (base_mod + t_mod).chunk( + 6, dim=chunk_dim + ) + if has_seq: + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + shift_msa.squeeze(2), + scale_msa.squeeze(2), + gate_msa.squeeze(2), + shift_mlp.squeeze(2), + scale_mlp.squeeze(2), + gate_mlp.squeeze(2), + ) + return shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp + + def project_self_attention( + self, x: torch.Tensor, freqs: torch.Tensor | dict[str, torch.Tensor] + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + q = self.self_attn.norm_q(self.self_attn.q(x)) + k = self.self_attn.norm_k(self.self_attn.k(x)) + v = self.self_attn.v(x) + if isinstance(freqs, dict): + b, s = x.shape[:2] + q = rope_apply( + q.view(b, s, self.num_heads, self.attn_head_dim), + freqs["grid_sizes"], + freqs["freqs"], + ).flatten(2) + k = rope_apply( + k.view(b, s, self.num_heads, self.attn_head_dim), + freqs["grid_sizes"], + freqs["freqs"], + ).flatten(2) + else: + q = apply_dense_rope(q, freqs, self.num_heads) + k = apply_dense_rope(k, freqs, self.num_heads) + return q, k, v + + def apply_cross_attention( + self, x: torch.Tensor, context: torch.Tensor, context_mask: torch.Tensor | None = None + ) -> torch.Tensor: + if context_mask is not None and context_mask.dim() == 3: + context_mask = context_mask.unsqueeze(1) + attn = self.cross_attn + b, n, d = x.size(0), attn.num_heads, attn.head_dim + q = attn.norm_q(attn.q(x)).view(b, -1, n * d) + k = attn.norm_k(attn.k(context)).view(b, -1, n * d) + v = attn.v(context).view(b, -1, n * d) + x = fastwam_masked_attention( + q=q, + k=k, + v=v, + num_heads=n, + ctx_mask=context_mask, + fp32_attention=self.fp32_attention, + ) + return attn.o(_linear_input(attn.o, x)) + + def project_self_attention_output(self, x: torch.Tensor) -> torch.Tensor: + return self.self_attn.o(_linear_input(self.self_attn.o, x)) + + def apply_norm1(self, x: torch.Tensor) -> torch.Tensor: + return _wan_layer_norm(self.norm1, x) + + def apply_norm2(self, x: torch.Tensor) -> torch.Tensor: + return _wan_layer_norm(self.norm2, x) + + def apply_norm3(self, x: torch.Tensor) -> torch.Tensor: + return _wan_layer_norm(self.norm3, x) + + def forward( + self, + x: torch.Tensor, + context: torch.Tensor, + t_mod: torch.Tensor, + freqs: torch.Tensor, + context_mask: torch.Tensor | None = None, + self_attn_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.split_modulation(self, t_mod) + residual_x = x + attn_input = modulate(self.apply_norm1(x), shift_msa, scale_msa) + q, k, v = self.project_self_attention(attn_input, freqs) + y = fastwam_masked_attention( + q=q, + k=k, + v=v, + num_heads=self.num_heads, + ctx_mask=self_attn_mask, + fp32_attention=self.fp32_attention, + ) + x = residual_x + gate_msa * self.project_self_attention_output(y) + x = x + self.apply_cross_attention(self.apply_norm3(x), context, context_mask=context_mask) + mlp_input = modulate(self.apply_norm2(x), shift_mlp, scale_mlp) + return x + gate_mlp * self.ffn(mlp_input) + + +class _FastWAMProjectedAttention(nn.Module): + def __init__(self, hidden_dim: int, attention_dim: int, num_heads: int, eps: float): + super().__init__() + self.dim = hidden_dim + self.num_heads = num_heads + self.head_dim = attention_dim // num_heads + self.q = nn.Linear(hidden_dim, attention_dim) + self.k = nn.Linear(hidden_dim, attention_dim) + self.v = nn.Linear(hidden_dim, attention_dim) + self.o = nn.Linear(attention_dim, hidden_dim) + self.norm_q = WanRMSNorm(attention_dim, eps=eps) + self.norm_k = WanRMSNorm(attention_dim, eps=eps) + + +class WanVideoDiT(WanModel): + def __init__( + self, + hidden_dim: int, + in_dim: int, + ffn_dim: int, + out_dim: int, + text_dim: int, + freq_dim: int, + eps: float, + patch_size: tuple[int, int, int], + num_heads: int, + attn_head_dim: int, + num_layers: int, + has_image_input: bool = False, + has_image_pos_emb: bool = False, + has_ref_conv: bool = False, + add_control_adapter: bool = False, + in_dim_control_adapter: int = 24, + seperated_timestep: bool = False, + require_vae_embedding: bool = False, + require_clip_embedding: bool = False, + fuse_vae_embedding_in_latents: bool = True, + action_conditioned: bool = False, + action_dim: int = 7, + action_group_causal_mask_mode="causal", + video_attention_mask_mode: str = "bidirectional", + use_gradient_checkpointing: bool = False, + fp32_attention: bool = True, + ): + del in_dim_control_adapter + if has_image_input: + raise ValueError("FastWAM currently expects Wan2.2 TI2V latents with fused image conditioning.") + if has_image_pos_emb: + raise ValueError("FastWAM does not support extra image positional embeddings in WanVideoDiT.") + if has_ref_conv: + raise ValueError("FastWAM does not support reference convolutions in WanVideoDiT.") + if add_control_adapter: + raise ValueError("FastWAM does not support control adapters in WanVideoDiT.") + if require_clip_embedding: + raise ValueError("FastWAM does not support CLIP embedding conditioning in WanVideoDiT.") + if require_vae_embedding or not fuse_vae_embedding_in_latents: + raise ValueError("FastWAM expects VAE conditioning to be fused in latents.") + if attn_head_dim != hidden_dim // num_heads: + raise ValueError( + "`attn_head_dim` must match the upstream Wan head dimension `hidden_dim // num_heads`; " + f"got {attn_head_dim} vs {hidden_dim // num_heads}." + ) + + super().__init__( + model_type="ti2v", + patch_size=patch_size, + text_len=512, + in_dim=in_dim, + dim=hidden_dim, + ffn_dim=ffn_dim, + freq_dim=freq_dim, + text_dim=text_dim, + out_dim=out_dim, + num_heads=num_heads, + num_layers=num_layers, + qk_norm=True, + cross_attn_norm=True, + eps=eps, + ) + self.blocks = torch.nn.ModuleList( + [ + FastWAMAttentionBlock( + hidden_dim=hidden_dim, + attn_head_dim=attn_head_dim, + num_heads=num_heads, + ffn_dim=ffn_dim, + eps=eps, + fp32_attention=fp32_attention, + ) + for _ in range(num_layers) + ] + ) + self.init_weights() + + self.hidden_dim = hidden_dim + self.attn_head_dim = attn_head_dim + self.seperated_timestep = seperated_timestep + self.fuse_vae_embedding_in_latents = fuse_vae_embedding_in_latents + self.video_attention_mask_mode = str(video_attention_mask_mode) + self.action_conditioned = action_conditioned + self.action_dim = action_dim + self.fp32_attention = bool(fp32_attention) + + if self.action_conditioned: + self.action_embedding = torch.nn.Linear(action_dim, hidden_dim) + self.action_group_causal_mask_mode = action_group_causal_mask_mode + + self.use_gradient_checkpointing = use_gradient_checkpointing + if self.use_gradient_checkpointing: + logger.info( + "Using gradient checkpointing for DiT blocks. This will save memory but use more computation." + ) + + def patchify(self, x: torch.Tensor): + return self.patch_embedding(x) + + def _validate_forward_inputs( + self, + x: torch.Tensor, + timestep: torch.Tensor, + context: torch.Tensor, + context_mask: torch.Tensor | None, + action: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if x.ndim != 5: + raise ValueError(f"`latents` must be 5D [B, C, T, H, W], got shape {tuple(x.shape)}") + num_latent_frames = x.shape[2] + if context.ndim != 3: + raise ValueError(f"`context` must be 3D [B, L, D], got shape {tuple(context.shape)}") + if timestep.ndim != 1: + raise ValueError(f"`timestep` must be 1D [B] or [1], got shape {tuple(timestep.shape)}") + if self.action_conditioned: + allow_text_only_single_frame = num_latent_frames == 1 and action is None + if not allow_text_only_single_frame: + if action is None: + raise ValueError("Action input is required for action-conditioned model.") + if action.ndim != 3: + raise ValueError( + f"`action` must be 3D [B, action_horizon, action_dim], got shape {tuple(action.shape)}" + ) + if action.shape[2] != self.action_dim: + raise ValueError( + f"`action` last dimension must be {self.action_dim}, got {action.shape[2]}" + ) + if num_latent_frames <= 1: + raise ValueError( + f"video length must be > 1 for action-conditioned model, got {num_latent_frames}" + ) + if action.shape[1] % (num_latent_frames - 1) != 0: + raise ValueError( + "action horizon must be divisible by (num_latent_frames - 1), " + f"got action_horizon={action.shape[1]}" + ) + if context_mask is None: + context_mask = torch.ones( + (context.shape[0], context.shape[1]), dtype=torch.bool, device=context.device + ) + else: + if context_mask.ndim != 2: + raise ValueError(f"`context_mask` must be 2D [B, L], got shape {tuple(context_mask.shape)}") + if context_mask.shape[0] != context.shape[0] or context_mask.shape[1] != context.shape[1]: + raise ValueError( + "`context_mask` shape must match `context` shape [B, L], " + f"got {tuple(context_mask.shape)} vs {tuple(context.shape)}" + ) + + batch_size = x.shape[0] + if batch_size != context.shape[0]: + if not self.training and batch_size == 1: + x = x.expand(context.shape[0], -1, -1, -1, -1) + batch_size = context.shape[0] + else: + raise ValueError( + f"Batch mismatch between latents and context: {batch_size} vs {context.shape[0]}." + ) + + if timestep.shape[0] not in (1, batch_size): + raise ValueError( + f"`timestep` length must be 1 or batch_size({batch_size}), got {timestep.shape[0]}" + ) + if timestep.shape[0] == 1 and batch_size > 1: + if self.training: + raise ValueError("During training, timestep length must match batch_size.") + timestep = timestep.expand(batch_size) + return x, timestep, context_mask + + def build_video_to_video_mask( + self, + video_seq_len: int, + video_tokens_per_frame: int, + device: torch.device, + ) -> torch.Tensor: + if video_seq_len <= 0: + raise ValueError(f"`video_seq_len` must be positive, got {video_seq_len}") + if video_tokens_per_frame <= 0: + raise ValueError(f"`video_tokens_per_frame` must be positive, got {video_tokens_per_frame}") + + if self.video_attention_mask_mode == "bidirectional": + return torch.ones((video_seq_len, video_seq_len), dtype=torch.bool, device=device) + + if self.video_attention_mask_mode == "per_frame_causal": + if video_seq_len % video_tokens_per_frame != 0: + raise ValueError( + "`video_seq_len` must be divisible by `video_tokens_per_frame` in `per_frame_causal` mode, " + f"got {video_seq_len} and {video_tokens_per_frame}" + ) + num_video_frames = video_seq_len // video_tokens_per_frame + frame_causal = torch.tril( + torch.ones((num_video_frames, num_video_frames), dtype=torch.bool, device=device) + ) + return frame_causal.repeat_interleave(video_tokens_per_frame, dim=0).repeat_interleave( + video_tokens_per_frame, dim=1 + ) + + if self.video_attention_mask_mode == "first_frame_causal": + video_mask = torch.ones((video_seq_len, video_seq_len), dtype=torch.bool, device=device) + first_frame_tokens = min(video_tokens_per_frame, video_seq_len) + video_mask[:first_frame_tokens, first_frame_tokens:] = False + return video_mask + + raise ValueError(f"Unsupported video attention mask mode: {self.video_attention_mask_mode}") + + def pre_dit( + self, + x: torch.Tensor, + timestep: torch.Tensor, + context: torch.Tensor, + context_mask: torch.Tensor | None = None, + action: torch.Tensor | None = None, + fuse_vae_embedding_in_latents: bool = False, + ) -> dict[str, Any]: + x, timestep, context_mask = self._validate_forward_inputs( + x=x, + timestep=timestep, + context=context, + context_mask=context_mask, + action=action, + ) + model_dtype = self.patch_embedding.weight.dtype + x = x.to(dtype=model_dtype) + context = context.to(dtype=model_dtype) + if action is not None: + action = action.to(dtype=model_dtype) + + batch_size = x.shape[0] + patch_h = int(self.patch_size[1]) + patch_w = int(self.patch_size[2]) + if x.shape[3] % patch_h != 0 or x.shape[4] % patch_w != 0: + raise ValueError( + "Latent spatial shape must be divisible by DiT patch size, " + f"got HxW=({x.shape[3]}, {x.shape[4]}), patch=({patch_h}, {patch_w})" + ) + tokens_per_frame = (x.shape[3] // patch_h) * (x.shape[4] // patch_w) + + if not (self.seperated_timestep and fuse_vae_embedding_in_latents): + raise NotImplementedError( + "FastWAM currently requires separated timesteps with fused VAE latents." + ) + + token_timesteps = torch.ones( + (batch_size, x.shape[2], tokens_per_frame), + dtype=model_dtype, + device=timestep.device, + ) * timestep.to(dtype=model_dtype).view(batch_size, 1, 1) + token_timesteps[:, 0, :] = 0 + token_timesteps = token_timesteps.reshape(batch_size, -1) + # Wan keeps the time embedding in fp32: the AdaLN modulation in the vendored + # Head/Block asserts e.dtype == float32 (numerical stability of the scale/shift). + # Upstream guarantees this via an fp32 autocast region, so it holds even when the + # model runs in bf16. Mirror that here, then cast the per-block modulation back to + # model_dtype so the bf16 attention blocks are not upcast to fp32. + with torch.amp.autocast("cuda", dtype=torch.float32): + token_t_emb = sinusoidal_embedding_1d(self.freq_dim, token_timesteps.reshape(-1)).float() + t = self.time_embedding(token_t_emb).reshape(batch_size, -1, self.hidden_dim) + t_mod = self.time_projection(t).unflatten(2, (6, self.hidden_dim)) + t_mod = t_mod.to(dtype=model_dtype) + + x = self.patchify(x) + f, h, w = x.shape[2:] + + context = self.text_embedding(context) + context_len = context.shape[1] + if self.action_conditioned and action is not None: + action_len = action.shape[1] + action_emb = self.action_embedding(action) + action_pos_embed = sinusoidal_embedding_1d( + self.hidden_dim, torch.arange(action_len, device=action_emb.device) + ).to(dtype=action_emb.dtype) + action_emb = action_emb + action_pos_embed.unsqueeze(0) + context = torch.cat([context, action_emb], dim=1) + + num_temporal_groups = f - 1 + if num_temporal_groups <= 0: + raise ValueError( + "Action-conditioned context mask requires at least 2 latent frames when `action` is provided." + ) + if action_emb.shape[1] % num_temporal_groups != 0: + raise ValueError( + f"Action embedding length {action_emb.shape[1]} must be divisible by " + f"number of temporal groups {num_temporal_groups}" + ) + action_group_mask = create_group_causal_attn_mask( + num_temporal_groups=num_temporal_groups, + num_query_per_group=tokens_per_frame, + num_key_per_group=action_len // num_temporal_groups, + mode=self.action_group_causal_mask_mode, + ).to(context.device) + + seq_len = f * h * w + final_context_mask = torch.zeros( + (batch_size, seq_len, context.shape[1]), dtype=torch.bool, device=context.device + ) + final_context_mask[:, :, :context_len] = context_mask.unsqueeze(1).expand(-1, seq_len, -1) + final_context_mask[:, tokens_per_frame:, context_len:] = action_group_mask.unsqueeze(0).expand( + batch_size, -1, -1 + ) + context_mask = final_context_mask + elif self.action_conditioned and action is None: + if f != 1: + raise ValueError( + "Action-conditioned model requires `action` unless running single-frame text-only mode " + "with num_latent_frames=1." + ) + context_mask = context_mask.unsqueeze(1).expand(-1, f * h * w, -1) + else: + context_mask = context_mask.unsqueeze(1).expand(-1, f * h * w, -1) + + x_tokens = rearrange(x, "b c f h w -> b (f h w) c").contiguous() + grid_sizes = torch.tensor([[f, h, w]] * batch_size, dtype=torch.long, device=x_tokens.device) + freqs = {"grid_sizes": grid_sizes, "freqs": self.freqs.to(x_tokens.device)} + + return { + "tokens": x_tokens, + "freqs": freqs, + "t": t, + "t_mod": t_mod, + "context": context, + "context_mask": context_mask, + "meta": { + "grid_sizes": grid_sizes, + "tokens_per_frame": tokens_per_frame, + "batch_size": batch_size, + }, + } + + def post_dit(self, x_tokens: torch.Tensor, pre_state: dict[str, Any]) -> torch.Tensor: + x = self.head(x_tokens, pre_state["t"]) + return torch.stack(super().unpatchify(x, pre_state["meta"]["grid_sizes"])) + + def forward( + self, + x: torch.Tensor, + timestep: torch.Tensor, + context: torch.Tensor, + context_mask: torch.Tensor | None = None, + action: torch.Tensor | None = None, + fuse_vae_embedding_in_latents: bool = False, + ): + pre_state = self.pre_dit( + x=x, + timestep=timestep, + context=context, + context_mask=context_mask, + action=action, + fuse_vae_embedding_in_latents=fuse_vae_embedding_in_latents, + ) + x_tokens = pre_state["tokens"] + context_emb = pre_state["context"] + t_mod = pre_state["t_mod"] + freqs = pre_state["freqs"] + context_attn_mask = pre_state["context_mask"] + self_attn_mask = ( + self.build_video_to_video_mask( + video_seq_len=x_tokens.shape[1], + video_tokens_per_frame=int(pre_state["meta"]["tokens_per_frame"]), + device=x_tokens.device, + ) + if self.video_attention_mask_mode != "bidirectional" + else None + ) + + for block in self.blocks: + if self.use_gradient_checkpointing: + x_tokens = gradient_checkpoint_forward( + block, + self.use_gradient_checkpointing, + x_tokens, + context_emb, + t_mod, + freqs, + context_mask=context_attn_mask, + self_attn_mask=self_attn_mask, + ) + else: + x_tokens = block( + x_tokens, + context_emb, + t_mod, + freqs, + context_mask=context_attn_mask, + self_attn_mask=self_attn_mask, + ) + + return self.post_dit(x_tokens, pre_state) diff --git a/tests/policies/fastwam/test_fastwam_policy.py b/tests/policies/fastwam/test_fastwam_policy.py new file mode 100644 index 000000000..05e86b7f4 --- /dev/null +++ b/tests/policies/fastwam/test_fastwam_policy.py @@ -0,0 +1,391 @@ +#!/usr/bin/env python + +# 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. + +import json + +import pytest +import torch +from safetensors import safe_open +from torch import nn + +pytest.importorskip("transformers", reason="fastwam requires the `fastwam` extra (transformers)") +pytest.importorskip("diffusers", reason="fastwam requires the `fastwam` extra (diffusers)") + +from lerobot.configs import FeatureType, PolicyFeature, PreTrainedConfig +from lerobot.policies import FastWAMConfig, get_policy_class, make_policy_config, make_pre_post_processors +from lerobot.policies.fastwam.modeling_fastwam import FastWAMPolicy +from lerobot.policies.fastwam.processor_fastwam import FastWAMActionToggleProcessorStep +from lerobot.utils.constants import ACTION, OBS_STATE + + +class FakeFastWAMCore(nn.Module): + def __init__(self): + super().__init__() + self.dit = nn.Linear(2, 2) + + def training_loss(self, sample): + assert sample["video"].ndim == 5 + assert sample["context"].ndim == 3 + return sample[ACTION].sum() * 0.0 + torch.tensor(1.0), {"loss_action": 1.0} + + def infer_action(self, **kwargs): + return {"action": torch.ones(1, kwargs["action_horizon"], 3)} + + +def test_fastwam_is_registered_and_publicly_exported(): + cfg = make_policy_config( + "fastwam", + action_dim=3, + proprio_dim=2, + action_horizon=4, + n_action_steps=2, + num_video_frames=5, + action_video_freq_ratio=1, + base_model_id=None, + ) + + assert isinstance(cfg, FastWAMConfig) + assert cfg.type == "fastwam" + assert get_policy_class("fastwam") is FastWAMPolicy + + +def test_config_validates_features_model_ids_and_saved_auto_route(tmp_path): + cfg = FastWAMConfig() + cfg.save_pretrained(tmp_path) + saved = json.loads((tmp_path / "config.json").read_text()) + + assert saved["pretrained_path"] is None + assert cfg.image_features["observation.images.image"].type == FeatureType.VISUAL + assert cfg.action_feature.shape == (7,) + assert cfg.robot_state_feature.shape == (8,) + with pytest.raises(ValueError, match="image feature"): + FastWAMConfig(input_features={OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(8,))}) + assert FastWAMConfig(tokenizer_model_id="somebody/other-tokenizer").tokenizer_model_id == ( + "somebody/other-tokenizer" + ) + + +def test_preprocessor_passes_images_through_and_postprocessor_toggles_actions(tmp_path): + cfg = FastWAMConfig( + action_dim=3, + proprio_dim=2, + action_horizon=4, + n_action_steps=2, + num_video_frames=5, + action_video_freq_ratio=1, + image_size=(2, 2), + device="cpu", + toggle_action_dimensions=[-1], + input_features={ + "observation.images.image": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 2, 2)), + OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(2,)), + }, + output_features={ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(3,))}, + base_model_id=None, + ) + dataset_stats = { + "observation.images.image": { + "mean": torch.full((3, 1, 1), 0.2), + "std": torch.full((3, 1, 1), 0.1), + }, + OBS_STATE: { + "mean": torch.tensor([1.0, 3.0]), + "std": torch.tensor([2.0, 4.0]), + }, + ACTION: { + "mean": torch.zeros(3), + "std": torch.ones(3), + }, + } + + preprocessor, postprocessor = make_pre_post_processors(cfg, dataset_stats=dataset_stats) + processed = preprocessor( + { + "observation.images.image": torch.tensor( + [ + [[0.0, 0.5], [1.0, 0.5]], + [[0.0, 0.5], [1.0, 0.5]], + [[0.0, 0.5], [1.0, 0.5]], + ] + ), + OBS_STATE: torch.tensor([3.0, 7.0]), + } + ) + preprocessor.save_pretrained(tmp_path, config_filename="policy_preprocessor.json") + postprocessor.save_pretrained(tmp_path, config_filename="policy_postprocessor.json") + _, loaded_postprocessor = make_pre_post_processors(cfg, pretrained_path=str(tmp_path)) + + # VISUAL normalization is IDENTITY + expected_image = torch.tensor( + [[[[0.0, 0.5], [1.0, 0.5]], [[0.0, 0.5], [1.0, 0.5]], [[0.0, 0.5], [1.0, 0.5]]]] + ) + assert preprocessor.name == "policy_preprocessor" + assert postprocessor.name == "policy_postprocessor" + assert torch.allclose(processed["observation.images.image"], expected_image) + assert torch.allclose(processed[OBS_STATE], torch.tensor([[1.0, 1.0]])) + assert torch.equal(dataset_stats["observation.images.image"]["mean"], torch.full((3, 1, 1), 0.2)) + assert any(isinstance(step, FastWAMActionToggleProcessorStep) for step in loaded_postprocessor.steps) + assert torch.equal( + loaded_postprocessor(torch.tensor([[0.25, 0.5, 1.0]])), torch.tensor([[0.25, 0.5, -1.0]]) + ) + + +def test_policy_forward_and_predict_action_adapt_lerobot_batches(monkeypatch): + captured = [] + + class CapturingCore(FakeFastWAMCore): + def infer_action(self, **kwargs): + captured.append( + { + "image_shape": tuple(kwargs["input_image"].shape), + "proprio_shape": tuple(kwargs["proprio"].shape), + "prompt": kwargs["prompt"], + } + ) + return {"action": torch.full((1, kwargs["action_horizon"], 3), float(len(captured)))} + + monkeypatch.setattr(FastWAMPolicy, "_build_core_model", lambda self, config: CapturingCore()) + cfg = FastWAMConfig( + action_dim=3, + proprio_dim=2, + action_horizon=4, + n_action_steps=2, + num_video_frames=5, + action_video_freq_ratio=1, + image_size=(16, 16), + input_features={ + "observation.images.image": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 16, 16)), + OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(2,)), + }, + output_features={ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(3,))}, + base_model_id=None, + ) + policy = FastWAMPolicy(cfg) + + loss, metrics = policy.forward( + { + "observation.images.image": torch.zeros(1, 3, 16, 16), + OBS_STATE: torch.zeros(1, 2), + ACTION: torch.zeros(1, 4, 3), + "context": torch.zeros(1, 5, 4096), + "context_mask": torch.ones(1, 5, dtype=torch.bool), + } + ) + action = policy.predict_action_chunk( + { + "observation.images.image": torch.stack( + [ + torch.zeros(3, 16, 16), + torch.ones(3, 16, 16), + ] + ), + OBS_STATE: torch.tensor([[0.0, 1.0], [2.0, 3.0]]), + "task": ["task 0", "task 1"], + } + ) + + assert loss.item() == 1.0 + assert metrics["loss_action"] == 1.0 + assert action.shape == (2, 4, 3) + assert action[:, 0, 0].tolist() == [1.0, 2.0] + assert [item["image_shape"] for item in captured] == [(1, 3, 16, 16), (1, 3, 16, 16)] + assert [item["proprio_shape"] for item in captured] == [(1, 2), (1, 2)] + assert [item["prompt"] for item in captured] == [ + cfg.prompt_template.format(task="task 0"), + cfg.prompt_template.format(task="task 1"), + ] + + +class CoreWithFrozenComponents(FakeFastWAMCore): + """Fake core mirroring the real one: frozen VAE / text encoder held as + *unregistered* attributes (via `object.__setattr__`) so they are excluded from + `state_dict()` and the saved checkpoint, but still moved by the `_apply` override.""" + + def __init__(self): + super().__init__() + object.__setattr__(self, "vae", nn.Linear(2, 2)) + object.__setattr__(self, "text_encoder", nn.Linear(2, 2)) + self.vae.requires_grad_(False) + self.text_encoder.requires_grad_(False) + + def _apply(self, fn, *args, **kwargs): + super()._apply(fn, *args, **kwargs) + self.vae._apply(fn) + self.text_encoder._apply(fn) + return self + + +def test_from_pretrained_uses_base_loader_and_skips_wan_backbone(monkeypatch, tmp_path): + cfg = FastWAMConfig( + action_dim=3, + proprio_dim=2, + action_horizon=4, + n_action_steps=2, + num_video_frames=5, + action_video_freq_ratio=1, + base_model_id=None, + ) + + def build_core(self, config): + core = CoreWithFrozenComponents() + with torch.no_grad(): + core.dit.weight.fill_(0.5) + return core + + monkeypatch.setattr(FastWAMPolicy, "_build_core_model", build_core) + + reference = FastWAMPolicy(cfg) + with torch.no_grad(): + reference.model.dit.weight.fill_(1.25) # a distinctive, trained-looking weight + reference.save_pretrained(tmp_path) + + # Building from Wan2.2 must never happen on a checkpoint load. + def fail_if_wan_pretrained_is_loaded(*args, **kwargs): + raise AssertionError("from_pretrained must not initialize or download the Wan2.2 backbone") + + monkeypatch.setattr( + "lerobot.policies.fastwam.wan.modular.FastWAM.from_wan22_pretrained", + fail_if_wan_pretrained_is_loaded, + ) + + policy = FastWAMPolicy.from_pretrained(tmp_path) + + assert isinstance(policy.model, CoreWithFrozenComponents) + # The bundled checkpoint weights overwrote the freshly built (0.5) DiT weights. + assert torch.allclose(policy.model.dit.weight, torch.full_like(policy.model.dit.weight, 1.25)) + + +def test_save_pretrained_excludes_frozen_components(monkeypatch, tmp_path): + cfg = FastWAMConfig( + action_dim=3, + proprio_dim=2, + action_horizon=4, + n_action_steps=2, + num_video_frames=5, + action_video_freq_ratio=1, + base_model_id=None, + ) + monkeypatch.setattr(FastWAMPolicy, "_build_core_model", lambda self, config: CoreWithFrozenComponents()) + policy = FastWAMPolicy(cfg) + + save_dir = tmp_path / "saved" + policy.save_pretrained(save_dir) + + assert (save_dir / "model.safetensors").is_file() + # No Wan sidecar files either: the frozen backbone comes from the diffusers repo. + assert not (save_dir / "Wan2.2_VAE.safetensors").exists() + assert not (save_dir / "google").exists() + + with safe_open(save_dir / "model.safetensors", framework="pt") as f: + keys = set(f.keys()) + # Lean checkpoint: only the trainable DiT is saved; the frozen VAE / UMT5 text + # encoder are excluded (loaded from the diffusers/transformers repos at init). + assert any(key.startswith("model.dit.") for key in keys) + assert not any(key.startswith("model.vae.") for key in keys) + assert not any(key.startswith("model.text_encoder.") for key in keys) + + +def test_frozen_components_excluded_from_params_but_follow_device_moves(monkeypatch): + cfg = FastWAMConfig( + action_dim=3, + proprio_dim=2, + action_horizon=4, + n_action_steps=2, + num_video_frames=5, + action_video_freq_ratio=1, + base_model_id=None, + ) + monkeypatch.setattr(FastWAMPolicy, "_build_core_model", lambda self, config: CoreWithFrozenComponents()) + policy = FastWAMPolicy(cfg) + + # Unregistered: excluded from state_dict and from the optimizer's parameter set. + sd = policy.state_dict() + assert not any(k.startswith("model.vae.") or k.startswith("model.text_encoder.") for k in sd) + param_names = [n for n, _ in policy.named_parameters()] + assert not any("vae" in n or "text_encoder" in n for n in param_names) + + # ...but the `_apply` override still carries them through `.to()` (dtype stands in + # for device on a CPU box), so they never strand off the rest of the model. + policy.to(torch.float64) + assert policy.model.dit.weight.dtype == torch.float64 # registered + assert policy.model.vae.weight.dtype == torch.float64 # unregistered, moved via _apply + assert policy.model.text_encoder.weight.dtype == torch.float64 + + +def test_pretrained_config_round_trips_fastwam_features(tmp_path): + cfg = FastWAMConfig(action_dim=7, proprio_dim=8, image_size=(224, 448), base_model_id=None) + cfg.save_pretrained(tmp_path) + + loaded = PreTrainedConfig.from_pretrained(tmp_path) + + assert loaded.type == "fastwam" + assert loaded.image_features["observation.images.image"].type == FeatureType.VISUAL + assert loaded.action_feature.shape == (7,) + assert loaded.robot_state_feature.shape == (8,) + + +def test_vae_adapter_empty_build_encode_decode_shapes(): + """Offline glue check of the diffusers-backed VAE adapter (random weights). + + Validates the encode/decode contract — 48 latent channels, 16x spatial / 4x + temporal compression, list-or-batch input, scaling round-trip — without any + weight download. (Numerical fidelity vs the original Wan VAE is a separate, + GPU + real-weights verification step.) + """ + pytest.importorskip("diffusers") + from diffusers import AutoencoderKLWan + + from lerobot.policies.fastwam.wan import WanVideoVAE38 + + # Production always loads a real pretrained VAE from the diffusers repo; here we + # build the same architecture with random weights and dummy standardization stats + # to exercise the adapter's shape/scaling contract offline (fidelity is checked + # separately, with real weights, on GPU). + arch = { + "base_dim": 160, + "decoder_base_dim": 256, + "z_dim": 48, + "dim_mult": [1, 2, 4, 4], + "num_res_blocks": 2, + "attn_scales": [], + "temporal_downsample": [False, True, True], + "dropout": 0.0, + "is_residual": True, + "in_channels": 12, + "out_channels": 12, + "patch_size": 2, + "scale_factor_spatial": 16, + "scale_factor_temporal": 4, + "clip_output": False, + "latents_mean": [0.0] * 48, + "latents_std": [1.0] * 48, + } + raw = AutoencoderKLWan.from_config(arch) + vae = WanVideoVAE38(dtype=torch.float32, device="cpu", pretrained=raw) + assert vae.z_dim == 48 + assert vae.upsampling_factor == 16 + assert vae.temporal_downsample_factor == 4 + + video = torch.rand(1, 3, 5, 32, 32) * 2 - 1 # [B,C,T,H,W] in [-1,1] + latents = vae.encode(video) + assert latents.shape == (1, 48, 2, 2, 2) # T'=(5-1)//4+1, H'=W'=32//16 + + decoded = vae.decode(latents) + assert decoded.shape[0] == 1 and decoded.shape[1] == 3 and decoded.shape[-2:] == (32, 32) + assert decoded.min() >= -1.0 and decoded.max() <= 1.0 + + # list input is accepted and equals the batched path + assert torch.equal(vae.encode([video[0]]), latents) diff --git a/uv.lock b/uv.lock index 5a76fcbf8..076d021f2 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 2 +revision = 3 requires-python = ">=3.12" resolution-markers = [ "(python_full_version >= '3.15' and platform_machine == 'AMD64' and sys_platform == 'linux') or (python_full_version >= '3.15' and platform_machine == 'x86_64' and sys_platform == 'linux')", @@ -2955,6 +2955,10 @@ eo1 = [ evaluation = [ { name = "av" }, ] +fastwam = [ + { name = "diffusers" }, + { name = "transformers" }, +] feetech = [ { name = "deepdiff" }, { name = "feetech-servo-sdk" }, @@ -3261,11 +3265,13 @@ requires-dist = [ { name = "lerobot", extras = ["deepdiff-dep"], marker = "extra == 'hardware'" }, { name = "lerobot", extras = ["dev"], marker = "extra == 'all'" }, { name = "lerobot", extras = ["diffusers-dep"], marker = "extra == 'diffusion'" }, + { name = "lerobot", extras = ["diffusers-dep"], marker = "extra == 'fastwam'" }, { name = "lerobot", extras = ["diffusers-dep"], marker = "extra == 'groot'" }, { name = "lerobot", extras = ["diffusers-dep"], marker = "extra == 'multi-task-dit'" }, { name = "lerobot", extras = ["diffusers-dep"], marker = "extra == 'vla-jepa'" }, { name = "lerobot", extras = ["diffusion"], marker = "extra == 'all'" }, { name = "lerobot", extras = ["dynamixel"], marker = "extra == 'all'" }, + { name = "lerobot", extras = ["fastwam"], marker = "extra == 'all'" }, { name = "lerobot", extras = ["feetech"], marker = "extra == 'all'" }, { name = "lerobot", extras = ["feetech"], marker = "extra == 'hopejr'" }, { name = "lerobot", extras = ["feetech"], marker = "extra == 'lekiwi'" }, @@ -3335,6 +3341,7 @@ requires-dist = [ { name = "lerobot", extras = ["training"], marker = "extra == 'all'" }, { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'annotations'" }, { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'eo1'" }, + { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'fastwam'" }, { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'groot'" }, { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'hilserl'" }, { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'libero'" }, @@ -3417,7 +3424,7 @@ requires-dist = [ { name = "transformers", marker = "extra == 'transformers-dep'", specifier = ">=5.4.0,<5.6.0" }, { name = "wandb", marker = "extra == 'training'", specifier = ">=0.24.0,<0.28.0" }, ] -provides-extras = ["dataset", "training", "hardware", "viz", "core-scripts", "evaluation", "dataset-viz", "av-dep", "pygame-dep", "placo-dep", "transformers-dep", "grpcio-dep", "accelerate-dep", "can-dep", "peft-dep", "scipy-dep", "diffusers-dep", "qwen-vl-utils-dep", "matplotlib-dep", "pyserial-dep", "deepdiff-dep", "pynput-dep", "pyzmq-dep", "motorbridge-dep", "motorbridge-smart-servo-dep", "feetech", "dynamixel", "damiao", "robstride", "openarms", "gamepad", "hopejr", "lekiwi", "unitree-g1", "reachy2", "rebot", "kinematics", "intelrealsense", "phone", "diffusion", "wallx", "pi", "molmoact2", "smolvla", "multi-task-dit", "groot", "sarm", "robometer", "topreward", "xvla", "eo1", "hilserl", "vla-jepa", "async", "peft", "annotations", "dev", "notebook", "test", "video-benchmark", "aloha", "pusht", "libero", "metaworld", "all"] +provides-extras = ["dataset", "training", "hardware", "viz", "core-scripts", "evaluation", "dataset-viz", "av-dep", "pygame-dep", "placo-dep", "transformers-dep", "grpcio-dep", "accelerate-dep", "can-dep", "peft-dep", "scipy-dep", "diffusers-dep", "qwen-vl-utils-dep", "matplotlib-dep", "pyserial-dep", "deepdiff-dep", "pynput-dep", "pyzmq-dep", "motorbridge-dep", "motorbridge-smart-servo-dep", "feetech", "dynamixel", "damiao", "robstride", "openarms", "gamepad", "hopejr", "lekiwi", "unitree-g1", "reachy2", "rebot", "kinematics", "intelrealsense", "phone", "diffusion", "wallx", "pi", "molmoact2", "smolvla", "multi-task-dit", "groot", "sarm", "robometer", "topreward", "xvla", "eo1", "fastwam", "hilserl", "vla-jepa", "async", "peft", "annotations", "dev", "notebook", "test", "video-benchmark", "aloha", "pusht", "libero", "metaworld", "all"] [[package]] name = "librt" From e6237338619ae303ff0ce7718fb143db98cbd10c Mon Sep 17 00:00:00 2001 From: Nicolas Rabault Date: Wed, 1 Jul 2026 17:05:43 +0200 Subject: [PATCH 26/38] perf(tests): cache draccus docstring extraction (#3903) draccus re-parses each config class's source on every parse() to extract field help text (~2.5s for TrainPipelineConfig). Memoize it for the test session; the source is constant within a run. Fast Tests test time: 664s -> 404s (-39%). --- tests/conftest.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/tests/conftest.py b/tests/conftest.py index cadeaf0d3..31ff07b8e 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -14,14 +14,23 @@ # See the License for the specific language governing permissions and # limitations under the License. +import functools import traceback +import draccus.wrappers.docstring as _draccus_docstring import pytest from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature from lerobot.utils.import_utils import is_package_available from tests.utils import DEVICE +# On every `draccus.parse()`, draccus rebuilds each dataclass field's help text by +# re-reading and re-parsing the class source (draccus.wrappers.docstring). For a config +# as large as TrainPipelineConfig this costs ~2.5s per parse — negligible for the single +# parse a CLI does, but tests parse configs hundreds of times. The source can't change +# within a run, so memoize it for the whole test session. +_draccus_docstring.get_attribute_docstring = functools.cache(_draccus_docstring.get_attribute_docstring) + # Import fixture modules as plugins. # Fixtures that depend on optional packages are only registered when those packages are available, # so that tests can be collected and run even with a minimal install. From 052d329470ea8d5c98a4b4bd1f6c18abd0ac7c34 Mon Sep 17 00:00:00 2001 From: Steven Palma Date: Wed, 1 Jul 2026 18:39:32 +0200 Subject: [PATCH 27/38] feat(visualization): add foxglove support (#3902) * Add Foxglove display mode for teleoperate Add a --display_mode flag (rerun|foxglove) to lerobot-teleoperate. When set to foxglove, stream observations/actions over a Foxglove WebSocket server: images as RawImage/CompressedImage, scalars as typed JSON channels with schemas generated from the feature names (sanitized so paths don't need quoting). Adds a `foxglove` extra. * Add Foxglove display mode to lerobot-record Wire the --display_mode flag (rerun|foxglove) into lerobot-record, matching lerobot-teleoperate: route init/log through the backend-agnostic dispatchers and stop the visualization backend on exit. * update foxglove-sdk to 0.25.1 * Use static lerobot.Scalars schema for Foxglove state topics Replace the per-topic JSON schema derived from feature names with a single static lerobot.Scalars schema: a scalars array of {label, value} objects. The same schema fits any robot regardless of which observation/action features it reports, and the label field lets Foxglove name each series automatically so one filtered path plots every feature. * add foxglove option to dataset viz * Make Foxglove dataset playback loop the sole frame emitter Address review: the listener no longer emits frames, it only mutates playback state and queues a one-shot seek index that the playback loop services. The loop is now the only caller of emit_frame, so concurrent random access into the on-disk dataset / video decoder never overlaps. Also remove the dead server_holder and tighten the _foxglove_safe_name docstring to state what it does and why. * Label Foxglove dataset scalars with feature dimension names Use the dataset's per-dimension feature names (e.g. joint names) as the Foxglove series labels for /observation/state and /action/state instead of bare indices. LeRobot stores `names` inconsistently (flat list, {category: [...]}, or {name: index}), so _feature_dim_names handles each and falls back to indices on any unknown format or length mismatch. * Make Foxglove server host bindable and refactor topic/channel handling Pass display_ip through as the Foxglove WebSocket bind host (127.0.0.1 for local only, 0.0.0.0 for all interfaces) instead of always binding locally. In lerobot-dataset-viz, fold the separate --port into --web-port so one flag covers both the Rerun web viewer and the Foxglove server port. Add a _foxglove_topic() helper and thread a per-topic channel cache through the log helpers so dataset playback stays self-contained instead of mutating the module-global cache. Promote SUCCESS to constants.py. * feat(viz): add support for foxglove in rollout + add to viz tag * fix(docs): remove misleading installation note * fix(visualization): no duplicated prefix, consolidated norm + warnings log * chore(viz): minor improvements * refactor(viz): split files + autoplay + updated docs + added minimal tests * fix(viz): right tags + warning * feat(deprecated ws-port): removing rerun's depreacted ws-port parameter in dataset visualization * chore(web ports): adding global variables for default foxglove/rerun web ports * feat(depth): adding depth support to foxglove visualizer. Because of foxglove limitations (min and max values on RawImage cannot be set from the SDK), depth is normalized between [0,1] when a depth range is provided. * fix(rerun depth range): making rerun depth range computation safe against missing stats * chore(foxglove depth): make it simple, and make it work. * fix(scaling): fixing depth frames scaling --------- Co-authored-by: Roman Shtylman Co-authored-by: Caroline Pascal --- docs/source/il_robots.mdx | 10 +- docs/source/using_dataset_tools.mdx | 2 + pyproject.toml | 1 + src/lerobot/rollout/configs.py | 9 +- src/lerobot/rollout/strategies/core.py | 7 +- src/lerobot/rollout/strategies/episodic.py | 7 +- src/lerobot/scripts/lerobot_dataset_viz.py | 132 +++- src/lerobot/scripts/lerobot_record.py | 35 +- src/lerobot/scripts/lerobot_rollout.py | 16 +- src/lerobot/scripts/lerobot_teleoperate.py | 48 +- src/lerobot/utils/constants.py | 1 + src/lerobot/utils/foxglove_visualization.py | 651 ++++++++++++++++++++ src/lerobot/utils/rerun_visualization.py | 184 ++++++ src/lerobot/utils/visualization_utils.py | 186 ++---- tests/utils/test_foxglove_visualization.py | 101 +++ tests/utils/test_rerun_visualization.py | 310 ++++++++++ tests/utils/test_visualization_utils.py | 300 +-------- uv.lock | 25 + 18 files changed, 1532 insertions(+), 493 deletions(-) create mode 100644 src/lerobot/utils/foxglove_visualization.py create mode 100644 src/lerobot/utils/rerun_visualization.py create mode 100644 tests/utils/test_foxglove_visualization.py create mode 100644 tests/utils/test_rerun_visualization.py diff --git a/docs/source/il_robots.mdx b/docs/source/il_robots.mdx index 178db13bb..64a39e29c 100644 --- a/docs/source/il_robots.mdx +++ b/docs/source/il_robots.mdx @@ -126,7 +126,7 @@ import time from lerobot.teleoperators.so_leader import SO101Leader, SO101LeaderConfig from lerobot.robots.so_follower import SO101Follower, SO101FollowerConfig from lerobot.cameras.opencv import OpenCVCameraConfig -from lerobot.utils.visualization_utils import init_rerun, log_rerun_data, shutdown_rerun +from lerobot.utils.visualization_utils import init_visualization, log_visualization_data, shutdown_visualization robot_config = SO101FollowerConfig( port="/dev/tty.usbmodem5AB90687491", @@ -142,7 +142,7 @@ teleop_config = SO101LeaderConfig( id="my_leader_arm", ) -init_rerun(session_name="teleoperation") +init_visualization("rerun", session_name="teleoperation") # pass "foxglove" to stream to Foxglove instead robot = SO101Follower(robot_config) teleop_device = SO101Leader(teleop_config) @@ -158,7 +158,7 @@ while True: observation = robot.get_observation() action = teleop_device.get_action() robot.send_action(action) - log_rerun_data(observation=observation, action=action) + log_visualization_data("rerun", observation=observation, action=action) elapsed_time = time.perf_counter() - start_time sleep_time = TIME_PER_FRAME - elapsed_time @@ -223,7 +223,7 @@ from lerobot.teleoperators.so_leader.config_so_leader import SO101LeaderConfig from lerobot.teleoperators.so_leader.so_leader import SO101Leader from lerobot.common.control_utils import init_keyboard_listener from lerobot.utils.utils import log_say -from lerobot.utils.visualization_utils import init_rerun +from lerobot.utils.visualization_utils import init_visualization from lerobot.scripts.lerobot_record import record_loop from lerobot.processor import make_default_processors @@ -270,7 +270,7 @@ def main(): # Initialize the keyboard listener and rerun visualization _, events = init_keyboard_listener() - init_rerun(session_name="recording") + init_visualization("rerun", session_name="recording") # Connect the robot and teleoperator robot.connect() diff --git a/docs/source/using_dataset_tools.mdx b/docs/source/using_dataset_tools.mdx index e9299d298..a6dcdb1a7 100644 --- a/docs/source/using_dataset_tools.mdx +++ b/docs/source/using_dataset_tools.mdx @@ -265,6 +265,8 @@ lerobot-dataset-viz \ Once executed, the tool opens `rerun.io` and displays the camera streams, robot states, and actions for the selected episode. +To use [Foxglove](https://foxglove.dev) instead of Rerun, install the extra add `--display-mode foxglove`. This starts a WebSocket server (connect the Foxglove app to `ws://127.0.0.1:8765`) that serves the episode as a seekable timeline you can play/pause and scrub. + For advanced usage—including visualizing datasets stored on a remote server—run: ```bash diff --git a/pyproject.toml b/pyproject.toml index 08ba7fc45..03487e8c4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -125,6 +125,7 @@ hardware = [ ] viz = [ "rerun-sdk>=0.24.0,<0.34.0", + "foxglove-sdk>=0.25.1,<0.26.0", ] # ── User-facing composite extras (map to CLI scripts) ───── # lerobot-record, lerobot-replay, lerobot-calibrate, lerobot-teleoperate, etc. diff --git a/src/lerobot/rollout/configs.py b/src/lerobot/rollout/configs.py index 60c47cfba..639e2ba29 100644 --- a/src/lerobot/rollout/configs.py +++ b/src/lerobot/rollout/configs.py @@ -226,11 +226,14 @@ class RolloutConfig: device: str | None = None task: str = "" display_data: bool = False - # Display data on a remote Rerun server + # Visualization backend used when display_data is True: "rerun" or "foxglove". + display_mode: str = "rerun" + # For "rerun": IP of a remote server to send to. For "foxglove": interface to bind the WebSocket + # server to (127.0.0.1 for local only, 0.0.0.0 for all interfaces). display_ip: str | None = None - # Port of the remote Rerun server + # For "rerun": port of the remote server. For "foxglove": port to bind the WebSocket server to. display_port: int | None = None - # Whether to display compressed images in Rerun + # Whether to display compressed (JPEG) images instead of raw frames display_compressed_images: bool = False # Use vocal synthesis to read events play_sounds: bool = True diff --git a/src/lerobot/rollout/strategies/core.py b/src/lerobot/rollout/strategies/core.py index 9c897522f..460ad12e5 100644 --- a/src/lerobot/rollout/strategies/core.py +++ b/src/lerobot/rollout/strategies/core.py @@ -26,7 +26,7 @@ from lerobot.utils.action_interpolator import ActionInterpolator from lerobot.utils.constants import OBS_STR from lerobot.utils.feature_utils import build_dataset_frame from lerobot.utils.robot_utils import precise_sleep -from lerobot.utils.visualization_utils import log_rerun_data +from lerobot.utils.visualization_utils import log_visualization_data from ..inference import InferenceEngine @@ -162,11 +162,12 @@ class RolloutStrategy(abc.ABC): action_dict: dict | None, runtime_ctx: RuntimeContext, ) -> None: - """Log observation/action telemetry to Rerun if display_data is enabled.""" + """Log observation/action telemetry to the visualization backend if display_data is enabled.""" cfg = runtime_ctx.cfg if not cfg.display_data: return - log_rerun_data( + log_visualization_data( + cfg.display_mode, observation=obs_processed, action=action_dict, compress_images=cfg.display_compressed_images, diff --git a/src/lerobot/rollout/strategies/episodic.py b/src/lerobot/rollout/strategies/episodic.py index e70e66787..15b9bb971 100644 --- a/src/lerobot/rollout/strategies/episodic.py +++ b/src/lerobot/rollout/strategies/episodic.py @@ -44,7 +44,7 @@ from lerobot.utils.feature_utils import build_dataset_frame from lerobot.utils.keyboard_input import init_keyboard_listener from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.utils import log_say -from lerobot.utils.visualization_utils import log_rerun_data +from lerobot.utils.visualization_utils import log_visualization_data from ..configs import EpisodicStrategyConfig from ..context import RolloutContext @@ -171,6 +171,7 @@ class EpisodicStrategy(RolloutStrategy): fps=fps, control_time_s=reset_time_s, display_data=cfg.display_data, + display_mode=cfg.display_mode, display_compressed=display_compressed, ) @@ -259,6 +260,7 @@ class EpisodicStrategy(RolloutStrategy): fps: float, control_time_s: float, display_data: bool, + display_mode: str, display_compressed: bool, ) -> None: """Reset-phase loop: teleop drives the robot if available, no recording.""" @@ -288,7 +290,8 @@ class EpisodicStrategy(RolloutStrategy): if display_data: obs_processed = processors.robot_observation_processor(obs) - log_rerun_data( + log_visualization_data( + display_mode, observation=obs_processed, action=act_teleop, compress_images=display_compressed, diff --git a/src/lerobot/scripts/lerobot_dataset_viz.py b/src/lerobot/scripts/lerobot_dataset_viz.py index 22a7208d4..ee25583a0 100644 --- a/src/lerobot/scripts/lerobot_dataset_viz.py +++ b/src/lerobot/scripts/lerobot_dataset_viz.py @@ -59,6 +59,18 @@ distant$ lerobot-dataset-viz \ local$ rerun rerun+http://IP:GRPC_PORT/proxy ``` +- Visualize data in Foxglove with a seekable, scrubbable timeline: +``` +local$ lerobot-dataset-viz \ + --repo-id lerobot/pusht \ + --episode-index 0 \ + --display-mode foxglove + +# then open the Foxglove app and connect to ws://127.0.0.1:8765 +``` +This starts a Foxglove WebSocket server that serves the episode on demand from the on-disk dataset, +so you can play/pause and scrub anywhere in the episode using Foxglove's playback controls. + """ import argparse @@ -73,9 +85,12 @@ import torch.utils.data import tqdm from lerobot.datasets import LeRobotDataset -from lerobot.utils.constants import ACTION, DONE, OBS_STATE, REWARD +from lerobot.utils.constants import ACTION, DONE, OBS_STATE, REWARD, SUCCESS from lerobot.utils.utils import init_logging +DEFAULT_FOXGLOVE_PORT = 8765 +DEFAULT_RERUN_PORT = 9090 + def get_feature_names(dataset: LeRobotDataset, key: str) -> list[str]: """Return per-dimension names for a feature from the dataset metadata. @@ -108,6 +123,12 @@ def to_hwc_uint8_numpy(chw_float32_torch: torch.Tensor) -> np.ndarray: return hwc_uint8_numpy +def to_hwc_float32_numpy(chw_float32_torch: torch.Tensor) -> np.ndarray: + check_chw_float32(chw_float32_torch) + hwc_float32_numpy = chw_float32_torch.permute(1, 2, 0).numpy() + return hwc_float32_numpy + + def build_blueprint_from_dataset(dataset: LeRobotDataset): """Build a Rerun blueprint laying out camera images and time series for the given dataset. @@ -126,32 +147,43 @@ def build_blueprint_from_dataset(dataset: LeRobotDataset): names = get_feature_names(dataset, key) styling = rr.SeriesLines(names=names) views.append(rrb.TimeSeriesView(origin=origin, name=origin, overrides={origin: styling})) - for key in (DONE, REWARD, "next.success"): + for key in (DONE, REWARD, SUCCESS): if key in dataset.features: views.append(rrb.TimeSeriesView(origin=key, name=key)) return rrb.Blueprint(rrb.Grid(*views)) -def to_hwc_uint16_numpy(chw_float32_torch: torch.Tensor) -> np.ndarray: - check_chw_float32(chw_float32_torch) - hwc_uint16_numpy = chw_float32_torch.round().type(torch.uint16).permute(1, 2, 0).numpy() - return hwc_uint16_numpy - - def visualize_dataset( dataset: LeRobotDataset, episode_index: int, batch_size: int = 32, num_workers: int = 0, mode: str = "local", - web_port: int = 9090, + web_port: int | None = None, grpc_port: int = 9876, save: bool = False, output_dir: Path | None = None, display_compressed_images: bool = False, + display_mode: str = "rerun", + host: str = "127.0.0.1", + autoplay: bool = True, **kwargs, ) -> Path | None: + if display_mode == "foxglove": + from lerobot.utils.foxglove_visualization import serve_foxglove_dataset_playback + + logging.info("Starting Foxglove server") + serve_foxglove_dataset_playback( + dataset, + episode_index, + host=host, + port=web_port if web_port is not None else DEFAULT_FOXGLOVE_PORT, + compress_images=display_compressed_images, + autoplay=autoplay, + ) + return None + if save: assert output_dir is not None, ( "Set an output directory where to write .rrd files with `--output-dir path/to/directory`." @@ -188,14 +220,20 @@ def visualize_dataset( if mode == "distant": server_uri = rr.serve_grpc(grpc_port=grpc_port) logging.info(f"Connect to a Rerun Server: rerun rerun+http://IP:{grpc_port}/proxy") - rr.serve_web_viewer(open_browser=False, web_port=web_port, connect_to=server_uri) + rr.serve_web_viewer( + open_browser=False, + web_port=web_port if web_port is not None else DEFAULT_RERUN_PORT, + connect_to=server_uri, + ) logging.info("Logging to Rerun") # Use the dataset's q01/q99 depth statistics for robust depth range bounds depth_ranges = {} for key in dataset.meta.depth_keys: - stats = dataset.meta.stats[key] + stats = (dataset.meta.stats or {}).get(key) + if not stats: + continue lo = stats["q01"] if "q01" in stats else stats["min"] hi = stats["q99"] if "q99" in stats else stats["max"] depth_ranges[key] = (float(np.asarray(lo).item()), float(np.asarray(hi).item())) @@ -213,11 +251,11 @@ def visualize_dataset( # display each camera image (or depth map) for key in dataset.meta.camera_keys: if key in dataset.meta.depth_keys: - depth = to_hwc_uint16_numpy(batch[key][i]) + depth = to_hwc_float32_numpy(batch[key][i]) depth_entity = rr.DepthImage( depth, colormap=rr.components.Colormap.Viridis, - depth_range=depth_ranges[key], + depth_range=depth_ranges.get(key), ) rr.log(key, entity=depth_entity) else: @@ -239,8 +277,8 @@ def visualize_dataset( if REWARD in batch: rr.log(REWARD, rr.Scalars(batch[REWARD][i].item())) - if "next.success" in batch: - rr.log("next.success", rr.Scalars(batch["next.success"][i].item())) + if SUCCESS in batch: + rr.log(SUCCESS, rr.Scalars(batch[SUCCESS][i].item())) # save .rrd locally if mode == "local" and save: @@ -312,13 +350,11 @@ def main(): parser.add_argument( "--web-port", type=int, - default=9090, - help="Web port for rerun.io when `--mode distant` is set.", - ) - parser.add_argument( - "--ws-port", - type=int, - help="deprecated, please use --grpc-port instead.", + default=None, + help=( + "Web/WebSocket port. For rerun `--mode distant` it is the web viewer port (default 9090); " + "for `--display-mode foxglove` it is the server bind port (default 8765)." + ), ) parser.add_argument( "--grpc-port", @@ -351,24 +387,56 @@ def main(): parser.add_argument( "--display-compressed-images", action="store_true", - help="If set, display compressed images in Rerun instead of uncompressed ones.", + help="If set, display compressed (JPEG) images instead of uncompressed ones.", + ) + + parser.add_argument( + "--display-mode", + type=str, + default="rerun", + choices=["rerun", "foxglove"], + help=( + "Visualization backend. 'rerun' uses the Rerun viewer (--mode/--save/--*-port apply). " + "'foxglove' starts a Foxglove WebSocket server that serves the episode as a seekable, " + "scrubbable timeline; connect the Foxglove app to ws://HOST:PORT (--host/--web-port)." + ), + ) + parser.add_argument( + "--host", + type=str, + default="127.0.0.1", + help=( + "Host to bind the Foxglove WebSocket server to when `--display-mode foxglove` is set " + "(127.0.0.1 for local only, 0.0.0.0 for all interfaces)." + ), + ) + parser.add_argument( + "--no-autoplay", + dest="autoplay", + action="store_false", + help=( + "For `--display-mode foxglove`: don't start playing automatically when a client " + "connects; wait for play to be pressed in the Foxglove app instead." + ), ) args = parser.parse_args() + + if args.display_mode == "foxglove": + rerun_only = ("mode", "save", "output_dir", "grpc_port", "batch_size", "num_workers") + ignored = [name for name in rerun_only if getattr(args, name) != parser.get_default(name)] + if ignored: + logging.warning( + "These flags only apply to `--display-mode rerun` and are ignored with " + "`--display-mode foxglove`: %s.", + ", ".join(f"--{name.replace('_', '-')}" for name in ignored), + ) + kwargs = vars(args) repo_id = kwargs.pop("repo_id") root = kwargs.pop("root") tolerance_s = kwargs.pop("tolerance_s") - if kwargs["ws_port"] is not None: - logging.warning( - "--ws-port is deprecated and will be removed in future versions. Please use --grpc-port instead." - ) - logging.warning("Setting grpc_port to ws_port value.") - kwargs["grpc_port"] = kwargs.pop("ws_port") - else: - kwargs.pop("ws_port") # Always remove ws_port from kwargs - init_logging() logging.info("Loading dataset") dataset = LeRobotDataset(repo_id, episodes=[args.episode_index], root=root, tolerance_s=tolerance_s) diff --git a/src/lerobot/scripts/lerobot_record.py b/src/lerobot/scripts/lerobot_record.py index b759d86e0..4233f0bc2 100644 --- a/src/lerobot/scripts/lerobot_record.py +++ b/src/lerobot/scripts/lerobot_record.py @@ -38,6 +38,9 @@ lerobot-record \\ --display_data=true ``` +To stream the data to Foxglove instead of Rerun, add ``--display_mode=foxglove`` (then connect the +Foxglove app to ``ws://127.0.0.1:8765``; override the port with ``--display_port=``). + Example recording with bimanual so100: ```shell lerobot-record \\ @@ -157,7 +160,11 @@ from lerobot.utils.utils import ( init_logging, log_say, ) -from lerobot.utils.visualization_utils import init_rerun, log_rerun_data +from lerobot.utils.visualization_utils import ( + init_visualization, + log_visualization_data, + shutdown_visualization, +) @dataclass @@ -168,11 +175,14 @@ class RecordConfig: teleop: TeleoperatorConfig | None = None # Display all cameras on screen display_data: bool = False - # Display data on a remote Rerun server + # Visualization backend used when display_data is True: "rerun" or "foxglove". + display_mode: str = "rerun" + # For "rerun": IP of a remote server to send to. For "foxglove": interface to bind the WebSocket + # server to (127.0.0.1 for local only, 0.0.0.0 for all interfaces). display_ip: str | None = None - # Port of the remote Rerun server + # For "rerun": port of the remote server. For "foxglove": port to bind the WebSocket server to. display_port: int | None = None - # Whether to display compressed images in Rerun + # Whether to display compressed (JPEG) images instead of raw frames display_compressed_images: bool = False # Use vocal synthesis to read events. play_sounds: bool = True @@ -233,6 +243,7 @@ def record_loop( control_time_s: int | None = None, single_task: str | None = None, display_data: bool = False, + display_mode: str = "rerun", display_compressed_images: bool = False, ): if dataset is not None and dataset.fps != fps: @@ -327,8 +338,11 @@ def record_loop( dataset.add_frame(frame) if display_data: - log_rerun_data( - observation=obs_processed, action=action_values, compress_images=display_compressed_images + log_visualization_data( + display_mode, + observation=obs_processed, + action=action_values, + compress_images=display_compressed_images, ) dt_s = time.perf_counter() - start_loop_t @@ -354,7 +368,9 @@ def record( init_logging() logging.info(pformat(asdict(cfg))) if cfg.display_data: - init_rerun(session_name="recording", ip=cfg.display_ip, port=cfg.display_port) + init_visualization( + cfg.display_mode, session_name="recording", ip=cfg.display_ip, port=cfg.display_port + ) display_compressed_images = ( True if (cfg.display_data and cfg.display_ip is not None and cfg.display_port is not None) @@ -464,6 +480,7 @@ def record( control_time_s=cfg.dataset.episode_time_s, single_task=cfg.dataset.single_task, display_data=cfg.display_data, + display_mode=cfg.display_mode, display_compressed_images=display_compressed_images, ) @@ -485,6 +502,7 @@ def record( control_time_s=cfg.dataset.reset_time_s, single_task=cfg.dataset.single_task, display_data=cfg.display_data, + display_mode=cfg.display_mode, ) if events["rerecord_episode"]: @@ -510,6 +528,9 @@ def record( if listener is not None: listener.stop() + if cfg.display_data: + shutdown_visualization(cfg.display_mode) + if cfg.dataset.push_to_hub: if dataset and dataset.num_episodes > 0: dataset.push_to_hub(tags=cfg.dataset.tags, private=cfg.dataset.private) diff --git a/src/lerobot/scripts/lerobot_rollout.py b/src/lerobot/scripts/lerobot_rollout.py index daee87bbe..879070721 100644 --- a/src/lerobot/scripts/lerobot_rollout.py +++ b/src/lerobot/scripts/lerobot_rollout.py @@ -145,6 +145,9 @@ Usage examples --dataset.rgb_encoder.vcodec=h264 \\ --dataset.rgb_encoder.preset=fast \\ --dataset.rgb_encoder.extra_options={"tune": "film", "profile:v": "high", "bf": 2} + + # Stream to Foxglove instead of Rerun: + # add --display_mode=foxglove, then connect the Foxglove app to ws://127.0.0.1:8765. """ import logging @@ -190,7 +193,7 @@ from lerobot.teleoperators import ( # noqa: F401 from lerobot.utils.import_utils import register_third_party_plugins from lerobot.utils.process import ProcessSignalHandler from lerobot.utils.utils import init_logging -from lerobot.utils.visualization_utils import init_rerun +from lerobot.utils.visualization_utils import init_visualization, shutdown_visualization logger = logging.getLogger(__name__) @@ -201,8 +204,13 @@ def rollout(cfg: RolloutConfig): init_logging() if cfg.display_data: - logger.info("Initializing Rerun visualization (ip=%s, port=%s)", cfg.display_ip, cfg.display_port) - init_rerun(session_name="rollout", ip=cfg.display_ip, port=cfg.display_port) + logger.info( + "Initializing %s visualization (ip=%s, port=%s)", + cfg.display_mode, + cfg.display_ip, + cfg.display_port, + ) + init_visualization(cfg.display_mode, session_name="rollout", ip=cfg.display_ip, port=cfg.display_port) signal_handler = ProcessSignalHandler(use_threads=True, display_pid=False) shutdown_event = signal_handler.shutdown_event @@ -227,6 +235,8 @@ def rollout(cfg: RolloutConfig): logger.info("Interrupted by user") finally: strategy.teardown(ctx) + if cfg.display_data: + shutdown_visualization(cfg.display_mode) logger.info("Rollout finished") diff --git a/src/lerobot/scripts/lerobot_teleoperate.py b/src/lerobot/scripts/lerobot_teleoperate.py index 5d806200d..30f13987e 100644 --- a/src/lerobot/scripts/lerobot_teleoperate.py +++ b/src/lerobot/scripts/lerobot_teleoperate.py @@ -31,6 +31,22 @@ lerobot-teleoperate \ --display_data=true ``` +To stream the data to Foxglove instead of Rerun, add ``--display_mode=foxglove`` +(then connect the Foxglove app to ``ws://127.0.0.1:8765``; override the port with ``--display_port=``): + +```shell +lerobot-teleoperate \ + --robot.type=so101_follower \ + --robot.port=/dev/tty.usbmodem58760431541 \ + --robot.cameras="{ front: {type: opencv, index_or_path: 0, width: 1920, height: 1080, fps: 30}}" \ + --robot.id=black \ + --teleop.type=so101_leader \ + --teleop.port=/dev/tty.usbmodem58760431551 \ + --teleop.id=blue \ + --display_data=true \ + --display_mode=foxglove +``` + Example teleoperation with bimanual so100: ```shell @@ -108,7 +124,11 @@ from lerobot.teleoperators import ( # noqa: F401 from lerobot.utils.import_utils import register_third_party_plugins from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.utils import init_logging, move_cursor_up -from lerobot.utils.visualization_utils import init_rerun, log_rerun_data, shutdown_rerun +from lerobot.utils.visualization_utils import ( + init_visualization, + log_visualization_data, + shutdown_visualization, +) @dataclass @@ -121,11 +141,14 @@ class TeleoperateConfig: teleop_time_s: float | None = None # Display all cameras on screen display_data: bool = False - # Display data on a remote Rerun server + # Visualization backend used when display_data is True: "rerun" or "foxglove". + display_mode: str = "rerun" + # For "rerun": IP of a remote server to send to. For "foxglove": interface to bind the WebSocket + # server to (127.0.0.1 for local only, 0.0.0.0 for all interfaces). display_ip: str | None = None - # Port of the remote Rerun server + # For "rerun": port of the remote server. For "foxglove": port to bind the WebSocket server to. display_port: int | None = None - # Whether to display compressed images in Rerun + # Whether to display compressed (JPEG) images instead of raw frames display_compressed_images: bool = False @@ -137,6 +160,7 @@ def teleop_loop( robot_action_processor: RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction], robot_observation_processor: RobotProcessorPipeline[RobotObservation, RobotObservation], display_data: bool = False, + display_mode: str = "rerun", duration: float | None = None, display_compressed_images: bool = False, ): @@ -149,8 +173,10 @@ def teleop_loop( teleop: The teleoperator device instance providing control actions. robot: The robot instance being controlled. fps: The target frequency for the control loop in frames per second. - display_data: If True, fetches robot observations and displays them in the console and Rerun. - display_compressed_images: If True, compresses images before sending them to Rerun for display. + display_data: If True, fetches robot observations and displays them in the console and the + visualization backend. + display_mode: Visualization backend to use when display_data is True ("rerun" or "foxglove"). + display_compressed_images: If True, compresses images before sending them to the backend for display. duration: The maximum duration of the teleoperation loop in seconds. If None, the loop runs indefinitely. teleop_action_processor: An optional pipeline to process raw actions from the teleoperator. robot_action_processor: An optional pipeline to process actions before they are sent to the robot. @@ -187,7 +213,8 @@ def teleop_loop( # Process robot observation through pipeline obs_transition = robot_observation_processor(obs) - log_rerun_data( + log_visualization_data( + display_mode, observation=obs_transition, action=teleop_action, compress_images=display_compressed_images, @@ -215,7 +242,9 @@ def teleoperate(cfg: TeleoperateConfig): init_logging() logging.info(pformat(asdict(cfg))) if cfg.display_data: - init_rerun(session_name="teleoperation", ip=cfg.display_ip, port=cfg.display_port) + init_visualization( + cfg.display_mode, session_name="teleoperation", ip=cfg.display_ip, port=cfg.display_port + ) display_compressed_images = ( True if (cfg.display_data and cfg.display_ip is not None and cfg.display_port is not None) @@ -235,6 +264,7 @@ def teleoperate(cfg: TeleoperateConfig): robot=robot, fps=cfg.fps, display_data=cfg.display_data, + display_mode=cfg.display_mode, duration=cfg.teleop_time_s, teleop_action_processor=teleop_action_processor, robot_action_processor=robot_action_processor, @@ -245,7 +275,7 @@ def teleoperate(cfg: TeleoperateConfig): pass finally: if cfg.display_data: - shutdown_rerun() + shutdown_visualization(cfg.display_mode) teleop.disconnect() robot.disconnect() diff --git a/src/lerobot/utils/constants.py b/src/lerobot/utils/constants.py index 482394ff6..8f735fe6d 100644 --- a/src/lerobot/utils/constants.py +++ b/src/lerobot/utils/constants.py @@ -37,6 +37,7 @@ ACTION_TOKEN_MASK = ACTION + ".token_mask" REWARD = "next.reward" TRUNCATED = "next.truncated" DONE = "next.done" +SUCCESS = "next.success" INFO = "info" ROBOTS = "robots" diff --git a/src/lerobot/utils/foxglove_visualization.py b/src/lerobot/utils/foxglove_visualization.py new file mode 100644 index 000000000..fc4136e12 --- /dev/null +++ b/src/lerobot/utils/foxglove_visualization.py @@ -0,0 +1,651 @@ +# 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. + +"""Foxglove visualization backend. + +Live control-loop streaming (:func:`log_foxglove_data`) and seekable dataset playback +(:func:`serve_foxglove_dataset_playback`) over a Foxglove WebSocket server. Callers usually select a +backend at runtime through the dispatch in :mod:`lerobot.utils.visualization_utils` rather than +importing from here directly. Requires the ``viz`` extra (``pip install 'lerobot[viz]'``). +""" + +import logging +import numbers +import time + +import cv2 +import numpy as np + +from lerobot.types import RobotAction, RobotObservation + +from .constants import ( + ACTION, + ACTION_PREFIX, + DONE, + OBS_IMAGES, + OBS_PREFIX, + OBS_STATE, + OBS_STR, + REWARD, + SUCCESS, + TRUNCATED, +) +from .import_utils import require_package + +# Static schema shared by all scalar topics. Each message carries a flat list of ``{label, value}`` +# pairs rather than one field per feature, so the same schema fits any robot regardless of which +# observation/action features it reports. The ``label`` field name is what Foxglove looks for to name +# each series automatically, so a single filtered path plots every feature, e.g. +# ``/observation/state.scalars[:]``. +_SCALARS_SCHEMA = { + "type": "object", + "title": "lerobot.Scalars", + "properties": { + "scalars": { + "type": "array", + "items": { + "type": "object", + "properties": { + "label": {"type": "string"}, + "value": {"type": "number"}, + }, + }, + } + }, +} + + +def _is_scalar(x): + return isinstance(x, (float | numbers.Real | np.integer | np.floating)) or ( + isinstance(x, np.ndarray) and x.ndim == 0 + ) + + +def init_foxglove(host: str = "127.0.0.1", port: int | None = 8765) -> None: + """ + Starts a Foxglove WebSocket server for visualizing the control loop. + + Connect to it from the Foxglove app at ``ws://:``. Calling this + more than once is a no-op while a server is already running. + + Args: + host: Host interface to bind the WebSocket server to. + port: Port to bind the WebSocket server to (defaults to 8765). + """ + + require_package("foxglove-sdk", extra="viz", import_name="foxglove") + import foxglove + + # Live-stream state lives as attributes on ``log_foxglove_data``: + # ``.server`` is the shared WebSocket server and + # ``.channels`` caches one Foxglove channel per topic + if getattr(log_foxglove_data, "server", None) is not None: + return + log_foxglove_data.server = foxglove.start_server(host=host, port=port or 8765) + log_foxglove_data.channels = {} + + +def shutdown_foxglove() -> None: + """Stops the Foxglove WebSocket server and clears cached channels.""" + + server = getattr(log_foxglove_data, "server", None) + if server is not None: + server.stop() + log_foxglove_data.server = None + log_foxglove_data.channels = {} + + +def _foxglove_safe_name(name: str) -> str: + """Replace ``.`` with ``_`` so a feature name is a single Foxglove topic-path segment. + + Foxglove treats ``.`` as a path separator, so an unsanitized name like ``observation.images.front`` + would split into nested segments instead of naming one topic. + """ + + return name.replace(".", "_") + + +def _foxglove_topic(key: str, *, is_image: bool = False) -> str: + """Build the Foxglove topic for a feature ``key``. + + Camera features map to a per-source image topic (``/observation/images/``); scalar features + share one aggregate topic per source: ``/observation/state`` for observations, ``/action/state`` + for actions. + """ + + if is_image: + name = str(key) + for prefix in (f"{OBS_IMAGES}.", OBS_PREFIX): + if name.startswith(prefix): + name = name[len(prefix) :] + break + return f"/{OBS_STR}/images/{_foxglove_safe_name(name)}" + source = ACTION if (str(key).startswith(ACTION_PREFIX) or str(key) == ACTION) else OBS_STR + return f"/{source}/state" + + +def _log_foxglove_scalars( + topic: str, values: dict[str, float], *, channels: dict | None = None, log_time: int | None = None +) -> None: + """Log scalars on a typed JSON channel using the static :data:`_SCALARS_SCHEMA`. + + ``values`` is an ordered mapping of feature name to value; it is emitted as a ``scalars`` array of + ``{label, value}`` objects. Insertion order is preserved so series stay stable across messages. + + ``channels`` is the per-topic channel cache to reuse (defaults to the live-stream cache on + :func:`log_foxglove_data`; dataset playback passes its own local cache to stay self-contained). + ``log_time`` is the message time in nanoseconds; when ``None`` the server's receive time is used. + """ + + if not values: + return + + import foxglove + + if channels is None: + channels = log_foxglove_data.channels + channel = channels.get(topic) + if channel is None: + channel = channels[topic] = foxglove.Channel(topic, schema=_SCALARS_SCHEMA, message_encoding="json") + msg = {"scalars": [{"label": label, "value": value} for label, value in values.items()]} + if log_time is None: + channel.log(msg) + else: + channel.log(msg, log_time=log_time) + + +def _labeled_scalars(name: str, values, labels: list[str] | None = None) -> dict[str, float]: + """Expand a 1D sequence into ``{label: value}`` entries with a consistent fallback.""" + + flat = [float(v) for v in values] + if labels is None or len(labels) != len(flat): + labels = [f"{name}_{i}" for i in range(len(flat))] + return dict(zip(labels, flat, strict=True)) + + +def _log_foxglove_image( + topic: str, + frame_id: str, + arr: np.ndarray, + *, + compress_images: bool, + channels: dict | None = None, + log_time: int | None = None, + depth_range: tuple[float, float] | None = None, + raw_depth_values: bool = False, +) -> None: + """Log an image on a cached per-topic channel. + + The encoding is chosen from the channel count and dtype: a single-channel ``float`` or ``uint16`` + frame is a depth map (``32FC1``/``16UC1``), single-channel ``uint8`` is ``mono8``, 3 => ``rgb8`` + (float input assumed in [0, 1], cast to uint8), 4 => ``rgba8``; other counts are skipped with a + warning. When ``compress_images`` is set, ``rgb8`` is JPEG-encoded instead. + + Args: + topic: Foxglove topic to log on. + frame_id: Frame id stamped on the message. + arr: Image as HWC or CHW (CHW is transposed to HWC), any dtype. + compress_images: JPEG-encode ``rgb8`` frames; ignored for other encodings. + channels: Per-topic channel cache to reuse (see :func:`_log_foxglove_scalars`). + log_time: Message time in nanoseconds, also written to the header timestamp; when ``None`` + the server's receive time is used. + depth_range: ``(lo, hi)`` clip bounds in a depth frame's own input units. Depth frames + (``32FC1``/``16UC1``) are rescaled onto Foxglove's default display max for their encoding + (``1.0`` / ``10000``) so they show with sensible contrast; ``depth_range`` sets the source + range, else the frame's own min/max is used. Ignored for ``mono8``/``rgb8``/``rgba8``. + raw_depth_values: If True, depth values are not rescaled and are logged as is. + """ + + from foxglove.channels import CompressedImageChannel, RawImageChannel + from foxglove.messages import CompressedImage, RawImage, Timestamp + + if channels is None: + channels = log_foxglove_data.channels + time_ns = time.time_ns() if log_time is None else log_time + timestamp = Timestamp(sec=time_ns // 1_000_000_000, nsec=time_ns % 1_000_000_000) + log_kwargs = {} if log_time is None else {"log_time": log_time} + + # Convert CHW -> HWC when needed (mirrors log_rerun_data). + if arr.ndim == 3 and arr.shape[0] in (1, 3, 4) and arr.shape[-1] not in (1, 3, 4): + arr = np.transpose(arr, (1, 2, 0)) + height, width = arr.shape[0], arr.shape[1] + n_channels = 1 if arr.ndim == 2 else arr.shape[2] + + if n_channels == 1 and arr.dtype != np.uint8: + # Depth map: infer the encoding from the dtype. + encoding, target_dtype, value_max = ( + ("32FC1", np.float32, 1.0) + if np.issubdtype(arr.dtype, np.floating) + else ("16UC1", np.uint16, 10000.0) + ) + if not raw_depth_values: + # Rescale onto the encoding's display max with respect to the given depth_range. + lo, hi = depth_range if depth_range is not None else (float(arr.min()), float(arr.max())) + arr = arr.clip(lo, hi).astype(np.float32) + arr = (arr - lo) / ((hi - lo) if hi > lo else 1.0) * value_max + arr = np.ascontiguousarray(arr, dtype=target_dtype) + else: + if n_channels == 3 and np.issubdtype(arr.dtype, np.floating): + arr = (arr * 255.0).clip(0, 255) + arr = np.ascontiguousarray(arr, dtype=np.uint8) + + if compress_images and n_channels == 3: + buf_src = cv2.cvtColor(arr, cv2.COLOR_RGB2BGR) + _, buf = cv2.imencode(".jpg", buf_src) + channel = channels.get(topic) + if channel is None: + channel = channels[topic] = CompressedImageChannel(topic=topic) + channel.log( + CompressedImage(timestamp=timestamp, frame_id=frame_id, data=buf.tobytes(), format="jpeg"), + **log_kwargs, + ) + return + + encoding = {1: "mono8", 3: "rgb8", 4: "rgba8"}.get(n_channels) + if encoding is None: + logging.warning( + "Foxglove: skipping image on topic '%s' with unsupported shape %s (%d channels); " + "expected 1 (mono8/16UC1/32FC1), 3 (rgb8), or 4 (rgba8) channels.", + topic, + tuple(arr.shape), + n_channels, + ) + return + + channel = channels.get(topic) + if channel is None: + channel = channels[topic] = RawImageChannel(topic=topic) + channel.log( + RawImage( + timestamp=timestamp, + frame_id=frame_id, + width=width, + height=height, + encoding=encoding, + step=width * n_channels * arr.itemsize, + data=arr.tobytes(), + ), + **log_kwargs, + ) + + +def log_foxglove_data( + observation: RobotObservation | None = None, + action: RobotAction | None = None, + compress_images: bool = False, +) -> None: + """ + Logs observation and action data to a Foxglove WebSocket server for real-time visualization. + + Mirrors ``log_rerun_data`` but emits Foxglove messages over the server started by + :func:`init_foxglove`. Data is mapped as follows: + - Scalars (and elements of 1D arrays) are accumulated per source and logged on the + ``/observation/state`` and ``/action/state`` topics as typed JSON messages using the static + ``lerobot.Scalars`` schema: a ``scalars`` array of ``{label, value}`` objects (see + :data:`_SCALARS_SCHEMA`). The ``label`` field lets Foxglove name each series automatically, so + ``/observation/state.scalars[:].value`` plots every feature at once. + - 3D NumPy arrays that resemble images are transposed from CHW to HWC when needed and logged on a + per-source topic (e.g. ``/observation/images/front``) as a ``RawImage`` (or a JPEG + ``CompressedImage`` when ``compress_images`` is True). + + Args: + observation: An optional dictionary containing observation data to log. + action: An optional dictionary containing action data to log. + compress_images: Whether to JPEG-compress images before logging to save bandwidth in exchange + for CPU and quality. + """ + + require_package("foxglove-sdk", extra="viz", import_name="foxglove") + + if getattr(log_foxglove_data, "server", None) is None: + raise RuntimeError("init_foxglove() must be called before log_foxglove_data().") + + now = time.time_ns() + + if observation: + obs_scalars: dict[str, float] = {} + for k, v in observation.items(): + if v is None: + continue + key = k[len(OBS_PREFIX) :] if str(k).startswith(OBS_PREFIX) else str(k) + if _is_scalar(v): + obs_scalars[key] = float(v) + elif isinstance(v, np.ndarray): + if v.ndim == 1: + obs_scalars.update(_labeled_scalars(key, v)) + else: + _log_foxglove_image( + _foxglove_topic(k, is_image=True), + key, + v, + compress_images=compress_images, + log_time=now, + ) + _log_foxglove_scalars(_foxglove_topic(OBS_STATE), obs_scalars, log_time=now) + + if action: + action_scalars: dict[str, float] = {} + for k, v in action.items(): + if v is None: + continue + key = k[len(ACTION_PREFIX) :] if str(k).startswith(ACTION_PREFIX) else str(k) + if _is_scalar(v): + action_scalars[key] = float(v) + elif isinstance(v, np.ndarray): + action_scalars.update(_labeled_scalars(key, v.flatten())) + _log_foxglove_scalars(_foxglove_topic(ACTION), action_scalars, log_time=now) + + +# ── Dataset playback over a Foxglove WebSocket server ───────────────────── +# A LeRobotDataset is random-access on disk, so rather than fire-and-forget a forward stream we +# advertise a seekable timeline and serve frames on demand for whatever time the user scrubs/plays +# to in the Foxglove app. This relies on the SDK's PlaybackControl capability. + + +def _feature_dim_names(feature: dict | None) -> list[str] | None: + """Best-effort per-dimension series labels for a 1D feature, or ``None`` to fall back to indices. + + LeRobot records a feature's ``names`` inconsistently: a flat list (``["x", "y"]``), a category + mapping (``{"motors": ["motor_0", "motor_1"]}``), or a name->index mapping + (``{"delta_x": 0, "delta_y": 1}``). Each is handled, but labels are only returned when their count + matches the feature's 1D shape, so a malformed/mismatched ``names`` can't silently mislabel series. + """ + + if not feature: + return None + shape = feature.get("shape") + dim = shape[0] if shape and len(shape) == 1 else None + names = feature.get("names") + labels: list[str] | None = None + if isinstance(names, dict): + values = list(names.values()) + if values and all(isinstance(v, (list, tuple)) for v in values): + labels = [str(n) for group in values for n in group] + elif values and all(isinstance(v, int) and not isinstance(v, bool) for v in values): + labels = [name for name, _ in sorted(names.items(), key=lambda kv: kv[1])] + elif isinstance(names, (list, tuple)): + labels = [str(n) for n in names] + if labels is not None and dim is not None and len(labels) == dim: + return labels + return None + + +def _frame_to_scalars(sample: dict, key: str, labels: list[str] | None = None) -> dict[str, float]: + """Flatten a frame's vector/scalar feature ``key`` into ``{label: value}`` entries. + + ``labels`` provides one name per dimension (from the dataset's feature metadata); when absent or + the wrong length, dimensions fall back to ``{name}_{i}`` (the short feature name), matching the + live stream so series names agree. A scalar feature becomes a single entry. Missing or ``None`` + features yield an empty mapping. + """ + + v = sample.get(key) + if v is None: + return {} + arr = v.numpy() if hasattr(v, "numpy") else np.asarray(v) + if key.startswith(OBS_PREFIX): + name = key[len(OBS_PREFIX) :] + elif key.startswith(ACTION_PREFIX): + name = key[len(ACTION_PREFIX) :] + else: + name = key + if arr.ndim == 0: + return {name: float(arr)} + return _labeled_scalars(name, arr.flatten(), labels) + + +def serve_foxglove_dataset_playback( + dataset, + episode_index: int, + *, + host: str = "127.0.0.1", + port: int = 8765, + compress_images: bool = False, + autoplay: bool = True, +) -> None: + """Serve a single dataset episode to Foxglove as a seekable, scrubbable timeline. + + Starts a Foxglove WebSocket server advertising the ``PlaybackControl`` capability over the + episode's time range. The Foxglove app drives play/pause/seek/speed; a background thread and a + ``ServerListener`` read frames from the on-disk ``dataset`` on demand and log them stamped at + their dataset timestamps, so the user can scrub anywhere in the episode. Blocks until interrupted. + + Args: + dataset: A ``LeRobotDataset`` loaded for the single episode to visualize. + episode_index: Index of the episode being visualized (used only for the session name). + host: Host interface to bind the WebSocket server to. + port: Port to bind the WebSocket server to. + compress_images: Whether to JPEG-compress camera frames before logging. + autoplay: If True, start playing automatically as soon as a client connects, instead of + waiting for the user to press play in the Foxglove app. + """ + + require_package("foxglove-sdk", extra="viz", import_name="foxglove") + import bisect + import threading + + import foxglove + from foxglove.websocket import ( + Capability, + PlaybackCommand, + PlaybackControlRequest, + PlaybackState, + PlaybackStatus, + ServerListener, + ) + + # Per-frame timestamps in nanoseconds (read straight from the table, no video decode). + times_ns = [int(round(float(t) * 1e9)) for t in dataset.hf_dataset["timestamp"]] + n_frames = len(times_ns) + if n_frames == 0: + raise ValueError("Cannot visualize an empty episode.") + first_ns, last_ns = times_ns[0], times_ns[-1] + camera_keys = list(dataset.meta.camera_keys) + # Dataset-wide q01/q99 depth bounds (fallback min/max) used to normalize depth to [0, 1]. + depth_ranges: dict[str, tuple[float, float]] = {} + for key in dataset.meta.depth_keys: + stats = (dataset.meta.stats or {}).get(key) + if not stats: + continue + lo = stats["q01"] if "q01" in stats else stats["min"] + hi = stats["q99"] if "q99" in stats else stats["max"] + depth_ranges[key] = (float(np.asarray(lo).item()), float(np.asarray(hi).item())) + # Per-dimension series labels from the dataset metadata (e.g. joint names), computed once. + scalar_labels = { + OBS_STATE: _feature_dim_names(dataset.meta.features.get(OBS_STATE)), + ACTION: _feature_dim_names(dataset.meta.features.get(ACTION)), + } + # Local channel cache so the playback server is self-contained and doesn't touch the live-stream cache. + channels: dict = {} + + def emit_frame(i: int) -> None: + """Log every channel for frame ``i`` stamped at its dataset timestamp.""" + sample = dataset[i] + log_time = times_ns[i] + for key in camera_keys: + arr = sample.get(key) + if arr is None: + continue + arr = arr.numpy() if hasattr(arr, "numpy") else np.asarray(arr) + _log_foxglove_image( + _foxglove_topic(key, is_image=True), + key, + arr, + compress_images=compress_images, + channels=channels, + log_time=log_time, + depth_range=depth_ranges.get(key), + raw_depth_values=True, + ) + _log_foxglove_scalars( + _foxglove_topic(OBS_STATE), + _frame_to_scalars(sample, OBS_STATE, scalar_labels[OBS_STATE]), + channels=channels, + log_time=log_time, + ) + _log_foxglove_scalars( + _foxglove_topic(ACTION), + _frame_to_scalars(sample, ACTION, scalar_labels[ACTION]), + channels=channels, + log_time=log_time, + ) + episode_scalars = {} + for feat, label in ( + (DONE, "done"), + (TRUNCATED, "truncated"), + (REWARD, "reward"), + (SUCCESS, "success"), + ): + v = sample.get(feat) + if v is not None: + episode_scalars[label] = float(v) + _log_foxglove_scalars("/episode/state", episode_scalars, channels=channels, log_time=log_time) + + lock = threading.Lock() + stop_event = threading.Event() + # Shared playback state, guarded by ``lock``. ``seek_idx`` is a one-shot request set by the + # listener and serviced by the playback loop, which is the *only* thread that emits frames (so + # concurrent random access into the on-disk dataset / video decoder never overlaps). + state = { + "status": PlaybackStatus.Paused, + "cursor": first_ns, + "speed": 1.0, + "last_idx": -1, + "seek_idx": None, + } + + def index_at(t_ns: int) -> int: + return max(0, min(n_frames - 1, bisect.bisect_right(times_ns, t_ns) - 1)) + + # One-shot latch so autoplay fires only on the first client subscription. + autoplay_started = threading.Event() + + class _PlaybackListener(ServerListener): + def on_subscribe(self, client, channel): + # Start playing automatically once a client actually connects (subscribes). Using the + # subscribe hook, rather than starting in Playing up front, means the timeline doesn't + # advance before anyone is watching. Fires once; the user can still pause/seek after. + if not autoplay: + return + with lock: + if autoplay_started.is_set() or state["status"] != PlaybackStatus.Paused: + return + autoplay_started.set() + state["status"] = PlaybackStatus.Playing + cursor, speed = state["cursor"], state["speed"] + server.broadcast_playback_state(PlaybackState(PlaybackStatus.Playing, cursor, speed, False, "")) + + def on_playback_control_request(self, req: PlaybackControlRequest): + # Only mutate state here; the playback loop performs all frame emission. + with lock: + did_seek = False + if req.seek_time is not None: + cursor = max(first_ns, min(last_ns, req.seek_time)) + state["cursor"] = cursor + state["last_idx"] = state["seek_idx"] = index_at(cursor) + did_seek = True + if req.playback_speed and req.playback_speed > 0: + state["speed"] = req.playback_speed + if req.playback_command == PlaybackCommand.Play: + # Restarting from the end replays from the beginning. + if state["cursor"] >= last_ns: + state["cursor"] = first_ns + state["last_idx"] = state["seek_idx"] = 0 + did_seek = True + state["status"] = PlaybackStatus.Playing + elif req.playback_command == PlaybackCommand.Pause: + state["status"] = PlaybackStatus.Paused + status, cursor, speed = state["status"], state["cursor"], state["speed"] + request_id = req.request_id or "" + return PlaybackState(status, cursor, speed, did_seek, request_id) + + server = foxglove.start_server( + name=f"{dataset.repo_id}/episode_{episode_index}", + host=host, + port=port, + capabilities=[Capability.PlaybackControl, Capability.Time], + server_listener=_PlaybackListener(), + playback_time_range=(first_ns, last_ns), + ) + + def playback_loop() -> None: + # Cap how far the cursor may advance in a single tick. A slow frame decode (or any stall) + # would otherwise make ``dt`` huge and produce one enormous catch-up batch; clamping it makes + # playback trail wall-clock under a slow decoder while each tick emits a bounded frame range. + max_tick_dt_s = 0.25 + prev = time.monotonic() + while not stop_event.is_set(): + time.sleep(1.0 / 60.0) + ended = False + speed = 1.0 + with lock: + now = time.monotonic() + dt = min(now - prev, max_tick_dt_s) + prev = now + # A queued seek is always serviced, even while paused, so scrubbing updates the view. + work = [] + seek_idx = state["seek_idx"] + if seek_idx is not None: + state["seek_idx"] = None + work.append(seek_idx) + if state["status"] == PlaybackStatus.Playing: + cursor = state["cursor"] + int(dt * 1e9 * state["speed"]) + start_idx = state["last_idx"] + 1 + if cursor >= last_ns: + cursor, target, ended = last_ns, n_frames - 1, True + else: + target = index_at(cursor) + state["cursor"] = cursor + work.extend(range(start_idx, target + 1)) + # cursor only grows while playing (seeks reset last_idx in the listener), so + # target >= last_idx here; a plain assignment is correct and clearer than max(). + state["last_idx"] = target + if ended: + state["status"] = PlaybackStatus.Ended + if not work: + continue + cursor, speed = state["cursor"], state["speed"] + # Emit outside the lock; this is the only thread that calls emit_frame. Re-check + # stop_event between frames so shutdown stays responsive even mid-batch. + for i in work: + if stop_event.is_set(): + break + emit_frame(i) + server.broadcast_time(cursor) + if ended: + server.broadcast_playback_state(PlaybackState(PlaybackStatus.Ended, cursor, speed, False, "")) + + # Emit the first frame so channels are advertised (done before the loop starts, so emission stays + # single-threaded). Late-connecting clients re-receive frames once they seek/play. + emit_frame(0) + with lock: + state["last_idx"] = 0 + server.broadcast_time(first_ns) + server.broadcast_playback_state(PlaybackState(PlaybackStatus.Paused, first_ns, 1.0, True, "")) + + thread = threading.Thread(target=playback_loop, name="foxglove-playback", daemon=True) + thread.start() + + print(f"Foxglove server running. Connect the Foxglove app to ws://{host}:{port}") + print("Use the playback controls in Foxglove to play/pause and scrub the episode. Ctrl-C to exit.") + try: + while not stop_event.is_set(): + time.sleep(0.5) + except KeyboardInterrupt: + print("Ctrl-C received. Exiting.") + finally: + stop_event.set() + thread.join(timeout=2.0) + server.stop() + channels.clear() diff --git a/src/lerobot/utils/rerun_visualization.py b/src/lerobot/utils/rerun_visualization.py new file mode 100644 index 000000000..af04b18f7 --- /dev/null +++ b/src/lerobot/utils/rerun_visualization.py @@ -0,0 +1,184 @@ +# 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. + +"""Rerun visualization backend. + +Live control-loop streaming to the Rerun viewer (:func:`log_rerun_data`). Callers usually select a +backend at runtime through the dispatch in :mod:`lerobot.utils.visualization_utils` rather than +importing from here directly. Requires the ``viz`` extra (``pip install 'lerobot[viz]'``). +""" + +import numbers +import os + +import numpy as np + +from lerobot.types import RobotAction, RobotObservation + +from .constants import ACTION, ACTION_PREFIX, OBS_PREFIX, OBS_STR +from .import_utils import require_package + + +def _is_scalar(x): + return isinstance(x, (float | numbers.Real | np.integer | np.floating)) or ( + isinstance(x, np.ndarray) and x.ndim == 0 + ) + + +def init_rerun( + session_name: str = "lerobot_control_loop", ip: str | None = None, port: int | None = None +) -> None: + """ + Initializes the Rerun SDK for visualizing the control loop. + + Args: + session_name: Name of the Rerun session. + ip: Optional IP for connecting to a Rerun server. + port: Optional port for connecting to a Rerun server. + """ + + require_package("rerun-sdk", extra="viz", import_name="rerun") + import rerun as rr + + log_rerun_data.blueprint = None # Reset blueprint cache for new session + + batch_size = os.getenv("RERUN_FLUSH_NUM_BYTES", "8000") + os.environ["RERUN_FLUSH_NUM_BYTES"] = batch_size + rr.init(session_name) + memory_limit = os.getenv("LEROBOT_RERUN_MEMORY_LIMIT", "10%") + if ip and port: + rr.connect_grpc(url=f"rerun+http://{ip}:{port}/proxy") + else: + rr.spawn(memory_limit=memory_limit) + + +def shutdown_rerun() -> None: + """Shuts down the Rerun SDK gracefully.""" + + require_package("rerun-sdk", extra="viz", import_name="rerun") + import rerun as rr + + rr.rerun_shutdown() + + +def _build_blueprint(observation_paths: set[str], action_paths: set[str], image_paths: set[str]): + """Build a Rerun blueprint laying out camera images, observation and action scalars in separate views. + + Camera images, observation and action scalars are arranged in a grid. + """ + + # Safe + zero-overhead: `log_rerun_data` already ran the `require_package` guard and imported rerun. + import rerun.blueprint as rrb + + views = [rrb.Spatial2DView(origin=path, name=path) for path in sorted(image_paths)] + + if observation_paths: + views.append(rrb.TimeSeriesView(name="observation", contents=sorted(observation_paths))) + if action_paths: + views.append(rrb.TimeSeriesView(name="action", contents=sorted(action_paths))) + + return rrb.Blueprint(rrb.Grid(*views)) + + +def _ensure_blueprint(observation_paths: set[str], action_paths: set[str], image_paths: set[str]) -> None: + """Build and send the blueprint once, from the first observation and action data.""" + if getattr(log_rerun_data, "blueprint", None) is not None: + return + + if not (observation_paths or action_paths or image_paths): + return + + # Safe + zero-overhead: `log_rerun_data` already ran the `require_package` guard and imported rerun. + import rerun as rr + + blueprint = _build_blueprint(observation_paths, action_paths, image_paths) + log_rerun_data.blueprint = blueprint + rr.send_blueprint(blueprint) + + +def log_rerun_data( + observation: RobotObservation | None = None, + action: RobotAction | None = None, + compress_images: bool = False, +) -> None: + """ + Logs observation and action data to Rerun for real-time visualization. + + This function iterates through the provided observation and action dictionaries and sends their contents + to the Rerun viewer. It handles different data types appropriately: + - Scalars values (floats, ints) are logged as `rr.Scalars`. + - 3D NumPy arrays that resemble images (e.g., with 1, 3, or 4 channels first) are transposed + from CHW to HWC format, (optionally) compressed to JPEG and logged as `rr.Image` or `rr.EncodedImage`. + - 1D NumPy arrays are logged as a single `rr.Scalars` batch under one entity path, so that every + dimension shares the same view instead of being split across one view per element. + - Multi-dimensional **action** arrays are flattened and logged as a single `rr.Scalars` batch. + + Keys are automatically namespaced with "observation." or "action." if not already present. + + On the first call, a blueprint is built and sent so observation and action scalars get separate + time-series views and each image gets its own spatial view. + + Args: + observation: An optional dictionary containing observation data to log. + action: An optional dictionary containing action data to log. + compress_images: Whether to compress images before logging to save bandwidth & memory in exchange for cpu and quality. + """ + + require_package("rerun-sdk", extra="viz", import_name="rerun") + import rerun as rr + + observation_paths: set[str] = set() + action_paths: set[str] = set() + image_paths: set[str] = set() + + if observation: + for k, v in observation.items(): + if v is None: + continue + key = k if str(k).startswith(OBS_PREFIX) else f"{OBS_STR}.{k}" + + if _is_scalar(v): + rr.log(key, rr.Scalars(float(v))) + observation_paths.add(key) + elif isinstance(v, np.ndarray): + arr = v + # Convert CHW -> HWC when needed + if arr.ndim == 3 and arr.shape[0] in (1, 3, 4) and arr.shape[-1] not in (1, 3, 4): + arr = np.transpose(arr, (1, 2, 0)) + if arr.ndim == 1: + rr.log(key, rr.Scalars(arr.astype(float))) + observation_paths.add(key) + else: + if arr.shape[-1] == 1: + img_entity = rr.DepthImage(arr, colormap=rr.components.Colormap.Viridis) + else: + img_entity = rr.Image(arr).compress() if compress_images else rr.Image(arr) + rr.log(key, entity=img_entity, static=True) + image_paths.add(key) + + if action: + for k, v in action.items(): + if v is None: + continue + key = k if str(k).startswith(ACTION_PREFIX) else f"{ACTION}.{k}" + + if _is_scalar(v): + rr.log(key, rr.Scalars(float(v))) + action_paths.add(key) + elif isinstance(v, np.ndarray): + # Flatten any (incl. higher-dimensional) array into a single batched Scalars + rr.log(key, rr.Scalars(v.reshape(-1).astype(float))) + action_paths.add(key) + + _ensure_blueprint(observation_paths, action_paths, image_paths) diff --git a/src/lerobot/utils/visualization_utils.py b/src/lerobot/utils/visualization_utils.py index a0f07f0c7..09a89b20a 100644 --- a/src/lerobot/utils/visualization_utils.py +++ b/src/lerobot/utils/visualization_utils.py @@ -12,166 +12,68 @@ # See the License for the specific language governing permissions and # limitations under the License. -import numbers -import os +"""Backend-agnostic visualization dispatch. -import numpy as np +Selects a visualization backend at runtime via a display-mode string (e.g. a ``--display_mode`` CLI +flag) so callers never branch on the backend. The concrete implementations live in +:mod:`lerobot.utils.rerun_visualization` and :mod:`lerobot.utils.foxglove_visualization`; importing +this module does not import ``rerun`` or ``foxglove`` (each backend imports its SDK lazily behind a +``require_package`` guard). +""" from lerobot.types import RobotAction, RobotObservation -from .constants import ACTION, ACTION_PREFIX, OBS_PREFIX, OBS_STR -from .import_utils import require_package +from .foxglove_visualization import init_foxglove, log_foxglove_data, shutdown_foxglove +from .rerun_visualization import init_rerun, log_rerun_data, shutdown_rerun + +# Visualization backends selectable at runtime via a display-mode string (e.g. a --display_mode flag). +VISUALIZATION_MODES = ("rerun", "foxglove") -def init_rerun( - session_name: str = "lerobot_control_loop", ip: str | None = None, port: int | None = None +def init_visualization( + display_mode: str, + *, + session_name: str = "lerobot_control_loop", + ip: str | None = None, + port: int | None = None, ) -> None: - """ - Initializes the Rerun SDK for visualizing the control loop. + """Initializes the visualization backend selected by ``display_mode``. - Args: - session_name: Name of the Rerun session. - ip: Optional IP for connecting to a Rerun server. - port: Optional port for connecting to a Rerun server. + For ``"rerun"``, ``ip``/``port`` point at an optional remote Rerun server. For ``"foxglove"``, + ``ip`` is the interface to bind the WebSocket server to (``127.0.0.1`` for local only, ``0.0.0.0`` + for all interfaces) and ``port`` is its port. """ - require_package("rerun-sdk", extra="viz", import_name="rerun") - import rerun as rr - - log_rerun_data.blueprint = None # Reset blueprint cache for new session - - batch_size = os.getenv("RERUN_FLUSH_NUM_BYTES", "8000") - os.environ["RERUN_FLUSH_NUM_BYTES"] = batch_size - rr.init(session_name) - memory_limit = os.getenv("LEROBOT_RERUN_MEMORY_LIMIT", "10%") - if ip and port: - rr.connect_grpc(url=f"rerun+http://{ip}:{port}/proxy") + if display_mode == "rerun": + init_rerun(session_name=session_name, ip=ip, port=port) + elif display_mode == "foxglove": + init_foxglove(host=ip or "127.0.0.1", port=port) else: - rr.spawn(memory_limit=memory_limit) + raise ValueError(f"Unknown display_mode '{display_mode}'. Expected one of {VISUALIZATION_MODES}.") -def shutdown_rerun() -> None: - """Shuts down the Rerun SDK gracefully.""" - - require_package("rerun-sdk", extra="viz", import_name="rerun") - import rerun as rr - - rr.rerun_shutdown() - - -def _is_scalar(x): - return isinstance(x, (float | numbers.Real | np.integer | np.floating)) or ( - isinstance(x, np.ndarray) and x.ndim == 0 - ) - - -def _build_blueprint(observation_paths: set[str], action_paths: set[str], image_paths: set[str]): - """Build a Rerun blueprint laying out camera images, observation and action scalars in separate views. - - Camera images, observation and action scalars are arranged in a grid. - """ - - # Safe + zero-overhead: `log_rerun_data` already ran the `require_package` guard and imported rerun. - import rerun.blueprint as rrb - - views = [rrb.Spatial2DView(origin=path, name=path) for path in sorted(image_paths)] - - if observation_paths: - views.append(rrb.TimeSeriesView(name="observation", contents=sorted(observation_paths))) - if action_paths: - views.append(rrb.TimeSeriesView(name="action", contents=sorted(action_paths))) - - return rrb.Blueprint(rrb.Grid(*views)) - - -def _ensure_blueprint(observation_paths: set[str], action_paths: set[str], image_paths: set[str]) -> None: - """Build and send the blueprint once, from the first observation and action data.""" - if getattr(log_rerun_data, "blueprint", None) is not None: - return - - if not (observation_paths or action_paths or image_paths): - return - - # Safe + zero-overhead: `log_rerun_data` already ran the `require_package` guard and imported rerun. - import rerun as rr - - blueprint = _build_blueprint(observation_paths, action_paths, image_paths) - log_rerun_data.blueprint = blueprint - rr.send_blueprint(blueprint) - - -def log_rerun_data( +def log_visualization_data( + display_mode: str, observation: RobotObservation | None = None, action: RobotAction | None = None, compress_images: bool = False, ) -> None: - """ - Logs observation and action data to Rerun for real-time visualization. + """Logs observation/action data to the backend selected by ``display_mode``.""" - This function iterates through the provided observation and action dictionaries and sends their contents - to the Rerun viewer. It handles different data types appropriately: - - Scalars values (floats, ints) are logged as `rr.Scalars`. - - 3D NumPy arrays that resemble images (e.g., with 1, 3, or 4 channels first) are transposed - from CHW to HWC format, (optionally) compressed to JPEG and logged as `rr.Image` or `rr.EncodedImage`. - - 1D NumPy arrays are logged as a single `rr.Scalars` batch under one entity path, so that every - dimension shares the same view instead of being split across one view per element. - - Multi-dimensional **action** arrays are flattened and logged as a single `rr.Scalars` batch. + if display_mode == "rerun": + log_rerun_data(observation=observation, action=action, compress_images=compress_images) + elif display_mode == "foxglove": + log_foxglove_data(observation=observation, action=action, compress_images=compress_images) + else: + raise ValueError(f"Unknown display_mode '{display_mode}'. Expected one of {VISUALIZATION_MODES}.") - Keys are automatically namespaced with "observation." or "action." if not already present. - On the first call, a blueprint is built and sent so observation and action scalars get separate - time-series views and each image gets its own spatial view. +def shutdown_visualization(display_mode: str) -> None: + """Shuts down the backend selected by ``display_mode``.""" - Args: - observation: An optional dictionary containing observation data to log. - action: An optional dictionary containing action data to log. - compress_images: Whether to compress images before logging to save bandwidth & memory in exchange for cpu and quality. - """ - - require_package("rerun-sdk", extra="viz", import_name="rerun") - import rerun as rr - - observation_paths: set[str] = set() - action_paths: set[str] = set() - image_paths: set[str] = set() - - if observation: - for k, v in observation.items(): - if v is None: - continue - key = k if str(k).startswith(OBS_PREFIX) else f"{OBS_STR}.{k}" - - if _is_scalar(v): - rr.log(key, rr.Scalars(float(v))) - observation_paths.add(key) - elif isinstance(v, np.ndarray): - arr = v - # Convert CHW -> HWC when needed - if arr.ndim == 3 and arr.shape[0] in (1, 3, 4) and arr.shape[-1] not in (1, 3, 4): - arr = np.transpose(arr, (1, 2, 0)) - if arr.ndim == 1: - rr.log(key, rr.Scalars(arr.astype(float))) - observation_paths.add(key) - else: - if arr.shape[-1] == 1: - img_entity = rr.DepthImage(arr, colormap=rr.components.Colormap.Viridis) - else: - img_entity = rr.Image(arr).compress() if compress_images else rr.Image(arr) - rr.log(key, entity=img_entity, static=True) - image_paths.add(key) - - if action: - for k, v in action.items(): - if v is None: - continue - key = k if str(k).startswith(ACTION_PREFIX) else f"{ACTION}.{k}" - - if _is_scalar(v): - rr.log(key, rr.Scalars(float(v))) - action_paths.add(key) - elif isinstance(v, np.ndarray): - # Flatten any (incl. higher-dimensional) array into a single batched Scalars - rr.log(key, rr.Scalars(v.reshape(-1).astype(float))) - action_paths.add(key) - - _ensure_blueprint(observation_paths, action_paths, image_paths) + if display_mode == "rerun": + shutdown_rerun() + elif display_mode == "foxglove": + shutdown_foxglove() + else: + raise ValueError(f"Unknown display_mode '{display_mode}'. Expected one of {VISUALIZATION_MODES}.") diff --git a/tests/utils/test_foxglove_visualization.py b/tests/utils/test_foxglove_visualization.py new file mode 100644 index 000000000..f3e80315c --- /dev/null +++ b/tests/utils/test_foxglove_visualization.py @@ -0,0 +1,101 @@ +#!/usr/bin/env python + +# Copyright 2025 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. + +"""Tests for the Foxglove backend's pure helpers. + +These cover topic naming, series labelling and feature-name parsing. They import +``foxglove_visualization`` directly and need NO ``foxglove`` extra: the SDK is imported lazily inside +the functions that talk to the server, so the helpers below run in the base test tier. +""" + +import numpy as np + +from lerobot.utils import foxglove_visualization as fv +from lerobot.utils.constants import ACTION, OBS_STATE + + +def test_foxglove_safe_name_collapses_dots(): + assert fv._foxglove_safe_name("observation.images.front") == "observation_images_front" + assert fv._foxglove_safe_name("plain") == "plain" + + +def test_foxglove_topic_image_strips_prefix_without_doubling_images(): + # Fully-qualified camera key -> single clean segment (no doubled "images"). + assert fv._foxglove_topic("observation.images.front", is_image=True) == "/observation/images/front" + # A nested camera name keeps its structure via safe-name collapsing. + assert ( + fv._foxglove_topic("observation.images.wrist.left", is_image=True) == "/observation/images/wrist_left" + ) + # Bare camera name (as real robots emit). + assert fv._foxglove_topic("front", is_image=True) == "/observation/images/front" + + +def test_foxglove_topic_scalar_sources(): + assert fv._foxglove_topic(OBS_STATE) == "/observation/state" + assert fv._foxglove_topic("observation.environment_state") == "/observation/state" + assert fv._foxglove_topic(ACTION) == "/action/state" + assert fv._foxglove_topic("action.delta") == "/action/state" + + +def test_labeled_scalars_uses_labels_then_index_fallback(): + assert fv._labeled_scalars("state", np.array([1.0, 2.0, 3.0])) == { + "state_0": 1.0, + "state_1": 2.0, + "state_2": 3.0, + } + assert fv._labeled_scalars("state", [1.0, 2.0], ["pan", "lift"]) == {"pan": 1.0, "lift": 2.0} + # Wrong-length labels fall back to index naming (never silently mislabels). + assert fv._labeled_scalars("q", [1.0, 2.0], ["only_one"]) == {"q_0": 1.0, "q_1": 2.0} + + +def test_frame_to_scalars_matches_live_labeling_and_handles_scalar(): + frame = {OBS_STATE: np.array([1.0, 2.0])} + # No metadata -> {short_name}_{i}, identical to the live-stream fallback. + assert fv._frame_to_scalars(frame, OBS_STATE) == fv._labeled_scalars("state", np.array([1.0, 2.0])) + assert fv._frame_to_scalars(frame, OBS_STATE) == {"state_0": 1.0, "state_1": 2.0} + # Metadata labels are honored. + assert fv._frame_to_scalars(frame, OBS_STATE, ["pan", "lift"]) == {"pan": 1.0, "lift": 2.0} + # A 0-d scalar becomes a single entry named by the short feature name. + assert fv._frame_to_scalars({ACTION: np.array(5.0)}, ACTION) == {"action": 5.0} + # A missing feature yields an empty mapping. + assert fv._frame_to_scalars({}, OBS_STATE) == {} + + +def test_feature_dim_names_formats(): + # Flat list of names. + assert fv._feature_dim_names({"shape": [2], "names": ["x", "y"]}) == ["x", "y"] + # Category mapping (dict of lists). + assert fv._feature_dim_names({"shape": [2], "names": {"motors": ["m0", "m1"]}}) == ["m0", "m1"] + # name -> index mapping (returned sorted by index). + assert fv._feature_dim_names({"shape": [2], "names": {"delta_x": 0, "delta_y": 1}}) == [ + "delta_x", + "delta_y", + ] + # Bool values must NOT be treated as an index map (bool is a subclass of int). + assert fv._feature_dim_names({"shape": [2], "names": {"a": True, "b": False}}) is None + # Mismatched length -> None (won't silently mislabel). + assert fv._feature_dim_names({"shape": [3], "names": ["x", "y"]}) is None + # Missing / absent names -> None. + assert fv._feature_dim_names(None) is None + assert fv._feature_dim_names({"shape": [2]}) is None + + +def test_is_scalar(): + assert fv._is_scalar(1.0) + assert fv._is_scalar(np.float32(2.0)) + assert fv._is_scalar(np.array(3.0)) # 0-d array + assert not fv._is_scalar(np.array([1.0, 2.0])) + assert not fv._is_scalar("x") diff --git a/tests/utils/test_rerun_visualization.py b/tests/utils/test_rerun_visualization.py new file mode 100644 index 000000000..e3d205dee --- /dev/null +++ b/tests/utils/test_rerun_visualization.py @@ -0,0 +1,310 @@ +#!/usr/bin/env python + +# Copyright 2025 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. + +import importlib +import sys +from types import SimpleNamespace + +import numpy as np +import pytest + +pytest.importorskip("rerun", reason="rerun-sdk is required (install lerobot[viz])") + +from lerobot.types import TransitionKey +from lerobot.utils.constants import OBS_STATE + + +@pytest.fixture +def mock_rerun(monkeypatch): + """ + Provide a mock `rerun` module (and `rerun.blueprint` submodule) so tests don't + depend on the real library. Also reload the module-under-test so it binds to + this mock `rr`. + """ + calls = [] + blueprints = [] + + class DummyScalar: + def __init__(self, value): + # Scalars may be built from a single float or from a 1D array batch. + self.value = value + + class DummyImage: + def __init__(self, arr): + self.arr = arr + + def compress(self, *a, **k): + return self + + class DummyDepthImage: + def __init__(self, arr, colormap=None): + self.arr = arr + self.colormap = colormap + + def dummy_log(key, obj=None, **kwargs): + # Accept either positional `obj` or keyword `entity` and record remaining kwargs. + if obj is None and "entity" in kwargs: + obj = kwargs.pop("entity") + calls.append((key, obj, kwargs)) + + def dummy_send_blueprint(blueprint, *a, **k): + blueprints.append(blueprint) + + # Mock the `rerun.blueprint` submodule used to build the layout. + dummy_rrb = SimpleNamespace( + Spatial2DView=lambda origin=None, name=None: SimpleNamespace( + kind="Spatial2DView", origin=origin, name=name + ), + TimeSeriesView=lambda name=None, contents=None: SimpleNamespace( + kind="TimeSeriesView", name=name, contents=contents + ), + Grid=lambda *views: SimpleNamespace(kind="Grid", views=list(views)), + Blueprint=lambda root: SimpleNamespace(kind="Blueprint", root=root), + ) + + dummy_rr = SimpleNamespace( + __name__="rerun", + __package__="rerun", + __spec__=SimpleNamespace(name="rerun", submodule_search_locations=None), + Scalars=DummyScalar, + Image=DummyImage, + DepthImage=DummyDepthImage, + components=SimpleNamespace(Colormap=SimpleNamespace(Viridis="viridis")), + log=dummy_log, + send_blueprint=dummy_send_blueprint, + init=lambda *a, **k: None, + spawn=lambda *a, **k: None, + blueprint=dummy_rrb, + ) + + # Inject fake modules into sys.modules (both `rerun` and `rerun.blueprint`). + monkeypatch.setitem(sys.modules, "rerun", dummy_rr) + monkeypatch.setitem(sys.modules, "rerun.blueprint", dummy_rrb) + + # Now import and reload the module under test, to bind to our rerun mock + import lerobot.utils.rerun_visualization as rv + + importlib.reload(rv) + + # Expose the reloaded module, the call recorder and the captured blueprints + yield rv, calls, blueprints + + +def _keys(calls): + """Helper to extract just the keys logged to rr.log""" + return [k for (k, _obj, _kw) in calls] + + +def _obj_for(calls, key): + """Find the first object logged under a given key.""" + for k, obj, _kw in calls: + if k == key: + return obj + raise KeyError(f"Key {key} not found in calls: {calls}") + + +def _kwargs_for(calls, key): + for k, _obj, kw in calls: + if k == key: + return kw + raise KeyError(f"Key {key} not found in calls: {calls}") + + +def _views_by_kind(blueprint, kind): + """Return the views of a given kind from the (single) blueprint's grid.""" + return [v for v in blueprint.root.views if v.kind == kind] + + +def test_log_rerun_data_envtransition_scalars_and_image(mock_rerun): + rv, calls, blueprints = mock_rerun + + # Build EnvTransition dict + obs = { + f"{OBS_STATE}.temperature": np.float32(25.0), + # CHW image should be converted to HWC for rr.Image + "observation.camera": np.zeros((3, 10, 20), dtype=np.uint8), + } + act = { + "action.throttle": 0.7, + # 1D array should be logged as a single Scalars batch under one entity path + "action.vector": np.array([1.0, 2.0], dtype=np.float32), + } + transition = { + TransitionKey.OBSERVATION: obs, + TransitionKey.ACTION: act, + } + + # Extract observation and action data from transition like in the real call sites + obs_data = transition.get(TransitionKey.OBSERVATION, {}) + action_data = transition.get(TransitionKey.ACTION, {}) + rv.log_rerun_data(observation=obs_data, action=action_data) + + # We expect: + # - observation.state.temperature -> Scalars + # - observation.camera -> Image (HWC) with static=True + # - action.throttle -> Scalars + # - action.vector -> single Scalars batch (no per-element suffix) + expected_keys = { + f"{OBS_STATE}.temperature", + "observation.camera", + "action.throttle", + "action.vector", + } + assert set(_keys(calls)) == expected_keys + + # Check scalar types and values + temp_obj = _obj_for(calls, f"{OBS_STATE}.temperature") + assert type(temp_obj).__name__ == "DummyScalar" + assert float(temp_obj.value) == pytest.approx(25.0) + + throttle_obj = _obj_for(calls, "action.throttle") + assert type(throttle_obj).__name__ == "DummyScalar" + assert float(throttle_obj.value) == pytest.approx(0.7) + + # 1D vector logged as a single batched Scalars under one entity path + vec = _obj_for(calls, "action.vector") + assert type(vec).__name__ == "DummyScalar" + np.testing.assert_allclose(np.asarray(vec.value), [1.0, 2.0]) + + # Check image handling: CHW -> HWC + img_obj = _obj_for(calls, "observation.camera") + assert type(img_obj).__name__ == "DummyImage" + assert img_obj.arr.shape == (10, 20, 3) # transposed + assert _kwargs_for(calls, "observation.camera").get("static", False) is True # static=True for images + + # A blueprint should have been built and sent exactly once, and cached on the function. + assert len(blueprints) == 1 + assert rv.log_rerun_data.blueprint is blueprints[0] + + bp = blueprints[0] + # One spatial view per image path + spatial_views = _views_by_kind(bp, "Spatial2DView") + assert {v.origin for v in spatial_views} == {"observation.camera"} + + # One time-series view each for observation and action scalars + ts_views = {v.name: v for v in _views_by_kind(bp, "TimeSeriesView")} + assert set(ts_views) == {"observation", "action"} + assert ts_views["observation"].contents == [f"{OBS_STATE}.temperature"] + assert ts_views["action"].contents == ["action.throttle", "action.vector"] + + +def test_log_rerun_data_plain_list_ordering_and_prefixes(mock_rerun): + rv, calls, blueprints = mock_rerun + + # First dict without prefixes treated as observation + # Second dict without prefixes treated as action + obs_plain = { + "temp": 1.5, + # Already HWC image => should stay as-is + "img": np.zeros((5, 6, 3), dtype=np.uint8), + "none": None, # should be skipped + } + act_plain = { + "throttle": 0.3, + "vec": np.array([9, 8, 7], dtype=np.float32), + } + + # Extract observation and action data from list like the old function logic did + # First dict was treated as observation, second as action + rv.log_rerun_data(observation=obs_plain, action=act_plain) + + # Expected keys with auto-prefixes. The 1D vector is a single batched Scalars. + expected = { + "observation.temp", + "observation.img", + "action.throttle", + "action.vec", + } + logged = set(_keys(calls)) + assert logged == expected + + # Scalars + t = _obj_for(calls, "observation.temp") + assert type(t).__name__ == "DummyScalar" + assert float(t.value) == pytest.approx(1.5) + + throttle = _obj_for(calls, "action.throttle") + assert type(throttle).__name__ == "DummyScalar" + assert float(throttle.value) == pytest.approx(0.3) + + # Image stays HWC + img = _obj_for(calls, "observation.img") + assert type(img).__name__ == "DummyImage" + assert img.arr.shape == (5, 6, 3) + assert _kwargs_for(calls, "observation.img").get("static", False) is True + + # Vector logged as a single batched Scalars under one entity path + vec = _obj_for(calls, "action.vec") + assert type(vec).__name__ == "DummyScalar" + np.testing.assert_allclose(np.asarray(vec.value), [9, 8, 7]) + + # Blueprint sent once with the expected view layout + assert len(blueprints) == 1 + bp = blueprints[0] + spatial_views = _views_by_kind(bp, "Spatial2DView") + assert {v.origin for v in spatial_views} == {"observation.img"} + ts_views = {v.name: v for v in _views_by_kind(bp, "TimeSeriesView")} + assert ts_views["observation"].contents == ["observation.temp"] + assert ts_views["action"].contents == ["action.throttle", "action.vec"] + + +def test_log_rerun_data_kwargs_only(mock_rerun): + rv, calls, blueprints = mock_rerun + + rv.log_rerun_data( + observation={"observation.temp": 10.0, "observation.gray": np.zeros((8, 8, 1), dtype=np.uint8)}, + action={"action.a": 1.0}, + ) + + keys = set(_keys(calls)) + assert "observation.temp" in keys + assert "observation.gray" in keys + assert "action.a" in keys + + temp = _obj_for(calls, "observation.temp") + assert type(temp).__name__ == "DummyScalar" + assert float(temp.value) == pytest.approx(10.0) + + img = _obj_for(calls, "observation.gray") + assert type(img).__name__ == "DummyDepthImage" # single-channel -> DepthImage + assert img.arr.shape == (8, 8, 1) # remains HWC + assert _kwargs_for(calls, "observation.gray").get("static", False) is True + + a = _obj_for(calls, "action.a") + assert type(a).__name__ == "DummyScalar" + assert float(a.value) == pytest.approx(1.0) + + # Blueprint sent once, with a spatial view for the image and time-series views for scalars + assert len(blueprints) == 1 + bp = blueprints[0] + assert {v.origin for v in _views_by_kind(bp, "Spatial2DView")} == {"observation.gray"} + ts_views = {v.name: v for v in _views_by_kind(bp, "TimeSeriesView")} + assert ts_views["observation"].contents == ["observation.temp"] + assert ts_views["action"].contents == ["action.a"] + + +def test_log_rerun_data_blueprint_sent_only_once(mock_rerun): + """The blueprint is built from the first call and not resent on subsequent calls.""" + rv, calls, blueprints = mock_rerun + + rv.log_rerun_data(observation={"temp": 1.0}, action={"a": 2.0}) + assert len(blueprints) == 1 + first_blueprint = rv.log_rerun_data.blueprint + + rv.log_rerun_data(observation={"temp": 3.0}, action={"a": 4.0}) + # Still only one blueprint, and the cached one is unchanged. + assert len(blueprints) == 1 + assert rv.log_rerun_data.blueprint is first_blueprint diff --git a/tests/utils/test_visualization_utils.py b/tests/utils/test_visualization_utils.py index f62a697cd..8dfc05edf 100644 --- a/tests/utils/test_visualization_utils.py +++ b/tests/utils/test_visualization_utils.py @@ -14,297 +14,23 @@ # See the License for the specific language governing permissions and # limitations under the License. -import importlib -import sys -from types import SimpleNamespace +"""Tests for the backend-agnostic visualization dispatch. + +These exercise the display-mode routing/validation only; they need neither ``rerun`` nor +``foxglove`` installed since the unknown-mode branch raises before touching any backend. Backend +behavior is covered in ``test_rerun_visualization.py`` and ``test_foxglove_visualization.py``. +""" -import numpy as np import pytest -pytest.importorskip("rerun", reason="rerun-sdk is required (install lerobot[viz])") - -from lerobot.types import TransitionKey -from lerobot.utils.constants import OBS_STATE +from lerobot.utils import visualization_utils as vu -@pytest.fixture -def mock_rerun(monkeypatch): - """ - Provide a mock `rerun` module (and `rerun.blueprint` submodule) so tests don't - depend on the real library. Also reload the module-under-test so it binds to - this mock `rr`. - """ - calls = [] - blueprints = [] - - class DummyScalar: - def __init__(self, value): - # Scalars may be built from a single float or from a 1D array batch. - self.value = value - - class DummyImage: - def __init__(self, arr): - self.arr = arr - - def compress(self, *a, **k): - return self - - class DummyDepthImage: - def __init__(self, arr, colormap=None): - self.arr = arr - self.colormap = colormap - - def dummy_log(key, obj=None, **kwargs): - # Accept either positional `obj` or keyword `entity` and record remaining kwargs. - if obj is None and "entity" in kwargs: - obj = kwargs.pop("entity") - calls.append((key, obj, kwargs)) - - def dummy_send_blueprint(blueprint, *a, **k): - blueprints.append(blueprint) - - # Mock the `rerun.blueprint` submodule used to build the layout. - dummy_rrb = SimpleNamespace( - Spatial2DView=lambda origin=None, name=None: SimpleNamespace( - kind="Spatial2DView", origin=origin, name=name - ), - TimeSeriesView=lambda name=None, contents=None: SimpleNamespace( - kind="TimeSeriesView", name=name, contents=contents - ), - Grid=lambda *views: SimpleNamespace(kind="Grid", views=list(views)), - Blueprint=lambda root: SimpleNamespace(kind="Blueprint", root=root), - ) - - dummy_rr = SimpleNamespace( - __name__="rerun", - __package__="rerun", - __spec__=SimpleNamespace(name="rerun", submodule_search_locations=None), - Scalars=DummyScalar, - Image=DummyImage, - DepthImage=DummyDepthImage, - components=SimpleNamespace(Colormap=SimpleNamespace(Viridis="viridis")), - log=dummy_log, - send_blueprint=dummy_send_blueprint, - init=lambda *a, **k: None, - spawn=lambda *a, **k: None, - blueprint=dummy_rrb, - ) - - # Inject fake modules into sys.modules (both `rerun` and `rerun.blueprint`). - monkeypatch.setitem(sys.modules, "rerun", dummy_rr) - monkeypatch.setitem(sys.modules, "rerun.blueprint", dummy_rrb) - - # Now import and reload the module under test, to bind to our rerun mock - import lerobot.utils.visualization_utils as vu - - importlib.reload(vu) - - # Expose the reloaded module, the call recorder and the captured blueprints - yield vu, calls, blueprints +def test_visualization_modes(): + assert vu.VISUALIZATION_MODES == ("rerun", "foxglove") -def _keys(calls): - """Helper to extract just the keys logged to rr.log""" - return [k for (k, _obj, _kw) in calls] - - -def _obj_for(calls, key): - """Find the first object logged under a given key.""" - for k, obj, _kw in calls: - if k == key: - return obj - raise KeyError(f"Key {key} not found in calls: {calls}") - - -def _kwargs_for(calls, key): - for k, _obj, kw in calls: - if k == key: - return kw - raise KeyError(f"Key {key} not found in calls: {calls}") - - -def _views_by_kind(blueprint, kind): - """Return the views of a given kind from the (single) blueprint's grid.""" - return [v for v in blueprint.root.views if v.kind == kind] - - -def test_log_rerun_data_envtransition_scalars_and_image(mock_rerun): - vu, calls, blueprints = mock_rerun - - # Build EnvTransition dict - obs = { - f"{OBS_STATE}.temperature": np.float32(25.0), - # CHW image should be converted to HWC for rr.Image - "observation.camera": np.zeros((3, 10, 20), dtype=np.uint8), - } - act = { - "action.throttle": 0.7, - # 1D array should be logged as a single Scalars batch under one entity path - "action.vector": np.array([1.0, 2.0], dtype=np.float32), - } - transition = { - TransitionKey.OBSERVATION: obs, - TransitionKey.ACTION: act, - } - - # Extract observation and action data from transition like in the real call sites - obs_data = transition.get(TransitionKey.OBSERVATION, {}) - action_data = transition.get(TransitionKey.ACTION, {}) - vu.log_rerun_data(observation=obs_data, action=action_data) - - # We expect: - # - observation.state.temperature -> Scalars - # - observation.camera -> Image (HWC) with static=True - # - action.throttle -> Scalars - # - action.vector -> single Scalars batch (no per-element suffix) - expected_keys = { - f"{OBS_STATE}.temperature", - "observation.camera", - "action.throttle", - "action.vector", - } - assert set(_keys(calls)) == expected_keys - - # Check scalar types and values - temp_obj = _obj_for(calls, f"{OBS_STATE}.temperature") - assert type(temp_obj).__name__ == "DummyScalar" - assert float(temp_obj.value) == pytest.approx(25.0) - - throttle_obj = _obj_for(calls, "action.throttle") - assert type(throttle_obj).__name__ == "DummyScalar" - assert float(throttle_obj.value) == pytest.approx(0.7) - - # 1D vector logged as a single batched Scalars under one entity path - vec = _obj_for(calls, "action.vector") - assert type(vec).__name__ == "DummyScalar" - np.testing.assert_allclose(np.asarray(vec.value), [1.0, 2.0]) - - # Check image handling: CHW -> HWC - img_obj = _obj_for(calls, "observation.camera") - assert type(img_obj).__name__ == "DummyImage" - assert img_obj.arr.shape == (10, 20, 3) # transposed - assert _kwargs_for(calls, "observation.camera").get("static", False) is True # static=True for images - - # A blueprint should have been built and sent exactly once, and cached on the function. - assert len(blueprints) == 1 - assert vu.log_rerun_data.blueprint is blueprints[0] - - bp = blueprints[0] - # One spatial view per image path - spatial_views = _views_by_kind(bp, "Spatial2DView") - assert {v.origin for v in spatial_views} == {"observation.camera"} - - # One time-series view each for observation and action scalars - ts_views = {v.name: v for v in _views_by_kind(bp, "TimeSeriesView")} - assert set(ts_views) == {"observation", "action"} - assert ts_views["observation"].contents == [f"{OBS_STATE}.temperature"] - assert ts_views["action"].contents == ["action.throttle", "action.vector"] - - -def test_log_rerun_data_plain_list_ordering_and_prefixes(mock_rerun): - vu, calls, blueprints = mock_rerun - - # First dict without prefixes treated as observation - # Second dict without prefixes treated as action - obs_plain = { - "temp": 1.5, - # Already HWC image => should stay as-is - "img": np.zeros((5, 6, 3), dtype=np.uint8), - "none": None, # should be skipped - } - act_plain = { - "throttle": 0.3, - "vec": np.array([9, 8, 7], dtype=np.float32), - } - - # Extract observation and action data from list like the old function logic did - # First dict was treated as observation, second as action - vu.log_rerun_data(observation=obs_plain, action=act_plain) - - # Expected keys with auto-prefixes. The 1D vector is a single batched Scalars. - expected = { - "observation.temp", - "observation.img", - "action.throttle", - "action.vec", - } - logged = set(_keys(calls)) - assert logged == expected - - # Scalars - t = _obj_for(calls, "observation.temp") - assert type(t).__name__ == "DummyScalar" - assert float(t.value) == pytest.approx(1.5) - - throttle = _obj_for(calls, "action.throttle") - assert type(throttle).__name__ == "DummyScalar" - assert float(throttle.value) == pytest.approx(0.3) - - # Image stays HWC - img = _obj_for(calls, "observation.img") - assert type(img).__name__ == "DummyImage" - assert img.arr.shape == (5, 6, 3) - assert _kwargs_for(calls, "observation.img").get("static", False) is True - - # Vector logged as a single batched Scalars under one entity path - vec = _obj_for(calls, "action.vec") - assert type(vec).__name__ == "DummyScalar" - np.testing.assert_allclose(np.asarray(vec.value), [9, 8, 7]) - - # Blueprint sent once with the expected view layout - assert len(blueprints) == 1 - bp = blueprints[0] - spatial_views = _views_by_kind(bp, "Spatial2DView") - assert {v.origin for v in spatial_views} == {"observation.img"} - ts_views = {v.name: v for v in _views_by_kind(bp, "TimeSeriesView")} - assert ts_views["observation"].contents == ["observation.temp"] - assert ts_views["action"].contents == ["action.throttle", "action.vec"] - - -def test_log_rerun_data_kwargs_only(mock_rerun): - vu, calls, blueprints = mock_rerun - - vu.log_rerun_data( - observation={"observation.temp": 10.0, "observation.gray": np.zeros((8, 8, 1), dtype=np.uint8)}, - action={"action.a": 1.0}, - ) - - keys = set(_keys(calls)) - assert "observation.temp" in keys - assert "observation.gray" in keys - assert "action.a" in keys - - temp = _obj_for(calls, "observation.temp") - assert type(temp).__name__ == "DummyScalar" - assert float(temp.value) == pytest.approx(10.0) - - img = _obj_for(calls, "observation.gray") - assert type(img).__name__ == "DummyDepthImage" # single-channel -> DepthImage - assert img.arr.shape == (8, 8, 1) # remains HWC - assert _kwargs_for(calls, "observation.gray").get("static", False) is True - - a = _obj_for(calls, "action.a") - assert type(a).__name__ == "DummyScalar" - assert float(a.value) == pytest.approx(1.0) - - # Blueprint sent once, with a spatial view for the image and time-series views for scalars - assert len(blueprints) == 1 - bp = blueprints[0] - assert {v.origin for v in _views_by_kind(bp, "Spatial2DView")} == {"observation.gray"} - ts_views = {v.name: v for v in _views_by_kind(bp, "TimeSeriesView")} - assert ts_views["observation"].contents == ["observation.temp"] - assert ts_views["action"].contents == ["action.a"] - - -def test_log_rerun_data_blueprint_sent_only_once(mock_rerun): - """The blueprint is built from the first call and not resent on subsequent calls.""" - vu, calls, blueprints = mock_rerun - - vu.log_rerun_data(observation={"temp": 1.0}, action={"a": 2.0}) - assert len(blueprints) == 1 - first_blueprint = vu.log_rerun_data.blueprint - - vu.log_rerun_data(observation={"temp": 3.0}, action={"a": 4.0}) - # Still only one blueprint, and the cached one is unchanged. - assert len(blueprints) == 1 - assert vu.log_rerun_data.blueprint is first_blueprint +@pytest.mark.parametrize("func", ["init_visualization", "log_visualization_data", "shutdown_visualization"]) +def test_dispatch_rejects_unknown_mode(func): + with pytest.raises(ValueError, match="Unknown display_mode"): + getattr(vu, func)("bogus") diff --git a/uv.lock b/uv.lock index 076d021f2..22afec6a6 100644 --- a/uv.lock +++ b/uv.lock @@ -1550,6 +1550,26 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/2c/47/c99d5268f354002ce80f8d029cd9d7d872969da1de8b93d32de4dc56d6f4/fonttools-4.63.0-py3-none-any.whl", hash = "sha256:445af2eab030a16b9171ea8bdda7ebf7d96bda2df88ee182a464252f6e05e20d", size = 1164562, upload-time = "2026-05-14T12:04:29.092Z" }, ] +[[package]] +name = "foxglove-sdk" +version = "0.25.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/c1/a7/86a252782ea0d9baf1357369ad1bbf1ed644768702b0266a3fa3a05361d0/foxglove_sdk-0.25.1.tar.gz", hash = "sha256:8230f3c32ea3ab715818687377491594ec9c7e58e6b0ed8ed91aadf937ce706b", size = 547778, upload-time = "2026-06-02T03:13:18.942Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/58/15/59f02e8201b8da09ce05d8774820c29efc9149862b70ee6b3a27968e791a/foxglove_sdk-0.25.1-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:5af9f9a691eefbe6e0a47875ff2f7d0fc36607f0920e8690bbdc2dfd4fb22451", size = 17911538, upload-time = "2026-06-02T03:13:12.493Z" }, + { url = "https://files.pythonhosted.org/packages/27/ed/16d809fab24cbfdf97c15c9cdd80eabfeb447ca545ede426950d62bac848/foxglove_sdk-0.25.1-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:3e908bd87d1926a05c785779d8252db6b87eef685f284ec1cf46ee501645d08e", size = 16452309, upload-time = "2026-06-02T03:13:10.607Z" }, + { url = "https://files.pythonhosted.org/packages/d6/c3/f95874935a3436841487df1f0202de4d20eabc0adb6b79c94c531bbe7eb3/foxglove_sdk-0.25.1-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:968e32c8668d172f6b546c8e7af658ed35a21ec165adc3bacf53a04dda159f12", size = 2355680, upload-time = "2026-06-02T02:34:01.668Z" }, + { url = "https://files.pythonhosted.org/packages/38/da/ad22d8d6e3fedde9fc0c49aa8b20394e5e0bc44ab3fba564c77a64ddc7e2/foxglove_sdk-0.25.1-cp310-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:3f75374fedafe259c40b19bc645589d9453708eab679a5b07c603035f936d29a", size = 2274075, upload-time = "2026-06-02T02:34:07.212Z" }, + { url = "https://files.pythonhosted.org/packages/a3/fa/1254adb5e72eff507695473e9c82d0e90395b61463e5353762250db30d3d/foxglove_sdk-0.25.1-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7b6d3af517a00342bf7b08a4a65b043f3eafaa197138752b6fbd704fb91043fa", size = 2282160, upload-time = "2026-06-02T02:34:08.812Z" }, + { url = "https://files.pythonhosted.org/packages/7c/e4/2b22ef06ba4058494c7aa35974d138f8f1ae4cf5273f77d69c9dc3a99b45/foxglove_sdk-0.25.1-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:aed27c0f03a45fd6abdd566498bfee2672391602bcff32c827b8e3a6d8f67ab1", size = 22685338, upload-time = "2026-06-02T02:34:04.688Z" }, + { url = "https://files.pythonhosted.org/packages/35/7c/58324c99b80eef0b674c8d4f5c2e07c66fd1480a27a8f0d4d79371805111/foxglove_sdk-0.25.1-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:419dd8308e3f91e2ae487b727f1bf1804642990876163b2a353db4a1b1de1425", size = 19326096, upload-time = "2026-06-02T02:34:10.939Z" }, + { url = "https://files.pythonhosted.org/packages/fe/9c/3452d92959e05fc6b1c1e5f032605d55623aeb6704357d20408f8781bc84/foxglove_sdk-0.25.1-cp310-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:0fcb36e628ab3d9043e193f12ad4dbbb955fe18616aac7ef5bca82c52910f108", size = 2539020, upload-time = "2026-06-02T03:13:14.365Z" }, + { url = "https://files.pythonhosted.org/packages/b5/af/57fa58525d3acb5c5480a6f0ef86450b1a0ccae2b21248edb1376073ce55/foxglove_sdk-0.25.1-cp310-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:7909fd9f94935935dd8813702d84ffdbfebeb3866673c618ce35e8cfedd03029", size = 2550999, upload-time = "2026-06-02T03:13:15.715Z" }, + { url = "https://files.pythonhosted.org/packages/90/78/f74bb167186c965d475ff360fa6eb7441d5ac6c6239d60f542f63984f849/foxglove_sdk-0.25.1-cp310-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:69d5966213b5212b8841b4004fe582db924a74f1610d8452ad890f6931702926", size = 2560166, upload-time = "2026-06-02T03:13:17.254Z" }, + { url = "https://files.pythonhosted.org/packages/81/83/1c4c6d04fbd4784fe44fb2da021db1adf1f03a371f1e5679a383c1173235/foxglove_sdk-0.25.1-cp310-abi3-win32.whl", hash = "sha256:2a1121a5c74590ff6e61628c4a46dc57d392d290b4beeb29d6852933da56224a", size = 1618124, upload-time = "2026-06-02T03:13:20.158Z" }, + { url = "https://files.pythonhosted.org/packages/5f/4d/bdb9e252a41a951eb53908ac9cb965b7480c3ba649174f5398d4fcf0ca1d/foxglove_sdk-0.25.1-cp310-abi3-win_amd64.whl", hash = "sha256:6ed3ad0d3e72cd7875e7e293709c5ff90494fe14f1b48a336baffc313a7272cc", size = 16588452, upload-time = "2026-06-02T03:13:21.636Z" }, +] + [[package]] name = "fqdn" version = "1.5.1" @@ -2811,6 +2831,7 @@ all = [ { name = "faker" }, { name = "fastapi" }, { name = "feetech-servo-sdk" }, + { name = "foxglove-sdk" }, { name = "grpcio" }, { name = "grpcio-tools" }, { name = "gym-aloha" }, @@ -2895,6 +2916,7 @@ core-scripts = [ { name = "av" }, { name = "datasets" }, { name = "deepdiff" }, + { name = "foxglove-sdk" }, { name = "jsonlines" }, { name = "pandas" }, { name = "pyarrow" }, @@ -2917,6 +2939,7 @@ dataset = [ dataset-viz = [ { name = "av" }, { name = "datasets" }, + { name = "foxglove-sdk" }, { name = "jsonlines" }, { name = "pandas" }, { name = "pyarrow" }, @@ -3187,6 +3210,7 @@ video-benchmark = [ { name = "scikit-image" }, ] viz = [ + { name = "foxglove-sdk" }, { name = "rerun-sdk" }, ] vla-jepa = [ @@ -3226,6 +3250,7 @@ requires-dist = [ { name = "fastapi", marker = "extra == 'phone'", specifier = "<1.0" }, { name = "feetech-servo-sdk", marker = "extra == 'feetech'", specifier = ">=1.0.0,<2.0.0" }, { name = "flash-attn", marker = "sys_platform != 'darwin' and extra == 'groot'", specifier = ">=2.5.9,<3.0.0" }, + { name = "foxglove-sdk", marker = "extra == 'viz'", specifier = ">=0.25.1,<0.26.0" }, { name = "grpcio", marker = "extra == 'grpcio-dep'", specifier = ">=1.73.1,<2.0.0" }, { name = "grpcio", marker = "extra == 'reachy2'", specifier = "<=1.73.1" }, { name = "grpcio-tools", marker = "extra == 'dev'", specifier = ">=1.73.1,<2.0.0" }, From b961d2a8c52d908239db3b082ebf4dc299f3e9ba Mon Sep 17 00:00:00 2001 From: Caroline Pascal Date: Thu, 2 Jul 2026 11:03:41 +0200 Subject: [PATCH 28/38] feat(libaom-av1): adding support for libaom-av1 codec (#3898) --- src/lerobot/configs/video.py | 10 +++++++++- tests/datasets/test_datasets.py | 3 ++- 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/src/lerobot/configs/video.py b/src/lerobot/configs/video.py index 3ea834508..e253526f6 100644 --- a/src/lerobot/configs/video.py +++ b/src/lerobot/configs/video.py @@ -36,7 +36,9 @@ HW_VIDEO_CODECS = [ "h264_vaapi", # Linux Intel/AMD "h264_qsv", # Intel Quick Sync ] -VALID_VIDEO_CODECS: frozenset[str] = frozenset({"h264", "hevc", "libsvtav1", "auto", *HW_VIDEO_CODECS}) +VALID_VIDEO_CODECS: frozenset[str] = frozenset( + {"h264", "hevc", "libsvtav1", "libaom-av1", "auto", *HW_VIDEO_CODECS} +) # Aliases for legacy video codec names. VIDEO_CODECS_ALIASES: dict[str, str] = {"av1": "libsvtav1"} @@ -220,6 +222,12 @@ class VideoEncoderConfig: if self.fast_decode: opts["tune"] = "fastdecode" set_if("threads", encoder_threads) + elif self.vcodec == "libaom-av1": + set_if("crf", self.crf) + set_if("preset", self.preset) + if encoder_threads is not None: + opts["threads"] = encoder_threads + opts["row-mt"] = 1 elif self.vcodec in ("h264_videotoolbox", "hevc_videotoolbox"): if self.crf is not None: opts["q:v"] = max(1, min(100, 100 - self.crf * 2)) diff --git a/tests/datasets/test_datasets.py b/tests/datasets/test_datasets.py index 225479814..97f799d9f 100644 --- a/tests/datasets/test_datasets.py +++ b/tests/datasets/test_datasets.py @@ -1531,6 +1531,7 @@ def test_valid_video_codecs_constant(): assert "h264" in VALID_VIDEO_CODECS assert "hevc" in VALID_VIDEO_CODECS assert "libsvtav1" in VALID_VIDEO_CODECS + assert "libaom-av1" in VALID_VIDEO_CODECS assert "auto" in VALID_VIDEO_CODECS assert "h264_videotoolbox" in VALID_VIDEO_CODECS assert "h264_nvenc" in VALID_VIDEO_CODECS @@ -1538,7 +1539,7 @@ def test_valid_video_codecs_constant(): assert "h264_qsv" in VALID_VIDEO_CODECS assert "hevc_videotoolbox" in VALID_VIDEO_CODECS assert "hevc_nvenc" in VALID_VIDEO_CODECS - assert len(VALID_VIDEO_CODECS) == 10 + assert len(VALID_VIDEO_CODECS) == 11 def test_delta_timestamps_with_episodes_filter(tmp_path, empty_lerobot_dataset_factory): From c746ca2df2724afc677644160c048a1beb5d858a Mon Sep 17 00:00:00 2001 From: Caroline Pascal Date: Thu, 2 Jul 2026 11:53:13 +0200 Subject: [PATCH 29/38] fix(depth unit): adding input depth unit storage in the dataset metadata (#3899) * fix(depth unit): storing raw depth units in the dataset metadata for correct depth statistics and depth raw frames handling. The unit is stored as a string ("m","mm") under "depth_unit" at the same level as "is_depth_map". Unit is inferred from the depth frame type. * feat(raw frame unit): adapting dataset reader so that raw depth frames are scaled according to the requested unit * feat(stats units): rescaling stats when loading a dataset so that the stats are given in the requested unit * tests(unit): adapting and extending depth tests to units manipulations * chore(format): formating code * feat(warning): adding a warning when depth unit is not specified in the dataset * chore(infer_depth_unit): moving the depth unit inference utility in a more accessible location * feat(rerun unit): adding correct depth unit display for rerun (foxglove does not support units yet) * feat(unit getter): adding a proper output_depth_unit getter to LeRobotDataset for cleaner integration * fix(streaming dataset): extending support for depth units to streaming datasets * test(rerun): fixing rerun tests --- src/lerobot/configs/__init__.py | 6 ++ src/lerobot/configs/video.py | 11 +++ src/lerobot/datasets/compute_stats.py | 2 +- src/lerobot/datasets/dataset_metadata.py | 32 +++++++- src/lerobot/datasets/dataset_reader.py | 22 +++++- src/lerobot/datasets/dataset_writer.py | 10 +++ src/lerobot/datasets/depth_utils.py | 19 ++--- src/lerobot/datasets/lerobot_dataset.py | 6 ++ src/lerobot/datasets/streaming_dataset.py | 26 ++++++- src/lerobot/scripts/lerobot_dataset_viz.py | 5 ++ src/lerobot/utils/rerun_visualization.py | 9 ++- tests/datasets/test_depth.py | 89 ++++++++++++++++++++++ tests/fixtures/dataset_factories.py | 8 ++ tests/utils/test_rerun_visualization.py | 3 +- 14 files changed, 229 insertions(+), 19 deletions(-) diff --git a/src/lerobot/configs/__init__.py b/src/lerobot/configs/__init__.py index 168b367db..c32e3368b 100644 --- a/src/lerobot/configs/__init__.py +++ b/src/lerobot/configs/__init__.py @@ -34,6 +34,8 @@ from .types import ( ) from .video import ( DEFAULT_DEPTH_UNIT, + DEPTH_METER_UNIT, + DEPTH_MILLIMETER_UNIT, VALID_VIDEO_CODECS, VIDEO_ENCODER_INFO_KEYS, DepthEncoderConfig, @@ -41,6 +43,7 @@ from .video import ( VideoEncoderConfig, depth_encoder_defaults, encoder_config_from_video_info, + infer_depth_unit, rgb_encoder_defaults, ) @@ -70,8 +73,11 @@ __all__ = [ "depth_encoder_defaults", # Factories "encoder_config_from_video_info", + "infer_depth_unit", # Constants "DEFAULT_DEPTH_UNIT", + "DEPTH_METER_UNIT", + "DEPTH_MILLIMETER_UNIT", "VALID_VIDEO_CODECS", "VIDEO_ENCODER_INFO_KEYS", ] diff --git a/src/lerobot/configs/video.py b/src/lerobot/configs/video.py index e253526f6..20e46a387 100644 --- a/src/lerobot/configs/video.py +++ b/src/lerobot/configs/video.py @@ -22,6 +22,8 @@ import logging from dataclasses import dataclass, field from typing import Any, ClassVar, Self +import numpy as np + from lerobot.utils.import_utils import require_package logger = logging.getLogger(__name__) @@ -67,6 +69,15 @@ DEPTH_METER_UNIT: str = "m" DEPTH_MILLIMETER_UNIT: str = "mm" DEFAULT_DEPTH_UNIT: str = DEPTH_MILLIMETER_UNIT + +def infer_depth_unit(dtype: np.dtype | type) -> str: + """Infer the physical unit of raw depth frames from their dtype. + + Floating-point frames are assumed to be in metres, integer frames in millimetres. + """ + return DEPTH_METER_UNIT if np.issubdtype(np.dtype(dtype), np.floating) else DEPTH_MILLIMETER_UNIT + + # Depth-specific tuning fields persisted under ``features[*]["info"]`` as ``video.``. DEPTH_ENCODER_INFO_FIELD_NAMES: frozenset[str] = frozenset({"depth_min", "depth_max", "shift", "use_log"}) diff --git a/src/lerobot/datasets/compute_stats.py b/src/lerobot/datasets/compute_stats.py index 88f7ea226..02ecd81a4 100644 --- a/src/lerobot/datasets/compute_stats.py +++ b/src/lerobot/datasets/compute_stats.py @@ -509,7 +509,7 @@ def compute_episode_stats( For 'image'/'video' features, stats are computed per channel and kept with a leading channel axis (e.g. shape (3, 1, 1) for RGB). RGB stats are divided by 255 to land in [0, 1]; depth maps (features flagged with ``is_depth_map``) skip - this rescaling and remain in their stored units. + this rescaling and remain in their stored units (stored in ``depth_unit``). """ if quantile_list is None: quantile_list = DEFAULT_QUANTILES diff --git a/src/lerobot/datasets/dataset_metadata.py b/src/lerobot/datasets/dataset_metadata.py index ea329668c..6e19d14fb 100644 --- a/src/lerobot/datasets/dataset_metadata.py +++ b/src/lerobot/datasets/dataset_metadata.py @@ -26,12 +26,13 @@ import pyarrow as pa import pyarrow.parquet as pq from huggingface_hub import snapshot_download -from lerobot.configs import VideoEncoderConfig +from lerobot.configs import DEPTH_METER_UNIT, VideoEncoderConfig from lerobot.utils.constants import DEFAULT_FEATURES, HF_LEROBOT_HOME, HF_LEROBOT_HUB_CACHE from lerobot.utils.feature_utils import _validate_feature_names from lerobot.utils.utils import flatten_dict from .compute_stats import aggregate_stats +from .depth_utils import MM_PER_METRE from .feature_utils import create_empty_dataset_info from .io_utils import ( get_file_size_in_mb, @@ -358,6 +359,35 @@ class LeRobotDatasetMetadata: return [key for key, ft in self.features.items() if _is_depth(ft)] + def rescale_depth_stats(self, output_unit: str) -> None: + """Rescale depth feature stats in place from their recorded unit to ``output_unit``. + + Depth stats are stored in the unit the frames were recorded in + (``features[key]["info"]["depth_unit"]``), while frames are returned in + ``output_unit`` on read. This converts the unit-bearing stat entries so + stats match the frames consumers see. + """ + missing_unit_keys = [ + key for key in self.depth_keys if (self.features[key].get("info") or {}).get("depth_unit") is None + ] + if missing_unit_keys: + logging.warning( + f"Depth feature(s) {missing_unit_keys} have no recorded 'depth_unit' in their info. " + f"Depth maps and stats for these keys will be returned AS IS, with no unit conversion " + f"to the requested output unit {output_unit!r}. Re-record the dataset or set 'depth_unit' " + f"in the feature info (meta/info.json) to enable conversion." + ) + if self.stats is None: + return + for key in self.depth_keys: + stored_unit = (self.features[key].get("info") or {}).get("depth_unit") + if stored_unit is None or stored_unit == output_unit or key not in self.stats: + continue + factor = MM_PER_METRE if stored_unit == DEPTH_METER_UNIT else 1.0 / MM_PER_METRE + self.stats[key] = { + stat: value if stat == "count" else value * factor for stat, value in self.stats[key].items() + } + @property def camera_keys(self) -> list[str]: """Keys to access visual modalities (regardless of their storage method).""" diff --git a/src/lerobot/datasets/dataset_reader.py b/src/lerobot/datasets/dataset_reader.py index e8e07301e..f4e1f6a31 100644 --- a/src/lerobot/datasets/dataset_reader.py +++ b/src/lerobot/datasets/dataset_reader.py @@ -22,10 +22,14 @@ from pathlib import Path import datasets import torch -from lerobot.configs import DEFAULT_DEPTH_UNIT, DepthEncoderConfig +from lerobot.configs import ( + DEFAULT_DEPTH_UNIT, + DEPTH_METER_UNIT, + DepthEncoderConfig, +) from .dataset_metadata import LeRobotDatasetMetadata -from .depth_utils import dequantize_depth +from .depth_utils import MM_PER_METRE, dequantize_depth from .feature_utils import ( check_delta_timestamps, get_delta_indices, @@ -102,6 +106,13 @@ class DatasetReader: for vid_key in self._meta.depth_keys } + # Get the input unit of each depth feature stored as raw images. + self._image_depth_units: dict[str, str | None] = { + key: (self._meta.features[key].get("info") or {}).get("depth_unit") + for key in self._meta.depth_keys + if key in self._meta.image_keys + } + def set_image_transforms(self, image_transforms: Callable | None) -> None: """Replace the transform applied to visual observations.""" if image_transforms is not None and not callable(image_transforms): @@ -329,6 +340,13 @@ class DatasetReader: continue item[cam] = self._image_transforms(item[cam]) + # Convert depth features to the output unit. + for key, stored_unit in self._image_depth_units.items(): + if key in item and stored_unit is not None and stored_unit != self._depth_output_unit: + item[key] = ( + item[key] * MM_PER_METRE if stored_unit == DEPTH_METER_UNIT else item[key] / MM_PER_METRE + ) + # Add task as a string task_idx = item["task_index"].item() item["task"] = self._meta.tasks.iloc[task_idx].name diff --git a/src/lerobot/datasets/dataset_writer.py b/src/lerobot/datasets/dataset_writer.py index 1aee1497c..a6049312f 100644 --- a/src/lerobot/datasets/dataset_writer.py +++ b/src/lerobot/datasets/dataset_writer.py @@ -36,6 +36,7 @@ from lerobot.configs import ( RGBEncoderConfig, VideoEncoderConfig, depth_encoder_defaults, + infer_depth_unit, rgb_encoder_defaults, ) @@ -209,6 +210,15 @@ class DatasetWriter: self.episode_buffer["timestamp"].append(timestamp) self.episode_buffer["task"].append(frame.pop("task")) + # Record each depth feature's input unit once, inferred from the first frame's dtype. + if frame_index == 0: + for depth_key in self._meta.depth_keys: + if depth_key not in frame: + continue + info = self._meta.features[depth_key].setdefault("info", {}) + if info.get("depth_unit") is None: + info["depth_unit"] = infer_depth_unit(np.asarray(frame[depth_key]).dtype) + # Start streaming encoder on first frame of episode if frame_index == 0 and self._streaming_encoder is not None: self._streaming_encoder.start_episode( diff --git a/src/lerobot/datasets/depth_utils.py b/src/lerobot/datasets/depth_utils.py index 801c86a09..a4e187eb4 100644 --- a/src/lerobot/datasets/depth_utils.py +++ b/src/lerobot/datasets/depth_utils.py @@ -34,12 +34,13 @@ from lerobot.configs.video import ( DEPTH_METER_UNIT, DEPTH_MILLIMETER_UNIT, DEPTH_QMAX, + infer_depth_unit, ) from .image_writer import squeeze_single_channel from .pyav_utils import write_u16_plane -_MM_PER_METRE = 1000.0 +MM_PER_METRE = 1000.0 _UINT16_MAX = 65535 @@ -57,11 +58,7 @@ def _depth_input_to_float32_and_unit( input_unit: Literal["auto", DEPTH_METER_UNIT, DEPTH_MILLIMETER_UNIT], ) -> tuple[NDArray[np.float32], Literal[DEPTH_METER_UNIT, DEPTH_MILLIMETER_UNIT]]: """Convert depth to float32 in the chosen unit, and return the resolved unit.""" - resolved_unit = ( - (DEPTH_METER_UNIT if np.issubdtype(depth.dtype, np.floating) else DEPTH_MILLIMETER_UNIT) - if input_unit == "auto" - else input_unit - ) + resolved_unit = infer_depth_unit(depth.dtype) if input_unit == "auto" else input_unit return depth.astype(np.float32, order="K"), resolved_unit @@ -126,12 +123,12 @@ def quantize_depth( # Convert depth_min, depth_max, and shift to the resolved input unit. depth_min_u = ( - np.float32(depth_min) if resolved_unit == DEPTH_METER_UNIT else np.float32(depth_min * _MM_PER_METRE) + np.float32(depth_min) if resolved_unit == DEPTH_METER_UNIT else np.float32(depth_min * MM_PER_METRE) ) depth_max_u = ( - np.float32(depth_max) if resolved_unit == DEPTH_METER_UNIT else np.float32(depth_max * _MM_PER_METRE) + np.float32(depth_max) if resolved_unit == DEPTH_METER_UNIT else np.float32(depth_max * MM_PER_METRE) ) - shift_u = np.float32(shift) if resolved_unit == DEPTH_METER_UNIT else np.float32(shift * _MM_PER_METRE) + shift_u = np.float32(shift) if resolved_unit == DEPTH_METER_UNIT else np.float32(shift * MM_PER_METRE) # Normalization and quantization is performed in the resolved input unit. if use_log: @@ -236,7 +233,7 @@ def dequantize_depth( # mm path: round + clamp in float32, skipping the uint16 round-trip # when returning a tensor (torch.uint16 is poorly supported). - buf.mul_(_MM_PER_METRE).round_().clamp_(0.0, _UINT16_MAX) + buf.mul_(MM_PER_METRE).round_().clamp_(0.0, _UINT16_MAX) if output_tensor: return buf return buf.cpu().numpy().astype(np.uint16, copy=False) @@ -259,7 +256,7 @@ def dequantize_depth( if output_unit == DEPTH_METER_UNIT: return torch.from_numpy(buf) if output_tensor else buf - np.multiply(buf, _MM_PER_METRE, out=buf) + np.multiply(buf, MM_PER_METRE, out=buf) np.rint(buf, out=buf) np.clip(buf, 0.0, _UINT16_MAX, out=buf) if output_tensor: diff --git a/src/lerobot/datasets/lerobot_dataset.py b/src/lerobot/datasets/lerobot_dataset.py index f600f1804..aba86efe3 100644 --- a/src/lerobot/datasets/lerobot_dataset.py +++ b/src/lerobot/datasets/lerobot_dataset.py @@ -224,6 +224,7 @@ class LeRobotDataset(torch.utils.data.Dataset): ) self.root = self.meta.root self.revision = self.meta.revision + self.meta.rescale_depth_stats(self._depth_output_unit) if episodes is not None and any( episode >= self.meta.total_episodes or episode < 0 for episode in episodes @@ -350,6 +351,11 @@ class LeRobotDataset(torch.utils.data.Dataset): """Frames per second used during data collection.""" return self.meta.fps + @property + def depth_output_unit(self) -> str: + """Physical unit (``"m"`` or ``"mm"``) depth maps and statistics are returned in on read.""" + return self._depth_output_unit + @property def num_frames(self) -> int: """Number of frames in selected episodes.""" diff --git a/src/lerobot/datasets/streaming_dataset.py b/src/lerobot/datasets/streaming_dataset.py index 4c4ae59bf..14d4a52a4 100644 --- a/src/lerobot/datasets/streaming_dataset.py +++ b/src/lerobot/datasets/streaming_dataset.py @@ -22,11 +22,11 @@ import numpy as np import torch from datasets import load_dataset -from lerobot.configs import DEFAULT_DEPTH_UNIT, DepthEncoderConfig +from lerobot.configs import DEFAULT_DEPTH_UNIT, DEPTH_METER_UNIT, DepthEncoderConfig from lerobot.utils.constants import HF_LEROBOT_HOME, LOOKAHEAD_BACKTRACKTABLE, LOOKBACK_BACKTRACKTABLE from .dataset_metadata import CODEBASE_VERSION, LeRobotDatasetMetadata -from .depth_utils import dequantize_depth +from .depth_utils import MM_PER_METRE, dequantize_depth from .feature_utils import get_delta_indices from .io_utils import item_to_torch from .utils import ( @@ -310,6 +310,7 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): ) self.root = self.meta.root self.revision = self.meta.revision + self.meta.rescale_depth_stats(self._depth_output_unit) # Check version check_version_compatibility(self.repo_id, self.meta._version, CODEBASE_VERSION) @@ -318,6 +319,13 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): for vid_key in self.meta.depth_keys } + # Input unit of each depth feature stored as raw images (dequantized separately from videos). + self._image_depth_units: dict[str, str | None] = { + key: (self.meta.features[key].get("info") or {}).get("depth_unit") + for key in self.meta.depth_keys + if key in self.meta.image_keys + } + self.delta_timestamps = None self.delta_indices = None @@ -348,6 +356,11 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): def fps(self): return self.meta.fps + @property + def depth_output_unit(self) -> str: + """Physical unit (``"m"`` or ``"mm"``) depth maps are returned in on read.""" + return self._depth_output_unit + @staticmethod def _iter_random_indices( rng: np.random.Generator, buffer_size: int, random_batch_size=100 @@ -530,6 +543,15 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): for update in updates: result.update(update) + # Convert raw-image depth features to the output unit (video depth is already converted). + for key, stored_unit in self._image_depth_units.items(): + if key in result and stored_unit is not None and stored_unit != self._depth_output_unit: + result[key] = ( + result[key] * MM_PER_METRE + if stored_unit == DEPTH_METER_UNIT + else result[key] / MM_PER_METRE + ) + result["task"] = self.meta.tasks.iloc[item["task_index"]].name yield result diff --git a/src/lerobot/scripts/lerobot_dataset_viz.py b/src/lerobot/scripts/lerobot_dataset_viz.py index ee25583a0..f4be878ad 100644 --- a/src/lerobot/scripts/lerobot_dataset_viz.py +++ b/src/lerobot/scripts/lerobot_dataset_viz.py @@ -84,6 +84,7 @@ import torch import torch.utils.data import tqdm +from lerobot.configs import DEPTH_MILLIMETER_UNIT from lerobot.datasets import LeRobotDataset from lerobot.utils.constants import ACTION, DONE, OBS_STATE, REWARD, SUCCESS from lerobot.utils.utils import init_logging @@ -228,6 +229,9 @@ def visualize_dataset( logging.info("Logging to Rerun") + # Depth frames and stats are dequantized to the dataset's depth_output_unit on load. + depth_meter = 1000.0 if dataset.depth_output_unit == DEPTH_MILLIMETER_UNIT else 1.0 + # Use the dataset's q01/q99 depth statistics for robust depth range bounds depth_ranges = {} for key in dataset.meta.depth_keys: @@ -254,6 +258,7 @@ def visualize_dataset( depth = to_hwc_float32_numpy(batch[key][i]) depth_entity = rr.DepthImage( depth, + meter=depth_meter, colormap=rr.components.Colormap.Viridis, depth_range=depth_ranges.get(key), ) diff --git a/src/lerobot/utils/rerun_visualization.py b/src/lerobot/utils/rerun_visualization.py index af04b18f7..46f2c0b4b 100644 --- a/src/lerobot/utils/rerun_visualization.py +++ b/src/lerobot/utils/rerun_visualization.py @@ -24,6 +24,7 @@ import os import numpy as np +from lerobot.configs import DEPTH_MILLIMETER_UNIT, infer_depth_unit from lerobot.types import RobotAction, RobotObservation from .constants import ACTION, ACTION_PREFIX, OBS_PREFIX, OBS_STR @@ -161,7 +162,13 @@ def log_rerun_data( observation_paths.add(key) else: if arr.shape[-1] == 1: - img_entity = rr.DepthImage(arr, colormap=rr.components.Colormap.Viridis) + # At record time, the depth unit is inferred from the frame type. + depth_unit = infer_depth_unit(arr.dtype) + img_entity = rr.DepthImage( + arr, + meter=1000.0 if depth_unit == DEPTH_MILLIMETER_UNIT else 1.0, + colormap=rr.components.Colormap.Viridis, + ) else: img_entity = rr.Image(arr).compress() if compress_images else rr.Image(arr) rr.log(key, entity=img_entity, static=True) diff --git a/tests/datasets/test_depth.py b/tests/datasets/test_depth.py index a075fa6b5..5391cc558 100644 --- a/tests/datasets/test_depth.py +++ b/tests/datasets/test_depth.py @@ -32,6 +32,7 @@ from lerobot.configs.video import ( ) from lerobot.datasets.depth_utils import dequantize_depth, quantize_depth from lerobot.datasets.image_writer import image_array_to_pil_image, write_image +from lerobot.utils.constants import DEFAULT_FEATURES from tests.fixtures.constants import ( DEFAULT_FPS, DUMMY_CAMERA_FEATURES, @@ -245,3 +246,91 @@ class TestFeatureFileRouting: dataset.save_episode() dataset.finalize() + + +class TestDepthUnitMetadata: + """The depth unit is inferred once from dtype, stored in ``info``, and drives stats + reads.""" + + NUM_FRAMES = 4 + + def _record(self, root, features_factory, depth_dtype, value, use_videos): + from lerobot.datasets.lerobot_dataset import LeRobotDataset + + features = features_factory(camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH, use_videos=use_videos) + dataset = LeRobotDataset.create( + repo_id=DUMMY_REPO_ID, + fps=DEFAULT_FPS, + features=features, + root=root, + use_videos=use_videos, + streaming_encoding=use_videos, + ) + for _ in range(self.NUM_FRAMES): + frame: dict = {"task": "test"} + for key, ft in dataset.meta.features.items(): + if key in DEFAULT_FEATURES: + continue + if key in dataset.meta.depth_keys: + frame[key] = np.full(ft["shape"], value, dtype=depth_dtype) + elif key in dataset.meta.camera_keys: + frame[key] = np.random.randint(0, 256, ft["shape"], dtype=np.uint8) + else: + frame[key] = np.zeros(ft["shape"], dtype=np.float32) + dataset.add_frame(frame) + return dataset + + @pytest.mark.parametrize("use_videos", [False, True]) + @pytest.mark.parametrize( + ("depth_dtype", "value", "expected_unit"), + [(np.float32, 2.0, DEPTH_METER_UNIT), (np.uint16, 2000, DEPTH_MILLIMETER_UNIT)], + ) + def test_recorded_unit_inferred_persisted_and_kept_in_stats( + self, tmp_path, features_factory, use_videos, depth_dtype, value, expected_unit + ): + """Unit is inferred from the first frame's dtype, drives stats (raw, never canonicalized), and survives a reload.""" + from lerobot.datasets.lerobot_dataset import LeRobotDataset + + dataset = self._record(tmp_path / "ds", features_factory, depth_dtype, value, use_videos) + assert dataset.meta.features[DEPTH_KEY]["info"]["depth_unit"] == expected_unit + dataset.save_episode() + mean = float(np.asarray(dataset.meta.stats[DEPTH_KEY]["mean"]).reshape(-1)[0]) + np.testing.assert_allclose(mean, value, rtol=0.05) + dataset.finalize() + + reloaded = LeRobotDataset(repo_id=DUMMY_REPO_ID, root=tmp_path / "ds") + assert reloaded.meta.features[DEPTH_KEY]["info"]["depth_unit"] == expected_unit + + @pytest.mark.parametrize("use_videos", [False, True]) + @pytest.mark.parametrize( + ("output_unit", "expected"), + [(DEPTH_MILLIMETER_UNIT, 2000.0), (DEPTH_METER_UNIT, 2.0)], + ) + def test_read_honors_output_unit_for_frames_and_stats( + self, tmp_path, features_factory, use_videos, output_unit, expected + ): + """Reloading with a ``depth_output_unit`` converts metre frames (image mode) and rescales stats while preserving count.""" + from lerobot.datasets.lerobot_dataset import LeRobotDataset + + dataset = self._record(tmp_path / "ds", features_factory, np.float32, 2.0, use_videos=use_videos) + dataset.save_episode() + count = float(np.asarray(dataset.meta.stats[DEPTH_KEY]["count"]).reshape(-1)[0]) + dataset.finalize() + + read_dataset = LeRobotDataset( + repo_id=DUMMY_REPO_ID, root=tmp_path / "ds", depth_output_unit=output_unit + ) + stats = read_dataset.meta.stats[DEPTH_KEY] + np.testing.assert_allclose(float(np.asarray(stats["mean"]).reshape(-1)[0]), expected, rtol=0.05) + np.testing.assert_allclose(float(np.asarray(stats["count"]).reshape(-1)[0]), count) + + if not use_videos: + depth = read_dataset[0][DEPTH_KEY] + assert torch.allclose(depth, torch.full_like(depth, expected)) + + from lerobot.datasets.streaming_dataset import StreamingLeRobotDataset + + stream_dataset = StreamingLeRobotDataset( + repo_id=DUMMY_REPO_ID, root=tmp_path / "ds", depth_output_unit=output_unit + ) + stream_depth = next(iter(stream_dataset))[DEPTH_KEY] + assert torch.allclose(stream_depth, torch.full_like(stream_depth, expected)) diff --git a/tests/fixtures/dataset_factories.py b/tests/fixtures/dataset_factories.py index 100922f9c..5c0b0f524 100644 --- a/tests/fixtures/dataset_factories.py +++ b/tests/fixtures/dataset_factories.py @@ -26,6 +26,7 @@ import pytest import torch from datasets import Dataset +from lerobot.configs.video import infer_depth_unit from lerobot.datasets.dataset_metadata import CODEBASE_VERSION, LeRobotDatasetMetadata from lerobot.datasets.feature_utils import get_hf_features_from_features from lerobot.datasets.io_utils import flatten_dict, hf_transform_to_torch @@ -535,6 +536,13 @@ def lerobot_dataset_factory( chunks_size=chunks_size, **info_kwargs, ) + # This synthetic path skips add_frame, so record the depth unit the writer would + # have stored (dummy depth is uint16) to keep ``depth_unit`` present in info.json. + # Reassign a fresh info dict to avoid mutating the shared feature constants. + for ft in info.features.values(): + ft_info = ft.get("info") + if ft_info is not None and ft_info.get("is_depth_map") and "depth_unit" not in ft_info: + ft["info"] = {**ft_info, "depth_unit": infer_depth_unit(np.uint16)} if stats is None: stats = stats_factory(features=info.features) if tasks is None: diff --git a/tests/utils/test_rerun_visualization.py b/tests/utils/test_rerun_visualization.py index e3d205dee..d4c3e6999 100644 --- a/tests/utils/test_rerun_visualization.py +++ b/tests/utils/test_rerun_visualization.py @@ -50,8 +50,9 @@ def mock_rerun(monkeypatch): return self class DummyDepthImage: - def __init__(self, arr, colormap=None): + def __init__(self, arr, meter=None, colormap=None): self.arr = arr + self.meter = meter self.colormap = colormap def dummy_log(key, obj=None, **kwargs): From 7ae12124b06d870ff3c12194772e17aade289c24 Mon Sep 17 00:00:00 2001 From: Caroline Pascal Date: Thu, 2 Jul 2026 15:29:14 +0200 Subject: [PATCH 30/38] fix(save codec options): making sure codec options are always set via `set_if` (#3910) * fix(save codec options): making sure codec options are always safely set through `set_if` * tests(update): updating tests --- src/lerobot/configs/video.py | 12 ++++++------ tests/datasets/test_video_encoding.py | 4 +++- 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/src/lerobot/configs/video.py b/src/lerobot/configs/video.py index 20e46a387..4b956f30e 100644 --- a/src/lerobot/configs/video.py +++ b/src/lerobot/configs/video.py @@ -226,24 +226,24 @@ class VideoEncoderConfig: if encoder_threads is not None: svtav1_parts.append(f"lp={encoder_threads}") if svtav1_parts: - opts["svtav1-params"] = ":".join(svtav1_parts) + set_if("svtav1-params", ":".join(svtav1_parts)) elif self.vcodec in ("h264", "hevc"): set_if("crf", self.crf) set_if("preset", self.preset) if self.fast_decode: - opts["tune"] = "fastdecode" + set_if("tune", "fastdecode") set_if("threads", encoder_threads) elif self.vcodec == "libaom-av1": set_if("crf", self.crf) set_if("preset", self.preset) if encoder_threads is not None: - opts["threads"] = encoder_threads - opts["row-mt"] = 1 + set_if("threads", encoder_threads) + set_if("row-mt", 1) elif self.vcodec in ("h264_videotoolbox", "hevc_videotoolbox"): if self.crf is not None: - opts["q:v"] = max(1, min(100, 100 - self.crf * 2)) + set_if("q:v", max(1, min(100, 100 - self.crf * 2))) elif self.vcodec in ("h264_nvenc", "hevc_nvenc"): - opts["rc"] = 0 + set_if("rc", 0) set_if("qp", self.crf) set_if("preset", self.preset) elif self.vcodec == "h264_vaapi": diff --git a/tests/datasets/test_video_encoding.py b/tests/datasets/test_video_encoding.py index 80819d665..4bfb65ee1 100644 --- a/tests/datasets/test_video_encoding.py +++ b/tests/datasets/test_video_encoding.py @@ -345,7 +345,9 @@ class TestExtraOptions: opts = cfg.get_codec_options() assert opts["qp"] == 20 assert isinstance(opts["qp"], int) - assert cfg.get_codec_options(as_strings=True)["qp"] == "20" + str_opts = cfg.get_codec_options(as_strings=True) + assert str_opts["qp"] == "20" + assert all(isinstance(v, str) for v in str_opts.values()) @require_libsvtav1 def test_structured_fields_win_on_collision(self): From 07285677a3ef2f287a05c271f4b668f20e2ce57f Mon Sep 17 00:00:00 2001 From: Pepijn <138571049+pkooij@users.noreply.github.com> Date: Thu, 2 Jul 2026 19:15:19 +0200 Subject: [PATCH 31/38] fix(train): drive Accelerate mixed precision from policy.dtype (#3912) * fix(train): drive Accelerate mixed precision from policy.dtype `accelerator.autocast()` was always a no-op because `mixed_precision` was never set, so `--policy.dtype=bfloat16` only cast the model params (via the policy) while autocast-eligible ops still ran in fp32/tf32. Map the active policy's `dtype` onto Accelerate's `mixed_precision` (bfloat16 -> bf16, float16 -> fp16, float32 -> no) so autocast is active for bf16/fp16 and stays full precision for float32. Policies without a string `dtype` field fall back to Accelerate's launcher default, so existing behavior is preserved. * style(train): condense mixed-precision comment to one line --- src/lerobot/scripts/lerobot_train.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index 44c94a1eb..6e8458523 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/src/lerobot/scripts/lerobot_train.py @@ -211,8 +211,12 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): # Accelerate auto-detects the device based on the available hardware and ignores the policy.device setting. # Force the device to be CPU when the active config's device is set to CPU (works for both policy and reward model training). force_cpu = cfg.trainable_config.device == "cpu" + # Drive Accelerate's autocast from policy.dtype (bf16/fp16 activate it; float32/absent -> launcher default). + policy_dtype = getattr(cfg.trainable_config, "dtype", None) + mixed_precision = {"bfloat16": "bf16", "float16": "fp16", "float32": "no"}.get(policy_dtype) accelerator = Accelerator( step_scheduler_with_optimizer=False, + mixed_precision=mixed_precision, kwargs_handlers=[ddp_kwargs], cpu=force_cpu, ) From 911734ec9c55f7b2d75b01fba0310075d27b96f3 Mon Sep 17 00:00:00 2001 From: Nikodem Bartnik <39432165+NikodemBartnik@users.noreply.github.com> Date: Fri, 3 Jul 2026 11:39:16 +0200 Subject: [PATCH 32/38] Docs/improve HF jobs documentation (#3909) * improve hf jobs docs * Update docs/source/hardware_guide.mdx Co-authored-by: Nicolas Rabault Signed-off-by: Nikodem Bartnik <39432165+NikodemBartnik@users.noreply.github.com> --------- Signed-off-by: Nikodem Bartnik <39432165+NikodemBartnik@users.noreply.github.com> Co-authored-by: Nicolas Rabault --- docs/source/hardware_guide.mdx | 18 ++++---- docs/source/il_robots.mdx | 79 +--------------------------------- 2 files changed, 10 insertions(+), 87 deletions(-) diff --git a/docs/source/hardware_guide.mdx b/docs/source/hardware_guide.mdx index 5f236d3e8..79c2de98c 100644 --- a/docs/source/hardware_guide.mdx +++ b/docs/source/hardware_guide.mdx @@ -82,18 +82,18 @@ VRAM is the first filter. Within a tier, pick by budget and availability — the ### Hugging Face Jobs -[Hugging Face Jobs](https://huggingface.co/docs/hub/jobs) lets you run training on managed HF infrastructure, billed by the second. The repo publishes a ready-to-use image: **`huggingface/lerobot-gpu:latest`**, rebuilt **every night at 02:00 UTC from `main`** ([`docker_publish.yml`](https://github.com/huggingface/lerobot/blob/main/.github/workflows/docker_publish.yml)) — so it tracks the current state of the repo, not a tagged release. +[Hugging Face Jobs](https://huggingface.co/docs/hub/jobs) lets you run training on managed HF infrastructure, billed by the second, without owning a GPU. `lerobot-train` submits and streams the job for you — just add `--job.target=` to a normal training command: ```bash -hf jobs run --flavor a10g-large huggingface/lerobot-gpu:latest \ - bash -c "nvidia-smi && lerobot-train \ - --policy.type=act --dataset.repo_id=/ \ - --policy.repo_id=/act_ --batch_size=8 --steps=50000" +lerobot-train \ + --policy.type=act --dataset.repo_id=/ \ + --policy.repo_id=/act_ \ + --job.target=a10g-large ``` Notes: -- The leading `nvidia-smi` is a quick sanity check that CUDA is visible inside the container — useful to fail fast if the flavor or driver mismatched. -- The default Job timeout is 30 minutes; pass `--timeout 4h` (or longer) for real training. -- `--flavor` maps onto the table above: `t4-small`/`t4-medium` (T4, ACT only), `l4x1`/`l4x4` (L4 24 GB), `a10g-small/large/largex2/largex4` (A10G 24 GB scaled out), `a100-large` (A100). For the current full catalogue + pricing see [https://huggingface.co/docs/hub/jobs](https://huggingface.co/docs/hub/jobs). -- Prefer not to write the `hf jobs run` wrapper yourself? `lerobot-train` can submit the job for you: just add `--job.target=` to a normal training command and it handles dataset upload, log streaming, and the final model push. See the [imitation-learning training guide](./il_robots). +- Run `hf auth login` once before submitting, the job runs under your token. +- `--job.target` maps onto the table above: `t4-small`/`t4-medium` (T4, ACT only), `l4x1`/`l4x4` (L4 24 GB), `a10g-small/large/largex2/largex4` (A10G 24 GB scaled out), `a100-large` (A100). List the current catalogue with pricing via `hf jobs hardware`, or see [https://huggingface.co/docs/hub/jobs](https://huggingface.co/docs/hub/jobs). +- The job defaults to a `2d` (48h) timeout. Override it with `--job.timeout=4h` (or any other valid duration string) to shorten or extend the timeout. The job automatically stops when the command completes. +- For the full walkthrough — dataset upload, checkpoint streaming, resuming a run on a job — see the [imitation-learning training guide](./il_robots#train-using-hugging-face-jobs). diff --git a/docs/source/il_robots.mdx b/docs/source/il_robots.mdx index 64a39e29c..5893b93f4 100644 --- a/docs/source/il_robots.mdx +++ b/docs/source/il_robots.mdx @@ -532,84 +532,7 @@ If your local computer doesn't have a powerful GPU you could utilize Google Cola Hugging Face jobs let's you easily select hardware and run the training in the cloud. So if you don't have a powerful GPU or you need more VRAM or just want to train a model much faster use HF Jobs! It's pay as you go and you simply pay for each second of use, you can see the pricing and additional information [here](https://huggingface.co/docs/hub/jobs). -> **Tip:** if you just want to launch a standard training run, you can skip building the command below and use the integrated **Train on HF Jobs via `--job.target`** flow described further down — `lerobot-train` then submits the job, uploads a local-only dataset for you, and streams the logs. - -To run the training manually use this command: - - - -```bash -hf jobs run \ - --flavor a10g-small \ - --timeout 4h \ - --secrets HF_TOKEN \ - huggingface/lerobot-gpu:latest \ - -- \ - python -m lerobot.scripts.lerobot_train \ - --dataset.repo_id=username/dataset \ - --policy.type=act \ - --steps=5000 \ - --batch_size=16 \ - --policy.device=cuda \ - --policy.repo_id=username/your_policy \ - --log_freq=100 -``` - - - - -```python -from huggingface_hub import run_job, get_token - -run_name = "act_so101_hf_jobs" -dataset_id = "username/dataset" -user_hub_id = "username" - -command_args = [ - "python", "-m", "lerobot.scripts.lerobot_train", - "--dataset.repo_id", dataset_id, - "--policy.type", "act", - "--steps", "5000", - "--batch_size", "16", - "--num_workers", "4", - "--policy.device", "cuda", - "--log_freq", "100", - "--save_freq", "1000", - "--save_checkpoint", "true", - "--wandb.enable", "false", - "--policy.repo_id", f"{user_hub_id}/{run_name}" -] - -print(f"Submitting job '{run_name}' to Hugging Face Infrastructure...") - -job_info = run_job( - image="huggingface/lerobot-gpu:latest", - command=command_args, - flavor="a10g-small", - timeout="4h", - secrets={"HF_TOKEN": get_token()} -) - -print("\n🚀 Job successfully launched!") -print(f"🔹 Job ID: {job_info.id}") -print(f"🔗 Live UI Dashboard & Logs: {job_info.url}") -``` - - - - - -You can modify the `--flavor` to use different hardware, for example: `t4-small`, `a100-large`, `h200`. Use `hf jobs hardware` to see the full list with pricing. -Depending on the model you want to train and the hardware you selected you can also modify the `--batch_size` and `--number_of_workers`. -For longer training sessions increase the timeout. - -Once the training is started you can go to [Jobs](https://huggingface.co/settings/jobs) and see if your jobs is running as well as all the outputs. Sometimes it takes a few minutes to schedule your job so be patient. - -After training the model will be pushed to hub and you can use it as any other model with LeRobot. - -#### Train on HF Jobs via `--job.target` (integrated CLI) - -`lerobot-train` runs locally by default. To run on a HuggingFace GPU without constructing the Docker command yourself, pass `--job.target` with a hardware flavor name: +`lerobot-train` runs locally by default. To run on a HuggingFace GPU, pass `--job.target` with a hardware flavor name: ```bash lerobot-train \ From e275ea3960332543e2a9f441356775a53720543f Mon Sep 17 00:00:00 2001 From: Pepijn <138571049+pkooij@users.noreply.github.com> Date: Fri, 3 Jul 2026 13:32:38 +0200 Subject: [PATCH 33/38] LingBot-VA: video-action world model (#3731) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(policies): add LingBot-VA autoregressive video-action world model Port the LingBot-VA policy (Wan2.2 dual-stream video+action world model) into LeRobot, following the EO-1 / VLA-JEPA conventions. Covers inference, checkpoint conversion, and predicted-video saving (training is deferred to a follow-up PR). - Vendored Wan transformer/attention/flex/VAE/scheduler modules (key names preserved for near-identity conversion); torch SDPA default, flashattn/flex lazy-guarded. - LingBotVAConfig (registered "lingbot_va") + processor with fixed-quantile action unnormalization; full dual-stream sampling loop with CFG, two flow-matching schedulers and KV cache, mapped onto select_action with observed-keyframe feedback. - convert_lingbot_va_checkpoints.py (libero/robotwin variants): bundles the ~5B transformer, lazy-pulls the frozen VAE+UMT5 from the source repo. - Predicted-video plumbing in lerobot_eval (predicted_frames_callback; opt-in via --policy.save_predicted_video) and ConstantWithWarmupSchedulerConfig. - pyproject: widen diffusers-dep to <0.37, add lingbot_va + imageio-dep extras, add lingbot_va and (missing) eo1 to `all`. - Factory + policies/__init__ wiring, docs page + toctree, and tests. Note: the LIBERO success-rate correctness gate must be validated on a CUDA GPU with the converted checkpoint. * feat(lingbot_va): RoboTwin eef-pose eval, single-file model, Hub checkpoints Make the LingBot-VA port runnable on both LIBERO and RoboTwin and clean up the package to LeRobot conventions. - Consolidate all vendored Wan2.2 model code (transformer, attention, VAE helpers, flow-matching scheduler, grid utils, flex-attention) into a single modeling_lingbot_va.py; remove the separate wan_*/schedulers modules. - Move the fixed action (un)normalization quantiles out of the config and into the post-processor (LIBERO 7-DoF + RoboTwin 16-d eef); remove the conversion script in favour of ready-to-use LeRobot-format checkpoints on the Hub. - Fixes found via on-sim validation: undo LIBERO's 180-degree image flip (image_hflip), encode obs as a multi-frame streaming-VAE clip, reset the streaming VAE cache between episodes, run the transformer in config.dtype, lazy-load frozen VAE/UMT5 by subfolder with the text encoder on CPU. - RoboTwin: add an end-effector-pose action mode to RoboTwinEnv (16-d per-arm xyz+quat+gripper deltas composed onto the initial eef pose, executed via CuRobo IK) and the robotwin_tshape latent layout (full-res head + half-res wrists via a second streaming VAE) with the upstream RoboTwin action quantiles + camera mapping. - Predicted-video saving works for both benchmarks; docs + tests updated. * feat(lingbot_va): implement training / fine-tuning (flow-matching loss) - Implement LingBotVAPolicy.forward(): dual-stream flow-matching training loss (latent + action, timestep-weighted, action-masked) ported from upstream train.py; VAE-encodes camera clips, UMT5-encodes the task, noises both streams, runs the block-causal flex-attention training pass (forward_train). - training_loss_from_streams() core + _build_training_streams() data prep (action scatter into the 30-d space, multi-frame VAE encode incl. robotwin_tshape). - get_optim_params returns only trainable transformer params (LoRA/PEFT friendly); VAE/UMT5 stay frozen. Training needs attn_mode='flex'. - Add a tiny-config single-training-step test (forward->loss->backward->AdamW) and a Training/fine-tuning section in the docs. * fix(lingbot_va): CI quality gate + fast-test collection - Add tests/policies/lingbot_va/__init__.py so the test files don't clash by basename with tests/policies/vla_jepa/* under pytest's default import mode (fast-test collection error). - Fix vendored typos flagged by the typos hook (pach_scale->patch_scale, total_tolen-> total_token_len, stablized->stabilized) and a mypy union-attr in RoboTwinEnv._read_eef_pose. - Apply Prettier formatting to docs/source/lingbot_va.mdx. * docs(lingbot_va): document EEF action-channel schema + camera order * Update lingbot_va.mdx Signed-off-by: Pepijn <138571049+pkooij@users.noreply.github.com> * Update pyproject.toml Signed-off-by: Pepijn <138571049+pkooij@users.noreply.github.com> * Update pyproject.toml Signed-off-by: Pepijn <138571049+pkooij@users.noreply.github.com> * refactor(lingbot_va): drop hardcoded action quantiles; source from checkpoint The LIBERO/RoboTwin action (un)normalization quantiles were hardcoded as module constants in processor_lingbot_va.py. They are already serialized into each checkpoint's policy_postprocessor.json (via LingBotVAActionUnnormalizeStep.get_config) and restored on load by PolicyProcessorPipeline.from_pretrained, so the constants are dead at eval/load time for the released checkpoints (verified: libero_long/robotwin/base all carry their quantiles on the Hub). - Remove LIBERO_ACTION_Q01/Q99, ROBOTWIN_ACTION_Q01/Q99 and _default_action_quantiles. - make_lingbot_va_pre_post_processors now defaults a fresh (unconverted) build to a neutral [-1, 1] mapping (identity rescale); real per-benchmark stats come from the saved checkpoint (or postprocessor_overrides), analogous to dataset-stats normalization. - Update the config doc comment to point at the checkpoint as the source of truth. - Tests: replace the LIBERO-default assertion with a neutral-default check, and add a save_pretrained/from_pretrained round-trip guard for the quantile serialization. * docs(lingbot_va): trim verbose comments - configuration_lingbot_va.py: condense multi-line field comments to one-liners (keep the ── section headers). - processor_lingbot_va.py: shorten the action-quantile explanation block. - modeling_lingbot_va.py: drop the bare "# ----" separator rules, keeping the one-line section headers. No code changes. * docs(lingbot_va): trim provenance comments; default wan path to base repo - configuration_lingbot_va.py: drop the "──" decorations and the "(from transformer/config.json)" note; default wan_pretrained_path to robbyant/lingbot-va-base (has the frozen vae/text_encoder/tokenizer subfolders). - modeling_lingbot_va.py: remove the vendored-code banner and the "(upstream wan_va/...)" section-header provenance/dash decorations; condense the transformer-dtype comment to one line. No code changes. * refactor(lingbot_va): use built-in UnnormalizerProcessorStep for actions Replace the bespoke LingBotVAActionUnnormalizeStep with the standard UnnormalizerProcessorStep in QUANTILES mode, which computes the identical (action + 1) / 2 * (q99 - q01) + q01 mapping. The per-channel q01/q99 are stored as the step's saved state (a safetensors file) and restored on load; a fresh build has no action stats so the step is an identity passthrough. The 3 Hub checkpoints (lerobot/lingbot_va_{libero_long,robotwin,base}) have been re-uploaded with the new post-processor (policy_postprocessor.json + *_unnormalizer_processor.safetensors); reloading from the Hub round-trips q01/q99. - processor_lingbot_va.py: drop the custom step + registry; build the post-processor with UnnormalizerProcessorStep (explicit ACTION->QUANTILES norm_map so the preprocessor / training path is unchanged). - tests: assert the built-in step is used, identity-when-no-stats, correct quantile unnormalization, and a save_pretrained/from_pretrained stats round-trip. * docs(lingbot_va): point checkpoint paths at the lerobot org The LeRobot-format checkpoints moved from pepijn223/* to lerobot/* (libero_long, robotwin, base). Update the eval/train --policy.path examples accordingly. * docs(lingbot_va): condense processor normalization comments * fix(lingbot-va): align RoboTwin evaluation (#3784) Thank you for the RoboTwin fix, and alignment! * applying fixes * updating uv lock and linting * adjusting test to match expected values * cleaning up deps * cleaning up top level imports, styling, and deps guards * cleanup * moving wan utils and loading utils to `utils.py` * removing ftfy by replicating the prompt_clean function without it (we don't expect to have weird chars given in the prompt anyway) * removing unused function * guarding for scipy dep, renaming test to avoid collision * adding back accelerate for peak memory usage optim + justifying robotwin description dep --------- Signed-off-by: Pepijn <138571049+pkooij@users.noreply.github.com> Co-authored-by: pepijn223 Co-authored-by: Gangwei XU Co-authored-by: Maxime Ellerbach --- .dockerignore | 4 + docs/source/_toctree.yml | 2 + docs/source/lingbot_va.mdx | 187 +++ pyproject.toml | 2 + src/lerobot/envs/configs.py | 8 +- src/lerobot/envs/robotwin.py | 175 ++- src/lerobot/optim/schedulers.py | 22 + src/lerobot/policies/__init__.py | 2 + src/lerobot/policies/factory.py | 15 + src/lerobot/policies/lingbot_va/README.md | 1 + src/lerobot/policies/lingbot_va/__init__.py | 21 + .../lingbot_va/configuration_lingbot_va.py | 168 +++ .../lingbot_va/modeling_lingbot_va.py | 853 ++++++++++++ .../lingbot_va/processor_lingbot_va.py | 87 ++ src/lerobot/policies/lingbot_va/utils.py | 1138 +++++++++++++++++ src/lerobot/scripts/lerobot_eval.py | 101 +- .../test_configuration_lingbot_va.py | 78 ++ tests/policies/lingbot_va/test_factory.py | 38 + tests/policies/lingbot_va/test_modules.py | 128 ++ tests/policies/lingbot_va/test_processor.py | 88 ++ uv.lock | 11 +- 21 files changed, 3110 insertions(+), 19 deletions(-) create mode 100644 docs/source/lingbot_va.mdx create mode 120000 src/lerobot/policies/lingbot_va/README.md create mode 100644 src/lerobot/policies/lingbot_va/__init__.py create mode 100644 src/lerobot/policies/lingbot_va/configuration_lingbot_va.py create mode 100644 src/lerobot/policies/lingbot_va/modeling_lingbot_va.py create mode 100644 src/lerobot/policies/lingbot_va/processor_lingbot_va.py create mode 100644 src/lerobot/policies/lingbot_va/utils.py create mode 100644 tests/policies/lingbot_va/test_configuration_lingbot_va.py create mode 100644 tests/policies/lingbot_va/test_factory.py create mode 100644 tests/policies/lingbot_va/test_modules.py create mode 100644 tests/policies/lingbot_va/test_processor.py diff --git a/.dockerignore b/.dockerignore index c0d8a84b5..3295cc1b4 100644 --- a/.dockerignore +++ b/.dockerignore @@ -22,6 +22,10 @@ outputs rl media +# Local virtualenvs (the image provides its own) +.venv +venv + # Logging logs diff --git a/docs/source/_toctree.yml b/docs/source/_toctree.yml index dcd14e131..79f4bc124 100644 --- a/docs/source/_toctree.yml +++ b/docs/source/_toctree.yml @@ -69,6 +69,8 @@ title: VLA-JEPA - local: eo1 title: EO-1 + - local: lingbot_va + title: LingBot-VA - local: fastwam title: FastWAM - local: groot diff --git a/docs/source/lingbot_va.mdx b/docs/source/lingbot_va.mdx new file mode 100644 index 000000000..d33e90340 --- /dev/null +++ b/docs/source/lingbot_va.mdx @@ -0,0 +1,187 @@ +# LingBot-VA + +LingBot-VA is an **autoregressive video-action world-model policy** built on the **Wan2.2** +video-diffusion stack. It interleaves, in one autoregressive sequence, the prediction of +future **video latents** and **robot actions** ("VA" = Video-Action). The LeRobot +integration wires LingBot-VA into the standard training, evaluation and processor +interfaces. + +## Model Overview + +LingBot-VA is a **dual-stream "mixture-of-transformers"**: a video/latent stream +(`patch_embedding_mlp → blocks → proj_out`) and an action stream +(`action_embedder → blocks → action_proj_out`) share the same 30 transformer blocks and +text conditioning. + +| Component | Class | Role | +| ------------------------ | ----------------------- | ----------------------------------------------------------- | +| DiT backbone (trainable) | `WanTransformer3DModel` | ~5B-param dual-stream transformer. | +| VAE (frozen) | `AutoencoderKLWan` | Wan2.2 VAE, `z_dim=48`. Lazy-pulled from the source repo. | +| Text encoder (frozen) | `UMT5EncoderModel` | UMT5-XXL, `d_model=4096`. Lazy-pulled from the source repo. | + +At inference the policy runs an autoregressive loop per chunk: it denoises the video-latent +stream (CFG, ~20 steps) and the action stream (~50 steps) with two independent +flow-matching schedulers, maintaining a KV cache across chunks. Real observed keyframes are +fed back into the KV cache as the chunk is executed (closed-loop world modeling). + +### What the LeRobot Integration Covers + +- Standard `policy.type=lingbot_va` configuration through LeRobot. +- Ready-to-use LeRobot-format checkpoints on the Hub (converted from the released upstream ones). +- Autoregressive dual-stream inference behind the standard `select_action` interface + (single-environment eval, `--eval.batch_size=1`). +- Opt-in saving of the policy's **predicted (imagined) videos** during eval / training. +- Evaluation with `lerobot-eval` on LIBERO and RoboTwin. +- Training / fine-tuning via the dual-stream flow-matching loss (`policy.forward`), see below. + +## Installation + +1. Install LeRobot by following the [Installation Guide](./installation). +2. Install the LingBot-VA extra: + +```bash +pip install -e ".[lingbot_va]" +``` + +## Checkpoints + +The released upstream checkpoints have been converted to LeRobot format and pushed to the Hub: + +| Variant | LeRobot checkpoint | +| ---------------------- | -------------------------------- | +| LIBERO-Long post-train | `lerobot/lingbot_va_libero_long` | +| RoboTwin post-train | `lerobot/lingbot_va_robotwin` | +| Pretrained base | `lerobot/lingbot_va_base` | + +Only the trainable ~5B transformer is stored in the LeRobot +`model.safetensors`. The frozen VAE + UMT5 + tokenizer (~20 GB) are pulled from +`config.wan_pretrained_path` at load time (defaults to the source `robbyant/*` repo). The +UMT5-XXL text encoder runs on CPU by default (`config.text_encoder_device`) so the 5B +transformer + VAE fit on a single 24–32 GB GPU. + +## Evaluation (LIBERO) + +```bash +lerobot-eval \ + --policy.path=lerobot/lingbot_va_libero_long \ + --policy.device=cuda \ + --env.type=libero --env.task=libero_10 \ + --env.observation_height=128 --env.observation_width=128 \ + --eval.n_episodes=50 --eval.batch_size=1 \ + --output_dir=outputs/eval/lingbot_va_libero +``` + +LingBot-VA's streaming inference (KV cache + observed-keyframe feedback) is implemented for +single-environment eval; use `--eval.batch_size=1`. + +## Evaluation (RoboTwin) + +RoboTwin 2.0 needs the SAPIEN + CuRobo simulator stack. You can use the benchmark Docker image +(`docker/Dockerfile.benchmark.robotwin`, which also needs `warp-lang==1.3.1` and CuRobo built +with the GPU's compute capability in `TORCH_CUDA_ARCH_LIST`). RoboTwin uses **end-effector-pose +control**, so run with `--env.action_mode=ee`: the policy predicts per-arm `xyz+quaternion+gripper` +deltas (`robotwin_tshape` latent layout) that are composed onto the episode's initial eef pose and +executed via CuRobo IK. + +```bash +lerobot-eval \ + --policy.path=lerobot/lingbot_va_robotwin \ + --policy.device=cuda \ + --env.type=robotwin --env.task=beat_block_hammer --env.action_mode=ee \ + --eval.n_episodes=10 --eval.batch_size=1 \ + --output_dir=outputs/eval/lingbot_va_robotwin +``` + +### Saving predicted (imagined) videos + +Set `--policy.save_predicted_video=true` to additionally VAE-decode the predicted video +latents and write `pred_episode_*.mp4` next to the env-rendered `eval_episode_*.mp4` videos. +The same flag works for the periodic eval during `lerobot-train`. + +## Training / fine-tuning + +`LingBotVAPolicy.forward(batch)` implements the dual-stream **flow-matching** loss +(`latent_loss + action_loss`, timestep-weighted, action-masked) from the paper: it VAE-encodes +the camera clips into video latents, UMT5-encodes the task, noises both streams, runs the +transformer's block-causal training pass and returns `(loss, metrics)`. Optimizer preset is AdamW +with a linear-warmup-then-constant schedule (matching upstream). + +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. + +```bash +lerobot-train \ + --policy.path=lerobot/lingbot_va_libero_long --policy.attn_mode=flex \ + --policy.use_peft=true \ + --dataset.repo_id= \ + --batch_size=1 --steps=... --output_dir=outputs/train/lingbot_va +``` + +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. + +## Data format (action channels & camera order) + +LingBot-VA is an **end-effector (Cartesian) pose** policy, it predicts EEF poses + gripper, not +joint positions. Actions live in a fixed multi-embodiment **30-dim** layout; map your robot's +action dimensions into these channels and pad the rest with `0` (`used_action_channel_ids` selects +the channels a given checkpoint actually uses): + +| channels | meaning | +| -------- | ----------------------------------------------------- | +| 0–6 | Left-arm end-effector pose | +| 7–13 | Right-arm end-effector pose | +| 14–20 | Left-arm joints (unused by the released checkpoints) | +| 21–27 | Right-arm joints (unused by the released checkpoints) | +| 28 | Left gripper | +| 29 | Right gripper | + +- **LIBERO** uses channels `0–6`: a 6-DoF EEF delta (xyz + rotation) + gripper (single arm). +- **RoboTwin** uses channels `[0–6, 28, 7–13, 29]`: left EEF (xyz + quaternion) + left gripper + + right EEF + right gripper (16 dims). The env converts these poses to joint trajectories via + CuRobo IK — joints are never predicted. + +Joint-space datasets (or a different EEF convention) must be remapped into this schema before +fine-tuning these checkpoints. + +**Camera order is fixed and order-sensitive**, per-camera latents are concatenated spatially in +`obs_cam_keys` order, so the physical camera→slot mapping must match training: + +| benchmark | `obs_cam_keys` (in order) | `camera_layout` | +| --------- | ----------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------- | +| LIBERO | `observation.images.image` (agentview / 3rd-person), `observation.images.image2` (eye-in-hand wrist) | `width_concat` (latents concatenated on width) | +| RoboTwin | `observation.images.head_camera`, `observation.images.left_camera`, `observation.images.right_camera` | `robotwin_tshape` (full-res head below, two half-res wrists on top) | + +The first camera is the exterior/head view and the rest are wrist views. + +## Inference Hyperparameters (LIBERO) + +| Key | Value | +| -------------------------------------- | --------------------------------------------------------------------------------- | +| height × width | 128 × 128 | +| cameras | `observation.images.image` (agentview), `observation.images.image2` (eye-in-hand) | +| action channels used | 0–6 (7-DoF arm + gripper) | +| action_per_frame / frame_chunk_size | 4 / 4 | +| attn_window | 30 | +| video / action denoising steps | 20 / 50 | +| guidance_scale / action_guidance_scale | 5 / 1 | +| snr_shift / action_snr_shift | 5.0 / 0.05 | + +These are the defaults of `LingBotVAConfig`; override any of them via `--policy.=...`. + +## Notes + +- **Attention backend:** inference uses the `torch` SDPA backend (always available). The + `flashattn` and `flex` backends are optional; `flex` is only needed for training. +- **Model size:** the DiT is ~5B params and the frozen VAE+UMT5 add ~20 GB; inference needs + roughly 18–24 GB of VRAM. + +## License + +LingBot-VA is released under Apache-2.0. See the +[upstream repository](https://github.com/Robbyant/lingbot-va). diff --git a/pyproject.toml b/pyproject.toml index 03487e8c4..84fe07328 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -236,6 +236,7 @@ fastwam = [ ] hilserl = ["lerobot[transformers-dep]", "lerobot[dataset]", "gym-hil>=0.1.14,<0.2.0", "lerobot[grpcio-dep]", "lerobot[placo-dep]"] vla_jepa = ["lerobot[transformers-dep]", "lerobot[diffusers-dep]", "lerobot[qwen-vl-utils-dep]"] +lingbot_va = ["lerobot[transformers-dep]", "lerobot[diffusers-dep]", "lerobot[accelerate-dep]"] # Features async = ["lerobot[grpcio-dep]", "lerobot[matplotlib-dep]"] @@ -318,6 +319,7 @@ all = [ "lerobot[xvla]", "lerobot[hilserl]", "lerobot[vla_jepa]", + "lerobot[lingbot_va]", "lerobot[async]", "lerobot[dev]", "lerobot[test]", diff --git a/src/lerobot/envs/configs.py b/src/lerobot/envs/configs.py index 84c40472f..3624357e2 100644 --- a/src/lerobot/envs/configs.py +++ b/src/lerobot/envs/configs.py @@ -757,7 +757,7 @@ class RoboTwinEnvConfig(EnvConfig): task: str = "beat_block_hammer" # single task or comma-separated list fps: int = 25 - episode_length: int = 300 + episode_length: int = 1200 obs_type: str = "pixels_agent_pos" render_mode: str = "rgb_array" # Available cameras from RoboTwin's aloha-agilex embodiment: head_camera @@ -768,6 +768,9 @@ class RoboTwinEnvConfig(EnvConfig): # must equal what SAPIEN actually renders. observation_height: int = 240 observation_width: int = 320 + # "joint": 14-d joint-space control. "ee": 16-d end-effector-pose deltas executed via CuRobo IK + # (for world-model policies like LingBot-VA that predict per-arm xyz+quaternion+gripper poses). + action_mode: str = "joint" features: dict[str, PolicyFeature] = field( default_factory=lambda: { ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(14,)), @@ -784,6 +787,8 @@ class RoboTwinEnvConfig(EnvConfig): ) def __post_init__(self): + if self.action_mode == "ee": + self.features[ACTION] = PolicyFeature(type=FeatureType.ACTION, shape=(16,)) cam_list = [c.strip() for c in self.camera_names.split(",") if c.strip()] for cam in cam_list: self.features[f"pixels/{cam}"] = PolicyFeature( @@ -826,6 +831,7 @@ class RoboTwinEnvConfig(EnvConfig): observation_height=self.observation_height, observation_width=self.observation_width, episode_length=self.episode_length, + action_mode=self.action_mode, ) diff --git a/src/lerobot/envs/robotwin.py b/src/lerobot/envs/robotwin.py index 823f14fa0..5b03f337b 100644 --- a/src/lerobot/envs/robotwin.py +++ b/src/lerobot/envs/robotwin.py @@ -17,6 +17,7 @@ from __future__ import annotations import importlib import logging +import os from collections import defaultdict from collections.abc import Callable, Sequence from functools import partial @@ -28,9 +29,17 @@ import torch from gymnasium import spaces from lerobot.types import RobotObservation +from lerobot.utils.import_utils import _scipy_available from .utils import _LazyAsyncVectorEnv +# scipy is only used for end-effector-pose composition (``--env.action_mode=ee``); guard it so this +# module (and its base-env unit tests, which mock the RoboTwin runtime) imports without scipy installed. +if _scipy_available: + from scipy.spatial.transform import Rotation +else: + Rotation = None + logger = logging.getLogger(__name__) # Camera names as used by RoboTwin 2.0. The wrapper appends "_rgb" when looking @@ -41,10 +50,124 @@ ROBOTWIN_CAMERA_NAMES: tuple[str, ...] = ( "right_camera", ) -ACTION_DIM = 14 # 7 DOF × 2 arms +ACTION_DIM = 14 # 7 DOF × 2 arms (joint-space control mode) +# End-effector-pose control mode: per arm [x, y, z, qx, qy, qz, qw, gripper] = 8, dual-arm = 16. +# Used by world-model policies (e.g. LingBot-VA) that predict eef-pose deltas executed via CuRobo IK. +EEF_ACTION_DIM = 16 ACTION_LOW = -1.0 ACTION_HIGH = 1.0 -DEFAULT_EPISODE_LENGTH = 300 +DEFAULT_EPISODE_LENGTH = 1200 +OFFICIAL_INSTRUCTION_ENV = "LEROBOT_ROBOTWIN_OFFICIAL_INSTRUCTION" +OFFICIAL_INSTRUCTION_TYPE_ENV = "LEROBOT_ROBOTWIN_INSTRUCTION_TYPE" +OFFICIAL_INSTRUCTION_MAX_ENV = "LEROBOT_ROBOTWIN_INSTRUCTION_MAX" + + +def _compose_eef_pose(new_pose: np.ndarray, init_pose: np.ndarray) -> np.ndarray: + """Compose a single-arm predicted delta pose onto the initial pose. + + ``new_pose`` / ``init_pose`` are 8-vectors ``[x, y, z, qx, qy, qz, qw, gripper]``. Translation + is added, rotation is composed (``init_R * new_R``), and the gripper is taken from the + prediction. Mirrors ``add_eef_pose`` in the upstream LingBot-VA RoboTwin client. + """ + new_r = Rotation.from_quat(new_pose[3:7]) + init_r = Rotation.from_quat(init_pose[3:7]) + out_rot = (init_r * new_r).as_quat() + out_trans = new_pose[:3] + init_pose[:3] + return np.concatenate([out_trans, out_rot, new_pose[7:8]]) + + +def _add_init_eef_pose(delta_pose: np.ndarray, init_pose: np.ndarray) -> np.ndarray: + """Compose a dual-arm (16-d) predicted delta pose onto the initial eef pose, normalizing quats.""" + left = _compose_eef_pose(delta_pose[:8], init_pose[:8]) + right = _compose_eef_pose(delta_pose[8:], init_pose[8:]) + out = np.concatenate([left, right]) + # Normalize the two quaternions (indices 3:7 and 11:15) as the upstream client does. + out[3:7] = out[3:7] / (np.linalg.norm(out[3:7]) + 1e-8) + out[11:15] = out[11:15] / (np.linalg.norm(out[11:15]) + 1e-8) + return out + + +def _env_flag(name: str, default: bool = False) -> bool: + raw = os.environ.get(name) + if raw is None: + return default + return raw.strip().lower() in {"1", "true", "yes", "on"} + + +def _arm_for_block(block: Any) -> str: + return "left" if float(block.get_pose().p[0]) < 0 else "right" + + +def _robotwin_blocks_episode_info(task_name: str, env: Any) -> dict[str, str] | None: + """Infer the episode-info dict used by RoboTwin's official instruction generator for block ranking.""" + if task_name == "blocks_ranking_rgb": + return { + "{A}": "red block", + "{B}": "green block", + "{C}": "blue block", + "{a}": _arm_for_block(env.block1), + "{b}": _arm_for_block(env.block2), + "{c}": _arm_for_block(env.block3), + } + if task_name == "blocks_ranking_size": + return { + "{A}": "large block", + "{B}": "medium block", + "{C}": "small block", + "{a}": _arm_for_block(env.block1), + "{b}": _arm_for_block(env.block2), + "{c}": _arm_for_block(env.block3), + } + return None + + +def _generate_robotwin_official_instruction(task_name: str, env: Any) -> str: + """Generate language with RoboTwin's official task templates, matching its eval client.""" + fallback = task_name.replace("_", " ") + episode_info = _robotwin_blocks_episode_info(task_name, env) + if episode_info is None: + logger.warning( + "Official RoboTwin instruction is not implemented for task=%s; using %r.", task_name, fallback + ) + return fallback + + try: + # Part of the robotwin simulator repo, this is being pulled by the docker image running robotwin + # see https://github.com/RoboTwin-Platform/RoboTwin/tree/main/description + # Used to generate the official instructions + from description.utils.generate_episode_instructions import generate_episode_descriptions + except Exception: + logger.warning( + "Failed to import RoboTwin official instruction generator; using %r.", fallback, exc_info=True + ) + return fallback + + instruction_type = os.environ.get(OFFICIAL_INSTRUCTION_TYPE_ENV, "seen") + try: + max_descriptions = int(os.environ.get(OFFICIAL_INSTRUCTION_MAX_ENV, "1000000")) + except ValueError: + max_descriptions = 1000000 + + results = generate_episode_descriptions(task_name, [episode_info], max_descriptions=max_descriptions) + if not results: + logger.warning( + "RoboTwin generated no official instructions for task=%s; using %r.", task_name, fallback + ) + return fallback + + options = results[0].get(instruction_type) or results[0].get("seen") or results[0].get("unseen") + if not options: + logger.warning( + "RoboTwin generated no %s official instructions for task=%s; using %r.", + instruction_type, + task_name, + fallback, + ) + return fallback + + return str(np.random.choice(options)) + + # D435 dims from task_config/_camera_config.yml (what demo_clean.yml selects). DEFAULT_CAMERA_H = 240 DEFAULT_CAMERA_W = 320 @@ -234,6 +357,7 @@ class RoboTwinEnv(gym.Env): observation_width: int | None = None, episode_length: int = DEFAULT_EPISODE_LENGTH, render_mode: str = "rgb_array", + action_mode: str = "joint", ): super().__init__() self.task_name = task_name @@ -241,6 +365,13 @@ class RoboTwinEnv(gym.Env): self.task_description = task_name.replace("_", " ") self.episode_index = episode_index self._reset_stride = n_envs + # "joint": 14-d joint-space actions via take_action(action). "ee": 16-d end-effector-pose + # deltas (added onto the episode's initial eef pose) executed via take_action(.., "ee") + IK. + if action_mode not in ("joint", "ee"): + raise ValueError(f"action_mode must be 'joint' or 'ee'; got {action_mode!r}") + self.action_mode = action_mode + self._action_dim = EEF_ACTION_DIM if action_mode == "ee" else ACTION_DIM + self._init_eef_pose: np.ndarray | None = None self.camera_names = list(camera_names) # Default to D435 dims (the camera type baked into task_config/demo_clean.yml). # The YAML-driven lookup is deferred to reset() so construction doesn't @@ -271,7 +402,7 @@ class RoboTwinEnv(gym.Env): } ) self.action_space = spaces.Box( - low=ACTION_LOW, high=ACTION_HIGH, shape=(ACTION_DIM,), dtype=np.float32 + low=ACTION_LOW, high=ACTION_HIGH, shape=(self._action_dim,), dtype=np.float32 ) def _ensure_env(self) -> None: @@ -317,6 +448,18 @@ class RoboTwinEnv(gym.Env): return {"pixels": images, "agent_pos": joint_state} + def _read_eef_pose(self) -> np.ndarray: + """Read the current 16-d dual-arm eef pose [left(xyz+quat)+grip, right(xyz+quat)+grip].""" + assert self._env is not None, "_read_eef_pose called before _ensure_env()" + ep = self._env.get_obs()["endpose"] + pose = ( + list(ep["left_endpose"]) + + [ep["left_gripper"]] + + list(ep["right_endpose"]) + + [ep["right_gripper"]] + ) + return np.asarray(pose, dtype=np.float64) + def reset(self, seed: int | None = None, **kwargs) -> tuple[RobotObservation, dict]: self._ensure_env() super().reset(seed=seed) @@ -330,16 +473,32 @@ class RoboTwinEnv(gym.Env): self.episode_index += self._reset_stride self._step_count = 0 + use_official_instruction = self.task_name in {"blocks_ranking_rgb", "blocks_ranking_size"} + if _env_flag(OFFICIAL_INSTRUCTION_ENV, default=use_official_instruction): + self.task_description = _generate_robotwin_official_instruction(self.task_name, self._env) + if hasattr(self._env, "set_instruction"): + self._env.set_instruction(instruction=self.task_description) + logger.info("RoboTwin official instruction | task=%s | %s", self.task_name, self.task_description) + else: + self.task_description = self.task_name.replace("_", " ") + + # In eef mode the policy predicts pose deltas relative to the initial eef pose. + if self.action_mode == "ee": + self._init_eef_pose = self._read_eef_pose() + obs = self._get_obs() return obs, {"is_success": False, "task": self.task_name} def step(self, action: np.ndarray) -> tuple[RobotObservation, float, bool, bool, dict[str, Any]]: assert self._env is not None, "step() called before reset()" - if action.ndim != 1 or action.shape[0] != ACTION_DIM: - raise ValueError(f"Expected 1-D action of shape ({ACTION_DIM},), got {action.shape}") + if action.ndim != 1 or action.shape[0] != self._action_dim: + raise ValueError(f"Expected 1-D action of shape ({self._action_dim},), got {action.shape}") with torch.enable_grad(): - if hasattr(self._env, "take_action"): + if self.action_mode == "ee": + ee_action = _add_init_eef_pose(np.asarray(action, dtype=np.float64), self._init_eef_pose) + self._env.take_action(ee_action, action_type="ee") + elif hasattr(self._env, "take_action"): self._env.take_action(action) else: self._env.step(action) @@ -398,6 +557,7 @@ def _make_env_fns( observation_height: int, observation_width: int, episode_length: int, + action_mode: str = "joint", ) -> list[Callable[[], RoboTwinEnv]]: """Return n_envs factory callables for a single task.""" @@ -410,6 +570,7 @@ def _make_env_fns( observation_height=observation_height, observation_width=observation_width, episode_length=episode_length, + action_mode=action_mode, ) return [partial(_make_one, i) for i in range(n_envs)] @@ -423,6 +584,7 @@ def create_robotwin_envs( observation_height: int = DEFAULT_CAMERA_H, observation_width: int = DEFAULT_CAMERA_W, episode_length: int = DEFAULT_EPISODE_LENGTH, + action_mode: str = "joint", ) -> dict[str, dict[int, Any]]: """Create vectorized RoboTwin 2.0 environments. @@ -473,6 +635,7 @@ def create_robotwin_envs( observation_height=observation_height, observation_width=observation_width, episode_length=episode_length, + action_mode=action_mode, ) if is_async: lazy = _LazyAsyncVectorEnv(fns, cached_obs_space, cached_act_space, cached_metadata) diff --git a/src/lerobot/optim/schedulers.py b/src/lerobot/optim/schedulers.py index 250650089..74111e7ef 100644 --- a/src/lerobot/optim/schedulers.py +++ b/src/lerobot/optim/schedulers.py @@ -83,6 +83,28 @@ class VQBeTSchedulerConfig(LRSchedulerConfig): return LambdaLR(optimizer, lr_lambda, -1) +@LRSchedulerConfig.register_subclass("constant_with_warmup") +@dataclass +class ConstantWithWarmupSchedulerConfig(LRSchedulerConfig): + """Linear warmup followed by a constant learning rate. + + Mirrors the ``warmup_constant_lambda`` used by LingBot-VA (upstream ``wan_va/train.py``): + the LR ramps linearly from 0 to the peak over ``num_warmup_steps`` steps, then stays flat. + """ + + num_warmup_steps: int = 1000 + + def build(self, optimizer: Optimizer, num_training_steps: int) -> LambdaLR: + warmup_steps = self.num_warmup_steps or 0 + + def lr_lambda(current_step): + if current_step < warmup_steps: + return float(current_step) / float(max(1, warmup_steps)) + return 1.0 + + return LambdaLR(optimizer, lr_lambda, -1) + + @LRSchedulerConfig.register_subclass("cosine_decay_with_warmup") @dataclass class CosineDecayWithWarmupSchedulerConfig(LRSchedulerConfig): diff --git a/src/lerobot/policies/__init__.py b/src/lerobot/policies/__init__.py index 4daa6abc5..494427692 100644 --- a/src/lerobot/policies/__init__.py +++ b/src/lerobot/policies/__init__.py @@ -21,6 +21,7 @@ from .factory import get_policy_class, make_policy, make_policy_config, make_pre from .fastwam.configuration_fastwam import FastWAMConfig as FastWAMConfig from .gaussian_actor.configuration_gaussian_actor import GaussianActorConfig as GaussianActorConfig from .groot.configuration_groot import GrootConfig as GrootConfig +from .lingbot_va.configuration_lingbot_va import LingBotVAConfig as LingBotVAConfig from .molmoact2.configuration_molmoact2 import MolmoAct2Config as MolmoAct2Config from .multi_task_dit.configuration_multi_task_dit import MultiTaskDiTConfig as MultiTaskDiTConfig from .pi0.configuration_pi0 import PI0Config as PI0Config @@ -46,6 +47,7 @@ __all__ = [ "FastWAMConfig", "GaussianActorConfig", "GrootConfig", + "LingBotVAConfig", "MolmoAct2Config", "MultiTaskDiTConfig", "PI0Config", diff --git a/src/lerobot/policies/factory.py b/src/lerobot/policies/factory.py index f5acb0170..12871218e 100644 --- a/src/lerobot/policies/factory.py +++ b/src/lerobot/policies/factory.py @@ -50,6 +50,7 @@ from .eo1.configuration_eo1 import EO1Config 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 @@ -163,6 +164,10 @@ def get_policy_class(name: str) -> type[PreTrainedPolicy]: 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 @@ -223,6 +228,8 @@ def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig: 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) else: @@ -458,6 +465,14 @@ def make_pre_post_processors( 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 diff --git a/src/lerobot/policies/lingbot_va/README.md b/src/lerobot/policies/lingbot_va/README.md new file mode 120000 index 000000000..2ec3c82af --- /dev/null +++ b/src/lerobot/policies/lingbot_va/README.md @@ -0,0 +1 @@ +../../../../docs/source/lingbot_va.mdx \ No newline at end of file diff --git a/src/lerobot/policies/lingbot_va/__init__.py b/src/lerobot/policies/lingbot_va/__init__.py new file mode 100644 index 000000000..db2ab17d1 --- /dev/null +++ b/src/lerobot/policies/lingbot_va/__init__.py @@ -0,0 +1,21 @@ +#!/usr/bin/env python + +# Copyright 2026 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. + +from .configuration_lingbot_va import LingBotVAConfig +from .modeling_lingbot_va import LingBotVAPolicy +from .processor_lingbot_va import make_lingbot_va_pre_post_processors + +__all__ = ["LingBotVAConfig", "LingBotVAPolicy", "make_lingbot_va_pre_post_processors"] diff --git a/src/lerobot/policies/lingbot_va/configuration_lingbot_va.py b/src/lerobot/policies/lingbot_va/configuration_lingbot_va.py new file mode 100644 index 000000000..424ea7c63 --- /dev/null +++ b/src/lerobot/policies/lingbot_va/configuration_lingbot_va.py @@ -0,0 +1,168 @@ +# Copyright 2026 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. + +"""Configuration for the LingBot-VA policy. + +LingBot-VA is an autoregressive video-action world-model policy built on the Wan2.2 +video-diffusion stack. It interleaves prediction of future video latents and robot +actions in a single dual-stream transformer. See ``docs/source/lingbot_va.mdx`` and the +upstream repository (https://github.com/Robbyant/lingbot-va). + +Defaults below match the upstream LIBERO configuration (``wan_va/configs/va_libero_cfg.py``) +and the ``transformer/config.json`` of the released checkpoints. +""" + +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.utils.constants import ACTION + + +@PreTrainedConfig.register_subclass("lingbot_va") +@dataclass +class LingBotVAConfig(PreTrainedConfig): + """Configuration for the native LingBot-VA policy integration in LeRobot.""" + + # Wan transformer architecture + patch_size: tuple[int, int, int] = (1, 2, 2) + num_attention_heads: int = 24 + attention_head_dim: int = 128 + in_channels: int = 48 + out_channels: int = 48 + action_dim: int = 30 + text_dim: int = 4096 + freq_dim: int = 256 + ffn_dim: int = 14336 + num_layers: int = 30 + cross_attn_norm: bool = True + eps: float = 1e-6 + rope_max_seq_len: int = 1024 + # "flex" = training only (needs recent torch); inference uses "torch" SDPA or "flashattn". + attn_mode: str = "torch" + + # Frozen sub-models (VAE + UMT5 text encoder + tokenizer) + # ~20 GB of frozen weights, NOT bundled in the checkpoint; lazily pulled from this HF repo / + # local dir (must hold diffusers-style ``vae/``, ``text_encoder/``, ``tokenizer/`` sub-folders). + wan_pretrained_path: str = "robbyant/lingbot-va-base" + dtype: str = "bfloat16" # transformer / VAE / text-encoder dtype: "bfloat16", "float16", "float32" + # Frozen UMT5-XXL encoder device; "cpu" frees ~11 GB VRAM (it runs once per episode). + text_encoder_device: str = "cpu" + + # Observation cameras (order matters: latents are concatenated on width; LIBERO defaults) + obs_cam_keys: list[str] = field( + default_factory=lambda: ["observation.images.image", "observation.images.image2"] + ) + # Undo the LIBERO env processor's extra horizontal flip to match the model's training orientation. + image_hflip: bool = False + # Camera latent layout: "width_concat" (cameras concatenated on width; LIBERO) or + # "robotwin_tshape" (full-res head + half-res wrists in a "T"; RoboTwin). + camera_layout: str = "width_concat" + + # Inference hyperparameters (LIBERO defaults) + n_obs_steps: int = 1 + height: int = 128 + width: int = 128 + action_per_frame: int = 4 + frame_chunk_size: int = 4 + attn_window: int = 30 + num_inference_steps: int = 20 + video_exec_step: int = -1 + action_num_inference_steps: int = 50 + guidance_scale: float = 5.0 + action_guidance_scale: float = 1.0 + snr_shift: float = 5.0 + action_snr_shift: float = 0.05 + max_sequence_length: int = 512 # UMT5 prompt length + + # Subset of the 30-d action space used by the benchmark (LIBERO = 7-DoF). The action + # (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))) + + # Opt-in: VAE-decode predicted video latents to ``self.last_predicted_frames`` for saving MP4s. + save_predicted_video: bool = False + + # Normalization: IDENTITY here; images are scaled + VAE-encoded and actions are + # quantile-(un)normalized inside the policy / dedicated processor steps. + normalization_mapping: dict[str, NormalizationMode] = field( + default_factory=lambda: { + "VISUAL": NormalizationMode.IDENTITY, + "STATE": NormalizationMode.IDENTITY, + "ACTION": NormalizationMode.IDENTITY, + } + ) + + # Optimizer / scheduler (training; AdamW + warmup-constant per upstream train.py) + optimizer_lr: float = 1e-5 + optimizer_betas: tuple[float, float] = (0.9, 0.95) + optimizer_eps: float = 1e-8 + optimizer_weight_decay: float = 1e-4 + optimizer_grad_clip_norm: float = 1.0 + scheduler_warmup_steps: int = 1000 + + def __post_init__(self): + super().__post_init__() + if self.attn_mode not in ("torch", "flashattn", "flex"): + raise ValueError(f"attn_mode must be one of 'torch', 'flashattn', 'flex'; got {self.attn_mode!r}") + + @property + def chunk_size(self) -> int: + """Number of single-step actions produced per autoregressive chunk.""" + return self.frame_chunk_size * self.action_per_frame + + @property + def n_action_steps(self) -> int: + """Number of actions executed before refilling (the whole chunk).""" + return self.chunk_size + + def validate_features(self) -> None: + image_features = [key for key, feat in self.input_features.items() if feat.type == FeatureType.VISUAL] + if not image_features: + raise ValueError( + "LingBot-VA requires at least one visual input feature. " + "No features of type FeatureType.VISUAL found in input_features." + ) + if ACTION not in self.output_features: + self.output_features[ACTION] = PolicyFeature( + type=FeatureType.ACTION, shape=(len(self.used_action_channel_ids),) + ) + + def get_optimizer_preset(self) -> AdamWConfig: + return AdamWConfig( + lr=self.optimizer_lr, + betas=self.optimizer_betas, + eps=self.optimizer_eps, + weight_decay=self.optimizer_weight_decay, + grad_clip_norm=self.optimizer_grad_clip_norm, + ) + + def get_scheduler_preset(self) -> LRSchedulerConfig | None: + # Upstream uses a linear warmup followed by a constant LR (warmup_constant_lambda). + return ConstantWithWarmupSchedulerConfig(num_warmup_steps=self.scheduler_warmup_steps) + + @property + def observation_delta_indices(self) -> list[int]: + temporal_downsample = 4 + stride = max(1, self.action_per_frame // temporal_downsample) + return list(range(0, self.frame_chunk_size * temporal_downsample * stride, stride)) + + @property + def action_delta_indices(self) -> list[int]: + return list(range(self.chunk_size)) + + @property + def reward_delta_indices(self) -> None: + return None diff --git a/src/lerobot/policies/lingbot_va/modeling_lingbot_va.py b/src/lerobot/policies/lingbot_va/modeling_lingbot_va.py new file mode 100644 index 000000000..0f70ad290 --- /dev/null +++ b/src/lerobot/policies/lingbot_va/modeling_lingbot_va.py @@ -0,0 +1,853 @@ +# Copyright 2024-2025 The Robbyant Team Authors. All rights reserved. +# Copyright 2026 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. + +"""LingBot-VA policy: an autoregressive video-action world model on the Wan2.2 stack. + +The sampling loop is a faithful re-implementation of the upstream streaming server +(``wan_va/wan_va_server.py``) and LIBERO client (``evaluation/libero/client.py``), adapted +to LeRobot's ``select_action`` interface: + + * the trainable dual-stream transformer is owned as a sub-module and round-trips in the + single ``model.safetensors`` checkpoint; + * the frozen Wan VAE + UMT5 text encoder + tokenizer are *lazily pulled* from + ``config.wan_pretrained_path`` (not bundled), so the LeRobot checkpoint stays small; + * ``predict_action_chunk`` runs one autoregressive chunk (video stream then action + stream, each with CFG and its own flow-matching scheduler) and updates the KV cache; + * ``select_action`` drains a per-step action queue and records the real observed + keyframes that are fed back into the KV cache when the queue is refilled. + +NOTE: The streaming path is written for single-environment eval (``--eval.batch_size=1``). +""" + +from collections import deque + +import torch +import torch.nn.functional as F # noqa: N812 +from einops import rearrange +from torch import Tensor + +from lerobot.policies.pretrained import PreTrainedPolicy +from lerobot.utils.constants import ACTION +from lerobot.utils.import_utils import require_package + +from .configuration_lingbot_va import LingBotVAConfig +from .utils import ( + FlowMatchScheduler, + WanTransformer3DModel, + WanVAEStreamingWrapper, + _sample_timestep_id, + _torch_dtype, + clean_prompt, + data_seq_to_patch, + denormalize_latents, + get_mesh_id, + load_text_encoder, + load_tokenizer, + load_vae, +) + + +class LingBotVAPolicy(PreTrainedPolicy): + """LeRobot wrapper for the LingBot-VA autoregressive video-action world model.""" + + config_class = LingBotVAConfig + name = "lingbot_va" + + def __init__(self, config: LingBotVAConfig, **kwargs): + require_package("diffusers", extra="lingbot_va") + require_package("transformers", extra="lingbot_va") + super().__init__(config) + config.validate_features() + self.config = config + + self.dtype = _torch_dtype(config.dtype) + + # Trainable dual-stream transformer (the only sub-module saved in the LeRobot checkpoint). + self.transformer = WanTransformer3DModel( + patch_size=tuple(config.patch_size), + num_attention_heads=config.num_attention_heads, + attention_head_dim=config.attention_head_dim, + in_channels=config.in_channels, + out_channels=config.out_channels, + action_dim=config.action_dim, + text_dim=config.text_dim, + freq_dim=config.freq_dim, + ffn_dim=config.ffn_dim, + num_layers=config.num_layers, + cross_attn_norm=config.cross_attn_norm, + eps=config.eps, + rope_max_seq_len=config.rope_max_seq_len, + attn_mode=config.attn_mode, + ) + # Run the transformer in config.dtype (bf16); norm/modulation paths upcast to fp32 internally. + self.transformer = self.transformer.to(self.dtype) + + # Frozen modules are stored OUTSIDE the nn.Module registry (plain dict) so they are + # neither saved into model.safetensors nor moved by ``.to()``. They are lazily loaded + # from ``config.wan_pretrained_path`` the first time inference runs. + self._frozen: dict = {} + + self.last_predicted_frames: Tensor | None = None + self.last_predicted_latents: Tensor | None = None + self.reset() + + # Frozen-module lazy loading (VAE + UMT5 + tokenizer) + def _ensure_frozen_modules(self): + if self._frozen: + return + path = self.config.wan_pretrained_path + device = self.config.device + + # The frozen modules always live in ``vae/``, ``text_encoder/`` and ``tokenizer/`` + # sub-folders -- both in the released diffusers-style HF repos and in the local + # ``--bundle-frozen`` output dir. ``from_pretrained(path, subfolder=...)`` resolves + # them for either a HF repo id or a local directory. + vae = load_vae(path, torch_dtype=self.dtype, torch_device=device, subfolder="vae") + # The UMT5-XXL text encoder (~11 GB) runs once per episode; keep it on its own + # (CPU by default) device so the 5B transformer + VAE fit on a single GPU. + text_encoder = load_text_encoder( + path, + torch_dtype=self.dtype, + torch_device=self.config.text_encoder_device, + subfolder="text_encoder", + ) + tokenizer = load_tokenizer(path, subfolder="tokenizer") + self._frozen = { + "vae": vae.eval(), + "streaming_vae": WanVAEStreamingWrapper(vae), + "text_encoder": text_encoder.eval(), + "tokenizer": tokenizer, + } + # RoboTwin's T-shape layout encodes the half-resolution wrist cameras through a second + # streaming VAE (separate causal cache) alongside the full-res head camera. + if self.config.camera_layout == "robotwin_tshape": + vae_half = load_vae(path, torch_dtype=self.dtype, torch_device=device, subfolder="vae") + self._frozen["streaming_vae_half"] = WanVAEStreamingWrapper(vae_half.eval()) + + @property + def _vae(self): + return self._frozen["vae"] + + @property + def _streaming_vae(self): + return self._frozen["streaming_vae"] + + # PreTrainedPolicy API + def get_optim_params(self) -> dict: + # Only the transformer is trainable; the VAE / text encoder stay frozen (kept outside the + # nn.Module registry). With PEFT/LoRA this naturally returns just the adapter params. + return [p for p in self.transformer.parameters() if p.requires_grad] + + def reset(self): + """Reset all per-episode streaming state (KV cache, queues, frame counter).""" + cfg = self.config + self._action_queue: deque = deque(maxlen=cfg.n_action_steps) + self._obs_buffer: list = [] # raw keyframe obs (one per env substep) observed this chunk + self._executed_actions: Tensor | None = ( + None # last chunk's actions (model-normalized) for KV feedback + ) + self._started = False # first select_action call uses the obs as the conditioning frame + self._exec_step = 0 # index of the action being executed within the current chunk + self._prev_j = 0 # sub-step index (within a predicted frame) of the last executed action + # Sample one keyframe every ``action_per_frame / temporal_downsample`` executed sub-steps so + # that exactly ``frame_chunk_size * temporal_downsample`` frames are VAE-encoded per chunk + # (the Wan2.2 VAE temporal downsample is 4 -> ``frame_chunk_size`` latent frames). + self._keyframe_stride = max(1, cfg.action_per_frame // 4) + self._frame_st_id = 0 + self._first_chunk = True + self._prompt: str | None = None + self._prompt_embeds = None + self._negative_prompt_embeds = None + self.last_predicted_frames = None + self.last_predicted_latents = None + self._use_cfg = (cfg.guidance_scale > 1) or (cfg.action_guidance_scale > 1) + # Two independent flow-matching schedulers (video latent + action streams). + self._scheduler = FlowMatchScheduler(shift=cfg.snr_shift, sigma_min=0.0, extra_one_step=True) + self._action_scheduler = FlowMatchScheduler( + shift=cfg.action_snr_shift, sigma_min=0.0, extra_one_step=True + ) + self._scheduler.set_timesteps(1000, training=True) + self._action_scheduler.set_timesteps(1000, training=True) + self._cache_initialised = False + # Clear KV cache on the (already-built) transformer, if present. + if hasattr(self, "transformer"): + self.transformer.clear_cache("pos") + # Reset the causal streaming-VAE feat cache between episodes (mirrors upstream ``_reset``). + # Without this the encoder carries over the previous episode's temporal state, corrupting the + # latent frame counts on the next episode's first encode. + if self._frozen: + self._frozen["streaming_vae"].clear_cache() + if "streaming_vae_half" in self._frozen: + self._frozen["streaming_vae_half"].clear_cache() + + # Training (flow-matching dual-stream loss). Requires attn_mode="flex". + def _ensure_train_schedulers(self): + if getattr(self, "_train_sched_latent", None) is None: + cfg = self.config + self._train_sched_latent = FlowMatchScheduler( + shift=cfg.snr_shift, sigma_min=0.0, extra_one_step=True + ) + self._train_sched_latent.set_timesteps(1000, training=True) + self._train_sched_action = FlowMatchScheduler( + shift=cfg.action_snr_shift, sigma_min=0.0, extra_one_step=True + ) + self._train_sched_action.set_timesteps(1000, training=True) + + @torch.no_grad() + def _add_noise_stream(self, latent, scheduler, action_mask, action_mode, noisy_cond_prob): + """Flow-matching noising of one stream (port of upstream ``Trainer._add_noise``).""" + device = latent.device + b, _c, f, _h, _w = latent.shape + p = self.config.patch_size + patch_f, patch_h, patch_w = (1, 1, 1) if action_mode else (p[0], p[1], p[2]) + + ts_ids = _sample_timestep_id(f, num_train_timesteps=scheduler.num_train_timesteps) + noise = torch.zeros_like(latent).normal_() + timesteps = scheduler.timesteps[ts_ids].to(device) + noisy_latents = scheduler.add_noise(latent, noise, timesteps, t_dim=2) + targets = scheduler.training_target(latent, noise, timesteps) + + grid_id = ( + get_mesh_id( + latent.shape[-3] // patch_f, + latent.shape[-2] // patch_h, + latent.shape[-1] // patch_w, + t=1 if action_mode else 0, + f_w=1, + f_shift=0, + action=action_mode, + ) + .to(device)[None] + .repeat(b, 1, 1) + ) + + if torch.rand(1).item() < noisy_cond_prob: + cond_ids = _sample_timestep_id( + f, min_timestep_bd=0.5, max_timestep_bd=1.0, num_train_timesteps=scheduler.num_train_timesteps + ) + cond_noise = torch.zeros_like(latent).normal_() + cond_timesteps = scheduler.timesteps[cond_ids].to(device) + latent = scheduler.add_noise(latent, cond_noise, cond_timesteps, t_dim=2) + else: + cond_timesteps = torch.zeros_like(timesteps) + + if action_mask is not None: + noisy_latents = noisy_latents * action_mask.float() + targets = targets * action_mask.float() + latent = latent * action_mask.float() + + return { + "timesteps": timesteps[None].repeat(b, 1), + "noisy_latents": noisy_latents, + "targets": targets, + "latent": latent, + "cond_timesteps": cond_timesteps[None].repeat(b, 1), + "grid_id": grid_id, + } + + def _flow_matching_loss(self, input_dict, pred): + """Dual-stream flow-matching loss (port of upstream ``Trainer.compute_loss``).""" + 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]) + latent_pred = data_seq_to_patch( + self.config.patch_size, + latent_pred, + ld["targets"].shape[-3], + ld["targets"].shape[-2], + ld["targets"].shape[-1], + batch_size=latent_pred.shape[0], + ) + bn, fn = ld["timesteps"].shape + lw = self._train_sched_latent.training_weight(ld["timesteps"].flatten()).reshape(bn, fn) + aw = self._train_sched_action.training_weight(ad["timesteps"].flatten()).reshape(bn, fn) + + latent_loss = F.mse_loss(latent_pred.float(), ld["targets"].float().detach(), reduction="none") + 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() + + amask = ad["actions_mask"].float() + action_loss = F.mse_loss(action_pred.float(), ad["targets"].float().detach(), reduction="none") + action_loss = ( + (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 + + def training_loss_from_streams(self, latents, actions, actions_mask, text_emb): + """Core dual-stream training loss given prepared latents / actions / text embeddings. + + ``latents``: ``[B, in_channels, F, h, w]`` (normalized video latents). + ``actions`` / ``actions_mask``: ``[B, action_dim, F, action_per_frame, 1]``. + ``text_emb``: ``[B, seq_len, text_dim]``. Returns ``(loss, {latent_loss, action_loss})``. + """ + if self.config.attn_mode != "flex": + raise ValueError( + "LingBot-VA training requires attn_mode='flex' (block-causal flow-matching masks). " + "Load/convert the policy with --policy.attn_mode=flex for training/fine-tuning." + ) + self._ensure_train_schedulers() + latent_dict = self._add_noise_stream( + latents, self._train_sched_latent, action_mask=None, action_mode=False, noisy_cond_prob=0.5 + ) + action_dict = self._add_noise_stream( + actions, self._train_sched_action, action_mask=actions_mask, action_mode=True, noisy_cond_prob=0.0 + ) + latent_dict["text_emb"] = text_emb + action_dict["text_emb"] = text_emb + action_dict["actions_mask"] = actions_mask + input_dict = { + "latent_dict": latent_dict, + "action_dict": action_dict, + "chunk_size": int(torch.randint(1, 5, (1,)).item()), + "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) + loss = latent_loss + action_loss + return loss, {"latent_loss": latent_loss.item(), "action_loss": action_loss.item()} + + def forward(self, batch: dict[str, Tensor]) -> 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'``. + """ + 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) + + @torch.no_grad() + def _build_training_streams(self, batch): + """Build (latents, actions, actions_mask, text_emb) from a LeRobot training batch. + + Camera frames per ``obs_cam_keys`` are expected as a temporal clip ``[B, C, T, H, W]`` (or + ``[B, T, C, H, W]``); they are VAE-encoded into ``F = T / temporal_downsample`` latent frames. + Actions ``[B, F*action_per_frame, n_used]`` are scattered into the model's ``action_dim`` space. + """ + cfg = self.config + device = cfg.device + # text embeddings + task = batch.get("task") + if isinstance(task, str): + task = [task] + text_emb = self._get_t5_prompt_embeds(list(task), cfg.max_sequence_length) + + # video latents (VAE-encode the camera clips) + latents = self._encode_training_latents(batch) + + # actions -> [B, action_dim, F, action_per_frame, 1] + act = batch[ACTION].to(device) # [B, F*apf, n_used] + b = act.shape[0] + used = cfg.used_action_channel_ids + apf, fc = cfg.action_per_frame, cfg.frame_chunk_size + act = act[:, : fc * apf].reshape(b, fc, apf, len(used)).permute(0, 3, 1, 2) # [B, n_used, F, apf] + full = act.new_zeros(b, cfg.action_dim, fc, apf) + idx = torch.as_tensor(used, device=device) + full[:, idx] = act + actions = full.unsqueeze(-1).to(self.dtype) # [B, action_dim, F, apf, 1] + mask = torch.zeros(cfg.action_dim, device=device, dtype=self.dtype) + mask[idx] = 1.0 + actions_mask = mask.view(1, -1, 1, 1, 1).expand_as(actions) + return latents, actions, actions_mask, text_emb + + @torch.no_grad() + def _encode_training_latents(self, batch) -> Tensor: + """VAE-encode the per-camera training clips into normalized video latents [B, C, F, h, w].""" + vae_device = next(self._vae.parameters()).device + + def _clip(key): + x = batch[key].to(vae_device) + if x.dim() == 4: # [B, C, H, W] -> single frame clip + x = x.unsqueeze(2) + elif x.shape[1] not in (1, 3) and x.shape[2] in (1, 3): # [B, T, C, H, W] -> [B, C, T, H, W] + x = x.permute(0, 2, 1, 3, 4) + return x.contiguous() + + def _encode(x, size): + b, c, t = x.shape[:3] + x = F.interpolate(x.flatten(0, 1).float(), size=size, mode="bilinear", align_corners=False) + x = (x.view(b, c, t, *size) * 2.0 - 1.0).to(self.dtype) + mu = self._vae.encode(x).latent_dist.mode() # [B, z_dim, F, h, w] + mean = torch.tensor(self._vae.config.latents_mean).view(1, -1, 1, 1, 1).to(mu.device) + inv_std = (1.0 / torch.tensor(self._vae.config.latents_std)).view(1, -1, 1, 1, 1).to(mu.device) + return ((mu.float() - mean) * inv_std).to(mu) + + keys = self.config.obs_cam_keys + if self.config.camera_layout == "robotwin_tshape": + h, w = self.config.height, self.config.width + head = _encode(_clip(keys[0]), (h, w)) + left = _encode(_clip(keys[1]), (h // 2, w // 2)) + right = _encode(_clip(keys[2]), (h // 2, w // 2)) + return torch.cat([torch.cat([left, right], dim=-1), head], dim=-2).to(self.config.device) + per_cam = [_encode(_clip(k), (self.config.height, self.config.width)) for k in keys] + return torch.cat(per_cam, dim=-1).to(self.config.device) + + @torch.no_grad() + def select_action(self, batch: dict[str, Tensor], **kwargs) -> Tensor: + """Return one action, refilling the chunk (and feeding back observed keyframes) as needed. + + Mirrors the upstream LIBERO client loop (``evaluation/libero/client.py``): the first obs is + the conditioning frame; every observation produced afterwards is buffered as a keyframe and, + once the chunk's actions are exhausted, the buffered frames + executed actions are fed back + into the KV cache before the next chunk is predicted. + """ + self.eval() + self._ensure_frozen_modules() + self._maybe_init_prompt(batch) + + if not self._started: + # First call: this observation conditions the first chunk (it is *not* a keyframe). + self._started = True + actions = self.predict_action_chunk(batch) # [B, chunk_size, n_used] + self._action_queue.extend(actions.transpose(0, 1)) # [chunk_size, B, n_used] + self._obs_buffer = [] + self._exec_step = 0 + else: + # This observation is the result of the previously executed action -> a candidate + # keyframe. Buffer it on the sub-step boundary the upstream client samples on. + if (self._prev_j + 1) % self._keyframe_stride == 0: + self._obs_buffer.append(self._extract_raw_obs(batch)) + if len(self._action_queue) == 0: + # All actions for the current chunk have been executed; feed the observed + # keyframes + executed actions back and predict the next chunk. + actions = self.predict_action_chunk(None) + self._action_queue.extend(actions.transpose(0, 1)) + self._exec_step = 0 + + self._prev_j = self._exec_step % self.config.action_per_frame + self._exec_step += 1 + return self._action_queue.popleft() + + @torch.no_grad() + def predict_action_chunk(self, batch: dict[str, Tensor], **kwargs) -> Tensor: + """Run one autoregressive chunk and return actions ``[B, chunk_size, n_used]`` (normalized).""" + self.eval() + self._ensure_frozen_modules() + self._maybe_init_prompt(batch) + + is_first = self._first_chunk + if is_first: + init_latent = self._encode_frames([self._extract_raw_obs(batch)]) + self._init_latent = init_latent + self._init_streaming_cache(init_latent) + self._obs_buffer = [] # frame 0 (the init obs) conditions the chunk; it is not fed back + actions, latents = self._infer(init_latent, frame_st_id=0) + self._first_chunk = False + else: + # Feed the real observed keyframes + the executed actions back into the KV cache. + self._compute_kv_cache(self._obs_buffer, self._executed_actions) + self._obs_buffer = [] + actions, latents = self._infer(None, frame_st_id=self._frame_st_id) + + # actions: [B, action_dim, F, action_per_frame, 1] (model-normalized). Keep for KV feedback. + self._executed_actions = actions + + if self.config.save_predicted_video: + # Match upstream LingBot-VA visualization: collect chunk latents and decode the + # concatenated latent sequence once after the rollout finishes. + self.last_predicted_frames = None + self.last_predicted_latents = latents.detach().to("cpu") + + # On the first chunk, frame 0 is the conditioning frame (already "known"): the upstream + # LIBERO client skips it (start_idx=1), so we drop the first frame's actions here. + used = self.config.used_action_channel_ids + a = actions[:, used] # [B, n_used, F, action_per_frame, 1] + if is_first: + a = a[:, :, 1:] # drop frame 0 -> (F-1) frames of actions + a = a.squeeze(-1).flatten(2) # [B, n_used, n_steps] + a = a.transpose(1, 2).contiguous() # [B, n_steps, n_used] + return a.to(torch.float32) + + # Prompt / text encoding + def _maybe_init_prompt(self, batch): + if self._prompt_embeds is not None or batch is None: + return + task = batch.get("task") + prompt = task[0] if isinstance(task, list | tuple) else task + self._prompt = prompt or "" + self._prompt_embeds, self._negative_prompt_embeds = self._encode_prompt(self._prompt) + + def _get_t5_prompt_embeds(self, prompt, max_sequence_length): + tokenizer = self._frozen["tokenizer"] + text_encoder = self._frozen["text_encoder"] + device = self.config.device + + prompt = [prompt] if isinstance(prompt, str) else prompt + prompt = [clean_prompt(u) for u in prompt] + + text_inputs = tokenizer( + prompt, + padding="max_length", + max_length=max_sequence_length, + truncation=True, + add_special_tokens=True, + return_attention_mask=True, + return_tensors="pt", + ) + text_input_ids, mask = text_inputs.input_ids, text_inputs.attention_mask + seq_lens = mask.gt(0).sum(dim=1).long() + + te_device = next(text_encoder.parameters()).device + prompt_embeds = text_encoder(text_input_ids.to(te_device), mask.to(te_device)).last_hidden_state + prompt_embeds = prompt_embeds.to(dtype=self.dtype, device=device) + prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens, strict=False)] + prompt_embeds = torch.stack( + [torch.cat([u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))]) for u in prompt_embeds], + dim=0, + ) + return prompt_embeds.to(device) + + def _encode_prompt(self, prompt): + max_len = self.config.max_sequence_length + prompt_embeds = self._get_t5_prompt_embeds(prompt, max_len) + negative_prompt_embeds = None + if self._use_cfg: + negative_prompt_embeds = self._get_t5_prompt_embeds("", max_len) + return prompt_embeds, negative_prompt_embeds + + # Observation (image) encoding -> normalized video latents + def _extract_raw_obs(self, batch) -> dict[str, Tensor]: + """Snapshot the configured camera images from a batch (kept raw for later VAE encoding).""" + return {k: batch[k].detach() for k in self.config.obs_cam_keys} + + def _camera_frame(self, raw_obs, key, size=None) -> Tensor: + """Return a single-frame camera tensor [1, C, 1, H, W] resized + scaled to [-1, 1].""" + img = raw_obs[key] + if img.dim() == 3: # [C, H, W] + img = img.unsqueeze(0) + # LeRobot images arrive as float in [0, 1], shape [B, C, H, W]. + img = img.to(self.config.device, torch.float32) + if self.config.image_hflip: + img = torch.flip(img, dims=[-1]) # undo the env processor's horizontal flip + if size is None: + size = (self.config.height, self.config.width) + img = F.interpolate(img, size=size, mode="bilinear", align_corners=False) + img = img * 2.0 - 1.0 + return img.unsqueeze(2).to(self.dtype) # [1, C, F=1, H, W] + + def _normalize_vae_latent(self, enc_out: Tensor) -> Tensor: + """Take the mean of a VAE encoder output and channel-normalize it (matches upstream).""" + mu, _logvar = torch.chunk(enc_out, 2, dim=1) + latents_mean = torch.tensor(self._vae.config.latents_mean).to(mu.device) + latents_std = torch.tensor(self._vae.config.latents_std).to(mu.device) + mean = latents_mean.view(1, -1, 1, 1, 1) + inv_std = (1.0 / latents_std).view(1, -1, 1, 1, 1) + return ((mu.float() - mean) * inv_std).to(mu) + + @torch.no_grad() + def _encode_frames(self, raw_frames: list) -> Tensor: + """VAE-encode a temporal clip of observed frames and concat the per-camera latents on width. + + ``raw_frames`` is a list of per-frame obs dicts (one per env sub-step). Each configured + camera is stacked along the temporal axis into a ``[1, C, F, H, W]`` clip and encoded in a + single streaming ``encode_chunk`` call so the VAE temporal downsample (x4) collapses the F + input frames into ``F / 4`` latent frames, with the causal ``feat_cache`` carried across + chunks (mirrors upstream ``_encode_obs``). + """ + vae_device = next(self._vae.parameters()).device + if self.config.camera_layout == "robotwin_tshape": + return self._encode_frames_tshape(raw_frames, vae_device) + per_cam_videos = [] + for k in self.config.obs_cam_keys: + frames = [self._camera_frame(fb, k) for fb in raw_frames] + per_cam_videos.append(torch.cat(frames, dim=2)) # [1, C, F, H, W] + videos = torch.cat(per_cam_videos, dim=0) # [num_cam, C, F, H, W] + enc_out = self._streaming_vae.encode_chunk(videos.to(vae_device).to(self.dtype)) + mu_norm = self._normalize_vae_latent(enc_out) + # Concatenate the per-camera latents along width. + video_latent = torch.cat(mu_norm.split(1, dim=0), dim=-1) + return video_latent.to(self.config.device) + + @torch.no_grad() + def _encode_frames_tshape(self, raw_frames: list, vae_device) -> Tensor: + """RoboTwin T-shape latent assembly: full-res head + half-res wrists (second streaming VAE). + + The two wrist latents are concatenated on width and stacked (on the height axis) on top of + the head latent, mirroring upstream ``_encode_obs`` for ``env_type='robotwin_tshape'``. + """ + cfg = self.config + h, w = cfg.height, cfg.width + head_key, left_key, right_key = cfg.obs_cam_keys[0], cfg.obs_cam_keys[1], cfg.obs_cam_keys[2] + head = torch.cat([self._camera_frame(fb, head_key, size=(h, w)) for fb in raw_frames], dim=2) + left = torch.cat( + [self._camera_frame(fb, left_key, size=(h // 2, w // 2)) for fb in raw_frames], dim=2 + ) + right = torch.cat( + [self._camera_frame(fb, right_key, size=(h // 2, w // 2)) for fb in raw_frames], dim=2 + ) + wrists = torch.cat([left, right], dim=0) # [2, C, F, H/2, W/2] + enc_high = self._streaming_vae.encode_chunk(head.to(vae_device).to(self.dtype)) + enc_lr = self._frozen["streaming_vae_half"].encode_chunk(wrists.to(vae_device).to(self.dtype)) + # wrists side-by-side on width, then stacked on top of the head latent on the height axis. + enc_out = torch.cat([torch.cat(enc_lr.split(1, dim=0), dim=-1), enc_high], dim=-2) + video_latent = self._normalize_vae_latent(enc_out) + return video_latent.to(self.config.device) + + # KV cache management + @property + def _latent_hw(self): + if self.config.camera_layout == "robotwin_tshape": + # head (full) on the bottom, two half-res wrists side-by-side on top -> 1.5x height. + return ((self.config.height // 16) * 3) // 2, self.config.width // 16 + h = self.config.height // 16 + w = (self.config.width // 16) * len(self.config.obs_cam_keys) + return h, w + + def _init_streaming_cache(self, init_latent): + cfg = self.config + latent_h, latent_w = self._latent_hw + p = cfg.patch_size + latent_token_per_chunk = (cfg.frame_chunk_size * latent_h * latent_w) // (p[0] * p[1] * p[2]) + action_token_per_chunk = cfg.frame_chunk_size * cfg.action_per_frame + self.transformer.create_empty_cache( + "pos", + cfg.attn_window, + latent_token_per_chunk, + action_token_per_chunk, + device=self.config.device, + dtype=self.dtype, + batch_size=2 if self._use_cfg else 1, + ) + self._cache_initialised = True + + def _repeat_input_for_cfg(self, input_dict): + if self._use_cfg: + input_dict["noisy_latents"] = input_dict["noisy_latents"].repeat(2, 1, 1, 1, 1) + input_dict["text_emb"] = torch.cat( + [ + self._prompt_embeds.to(self.dtype).clone(), + self._negative_prompt_embeds.to(self.dtype).clone(), + ], + dim=0, + ) + input_dict["grid_id"] = input_dict["grid_id"][None].repeat(2, 1, 1) + input_dict["timesteps"] = input_dict["timesteps"][None].repeat(2, 1) + else: + input_dict["grid_id"] = input_dict["grid_id"][None] + input_dict["timesteps"] = input_dict["timesteps"][None] + return input_dict + + def _prepare_latent_input( + self, + latent_model_input, + action_model_input, + latent_t=0, + action_t=0, + latent_cond=None, + action_cond=None, + frame_st_id=0, + ): + cfg = self.config + device = self.config.device + p = cfg.patch_size + out = {} + if latent_model_input is not None: + out["latent_res_lst"] = { + "noisy_latents": latent_model_input, + "timesteps": torch.ones([latent_model_input.shape[2]], dtype=torch.float32, device=device) + * latent_t, + "grid_id": get_mesh_id( + latent_model_input.shape[-3] // p[0], + latent_model_input.shape[-2] // p[1], + latent_model_input.shape[-1] // p[2], + 0, + 1, + frame_st_id, + ).to(device), + "text_emb": self._prompt_embeds.to(self.dtype).clone(), + } + if latent_cond is not None: + out["latent_res_lst"]["noisy_latents"][:, :, 0:1] = latent_cond[:, :, 0:1] + out["latent_res_lst"]["timesteps"][0:1] *= 0 + if action_model_input is not None: + out["action_res_lst"] = { + "noisy_latents": action_model_input, + "timesteps": torch.ones([action_model_input.shape[2]], dtype=torch.float32, device=device) + * action_t, + "grid_id": get_mesh_id( + action_model_input.shape[-3], + action_model_input.shape[-2], + action_model_input.shape[-1], + 1, + 1, + frame_st_id, + action=True, + ).to(device), + "text_emb": self._prompt_embeds.to(self.dtype).clone(), + } + if action_cond is not None: + out["action_res_lst"]["noisy_latents"][:, :, 0:1] = action_cond[:, :, 0:1] + out["action_res_lst"]["timesteps"][0:1] *= 0 + out["action_res_lst"]["noisy_latents"][:, ~self._action_mask] *= 0 + return out + + @property + def _action_mask(self): + mask = torch.zeros([self.config.action_dim], dtype=torch.bool) + mask[self.config.used_action_channel_ids] = True + return mask + + # Action conditioning (executed action history) (de)normalization + def _preprocess_action_state(self, action_norm: Tensor) -> Tensor: + """Build the action-conditioning tensor from the already-normalized executed actions. + + ``action_norm`` is the model-space action chunk ``[B, action_dim, F, action_per_frame, 1]``. + Upstream re-derives the conditioning from the raw executed action via quantile norm; here + the executed actions are already in the model-normalized space, so we pass them through. + """ + return action_norm.to(self.config.device, self.dtype) + + def _compute_kv_cache(self, obs_buffer, executed_actions): + """Feed real observed keyframes + executed actions back into the KV cache.""" + if not obs_buffer or executed_actions is None: + return + self.transformer.clear_pred_cache("pos") + # Encode the buffered keyframe clip in one streaming call (carries the causal VAE cache). + latent_model_input = self._encode_frames(obs_buffer) + # On the first feedback, prepend the init latent so the latent/action frame counts align + # (upstream prepends ``init_latent`` to the observed keyframes when frame_st_id == 0). + if self._frame_st_id == 0 and getattr(self, "_init_latent", None) is not None: + latent_model_input = torch.cat([self._init_latent, latent_model_input], dim=2) + action_model_input = self._preprocess_action_state(executed_actions) + action_model_input = action_model_input.to(latent_model_input) + input_dict = self._prepare_latent_input( + latent_model_input, action_model_input, frame_st_id=self._frame_st_id + ) + with torch.no_grad(): + self.transformer( + self._repeat_input_for_cfg(input_dict["latent_res_lst"]), + update_cache=2, + cache_name="pos", + action_mode=False, + ) + self.transformer( + self._repeat_input_for_cfg(input_dict["action_res_lst"]), + update_cache=2, + cache_name="pos", + action_mode=True, + ) + self._frame_st_id += latent_model_input.shape[2] + + # The core dual-stream denoising loop (one chunk) + @torch.no_grad() + def _infer(self, init_latent, frame_st_id=0): + cfg = self.config + device = self.config.device + latent_h, latent_w = self._latent_hw + frame_chunk_size = cfg.frame_chunk_size + + latents = torch.randn(1, 48, frame_chunk_size, latent_h, latent_w, device=device, dtype=self.dtype) + actions = torch.randn( + 1, cfg.action_dim, frame_chunk_size, cfg.action_per_frame, 1, device=device, dtype=self.dtype + ) + + self._scheduler.set_timesteps(cfg.num_inference_steps) + self._action_scheduler.set_timesteps(cfg.action_num_inference_steps) + timesteps = F.pad(self._scheduler.timesteps, (0, 1), mode="constant", value=0) + if cfg.video_exec_step != -1: + timesteps = timesteps[: cfg.video_exec_step] + action_timesteps = F.pad(self._action_scheduler.timesteps, (0, 1), mode="constant", value=0) + + # 1. Video-latent denoising loop + for i, t in enumerate(timesteps): + last_step = i == len(timesteps) - 1 + latent_cond = ( + init_latent[:, :, 0:1].to(self.dtype) + if frame_st_id == 0 and init_latent is not None + else None + ) + input_dict = self._prepare_latent_input( + latents, None, t, t, latent_cond, None, frame_st_id=frame_st_id + ) + video_noise_pred = self.transformer( + self._repeat_input_for_cfg(input_dict["latent_res_lst"]), + update_cache=1 if last_step else 0, + cache_name="pos", + action_mode=False, + ) + if not last_step or cfg.video_exec_step != -1: + video_noise_pred = data_seq_to_patch( + cfg.patch_size, + video_noise_pred, + frame_chunk_size, + latent_h, + latent_w, + batch_size=2 if self._use_cfg else 1, + ) + if cfg.guidance_scale > 1: + video_noise_pred = video_noise_pred[1:] + cfg.guidance_scale * ( + video_noise_pred[:1] - video_noise_pred[1:] + ) + else: + video_noise_pred = video_noise_pred[:1] + latents = self._scheduler.step(video_noise_pred, t, latents, return_dict=False) + if frame_st_id == 0 and latent_cond is not None: + latents[:, :, 0:1] = latent_cond + + # 2. Action denoising loop + for i, t in enumerate(action_timesteps): + last_step = i == len(action_timesteps) - 1 + action_cond = ( + torch.zeros([1, cfg.action_dim, 1, cfg.action_per_frame, 1], device=device, dtype=self.dtype) + if frame_st_id == 0 + else None + ) + input_dict = self._prepare_latent_input( + None, actions, t, t, None, action_cond, frame_st_id=frame_st_id + ) + action_noise_pred = self.transformer( + self._repeat_input_for_cfg(input_dict["action_res_lst"]), + update_cache=1 if last_step else 0, + cache_name="pos", + action_mode=True, + ) + if not last_step: + action_noise_pred = rearrange(action_noise_pred, "b (f n) c -> b c f n 1", f=frame_chunk_size) + if cfg.action_guidance_scale > 1: + action_noise_pred = action_noise_pred[1:] + cfg.action_guidance_scale * ( + action_noise_pred[:1] - action_noise_pred[1:] + ) + else: + action_noise_pred = action_noise_pred[:1] + actions = self._action_scheduler.step(action_noise_pred, t, actions, return_dict=False) + if frame_st_id == 0 and action_cond is not None: + actions[:, :, 0:1] = action_cond + + actions[:, ~self._action_mask] *= 0 + return actions, latents + + # Predicted-video decoding (opt-in) + @torch.no_grad() + def decode_predicted_latents(self, latents) -> Tensor: + """Decode a concatenated predicted-latent sequence into ``[T, H, W, 3]`` uint8 frames.""" + return self._decode_predicted_video(latents) + + @torch.no_grad() + def _decode_predicted_video(self, latents) -> Tensor: + """VAE-decode predicted latents into a uint8 frame stack ``[T, H, W, 3]`` on CPU.""" + vae = self._vae + z_dim = vae.config.z_dim + vae_device = next(vae.parameters()).device + latents = latents.to(device=vae_device, dtype=vae.dtype) + latents = denormalize_latents(latents, vae.config.latents_mean, vae.config.latents_std, z_dim) + video = vae.decode(latents, return_dict=False)[0] # [B, C, F, H, W] in [-1, 1] + video = (video.float().clamp(-1, 1) + 1.0) / 2.0 + video = (video[0].permute(1, 2, 3, 0) * 255.0).round().to(torch.uint8) # [F, H, W, C] + return video.cpu() diff --git a/src/lerobot/policies/lingbot_va/processor_lingbot_va.py b/src/lerobot/policies/lingbot_va/processor_lingbot_va.py new file mode 100644 index 000000000..119f77d5b --- /dev/null +++ b/src/lerobot/policies/lingbot_va/processor_lingbot_va.py @@ -0,0 +1,87 @@ +# Copyright 2026 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. + +"""Pre/post-processor pipelines for the LingBot-VA policy. + +The preprocessor passes inputs through (IDENTITY) and the postprocessor maps the policy's +``[-1, 1]`` actions back to physical units with the built-in ``UnnormalizerProcessorStep`` +(QUANTILES) using per-channel q01/q99 restored from the checkpoint. +""" + +from typing import Any + +import torch + +from lerobot.configs.types import FeatureType, NormalizationMode +from lerobot.processor import ( + AddBatchDimensionProcessorStep, + DeviceProcessorStep, + NormalizerProcessorStep, + PolicyAction, + PolicyProcessorPipeline, + ProcessorStep, + RenameObservationsProcessorStep, + UnnormalizerProcessorStep, +) +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 + + +def make_lingbot_va_pre_post_processors( + config: LingBotVAConfig, + dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None, +) -> tuple[ + PolicyProcessorPipeline[dict[str, Any], dict[str, Any]], + PolicyProcessorPipeline[PolicyAction, PolicyAction], +]: + """Build the pre/post processor pipelines for LingBot-VA.""" + + input_steps: list[ProcessorStep] = [ + 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. + output_steps: list[ProcessorStep] = [ + UnnormalizerProcessorStep( + features=config.output_features, + norm_map={FeatureType.ACTION: NormalizationMode.QUANTILES}, + 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, + ), + ) diff --git a/src/lerobot/policies/lingbot_va/utils.py b/src/lerobot/policies/lingbot_va/utils.py new file mode 100644 index 000000000..86f885549 --- /dev/null +++ b/src/lerobot/policies/lingbot_va/utils.py @@ -0,0 +1,1138 @@ +# Copyright 2024-2025 The Robbyant Team Authors. All rights reserved. +# Copyright 2026 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. + +"""Vendored Wan2.2 model code and plumbing for the LingBot-VA policy. + +Everything the policy builds on lives here: grid/patch reshaping, attention backends, +VAE (de)normalization + frozen-component loaders, the flow-matching scheduler, and the +dual-stream Wan transformer (``WanTransformer3DModel`` and its sub-modules). Only the +LeRobot-facing ``LingBotVAPolicy`` orchestrator stays in ``modeling_lingbot_va.py``; this +module imports nothing from it (one-directional dependency). +""" + +import html +import math +import re +from copy import deepcopy +from functools import partial +from typing import TYPE_CHECKING + +import torch +import torch.nn as nn +import torch.nn.functional as F # noqa: N812 +from einops import rearrange + +from lerobot.utils.import_utils import _diffusers_available, _transformers_available + +if TYPE_CHECKING or _diffusers_available: + from diffusers import AutoencoderKLWan + from diffusers.configuration_utils import ConfigMixin, register_to_config + from diffusers.models.attention import FeedForward + from diffusers.models.embeddings import PixArtAlphaTextProjection, TimestepEmbedding, Timesteps + from diffusers.models.modeling_utils import ModelMixin + from diffusers.models.normalization import FP32LayerNorm +else: + AutoencoderKLWan = FeedForward = PixArtAlphaTextProjection = None + TimestepEmbedding = Timesteps = FP32LayerNorm = None + + class ModelMixin: + pass + + class ConfigMixin: + pass + + def register_to_config(func): + return func + + +if TYPE_CHECKING or _transformers_available: + from transformers import T5TokenizerFast, UMT5EncoderModel +else: + T5TokenizerFast = UMT5EncoderModel = None + + +# Grid-id / patch utilities +def data_seq_to_patch(patch_size, data_seq, latent_num_frames, latent_height, latent_width, batch_size=1): + """Reshape a flattened patch sequence back into a ``(B, C, F, H, W)`` latent grid.""" + p_t, p_h, p_w = patch_size + post_patch_num_frames = latent_num_frames // p_t + post_patch_height = latent_height // p_h + post_patch_width = latent_width // p_w + + data_patch = data_seq.reshape( + batch_size, post_patch_num_frames, post_patch_height, post_patch_width, p_t, p_h, p_w, -1 + ) + data_patch = data_patch.permute(0, 7, 1, 4, 2, 5, 3, 6) + data_patch = data_patch.flatten(6, 7).flatten(4, 5).flatten(2, 3) + return data_patch + + +def get_mesh_id(f, h, w, t, f_w=1, f_shift=0, action=False): + """Build the (frame, height, width, stream) grid ids used to index the rotary embedding.""" + f_idx = torch.arange(f_shift, f + f_shift) * f_w + h_idx = torch.arange(h) + w_idx = torch.arange(w) + ff, hh, ww = torch.meshgrid(f_idx, h_idx, w_idx, indexing="ij") + if action: + ff_offset = (torch.ones([h]).cumsum(0) / (h + 1)).view(1, -1, 1) + ff = ff + ff_offset + hh = torch.ones_like(hh) * -1 + ww = torch.ones_like(ww) * -1 + + grid_id = torch.cat([ff.unsqueeze(0), hh.unsqueeze(0), ww.unsqueeze(0)], dim=0).flatten(1) + grid_id = torch.cat([grid_id, torch.full_like(grid_id[:1], t)], dim=0) + return grid_id + + +# Attention backends +def custom_sdpa(q, k, v): + """Scaled-dot-product attention operating on ``(B, S, H, D)`` tensors.""" + out = F.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)) + return out.transpose(1, 2) + + +def _load_flash_attn_func(): + try: + from flash_attn_interface import flash_attn_func + except ImportError: + try: + from flash_attn import flash_attn_func + except ImportError as e: + raise ImportError( + "attn_mode='flashattn' requires the `flash_attn` package, which is not installed. " + "Install it, or use attn_mode='torch' (the default)." + ) from e + return flash_attn_func + + +# Wan2.2 VAE helpers (stock diffusers ``AutoencoderKLWan``) +def _vae_patchify(x, patch_size): + if patch_size is None or patch_size == 1: + return x + batch_size, channels, frames, height, width = x.shape + x = x.view( + batch_size, channels, frames, height // patch_size, patch_size, width // patch_size, patch_size + ) + x = x.permute(0, 1, 6, 4, 2, 3, 5).contiguous() + x = x.view( + batch_size, channels * patch_size * patch_size, frames, height // patch_size, width // patch_size + ) + return x + + +def denormalize_latents(latents: torch.Tensor, latents_mean, latents_std, z_dim) -> torch.Tensor: + """Inverse of the encode-time latent normalization, for VAE-decoding predicted latents.""" + mean = torch.tensor(latents_mean).view(1, z_dim, 1, 1, 1).to(latents.device, latents.dtype) + inv_std = 1.0 / torch.tensor(latents_std).view(1, z_dim, 1, 1, 1).to(latents.device, latents.dtype) + return latents / inv_std + mean + + +def load_vae(vae_path, torch_dtype, torch_device, subfolder=None): + vae = AutoencoderKLWan.from_pretrained(vae_path, subfolder=subfolder, torch_dtype=torch_dtype) + return vae.to(torch_device) + + +def load_text_encoder(text_encoder_path, torch_dtype, torch_device, subfolder=None): + text_encoder = UMT5EncoderModel.from_pretrained( + text_encoder_path, subfolder=subfolder, torch_dtype=torch_dtype + ) + return text_encoder.to(torch_device) + + +def load_tokenizer(tokenizer_path, subfolder=None): + return T5TokenizerFast.from_pretrained(tokenizer_path, subfolder=subfolder) + + +# Misc +def clean_prompt(text: str) -> str: + """Normalize a task prompt (HTML-unescape + whitespace collapse). + + Mirrors diffusers' Wan ``prompt_clean`` minus ``ftfy.fix_text``, + which is a no-op for the ASCII task strings used here, so we avoid the extra ``ftfy`` dep. + """ + text = html.unescape(html.unescape(text)).strip() + return re.sub(r"\s+", " ", text).strip() + + +def _torch_dtype(name: str) -> torch.dtype: + return {"bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32}[name] + + +def _sample_timestep_id( + batch_size: int = 1, + min_timestep_bd: float = 0.0, + max_timestep_bd: float = 1.0, + num_train_timesteps: int = 1000, +) -> torch.Tensor: + """Sample per-frame flow-matching timestep ids (upstream ``utils.sample_timestep_id``).""" + u = torch.rand(size=[batch_size]) * (max_timestep_bd - min_timestep_bd) + min_timestep_bd + return (u * num_train_timesteps).clamp(min=0, max=num_train_timesteps - 1).to(torch.int64) + + +# Flow-matching scheduler +# LingBot-VA uses two independent instances at inference (one for the video-latent stream, +# one for the action stream), each with its own ``shift`` and number of denoising steps. +class FlowMatchScheduler: + def __init__( + self, + num_inference_steps=100, + num_train_timesteps=1000, + shift=3.0, + sigma_max=1.0, + sigma_min=0.003 / 1.002, + inverse_timesteps=False, + extra_one_step=False, + reverse_sigmas=False, + exponential_shift=False, + exponential_shift_mu=None, + shift_terminal=None, + ): + self.num_train_timesteps = num_train_timesteps + self.shift = shift + self.sigma_max = sigma_max + self.sigma_min = sigma_min + self.inverse_timesteps = inverse_timesteps + self.extra_one_step = extra_one_step + self.reverse_sigmas = reverse_sigmas + self.exponential_shift = exponential_shift + self.exponential_shift_mu = exponential_shift_mu + self.shift_terminal = shift_terminal + self.set_timesteps(num_inference_steps) + + def set_timesteps( + self, + num_inference_steps=100, + denoising_strength=1.0, + training=False, + shift=None, + dynamic_shift_len=None, + ): + if shift is not None: + self.shift = shift + sigma_start = self.sigma_min + (self.sigma_max - self.sigma_min) * denoising_strength + if self.extra_one_step: + self.sigmas = torch.linspace(sigma_start, self.sigma_min, num_inference_steps + 1)[:-1] + else: + self.sigmas = torch.linspace(sigma_start, self.sigma_min, num_inference_steps) + if self.inverse_timesteps: + self.sigmas = torch.flip(self.sigmas, dims=[0]) + if self.exponential_shift: + mu = ( + self.calculate_shift(dynamic_shift_len) + if dynamic_shift_len is not None + else self.exponential_shift_mu + ) + self.sigmas = math.exp(mu) / (math.exp(mu) + (1 / self.sigmas - 1)) + else: + self.sigmas = self.shift * self.sigmas / (1 + (self.shift - 1) * self.sigmas) + if self.shift_terminal is not None: + one_minus_z = 1 - self.sigmas + scale_factor = one_minus_z[-1] / (1 - self.shift_terminal) + self.sigmas = 1 - (one_minus_z / scale_factor) + if self.reverse_sigmas: + self.sigmas = 1 - self.sigmas + self.timesteps = self.sigmas * self.num_train_timesteps + if training: + x = self.timesteps + y = torch.exp(-2 * ((x - num_inference_steps / 2) / num_inference_steps) ** 2) + y_shifted = y - y.min() + bsmntw_weighing = y_shifted * (num_inference_steps / y_shifted.sum()) + self.linear_timesteps_weights = bsmntw_weighing + self.training = True + else: + self.training = False + + def step(self, model_output, timestep, sample, to_final=False, **kwargs): + if isinstance(timestep, torch.Tensor): + timestep = timestep.cpu() + timestep_id = torch.argmin((self.timesteps - timestep).abs()) + sigma = self.sigmas[timestep_id] + if to_final or timestep_id + 1 >= len(self.timesteps): + sigma_ = 1 if (self.inverse_timesteps or self.reverse_sigmas) else 0 + else: + sigma_ = self.sigmas[timestep_id + 1] + prev_sample = sample + model_output * (sigma_ - sigma) + return prev_sample + + def add_noise(self, original_samples, noise, timestep, t_dim=2): + if isinstance(timestep, torch.Tensor): + timestep = timestep.cpu() + timestep = timestep[None] + timestep_id = torch.argmin((self.timesteps[:, None] - timestep).abs(), dim=0) + shape = [1] * noise.ndim + shape[t_dim] = timestep_id.shape[0] + sigma = self.sigmas[timestep_id].to(original_samples).view(shape) + sample = (1 - sigma) * original_samples + sigma * noise + return sample + + def training_target(self, sample, noise, timestep): + target = noise - sample + return target + + def training_weight(self, timestep): + timestep_id = torch.argmin( + (self.timesteps[:, None].to(timestep.device) - timestep[None]).abs(), dim=0 + ) + weights = self.linear_timesteps_weights.to(timestep.device)[timestep_id].to(timestep.device) + return weights + + def calculate_shift( + self, + image_seq_len, + base_seq_len: int = 256, + max_seq_len: int = 8192, + base_shift: float = 0.5, + max_shift: float = 0.9, + ): + m = (max_shift - base_shift) / (max_seq_len - base_seq_len) + b = base_shift - m * base_seq_len + mu = image_seq_len * m + b + return mu + + +class FlexAttnFunc(nn.Module): + """Flex-attention backend (training only; ``attn_mode='flex'``). + + Builds the block-causal / window / noise-vs-clean masks used by the dual-stream + flow-matching training. Inference uses the ``torch`` SDPA backend. The flex-attention + APIs and their ``torch.compile`` wrappers are imported/initialised lazily so importing + this module never requires a flex-attention-capable PyTorch build. + """ + + flex_attn = None + compiled_create_block_mask = None + attention_mask = None + cross_attention_mask = None + + def __init__(self, is_cross=False) -> None: + super().__init__() + self.is_cross = is_cross + + @classmethod + def _ensure_compiled(cls): + if cls.flex_attn is None: + from torch.nn.attention.flex_attention import create_block_mask, flex_attention + + cls.flex_attn = torch.compile(flex_attention, dynamic=True) + cls.compiled_create_block_mask = torch.compile(create_block_mask) + + def forward(self, query, key, value, dtype=torch.bfloat16): + self._ensure_compiled() + q_varlen = rearrange(query[0], "s n d -> 1 n s d") + k_varlen = rearrange(key[0], "s n d -> 1 n s d") + v_varlen = rearrange(value[0], "s n d -> 1 n s d") + + half_dtypes = (torch.float16, torch.bfloat16) + if dtype not in half_dtypes: + raise ValueError(f"Flex attention requires a half-precision dtype, got {dtype}.") + + def half(x): + return x if x.dtype in half_dtypes else x.to(dtype) + + q_varlen = half(q_varlen) + k_varlen = half(k_varlen) + v_varlen = half(v_varlen) + q_varlen = q_varlen.to(v_varlen.dtype) + k_varlen = k_varlen.to(v_varlen.dtype) + + block_mask = FlexAttnFunc.cross_attention_mask if self.is_cross else FlexAttnFunc.attention_mask + + x_out = FlexAttnFunc.flex_attn( + q_varlen, + k_varlen, + v_varlen, + block_mask=block_mask, + kernel_options={ + "BLOCK_M": 64, + "BLOCK_N": 64, + "BLOCK_M1": 32, + "BLOCK_N1": 64, + "BLOCK_M2": 64, + "BLOCK_N2": 32, + }, + ) + + x_out = rearrange(x_out, "b n s d -> b s n d") + return x_out + + @staticmethod + @torch.no_grad() + def init_mask(latent_shape, action_shape, padded_length, chunk_size, window_size, patch_size, device): + FlexAttnFunc._ensure_compiled() + torch._inductor.config.realize_opcount_threshold = 100 + b, _, l_f, l_h, l_w = latent_shape + _, _, a_f, a_h, a_w = action_shape + + latent_seq_id = ( + torch.arange(b)[:, None, None, None] + .expand(-1, l_f // patch_size[0], l_h // patch_size[1], l_w // patch_size[2]) + .flatten() + ) + action_seq_id = torch.arange(b)[:, None, None, None].expand(-1, a_f, a_h, a_w).flatten() + seq_ids = torch.cat([latent_seq_id] * 2 + [action_seq_id] * 2) + + latent_frame_id = ( + torch.arange(l_f)[None, :, None, None] + .expand(b, -1, l_h // patch_size[1], l_w // patch_size[2])[None] + .flatten() + ) + action_frame_id = torch.arange(a_f)[None, :, None, None].expand(b, -1, a_h, a_w)[None].flatten() + frame_ids = torch.cat( + [latent_frame_id // chunk_size * 2] * 2 + [action_frame_id // chunk_size * 2 + 1] * 2 + ) + + noise_ids = torch.cat( + [ + torch.zeros_like(latent_frame_id), + torch.ones_like(latent_frame_id), + torch.zeros_like(action_frame_id), + torch.ones_like(action_frame_id), + ] + ) + + seq_ids = F.pad(seq_ids, (0, padded_length), value=-1) + frame_ids = F.pad(frame_ids, (0, padded_length), value=-1) + noise_ids = F.pad(noise_ids, (0, padded_length), value=-1) + + mask_mod = FlexAttnFunc._get_mask_mod( + seq_ids.long().to(device), frame_ids.long().to(device), noise_ids.long().to(device), window_size + ) + block_mask = FlexAttnFunc.compiled_create_block_mask( + mask_mod, 1, 1, len(seq_ids), len(seq_ids), device=device, _compile=True + ) + FlexAttnFunc.attention_mask = block_mask + + text_seq_ids = torch.arange(b)[:, None].expand(-1, 512).flatten() + mask_mod_cross = FlexAttnFunc._get_cross_mask_mod( + seq_ids.long().to(device), text_seq_ids.long().to(device) + ) + block_mask_cross = FlexAttnFunc.compiled_create_block_mask( + mask_mod_cross, 1, 1, len(seq_ids), len(text_seq_ids), device=device, _compile=True + ) + FlexAttnFunc.cross_attention_mask = block_mask_cross + + @staticmethod + @torch.no_grad() + def _get_cross_mask_mod(seq_ids, text_seq_ids): + def seq_mask(b, h, q_idx, kv_idx): + return ( + (seq_ids[q_idx] == text_seq_ids[kv_idx]) & (seq_ids[q_idx] >= 0) & (text_seq_ids[kv_idx] >= 0) + ) + + return seq_mask + + @staticmethod + @torch.no_grad() + def _get_mask_mod(seq_ids, frame_ids, noise_ids, window_size): + from torch.nn.attention.flex_attention import and_masks, or_masks + + def seq_mask(b, h, q_idx, kv_idx): + return (seq_ids[q_idx] == seq_ids[kv_idx]) & (seq_ids[q_idx] >= 0) & (seq_ids[kv_idx] >= 0) + + def block_causal_mask(b, h, q_idx, kv_idx): + return frame_ids[kv_idx] <= frame_ids[q_idx] + + def block_causal_mask_exclude_self(b, h, q_idx, kv_idx): + return frame_ids[kv_idx] < frame_ids[q_idx] + + def block_self_mask(b, h, q_idx, kv_idx): + return frame_ids[kv_idx] == frame_ids[q_idx] + + def clean2clean_mask(b, h, q_idx, kv_idx): + return (noise_ids[q_idx] == 1) & (noise_ids[kv_idx] == 1) + + def noise2clean_mask(b, h, q_idx, kv_idx): + return (noise_ids[q_idx] == 0) & (noise_ids[kv_idx] == 1) + + def noise2noise_mask(b, h, q_idx, kv_idx): + return (noise_ids[q_idx] == 0) & (noise_ids[kv_idx] == 0) + + def block_window_mask(b, h, q_idx, kv_idx, window_size: int): + return (frame_ids[q_idx] - frame_ids[kv_idx]).abs() <= window_size + + mask_list = [] + mask_list.append(and_masks(clean2clean_mask, block_causal_mask)) + mask_list.append(and_masks(noise2clean_mask, block_causal_mask_exclude_self)) + mask_list.append(and_masks(noise2noise_mask, block_self_mask)) + mask = or_masks(*mask_list) + mask = and_masks(mask, seq_mask) + mask = and_masks(mask, partial(block_window_mask, window_size=window_size)) + return mask + + +class WanRotaryPosEmbed(nn.Module): + """Rotary position embedding with separate frequency bases for frame / height / width.""" + + def __init__(self, attention_head_dim: int, patch_size, max_seq_len: int, theta: float = 10000.0): + super().__init__() + + self.attention_head_dim = attention_head_dim + self.patch_size = patch_size + self.max_seq_len = max_seq_len + self.theta = theta + + self.f_dim = self.attention_head_dim - 2 * (self.attention_head_dim // 3) + self.h_dim = self.attention_head_dim // 3 + self.w_dim = self.attention_head_dim // 3 + + f_freqs_base, h_freqs_base, w_freqs_base = self._precompute_freqs_base() + self.f_freqs_base = f_freqs_base + self.h_freqs_base = h_freqs_base + self.w_freqs_base = w_freqs_base + + def _precompute_freqs_base(self): + f_freqs_base = 1.0 / ( + self.theta ** (torch.arange(0, self.f_dim, 2)[: (self.f_dim // 2)].double() / self.f_dim) + ) + h_freqs_base = 1.0 / ( + self.theta ** (torch.arange(0, self.h_dim, 2)[: (self.h_dim // 2)].double() / self.h_dim) + ) + w_freqs_base = 1.0 / ( + self.theta ** (torch.arange(0, self.w_dim, 2)[: (self.w_dim // 2)].double() / self.w_dim) + ) + return f_freqs_base, h_freqs_base, w_freqs_base + + def forward(self, grid_ids): + with torch.no_grad(): + f_freqs = grid_ids[:, 0, :].unsqueeze(-1) * self.f_freqs_base.to(grid_ids.device) + h_freqs = grid_ids[:, 1, :].unsqueeze(-1) * self.h_freqs_base.to(grid_ids.device) + w_freqs = grid_ids[:, 2, :].unsqueeze(-1) * self.w_freqs_base.to(grid_ids.device) + freqs = torch.cat([f_freqs, h_freqs, w_freqs], dim=-1).float() + freqs_cis = torch.polar(torch.ones_like(freqs), freqs) + + return freqs_cis + + +class WanAttention(nn.Module): + """Self/cross attention with KV-caching for autoregressive streaming inference. + + Backends: ``torch`` (default SDPA), ``flashattn`` (optional), ``flex`` (training masks). + """ + + def __init__( + self, + dim, + heads=8, + dim_head=64, + eps=1e-5, + dropout=0.0, + cross_attention_dim_head=None, + attn_mode="torch", + ): + super().__init__() + if attn_mode == "torch": + self.attn_op = custom_sdpa + elif attn_mode == "flashattn": + self.attn_op = _load_flash_attn_func() + elif attn_mode == "flex": + self.attn_op = FlexAttnFunc(cross_attention_dim_head is not None) + else: + raise ValueError( + f"Unsupported attention mode: {attn_mode}, only support 'torch', 'flashattn' and 'flex'" + ) + + self.inner_dim = dim_head * heads + self.heads = heads + self.cross_attention_dim_head = cross_attention_dim_head + self.kv_inner_dim = ( + self.inner_dim if cross_attention_dim_head is None else cross_attention_dim_head * heads + ) + + self.to_q = nn.Linear(dim, self.inner_dim, bias=True) + self.to_k = nn.Linear(dim, self.kv_inner_dim, bias=True) + self.to_v = nn.Linear(dim, self.kv_inner_dim, bias=True) + self.to_out = nn.ModuleList([nn.Linear(self.inner_dim, dim, bias=True), nn.Dropout(dropout)]) + self.norm_q = nn.RMSNorm(dim_head * heads, eps=eps, elementwise_affine=True) + self.norm_k = nn.RMSNorm(dim_head * heads, eps=eps, elementwise_affine=True) + # KV cache only lives on self-attention modules (cross_attention_dim_head is None). + self.attn_caches = {} if cross_attention_dim_head is None else None + + def clear_pred_cache(self, cache_name): + if self.attn_caches is None: + return + cache = self.attn_caches[cache_name] + is_pred = cache["is_pred"] + cache["mask"][is_pred] = False + + def clear_cache(self, cache_name): + if self.attn_caches is None: + return + self.attn_caches[cache_name] = None + + def init_kv_cache(self, cache_name, total_token_len, num_head, head_dim, device, dtype, batch_size): + if self.attn_caches is None: + return + self.attn_caches[cache_name] = { + "k": torch.empty([batch_size, total_token_len, num_head, head_dim], device=device, dtype=dtype), + "v": torch.empty([batch_size, total_token_len, num_head, head_dim], device=device, dtype=dtype), + "id": torch.full((total_token_len,), -1, device=device), + "mask": torch.zeros((total_token_len,), dtype=torch.bool, device=device), + "is_pred": torch.zeros((total_token_len,), dtype=torch.bool, device=device), + } + + def allocate_slots(self, cache_name, key_size): + cache = self.attn_caches[cache_name] + mask = cache["mask"] + ids = cache["id"] + free = (~mask).nonzero(as_tuple=False).squeeze(-1) + + if free.numel() < key_size: + used = mask.nonzero(as_tuple=False).squeeze(-1) + + used_ids = ids[used] + order = torch.argsort(used_ids) + need = key_size - free.numel() + to_free = used[order[:need]] + + mask[to_free] = False + ids[to_free] = -1 + free = (~mask).nonzero(as_tuple=False).squeeze(-1) + + if free.numel() < key_size: + raise RuntimeError(f"KV cache exhausted: need {key_size} free slots, have {free.numel()}.") + return free[:key_size] + + def _next_cache_id(self, cache_name): + ids = self.attn_caches[cache_name]["id"] + mask = self.attn_caches[cache_name]["mask"] + + if mask.any(): + return ids[mask].max() + 1 + else: + return torch.tensor(0, device=ids.device, dtype=ids.dtype) + + def update_cache(self, cache_name, key, value, is_pred): + cache = self.attn_caches[cache_name] + + key_size = key.shape[1] + slots = self.allocate_slots(cache_name, key_size) + + new_id = self._next_cache_id(cache_name) + + cache["k"][:, slots] = key + cache["v"][:, slots] = value + cache["mask"][slots] = True + cache["id"][slots] = new_id + cache["is_pred"][slots] = is_pred + return slots + + def restore_cache(self, cache_name, slots): + self.attn_caches[cache_name]["mask"][slots] = False + + def forward(self, q, k, v, rotary_emb, update_cache=0, cache_name="pos"): + kv_cache = ( + self.attn_caches[cache_name] + if (self.attn_caches is not None) and (cache_name in self.attn_caches) + else None + ) + + query, key, value = self.to_q(q), self.to_k(k), self.to_v(v) + query = self.norm_q(query) + query = query.unflatten(2, (self.heads, -1)) + key = self.norm_k(key) + key = key.unflatten(2, (self.heads, -1)) + value = value.unflatten(2, (self.heads, -1)) + if rotary_emb is not None: + + def apply_rotary_emb(x, freqs): + x_out = torch.view_as_complex( + x.to(torch.float64).reshape(x.shape[0], x.shape[1], x.shape[2], -1, 2) + ) + x_out = torch.view_as_real(x_out * freqs).flatten(3) + return x_out.to(x.dtype) + + query = apply_rotary_emb(query, rotary_emb) + key = apply_rotary_emb(key, rotary_emb) + slots = None + if kv_cache is not None and kv_cache["k"] is not None: + slots = self.update_cache(cache_name, key, value, is_pred=(update_cache == 1)) + key_pool = self.attn_caches[cache_name]["k"] + value_pool = self.attn_caches[cache_name]["v"] + mask = self.attn_caches[cache_name]["mask"] + valid = mask.nonzero(as_tuple=False).squeeze(-1) + key = key_pool[:, valid] + value = value_pool[:, valid] + + hidden_states = self.attn_op(query, key, value) + + if update_cache == 0 and kv_cache is not None and kv_cache["k"] is not None: + self.restore_cache(cache_name, slots) + + hidden_states = hidden_states.flatten(2, 3) + hidden_states = hidden_states.type_as(query) + hidden_states = self.to_out[0](hidden_states) + hidden_states = self.to_out[1](hidden_states) + return hidden_states + + +# Dual-stream Wan2.2 transformer +class WanTimeTextImageEmbedding(nn.Module): + def __init__(self, dim, time_freq_dim, time_proj_dim, text_embed_dim, pos_embed_seq_len): + super().__init__() + + self.timesteps_proj = Timesteps( + num_channels=time_freq_dim, flip_sin_to_cos=True, downscale_freq_shift=0 + ) + self.time_embedder = TimestepEmbedding(in_channels=time_freq_dim, time_embed_dim=dim) + self.act_fn = nn.SiLU() + self.time_proj = nn.Linear(dim, time_proj_dim) + self.text_embedder = PixArtAlphaTextProjection(text_embed_dim, dim, act_fn="gelu_tanh") + + def forward(self, timestep: torch.Tensor, dtype=None): + b, seq_len = timestep.shape + timestep = timestep.reshape(-1) + timestep = self.timesteps_proj(timestep) + time_embedder_dtype = self.time_embedder.linear_1.weight.dtype + if timestep.dtype != time_embedder_dtype and time_embedder_dtype != torch.int8: + timestep = timestep.to(time_embedder_dtype) + temb = self.time_embedder(timestep).to(dtype=dtype) + timestep_proj = self.time_proj(self.act_fn(temb)) + return temb.reshape(b, seq_len, -1), timestep_proj.reshape(b, seq_len, -1) + + +class WanTransformerBlock(nn.Module): + def __init__(self, dim, ffn_dim, num_heads, cross_attn_norm=False, eps=1e-6, attn_mode: str = "torch"): + super().__init__() + self.attn_mode = attn_mode + + # 1. Self-attention + self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) + self.attn1 = WanAttention( + dim=dim, + heads=num_heads, + dim_head=dim // num_heads, + eps=eps, + cross_attention_dim_head=None, + attn_mode=attn_mode, + ) + + # 2. Cross-attention + self.attn2 = WanAttention( + dim=dim, + heads=num_heads, + dim_head=dim // num_heads, + eps=eps, + cross_attention_dim_head=dim // num_heads, + attn_mode=attn_mode, + ) + self.norm2 = FP32LayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity() + + # 3. Feed-forward + self.ffn = FeedForward(dim, inner_dim=ffn_dim, activation_fn="gelu-approximate") + self.norm3 = FP32LayerNorm(dim, eps, elementwise_affine=False) + + self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + + def forward( + self, hidden_states, encoder_hidden_states, temb, rotary_emb, update_cache=0, cache_name="pos" + ) -> torch.Tensor: + temb_scale_shift_table = self.scale_shift_table[None] + temb.float() + shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = rearrange( + temb_scale_shift_table, "b l n c -> b n l c" + ).chunk(6, dim=1) + shift_msa = shift_msa.squeeze(1) + scale_msa = scale_msa.squeeze(1) + gate_msa = gate_msa.squeeze(1) + c_shift_msa = c_shift_msa.squeeze(1) + c_scale_msa = c_scale_msa.squeeze(1) + c_gate_msa = c_gate_msa.squeeze(1) + # 1. Self-attention + norm_hidden_states = (self.norm1(hidden_states.float()) * (1.0 + scale_msa) + shift_msa).type_as( + hidden_states + ) + attn_output = self.attn1( + norm_hidden_states, + norm_hidden_states, + norm_hidden_states, + rotary_emb, + update_cache=update_cache, + cache_name=cache_name, + ) + hidden_states = (hidden_states.float() + attn_output * gate_msa).type_as(hidden_states) + + # 2. Cross-attention + norm_hidden_states = self.norm2(hidden_states.float()).type_as(hidden_states) + attn_output = self.attn2( + norm_hidden_states, + encoder_hidden_states, + encoder_hidden_states, + None, + update_cache=0, + cache_name=cache_name, + ) + hidden_states = hidden_states + attn_output + + # 3. Feed-forward + norm_hidden_states = (self.norm3(hidden_states.float()) * (1.0 + c_scale_msa) + c_shift_msa).type_as( + hidden_states + ) + + ff_output = self.ffn(norm_hidden_states) + + hidden_states = (hidden_states.float() + ff_output.float() * c_gate_msa).type_as(hidden_states) + return hidden_states + + +class WanTransformer3DModel(ModelMixin, ConfigMixin): + """Dual-stream (video + action) Wan2.2 DiT backbone with autoregressive KV caching.""" + + _supports_gradient_checkpointing = True + _skip_layerwise_casting_patterns = [ + "patch_embedding_mlp", + "condition_embedder", + "condition_embedder_action", + "norm", + ] + _no_split_modules = ["WanTransformerBlock"] + _keep_in_fp32_modules = [ + "time_embedder", + "scale_shift_table", + "scale_shift_table_action", + "norm1", + "action_norm1", + "text_norm1", + "norm2", + "action_norm2", + "text_norm2", + "norm3", + "action_norm3", + "text_norm3", + ] + _keys_to_ignore_on_load_unexpected = ["norm_added_q"] + _repeated_blocks = ["WanTransformerBlock"] + + @register_to_config + def __init__( + self, + patch_size=(1, 2, 2), + num_attention_heads=24, + attention_head_dim=128, + in_channels=48, + out_channels=48, + action_dim=30, + text_dim=4096, + freq_dim=256, + ffn_dim=14336, + num_layers=30, + cross_attn_norm=True, + eps=1e-06, + rope_max_seq_len=1024, + pos_embed_seq_len=None, + attn_mode="torch", + ): + super().__init__() + self.patch_size = patch_size + self.num_attention_heads = num_attention_heads + self.attention_head_dim = attention_head_dim + inner_dim = num_attention_heads * attention_head_dim + self.rope = WanRotaryPosEmbed(attention_head_dim, patch_size, rope_max_seq_len) + self.patch_embedding_mlp = nn.Linear( + in_channels * patch_size[0] * patch_size[1] * patch_size[2], inner_dim + ) + self.action_embedder = nn.Linear(action_dim, inner_dim) + self.condition_embedder = WanTimeTextImageEmbedding( + dim=inner_dim, + time_freq_dim=freq_dim, + time_proj_dim=inner_dim * 6, + text_embed_dim=text_dim, + pos_embed_seq_len=pos_embed_seq_len, + ) + self.condition_embedder_action = deepcopy(self.condition_embedder) + + self.blocks = nn.ModuleList( + [ + WanTransformerBlock( + inner_dim, ffn_dim, num_attention_heads, cross_attn_norm, eps, attn_mode=attn_mode + ) + for _ in range(num_layers) + ] + ) + + self.norm_out = FP32LayerNorm(inner_dim, eps, elementwise_affine=False) + self.proj_out = nn.Linear(inner_dim, out_channels * math.prod(patch_size)) + self.action_proj_out = nn.Linear(inner_dim, action_dim) + self.scale_shift_table = nn.Parameter(torch.randn(1, 2, inner_dim) / inner_dim**0.5) + + # KV-cache management for autoregressive streaming inference + def clear_cache(self, cache_name): + for block in self.blocks: + block.attn1.clear_cache(cache_name) + + def clear_pred_cache(self, cache_name): + for block in self.blocks: + block.attn1.clear_pred_cache(cache_name) + + def create_empty_cache( + self, + cache_name, + attn_window, + latent_token_per_chunk, + action_token_per_chunk, + device, + dtype, + batch_size, + ): + total_token_len = (attn_window // 2) * latent_token_per_chunk + ( + attn_window // 2 + ) * action_token_per_chunk + for block in self.blocks: + block.attn1.init_kv_cache( + cache_name, + total_token_len, + self.num_attention_heads, + self.attention_head_dim, + device, + dtype, + batch_size, + ) + + # Embedding helpers (shared by train + inference paths) + def _input_embed(self, latents, input_type="latent"): + if input_type == "latent": + hidden_states = rearrange( + latents, + "b c (f p1) (h p2) (w p3) -> b (f h w) (c p1 p2 p3)", + p1=self.patch_size[0], + p2=self.patch_size[1], + p3=self.patch_size[2], + ) + hidden_states = self.patch_embedding_mlp(hidden_states) + elif input_type == "action": + hidden_states = rearrange(latents, "b c f h w -> b (f h w) c") + hidden_states = self.action_embedder(hidden_states) + elif input_type == "text": + hidden_states = self.condition_embedder.text_embedder(latents) + else: + raise ValueError(f"Unsupported input type: {input_type}") + return hidden_states + + def _time_embed(self, timesteps, h, w, dtype, action_mode=False): + patch_scale_h, patch_scale_w = (1, 1) if action_mode else (self.patch_size[1], self.patch_size[2]) + latent_time_steps = torch.repeat_interleave( + timesteps, (h // patch_scale_h) * (w // patch_scale_w), dim=1 + ) + current_condition_embedder = ( + self.condition_embedder_action if action_mode else self.condition_embedder + ) + temb, timestep_proj = current_condition_embedder(latent_time_steps, dtype=dtype) + timestep_proj = timestep_proj.unflatten(2, (6, -1)) # B L 6 C + return temb, timestep_proj + + # Dual-stream training forward (flow matching). Requires attn_mode='flex'. + def forward_train(self, input_dict): + input_dict["latent_dict"]["noisy_latents"] = input_dict["latent_dict"]["noisy_latents"].to( + torch.bfloat16 + ) + input_dict["latent_dict"]["latent"] = input_dict["latent_dict"]["latent"].to(torch.bfloat16) + input_dict["action_dict"]["noisy_latents"] = input_dict["action_dict"]["noisy_latents"].to( + torch.bfloat16 + ) + input_dict["action_dict"]["latent"] = input_dict["action_dict"]["latent"].to(torch.bfloat16) + + latent_dict = input_dict["latent_dict"] + action_dict = input_dict["action_dict"] + batch_size = latent_dict["noisy_latents"].shape[0] + + latent_hidden_states = self._input_embed(latent_dict["noisy_latents"], input_type="latent").flatten( + 0, 1 + )[None] + action_hidden_states = self._input_embed(action_dict["noisy_latents"], input_type="action").flatten( + 0, 1 + )[None] + text_hidden_states = self._input_embed(latent_dict["text_emb"], input_type="text") + + text_hidden_states = text_hidden_states.flatten(0, 1)[None] + + condition_latent_hidden_states = self._input_embed( + latent_dict["latent"], input_type="latent" + ).flatten(0, 1)[None] + condition_action_hidden_states = self._input_embed( + action_dict["latent"], input_type="action" + ).flatten(0, 1)[None] + + hidden_states = torch.cat( + [ + latent_hidden_states, + condition_latent_hidden_states, + action_hidden_states, + condition_action_hidden_states, + ], + dim=1, + ) + + latent_grid_id = latent_dict["grid_id"].permute(1, 0, 2).flatten(1)[None] + action_grid_id = action_dict["grid_id"].permute(1, 0, 2).flatten(1)[None] + full_grid_id = torch.cat([latent_grid_id] * 2 + [action_grid_id] * 2, dim=2) + + rotary_emb = self.rope(full_grid_id)[:, :, None] + + latent_time_steps = torch.cat( + [latent_dict["timesteps"].flatten(0, 1), latent_dict["cond_timesteps"].flatten(0, 1)] + )[None] + action_time_steps = torch.cat( + [action_dict["timesteps"].flatten(0, 1), action_dict["cond_timesteps"].flatten(0, 1)] + )[None] + latent_temb, latent_timestep_proj = self._time_embed( + latent_time_steps, + latent_dict["noisy_latents"].shape[-2], + latent_dict["noisy_latents"].shape[-1], + dtype=hidden_states.dtype, + action_mode=False, + ) + action_temb, action_timestep_proj = self._time_embed( + action_time_steps, + action_dict["noisy_latents"].shape[-2], + action_dict["noisy_latents"].shape[-1], + dtype=hidden_states.dtype, + action_mode=True, + ) + temb = torch.cat([latent_temb, action_temb], dim=1) + timestep_proj = torch.cat([latent_timestep_proj, action_timestep_proj], dim=1) + + total_length = hidden_states.shape[1] + padded_length = (128 - total_length % 128) % 128 + hidden_states = F.pad(hidden_states, (0, 0, 0, padded_length)) + rotary_emb = F.pad(rotary_emb, (0, 0, 0, 0, 0, padded_length)) + temb = F.pad(temb, (0, 0, 0, padded_length)) + timestep_proj = F.pad(timestep_proj, (0, 0, 0, 0, 0, padded_length)) + + split_list = [ + latent_hidden_states.shape[1], + condition_latent_hidden_states.shape[1], + action_hidden_states.shape[1], + condition_action_hidden_states.shape[1], + padded_length, + ] + + FlexAttnFunc.init_mask( + latent_dict["noisy_latents"].shape, + action_dict["noisy_latents"].shape, + padded_length, + input_dict["chunk_size"], + window_size=input_dict["window_size"], + patch_size=self.patch_size, + device=hidden_states.device, + ) + + for block in self.blocks: + hidden_states = block( + hidden_states, text_hidden_states, timestep_proj, rotary_emb, update_cache=False + ) + temb_scale_shift_table = self.scale_shift_table[None] + temb[:, :, None, ...] + shift, scale = rearrange(temb_scale_shift_table, "b l n c -> b n l c").chunk(2, dim=1) + shift = shift.to(hidden_states.device).squeeze(1) + scale = scale.to(hidden_states.device).squeeze(1) + hidden_states = (self.norm_out(hidden_states.float()) * (1.0 + scale) + shift).type_as(hidden_states) + latent_hidden_states, _, action_hidden_states, _, _ = torch.split(hidden_states, split_list, dim=1) + latent_hidden_states = self.proj_out(latent_hidden_states) + latent_hidden_states = rearrange( + latent_hidden_states, "1 (b l) (n c) -> b (l n) c", n=math.prod(self.patch_size), b=batch_size + ) + action_hidden_states = self.action_proj_out(action_hidden_states) + action_hidden_states = rearrange(action_hidden_states, "1 (b l) c -> b l c", b=batch_size) + + return latent_hidden_states, action_hidden_states + + # Single-stream inference forward (one denoising step for one stream) + def forward(self, input_dict, update_cache=0, cache_name="pos", action_mode=False, train_mode=False): + if train_mode: + return self.forward_train(input_dict) + if action_mode: # action input emb + latent_hidden_states = rearrange(input_dict["noisy_latents"], "b c f h w -> b (f h w) c") + latent_hidden_states = self.action_embedder(latent_hidden_states) # B L1 C + else: # latent input emb + latent_hidden_states = rearrange( + input_dict["noisy_latents"], + "b c (f p1) (h p2) (w p3) -> b (f h w) (c p1 p2 p3)", + p1=self.patch_size[0], + p2=self.patch_size[1], + p3=self.patch_size[2], + ) + latent_hidden_states = self.patch_embedding_mlp(latent_hidden_states) + text_hidden_states = self.condition_embedder.text_embedder(input_dict["text_emb"]) # B L2 C + + latent_grid_id = input_dict["grid_id"] + rotary_emb = self.rope(latent_grid_id)[:, :, None] # 1 L 1 C + patch_scale_h, patch_scale_w = (1, 1) if action_mode else (self.patch_size[1], self.patch_size[2]) + + latent_time_steps = torch.repeat_interleave( + input_dict["timesteps"], + (input_dict["noisy_latents"].shape[-2] // patch_scale_h) + * (input_dict["noisy_latents"].shape[-1] // patch_scale_w), + dim=1, + ) # L + current_condition_embedder = ( + self.condition_embedder_action if action_mode else self.condition_embedder + ) + temb, timestep_proj = current_condition_embedder(latent_time_steps, dtype=latent_hidden_states.dtype) + timestep_proj = timestep_proj.unflatten(2, (6, -1)) # B L 6 C + + for block in self.blocks: + latent_hidden_states = block( + latent_hidden_states, + text_hidden_states, + timestep_proj, + rotary_emb, + update_cache=update_cache, + cache_name=cache_name, + ) + temb_scale_shift_table = self.scale_shift_table[None] + temb[:, :, None, ...] + shift, scale = rearrange(temb_scale_shift_table, "b l n c -> b n l c").chunk(2, dim=1) + shift = shift.to(latent_hidden_states.device).squeeze(1) + scale = scale.to(latent_hidden_states.device).squeeze(1) + latent_hidden_states = (self.norm_out(latent_hidden_states.float()) * (1.0 + scale) + shift).type_as( + latent_hidden_states + ) + + if action_mode: + latent_hidden_states = self.action_proj_out(latent_hidden_states) + else: + latent_hidden_states = self.proj_out(latent_hidden_states) + latent_hidden_states = rearrange( + latent_hidden_states, "b l (n c) -> b (l n) c", n=math.prod(self.patch_size) + ) + + return latent_hidden_states + + +class WanVAEStreamingWrapper: + """Wraps an ``AutoencoderKLWan`` encoder to support causal streaming encoding across chunks.""" + + def __init__(self, vae_model): + self.vae = vae_model + self.encoder = vae_model.encoder + self.quant_conv = vae_model.quant_conv + + if hasattr(self.vae, "_cached_conv_counts"): + self.enc_conv_num = self.vae._cached_conv_counts["encoder"] + else: + count = 0 + for m in self.encoder.modules(): + if m.__class__.__name__ == "WanCausalConv3d": + count += 1 + self.enc_conv_num = count + + self.clear_cache() + + def clear_cache(self): + self.feat_cache = [None] * self.enc_conv_num + + def encode_chunk(self, x_chunk): + if hasattr(self.vae.config, "patch_size") and self.vae.config.patch_size is not None: + x_chunk = _vae_patchify(x_chunk, self.vae.config.patch_size) + feat_idx = [0] + out = self.encoder(x_chunk, feat_cache=self.feat_cache, feat_idx=feat_idx) + enc = self.quant_conv(out) + return enc diff --git a/src/lerobot/scripts/lerobot_eval.py b/src/lerobot/scripts/lerobot_eval.py index 1ec4ea75f..722763d6e 100644 --- a/src/lerobot/scripts/lerobot_eval.py +++ b/src/lerobot/scripts/lerobot_eval.py @@ -169,6 +169,7 @@ def rollout( env_features: dict | None = None, recording_repo_id: str | None = None, recording_private: bool = False, + predicted_latents_callback: Callable[[PreTrainedPolicy], None] | None = None, ) -> dict: """Run a batched policy rollout once through a batch of environments. @@ -198,6 +199,9 @@ def rollout( are returned optionally because they typically take more memory to cache. Defaults to False. render_callback: Optional rendering callback to be used after the environments are reset, and after every step. + predicted_latents_callback: Optional callback invoked after every ``select_action`` with the policy + itself. World-model policies (e.g. LingBot-VA) stash predicted video latents on + ``policy.last_predicted_latents``; this lets the caller concatenate chunks and decode once. Returns: The dictionary described above. """ @@ -276,6 +280,8 @@ def rollout( observation = preprocessor(observation) with torch.inference_mode(): action = policy.select_action(observation) + if predicted_latents_callback is not None: + predicted_latents_callback(policy) action = postprocessor(action) action_transition = {ACTION: action} @@ -295,12 +301,22 @@ def rollout( # available if none of the envs finished. if "final_info" in info: final_info = info["final_info"] - if not isinstance(final_info, dict): - raise RuntimeError( - "Unsupported `final_info` format: expected dict (Gymnasium >= 1.0). " - "You're likely using an older version of gymnasium (< 1.0). Please upgrade." + if isinstance(final_info, dict): + is_success = final_info.get("is_success", [False] * env.num_envs) + successes = ( + is_success.tolist() + if hasattr(is_success, "tolist") + else [bool(is_success)] * env.num_envs ) - successes = final_info["is_success"].tolist() + else: + # Gymnasium < 1.0 returns final_info as a per-env sequence/object array, + # with entries set to a dict only for envs that just finished. + successes = [] + for item in final_info: + if isinstance(item, dict) and "is_success" in item: + successes.append(bool(item["is_success"])) + else: + successes.append(False) elif "is_success" in info: is_success = info["is_success"] successes = ( @@ -400,6 +416,7 @@ def eval_policy( env_features: dict | None = None, recording_repo_id: str | None = None, recording_private: bool = False, + save_predicted_video: bool = False, ) -> dict: """ Args: @@ -418,6 +435,11 @@ def eval_policy( if max_episodes_rendered > 0 and not videos_dir: raise ValueError("If max_episodes_rendered > 0, videos_dir must be provided.") + # World-model policies (e.g. LingBot-VA) opt into predicted-video saving via their config. + save_predicted_video = save_predicted_video or bool( + getattr(getattr(policy, "config", None), "save_predicted_video", False) + ) + if not isinstance(policy, PreTrainedPolicy): exc = ValueError( f"Policy of type 'PreTrainedPolicy' is expected, but type '{type(policy)}' was provided." @@ -461,6 +483,22 @@ def eval_policy( if max_episodes_rendered > 0: video_paths: list[str] = [] + if save_predicted_video: + if not videos_dir: + raise ValueError("If save_predicted_video is True, videos_dir must be provided.") + predicted_video_paths: list[str] = [] + n_predicted_rendered = 0 + + # Collect predicted-video latents across a rollout (world-model policies only). The latents are + # concatenated and decoded once after the rollout, matching upstream LingBot-VA's visualization path. + def collect_predicted_latents(policy: PreTrainedPolicy): + latents = getattr(policy, "last_predicted_latents", None) + if latents is not None: + pred_latents.append( + latents.detach().to("cpu") if hasattr(latents, "detach") else torch.as_tensor(latents).cpu() + ) + policy.last_predicted_latents = None + if return_episode_data: episode_data: dict | None = None @@ -472,6 +510,9 @@ def eval_policy( if max_episodes_rendered > 0: ep_frames: list[np.ndarray] = [] + if save_predicted_video: + pred_latents: list[torch.Tensor] = [] + if start_seed is None: seeds = None else: @@ -492,6 +533,7 @@ def eval_policy( env_features=env_features, recording_repo_id=recording_repo_id, recording_private=recording_private, + predicted_latents_callback=collect_predicted_latents if save_predicted_video else None, ) # Figure out where in each rollout sequence the first done condition was encountered (results after @@ -557,6 +599,35 @@ def eval_policy( threads.append(thread) n_episodes_rendered += 1 + # Maybe save the policy's predicted (imagined) video for this batch's rollout. + if save_predicted_video and len(pred_latents) > 0: + predicted_latent = torch.cat(pred_latents, dim=2) + decoder = getattr(policy, "decode_predicted_latents", None) or getattr( + policy, "_decode_predicted_video", None + ) + if decoder is None: + raise AttributeError( + "Policy config requested predicted-video saving, but the policy does not expose " + "`decode_predicted_latents` or `_decode_predicted_video`." + ) + predicted_video = decoder(predicted_latent) + if hasattr(predicted_video, "detach"): + predicted_video = predicted_video.detach().to("cpu").numpy() + videos_dir.mkdir(parents=True, exist_ok=True) + predicted_video_path = videos_dir / f"pred_episode_{n_predicted_rendered}.mp4" + predicted_video_paths.append(str(predicted_video_path)) + thread = threading.Thread( + target=write_video, + args=( + str(predicted_video_path), + predicted_video, + env.unwrapped.metadata["render_fps"], + ), + ) + thread.start() + threads.append(thread) + n_predicted_rendered += 1 + progbar.set_postfix( {"running_success_rate": f"{np.mean(all_successes[:n_episodes]).item() * 100:.1f}%"} ) @@ -600,6 +671,9 @@ def eval_policy( if max_episodes_rendered > 0: info["video_paths"] = video_paths + if save_predicted_video: + info["predicted_video_paths"] = predicted_video_paths + return info @@ -740,9 +814,10 @@ class TaskMetrics(TypedDict): max_rewards: list[float] successes: list[bool] video_paths: list[str] + predicted_video_paths: list[str] -ACC_KEYS = ("sum_rewards", "max_rewards", "successes", "video_paths") +ACC_KEYS = ("sum_rewards", "max_rewards", "successes", "video_paths", "predicted_video_paths") def eval_one( @@ -791,6 +866,7 @@ def eval_one( max_rewards=[ep["max_reward"] for ep in per_episode], successes=[ep["success"] for ep in per_episode], video_paths=task_result.get("video_paths", []), + predicted_video_paths=task_result.get("predicted_video_paths", []), ) @@ -851,6 +927,7 @@ def run_one( if max_episodes_rendered > 0: metrics.setdefault("video_paths", []) + metrics.setdefault("predicted_video_paths", []) return task_group, task_id, metrics @@ -908,11 +985,11 @@ def eval_policy_all( _append("sum_rewards", metrics.get("sum_rewards")) _append("max_rewards", metrics.get("max_rewards")) _append("successes", metrics.get("successes")) - # video_paths is list-like - paths = metrics.get("video_paths", []) - if paths: - group_acc[group]["video_paths"].extend(paths) - overall["video_paths"].extend(paths) + for key in ("video_paths", "predicted_video_paths"): + paths = metrics.get(key, []) + if paths: + group_acc[group][key].extend(paths) + overall[key].extend(paths) # Choose runner (sequential vs threaded) task_runner = partial( @@ -984,6 +1061,7 @@ def eval_policy_all( "pc_success": _agg_from_list(acc["successes"]) * 100 if acc["successes"] else float("nan"), "n_episodes": len(acc["sum_rewards"]), "video_paths": list(acc["video_paths"]), + "predicted_video_paths": list(acc["predicted_video_paths"]), } # overall aggregates @@ -995,6 +1073,7 @@ def eval_policy_all( "eval_s": time.time() - start_t, "eval_ep_s": (time.time() - start_t) / max(1, len(overall["sum_rewards"])), "video_paths": list(overall["video_paths"]), + "predicted_video_paths": list(overall["predicted_video_paths"]), } return { diff --git a/tests/policies/lingbot_va/test_configuration_lingbot_va.py b/tests/policies/lingbot_va/test_configuration_lingbot_va.py new file mode 100644 index 000000000..089d5ec02 --- /dev/null +++ b/tests/policies/lingbot_va/test_configuration_lingbot_va.py @@ -0,0 +1,78 @@ +#!/usr/bin/env python + +# Copyright 2026 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. + +from __future__ import annotations + +import pytest + +from lerobot.configs.policies import PreTrainedConfig +from lerobot.configs.types import FeatureType, PolicyFeature +from lerobot.policies.lingbot_va.configuration_lingbot_va import LingBotVAConfig +from lerobot.utils.constants import ACTION, OBS_IMAGES + + +def make_config(**overrides) -> LingBotVAConfig: + kwargs = {"device": "cpu"} + kwargs.update(overrides) + return LingBotVAConfig(**kwargs) + + +def test_registered_in_choice_registry() -> None: + assert "lingbot_va" in PreTrainedConfig.get_known_choices() + assert PreTrainedConfig.get_choice_class("lingbot_va") is LingBotVAConfig + + +def test_type_property() -> None: + assert make_config().type == "lingbot_va" + + +def test_chunk_size_and_action_steps() -> None: + cfg = make_config(frame_chunk_size=4, action_per_frame=4) + assert cfg.chunk_size == 16 + assert cfg.n_action_steps == 16 + assert cfg.action_delta_indices == list(range(16)) + assert cfg.observation_delta_indices == list(range(16)) + assert cfg.reward_delta_indices is None + + +def test_optimizer_and_scheduler_presets() -> None: + cfg = make_config() + opt = cfg.get_optimizer_preset() + assert opt.lr == cfg.optimizer_lr + sched = cfg.get_scheduler_preset() + assert sched.num_warmup_steps == cfg.scheduler_warmup_steps + + +def test_validate_features_sets_action_feature() -> None: + cfg = make_config() + cfg.input_features = {f"{OBS_IMAGES}.image": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 128, 128))} + cfg.output_features = {} + cfg.validate_features() + assert ACTION in cfg.output_features + assert cfg.output_features[ACTION].shape == (len(cfg.used_action_channel_ids),) + + +def test_validate_features_no_visual_raises() -> None: + cfg = make_config() + cfg.input_features = {} + cfg.output_features = {} + with pytest.raises(ValueError, match="at least one visual input feature"): + cfg.validate_features() + + +def test_invalid_attn_mode_raises() -> None: + with pytest.raises(ValueError, match="attn_mode"): + make_config(attn_mode="banana") diff --git a/tests/policies/lingbot_va/test_factory.py b/tests/policies/lingbot_va/test_factory.py new file mode 100644 index 000000000..1aec34df2 --- /dev/null +++ b/tests/policies/lingbot_va/test_factory.py @@ -0,0 +1,38 @@ +#!/usr/bin/env python + +# Copyright 2026 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. + +from __future__ import annotations + +import pytest + +from lerobot.policies.factory import make_policy_config +from lerobot.policies.lingbot_va.configuration_lingbot_va import LingBotVAConfig + + +def test_make_policy_config_returns_lingbot_va() -> None: + cfg = make_policy_config("lingbot_va", device="cpu") + assert isinstance(cfg, LingBotVAConfig) + + +def test_get_policy_class_resolves_lazily() -> None: + # Importing the policy class pulls in diffusers (Wan2.2 stack); skip if unavailable. + pytest.importorskip("diffusers") + pytest.importorskip("transformers") + from lerobot.policies.factory import get_policy_class + + cls = get_policy_class("lingbot_va") + assert cls.name == "lingbot_va" + assert cls.config_class is LingBotVAConfig diff --git a/tests/policies/lingbot_va/test_modules.py b/tests/policies/lingbot_va/test_modules.py new file mode 100644 index 000000000..3ca6b2ba1 --- /dev/null +++ b/tests/policies/lingbot_va/test_modules.py @@ -0,0 +1,128 @@ +#!/usr/bin/env python + +# Copyright 2026 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. + +"""Unit tests for the vendored LingBot-VA helper code (scheduler + grid utilities).""" + +from __future__ import annotations + +import pytest +import torch + +pytest.importorskip("diffusers") # the model code lives in modeling_lingbot_va, which imports diffusers + +from lerobot.policies.lingbot_va.modeling_lingbot_va import FlowMatchScheduler +from lerobot.policies.lingbot_va.utils import data_seq_to_patch, get_mesh_id + + +def test_flow_match_scheduler_timesteps_monotone_decreasing() -> None: + sch = FlowMatchScheduler(shift=5.0, sigma_min=0.0, extra_one_step=True) + sch.set_timesteps(20) + assert sch.timesteps.shape == (20,) + diffs = sch.timesteps[1:] - sch.timesteps[:-1] + assert torch.all(diffs <= 0) # decreasing + + +def test_flow_match_scheduler_step_preserves_shape() -> None: + sch = FlowMatchScheduler(shift=5.0, sigma_min=0.0, extra_one_step=True) + sch.set_timesteps(20) + sample = torch.zeros(1, 48, 4, 8, 16) + out = sch.step(torch.ones_like(sample), sch.timesteps[0], sample) + assert out.shape == sample.shape + + +def test_flow_match_scheduler_add_noise() -> None: + sch = FlowMatchScheduler(shift=5.0, sigma_min=0.0, extra_one_step=True) + sch.set_timesteps(20) + sample = torch.randn(1, 48, 4, 8, 16) + noise = torch.randn_like(sample) + noisy = sch.add_noise(sample, noise, sch.timesteps[:4], t_dim=2) + assert noisy.shape == sample.shape + + +def test_get_mesh_id_latent_shape() -> None: + grid = get_mesh_id(4, 8, 16, 0, 1, 0) + assert grid.shape == (4, 4 * 8 * 16) # (f, h, w, stream) x tokens + + +def test_get_mesh_id_action_shape() -> None: + grid = get_mesh_id(4, 4, 1, 1, 1, 0, action=True) + assert grid.shape == (4, 4 * 4 * 1) + # Action rows for h/w are sentinel -1. + assert torch.all(grid[1] < 0) + assert torch.all(grid[2] < 0) + + +def test_data_seq_to_patch_roundtrip_shape() -> None: + b, f, h, w, c = 1, 4, 8, 16, 48 + seq = torch.arange(b * f * h * w * c, dtype=torch.float32).reshape(b, f * h * w, c) + out = data_seq_to_patch((1, 2, 2), seq, f, h, w, batch_size=b) + assert out.shape == (b, c, f, h, w) + + +def test_training_step_reduces_loss_tiny_flex() -> None: + """End-to-end single training step (flow-matching loss -> backward -> AdamW) on a tiny config. + + Exercises the flex-attention training path; requires a CUDA GPU with flex-attention support. + """ + if not torch.cuda.is_available(): + import pytest + + pytest.skip("training step test requires a CUDA GPU (flex-attention)") + + from lerobot.configs.types import FeatureType, PolicyFeature + from lerobot.policies.lingbot_va.configuration_lingbot_va import LingBotVAConfig + from lerobot.policies.lingbot_va.modeling_lingbot_va import LingBotVAPolicy + from lerobot.utils.constants import ACTION, OBS_IMAGES + + cfg = LingBotVAConfig( + attn_mode="flex", + dtype="bfloat16", + in_channels=16, + out_channels=16, + action_dim=8, + text_dim=32, + freq_dim=64, + ffn_dim=64, + num_attention_heads=2, + attention_head_dim=24, + num_layers=2, + frame_chunk_size=2, + action_per_frame=4, + used_action_channel_ids=[0, 1, 2, 3], + obs_cam_keys=[f"{OBS_IMAGES}.image"], + device="cuda", + ) + cfg.input_features = {f"{OBS_IMAGES}.image": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 64, 64))} + cfg.output_features = {ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(4,))} + cfg.validate_features() + + policy = LingBotVAPolicy(cfg).to("cuda") + policy.train() + opt = torch.optim.AdamW(policy.get_optim_params(), lr=1e-4) + + b, fc, apf = 1, cfg.frame_chunk_size, cfg.action_per_frame + latents = torch.randn(b, cfg.in_channels, fc, 4, 4, device="cuda", dtype=torch.bfloat16) + actions = torch.randn(b, cfg.action_dim, fc, apf, 1, device="cuda", dtype=torch.bfloat16) + amask = torch.zeros(cfg.action_dim, device="cuda") + amask[cfg.used_action_channel_ids] = 1.0 + actions_mask = amask.view(1, -1, 1, 1, 1).expand_as(actions) + text_emb = torch.randn(b, cfg.max_sequence_length, cfg.text_dim, device="cuda", dtype=torch.bfloat16) + + loss, metrics = policy.training_loss_from_streams(latents, actions, actions_mask, text_emb) + assert torch.isfinite(loss) and {"latent_loss", "action_loss"} <= set(metrics) + loss.backward() + assert any(p.grad is not None and torch.isfinite(p.grad).all() for p in policy.get_optim_params()) + opt.step() diff --git a/tests/policies/lingbot_va/test_processor.py b/tests/policies/lingbot_va/test_processor.py new file mode 100644 index 000000000..e79771600 --- /dev/null +++ b/tests/policies/lingbot_va/test_processor.py @@ -0,0 +1,88 @@ +#!/usr/bin/env python + +# Copyright 2026 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. + +from __future__ import annotations + +import torch + +from lerobot.configs.types import FeatureType, PolicyFeature +from lerobot.policies.lingbot_va.configuration_lingbot_va import LingBotVAConfig +from lerobot.policies.lingbot_va.processor_lingbot_va import make_lingbot_va_pre_post_processors +from lerobot.processor import PolicyProcessorPipeline, UnnormalizerProcessorStep +from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action +from lerobot.utils.constants import ( + ACTION, + OBS_IMAGES, + POLICY_POSTPROCESSOR_DEFAULT_NAME, + POLICY_PREPROCESSOR_DEFAULT_NAME, +) + + +def _make_config() -> LingBotVAConfig: + cfg = LingBotVAConfig(device="cpu") + cfg.input_features = {f"{OBS_IMAGES}.image": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 128, 128))} + cfg.output_features = {} + cfg.validate_features() + return cfg + + +def test_make_pre_post_processors_names_and_steps() -> None: + cfg = _make_config() + pre, post = make_lingbot_va_pre_post_processors(cfg, dataset_stats=None) + assert pre.name == POLICY_PREPROCESSOR_DEFAULT_NAME + assert post.name == POLICY_POSTPROCESSOR_DEFAULT_NAME + # Actions are unnormalized by the standard built-in quantile unnormalizer. + assert any(isinstance(s, UnnormalizerProcessorStep) for s in post.steps) + + +def test_freshly_built_postprocessor_is_identity() -> None: + # Without action stats the quantile unnormalizer is a no-op (identity passthrough): the real + # per-benchmark q01/q99 are restored from the saved checkpoint on load, not hardcoded here. + cfg = _make_config() + _, post = make_lingbot_va_pre_post_processors(cfg, dataset_stats=None) + normed = torch.tensor([[0.3, -0.5, 1.0, -1.0, 0.0, 0.7, -0.2]]) + assert torch.allclose(post(normed), normed, atol=1e-6) + + +def test_postprocessor_quantile_unnormalization() -> None: + # QUANTILES unnormalize maps [-1, 1] -> [q01, q99]: -1 -> q01, +1 -> q99. + cfg = _make_config() + q01 = [-1.0, -0.5, 0.0, -1.0, -1.0, -1.0, -1.0] + q99 = [1.0, 0.5, 2.0, 1.0, 1.0, 1.0, 1.0] + stats = {ACTION: {"q01": q01, "q99": q99}} + _, post = make_lingbot_va_pre_post_processors(cfg, dataset_stats=stats) + out_lo = post(torch.full((1, 7), -1.0)) + out_hi = post(torch.full((1, 7), 1.0)) + assert torch.allclose(out_lo, torch.tensor(q01).unsqueeze(0), atol=1e-4) + assert torch.allclose(out_hi, torch.tensor(q99).unsqueeze(0), atol=1e-4) + + +def test_postprocessor_stats_survive_save_load(tmp_path) -> None: + # Regression guard for the Hub mechanism: the q01/q99 stats live in the saved post-processor + # state and must round-trip through save_pretrained / from_pretrained. + cfg = _make_config() + q01 = [-0.6, -0.8, -0.9, -0.1, -0.15, -0.25, -1.0] + q99 = [0.9, 0.85, 0.9, 0.17, 0.18, 0.34, 1.0] + _, post = make_lingbot_va_pre_post_processors(cfg, dataset_stats={ACTION: {"q01": q01, "q99": q99}}) + post.save_pretrained(tmp_path) + loaded = PolicyProcessorPipeline.from_pretrained( + tmp_path, + config_filename=f"{POLICY_POSTPROCESSOR_DEFAULT_NAME}.json", + to_transition=policy_action_to_transition, + to_output=transition_to_policy_action, + ) + out = loaded(torch.full((1, 7), -1.0)) + assert torch.allclose(out, torch.tensor(q01).unsqueeze(0), atol=1e-4) diff --git a/uv.lock b/uv.lock index 22afec6a6..260fb07f0 100644 --- a/uv.lock +++ b/uv.lock @@ -3057,6 +3057,11 @@ libero = [ { name = "torchcodec", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'AMD64' and sys_platform == 'linux') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'arm64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'win32'" }, { name = "transformers" }, ] +lingbot-va = [ + { name = "accelerate" }, + { name = "diffusers" }, + { name = "transformers" }, +] matplotlib-dep = [ { name = "contourpy" }, { name = "matplotlib" }, @@ -3265,6 +3270,7 @@ requires-dist = [ { name = "ipykernel", marker = "extra == 'notebook'", specifier = ">=6.0.0,<7.0.0" }, { name = "jsonlines", marker = "extra == 'dataset'", specifier = ">=4.0.0,<5.0.0" }, { name = "jupyter", marker = "extra == 'notebook'", specifier = ">=1.0.0,<2.0.0" }, + { name = "lerobot", extras = ["accelerate-dep"], marker = "extra == 'lingbot-va'" }, { name = "lerobot", extras = ["accelerate-dep"], marker = "extra == 'smolvla'" }, { name = "lerobot", extras = ["accelerate-dep"], marker = "extra == 'training'" }, { name = "lerobot", extras = ["aloha"], marker = "extra == 'all'" }, @@ -3292,6 +3298,7 @@ requires-dist = [ { name = "lerobot", extras = ["diffusers-dep"], marker = "extra == 'diffusion'" }, { name = "lerobot", extras = ["diffusers-dep"], marker = "extra == 'fastwam'" }, { name = "lerobot", extras = ["diffusers-dep"], marker = "extra == 'groot'" }, + { name = "lerobot", extras = ["diffusers-dep"], marker = "extra == 'lingbot-va'" }, { name = "lerobot", extras = ["diffusers-dep"], marker = "extra == 'multi-task-dit'" }, { name = "lerobot", extras = ["diffusers-dep"], marker = "extra == 'vla-jepa'" }, { name = "lerobot", extras = ["diffusion"], marker = "extra == 'all'" }, @@ -3312,6 +3319,7 @@ requires-dist = [ { name = "lerobot", extras = ["kinematics"], marker = "extra == 'all'" }, { name = "lerobot", extras = ["lekiwi"], marker = "extra == 'all'" }, { name = "lerobot", extras = ["libero"], marker = "sys_platform == 'linux' and extra == 'all'" }, + { name = "lerobot", extras = ["lingbot-va"], marker = "extra == 'all'" }, { name = "lerobot", extras = ["matplotlib-dep"], marker = "extra == 'async'" }, { name = "lerobot", extras = ["matplotlib-dep"], marker = "extra == 'sarm'" }, { name = "lerobot", extras = ["matplotlib-dep"], marker = "extra == 'unitree-g1'" }, @@ -3370,6 +3378,7 @@ requires-dist = [ { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'groot'" }, { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'hilserl'" }, { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'libero'" }, + { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'lingbot-va'" }, { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'molmoact2'" }, { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'multi-task-dit'" }, { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'peft'" }, @@ -3449,7 +3458,7 @@ requires-dist = [ { name = "transformers", marker = "extra == 'transformers-dep'", specifier = ">=5.4.0,<5.6.0" }, { name = "wandb", marker = "extra == 'training'", specifier = ">=0.24.0,<0.28.0" }, ] -provides-extras = ["dataset", "training", "hardware", "viz", "core-scripts", "evaluation", "dataset-viz", "av-dep", "pygame-dep", "placo-dep", "transformers-dep", "grpcio-dep", "accelerate-dep", "can-dep", "peft-dep", "scipy-dep", "diffusers-dep", "qwen-vl-utils-dep", "matplotlib-dep", "pyserial-dep", "deepdiff-dep", "pynput-dep", "pyzmq-dep", "motorbridge-dep", "motorbridge-smart-servo-dep", "feetech", "dynamixel", "damiao", "robstride", "openarms", "gamepad", "hopejr", "lekiwi", "unitree-g1", "reachy2", "rebot", "kinematics", "intelrealsense", "phone", "diffusion", "wallx", "pi", "molmoact2", "smolvla", "multi-task-dit", "groot", "sarm", "robometer", "topreward", "xvla", "eo1", "fastwam", "hilserl", "vla-jepa", "async", "peft", "annotations", "dev", "notebook", "test", "video-benchmark", "aloha", "pusht", "libero", "metaworld", "all"] +provides-extras = ["dataset", "training", "hardware", "viz", "core-scripts", "evaluation", "dataset-viz", "av-dep", "pygame-dep", "placo-dep", "transformers-dep", "grpcio-dep", "accelerate-dep", "can-dep", "peft-dep", "scipy-dep", "diffusers-dep", "qwen-vl-utils-dep", "matplotlib-dep", "pyserial-dep", "deepdiff-dep", "pynput-dep", "pyzmq-dep", "motorbridge-dep", "motorbridge-smart-servo-dep", "feetech", "dynamixel", "damiao", "robstride", "openarms", "gamepad", "hopejr", "lekiwi", "unitree-g1", "reachy2", "rebot", "kinematics", "intelrealsense", "phone", "diffusion", "wallx", "pi", "molmoact2", "smolvla", "multi-task-dit", "groot", "sarm", "robometer", "topreward", "xvla", "eo1", "fastwam", "hilserl", "vla-jepa", "lingbot-va", "async", "peft", "annotations", "dev", "notebook", "test", "video-benchmark", "aloha", "pusht", "libero", "metaworld", "all"] [[package]] name = "librt" From 708fa1d18952fc0b0c6810926f39347c7bd721df Mon Sep 17 00:00:00 2001 From: Steven Palma Date: Fri, 3 Jul 2026 21:15:09 +0200 Subject: [PATCH 34/38] feat(policies): add Gr00t N1.7 policy (#3922) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Add GR00T N1.7 support Add GR00T N1.7 policy configuration, checkpoint compatibility, processor parity, LIBERO documentation, and focused tests. Co-authored-by: Ryan Halabi * Move Groot processor compatibility into Groot loader * Restore GR00T Flash Attention install guidance * Allow Groot fake RTC chunk prefetch * Fix GR00T N1.7 RTC action decoding * Trim GR00T N1.7 RTC chunks to valid horizon * Ignore padded GR00T N1.7 RTC prefix rows * removed n1.5 dependency * removed remaining N1.5 traces * groot: auto-enable LIBERO gripper action transform for libero_sim GR00T N1.7 emits gripper in [0,1] but LIBERO expects [-1,1]. The decode transform existed but was never auto-enabled for embodiment_tag=libero_sim, so the policy scored 0% on LIBERO eval. Auto-set it in __post_init__ (still overridable). LIBERO Spatial eval: 0% -> 98%. * Reconnect GR00T relative action processors * groot: remove dead N1.5 code (eagle2_hg_model, flow_matching_action_head, action_encoder) N1.7 backbone is nvidia/Cosmos-Reason2-2B via Qwen3VLForConditionalGeneration, not Eagle2 — eagle2_hg_model/ had zero refs outside its own dir. GR00TN17ActionHead (groot_n1_7.py) re-implements MultiEmbodimentActionEncoder + CategorySpecificLinear + swish + SinusoidalPositionalEncoding locally, so flow_matching_action_head.py (N1.5 FlowmatchingActionHead) and its sole dependency action_encoder.py are dead. Verified: no src/ or tests/ reference. Removed (~2037 LOC): - eagle2_hg_model/ (4 files, ~1575 LOC) - action_head/flow_matching_action_head.py (408 LOC) - action_head/action_encoder.py (54 LOC) cross_attention_dit.py KEPT (DiT/AlternateVLDiT/SelfAttentionTransformer live in N1.7). * groot: reuse lerobot get_device_from_parameters instead of inline lookup modeling_groot.py duplicated next(self.parameters()).device twice. LeRobot ships get_device_from_parameters in policies/utils.py (used by diffusion, vqbet, tdmpc, gaussian_actor). Reuse it for consistency with the framework. * groot: fix stale Eagle VLM docstring in processor (N1.7 uses Qwen3-VL backbone) Addresses checker nit: processor_groot.py docstring still described the N1.5 Eagle VLM path with eagle_content/eagle_* keys that no longer exist in the code. * test(groot): add N1.7 original-vs-LeRobot output parity test Verifies the LeRobot GR00T N1.7 integration produces equivalent raw action_pred to NVIDIA Isaac-GR00T for the same checkpoint, inputs, seed, precision (fp32) and attention kernel (SDPA): max|diff|=8.9e-7 on the libero_sim embodiment (GR00T-N1.7-LIBERO/libero_10). The two impls pin incompatible transformers majors (orig 4.57.3 vs LeRobot 5.x) and cannot share a process, so the original outputs + exact collated inputs are produced out-of-process and loaded from an .npz. The test skips on CI / when the checkpoint or artifact are absent. * test(groot): parametrize N1.7 parity across all checkpoint embodiments Generalize the original-vs-LeRobot N1.7 output-parity test from a single libero_sim case to every embodiment tag in the checkpoint (libero_sim, oxe_droid, real_g1, the real_r1_pro_sharpa family, and the xdof family). Inputs are built generically from checkpoint metadata; the test discovers per-tag .npz artifacts and runs one parametrized case each, loading the LeRobot model once via a fixture. All 9 embodiments match the original to fp32 epsilon (max|diff| < 3e-6), confirming the integration is correct across the model's full embodiment space and not overfit to libero_sim. * test(groot): self-contained parity test + in-repo producer + docs - Rename test_groot_n1_7_vs_original.py -> test_groot_vs_original.py - Make the test self-contained: producer script (dump_original_n1_7.py) now lives next to the test; default artifact dir is repo-relative (tests/policies/groot/artifacts/), overridable via GROOT_N1_7_PARITY_DIR. The test only reads artifacts and skips if absent -- it never creates external dirs. - Heavy .npz artifacts (~6-9MB each) are gitignored and regenerated by the producer; never committed. - Drop the verbose 'MULTIPLE EMBODIMENTS' docstring block (kept a one-line note). - Document the parity procedure in the groot policy README (docs/source/policy_groot_README.md). - Rename test fn test_groot_n1_7_get_action_parity -> test_groot_get_action_parity. 9/9 embodiments still pass (max|diff| < 3e-6, fp32 eps). * docs(groot): drop WHY TWO ENVIRONMENTS block from parity test docstring * test(groot): move parity producer into utils/ package Mirror the tests/policies/pi0_pi05/utils convention: move dump_original_n1_7.py into a tests/policies/groot/utils/ package (with __init__.py) and update all path references in the test docstring/skip-message and the policy README. * test(groot): adopt test_groot_lerobot for GR00T N1.7, drop N1.5 The test loaded MODEL_PATH='aractingi/bimanual-handover-groot-10k', an N1.5 checkpoint (config base_model_path=nvidia/GR00T-N1.5-3B, no model_version). On load, model_version defaults to n1.7 while the base path infers n1.5, so the version-consistency guard in GrootConfig.__post_init__ raised ValueError and both test_lerobot_groot_inference and test_lerobot_groot_forward_pass failed. N1.5 is no longer a supported model_version. Adopt the test for N1.7: - MODEL_PATH -> nvidia/GR00T-N1.7-3B (root-level sharded safetensors; loads via GrootPolicy.from_pretrained as a base N1.7 model). - Embodiment tag 'gr1' (N1.5) -> 'gr1_unified' (valid N1.7 tag from the checkpoint embodiment_id.json), via a single EMBODIMENT_TAG constant. - DUMMY_ACTION_HORIZON 16 -> 40 to match N1.7's native action-chunk size. - Docstrings/labels updated to 'GR00T N1.7'. Both tests run and pass on CUDA; full tests/policies/groot/ suite is 73 passed / 0 failed / 0 skipped. * docs(groot): document the N1.5 removal and the N1.7 parity test - groot.mdx: breaking-change warning and migration path (pin lerobot==0.5.1 to keep N1.5, or move to N1.7); the dead `huggingface-cli download` is replaced with `hf download`. - policy_groot_README.md: N1.5 removal note, updated paper / model-card links, and the two-comparison (model parity + preprocessor parity) description of the original-vs-LeRobot test, including the raw-observation artifacts and recorded seed. * fix(groot): N1.7 backbone loading and DiT parameter-count logging - select_layer default tracks the N1.7-3B checkpoint value (16); real checkpoint loads still override it from config.json. - get_backbone_cls recognizes Cosmos-Reason2 / Qwen3-VL backbones by name and warns (instead of silently assuming) when an unrecognized backbone is loaded only on the strength of backbone_model_type='qwen'. - 'revision' pins the GR00T checkpoint repo only and is no longer forwarded into the unrelated backbone repo load; pin the backbone via transformers_loading_kwargs instead. - DiT / SelfAttentionTransformer parameter counts go through logging.debug instead of print(). * fix(groot): N1.7 config defaults, N1.5 rejection, and processor/model runtime fixes Covers the GR00T N1.7 source trio (configuration, processor, model wrapper). Config: - GrootConfig defaults are the N1.7 values; explicitly passed legacy N1.5-era values (chunk_size=50, max_state_dim=64, ...) are remapped with a warning instead of silently. - action_decode_transform gains an 'auto' sentinel so an explicit 'none' opt-out wins over the libero_sim default and survives save/load round-trips. - action_delta_indices is cached on the inputs that determine it. - Legacy N1.5 checkpoints/configs (tokenizer_assets_repo, model_type/ architectures/eagle backbone markers) are rejected with a single clear error pointing to lerobot==0.5.1. Processor: - GrootN17ActionDecodeStep handles the 2-D (B, D) actions delivered by sync select_action (relative eef/non-eef decode in eval/record flows). - Postprocessor falls back to dataset stats when a raw checkpoint lacks the configured embodiment tag; raw-state cache is per-instance, not process-global; caller overrides (device, rename_map) are honored on the raw-checkpoint branch. - Camera/modality-key mismatches warn (including the zero-match fallback); deprecated Qwen2VLImageProcessorFast replaced with Qwen2VLImageProcessor; removed N1.5 processor steps are stubbed to raise the removal guidance and the action-unpack step is re-registered as _v2. Model: - Flash-attention probe is diagnostic-only; forward raises on a missing loss; print() replaced with logging; N1.5 base-path mismatch includes the removal guidance. * fix(groot): skip normalization overrides for training * fix(groot): GPU/tensor N1.7 image preprocessing + resize to trained resolution GR00T training was dataloader-bound (0->100->0 GPU-utilization sawtooth). GrootN17VLMEncodeStep ran the Qwen3-VL image processor per frame on PIL images on the single CPU main-loop thread, and that cost is timed inside dataloading_s (preprocessor(batch) runs in the main process, not the dataloader workers), so adding workers cannot hide it. - Feed the torchvision-backed Qwen3-VL processor (C,H,W) uint8 tensors instead of a per-frame Image.fromarray PIL roundtrip, and run resize/normalize/patchify on config.device (GPU) when available. Bit-identical on CPU when no resize is configured; with a resize only the PIL->torchvision bicubic backend differs (<2/255 per pixel). The use_albumentations path stays PIL/cv2; reload on a box without the saved device falls back to CPU. - Default image_target_size/crop to the N1.7 backbone's training geometry (256x256 / 230x230) when a checkpoint ships no image sizing (checkpoint_assets is None, e.g. finetuning nvidia/GR00T-N1.7-3B via repo-id with a new embodiment). Previously image_target_size=None disabled the resize, so full-resolution frames were patchified into ~4.7x more vision tokens than the model was trained on -- inflating dataloading_s (patchify) and update_s (VLM sequence) and skewing the input distribution. Checkpoints that pin their own sizing are honored; the default constants are shared with GR00T_N1_7_DEFAULTS. Net: preprocessing leaves the CPU critical path and the VLM sees the resolution it was trained on -- faster training/inference and a correct train/serve distribution. Affects inference too (shared preprocessor); existing checkpoints still load (backward compatible) but must be retrained to gain the benefits. * refactor(groot): N1.7 style cleanup (utils, imports, flash-attn, config) Mechanical refactor of the GR00T N1.7 policy to match the repo's architecture and style standards. No change to policy algorithm/numerics; only UX/CLI and packaging changes. Tests are intentionally left untouched (out of scope) and need updating for the removed `model_version` field. Cleanup & consolidation: - Add `groot/utils.py` holding the pure, side-effect-free helpers (JSON I/O, value coercion, stat flattening, rot6d/SE3 math, language/batch prep) shared by the config and processor layers. - Remove dead code: the unused `resolve_groot_n1_7_backbone_model` cache-resolver cluster, `GR00TN17Config.to_filtered_dict/json`, and the `_copy_default` wrapper. Imports & execution guards: - Hoist nested imports to module top; relative imports within the package, absolute for external modules. The version-gated Qwen3-VL classes import under the single `_transformers_available` guard (transformers is pinned >=5.4, which ships them). - No import-time side effects: `_register_with_transformers()` now runs in `GR00TN17.__init__` (idempotent via `register(exist_ok=True)`), and the N1.5 step stubs register lazily before pipeline deserialization (idempotent via the registry, no run-once globals). - Gate optional deps at the point of use with `require_package(..., extra="groot")`. Dependencies & docs: - Drop `flash-attn` (and its build-only dep `ninja`) from the `groot` extra; default to SDPA (numerically equivalent) with opt-in via `--policy.use_flash_attention`. Un-comment `lerobot[groot]` in the `all` extra and regenerate `uv.lock`. - Rewrite the `groot.mdx` install section: flash-attn is a purely optional, user-managed optimization that LeRobot neither installs nor requires. Config & CLI: - Surface previously-frozen knobs on `GrootConfig` (plumbed into `GR00TN17Config`; no-ops at their defaults): inference — `num_inference_timesteps`, `rtc_ramp_rate`, `use_flash_attention`; fine-tuning — `tune_top_llm_layers` (partial-LLM tuning) and `tune_vlln` (previously hardwired to True). - Convert the single-valued `model_version` and `n1_7_backbone_model` fields to internal constants. - Keep `base_model_path`: it is NOT equivalent to `pretrained_path` (raw NVIDIA checkpoints have no LeRobot `type` field and load only via `base_model_path`) and is genuinely user-tunable. - Keep the deprecated Isaac-GR00T/N1.5 fields (and the dead LoRA fields) as a back-compat block so a v0.5.1 N1.5 `config.json` still parses under draccus and is rejected with the friendly N1.5 removal message instead of an opaque decode error. * Optimize GR00T N1.7 image preprocessing * Remove PIL fallback from GR00T preprocessing * Fix GROOT relative action training stats * Address GROOT relative action review feedback * Fix GROOT N1.7 relative action stats * Fix GROOT relative action training stats * Fix GROOT relative action padding and RTC leftovers * Reset rollout state after robot episode end * Revert "Reset rollout state after robot episode end" This reverts commit 1322f45aec088d3ca346640d995d28edcf71d00f. * Move GROOT relative stats out of train script * Guard GR00T relative action stepwise decode * Match GR00T N1.7 OSS preprocessing and relative actions * Apply LIBERO action decode override after loading * Format GR00T OSS parity changes * chore(policies): add guards, warnings and comments + recover tests n1.5 check * fix(style): pre-commit * fix(ci): guard dependecy checks * chore(groot): move cv2 to the top as its in the default install tag * chore(policies): add explicit dataset dependecy to gr00t implementation * fix(test): add guard * fix(groot): make N1.7 letterbox opt-in * feat(groot): activate checkpoint-configured N1.7 raw-state dropout during training Isaac-GR00T applies dual state regularization during fine-tuning: raw-state zeroing driven by the processor sidecar's state_dropout_prob (0.2 for the inspected N1.7 checkpoint) plus encoded-feature dropout. Baseline LeRobot kept the processor in deterministic mode, so the raw-state dropout never activated (RCA Tier-2 contributor to the LeRobot-trained SO-101 failures). - GrootN17PackInputsStep: runtime-only 'training' flag + state_dropout_prob; whole-sample state zeroing gated on torch.is_grad_enabled() so eval and no_grad validation paths are unaffected - sidecar loader reads state_dropout_prob from processor_config.json - state_dropout_prob serializes with the step; the training flag intentionally does not (reloaded pipelines default to eval, re-enabled only when processors are rebuilt with dataset_meta) - _set_groot_preprocessor_training toggles any dataclass step exposing a 'training' field on serialized-pipeline reloads Verification: tests/policies/groot/test_groot_state_dropout.py (4 passed) on RTX PRO 6000 / CUDA 13.3. * fix(groot): align N1.7 fine-tuning optimizer/scheduler/precision with Isaac-GR00T Evidence from the LeRobot-vs-OSS checkpoint comparison: the LeRobot/HF 8k checkpoint's DiT moved only ~19% as far from base as the OSS-trained one (0.0547 vs 0.285 relative L2) - undertrained because the scheduler decayed over a hardcoded 10k steps regardless of --steps, on top of beta1/clip mismatches. - AdamW betas (0.95, 0.999) -> (0.9, 0.999) and grad_clip_norm 10.0 -> 1.0 (Isaac defaults) - scheduler: hardcoded CosineDecayWithWarmup(10k decay, floor 10% peak) -> DiffuserSchedulerConfig HF cosine with ceil(max_steps * warmup_ratio) warmup, deriving num_training_steps from the outer --steps at runtime - model_params_fp32 (default true): keep master weights in FP32 and compute under BF16 autocast like the native N1.7 recipe (fixes optimizer-update numerics vs pure-BF16 params) - weight-decay grouping via transformers get_parameter_names: biases and norm parameters excluded from decay - restore the TF4 lm_head/embedding weight tie so the unused Qwen LM head stays frozen and deduplicated in checkpoints - action_mask kept in native dtype for the masked flow-matching loss - drop_n_last_frames: exclude episode tails that cannot supply a complete action chunk (Isaac sampler behavior) Verification: tests/policies/groot/test_groot_training_optim_contract.py (7 passed) + remaining groot suite 11 passed/5 skipped on RTX PRO 6000 / CUDA 13.3. Note: tests/policies/groot/test_groot_n1_7.py does not collect on the base branch (pre-existing ImportError, fixed in PR #37). * feat(groot): train-time random crop for N1.7 (eval keeps center crop) Isaac-GR00T crops a random crop_fraction window during training and the deterministic center window at eval, replaying the sampled window across all camera views of a sample. This contract is unchanged since the N1.5 release (gr00t/data/transform/video.py: "If mode is 'train', return a random crop transform. If mode is 'eval', return a center crop transform.") and mirrors LeRobot's own Diffusion/VQBeT crop_is_random pattern. The LeRobot N1.7 port used the eval center crop for training too, so the fine-tuned projector/DiT never sees frame borders and trains on a single fixed appearance point. Scope: crop geometry ONLY - no color jitter, no new dependencies. The random window is plain numpy slicing inside the existing cv2 eval transform: - _transform_n1_7_image_for_vlm_albumentations gains crop_position=(y, x) fractions; None keeps the center crop byte-identical to before (verified by test) - GrootN17VLMEncodeStep gains a runtime-only 'training' flag (never serialized; reloaded pipelines default to eval); training samples ONE window per sample and reuses it across (timestep, view) frames - Isaac's cross-view consistency - gated on torch.is_grad_enabled() so no_grad validation and frozen-eval paths are unaffected - wired via dataset_meta is not None in make_groot_pre_post_processors and the existing _set_groot_preprocessor_training on serialized reloads Verification: tests/policies/groot/test_groot_train_random_crop.py (8 passed: center-crop bit-exactness with crop_position=None, corner/center windows, cross-view replay, train!=eval, no_grad gating, seed reproducibility, serialization contract) + groot suite 23 passed / 5 skipped on RTX PRO 6000 / CUDA 13.3. * docs(groot): update Training & hardware Evaluation commands Replace the multi-GPU accelerate-launch Training snippet with the current single-command 'uv run lerobot-train' N1.7 recipe (relative actions excluding gripper, bf16, flash attention, chunk/n_action_steps=16, bs64/20k steps). Replace the bimanual 'Evaluate in your hardware setup' rollout example with the SO-101 follower RTC 'uv run lerobot-rollout' command (strategy.type=base, inference.type=rtc, wrist+front cameras, place-the-vial task). Docs-only; no source/test changes. * docs(groot): parameterize commands with env vars + fill LIBERO results - Introduce BASE_MODEL / DATASET_ID / REPO_ID / JOB_NAME / OUTPUT_DIR env vars in the training command and reuse OUTPUT_DIR + BASE_MODEL in the rollout cmd. - Fill the LIBERO benchmark table with GR00T-LeRobot success rates (Spatial 94%, Object 98%, Goal 93%, LIBERO 10/Long 90%; avg 93.75%), drop the OSS column and XX placeholders. LeRobot-focused. * docs(groot): drop export block, reference env vars directly Use $DATASET_ID / $BASE_MODEL / $REPO_ID / $OUTPUT_DIR / $JOB_NAME as bare placeholders in the commands without concrete export assignments. * docs(groot): keep BASE_MODEL export in training command * docs(groot): use literal HF repo IDs for dataset/policy repo_id Public-facing Hub references (--dataset.repo_id, --policy.repo_id) shown as concrete IDs; local-only values ($OUTPUT_DIR, $JOB_NAME) stay as placeholders. * docs(groot): add LIBERO training command example * docs(groot): remove LIBERO checkpoints subdirectory section * docs(groot): use $BASE_MODEL for base_model_path in LIBERO eval * docs(groot): drop hf download step from LIBERO eval, fix intro * docs(groot): restore suite checkpoint download intro sentence * docs(groot): remove checkpoint download note above LIBERO eval * docs(groot): update training and rollout commands with new parameters and dependencies * Add sample so101 training command * Remove sample so101 training command * docs(groot): remove optional Flash Attention setup instructions and update base model path for evaluation * docs(groot): update training command with image transformation parameters * docs(groot): add note on inference.queue_threshold value for stable inference * chore(style): pre-commit gr00t * docs(groot): update * chore(policies): minor details * fix(groot): license headers + test guards * chore(policies): fix tests * docs(groot): relative actions param doc * chore(policy): address some of the AI review items --------- Co-authored-by: Andrew Wrenn Co-authored-by: Ryan Halabi Co-authored-by: nv-sachdevkartik Co-authored-by: groot-validation Co-authored-by: johnnynunez Co-authored-by: lbenhorin --- README.md | 2 +- docs/source/_toctree.yml | 2 +- docs/source/envhub_isaaclab_arena.mdx | 2 +- docs/source/groot.mdx | 227 +- docs/source/policy_groot_README.md | 115 +- pyproject.toml | 5 +- src/lerobot/policies/factory.py | 35 +- .../groot/action_head/action_encoder.py | 54 - .../groot/action_head/cross_attention_dit.py | 87 +- .../action_head/flow_matching_action_head.py | 408 --- .../policies/groot/configuration_groot.py | 416 ++- .../configuration_eagle2_5_vl.py | 135 - .../image_processing_eagle2_5_vl_fast.py | 503 --- .../eagle2_hg_model/modeling_eagle2_5_vl.py | 396 --- .../eagle2_hg_model/processing_eagle2_5_vl.py | 541 --- src/lerobot/policies/groot/groot_n1.py | 380 --- src/lerobot/policies/groot/groot_n1_7.py | 951 ++++++ src/lerobot/policies/groot/modeling_groot.py | 410 ++- src/lerobot/policies/groot/processor_groot.py | 2572 ++++++++++++-- src/lerobot/policies/groot/utils.py | 293 +- src/lerobot/utils/import_utils.py | 1 + tests/policies/groot/test_groot_lerobot.py | 33 +- tests/policies/groot/test_groot_n1_7.py | 3025 +++++++++++++++++ .../groot/test_groot_n1_7_oss_parity.py | 259 ++ .../groot/test_groot_state_dropout.py | 100 + .../groot/test_groot_train_random_crop.py | 169 + .../test_groot_training_optim_contract.py | 125 + .../policies/groot/test_groot_vs_original.py | 557 +-- .../groot/utils/dump_original_n1_7.py | 212 ++ uv.lock | 52 +- 30 files changed, 8593 insertions(+), 3474 deletions(-) delete mode 100644 src/lerobot/policies/groot/action_head/action_encoder.py delete mode 100644 src/lerobot/policies/groot/action_head/flow_matching_action_head.py delete mode 100755 src/lerobot/policies/groot/eagle2_hg_model/configuration_eagle2_5_vl.py delete mode 100644 src/lerobot/policies/groot/eagle2_hg_model/image_processing_eagle2_5_vl_fast.py delete mode 100755 src/lerobot/policies/groot/eagle2_hg_model/modeling_eagle2_5_vl.py delete mode 100755 src/lerobot/policies/groot/eagle2_hg_model/processing_eagle2_5_vl.py delete mode 100644 src/lerobot/policies/groot/groot_n1.py create mode 100644 src/lerobot/policies/groot/groot_n1_7.py create mode 100644 tests/policies/groot/test_groot_n1_7.py create mode 100644 tests/policies/groot/test_groot_n1_7_oss_parity.py create mode 100644 tests/policies/groot/test_groot_state_dropout.py create mode 100644 tests/policies/groot/test_groot_train_random_crop.py create mode 100644 tests/policies/groot/test_groot_training_optim_contract.py create mode 100644 tests/policies/groot/utils/dump_original_n1_7.py diff --git a/README.md b/README.md index f72952102..acfad05fb 100644 --- a/README.md +++ b/README.md @@ -105,7 +105,7 @@ lerobot-train \ | -------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | **Imitation Learning** | [ACT](./docs/source/policy_act_README.md), [Diffusion](./docs/source/policy_diffusion_README.md), [VQ-BeT](./docs/source/policy_vqbet_README.md), [Multitask DiT Policy](./docs/source/policy_multi_task_dit_README.md) | | **Reinforcement Learning** | [HIL-SERL](./docs/source/hilserl.mdx), [TDMPC](./docs/source/policy_tdmpc_README.md) & QC-FQL (coming soon) | -| **VLAs Models** | [Pi0](./docs/source/pi0.mdx), [Pi0Fast](./docs/source/pi0fast.mdx), [Pi0.5](./docs/source/pi05.mdx), [GR00T N1.5](./docs/source/policy_groot_README.md), [SmolVLA](./docs/source/policy_smolvla_README.md), [XVLA](./docs/source/xvla.mdx), [EO-1](./docs/source/eo1.mdx), [MolmoAct2](./docs/source/molmoact2.mdx), [WALL-OSS](./docs/source/walloss.mdx) | +| **VLAs Models** | [Pi0](./docs/source/pi0.mdx), [Pi0Fast](./docs/source/pi0fast.mdx), [Pi0.5](./docs/source/pi05.mdx), [GR00T N1.7](./docs/source/policy_groot_README.md), [SmolVLA](./docs/source/policy_smolvla_README.md), [XVLA](./docs/source/xvla.mdx), [EO-1](./docs/source/eo1.mdx), [MolmoAct2](./docs/source/molmoact2.mdx), [WALL-OSS](./docs/source/walloss.mdx) | | **World Models** | [VLA-JEPA](./docs/source/vla_jepa.mdx) (more coming soon) | | **Reward Models** | [SARM](./docs/source/sarm.mdx), [TOPReward](./docs/source/topreward.mdx), [Robometer](./docs/source/robometer.mdx) | diff --git a/docs/source/_toctree.yml b/docs/source/_toctree.yml index 79f4bc124..d4bdf892e 100644 --- a/docs/source/_toctree.yml +++ b/docs/source/_toctree.yml @@ -74,7 +74,7 @@ - local: fastwam title: FastWAM - local: groot - title: NVIDIA GR00T N1.5 + title: NVIDIA GR00T - local: xvla title: X-VLA - local: multi_task_dit diff --git a/docs/source/envhub_isaaclab_arena.mdx b/docs/source/envhub_isaaclab_arena.mdx index b934240d6..cd077806d 100644 --- a/docs/source/envhub_isaaclab_arena.mdx +++ b/docs/source/envhub_isaaclab_arena.mdx @@ -193,7 +193,7 @@ To learn more about training policies with LeRobot, please refer to the training - [SmolVLA](./smolvla) - [Pi0.5](./pi05) -- [GR00T N1.5](./groot) +- [GR00T N1.7](./groot) Sample IsaacLab Arena datasets are available on HuggingFace Hub for experimentation: diff --git a/docs/source/groot.mdx b/docs/source/groot.mdx index 3ab202fb2..af6bfe1ab 100644 --- a/docs/source/groot.mdx +++ b/docs/source/groot.mdx @@ -1,16 +1,19 @@ -# GR00T N1.5 Policy +# GR00T Policy -GR00T N1.5 is an open foundation model from NVIDIA designed for generalized humanoid robot reasoning and skills. It is a cross-embodiment model that accepts multimodal input, including language and images, to perform manipulation tasks in diverse environments. +GR00T is an NVIDIA foundation model family for generalized humanoid robot reasoning and skills. It is a cross-embodiment policy that accepts multimodal input, including language, images, and proprioception, to perform manipulation tasks in diverse environments. -This document outlines the specifics of its integration and usage within the LeRobot framework. +LeRobot integrates GR00T N1.7 through the `groot` policy type. + +> [!WARNING] +> **Breaking change:** GR00T N1.5 support was removed from LeRobot, and current releases support GR00T N1.7 only. N1.5 checkpoints and configs are rejected with a migration note. To keep using an N1.5 checkpoint, pin the last release that supports it: `pip install 'lerobot==0.5.1'`. To use the current release, migrate to GR00T N1.7 (base model [`nvidia/GR00T-N1.7-3B`](https://huggingface.co/nvidia/GR00T-N1.7-3B)). ## Model Overview -NVIDIA Isaac GR00T N1.5 is an upgraded version of the GR00T N1 foundation model. It is built to improve generalization and language-following abilities for humanoid robots. +GR00T N1.7 uses a Cosmos-Reason2/Qwen3-VL backbone and provides checkpoints for SimplerEnv, DROID, and LIBERO. -Developers and researchers can post-train GR00T N1.5 with their own real or synthetic data to adapt it for specific humanoid robots or tasks. +Developers and researchers can post-train GR00T with their own real or synthetic data to adapt it for specific humanoid robots or tasks. -GR00T N1.5 (specifically the GR00T-N1.5-3B model) is built using pre-trained vision and language encoders. It utilizes a flow matching action transformer to model a chunk of actions, conditioned on vision, language, and proprioception. +GR00T uses pre-trained vision and language encoders with a flow matching action transformer to model a chunk of actions conditioned on vision, language, and proprioception. =2.2.1,<2.8.0" "torchvision>=0.21.0,<0.23.0" # --index-url https://download.pytorch.org/whl/cu1XX -pip install ninja "packaging>=24.2,<26.0" # flash attention dependencies -pip install "flash-attn>=2.5.9,<3.0.0" --no-build-isolation -python -c "import flash_attn; print(f'Flash Attention {flash_attn.__version__} imported successfully')" +pip install "lerobot[groot]" ``` -3. Install LeRobot by running: +For a source checkout: ```bash -pip install lerobot[groot] +pip install -e ".[groot]" ``` ## Usage -To use GR00T in your LeRobot configuration, specify the policy type as: +To use GR00T N1.7: -```python -policy.type=groot +```bash +--policy.type=groot ``` ## Training @@ -63,72 +57,171 @@ policy.type=groot Here's a complete training command for finetuning the base GR00T model on your own dataset: +This command is using the `new_embodiment` flag, which is used for the SO-101 robot, [read more about how GR00T handles different embodiments.](https://github.com/NVIDIA/Isaac-GR00T/blob/main/getting_started/policy.md#--embodiment-tag). + ```bash -# Using a multi-GPU setup -accelerate launch \ - --multi_gpu \ - --num_processes=$NUM_GPUS \ - $(which lerobot-train) \ - --output_dir=$OUTPUT_DIR \ - --save_checkpoint=true \ - --batch_size=$BATCH_SIZE \ - --steps=$NUM_STEPS \ - --save_freq=$SAVE_FREQ \ - --log_freq=$LOG_FREQ \ - --policy.push_to_hub=true \ +# install extra deps for training +pip install "lerobot[training]" + +hf auth login +wandb login + +export DATASET_NAME=your_data_set +export HF_USER=your_hf_username +export DATASET=$HF_USER/$DATASET_NAME +export REPO_ID="${DATASET}_GR00T17" #this is the model that will be uploaded to huggingface +export OUTPUT_DIR=outputs/train/$REPO_ID + +lerobot-train \ + --dataset.repo_id=$DATASET \ + --dataset.image_transforms.enable=true \ --policy.type=groot \ + --policy.device=cuda \ + --policy.base_model_path=nvidia/GR00T-N1.7-3B \ + --policy.embodiment_tag=new_embodiment \ + --policy.chunk_size=16 \ + --policy.n_action_steps=16 \ + --policy.use_relative_actions=true \ + --policy.relative_exclude_joints='["gripper"]' \ + --policy.use_bf16=true \ + --policy.push_to_hub=true \ --policy.repo_id=$REPO_ID \ - --policy.tune_diffusion_model=false \ - --dataset.repo_id=$DATASET_ID \ + --seed=42 \ + --batch_size=64 \ + --steps=20000 \ + --save_checkpoint=true \ + --save_freq=5000 \ + --use_policy_training_preset=true \ + --env_eval_freq=0 \ + --eval_steps=0 \ + --log_freq=10 \ + --output_dir=$OUTPUT_DIR \ + --job_name=$DATASET \ --wandb.enable=true \ - --wandb.disable_artifact=true \ - --job_name=$JOB_NAME + --wandb.disable_artifact=true + ``` ## Performance Results -### Libero Benchmark Results +### LIBERO Benchmark Results > [!NOTE] -> Follow our instructions for Libero usage: [Libero](./libero) +> Follow the [LIBERO](./libero) setup instructions before running `lerobot-eval`. -GR00T has demonstrated strong performance on the Libero benchmark suite. To compare and test its LeRobot implementation, we finetuned the GR00T N1.5 model for 30k steps on the Libero dataset and compared the results to the GR00T reference results. +GR00T N1.7 has demonstrated strong performance on the LIBERO benchmark suite. To reproduce LeRobot results, follow the instructions in the [LIBERO](./libero) section. -| Benchmark | LeRobot Implementation | GR00T Reference | -| ------------------ | ---------------------- | --------------- | -| **Libero Spatial** | 82.0% | 92.0% | -| **Libero Object** | 99.0% | 92.0% | -| **Libero Long** | 82.0% | 76.0% | -| **Average** | 87.0% | 87.0% | +### Train on LIBERO -These results demonstrate GR00T's strong generalization capabilities across diverse robotic manipulation tasks. To reproduce these results, you can follow the instructions in the [Libero](https://huggingface.co/docs/lerobot/libero) section. +Example training command for a LIBERO suite (here `libero_spatial`): + +```bash +IMAGE_TRANSFORMS='{ + "brightness": {"weight": 1.0, "type": "ColorJitter", "kwargs": {"brightness": [0.7, 1.3]}}, + "contrast": {"weight": 1.0, "type": "ColorJitter", "kwargs": {"contrast": [0.6, 1.4]}}, + "saturation": {"weight": 1.0, "type": "ColorJitter", "kwargs": {"saturation": [0.5, 1.5]}}, + "hue": {"weight": 1.0, "type": "ColorJitter", "kwargs": {"hue": [-0.08, 0.08]}} +}' + +lerobot-train \ + --dataset.repo_id=IPEC-COMMUNITY/libero_spatial_no_noops_1.0.0_lerobot \ + --dataset.root=/datasets/libero_spatial \ + --dataset.revision=main \ + --dataset.video_backend=pyav \ + --dataset.image_transforms.enable=true \ + --dataset.image_transforms.max_num_transforms=4 \ + --dataset.image_transforms.tfs="$IMAGE_TRANSFORMS" \ + --policy.type=groot \ + --policy.base_model_path=nvidia/GR00T-N1.7-3B \ + --policy.embodiment_tag=libero_sim \ + --policy.push_to_hub=false \ + --policy.use_relative_actions=false \ + --policy.max_steps=20000 \ + --batch_size=320 \ + --steps=20000 \ + --save_freq=2000 \ + --env_eval_freq=0 \ + --eval_steps=0 \ + --log_freq=10 \ + --wandb.enable=true \ + --wandb.project=lerobot \ + --wandb.mode=online \ + --wandb.disable_artifact=true \ + --num_workers=4 \ + --prefetch_factor=2 \ + --persistent_workers=true \ + --output_dir=$OUTPUT_DIR \ + --job_name=$JOB_NAME +``` + +This will follow the recipe found [here](https://github.com/NVIDIA/Isaac-GR00T/blob/main/examples/LIBERO/README.md). + +### GR00T N1.7 LIBERO Results + +Preliminary LeRobot integration results (GR00T-LeRobot, `eval.n_episodes >= 50` per suite): + +| Suite | Success rate | +| ---------------- | -----------: | +| LIBERO Spatial | 94% | +| LIBERO Object | 98% | +| LIBERO Goal | 93% | +| LIBERO 10 (Long) | 90% | +| **Average** | **93.75%** | + +```bash +export MODEL_ID=your_trained_model_on_huggingface + +lerobot-eval \ + --policy.type=groot \ + --policy.base_model_path=$MODEL_ID \ + --policy.embodiment_tag=libero_sim \ + --env.type=libero \ + --env.task=libero_spatial \ + --eval.n_episodes=50 +``` + +Use `eval.n_episodes >= 50` per suite when reporting success rates. ### Evaluate in your hardware setup Once you have trained your model using your parameters you can run inference in your downstream task. Follow the instructions in [Policy Deployment (lerobot-rollout)](./inference). For example: ```bash -lerobot-rollout\ - --strategy.type=sentry \ - --strategy.upload_every_n_episodes=5 \ - --robot.type=bi_so_follower \ - --robot.left_arm_port=/dev/ttyACM1 \ - --robot.right_arm_port=/dev/ttyACM0 \ - --robot.id=bimanual_follower \ - --robot.cameras='{ right: {"type": "opencv", "index_or_path": 0, "width": 640, "height": 480, "fps": 30}, - left: {"type": "opencv", "index_or_path": 2, "width": 640, "height": 480, "fps": 30}, - top: {"type": "opencv", "index_or_path": 4, "width": 640, "height": 480, "fps": 30}, - }' \ +# install extra deps for roullout and real hardware +pip install "lerobot[feetech,viz]" + +export MODEL_ID=your_trained_model_on_huggingface + +# make sure that camera index matches your setup! +# find index using `uv run lerobot-find-cameras opencv` +WRIST_CAM='wrist: {type: opencv, index_or_path: 2, width: 640, height: 480, fps: 30, fourcc: "MJPG"}' +FRONT_CAM='front: {type: opencv, index_or_path: 0, width: 640, height: 480, fps: 30, fourcc: "MJPG"}' +export ROBOT_CAMERAS="{ $WRIST_CAM, $FRONT_CAM }" +export ROBOT_ID=follower_robot +export ROBOT_PORT=/dev/ttyACM0 + +uv run lerobot-rollout \ + --strategy.type=base \ + --policy.path=$MODEL_ID \ + --policy.base_model_path=nvidia/GR00T-N1.7-3B \ + --policy.n_action_steps=8 \ + --robot.type=so101_follower \ + --robot.port=$ROBOT_PORT \ + --robot.id=$ROBOT_ID \ + --robot.cameras="$ROBOT_CAMERAS" \ + --task="place the vial in the rack" \ + --duration=60 \ + --device=cuda \ --display_data=true \ - --dataset.repo_id=/eval_groot-bimanual \ - --dataset.single_task="Grab and handover the red cube to the other arm" \ - --dataset.streaming_encoding=true \ - --dataset.encoder_threads=2 \ - # --dataset.rgb_encoder.vcodec=auto \ - --policy.path=/groot-bimanual \ # your trained model - --duration=600 + --inference.type=rtc \ + --inference.rtc.enabled=True \ # set to False if it causes inference instability + --inference.rtc.execution_horizon=8 \ + --inference.queue_threshold=0 ``` +> [!NOTE] +> Value of `inference.queue_threshold` should not exceed 5 to ensure stable inference. + ## License -This model follows NVIDIA's proprietary license, consistent with the original [GR00T repository](https://github.com/NVIDIA/Isaac-GR00T). Future versions (starting from N1.7) will follow **Apache 2.0 License**. +GR00T N1.7 is released under the [NVIDIA Open Model License Agreement](https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-open-model-license/). diff --git a/docs/source/policy_groot_README.md b/docs/source/policy_groot_README.md index efcd76ebe..4b256fb27 100644 --- a/docs/source/policy_groot_README.md +++ b/docs/source/policy_groot_README.md @@ -1,6 +1,13 @@ ## Research Paper -Paper: https://research.nvidia.com/labs/gear/gr00t-n1_5/ +GR00T N1 technical report (covers the GR00T N1.x family, including N1.7): https://arxiv.org/abs/2503.14734 + +GR00T N1.7 model card: https://huggingface.co/nvidia/GR00T-N1.7-3B + +GR00T N1.5 research page (earlier version): https://research.nvidia.com/labs/gear/gr00t-n1_5/ + +> GR00T N1.5 support was removed from LeRobot; the last release supporting it is `lerobot==0.5.1`. +> Current releases support GR00T N1.7 only. ## Repository @@ -24,4 +31,108 @@ Code: https://github.com/NVIDIA/Isaac-GR00T Blog: https://developer.nvidia.com/isaac/gr00t -Hugging Face Model: https://huggingface.co/nvidia/GR00T-N1.5-3B +Hugging Face Models: + +- GR00T N1.7: https://huggingface.co/nvidia/GR00T-N1.7-3B +- GR00T N1.7 LIBERO checkpoints: https://huggingface.co/nvidia/GR00T-N1.7-LIBERO + +
+Original-vs-LeRobot parity test + +## Original-vs-LeRobot parity test + +`tests/policies/groot/test_groot_vs_original.py` verifies this LeRobot +reimplementation of GR00T N1.7 (Qwen3-VL backbone + flow-matching action head) +against NVIDIA's original `gr00t` package with two comparisons, each parametrized +over every embodiment tag present in the checkpoint: + +1. **Model parity** — given byte-identical pre-processed inputs and the same + flow-matching seed (recorded in each artifact), both implementations must produce + the **same raw model output** (`get_action(...)["action_pred"]`, the normalized + flow-matching prediction). Output shapes must match exactly; any action-horizon + or action-dim mismatch fails the test. +2. **Preprocessor parity** — given the identical raw observations (per-camera + frames, state vectors, language instruction), LeRobot's own preprocessor pipeline + (real Qwen3-VL chat template / tokenizer / image packing + checkpoint-driven + state normalization, no mocks) must produce the **same collated model inputs** + (`input_ids`, `attention_mask`, `pixel_values`, `image_grid_thw`, `state`, + `embodiment_id`) as the original package's processor. + +### Why two environments + +The original `gr00t` package pins `transformers==4.57.3` (Python 3.10); this +integration requires `transformers>=5.x` (Qwen3-VL). Under 5.x, `PretrainedConfig` +is itself a defaulted dataclass, so the original config dataclasses fail to import +(`non-default argument follows default argument`). The two implementations therefore +**cannot be imported in the same Python process**. + +So the test uses a **producer / consumer** split across two venvs: + +1. **Producer** — `tests/policies/groot/utils/dump_original_n1_7.py`, run in the _original_ + gr00t venv. For each embodiment it builds dummy inputs generically from the + checkpoint metadata (state dims from `statistics.json`; camera/language keys from + the processor modality configs), runs the original model, and saves to one `.npz` + per tag: the raw observations (`raw::` keys), the exact collated inputs + (`in::` keys), the seed, and the raw `action_pred`. +2. **Consumer** — the pytest above, run in the _LeRobot_ venv. It discovers every + `.npz`; the model-parity case replays the byte-identical collated inputs through + the LeRobot model with the recorded seed and asserts the outputs match, and the + preprocessor-parity case replays the raw observations through LeRobot's full + preprocessor pipeline and asserts the collated tensors match. + +> Artifacts generated by older versions of the dump script contain no `raw::` +> fields; the preprocessor-parity case then **skips** with a regeneration hint. +> Re-run the producer to refresh them. + +### Fairness controls + +- **Same pre-processed inputs (model parity)** — the original processor's `input_ids`, + `pixel_values`, `image_grid_thw`, `attention_mask`, `state`, `embodiment_id` are + fed verbatim to the LeRobot model (no re-tokenization / re-normalization), so the + model comparison isolates the model. LeRobot's own tokenization / image packing is + covered separately by the preprocessor-parity case, which compares its output + against those same collated tensors from identical raw observations. +- **Same precision + attention kernel** — both sides run **fp32 + SDPA**. The + original defaults to `use_flash_attention=True` (flash_attention_2 + bf16); the + producer forces SDPA + fp32. (With the defaults the gap is ~3e-2 — pure + kernel/rounding noise, not an implementation difference.) +- **Same flow-matching seed** — fixed right before sampling on both sides; the + producer records it in each artifact (`--seed`, default 42) and the consumer + replays the recorded value. + +### How to run + +```bash +# Resolve a local checkpoint (GR00T-N1.7-LIBERO / libero_10) +CKPT=$(python - <<'PY' +import os +from huggingface_hub import snapshot_download +print(os.path.join(snapshot_download("nvidia/GR00T-N1.7-LIBERO", + allow_patterns=["libero_10/*"]), "libero_10")) +PY +) + +# 1) Produce the original-side artifacts for all embodiments (original gr00t venv, CUDA) +CUDA_VISIBLE_DEVICES=0 /path/to/Isaac-GR00T/.venv-original/bin/python \ + tests/policies/groot/utils/dump_original_n1_7.py \ + --ckpt "$CKPT" --out-dir tests/policies/groot/artifacts --device cuda --seed 42 + +# 2) Run the parity test (LeRobot venv) — one parametrized case per embodiment +CUDA_VISIBLE_DEVICES=0 GROOT_PARITY_DEVICE=cuda \ + uv run pytest tests/policies/groot/test_groot_vs_original.py -v -s +``` + +The `.npz` artifacts are local-only (gitignored, ~6–10 MB each) and are regenerated by +the producer; they are never committed. The tests **skip** (do not fail) on CI or +when the checkpoint / artifacts are absent. + +#### Env knobs (all optional) + +| Var | Default | Purpose | +| ----------------------------------------- | -------------------------------- | ------------------------------------- | +| `GROOT_N1_7_PARITY_DIR` | `tests/policies/groot/artifacts` | directory of per-tag `.npz` artifacts | +| `GROOT_N1_7_LIBERO_CKPT` | auto (HF cache) | override checkpoint dir | +| `GROOT_PARITY_DEVICE` | `cuda` if available | `cpu` or `cuda` | +| `GROOT_PARITY_ATOL` / `GROOT_PARITY_RTOL` | `1e-3` | comparison tolerance | + +
diff --git a/pyproject.toml b/pyproject.toml index 84fe07328..882dd0b6f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -219,11 +219,10 @@ groot = [ "lerobot[transformers-dep]", "lerobot[peft-dep]", "lerobot[diffusers-dep]", + "lerobot[dataset]", # NOTE: processor_groot builds a LeRobotDataset for relative-action training stats "dm-tree>=0.1.8,<1.0.0", "timm>=1.0.0,<1.1.0", "decord>=0.6.0,<1.0.0; (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "ninja>=1.11.1,<2.0.0", - "flash-attn>=2.5.9,<3.0.0 ; sys_platform != 'darwin'" ] sarm = ["lerobot[transformers-dep]", "pydantic>=2.0.0,<3.0.0", "faker>=33.0.0,<35.0.0", "lerobot[matplotlib-dep]", "lerobot[qwen-vl-utils-dep]"] robometer = ["lerobot[transformers-dep]", "lerobot[qwen-vl-utils-dep]", "lerobot[peft-dep]"] @@ -315,7 +314,7 @@ all = [ "lerobot[molmoact2]", "lerobot[smolvla]", "lerobot[fastwam]", - # "lerobot[groot]", TODO(Steven): Gr00t requires specific installation instructions for flash-attn + "lerobot[groot]", "lerobot[xvla]", "lerobot[hilserl]", "lerobot[vla_jepa]", diff --git a/src/lerobot/policies/factory.py b/src/lerobot/policies/factory.py index 12871218e..483804a57 100644 --- a/src/lerobot/policies/factory.py +++ b/src/lerobot/policies/factory.py @@ -295,26 +295,23 @@ def make_pre_post_processors( policy configuration type. """ if pretrained_path: - # TODO(Steven): Temporary patch, implement correctly the processors for Gr00t if isinstance(policy_cfg, GrootConfig): - # GROOT handles normalization in groot_pack_inputs_v3 step - # Need to override both stats AND normalize_min_max since saved config might be empty - preprocessor_overrides = {} - postprocessor_overrides = {} - preprocessor_overrides["groot_pack_inputs_v3"] = { - "stats": kwargs.get("dataset_stats"), - "normalize_min_max": True, - } + from .groot.processor_groot import make_groot_pre_post_processors_from_pretrained - # Also ensure postprocessing slices to env action dim and unnormalizes with dataset stats - env_action_dim = policy_cfg.output_features[ACTION].shape[0] - postprocessor_overrides["groot_action_unpack_unnormalize_v1"] = { - "stats": kwargs.get("dataset_stats"), - "normalize_min_max": True, - "env_action_dim": env_action_dim, - } - kwargs["preprocessor_overrides"] = preprocessor_overrides - kwargs["postprocessor_overrides"] = postprocessor_overrides + return make_groot_pre_post_processors_from_pretrained( + config=policy_cfg, + pretrained_path=pretrained_path, + dataset_stats=kwargs.get("dataset_stats"), + dataset_meta=kwargs.get("dataset_meta"), + preprocessor_overrides=kwargs.get("preprocessor_overrides"), + postprocessor_overrides=kwargs.get("postprocessor_overrides"), + preprocessor_config_filename=kwargs.get( + "preprocessor_config_filename", f"{POLICY_PREPROCESSOR_DEFAULT_NAME}.json" + ), + postprocessor_config_filename=kwargs.get( + "postprocessor_config_filename", f"{POLICY_POSTPROCESSOR_DEFAULT_NAME}.json" + ), + ) preprocessor = PolicyProcessorPipeline.from_pretrained( pretrained_model_name_or_path=pretrained_path, @@ -420,6 +417,7 @@ def make_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): @@ -570,6 +568,7 @@ def make_policy( set_dataset_feature_metadata = getattr(cfg, "set_dataset_feature_metadata", None) if callable(set_dataset_feature_metadata): set_dataset_feature_metadata(ds_meta.features) + cfg._runtime_dataset_meta = ds_meta kwargs["config"] = cfg diff --git a/src/lerobot/policies/groot/action_head/action_encoder.py b/src/lerobot/policies/groot/action_head/action_encoder.py deleted file mode 100644 index c6fa0a779..000000000 --- a/src/lerobot/policies/groot/action_head/action_encoder.py +++ /dev/null @@ -1,54 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# 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. - -import torch -import torch.nn as nn - - -def swish(x): - return x * torch.sigmoid(x) - - -class SinusoidalPositionalEncoding(nn.Module): - """ - Produces a sinusoidal encoding of shape (B, T, w) - given timesteps of shape (B, T). - """ - - def __init__(self, embedding_dim): - super().__init__() - self.embedding_dim = embedding_dim - - def forward(self, timesteps): - # timesteps: shape (B, T) - # We'll compute sin/cos frequencies across dim T - timesteps = timesteps.float() # ensure float - - b, t = timesteps.shape - device = timesteps.device - - half_dim = self.embedding_dim // 2 - # typical log space frequencies for sinusoidal encoding - exponent = -torch.arange(half_dim, dtype=torch.float, device=device) * ( - torch.log(torch.tensor(10000.0)) / half_dim - ) - # Expand timesteps to (B, T, 1) then multiply - freqs = timesteps.unsqueeze(-1) * exponent.exp() # (B, T, half_dim) - - sin = torch.sin(freqs) - cos = torch.cos(freqs) - enc = torch.cat([sin, cos], dim=-1) # (B, T, w) - - return enc diff --git a/src/lerobot/policies/groot/action_head/cross_attention_dit.py b/src/lerobot/policies/groot/action_head/cross_attention_dit.py index a4cd1a0b7..697459240 100755 --- a/src/lerobot/policies/groot/action_head/cross_attention_dit.py +++ b/src/lerobot/policies/groot/action_head/cross_attention_dit.py @@ -1,11 +1,12 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 +#!/usr/bin/env python + +# Copyright 2025 NVIDIA Corporation and 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 +# 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, @@ -14,6 +15,7 @@ # limitations under the License. +import logging from typing import TYPE_CHECKING import torch @@ -42,6 +44,9 @@ else: Timesteps = None +logger = logging.getLogger(__name__) + + class TimestepEncoder(nn.Module): def __init__(self, embedding_dim, compute_dtype=torch.float32): require_package("diffusers", extra="groot") @@ -181,8 +186,7 @@ class BasicTransformerBlock(nn.Module): attn_output = self.attn1( norm_hidden_states, encoder_hidden_states=encoder_hidden_states, - attention_mask=attention_mask, - # encoder_attention_mask=encoder_attention_mask, + attention_mask=encoder_attention_mask if encoder_hidden_states is not None else attention_mask, ) if self.final_dropout: attn_output = self.final_dropout(attn_output) @@ -266,8 +270,8 @@ class DiT(ModelMixin, ConfigMixin): self.norm_out = nn.LayerNorm(self.inner_dim, elementwise_affine=False, eps=1e-6) self.proj_out_1 = nn.Linear(self.inner_dim, 2 * self.inner_dim) self.proj_out_2 = nn.Linear(self.inner_dim, self.config.output_dim) - print( - "Total number of DiT parameters: ", + logger.debug( + "Total number of DiT parameters: %d", sum(p.numel() for p in self.parameters() if p.requires_grad), ) @@ -318,6 +322,71 @@ class DiT(ModelMixin, ConfigMixin): return self.proj_out_2(hidden_states) +class AlternateVLDiT(DiT): + """N1.7 DiT variant that alternates cross-attention over image and text tokens.""" + + def __init__(self, *args, attend_text_every_n_blocks: int = 2, **kwargs): + super().__init__(*args, **kwargs) + self.attend_text_every_n_blocks = attend_text_every_n_blocks + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + timestep: torch.LongTensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + return_all_hidden_states: bool = False, + image_mask: torch.Tensor | None = None, + backbone_attention_mask: torch.Tensor | None = None, + ): + if image_mask is None: + raise ValueError("image_mask is required for AlternateVLDiT.") + if backbone_attention_mask is None: + raise ValueError("backbone_attention_mask is required for AlternateVLDiT.") + + temb = self.timestep_encoder(timestep) + hidden_states = hidden_states.contiguous() + encoder_hidden_states = encoder_hidden_states.contiguous() + + image_attention_mask = image_mask & backbone_attention_mask + non_image_attention_mask = (~image_mask) & backbone_attention_mask + + all_hidden_states = [hidden_states] + if not self.config.interleave_self_attention: + raise ValueError("AlternateVLDiT requires interleave_self_attention=True.") + + for idx, block in enumerate(self.transformer_blocks): + if idx % 2 == 1: + hidden_states = block( + hidden_states, + attention_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + temb=temb, + ) + else: + curr_encoder_attention_mask = ( + non_image_attention_mask + if idx % (2 * self.attend_text_every_n_blocks) == 0 + else image_attention_mask + ) + hidden_states = block( + hidden_states, + attention_mask=None, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=curr_encoder_attention_mask, + temb=temb, + ) + all_hidden_states.append(hidden_states) + + conditioning = temb + shift, scale = self.proj_out_1(F.silu(conditioning)).chunk(2, dim=1) + hidden_states = self.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None] + if return_all_hidden_states: + return self.proj_out_2(hidden_states), all_hidden_states + return self.proj_out_2(hidden_states) + + class SelfAttentionTransformer(ModelMixin, ConfigMixin): _supports_gradient_checkpointing = True @@ -362,8 +431,8 @@ class SelfAttentionTransformer(ModelMixin, ConfigMixin): for _ in range(self.config.num_layers) ] ) - print( - "Total number of SelfAttentionTransformer parameters: ", + logger.debug( + "Total number of SelfAttentionTransformer parameters: %d", sum(p.numel() for p in self.parameters() if p.requires_grad), ) diff --git a/src/lerobot/policies/groot/action_head/flow_matching_action_head.py b/src/lerobot/policies/groot/action_head/flow_matching_action_head.py deleted file mode 100644 index 986820670..000000000 --- a/src/lerobot/policies/groot/action_head/flow_matching_action_head.py +++ /dev/null @@ -1,408 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# 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. - -from dataclasses import field -from typing import TYPE_CHECKING - -import torch -import torch.nn.functional as F # noqa: N812 -from torch import nn -from torch.distributions import Beta - -from lerobot.utils.import_utils import _transformers_available - -# Conditional import for type checking and lazy loading -if TYPE_CHECKING or _transformers_available: - from transformers import PretrainedConfig - from transformers.feature_extraction_utils import BatchFeature -else: - PretrainedConfig = object - BatchFeature = None - -from .action_encoder import ( - SinusoidalPositionalEncoding, - swish, -) -from .cross_attention_dit import DiT, SelfAttentionTransformer - - -class CategorySpecificLinear(nn.Module): - def __init__(self, num_categories, input_dim, hidden_dim): - super().__init__() - self.num_categories = num_categories - # For each category, we have separate weights and biases. - self.W = nn.Parameter(0.02 * torch.randn(num_categories, input_dim, hidden_dim)) - self.b = nn.Parameter(torch.zeros(num_categories, hidden_dim)) - - def forward(self, x, cat_ids): - selected_w = self.W[cat_ids] - selected_b = self.b[cat_ids] - return torch.bmm(x, selected_w) + selected_b.unsqueeze(1) - - -class CategorySpecificMLP(nn.Module): - def __init__(self, num_categories, input_dim, hidden_dim, output_dim): - super().__init__() - self.num_categories = num_categories - self.layer1 = CategorySpecificLinear(num_categories, input_dim, hidden_dim) - self.layer2 = CategorySpecificLinear(num_categories, hidden_dim, output_dim) - - def forward(self, x, cat_ids): - hidden = F.relu(self.layer1(x, cat_ids)) - return self.layer2(hidden, cat_ids) - - -class MultiEmbodimentActionEncoder(nn.Module): - def __init__(self, action_dim, hidden_size, num_embodiments): - super().__init__() - self.hidden_size = hidden_size - self.num_embodiments = num_embodiments - - # W1: R^{w x d}, W2: R^{w x 2w}, W3: R^{w x w} - self.W1 = CategorySpecificLinear(num_embodiments, action_dim, hidden_size) # (d -> w) - self.W2 = CategorySpecificLinear(num_embodiments, 2 * hidden_size, hidden_size) # (2w -> w) - self.W3 = CategorySpecificLinear(num_embodiments, hidden_size, hidden_size) # (w -> w) - self.pos_encoding = SinusoidalPositionalEncoding(hidden_size) - - def forward(self, actions, timesteps, cat_ids): - """ - actions: shape (B, T, action_dim) - timesteps: shape (B,) -- a single scalar per batch item - cat_ids: shape (B,) - returns: shape (B, T, hidden_size) - """ - b, t, _ = actions.shape - - # 1) Expand each batch's single scalar time 'tau' across all T steps - # so that shape => (B, T) - # e.g. if timesteps is (B,), replicate across T - if timesteps.dim() == 1 and timesteps.shape[0] == b: - # shape (B,) => (B,T) - timesteps = timesteps.unsqueeze(1).expand(-1, t) - else: - raise ValueError("Expected `timesteps` to have shape (B,) so we can replicate across T.") - - # 2) Standard action MLP step for shape => (B, T, w) - a_emb = self.W1(actions, cat_ids) - - # 3) Get the sinusoidal encoding (B, T, w) - tau_emb = self.pos_encoding(timesteps).to(dtype=a_emb.dtype) - - # 4) Concat along last dim => (B, T, 2w), then W2 => (B, T, w), swish - x = torch.cat([a_emb, tau_emb], dim=-1) - x = swish(self.W2(x, cat_ids)) - - # 5) Finally W3 => (B, T, w) - x = self.W3(x, cat_ids) - return x - - -class FlowmatchingActionHeadConfig(PretrainedConfig): - """NOTE: N1.5 uses XEmbFlowmatchingPolicyHeadConfig as action head""" - - add_pos_embed: bool = field(default=True, metadata={"help": "Whether to add positional embedding"}) - model_dtype: str = field(default="float32", metadata={"help": "Model data type."}) - diffusion_model_cfg: dict = field(default=None, metadata={"help": "Diffusion model configuration."}) - input_embedding_dim: int = field(default=1536, metadata={"help": "Input embedding channel dimension."}) - backbone_embedding_dim: int = field( - default=1536, metadata={"help": "Backbone embedding channel dimension."} - ) - - hidden_size: int = field(default=1024, metadata={"help": "Input embedding dimension."}) - max_seq_len: int = field(default=1024, metadata={"help": "Maximum Sequence Length"}) - action_dim: int = field(default=None, metadata={"help": "Action dimension."}) - action_horizon: int = field(default=None, metadata={"help": "Action horizon."}) - noise_beta_alpha: float = field(default=1.5, metadata={"help": ""}) - noise_beta_beta: float = field(default=1.0, metadata={"help": ""}) - noise_s: float = field(default=0.999, metadata={"help": "Flow matching noise Beta distribution s."}) - num_timestep_buckets: int = field( - default=1000, metadata={"help": "Number of timestep discretization buckets."} - ) - num_inference_timesteps: int = field( - default=None, - metadata={"help": "Number of inference steps for noise diffusion."}, - ) - max_num_embodiments: int = field(default=32, metadata={"help": "Number of embodiments."}) - tune_projector: bool = field(default=True, metadata={"help": "Whether to tune the projector."}) - tune_diffusion_model: bool = field( - default=True, metadata={"help": "Whether to tune the diffusion model."} - ) - load_pretrained_det_decode_layer_path: str = field( - default=None, metadata={"help": "Path to pretrained detection model."} - ) - detection_coeff: float = field(default=1.0, metadata={"help": "Detection coefficient."}) - - freeze_decode_layer: bool = field(default=False) - expand_batch: int = field(default=None) - use_vlln: bool = field(default=True) - - vl_self_attention_cfg: dict = field(default=None) - num_target_vision_tokens: int = field(default=32, metadata={"help": "Number of target vision tokens."}) - - def __init__(self, **kwargs): - super().__init__(**kwargs) - for key, value in kwargs.items(): - setattr(self, key, value) - - -class FlowmatchingActionHead(nn.Module): - config_class = FlowmatchingActionHeadConfig - supports_gradient_checkpointing = True - - def __init__( - self, - config: FlowmatchingActionHeadConfig, - ): - super().__init__() - self.hidden_size = config.hidden_size - self.input_embedding_dim = config.input_embedding_dim - - self.model = DiT(**config.diffusion_model_cfg) - self.action_dim = config.action_dim - self.action_horizon = config.action_horizon - self.num_inference_timesteps = config.num_inference_timesteps - - self.state_encoder = CategorySpecificMLP( - num_categories=config.max_num_embodiments, - input_dim=config.max_state_dim, - hidden_dim=self.hidden_size, - output_dim=self.input_embedding_dim, - ) - self.action_encoder = MultiEmbodimentActionEncoder( - action_dim=config.action_dim, - hidden_size=self.input_embedding_dim, - num_embodiments=config.max_num_embodiments, - ) - self.action_decoder = CategorySpecificMLP( - num_categories=config.max_num_embodiments, - input_dim=self.hidden_size, - hidden_dim=self.hidden_size, - output_dim=self.action_dim, - ) - self.future_tokens = nn.Embedding(config.num_target_vision_tokens, self.input_embedding_dim) - nn.init.normal_(self.future_tokens.weight, mean=0.0, std=0.02) - - self.vlln = nn.LayerNorm(config.backbone_embedding_dim) if config.use_vlln else nn.Identity() - self.vl_self_attention = ( - SelfAttentionTransformer(**config.vl_self_attention_cfg) if config.use_vlln else nn.Identity() - ) - - if config.add_pos_embed: - self.position_embedding = nn.Embedding(config.max_seq_len, self.input_embedding_dim) - nn.init.normal_(self.position_embedding.weight, mean=0.0, std=0.02) - - self._noise_beta_alpha = config.noise_beta_alpha - self._noise_beta_beta = config.noise_beta_beta - self._beta_dist = None - self.num_timestep_buckets = config.num_timestep_buckets - self.config = config - self.set_trainable_parameters(config.tune_projector, config.tune_diffusion_model) - - def set_trainable_parameters(self, tune_projector: bool, tune_diffusion_model: bool): - self.tune_projector = tune_projector - self.tune_diffusion_model = tune_diffusion_model - for p in self.parameters(): - p.requires_grad = True - if not tune_projector: - self.state_encoder.requires_grad_(False) - self.action_encoder.requires_grad_(False) - self.action_decoder.requires_grad_(False) - if self.config.add_pos_embed: - self.position_embedding.requires_grad_(False) - if not tune_diffusion_model: - self.model.requires_grad_(False) - print(f"Tune action head projector: {self.tune_projector}") - print(f"Tune action head diffusion model: {self.tune_diffusion_model}") - # Check if any parameters are still trainable. If not, print a warning. - if not tune_projector and not tune_diffusion_model: - for name, p in self.named_parameters(): - if p.requires_grad: - print(f"Action head trainable parameter: {name}") - if not any(p.requires_grad for p in self.parameters()): - print("Warning: No action head trainable parameters found.") - - def set_frozen_modules_to_eval_mode(self): - """ - Huggingface will call model.train() at each training_step. To ensure - the expected behaviors for modules like dropout, batchnorm, etc., we - need to call model.eval() for the frozen modules. - """ - if self.training: - if not self.tune_projector: - self.state_encoder.eval() - self.action_encoder.eval() - self.action_decoder.eval() - if self.config.add_pos_embed: - self.position_embedding.eval() - if not self.tune_diffusion_model: - self.model.eval() - - def sample_time(self, batch_size, device, dtype): - if self._beta_dist is None: - self._beta_dist = Beta(self._noise_beta_alpha, self._noise_beta_beta, validate_args=False) - sample = self._beta_dist.sample([batch_size]).to(device, dtype=dtype) - return (self.config.noise_s - sample) / self.config.noise_s - - def prepare_input(self, batch: dict) -> BatchFeature: - return BatchFeature(data=batch) - - def process_backbone_output(self, backbone_output: BatchFeature) -> BatchFeature: - backbone_features = backbone_output["backbone_features"] - backbone_features = self.vlln(backbone_features) - backbone_features = self.vl_self_attention(backbone_features) - backbone_output["backbone_features"] = backbone_features - return backbone_output - - def forward(self, backbone_output: BatchFeature, action_input: BatchFeature) -> BatchFeature: - # Set frozen modules to eval - self.set_frozen_modules_to_eval_mode() - - backbone_output = self.process_backbone_output(backbone_output) - - if self.config.expand_batch is not None: - for k, v in backbone_output.items(): - ndim = len(v.shape) - factors = [self.config.expand_batch] - while len(factors) < ndim: - factors.append(1) - factors = tuple(factors) - expanded = v.repeat(*factors) - backbone_output[k] = expanded - - for k, v in action_input.items(): - ndim = len(v.shape) - factors = [self.config.expand_batch] - while len(factors) < ndim: - factors.append(1) - factors = tuple(factors) - expanded = v.repeat(*factors) - action_input[k] = expanded - - # Get vision and language embeddings. - vl_embs = backbone_output.backbone_features - device = vl_embs.device - - # Get embodiment ID. - embodiment_id = action_input.embodiment_id - - # Embed state. - state_features = self.state_encoder(action_input.state, embodiment_id) - - # Embed noised action trajectory. - actions = action_input.action - noise = torch.randn(actions.shape, device=actions.device, dtype=actions.dtype) - t = self.sample_time(actions.shape[0], device=actions.device, dtype=actions.dtype) - t = t[:, None, None] # shape (B,1,1) for broadcast - - noisy_trajectory = (1 - t) * noise + t * actions - velocity = actions - noise - - # Convert (continuous) t -> discrete if needed - t_discretized = (t[:, 0, 0] * self.num_timestep_buckets).long() - action_features = self.action_encoder(noisy_trajectory, t_discretized, embodiment_id) - - # Maybe add position embedding. - if self.config.add_pos_embed: - pos_ids = torch.arange(action_features.shape[1], dtype=torch.long, device=device) - pos_embs = self.position_embedding(pos_ids).unsqueeze(0) - action_features = action_features + pos_embs - - # Join vision, language, state and action embedding along sequence dimension. - future_tokens = self.future_tokens.weight.unsqueeze(0).expand(vl_embs.shape[0], -1, -1) - sa_embs = torch.cat((state_features, future_tokens, action_features), dim=1) - - vl_attn_mask = backbone_output.backbone_attention_mask - - model_output = self.model( - hidden_states=sa_embs, - encoder_hidden_states=vl_embs, - encoder_attention_mask=vl_attn_mask, - timestep=t_discretized, - return_all_hidden_states=False, # NOTE (YL): not using flare now - ) - pred = self.action_decoder(model_output, embodiment_id) - pred_actions = pred[:, -actions.shape[1] :] - - # Slice out only the action portion of pred and target. - action_mask = action_input.action_mask - loss = F.mse_loss(pred_actions, velocity, reduction="none") * action_mask - loss = loss.sum() / action_mask.sum() - output_dict = { - "loss": loss, - } - return BatchFeature(data=output_dict) - - @torch.no_grad() - def get_action(self, backbone_output: BatchFeature, action_input: BatchFeature) -> BatchFeature: - backbone_output = self.process_backbone_output(backbone_output) - - # Get vision and language embeddings. - vl_embs = backbone_output.backbone_features - embodiment_id = action_input.embodiment_id - - # Embed state. - state_features = self.state_encoder(action_input.state, embodiment_id) - - # Set initial actions as the sampled noise. - batch_size = vl_embs.shape[0] - device = vl_embs.device - actions = torch.randn( - size=(batch_size, self.config.action_horizon, self.config.action_dim), - dtype=vl_embs.dtype, - device=device, - ) - - num_steps = self.num_inference_timesteps - dt = 1.0 / num_steps - - # Run denoising steps. - for t in range(num_steps): - t_cont = t / float(num_steps) # e.g. goes 0, 1/N, 2/N, ... - t_discretized = int(t_cont * self.num_timestep_buckets) - - # Embed noised action trajectory. - timesteps_tensor = torch.full(size=(batch_size,), fill_value=t_discretized, device=device) - action_features = self.action_encoder(actions, timesteps_tensor, embodiment_id) - # Maybe add position embedding. - if self.config.add_pos_embed: - pos_ids = torch.arange(action_features.shape[1], dtype=torch.long, device=device) - pos_embs = self.position_embedding(pos_ids).unsqueeze(0) - action_features = action_features + pos_embs - - # Join vision, language, state and action embedding along sequence dimension. - future_tokens = self.future_tokens.weight.unsqueeze(0).expand(vl_embs.shape[0], -1, -1) - sa_embs = torch.cat((state_features, future_tokens, action_features), dim=1) - - # Run model forward. - model_output = self.model( - hidden_states=sa_embs, - encoder_hidden_states=vl_embs, - timestep=timesteps_tensor, - ) - pred = self.action_decoder(model_output, embodiment_id) - - pred_velocity = pred[:, -self.action_horizon :] - - # Update actions using euler integration. - actions = actions + dt * pred_velocity - return BatchFeature(data={"action_pred": actions}) - - @property - def device(self): - return next(iter(self.parameters())).device - - @property - def dtype(self): - return next(iter(self.parameters())).dtype diff --git a/src/lerobot/policies/groot/configuration_groot.py b/src/lerobot/policies/groot/configuration_groot.py index 17cb631d7..97e08bb76 100644 --- a/src/lerobot/policies/groot/configuration_groot.py +++ b/src/lerobot/policies/groot/configuration_groot.py @@ -14,12 +14,229 @@ # See the License for the specific language governing permissions and # limitations under the License. +import logging +import math from dataclasses import dataclass, field +from pathlib import Path from lerobot.configs import FeatureType, NormalizationMode, PolicyFeature, PreTrainedConfig -from lerobot.optim import AdamWConfig, CosineDecayWithWarmupSchedulerConfig +from lerobot.optim import AdamWConfig, DiffuserSchedulerConfig from lerobot.utils.constants import ACTION, OBS_STATE +from .utils import read_json + +logger = logging.getLogger(__name__) + +GROOT_N1_7 = "n1.7" +# Legacy GR00T N1.5 identifier. N1.5 is NOT a supported model_version (it is +# intentionally absent from _GROOT_MODEL_VERSION_ALIASES so normalize_groot_model_version +# still rejects it). It is retained only so that infer_groot_model_version can recognise +# an N1.5 base path/checkpoint and the N1.7 config/loader can reject the mismatch. +GROOT_N1_5 = "n1.5" +# Canonical guidance appended to every error raised when an N1.5 checkpoint, config, +# or processor pipeline is detected. Keep this message in sync with docs/source/groot.mdx. +GROOT_N1_5_REMOVAL_GUIDANCE = ( + "GR00T N1.5 support was removed from LeRobot. " + "To keep using an N1.5 checkpoint, pin the last release that supports it: " + "`pip install 'lerobot==0.5.1'`. To use the current release, migrate to GR00T N1.7 " + "(model_version='n1.7', base model nvidia/GR00T-N1.7-3B)." +) +GROOT_N1_7_BASE_MODEL = "nvidia/GR00T-N1.7-3B" +GROOT_N1_7_BACKBONE_MODEL = "nvidia/Cosmos-Reason2-2B" +# Default GR00T N1.7 training resolution. Fallback if processor_config lacks sizing. Prevents mismatched +# full-res patchification by forcing a resize. Mirrored by GR00T_N1_7_DEFAULTS in groot_n1_7.py. +N1_7_DEFAULT_IMAGE_TARGET_SIZE = (256, 256) +N1_7_DEFAULT_IMAGE_CROP_SIZE = (230, 230) +GROOT_ACTION_DECODE_TRANSFORM_LIBERO = "libero" +# Sentinel meaning "the user did not pick an action decode transform": __post_init__ resolves it +# to the embodiment default ('libero' for 'libero_sim', otherwise None). It is distinct from an +# explicit 'none' (resolved to None) so an opt-out survives a draccus save/load round-trip. +GROOT_ACTION_DECODE_TRANSFORM_AUTO = "auto" + +_GROOT_MODEL_VERSION_ALIASES = { + "n1.7": GROOT_N1_7, + "n1_7": GROOT_N1_7, + "n1d7": GROOT_N1_7, + "n17": GROOT_N1_7, + "1.7": GROOT_N1_7, +} + +# Legacy N1.5 spellings, kept ONLY so they can be detected and rejected with +# GROOT_N1_5_REMOVAL_GUIDANCE (see GROOT_N1_5 above). Never map these to a supported version. +_GROOT_N1_5_VERSION_ALIASES = {"n1.5", "n1_5", "n1d5", "n15", "1.5"} + +_GROOT_ACTION_DECODE_TRANSFORM_ALIASES = { + GROOT_ACTION_DECODE_TRANSFORM_AUTO: GROOT_ACTION_DECODE_TRANSFORM_AUTO, + "none": None, + "": None, + GROOT_ACTION_DECODE_TRANSFORM_LIBERO: GROOT_ACTION_DECODE_TRANSFORM_LIBERO, +} + + +def normalize_groot_model_version(model_version: str) -> str: + normalized = _GROOT_MODEL_VERSION_ALIASES.get(model_version.lower()) + if normalized is None: + supported = GROOT_N1_7 + message = f"Unsupported GR00T model_version '{model_version}'. Supported versions: {supported}." + if model_version.lower() in _GROOT_N1_5_VERSION_ALIASES: + message = f"{message} {GROOT_N1_5_REMOVAL_GUIDANCE}" + raise ValueError(message) + return normalized + + +def normalize_groot_action_decode_transform(transform: str | None) -> str | None: + if transform is None: + return None + normalized = _GROOT_ACTION_DECODE_TRANSFORM_ALIASES.get(transform.lower()) + if normalized is None and transform.lower() not in _GROOT_ACTION_DECODE_TRANSFORM_ALIASES: + supported = ", ".join( + sorted(key for key, value in _GROOT_ACTION_DECODE_TRANSFORM_ALIASES.items() if value is not None) + ) + raise ValueError( + f"Unsupported GR00T N1.7 action decode transform '{transform}'. " + f"Supported transforms: none, {supported}." + ) + return normalized + + +def infer_groot_model_version(model_path: str | None) -> str | None: + if not model_path: + return None + model_path_lower = model_path.lower() + if "gr00t-n1.7" in model_path_lower or "gr00t_n1.7" in model_path_lower: + return GROOT_N1_7 + # Detect legacy N1.5 paths so the N1.7 config/loader can reject the mismatch. + # N1.5 is unsupported, but it must still be recognised here to fail loudly + # rather than silently treating an N1.5 checkpoint as N1.7. + if "gr00t-n1.5" in model_path_lower or "gr00t_n1.5" in model_path_lower: + return GROOT_N1_5 + config_version = _infer_groot_model_version_from_local_config(model_path) + if config_version is not None: + return config_version + return None + + +def is_raw_groot_n1_7_checkpoint(model_path: str | Path | None) -> bool: + if model_path is None: + return False + + path = Path(model_path).expanduser() + if path.is_dir(): + config_path = path / "config.json" + elif path.name == "config.json": + config_path = path + else: + return False + + config = read_json(config_path) + return "type" not in config and _infer_groot_model_version_from_config(config) == GROOT_N1_7 + + +def infer_groot_n1_7_embodiment_tag(model_path: str | Path | None) -> str | None: + if model_path is None: + return None + + processor_config_path = Path(model_path).expanduser() / "processor_config.json" + processor_config = read_json(processor_config_path) + + modality_configs = processor_config.get("processor_kwargs", {}).get("modality_configs", {}) + if not isinstance(modality_configs, dict): + return None + if "libero_sim" in modality_configs: + return "libero_sim" + if len(modality_configs) == 1: + return next(iter(modality_configs)) + return None + + +def infer_groot_n1_7_action_horizon( + model_path: str | Path | None, embodiment_tag: str | None = None +) -> int | None: + if model_path is None: + return None + + processor_config_path = Path(model_path).expanduser() / "processor_config.json" + processor_config = read_json(processor_config_path) + + processor_kwargs = processor_config.get("processor_kwargs", {}) + if not isinstance(processor_kwargs, dict): + return None + modality_configs = processor_kwargs.get("modality_configs", {}) + if not isinstance(modality_configs, dict): + return None + + if embodiment_tag is None: + embodiment_tag = infer_groot_n1_7_embodiment_tag(model_path) + if embodiment_tag is None: + return None + + embodiment_config = modality_configs.get(embodiment_tag, {}) + if not isinstance(embodiment_config, dict): + return None + action_config = embodiment_config.get("action", {}) + if not isinstance(action_config, dict): + return None + delta_indices = action_config.get("delta_indices", []) + if not isinstance(delta_indices, list): + return None + return len(delta_indices) or None + + +def infer_groot_n1_7_action_execution_horizon( + model_path: str | Path | None, embodiment_tag: str | None = None +) -> int | None: + action_horizon = infer_groot_n1_7_action_horizon(model_path, embodiment_tag) + if action_horizon is None: + return None + + if embodiment_tag is None: + embodiment_tag = infer_groot_n1_7_embodiment_tag(model_path) + if embodiment_tag == "libero_sim": + # NVIDIA's N1.7 LIBERO rollout wrapper replans after 8 of the 16 decoded + # actions. Keeping that execution cadence avoids stale open-loop chunks. + return min(action_horizon, 8) + return action_horizon + + +def _infer_groot_model_version_from_local_config(model_path: str) -> str | None: + path = Path(model_path).expanduser() + if path.is_dir(): + config_path = path / "config.json" + elif path.name == "config.json": + config_path = path + else: + return None + + return _infer_groot_model_version_from_config(read_json(config_path)) + + +def _infer_groot_model_version_from_config(config: dict) -> str | None: + model_version = config.get("model_version") + if isinstance(model_version, str): + if model_version.lower() in _GROOT_N1_5_VERSION_ALIASES: + return GROOT_N1_5 + try: + return normalize_groot_model_version(model_version) + except ValueError: + return None + + candidates = [config.get("model_type"), *(config.get("architectures") or [])] + for candidate in candidates: + if not isinstance(candidate, str): + continue + normalized = candidate.lower().replace("-", "_") + if normalized in {"gr00tn1d7", "gr00t_n1d7", "gr00t_n1_7"}: + return GROOT_N1_7 + if normalized in {"gr00t_n1_5", "gr00tn1_5", "gr00t_n15", "gr00t_n1d5", "gr00tn1d5"}: + return GROOT_N1_5 + if config.get("model_name") == GROOT_N1_7_BACKBONE_MODEL: + return GROOT_N1_7 + # The Eagle VLM backbone is specific to pre-N1.7 GR00T checkpoints (N1.7 uses Cosmos/Qwen3-VL). + backbone_cfg = config.get("backbone_cfg") + if isinstance(backbone_cfg, dict) and "eagle_path" in backbone_cfg: + return GROOT_N1_5 + return None + @PreTrainedConfig.register_subclass("groot") @dataclass @@ -28,35 +245,44 @@ class GrootConfig(PreTrainedConfig): # Basic policy settings n_obs_steps: int = 1 - chunk_size: int = 50 - n_action_steps: int = 50 + chunk_size: int = 40 + n_action_steps: int = 40 # Dimension settings (must match pretrained GR00T model expectations) # Maximum state dimension. Shorter states will be zero-padded. - max_state_dim: int = 64 + max_state_dim: int = 132 # Maximum action dimension. Shorter actions will be zero-padded. - max_action_dim: int = 32 + max_action_dim: int = 132 - # Normalization (start with identity, adjust as needed) + # GR00T normalizes state/action internally in its processor steps (min/max with + # q01/q99 percentiles, per embodiment), and the Qwen3-VL backbone's image processor + # handles image normalization. The policy therefore does NOT use LeRobot's + # NormalizerProcessorStep/UnnormalizerProcessorStep, so this mapping is intentionally + # IDENTITY for every feature and is not consulted by make_groot_pre_post_processors. normalization_mapping: dict[str, NormalizationMode] = field( default_factory=lambda: { "VISUAL": NormalizationMode.IDENTITY, - "STATE": NormalizationMode.MEAN_STD, - "ACTION": NormalizationMode.MEAN_STD, + "STATE": NormalizationMode.IDENTITY, + "ACTION": NormalizationMode.IDENTITY, } ) - # Image preprocessing (adjust to match Groot's expected input) - image_size: tuple[int, int] = (224, 224) + # Groot-specific model parameters - # Groot-specific model parameters (from groot_finetune_script.py) + # Path or HuggingFace model ID for the base GR00T N1.7 model whose backbone weights and + # checkpoint sidecars (statistics.json, processor_config.json, ...) are loaded. This is the + # model *source*, and is intentionally distinct from the inherited `pretrained_path`: + # `pretrained_path` (`--policy.path`) points at a saved LeRobot checkpoint directory whose + # `config.json` carries a `type` field, whereas a raw NVIDIA GR00T checkpoint has no such + # field and so can only be loaded through `base_model_path` (`--policy.base_model_path`). + # Defaults to GROOT_N1_7_BASE_MODEL when unset (resolved in __post_init__). + base_model_path: str | None = None - # Path or HuggingFace model ID for the base Groot model - base_model_path: str = "nvidia/GR00T-N1.5-3B" - - # HF repo ID (or local path) that hosts vocab.json and merges.txt for Eagle tokenizer. - tokenizer_assets_repo: str = "lerobot/eagle2hg-processor-groot-n1p5" + # Optional named action transform applied after raw N1.7 checkpoint decoding and before env.step(). + # 'auto' (default) resolves to the embodiment default ('libero' for 'libero_sim', otherwise no + # transform). Pass 'none' to explicitly disable the transform, including for 'libero_sim'. + action_decode_transform: str | None = GROOT_ACTION_DECODE_TRANSFORM_AUTO # Embodiment tag to use for training (e.g. 'new_embodiment', 'gr1') embodiment_tag: str = "new_embodiment" @@ -75,38 +301,67 @@ class GrootConfig(PreTrainedConfig): # Whether to fine-tune the diffusion model tune_diffusion_model: bool = True - # LoRA parameters (from groot_finetune_script.py) - # Rank for the LORA model. If 0, no LORA will be used. - lora_rank: int = 0 + # Whether to fine-tune the VL LayerNorm + VL self-attention projector in the action head. + tune_vlln: bool = True - # Alpha value for the LORA model - lora_alpha: int = 16 + # Number of top LLM backbone layers to fine-tune (0 = none). Lets you adapt just the final + # language layers without unfreezing the whole backbone; independent of `tune_llm`, which tunes + # the entire LLM. + tune_top_llm_layers: int = 0 - # Dropout rate for the LORA model - lora_dropout: float = 0.1 + # Inference-time knob: Number of flow-matching denoising steps used to decode an action chunk. + # Trades inference latency for action quality. + # None keeps the checkpoint value (GR00T N1.7 default: 4). + num_inference_timesteps: int | None = None - # Whether to use the full model for LORA - lora_full_model: bool = False + # Inference-time knob: Real-Time Chunking (RTC) overlap-blend ramp rate, used when the RTC engine + # supplies a previous-chunk prefix. Higher values blend the overlapping prefix more aggressively. + # None keeps the checkpoint value (GR00T N1.7 default: 6.0). + rtc_ramp_rate: float | None = None - # Training parameters (matching groot_finetune_script.py) + # Inference-time knob: Whether to request the flash-attention-2 kernel for the Qwen3-VL backbone. + # flash-attn is an optional, user-managed optimization; when it is absent (the default), + # the backbone transparently falls back to SDPA, which is numerically equivalent. + # Set to True only after installing a flash-attn build matching your torch/CUDA env. + use_flash_attention: bool = False + + # Enable GR00T-style state-relative action chunks (action chunk expressed relative to the current + # observation state). + use_relative_actions: bool = False + + # relative_exclude_joints names the action dimensions that stay absolute; the + # match is substring/case-insensitive against the dataset action feature names. With the empty + # default every dimension is treated as relative, including the gripper -- set e.g. ["gripper"] to + # keep the gripper absolute, matching the Isaac-GR00T single-arm + absolute-gripper convention. + relative_exclude_joints: list[str] = field(default_factory=list) + + # Training parameters optimizer_lr: float = 1e-4 - optimizer_betas: tuple[float, float] = (0.95, 0.999) + # Isaac-GR00T N1.7 fine-tunes with AdamW betas (0.9, 0.999). + optimizer_betas: tuple[float, float] = (0.9, 0.999) optimizer_eps: float = 1e-8 optimizer_weight_decay: float = 1e-5 warmup_ratio: float = 0.05 use_bf16: bool = True + # The native N1.7 fine-tuning recipe keeps model parameters in FP32 and computes under BF16 autocast. + model_params_fp32: bool = True - # Dataset parameters - # Video backend to use for training ('decord' or 'torchvision_av') + # TODO(Steven): Remove these deprecated fields in a future release. + # Deprecated Isaac-GR00T runner / GR00T N1.5 fields, plus the (never-wired) LoRA fields — all + # unused by the LeRobot N1.7 implementation except the `tokenizer_assets_repo` N1.5 tripwire and + # the `image_size` legacy remap in __post_init__. They are kept ONLY so a config.json saved by an + # earlier lerobot release (notably a GR00T N1.5 checkpoint) still parses under draccus — which + # rejects unknown fields — and is then rejected with a clear N1.5 removal message rather than an + # opaque draccus decoding error. + image_size: tuple[int, int] = (256, 256) # image sizing is handled by the backbone's image processor. + tokenizer_assets_repo: str | None = None + lora_rank: int = 0 + lora_alpha: int = 16 + lora_dropout: float = 0.1 + lora_full_model: bool = False video_backend: str = "decord" - - # Whether to balance dataset weights in mixture datasets balance_dataset_weights: bool = True - - # Whether to sample trajectories weighted by their length balance_trajectory_weights: bool = True - - # Optional dataset paths for delegating training to Isaac-GR00T runner dataset_paths: list[str] | None = None output_dir: str = "./tmp/gr00t" save_steps: int = 1000 @@ -117,6 +372,65 @@ class GrootConfig(PreTrainedConfig): resume: bool = False def __post_init__(self): + if self.tokenizer_assets_repo is not None: + raise ValueError( + "Config sets 'tokenizer_assets_repo', which only existed for GR00T N1.5; this looks " + f"like a legacy GR00T N1.5 checkpoint or config. {GROOT_N1_5_REMOVAL_GUIDANCE}" + ) + + self.action_decode_transform = normalize_groot_action_decode_transform(self.action_decode_transform) + if self.base_model_path is None: + self.base_model_path = GROOT_N1_7_BASE_MODEL + + # The N1.7 LIBERO checkpoints emit a [0, 1] gripper action, but the LIBERO + # simulator expects the OpenVLA/[-1, 1] sign convention. NVIDIA's rollout + # wrapper applies this conversion; mirror it here so eval on the + # 'libero_sim' embodiment grasps correctly instead of scoring 0% success. + # This matches the embodiment-specific handling already done for the + # action execution horizon (see infer_groot_n1_7_action_execution_horizon). + # Only the 'auto' sentinel resolves to the embodiment default; an explicit + # 'none' (normalized to None above) keeps the transform disabled. + if self.action_decode_transform == GROOT_ACTION_DECODE_TRANSFORM_AUTO: + self.action_decode_transform = ( + GROOT_ACTION_DECODE_TRANSFORM_LIBERO if self.embodiment_tag == "libero_sim" else None + ) + + # GR00T N1.5-era default values (e.g. --policy.chunk_size=50 from old commands or + # stale configs) are migrated to the values the N1.7 checkpoints expect, with a + # warning. The dataclass defaults are already the N1.7 values, so a plain + # GrootConfig() never triggers this. + legacy_default_remaps = ( + ("max_state_dim", 64, 132), + ("max_action_dim", 32, 132), + ("chunk_size", 50, 40), + ("n_action_steps", 50, 40), + ("image_size", (224, 224), (256, 256)), + ) + for field_name, legacy_value, n1_7_value in legacy_default_remaps: + current_value = getattr(self, field_name) + if isinstance(legacy_value, tuple): + current_value = tuple(current_value) + if current_value == legacy_value: + logger.warning( + "GrootConfig.%s=%s matches a legacy GR00T N1.5-era default; remapping it to %s, " + "the value expected by GR00T N1.7 checkpoints. Set a different value explicitly " + "if this is not what you want.", + field_name, + legacy_value, + n1_7_value, + ) + setattr(self, field_name, n1_7_value) + + inferred_version = infer_groot_model_version(self.base_model_path) + if inferred_version is not None and inferred_version != GROOT_N1_7: + message = ( + f"GR00T model_version '{GROOT_N1_7}' does not match base_model_path " + f"'{self.base_model_path}', which looks like '{inferred_version}'." + ) + if inferred_version == GROOT_N1_5: + message = f"{message} {GROOT_N1_5_REMOVAL_GUIDANCE}" + raise ValueError(message) + super().__post_init__() if self.n_action_steps > self.chunk_size: @@ -124,9 +438,6 @@ class GrootConfig(PreTrainedConfig): f"n_action_steps ({self.n_action_steps}) cannot exceed chunk_size ({self.chunk_size})" ) - # groot_repo_path is now optional since we ported the components - # No validation needed - def validate_features(self) -> None: """Validate and set up input/output features for Groot.""" image_features = [key for key, feat in self.input_features.items() if feat.type == FeatureType.VISUAL] @@ -173,15 +484,20 @@ class GrootConfig(PreTrainedConfig): betas=self.optimizer_betas, eps=self.optimizer_eps, weight_decay=self.optimizer_weight_decay, + grad_clip_norm=1.0, ) - def get_scheduler_preset(self) -> CosineDecayWithWarmupSchedulerConfig: - """Return scheduler configuration.""" - return CosineDecayWithWarmupSchedulerConfig( - num_warmup_steps=int(10000 * self.warmup_ratio), # 5% warmup by default - num_decay_steps=10000, # Adjust based on training steps - peak_lr=self.optimizer_lr, - decay_lr=self.optimizer_lr * 0.1, + def get_scheduler_preset(self) -> DiffuserSchedulerConfig: + """Return scheduler configuration. + + Isaac-GR00T uses the HF Trainer cosine schedule with ~5% warmup over the + actual training update count; DiffuserSchedulerConfig wraps the same + diffusers/transformers `get_scheduler("cosine")` implementation and + derives num_training_steps from the outer --steps value at runtime. + """ + return DiffuserSchedulerConfig( + name="cosine", + num_warmup_steps=math.ceil(self.max_steps * self.warmup_ratio), ) @property @@ -192,7 +508,15 @@ class GrootConfig(PreTrainedConfig): @property def action_delta_indices(self) -> list[int]: """Return indices for delta actions.""" - return list(range(min(self.chunk_size, 16))) + model_action_horizon = ( + infer_groot_n1_7_action_horizon(self.base_model_path, self.embodiment_tag) or 40 + ) + return list(range(min(self.chunk_size, model_action_horizon))) + + @property + def drop_n_last_frames(self) -> int: + """Exclude episode tails that cannot supply a complete N1.7 action chunk.""" + return max(0, len(self.action_delta_indices) - 1) @property def reward_delta_indices(self) -> None: diff --git a/src/lerobot/policies/groot/eagle2_hg_model/configuration_eagle2_5_vl.py b/src/lerobot/policies/groot/eagle2_hg_model/configuration_eagle2_5_vl.py deleted file mode 100755 index 526b4f7a2..000000000 --- a/src/lerobot/policies/groot/eagle2_hg_model/configuration_eagle2_5_vl.py +++ /dev/null @@ -1,135 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# 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. - - -import copy - -from transformers.configuration_utils import PretrainedConfig -from transformers.models.llama.configuration_llama import LlamaConfig -from transformers.models.qwen2.configuration_qwen2 import Qwen2Config -from transformers.models.qwen3.configuration_qwen3 import Qwen3Config -from transformers.models.siglip.configuration_siglip import SiglipVisionConfig -from transformers.utils import logging - -logger = logging.get_logger(__name__) - - -class Eagle25VLConfig(PretrainedConfig): - model_type = "eagle_2_5_vl" - is_composition = True - sub_configs = {"vision_config": SiglipVisionConfig, "text_config": Qwen2Config} - - def __init__( - self, - vision_config=None, - text_config=None, - use_backbone_lora=0, - use_llm_lora=0, - pad2square=False, - select_layer=-4, - force_image_size=None, - downsample_ratio=0.5, - template=None, - dynamic_image_size=False, - use_thumbnail=False, - loss_version="v1", - min_dynamic_tiles=1, - max_dynamic_tiles=6, - mlp_checkpoint=False, - initializer_range=0.02, - _attn_implementation="flash_attention_2", - _attn_implementation_autoset=False, - llm_config=None, - image_token_index=None, - use_pixel_shuffle=True, - mlp_connector_layers=2, - **kwargs, - ): - super().__init__(**kwargs) - - if vision_config is None: - vision_config = {"model_type": "siglip_vision_model"} - logger.info("vision_config is None. Initializing the InternVisionConfig with default values.") - - if text_config is None: - text_config = {"architectures": ["Qwen2ForCausalLM"]} - logger.info( - "text_config is None. Initializing the LlamaConfig config with default values (`LlamaConfig`)." - ) - - if vision_config["model_type"] == "siglip_vision_model": - self.vision_config = SiglipVisionConfig(**vision_config) - else: - raise ValueError("Unsupported model_type: {}".format(vision_config["model_type"])) - - if text_config["architectures"][0] == "LlamaForCausalLM": - self.text_config = LlamaConfig(**text_config) - elif text_config["architectures"][0] == "Qwen2ForCausalLM": - self.text_config = Qwen2Config(**text_config) - elif text_config["architectures"][0] == "Qwen3ForCausalLM": - self.text_config = Qwen3Config(**text_config) - else: - raise ValueError("Unsupported architecture: {}".format(text_config["architectures"][0])) - self.use_backbone_lora = use_backbone_lora - self.use_llm_lora = use_llm_lora - self.mlp_checkpoint = mlp_checkpoint - self.pad2square = pad2square - self.select_layer = select_layer - self.force_image_size = force_image_size - self.downsample_ratio = downsample_ratio - self.template = template - self.dynamic_image_size = dynamic_image_size - self.use_thumbnail = use_thumbnail - self.loss_version = loss_version - self.initializer_range = initializer_range - self.min_dynamic_tiles = min_dynamic_tiles - self.max_dynamic_tiles = max_dynamic_tiles - self.tie_word_embeddings = self.text_config.tie_word_embeddings - self._attn_implementation = _attn_implementation - self._attn_implementation_autoset = _attn_implementation_autoset - self.image_token_index = image_token_index - self.use_pixel_shuffle = use_pixel_shuffle - self.mlp_connector_layers = mlp_connector_layers - logger.info(f"min_dynamic_tiles: {self.min_dynamic_tiles}") - logger.info(f"max_dynamic_tiles: {self.max_dynamic_tiles}") - - def to_dict(self): - """ - Serializes this instance to a Python dictionary. Override the default [`~PretrainedConfig.to_dict`]. - - Returns: - `Dict[str, any]`: Dictionary of all the attributes that make up this configuration instance, - """ - output = copy.deepcopy(self.__dict__) - output["vision_config"] = self.vision_config.to_dict() - output["text_config"] = self.text_config.to_dict() - output["model_type"] = self.__class__.model_type - output["use_backbone_lora"] = self.use_backbone_lora - output["use_llm_lora"] = self.use_llm_lora - output["pad2square"] = self.pad2square - output["select_layer"] = self.select_layer - output["force_image_size"] = self.force_image_size - output["downsample_ratio"] = self.downsample_ratio - output["template"] = self.template - output["dynamic_image_size"] = self.dynamic_image_size - output["use_thumbnail"] = self.use_thumbnail - output["min_dynamic_tiles"] = self.min_dynamic_tiles - output["max_dynamic_tiles"] = self.max_dynamic_tiles - output["tie_word_embeddings"] = self.tie_word_embeddings - output["_attn_implementation"] = self._attn_implementation - output["_attn_implementation_autoset"] = self._attn_implementation_autoset - output["use_pixel_shuffle"] = self.use_pixel_shuffle - output["mlp_connector_layers"] = self.mlp_connector_layers - return output diff --git a/src/lerobot/policies/groot/eagle2_hg_model/image_processing_eagle2_5_vl_fast.py b/src/lerobot/policies/groot/eagle2_hg_model/image_processing_eagle2_5_vl_fast.py deleted file mode 100644 index 90e9dcecc..000000000 --- a/src/lerobot/policies/groot/eagle2_hg_model/image_processing_eagle2_5_vl_fast.py +++ /dev/null @@ -1,503 +0,0 @@ -# -------------------------------------------------------- -# NVIDIA -# Copyright (c) 2025 NVIDIA -# Licensed under The MIT License [see LICENSE for details] -# -------------------------------------------------------- - -from __future__ import annotations - -# copy from https://github.com/huggingface/transformers/blob/main/src/transformers/models/llava_onevision/image_processing_llava_onevision_fast.py -from transformers.image_processing_utils import ( - BatchFeature, - get_patch_output_size, -) -from transformers.image_processing_utils_fast import ( - BaseImageProcessorFast, - ImagesKwargs, - group_images_by_shape, - reorder_images, -) -from transformers.image_utils import ( - IMAGENET_STANDARD_MEAN, # 0.5, 0.5, 0.5 - IMAGENET_STANDARD_STD, # 0.5, 0.5, 0.5 - ChannelDimension, - ImageInput, - PILImageResampling, - SizeDict, - get_image_size, - make_flat_list_of_images, - validate_kwargs, -) -from transformers.processing_utils import Unpack -from transformers.utils import ( - TensorType, - add_start_docstrings, - is_torch_available, - is_torchvision_v2_available, -) -from transformers.video_utils import VideoInput - -if is_torch_available(): - import torch -if is_torchvision_v2_available(): - from torchvision.transforms.v2 import functional as F # noqa: N812 - from transformers.image_utils import pil_torch_interpolation_mapping -else: - from torchvision.transforms import functional as F # noqa: N812 - - -def crop(img: torch.Tensor, left: int, top: int, right: int, bottom: int) -> torch.Tensor: - """Crop the given numpy array. - - Args: - img (torch.Tensor): Image to be cropped. Format should be (C, H, W). - left (int): The left coordinate of the crop box. - top (int): The top coordinate of the crop box. - right (int): The right coordinate of the crop box. - bottom (int): The bottom coordinate of the crop box. - - Returns: - torch.Tensor: Cropped image. - """ - if not isinstance(img, torch.Tensor): - raise TypeError(f"img should be torch.Tensor. Got {type(img)}") - - if img.ndim not in [2, 3]: - raise ValueError(f"Image should have 2 or 3 dimensions. Got {img.ndim}") - - img_height = img.shape[1] - img_width = img.shape[2] - if top < 0 or left < 0 or bottom > img_height or right > img_width: - raise ValueError("Crop coordinates out of bounds") - - if top >= bottom or left >= right: - raise ValueError("Invalid crop coordinates") - - return img[:, top:bottom, left:right] - - -class Eagle25VLFastImageProcessorKwargs(ImagesKwargs): - max_dynamic_tiles: int | None - min_dynamic_tiles: int | None - use_thumbnail: bool | None - pad_during_tiling: bool | None - do_pad: bool | None - - -@add_start_docstrings( - "Constructs a fast ConvNeXT image processor. Based on [`SiglipImageProcessor`] with incorporation of processing each video frame.", - # BASE_IMAGE_PROCESSOR_FAST_DOCSTRING, TODO: this was depreciated from transformers remove! - """ - image_grid_pinpoints (`List[List[int]]`, *optional*): - A list of possible resolutions to use for processing high resolution images. The best resolution is selected - based on the original size of the image. Can be overridden by `image_grid_pinpoints` in the `preprocess` - method. Not used for processing videos. - do_pad (`bool`, *optional*): - Whether to pad the image. If `True`, will pad the patch dimension of the images in the batch to the largest - number of patches in the batch. Padding will be applied to the bottom and right with zeros. - """, -) -class Eagle25VLImageProcessorFast(BaseImageProcessorFast): - resample = PILImageResampling.BICUBIC - image_mean = IMAGENET_STANDARD_MEAN - image_std = IMAGENET_STANDARD_STD - size = {"height": 448, "width": 448} - default_to_square = False - crop_size = None - do_resize = True - do_center_crop = None - do_rescale = True - do_normalize = True - do_convert_rgb = True - do_pad = True - max_dynamic_tiles = 12 - min_dynamic_tiles = 1 - use_thumbnail = True - pad_during_tiling = False - valid_kwargs = Eagle25VLFastImageProcessorKwargs - model_input_names = ["pixel_values_videos"] - - def __init__(self, **kwargs: Unpack[Eagle25VLFastImageProcessorKwargs]): - super().__init__(**kwargs) - - @add_start_docstrings( - # BASE_IMAGE_PROCESSOR_FAST_DOCSTRING_PREPROCESS, TODO: this was depreciated from transformers remove! - """ - max_dynamic_tiles (`int`, *optional*): - The maximum number of dynamic tiles to use for processing high resolution images. - min_dynamic_tiles (`int`, *optional*): - The minimum number of dynamic tiles to use for processing high resolution images. - use_thumbnail (`bool`, *optional*): - Whether to use a thumbnail for processing high resolution images. - pad_during_tiling (`bool`, *optional*): - Whether to pad the image during tiling. - do_pad (`bool`, *optional*): - Whether to pad the image. If `True`, will pad the patch dimension of the images in the batch to the largest - number of patches in the batch. Padding will be applied to the bottom and right with zeros. - """, - ) - - # NOTE(YL): we will overload the preprocess method to add the image_flags - # def preprocess( - # self, images: ImageInput, **kwargs: Unpack[Eagle25VLFastImageProcessorKwargs] - # ) -> BatchFeature: - # return super().preprocess(images, **kwargs) - - def _prepare_images_structure( - self, - images: ImageInput, - expected_ndims: int = 3, - ) -> ImageInput: - """ - Prepare the images structure for processing. - - Args: - images (`ImageInput`): - The input images to process. - expected_ndims (`int`, *optional*, defaults to 3): - Expected number of dimensions for the images (added for transformers >=4.53.0 compatibility). - - Returns: - `ImageInput`: The images with a valid nesting. - """ - return make_flat_list_of_images(images) - - def _resize_for_patching( - self, - image: torch.Tensor, - target_resolution: tuple, - interpolation: F.InterpolationMode, - input_data_format: ChannelDimension, - ) -> torch.Tensor: - """ - Resizes an image to a target resolution while maintaining aspect ratio. - - Args: - image ("torch.Tensor"): - The input image. - target_resolution (tuple): - The target resolution (height, width) of the image. - interpolation (`InterpolationMode`): - Resampling filter to use if resizing the image. - input_data_format (`ChannelDimension` or `str`): - The channel dimension format of the input image. - - Returns: - "torch.Tensor": The resized and padded image. - """ - new_height, new_width = get_patch_output_size(image, target_resolution, input_data_format) - - # Resize the image - resized_image = F.resize(image, (new_height, new_width), interpolation=interpolation) - - return resized_image - - def find_closest_aspect_ratio(self, aspect_ratio, target_ratios, width, height, image_size): - """ - previous version mainly focus on ratio. - We also consider area ratio here. - """ - best_factor = float("-inf") - best_ratio = (1, 1) - area = width * height - for ratio in target_ratios: - target_aspect_ratio = ratio[0] / ratio[1] - # ratio_diff = abs(aspect_ratio - target_aspect_ratio) - # area_ratio = (ratio[0] * ratio[1] * image_size * image_size) / area - """ - new area > 60% of original image area is enough. - """ - factor_based_on_area_n_ratio = min( - (ratio[0] * ratio[1] * image_size * image_size) / area, 0.6 - ) * min(target_aspect_ratio / aspect_ratio, aspect_ratio / target_aspect_ratio) - - if factor_based_on_area_n_ratio > best_factor: - best_factor = factor_based_on_area_n_ratio - best_ratio = ratio - - return best_ratio - - def _pad_for_patching( - self, image: torch.Tensor, target_resolution: tuple, input_data_format: ChannelDimension - ) -> torch.Tensor: - """ - Pad an image to a target resolution while maintaining aspect ratio. - """ - target_height, target_width = target_resolution - new_height, new_width = get_patch_output_size(image, target_resolution, input_data_format) - - paste_x = (target_width - new_width) // 2 - paste_y = (target_height - new_height) // 2 - - padded_image = F.pad(image, padding=[paste_x, paste_y, paste_x, paste_y]) - - return padded_image - - def _get_image_patches( - self, - image: torch.Tensor, - min_num: int, - max_num: int, - size: tuple, - tile_size: int, - use_thumbnail: bool, - interpolation: F.InterpolationMode, - pad_during_tiling: bool, - ) -> list[torch.Tensor]: - image_size = get_image_size(image, channel_dim=ChannelDimension.FIRST) - orig_height, orig_width = image_size - aspect_ratio = orig_width / orig_height - - # calculate the existing image aspect ratio - target_ratios = { - (i, j) - for n in range(min_num, max_num + 1) - for i in range(1, n + 1) - for j in range(1, n + 1) - if i * j <= max_num and i * j >= min_num - } - target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1]) - - # find the closest aspect ratio to the target - target_aspect_ratio = self.find_closest_aspect_ratio( - aspect_ratio, target_ratios, orig_width, orig_height, tile_size - ) - - # calculate the target width and height - target_width = tile_size * target_aspect_ratio[0] - target_height = tile_size * target_aspect_ratio[1] - blocks = target_aspect_ratio[0] * target_aspect_ratio[1] - if pad_during_tiling: - resized_image = self._resize_for_patching( - image, - (target_height, target_width), - interpolation=interpolation, - input_data_format=ChannelDimension.FIRST, - ) - padded_image = self._pad_for_patching( - resized_image, - (target_height, target_width), - input_data_format=ChannelDimension.FIRST, - ) - image_used_to_split = padded_image - else: - image_used_to_split = F.resize(image, (target_height, target_width), interpolation=interpolation) - - processed_tiles = [] - for i in range(blocks): - box = ( - (i % (target_width // tile_size)) * tile_size, - (i // (target_width // tile_size)) * tile_size, - ((i % (target_width // tile_size)) + 1) * tile_size, - ((i // (target_width // tile_size)) + 1) * tile_size, - ) - # split the image - split_img = crop(image_used_to_split, box[0], box[1], box[2], box[3]) - processed_tiles.append(split_img) - assert len(processed_tiles) == blocks - - if use_thumbnail and len(processed_tiles) != 1: - thumbnail_img = F.resize(image, (tile_size, tile_size), interpolation=interpolation) - processed_tiles.append(thumbnail_img) - - return processed_tiles - - def _pad_for_batching( - self, - pixel_values: list[torch.Tensor], - ) -> list[torch.Tensor]: - """ - Pads images on the `num_of_patches` dimension with zeros to form a batch of same number of patches. - - Args: - pixel_values (`List[torch.Tensor]`): - An array of pixel values of each images of shape (`batch_size`, `num_patches`, `image_in_3D`) - - Returns: - List[`torch.Tensor`]: The padded images. - """ - max_patch = max(len(x) for x in pixel_values) - pixel_values = [ - torch.nn.functional.pad(image, pad=[0, 0, 0, 0, 0, 0, 0, max_patch - image.shape[0]]) - for image in pixel_values - ] - - return pixel_values - - def _preprocess( - self, - images: list[torch.Tensor], - do_resize: bool, - size: SizeDict, - max_dynamic_tiles: int, - min_dynamic_tiles: int, - use_thumbnail: bool, - pad_during_tiling: bool, - interpolation: F.InterpolationMode | None, - do_center_crop: bool, - crop_size: SizeDict, - do_rescale: bool, - rescale_factor: float, - do_normalize: bool, - image_mean: float | list[float] | None, - image_std: float | list[float] | None, - do_pad: bool, - return_tensors: str | TensorType | None, - pad_size: SizeDict | None = None, # Added for transformers >=4.53.0 compatibility - disable_grouping: bool | None = None, # Added for transformers >=4.53.0 compatibility - ) -> BatchFeature: - processed_images = [] - image_sizes = [] - # Determine the size tuple - if size and size.height and size.width: - size_tuple = (size.height, size.width) - else: - size_tuple = (size.shortest_edge, size.shortest_edge) - - # Determine the patch size - if crop_size and crop_size.height: - tile_size = crop_size.height - elif size and size.height: - tile_size = size.height - else: - tile_size = size.shortest_edge - - for image in images: - image_patches = self._get_image_patches( - image, - min_num=min_dynamic_tiles, - max_num=max_dynamic_tiles, - size=size_tuple, - tile_size=tile_size, - use_thumbnail=use_thumbnail, - interpolation=interpolation, - pad_during_tiling=pad_during_tiling, - ) - - # Group images by size for batched processing - processed_image_patches_grouped = {} - # Added for transformers >=4.53.0 compatibility - grouped_image_patches, grouped_image_patches_index = group_images_by_shape( - image_patches, - disable_grouping=disable_grouping, - ) - - for shape, stacked_image_patches in grouped_image_patches.items(): - if do_resize: - stacked_image_patches = self.resize( - image=stacked_image_patches, - size=size, - interpolation=interpolation, - ) - if do_center_crop: - stacked_image_patches = self.center_crop(stacked_image_patches, crop_size) - # Fused rescale and normalize - stacked_image_patches = self.rescale_and_normalize( - stacked_image_patches, - do_rescale, - rescale_factor, - do_normalize, - image_mean, - image_std, - ) - processed_image_patches_grouped[shape] = stacked_image_patches - processed_image_patches = reorder_images( - processed_image_patches_grouped, grouped_image_patches_index - ) - processed_image_patches = ( - torch.stack(processed_image_patches, dim=0) if return_tensors else processed_image_patches - ) - processed_images.append(processed_image_patches) - image_sizes.append(get_image_size(image, ChannelDimension.FIRST)) - - if do_pad: - processed_images = self._pad_for_batching(processed_images) - - # processed_images = torch.stack(processed_images, dim=0) if return_tensors else processed_images - processed_images = torch.cat(processed_images, dim=0) if return_tensors else processed_images - return BatchFeature( - data={"pixel_values": processed_images, "image_sizes": image_sizes}, - tensor_type=return_tensors, - ) - - def preprocess( - self, - images: ImageInput, - videos: VideoInput = None, - **kwargs: Unpack[Eagle25VLFastImageProcessorKwargs], - ) -> BatchFeature: - validate_kwargs( - captured_kwargs=kwargs.keys(), - valid_processor_keys=self.valid_kwargs.__annotations__.keys(), - ) - # Set default kwargs from self. This ensures that if a kwarg is not provided - # by the user, it gets its default value from the instance, or is set to None. - for kwarg_name in self.valid_kwargs.__annotations__: - kwargs.setdefault(kwarg_name, getattr(self, kwarg_name, None)) - - # Extract parameters that are only used for preparing the input images - do_convert_rgb = kwargs.pop("do_convert_rgb") - input_data_format = kwargs.pop("input_data_format") - device = kwargs.pop("device") - # Prepare input images - # transformers >= 4.53.0: uses _prepare_image_like_inputs instead of _prepare_input_images - if images is not None: - images = self._prepare_image_like_inputs( - images=images, - do_convert_rgb=do_convert_rgb, - input_data_format=input_data_format, - device=device, - ) - - if videos is not None: - videos = self._prepare_image_like_inputs( - images=videos, - do_convert_rgb=do_convert_rgb, - input_data_format=input_data_format, - device=device, - ) - - # Update kwargs that need further processing before being validated - kwargs = self._further_process_kwargs(**kwargs) - - # Validate kwargs - self._validate_preprocess_kwargs(**kwargs) - - # torch resize uses interpolation instead of resample - # Added for transformers >=4.53.0 compatibility - resample = kwargs.pop("resample", self.resample) - kwargs["interpolation"] = ( - pil_torch_interpolation_mapping[resample] - if isinstance(resample, PILImageResampling | int) - else resample - ) - - # Filter kwargs to only include those accepted by _preprocess - valid_preprocess_kwargs = { - "do_resize", - "size", - "max_dynamic_tiles", - "min_dynamic_tiles", - "use_thumbnail", - "pad_during_tiling", - "interpolation", - "do_center_crop", - "crop_size", - "do_rescale", - "rescale_factor", - "do_normalize", - "image_mean", - "image_std", - "do_pad", - "return_tensors", - "pad_size", - "disable_grouping", - } - filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_preprocess_kwargs} - if images is not None: - return self._preprocess(images, **filtered_kwargs) - elif videos is not None: - return self._preprocess(videos, **filtered_kwargs) - - -__all__ = ["Eagle25VLImageProcessorFast"] diff --git a/src/lerobot/policies/groot/eagle2_hg_model/modeling_eagle2_5_vl.py b/src/lerobot/policies/groot/eagle2_hg_model/modeling_eagle2_5_vl.py deleted file mode 100755 index 6e5532ea4..000000000 --- a/src/lerobot/policies/groot/eagle2_hg_model/modeling_eagle2_5_vl.py +++ /dev/null @@ -1,396 +0,0 @@ -# -------------------------------------------------------- -# NVIDIA -# Copyright (c) 2025 NVIDIA -# Licensed under The MIT License [see LICENSE for details] -# -------------------------------------------------------- - -import inspect - -import torch -import torch.utils.checkpoint as cp -from peft import LoraConfig, get_peft_model -from torch import nn -from torch.nn import CrossEntropyLoss -from transformers import GenerationConfig -from transformers.generation import GenerationMixin -from transformers.modeling_outputs import CausalLMOutputWithPast -from transformers.modeling_utils import PreTrainedModel -from transformers.models.llama.modeling_llama import LlamaForCausalLM -from transformers.models.qwen2.modeling_qwen2 import Qwen2ForCausalLM -from transformers.models.qwen3.modeling_qwen3 import Qwen3ForCausalLM -from transformers.models.siglip.modeling_siglip import SiglipVisionModel -from transformers.utils import add_start_docstrings, logging - -from .configuration_eagle2_5_vl import Eagle25VLConfig - -logger = logging.get_logger(__name__) - - -# copy from https://github.com/huggingface/transformers/blob/main/src/transformers/models/llava_onevision/modeling_llava_onevision.py#L241C1-L280C1 -EAGLE2_5_VL_START_DOCSTRING = r""" - This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the - library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads - etc.) - - This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. - Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage - and behavior. - - Parameters: - config ([`Eagle25VLConfig`]): - Model configuration class with all the parameters of the model. Initializing with a config file does not - load the weights associated with the model, only the configuration. Check out the - [`~PreTrainedModel.from_pretrained`] method to load the model weights. -""" - - -@add_start_docstrings( - "The bare Eagle2_5_VL Model outputting raw hidden-states without any specific head on top.", - EAGLE2_5_VL_START_DOCSTRING, -) -class Eagle25VLPreTrainedModel(PreTrainedModel): - config_class = Eagle25VLConfig - base_model_prefix = "model" - main_input_name = "input_ids" - supports_gradient_checkpointing = True - _no_split_modules = [ - "Qwen2DecoderLayer", - "LlamaDecoderLayer", - "Siglip2EncoderLayer", - "SiglipEncoderLayer", - ] - _skip_keys_device_placement = "past_key_values" - _supports_flash_attn = True - _supports_flash_attn_2 = True - _supports_cache_class = True - _supports_static_cache = True - _supports_quantized_cache = True - _supports_sdpa = True - - def _init_weights(self, module): - std = self.config.initializer_range - if isinstance(module, nn.Linear | nn.Conv2d): - module.weight.data.normal_(mean=0.0, std=std) - if module.bias is not None: - module.bias.data.zero_() - elif isinstance(module, nn.Embedding): - module.weight.data.normal_(mean=0.0, std=std) - if module.padding_idx is not None: - module.weight.data[module.padding_idx].zero_() - - -class Eagle25VLForConditionalGeneration(Eagle25VLPreTrainedModel, GenerationMixin): - config_class = Eagle25VLConfig - - def __init__(self, config: Eagle25VLConfig, vision_model=None, language_model=None): - super().__init__(config) - - image_size = config.force_image_size or config.vision_config.image_size - patch_size = config.vision_config.patch_size - self.patch_size = patch_size - if config.use_pixel_shuffle: - self.num_image_token = int((image_size // patch_size) ** 2 * (config.downsample_ratio**2)) - else: - self.num_image_token = int((image_size // patch_size) ** 2) - - self.select_layer = config.select_layer - self.downsample_ratio = config.downsample_ratio - self.loss_version = config.loss_version - self.mlp_checkpoint = config.mlp_checkpoint - self.use_pixel_shuffle = config.use_pixel_shuffle - self.mlp_connector_layers = config.mlp_connector_layers - logger.info(f"num_image_token: {self.num_image_token}") - logger.info(f"mlp_checkpoint: {self.mlp_checkpoint}") - if vision_model is not None: - self.vision_model = vision_model - else: - if config.vision_config.model_type == "siglip_vision_model": - config.vision_config._attn_implementation = "flash_attention_2" - self.vision_model = SiglipVisionModel(config.vision_config) - else: - raise NotImplementedError(f"{config.vision_config.model_type} is not implemented.") - - if language_model is not None: - self.language_model = language_model - else: - if config.text_config.architectures[0] == "LlamaForCausalLM": - self.language_model = LlamaForCausalLM(config.text_config) - elif config.text_config.architectures[0] == "Phi3ForCausalLM": - raise NotImplementedError("Phi3 is not implemented.") - # self.language_model = Phi3ForCausalLM(config.text_config) - elif config.text_config.architectures[0] == "Qwen2ForCausalLM": - assert config.text_config._attn_implementation == "flash_attention_2", ( - f"Qwen2 must use flash_attention_2 but got {config.text_config._attn_implementation}" - ) - self.language_model = Qwen2ForCausalLM(config.text_config) - elif config.text_config.architectures[0] == "Qwen3ForCausalLM": - self.language_model = Qwen3ForCausalLM(config.text_config) - else: - raise NotImplementedError(f"{config.text_config.architectures[0]} is not implemented.") - - vit_hidden_size = config.vision_config.hidden_size - llm_hidden_size = config.text_config.hidden_size - - if config.mlp_connector_layers == 2: - self.mlp1 = nn.Sequential( - nn.LayerNorm(vit_hidden_size * int(1 / self.downsample_ratio) ** 2), - nn.Linear(vit_hidden_size * int(1 / self.downsample_ratio) ** 2, llm_hidden_size), - nn.GELU(), - nn.Linear(llm_hidden_size, llm_hidden_size), - ) - elif config.mlp_connector_layers == 1 and config.use_pixel_shuffle: - self.mlp1 = nn.Sequential( - nn.Linear(vit_hidden_size * int(1 / self.downsample_ratio) ** 2, llm_hidden_size), - ) - elif config.mlp_connector_layers == 1 and not config.use_pixel_shuffle: - self.mlp1 = nn.Sequential( - nn.Linear(vit_hidden_size, llm_hidden_size), - ) - else: - raise NotImplementedError(f"{config.mlp_connector_layers} is not implemented.") - - self.image_token_index = config.image_token_index - self.neftune_alpha = None - - if config.use_backbone_lora: - self.wrap_backbone_lora(r=config.use_backbone_lora, lora_alpha=2 * config.use_backbone_lora) - - self.use_llm_lora = config.use_llm_lora - if config.use_llm_lora: - self.wrap_llm_lora(r=config.use_llm_lora, lora_alpha=2 * config.use_llm_lora) - - self.check_forward_kwargs() - - def check_forward_kwargs(self): - # We intentionally avoid using **kwargs in forward because Hugging Face Transformers - # has special handling for functions with **kwargs parameters that would affect - # how our model is processed during training and inference. - forward_params = inspect.signature(self.forward).parameters - assert not any(k.kind == inspect.Parameter.VAR_KEYWORD for k in forward_params.values()) - - def wrap_backbone_lora(self, r=128, lora_alpha=256, lora_dropout=0.05): - lora_config = LoraConfig( - r=r, - target_modules=[ - "self_attn.q_proj", - "self_attn.k_proj", - "self_attn.v_proj", - "self_attn.out_proj", - "mlp.fc1", - "mlp.fc2", - ], - lora_alpha=lora_alpha, - lora_dropout=lora_dropout, - ) - self.vision_model = get_peft_model(self.vision_model, lora_config) - self.vision_model.print_trainable_parameters() - - def wrap_llm_lora(self, r=128, lora_alpha=256, lora_dropout=0.05): - lora_config = LoraConfig( - r=r, - target_modules=[ - "self_attn.q_proj", - "self_attn.k_proj", - "self_attn.v_proj", - "self_attn.o_proj", - "mlp.gate_proj", - "mlp.down_proj", - "mlp.up_proj", - ], - lora_alpha=lora_alpha, - lora_dropout=lora_dropout, - task_type="CAUSAL_LM", - ) - self.language_model = get_peft_model(self.language_model, lora_config) - self.language_model.enable_input_require_grads() - self.language_model.print_trainable_parameters() - self.use_llm_lora = True - - def forward( - self, - pixel_values: torch.FloatTensor, - input_ids: torch.LongTensor = None, - attention_mask: torch.Tensor | None = None, - position_ids: torch.LongTensor | None = None, - image_flags: torch.LongTensor | None = None, - past_key_values: list[torch.FloatTensor] | None = None, - labels: torch.LongTensor | None = None, - use_cache: bool | None = None, - output_attentions: bool | None = None, - output_hidden_states: bool | None = None, - return_dict: bool | None = None, - num_tiles_list: list[torch.Tensor] | None = None, - ) -> tuple | CausalLMOutputWithPast: - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - input_embeds = self.language_model.get_input_embeddings()(input_ids) - - vit_embeds = self.extract_feature(pixel_values) - - if image_flags is not None: - image_flags = image_flags.view(-1) - vit_embeds = vit_embeds[image_flags == 1] - - b, n, c = input_embeds.shape - input_embeds = input_embeds.reshape(b * n, c) - - input_ids = input_ids.reshape(b * n) - selected = input_ids == self.image_token_index - try: - input_embeds[selected] = input_embeds[selected] * 0.0 + vit_embeds.reshape(-1, c) - except Exception as e: - vit_embeds = vit_embeds.reshape(-1, c) - print( - f"warning: {e}, input_embeds[selected].shape={input_embeds[selected].shape}, " - f"vit_embeds.shape={vit_embeds.shape}" - ) - n_token = selected.sum() - input_embeds[selected] = input_embeds[selected] * 0.0 + vit_embeds[:n_token] - - input_embeds = input_embeds.reshape(b, n, c) - - outputs = self.language_model( - inputs_embeds=input_embeds, - attention_mask=attention_mask, - position_ids=position_ids, - past_key_values=past_key_values, - use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - ) - logits = outputs.logits - - loss = None - if labels is not None: - # Shift so that tokens < n predict n - shift_logits = logits[..., :-1, :].contiguous() - shift_labels = labels[..., 1:].contiguous() - # Flatten the tokens - loss_fct = CrossEntropyLoss() - shift_logits = shift_logits.view(-1, self.language_model.config.vocab_size) - shift_labels = shift_labels.view(-1) - # Enable model parallelism - shift_labels = shift_labels.to(shift_logits.device) - loss = loss_fct(shift_logits, shift_labels) - - if not return_dict: - output = (logits,) + outputs[1:] - return (loss,) + output if loss is not None else output - - return CausalLMOutputWithPast( - loss=loss, - logits=logits, - past_key_values=outputs.past_key_values, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - ) - - def pixel_shuffle(self, x, scale_factor=0.5): - n, w, h, c = x.size() - # N, W, H, C --> N, W, H * scale, C // scale - x = x.view(n, w, int(h * scale_factor), int(c / scale_factor)) - # N, W, H * scale, C // scale --> N, H * scale, W, C // scale - x = x.permute(0, 2, 1, 3).contiguous() - # N, H * scale, W, C // scale --> N, H * scale, W * scale, C // (scale ** 2) - x = x.view(n, int(h * scale_factor), int(w * scale_factor), int(c / (scale_factor * scale_factor))) - - x = x.permute(0, 2, 1, 3).contiguous() - return x - - def extract_feature(self, pixel_values): - if self.select_layer == -1: - vit_embeds = self.vision_model( - pixel_values=pixel_values, output_hidden_states=False, return_dict=True - ) - if hasattr(vit_embeds, "last_hidden_state"): - vit_embeds = vit_embeds.last_hidden_state - - else: - vit_embeds = self.vision_model( - pixel_values=pixel_values, output_hidden_states=True, return_dict=True - ).hidden_states[self.select_layer] - - if self.use_pixel_shuffle: - h = w = int(vit_embeds.shape[1] ** 0.5) - vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], h, w, -1) - vit_embeds = self.pixel_shuffle( - vit_embeds, scale_factor=self.downsample_ratio - ) # torch.Size([B, 1024, 1024]) -> torch.Size([B, 16, 16, 4096]) - vit_embeds = vit_embeds.reshape( - vit_embeds.shape[0], -1, vit_embeds.shape[-1] - ) # torch.Size([B, 16, 16, 4096]) -> torch.Size([B, 256, 4096]) - - if self.mlp_checkpoint and vit_embeds.requires_grad: - vit_embeds = cp.checkpoint(self.mlp1, vit_embeds) - else: - vit_embeds = self.mlp1(vit_embeds) - - return vit_embeds - - @torch.no_grad() - def generate( - self, - pixel_values: torch.FloatTensor | None = None, - input_ids: torch.FloatTensor | None = None, - attention_mask: torch.LongTensor | None = None, - visual_features: torch.FloatTensor | None = None, - generation_config: GenerationConfig | None = None, - output_hidden_states: bool | None = None, - image_sizes: list[tuple[int, int]] | None = None, - **generate_kwargs, - ) -> torch.LongTensor: - if pixel_values is not None: - if visual_features is not None: - vit_embeds = visual_features - else: - vit_embeds = self.extract_feature(pixel_values) - - input_embeds = self.language_model.get_input_embeddings()(input_ids) - b, n, c = input_embeds.shape - input_embeds = input_embeds.reshape(b * n, c) - - input_ids = input_ids.reshape(b * n) - selected = input_ids == self.config.image_token_index - assert selected.sum() != 0 - input_embeds[selected] = vit_embeds.reshape(-1, c).to(input_embeds.device) - - input_embeds = input_embeds.reshape(b, n, c) - else: - input_embeds = self.language_model.get_input_embeddings()(input_ids) - - if "use_cache" not in generate_kwargs: - generate_kwargs["use_cache"] = True - - outputs = self.language_model.generate( - inputs_embeds=input_embeds, - attention_mask=attention_mask, - generation_config=generation_config, - output_hidden_states=output_hidden_states, - **generate_kwargs, - ) - - return outputs - - # Copied from transformers.models.llava_next.modeling_llava_next.LlavaNextForConditionalGeneration.get_input_embeddings - def get_input_embeddings(self): - return self.language_model.get_input_embeddings() - - # Copied from transformers.models.llava_next.modeling_llava_next.LlavaNextForConditionalGeneration.set_input_embeddings - def set_input_embeddings(self, value): - self.language_model.set_input_embeddings(value) - - # Copied from transformers.models.llava_next.modeling_llava_next.LlavaNextForConditionalGeneration.get_output_embeddings - def get_output_embeddings(self): - return self.language_model.get_output_embeddings() - - # Copied from transformers.models.llava_next.modeling_llava_next.LlavaNextForConditionalGeneration.set_output_embeddings - def set_output_embeddings(self, new_embeddings): - self.language_model.set_output_embeddings(new_embeddings) - - # Copied from transformers.models.llava_next.modeling_llava_next.LlavaNextForConditionalGeneration.set_decoder - def set_decoder(self, decoder): - self.language_model.set_decoder(decoder) - - # Copied from transformers.models.llava_next.modeling_llava_next.LlavaNextForConditionalGeneration.get_decoder - def get_decoder(self): - return self.language_model.get_decoder() diff --git a/src/lerobot/policies/groot/eagle2_hg_model/processing_eagle2_5_vl.py b/src/lerobot/policies/groot/eagle2_hg_model/processing_eagle2_5_vl.py deleted file mode 100755 index b36e70c47..000000000 --- a/src/lerobot/policies/groot/eagle2_hg_model/processing_eagle2_5_vl.py +++ /dev/null @@ -1,541 +0,0 @@ -# Copyright 2024 The HuggingFace Inc. team. -# -# 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. -""" -Processor class for Eagle25VL. -copy from https://github.com/huggingface/transformers/blob/main/src/transformers/models/llava_onevision/processing_llava_onevision.py -""" - -import base64 -import os -import re -from io import BytesIO - -import requests -import torch -from PIL import Image -from transformers.feature_extraction_utils import BatchFeature -from transformers.image_utils import ImageInput -from transformers.processing_utils import ProcessingKwargs, ProcessorMixin, Unpack -from transformers.tokenization_utils_base import PreTokenizedInput, TextInput -from transformers.utils import logging -from transformers.video_utils import VideoInput - -logger = logging.get_logger(__name__) - - -FRAME_FACTOR = 2 -FPS = 2.0 -FPS_MIN_FRAMES = 4 -FPS_MAX_FRAMES = 256 - - -def to_rgb(pil_image: Image.Image) -> Image.Image: - if pil_image.mode == "RGBA": - white_background = Image.new("RGB", pil_image.size, (255, 255, 255)) - white_background.paste(pil_image, mask=pil_image.split()[3]) # Use alpha channel as mask - return white_background - else: - return pil_image.convert("RGB") - - -def fetch_image(ele: dict[str, str | Image.Image]) -> Image.Image: - image = ele["image"] if "image" in ele else ele["image_url"] - image_obj = None - if isinstance(image, Image.Image): - image_obj = image - elif image.startswith("http://") or image.startswith("https://"): - response = requests.get(image, stream=True, timeout=10) - image_obj = Image.open(BytesIO(response.content)) - elif image.startswith("file://"): - image_obj = Image.open(image[7:]) - elif image.startswith("data:image"): - if "base64," in image: - _, base64_data = image.split("base64,", 1) - data = base64.b64decode(base64_data) - image_obj = Image.open(BytesIO(data)) - else: - image_obj = Image.open(image) - if image_obj is None: - raise ValueError( - f"Unrecognized image input, support local path, http url, base64 and PIL.Image, got {image}" - ) - image = to_rgb(image_obj) - if "scale_factor" in ele: - scale_factor = ele["scale_factor"] - image = image.resize((image.width * scale_factor, image.height * scale_factor), Image.BILINEAR) - return image - - -class Eagle25VLProcessorKwargs(ProcessingKwargs, total=False): - # see processing_utils.ProcessingKwargs documentation for usage. - _defaults = { - "text_kwargs": { - "padding": False, - }, - "images_kwargs": {}, - "videos_kwargs": {"max_dynamic_tiles": 1}, - } - - -class Eagle25VLProcessor(ProcessorMixin): - r""" - Constructs a Eagle25VL processor which wraps a Eagle25VL video processor, Eagle25VL image processor and a Eagle25VL tokenizer into a single processor. - - [`Eagle25VLProcessor`] offers all the functionalities of [`Eagle25VLVideoProcessor`], [`Eagle25VLImageProcessor`] and [`Eagle25VLTokenizer`]. See the - [`~Eagle25VLVideoProcessor.__call__`], [`~Eagle25VLProcessor.__call__`] and [`~Eagle25VLProcessor.decode`] for more information. - - Args: - image_processor ([`LlavaOnevisionImageProcessor`], *optional*): - The image processor is a required input. - tokenizer ([`LlamaTokenizerFast`], *optional*): - The tokenizer is a required input. - num_image_tokens (`int`, *optional*): - Number of image tokens for one imagethat will be returned by vision tower. - vision_feature_select_strategy (`str`, *optional*): - The feature selection strategy used to select the vision feature from the vision backbone. - Should be same as in model's config - chat_template (`str`, *optional*): A Jinja template which will be used to convert lists of messages - in a chat into a tokenizable string. - image_token (`str`, *optional*, defaults to `""`): - Special token used to denote image location. - video_token (`str`, *optional*, defaults to `"