Add files using upload-large-folder tool

.gitattributes

@@ -18,3 +18,35 @@
 *.mp4 filter=lfs diff=lfs merge=lfs -text
 *.arrow filter=lfs diff=lfs merge=lfs -text
 *.json !text !filter !merge !diff
+media/lerobot-logo-light.png filter=lfs diff=lfs merge=lfs -text
+media/aloha/follower_rotated.webp filter=lfs diff=lfs merge=lfs -text
+media/aloha/follower_zero.webp filter=lfs diff=lfs merge=lfs -text
+media/lerobot-logo-thumbnail.png filter=lfs diff=lfs merge=lfs -text
+media/wandb.png filter=lfs diff=lfs merge=lfs -text
+media/gym/pusht_diffusion.gif filter=lfs diff=lfs merge=lfs -text
+media/aloha/leader_rest.webp filter=lfs diff=lfs merge=lfs -text
+media/aloha/leader_zero.webp filter=lfs diff=lfs merge=lfs -text
+media/koch/follower_rest.webp filter=lfs diff=lfs merge=lfs -text
+media/koch/follower_zero.webp filter=lfs diff=lfs merge=lfs -text
+media/gym/aloha_act.gif filter=lfs diff=lfs merge=lfs -text
+media/aloha/leader_rotated.webp filter=lfs diff=lfs merge=lfs -text
+media/koch/follower_rotated.webp filter=lfs diff=lfs merge=lfs -text
+media/koch/leader_rest.webp filter=lfs diff=lfs merge=lfs -text
+media/gym/simxarm_tdmpc.gif filter=lfs diff=lfs merge=lfs -text
+media/koch/leader_rotated.webp filter=lfs diff=lfs merge=lfs -text
+media/aloha/follower_rest.webp filter=lfs diff=lfs merge=lfs -text
+media/koch/leader_zero.webp filter=lfs diff=lfs merge=lfs -text
+media/lekiwi/kiwi.webp filter=lfs diff=lfs merge=lfs -text
+media/lekiwi/mobile_calib_rest.webp filter=lfs diff=lfs merge=lfs -text
+media/lekiwi/mobile_calib_rotated.webp filter=lfs diff=lfs merge=lfs -text
+media/moss/follower_initial.webp filter=lfs diff=lfs merge=lfs -text
+media/lekiwi/motor_ids.webp filter=lfs diff=lfs merge=lfs -text
+media/moss/follower_rotated.webp filter=lfs diff=lfs merge=lfs -text
+media/lekiwi/mobile_calib_zero.webp filter=lfs diff=lfs merge=lfs -text
+media/moss/follower_rest.webp filter=lfs diff=lfs merge=lfs -text
+media/moss/leader_rotated.webp filter=lfs diff=lfs merge=lfs -text
+media/moss/follower_zero.webp filter=lfs diff=lfs merge=lfs -text
+media/moss/leader_zero.webp filter=lfs diff=lfs merge=lfs -text
+media/so100/follower_initial.webp filter=lfs diff=lfs merge=lfs -text
+media/so100/follower_rest.webp filter=lfs diff=lfs merge=lfs -text
+media/so100/leader_follower.webp filter=lfs diff=lfs merge=lfs -text

lerobot/common/policies/vqbet/configuration_vqbet.py

@@ -0,0 +1,200 @@
+#!/usr/bin/env python
+
+# Copyright 2024 Seungjae Lee and Yibin Wang and Haritheja Etukuru
+# and H. Jin Kim and Nur Muhammad Mahi Shafiullah and Lerrel Pinto
+# 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
+#
+# 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 dataclass, field
+
+from lerobot.common.optim.optimizers import AdamConfig
+from lerobot.common.optim.schedulers import VQBeTSchedulerConfig
+from lerobot.configs.policies import PreTrainedConfig
+from lerobot.configs.types import NormalizationMode
+
+
+@PreTrainedConfig.register_subclass("vqbet")
+@dataclass
+class VQBeTConfig(PreTrainedConfig):
+    """Configuration class for VQ-BeT.
+
+    Defaults are configured for training with PushT providing proprioceptive and single camera observations.
+
+    The parameters you will most likely need to change are the ones which depend on the environment / sensors.
+    Those are: `input_shapes` and `output_shapes`.
+
+    Notes on the inputs and outputs:
+        - "observation.state" is required as an input key.
+        - At least one key starting with "observation.image is required as an input.
+        - If there are multiple keys beginning with "observation.image" they are treated as multiple camera
+          views. Right now we only support all images having the same shape.
+        - "action" is required as an output key.
+
+    Args:
+        n_obs_steps: Number of environment steps worth of observations to pass to the policy (takes the
+            current step and additional steps going back).
+        n_action_pred_token: Total number of current token and future tokens that VQ-BeT predicts.
+        action_chunk_size: Action chunk size of each action prediction token.
+        input_shapes: A dictionary defining the shapes of the input data for the policy.
+            The key represents the input data name, and the value is a list indicating the dimensions
+            of the corresponding data. For example, "observation.image" refers to an input from
+            a camera with dimensions [3, 96, 96], indicating it has three color channels and 96x96 resolution.
+            Importantly, shapes doesnt include batch dimension or temporal dimension.
+        output_shapes: A dictionary defining the shapes of the output data for the policy.
+            The key represents the output data name, and the value is a list indicating the dimensions
+            of the corresponding data. For example, "action" refers to an output shape of [14], indicating
+            14-dimensional actions. Importantly, shapes doesnt include batch dimension or temporal dimension.
+        input_normalization_modes: A dictionary with key representing the modality (e.g. "observation.state"),
+            and the value specifies the normalization mode to apply. The two available modes are "mean_std"
+            which subtracts the mean and divides by the standard deviation and "min_max" which rescale in a
+            [-1, 1] range.
+        output_normalization_modes: Similar dictionary as `normalize_input_modes`, but to unnormalize to the
+            original scale. Note that this is also used for normalizing the training targets.
+        vision_backbone: Name of the torchvision resnet backbone to use for encoding images.
+        crop_shape: (H, W) shape to crop images to as a preprocessing step for the vision backbone. Must fit
+            within the image size. If None, no cropping is done.
+        crop_is_random: Whether the crop should be random at training time (it's always a center crop in eval
+            mode).
+        pretrained_backbone_weights: Pretrained weights from torchvision to initialize the backbone.
+            `None` means no pretrained weights.
+        use_group_norm: Whether to replace batch normalization with group normalization in the backbone.
+            The group sizes are set to be about 16 (to be precise, feature_dim // 16).
+        spatial_softmax_num_keypoints: Number of keypoints for SpatialSoftmax.
+        n_vqvae_training_steps: Number of optimization steps for training Residual VQ.
+        vqvae_n_embed: Number of embedding vectors in the RVQ dictionary (each layer).
+        vqvae_embedding_dim: Dimension of each embedding vector in the RVQ dictionary.
+        vqvae_enc_hidden_dim: Size of hidden dimensions of Encoder / Decoder part of Residaul VQ-VAE
+        gpt_block_size: Max block size of minGPT (should be larger than the number of input tokens)
+        gpt_input_dim: Size of output input of GPT. This is also used as the dimension of observation features.
+        gpt_output_dim: Size of output dimension of GPT. This is also used as a input dimension of offset / bin prediction headers.
+        gpt_n_layer: Number of layers of GPT
+        gpt_n_head: Number of headers of GPT
+        gpt_hidden_dim: Size of hidden dimensions of GPT
+        dropout: Dropout rate for GPT
+        mlp_hidden_dim: Size of hidden dimensions of offset header / bin prediction headers parts of VQ-BeT
+        offset_loss_weight:  A constant that is multiplied to the offset loss
+        primary_code_loss_weight: A constant that is multiplied to the primary code prediction loss
+        secondary_code_loss_weight: A constant that is multiplied to the secondary code prediction loss
+        bet_softmax_temperature: Sampling temperature of code for rollout with VQ-BeT
+        sequentially_select: Whether select code of primary / secondary as sequentially (pick primary code,
+            and then select secodnary code), or at the same time.
+    """
+
+    # Inputs / output structure.
+    n_obs_steps: int = 5
+    n_action_pred_token: int = 3
+    action_chunk_size: int = 5
+
+    normalization_mapping: dict[str, NormalizationMode] = field(
+        default_factory=lambda: {
+            "VISUAL": NormalizationMode.IDENTITY,
+            "STATE": NormalizationMode.MIN_MAX,
+            "ACTION": NormalizationMode.MIN_MAX,
+        }
+    )
+
+    # Architecture / modeling.
+    # Vision backbone.
+    vision_backbone: str = "resnet18"
+    crop_shape: tuple[int, int] | None = (84, 84)
+    crop_is_random: bool = True
+    pretrained_backbone_weights: str | None = None
+    use_group_norm: bool = True
+    spatial_softmax_num_keypoints: int = 32
+    # VQ-VAE
+    n_vqvae_training_steps: int = 20000
+    vqvae_n_embed: int = 16
+    vqvae_embedding_dim: int = 256
+    vqvae_enc_hidden_dim: int = 128
+    # VQ-BeT
+    gpt_block_size: int = 500
+    gpt_input_dim: int = 512
+    gpt_output_dim: int = 512
+    gpt_n_layer: int = 8
+    gpt_n_head: int = 8
+    gpt_hidden_dim: int = 512
+    dropout: float = 0.1
+    mlp_hidden_dim: int = 1024
+    offset_loss_weight: float = 10000.0
+    primary_code_loss_weight: float = 5.0
+    secondary_code_loss_weight: float = 0.5
+    bet_softmax_temperature: float = 0.1
+    sequentially_select: bool = False
+
+    # Training presets
+    optimizer_lr: float = 1e-4
+    optimizer_betas: tuple = (0.95, 0.999)
+    optimizer_eps: float = 1e-8
+    optimizer_weight_decay: float = 1e-6
+    optimizer_vqvae_lr: float = 1e-3
+    optimizer_vqvae_weight_decay: float = 1e-4
+    scheduler_warmup_steps: int = 500
+
+    def __post_init__(self):
+        super().__post_init__()
+
+        """Input validation (not exhaustive)."""
+        if not self.vision_backbone.startswith("resnet"):
+            raise ValueError(
+                f"`vision_backbone` must be one of the ResNet variants. Got {self.vision_backbone}."
+            )
+
+    def get_optimizer_preset(self) -> AdamConfig:
+        return AdamConfig(
+            lr=self.optimizer_lr,
+            betas=self.optimizer_betas,
+            eps=self.optimizer_eps,
+            weight_decay=self.optimizer_weight_decay,
+        )
+
+    def get_scheduler_preset(self) -> VQBeTSchedulerConfig:
+        return VQBeTSchedulerConfig(
+            num_warmup_steps=self.scheduler_warmup_steps,
+            num_vqvae_training_steps=self.n_vqvae_training_steps,
+        )
+
+    def validate_features(self) -> None:
+        # Note: this check was previously performed inside VQBeTRgbEncoder in the form of
+        # assert len(image_keys) == 1
+        if not len(self.image_features) == 1:
+            raise ValueError("You must provide only one image among the inputs.")
+
+        if self.crop_shape is not None:
+            for key, image_ft in self.image_features.items():
+                if self.crop_shape[0] > image_ft.shape[1] or self.crop_shape[1] > image_ft.shape[2]:
+                    raise ValueError(
+                        f"`crop_shape` should fit within the images shapes. Got {self.crop_shape} "
+                        f"for `crop_shape` and {image_ft.shape} for "
+                        f"`{key}`."
+                    )
+
+        # Check that all input images have the same shape.
+        first_image_key, first_image_ft = next(iter(self.image_features.items()))
+        for key, image_ft in self.image_features.items():
+            if image_ft.shape != first_image_ft.shape:
+                raise ValueError(
+                    f"`{key}` does not match `{first_image_key}`, but we expect all image shapes to match."
+                )
+
+    @property
+    def observation_delta_indices(self) -> list:
+        return list(range(1 - self.n_obs_steps, 1))
+
+    @property
+    def action_delta_indices(self) -> list:
+        return list(range(1 - self.n_obs_steps, self.n_action_pred_token + self.action_chunk_size - 1))
+
+    @property
+    def reward_delta_indices(self) -> None:
+        return None

lerobot/common/policies/vqbet/modeling_vqbet.py

@@ -0,0 +1,911 @@
+#!/usr/bin/env python
+
+# Copyright 2024 Seungjae Lee and Yibin Wang and Haritheja Etukuru
+# and H. Jin Kim and Nur Muhammad Mahi Shafiullah and Lerrel Pinto
+# 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
+#
+# 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 warnings
+from collections import deque
+from typing import Callable, List
+
+import einops
+import numpy as np
+import torch
+import torch.nn.functional as F  # noqa: N812
+import torchvision
+from torch import Tensor, nn
+
+from lerobot.common.policies.normalize import Normalize, Unnormalize
+from lerobot.common.policies.pretrained import PreTrainedPolicy
+from lerobot.common.policies.utils import get_device_from_parameters, get_output_shape, populate_queues
+from lerobot.common.policies.vqbet.configuration_vqbet import VQBeTConfig
+from lerobot.common.policies.vqbet.vqbet_utils import GPT, ResidualVQ
+
+# ruff: noqa: N806
+
+
+class VQBeTPolicy(PreTrainedPolicy):
+    """
+    VQ-BeT Policy as per "Behavior Generation with Latent Actions"
+    """
+
+    config_class = VQBeTConfig
+    name = "vqbet"
+
+    def __init__(
+        self,
+        config: VQBeTConfig | None = None,
+        dataset_stats: dict[str, dict[str, Tensor]] | None = None,
+    ):
+        """
+        Args:
+            config: Policy configuration class instance or None, in which case the default instantiation of
+                the configuration class is used.
+            dataset_stats: Dataset statistics to be used for normalization. If not passed here, it is expected
+                that they will be passed with a call to `load_state_dict` before the policy is used.
+        """
+        super().__init__(config)
+        config.validate_features()
+        self.config = config
+
+        self.normalize_inputs = Normalize(config.input_features, config.normalization_mapping, dataset_stats)
+        self.normalize_targets = Normalize(
+            config.output_features, config.normalization_mapping, dataset_stats
+        )
+        self.unnormalize_outputs = Unnormalize(
+            config.output_features, config.normalization_mapping, dataset_stats
+        )
+
+        self.vqbet = VQBeTModel(config)
+
+        self.reset()
+
+    def get_optim_params(self) -> dict:
+        vqvae_params = (
+            list(self.vqbet.action_head.vqvae_model.encoder.parameters())
+            + list(self.vqbet.action_head.vqvae_model.decoder.parameters())
+            + list(self.vqbet.action_head.vqvae_model.vq_layer.parameters())
+        )
+        decay_params, no_decay_params = self.vqbet.policy.configure_parameters()
+        decay_params = (
+            decay_params
+            + list(self.vqbet.rgb_encoder.parameters())
+            + list(self.vqbet.state_projector.parameters())
+            + list(self.vqbet.rgb_feature_projector.parameters())
+            + [self.vqbet.action_token]
+            + list(self.vqbet.action_head.map_to_cbet_preds_offset.parameters())
+        )
+
+        if self.config.sequentially_select:
+            decay_params = (
+                decay_params
+                + list(self.vqbet.action_head.map_to_cbet_preds_primary_bin.parameters())
+                + list(self.vqbet.action_head.map_to_cbet_preds_secondary_bin.parameters())
+            )
+        else:
+            decay_params = decay_params + list(self.vqbet.action_head.map_to_cbet_preds_bin.parameters())
+
+        return [
+            {
+                "params": decay_params,
+            },
+            {
+                "params": vqvae_params,
+                "weight_decay": self.config.optimizer_vqvae_weight_decay,
+                "lr": self.config.optimizer_vqvae_lr,
+            },
+            {
+                "params": no_decay_params,
+                "weight_decay": 0.0,
+            },
+        ]
+
+    def reset(self):
+        """
+        Clear observation and action queues. Should be called on `env.reset()`
+        queues are populated during rollout of the policy, they contain the n latest observations and actions
+        """
+        self._queues = {
+            "observation.images": deque(maxlen=self.config.n_obs_steps),
+            "observation.state": deque(maxlen=self.config.n_obs_steps),
+            "action": deque(maxlen=self.config.action_chunk_size),
+        }
+
+    @torch.no_grad
+    def select_action(self, batch: dict[str, Tensor]) -> Tensor:
+        """Select a single action given environment observations.
+
+        This method wraps `select_actions` in order to return one action at a time for execution in the
+        environment. It works by managing the actions in a queue and only calling `select_actions` when the
+        queue is empty.
+        """
+
+        batch = self.normalize_inputs(batch)
+        batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original
+        batch["observation.images"] = torch.stack([batch[key] for key in self.config.image_features], dim=-4)
+        # Note: It's important that this happens after stacking the images into a single key.
+        self._queues = populate_queues(self._queues, batch)
+
+        if not self.vqbet.action_head.vqvae_model.discretized.item():
+            warnings.warn(
+                "To evaluate in the environment, your VQ-BeT model should contain a pretrained Residual VQ.",
+                stacklevel=1,
+            )
+
+        if len(self._queues["action"]) == 0:
+            batch = {k: torch.stack(list(self._queues[k]), dim=1) for k in batch if k in self._queues}
+            actions = self.vqbet(batch, rollout=True)[:, : self.config.action_chunk_size]
+
+            # the dimension of returned action is (batch_size, action_chunk_size, action_dim)
+            actions = self.unnormalize_outputs({"action": actions})["action"]
+            # since the data in the action queue's dimension is (action_chunk_size, batch_size, action_dim), we transpose the action and fill the queue
+            self._queues["action"].extend(actions.transpose(0, 1))
+
+        action = self._queues["action"].popleft()
+        return action
+
+    def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict]:
+        """Run the batch through the model and compute the loss for training or validation."""
+        batch = self.normalize_inputs(batch)
+        batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original
+        batch["observation.images"] = torch.stack([batch[key] for key in self.config.image_features], dim=-4)
+        batch = self.normalize_targets(batch)
+        # VQ-BeT discretizes action using VQ-VAE before training BeT (please refer to section 3.2 in the VQ-BeT paper https://arxiv.org/pdf/2403.03181)
+        if not self.vqbet.action_head.vqvae_model.discretized.item():
+            # loss: total loss of training RVQ
+            # n_different_codes: how many of the total possible VQ codes are being used in single batch (how many of them have at least one encoder embedding as a nearest neighbor). This can be at most `vqvae_n_embed * number of layers of RVQ (=2)`.
+            # n_different_combinations: how many different code combinations are being used out of all possible combinations in single batch. This can be at most `vqvae_n_embed ^ number of layers of RVQ (=2)` (hint consider the RVQ as a decision tree).
+            loss, n_different_codes, n_different_combinations, recon_l1_error = (
+                self.vqbet.action_head.discretize(self.config.n_vqvae_training_steps, batch["action"])
+            )
+            return loss, {
+                "n_different_codes": n_different_codes,
+                "n_different_combinations": n_different_combinations,
+                "recon_l1_error": recon_l1_error,
+            }
+        # if Residual VQ is already trained, VQ-BeT trains its GPT and bin prediction head / offset prediction head parts.
+        _, loss_dict = self.vqbet(batch, rollout=False)
+        loss = loss_dict.pop("loss")
+
+        return loss, loss_dict
+
+
+class SpatialSoftmax(nn.Module):
+    """
+    Spatial Soft Argmax operation described in "Deep Spatial Autoencoders for Visuomotor Learning" by Finn et al.
+    (https://arxiv.org/pdf/1509.06113). A minimal port of the robomimic implementation.
+
+    At a high level, this takes 2D feature maps (from a convnet/ViT) and returns the "center of mass"
+    of activations of each channel, i.e., keypoints in the image space for the policy to focus on.
+
+    Example: take feature maps of size (512x10x12). We generate a grid of normalized coordinates (10x12x2):
+    -----------------------------------------------------
+    | (-1., -1.)   | (-0.82, -1.)   | ... | (1., -1.)   |
+    | (-1., -0.78) | (-0.82, -0.78) | ... | (1., -0.78) |
+    | ...          | ...            | ... | ...         |
+    | (-1., 1.)    | (-0.82, 1.)    | ... | (1., 1.)    |
+    -----------------------------------------------------
+    This is achieved by applying channel-wise softmax over the activations (512x120) and computing the dot
+    product with the coordinates (120x2) to get expected points of maximal activation (512x2).
+
+    The example above results in 512 keypoints (corresponding to the 512 input channels). We can optionally
+    provide num_kp != None to control the number of keypoints. This is achieved by a first applying a learnable
+    linear mapping (in_channels, H, W) -> (num_kp, H, W).
+    """
+
+    def __init__(self, input_shape, num_kp=None):
+        """
+        Args:
+            input_shape (list): (C, H, W) input feature map shape.
+            num_kp (int): number of keypoints in output. If None, output will have the same number of channels as input.
+        """
+        super().__init__()
+
+        assert len(input_shape) == 3
+        self._in_c, self._in_h, self._in_w = input_shape
+
+        if num_kp is not None:
+            self.nets = torch.nn.Conv2d(self._in_c, num_kp, kernel_size=1)
+            self._out_c = num_kp
+        else:
+            self.nets = None
+            self._out_c = self._in_c
+
+        # we could use torch.linspace directly but that seems to behave slightly differently than numpy
+        # and causes a small degradation in pc_success of pre-trained models.
+        pos_x, pos_y = np.meshgrid(np.linspace(-1.0, 1.0, self._in_w), np.linspace(-1.0, 1.0, self._in_h))
+        pos_x = torch.from_numpy(pos_x.reshape(self._in_h * self._in_w, 1)).float()
+        pos_y = torch.from_numpy(pos_y.reshape(self._in_h * self._in_w, 1)).float()
+        # register as buffer so it's moved to the correct device.
+        self.register_buffer("pos_grid", torch.cat([pos_x, pos_y], dim=1))
+
+    def forward(self, features: Tensor) -> Tensor:
+        """
+        Args:
+            features: (B, C, H, W) input feature maps.
+        Returns:
+            (B, K, 2) image-space coordinates of keypoints.
+        """
+        if self.nets is not None:
+            features = self.nets(features)
+
+        # [B, K, H, W] -> [B * K, H * W] where K is number of keypoints
+        features = features.reshape(-1, self._in_h * self._in_w)
+        # 2d softmax normalization
+        attention = F.softmax(features, dim=-1)
+        # [B * K, H * W] x [H * W, 2] -> [B * K, 2] for spatial coordinate mean in x and y dimensions
+        expected_xy = attention @ self.pos_grid
+        # reshape to [B, K, 2]
+        feature_keypoints = expected_xy.view(-1, self._out_c, 2)
+
+        return feature_keypoints
+
+
+class VQBeTModel(nn.Module):
+    """VQ-BeT: The underlying neural network for VQ-BeT
+
+    Note: In this code we use the terms `rgb_encoder`, 'policy', `action_head`. The meanings are as follows.
+        - The `rgb_encoder` process rgb-style image observations to one-dimensional embedding vectors
+        - A `policy` is a minGPT architecture, that takes observation sequences and action query tokens to generate `features`.
+        - These `features` pass through the action head, which passes through the code prediction, offset prediction head,
+        and finally generates a prediction for the action chunks.
+
+        -------------------------------** legend **-------------------------------
+        │   n = n_obs_steps, p = n_action_pred_token, c = action_chunk_size)   │
+        │   o_{t} : visual observation at timestep {t}                           │
+        │   s_{t} : state observation at timestep {t}                            │
+        │   a_{t} : action at timestep {t}                                       │
+        │   A_Q : action_query_token                                             │
+        --------------------------------------------------------------------------
+
+
+        Training Phase 1. Discretize action using Residual VQ (for config.n_vqvae_training_steps steps)
+
+
+        ┌─────────────────┐            ┌─────────────────┐            ┌─────────────────┐
+        │                 │            │                 │            │                 │
+        │   RVQ encoder   │    ─►      │     Residual    │    ─►      │   RVQ Decoder   │
+        │ (a_{t}~a_{t+p}) │            │  Code Quantizer │            │                 │
+        │                 │            │                 │            │                 │
+        └─────────────────┘            └─────────────────┘            └─────────────────┘
+
+        Training Phase 2.
+
+          timestep {t-n+1}   timestep {t-n+2}                timestep {t}
+            ┌─────┴─────┐     ┌─────┴─────┐                 ┌─────┴─────┐
+
+        o_{t-n+1}         o_{t-n+2}           ...         o_{t}
+            │                 │                             │
+            │ s_{t-n+1}       │ s_{t-n+2}         ...       │   s_{t}           p
+            │     │           │     │                       │     │     ┌───────┴───────┐
+            │     │    A_Q    │     │    A_Q          ...   │     │    A_Q     ...     A_Q
+            │     │     │     │     │     │                 │     │     │               │
+        ┌───▼─────▼─────▼─────▼─────▼─────▼─────────────────▼─────▼─────▼───────────────▼───┐
+        │                                                                                   │
+        │                                       GPT                                         │       =>    policy
+        │                                                                                   │
+        └───────────────▼─────────────────▼─────────────────────────────▼───────────────▼───┘
+                        │                 │                             │               │
+                    ┌───┴───┐         ┌───┴───┐                     ┌───┴───┐       ┌───┴───┐
+                  code    offset    code    offset                code    offset  code    offset
+                    ▼       │         ▼       │                     ▼       │       ▼       │       =>    action_head
+               RVQ Decoder  │    RVQ Decoder  │                RVQ Decoder  │  RVQ Decoder  │
+                    └── + ──┘         └── + ──┘                     └── + ──┘       └── + ──┘
+                        ▼                 ▼                             ▼               ▼
+                   action chunk      action chunk                  action chunk     action chunk
+                    a_{t-n+1} ~       a_{t-n+2} ~                   a_{t} ~     ...  a_{t+p-1} ~
+                     a_{t-n+c}         a_{t-n+c+1}                   a_{t+c-1}        a_{t+p+c-1}
+
+                                                                        ▼
+                                                      ONLY this chunk is used in rollout!
+    """
+
+    def __init__(self, config: VQBeTConfig):
+        super().__init__()
+        self.config = config
+
+        self.rgb_encoder = VQBeTRgbEncoder(config)
+        self.num_images = len(self.config.image_features)
+        # This action query token is used as a prompt for querying action chunks. Please refer to "A_Q" in the image above.
+        # Note: During the forward pass, this token is repeated as many times as needed. The authors also experimented with initializing the necessary number of tokens independently and observed inferior results.
+        self.action_token = nn.Parameter(torch.randn(1, 1, self.config.gpt_input_dim))
+
+        # To input state and observation features into GPT layers, we first project the features to fit the shape of input size of GPT.
+        self.state_projector = MLP(
+            config.robot_state_feature.shape[0], hidden_channels=[self.config.gpt_input_dim]
+        )
+        self.rgb_feature_projector = MLP(
+            self.rgb_encoder.feature_dim, hidden_channels=[self.config.gpt_input_dim]
+        )
+
+        # GPT part of VQ-BeT
+        self.policy = GPT(config)
+        # bin prediction head / offset prediction head part of VQ-BeT
+        self.action_head = VQBeTHead(config)
+
+        # Action tokens for: each observation step, the current action token, and all future action tokens.
+        num_tokens = self.config.n_action_pred_token + self.config.n_obs_steps - 1
+        self.register_buffer(
+            "select_target_actions_indices",
+            torch.row_stack([torch.arange(i, i + self.config.action_chunk_size) for i in range(num_tokens)]),
+        )
+
+    def forward(self, batch: dict[str, Tensor], rollout: bool) -> tuple[dict, dict]:
+        # Input validation.
+        assert set(batch).issuperset({"observation.state", "observation.images"})
+        batch_size, n_obs_steps = batch["observation.state"].shape[:2]
+        assert n_obs_steps == self.config.n_obs_steps
+
+        # Extract image feature (first combine batch and sequence dims).
+        img_features = self.rgb_encoder(
+            einops.rearrange(batch["observation.images"], "b s n ... -> (b s n) ...")
+        )
+        # Separate batch and sequence dims.
+        img_features = einops.rearrange(
+            img_features, "(b s n) ... -> b s n ...", b=batch_size, s=n_obs_steps, n=self.num_images
+        )
+
+        # Arrange prior and current observation step tokens as shown in the class docstring.
+        # First project features to token dimension.
+        rgb_tokens = self.rgb_feature_projector(
+            img_features
+        )  # (batch, obs_step, number of different cameras, projection dims)
+        input_tokens = [rgb_tokens[:, :, i] for i in range(rgb_tokens.size(2))]
+        input_tokens.append(
+            self.state_projector(batch["observation.state"])
+        )  # (batch, obs_step, projection dims)
+        input_tokens.append(einops.repeat(self.action_token, "1 1 d -> b n d", b=batch_size, n=n_obs_steps))
+        # Interleave tokens by stacking and rearranging.
+        input_tokens = torch.stack(input_tokens, dim=2)
+        input_tokens = einops.rearrange(input_tokens, "b n t d -> b (n t) d")
+
+        len_additional_action_token = self.config.n_action_pred_token - 1
+        future_action_tokens = self.action_token.repeat(batch_size, len_additional_action_token, 1)
+
+        # add additional action query tokens for predicting future action chunks
+        input_tokens = torch.cat([input_tokens, future_action_tokens], dim=1)
+
+        # get action features (pass through GPT)
+        features = self.policy(input_tokens)
+        # len(self.config.input_features) is the number of different observation modes.
+        # this line gets the index of action prompt tokens.
+        historical_act_pred_index = np.arange(0, n_obs_steps) * (len(self.config.input_features) + 1) + len(
+            self.config.input_features
+        )
+
+        # only extract the output tokens at the position of action query:
+        # Behavior Transformer (BeT), and VQ-BeT are both sequence-to-sequence prediction models,
+        # mapping sequential observation to sequential action (please refer to section 2.2 in BeT paper https://arxiv.org/pdf/2206.11251).
+        # Thus, it predicts a historical action sequence, in addition to current and future actions (predicting future actions : optional).
+        if len_additional_action_token > 0:
+            features = torch.cat(
+                [features[:, historical_act_pred_index], features[:, -len_additional_action_token:]], dim=1
+            )
+        else:
+            features = features[:, historical_act_pred_index]
+        # pass through action head
+        action_head_output = self.action_head(features)
+        # if rollout, VQ-BeT don't calculate loss
+        if rollout:
+            return action_head_output["predicted_action"][:, n_obs_steps - 1, :].reshape(
+                batch_size, self.config.action_chunk_size, -1
+            )
+        # else, it calculate overall loss (bin prediction loss, and offset loss)
+        else:
+            output = batch["action"][:, self.select_target_actions_indices]
+            loss = self.action_head.loss_fn(action_head_output, output, reduction="mean")
+            return action_head_output, loss
+
+
+class VQBeTHead(nn.Module):
+    def __init__(self, config: VQBeTConfig):
+        """
+        VQBeTHead takes output of GPT layers, and pass the feature through bin prediction head (`self.map_to_cbet_preds_bin`), and offset prediction head (`self.map_to_cbet_preds_offset`)
+
+        self.map_to_cbet_preds_bin: outputs probability of each code (for each layer).
+            The input dimension of `self.map_to_cbet_preds_bin` is same with the output of GPT,
+            and the output dimension of `self.map_to_cbet_preds_bin` is `self.vqvae_model.vqvae_num_layers (=fixed as 2) * self.config.vqvae_n_embed`.
+            if the agent select the code sequentially, we use self.map_to_cbet_preds_primary_bin and self.map_to_cbet_preds_secondary_bin instead of self._map_to_cbet_preds_bin.
+
+        self.map_to_cbet_preds_offset: output the predicted offsets for all the codes in all the layers.
+            The input dimension of ` self.map_to_cbet_preds_offset` is same with the output of GPT,
+            and the output dimension of ` self.map_to_cbet_preds_offset` is `self.vqvae_model.vqvae_num_layers (=fixed as 2) * self.config.vqvae_n_embed * config.action_chunk_size * config.action_feature.shape[0]`.
+        """
+
+        super().__init__()
+        self.config = config
+        # init vqvae
+        self.vqvae_model = VqVae(config)
+        if config.sequentially_select:
+            self.map_to_cbet_preds_primary_bin = MLP(
+                in_channels=config.gpt_output_dim,
+                hidden_channels=[self.config.vqvae_n_embed],
+            )
+            self.map_to_cbet_preds_secondary_bin = MLP(
+                in_channels=config.gpt_output_dim + self.config.vqvae_n_embed,
+                hidden_channels=[self.config.vqvae_n_embed],
+            )
+        else:
+            self.map_to_cbet_preds_bin = MLP(
+                in_channels=config.gpt_output_dim,
+                hidden_channels=[self.vqvae_model.vqvae_num_layers * self.config.vqvae_n_embed],
+            )
+        self.map_to_cbet_preds_offset = MLP(
+            in_channels=config.gpt_output_dim,
+            hidden_channels=[
+                self.vqvae_model.vqvae_num_layers
+                * self.config.vqvae_n_embed
+                * config.action_chunk_size
+                * config.action_feature.shape[0],
+            ],
+        )
+        # loss
+        self._focal_loss_fn = FocalLoss(gamma=2.0)
+
+    def discretize(self, n_vqvae_training_steps, actions):
+        # Resize the action sequence data to fit the action chunk size using a sliding window approach.
+        actions = torch.cat(
+            [
+                actions[:, j : j + self.config.action_chunk_size, :]
+                for j in range(actions.shape[1] + 1 - self.config.action_chunk_size)
+            ],
+            dim=0,
+        )
+        # `actions` is a tensor of shape (new_batch, action_chunk_size, action_dim) where new_batch is the number of possible chunks created from the original sequences using the sliding window.
+
+        loss, metric = self.vqvae_model.vqvae_forward(actions)
+        n_different_codes = sum(
+            [len(torch.unique(metric[2][:, i])) for i in range(self.vqvae_model.vqvae_num_layers)]
+        )
+        n_different_combinations = len(torch.unique(metric[2], dim=0))
+        recon_l1_error = metric[0].detach().cpu().item()
+        self.vqvae_model.optimized_steps += 1
+        # if we updated RVQ more than `n_vqvae_training_steps` steps, we freeze the RVQ part.
+        if self.vqvae_model.optimized_steps >= n_vqvae_training_steps:
+            self.vqvae_model.discretized = torch.tensor(True)
+            self.vqvae_model.vq_layer.freeze_codebook = torch.tensor(True)
+            print("Finished discretizing action data!")
+            self.vqvae_model.eval()
+            for param in self.vqvae_model.vq_layer.parameters():
+                param.requires_grad = False
+        return loss, n_different_codes, n_different_combinations, recon_l1_error
+
+    def forward(self, x, **kwargs) -> dict:
+        # N is the batch size, and T is number of action query tokens, which are process through same GPT
+        N, T, _ = x.shape
+        # we calculate N and T side parallelly. Thus, the dimensions would be
+        # (batch size * number of action query tokens, action chunk size, action dimension)
+        x = einops.rearrange(x, "N T WA -> (N T) WA")
+
+        # sample offsets
+        cbet_offsets = self.map_to_cbet_preds_offset(x)
+        cbet_offsets = einops.rearrange(
+            cbet_offsets,
+            "(NT) (G C WA) -> (NT) G C WA",
+            G=self.vqvae_model.vqvae_num_layers,
+            C=self.config.vqvae_n_embed,
+        )
+        # if self.config.sequentially_select is True, bin prediction head first sample the primary code, and then sample secondary code
+        if self.config.sequentially_select:
+            cbet_primary_logits = self.map_to_cbet_preds_primary_bin(x)
+
+            # select primary bin first
+            cbet_primary_probs = torch.softmax(
+                cbet_primary_logits / self.config.bet_softmax_temperature, dim=-1
+            )
+            NT, choices = cbet_primary_probs.shape
+            sampled_primary_centers = einops.rearrange(
+                torch.multinomial(cbet_primary_probs.view(-1, choices), num_samples=1),
+                "(NT) 1 -> NT",
+                NT=NT,
+            )
+
+            cbet_secondary_logits = self.map_to_cbet_preds_secondary_bin(
+                torch.cat(
+                    (x, F.one_hot(sampled_primary_centers, num_classes=self.config.vqvae_n_embed)),
+                    axis=1,
+                )
+            )
+            cbet_secondary_probs = torch.softmax(
+                cbet_secondary_logits / self.config.bet_softmax_temperature, dim=-1
+            )
+            sampled_secondary_centers = einops.rearrange(
+                torch.multinomial(cbet_secondary_probs.view(-1, choices), num_samples=1),
+                "(NT) 1 -> NT",
+                NT=NT,
+            )
+            sampled_centers = torch.stack((sampled_primary_centers, sampled_secondary_centers), axis=1)
+            cbet_logits = torch.stack([cbet_primary_logits, cbet_secondary_logits], dim=1)
+        # if self.config.sequentially_select is False, bin prediction head samples primary and secondary code at once.
+        else:
+            cbet_logits = self.map_to_cbet_preds_bin(x)
+            cbet_logits = einops.rearrange(
+                cbet_logits, "(NT) (G C) -> (NT) G C", G=self.vqvae_model.vqvae_num_layers
+            )
+            cbet_probs = torch.softmax(cbet_logits / self.config.bet_softmax_temperature, dim=-1)
+            NT, G, choices = cbet_probs.shape
+            sampled_centers = einops.rearrange(
+                torch.multinomial(cbet_probs.view(-1, choices), num_samples=1),
+                "(NT G) 1 -> NT G",
+                NT=NT,
+            )
+
+        device = get_device_from_parameters(self)
+        indices = (
+            torch.arange(NT, device=device).unsqueeze(1),
+            torch.arange(self.vqvae_model.vqvae_num_layers, device=device).unsqueeze(0),
+            sampled_centers,
+        )
+        # Use advanced indexing to sample the values (Extract the only offsets corresponding to the sampled codes.)
+        sampled_offsets = cbet_offsets[indices]
+        # Then, sum the offsets over the RVQ layers to get a net offset for the bin prediction
+        sampled_offsets = sampled_offsets.sum(dim=1)
+        with torch.no_grad():
+            # Get the centroids (= vectors corresponding to the codes) of each layer to pass it through RVQ decoder
+            return_decoder_input = self.vqvae_model.get_embeddings_from_code(sampled_centers).clone().detach()
+            # pass the centroids through decoder to get actions.
+            decoded_action = self.vqvae_model.get_action_from_latent(return_decoder_input).clone().detach()
+        # reshaped extracted offset to match with decoded centroids
+        sampled_offsets = einops.rearrange(
+            sampled_offsets, "NT (W A) -> NT W A", W=self.config.action_chunk_size
+        )
+        # add offset and decoded centroids
+        predicted_action = decoded_action + sampled_offsets
+        predicted_action = einops.rearrange(
+            predicted_action,
+            "(N T) W A -> N T (W A)",
+            N=N,
+            T=T,
+            W=self.config.action_chunk_size,
+        )
+
+        return {
+            "cbet_logits": cbet_logits,
+            "predicted_action": predicted_action,
+            "sampled_centers": sampled_centers,
+            "decoded_action": decoded_action,
+        }
+
+    def loss_fn(self, pred, target, **kwargs):
+        """
+        for given ground truth action values (target), and prediction (pred) this function calculates the overall loss.
+
+        predicted_action: predicted action chunk (offset + decoded centroids)
+        sampled_centers: sampled centroids (code of RVQ)
+        decoded_action: decoded action, which is produced by passing sampled_centers through RVQ decoder
+        NT: batch size * T
+        T: number of action query tokens, which are process through same GPT
+        cbet_logits: probability of all codes in each layer
+        """
+        action_seq = target
+        predicted_action = pred["predicted_action"]
+        sampled_centers = pred["sampled_centers"]
+        decoded_action = pred["decoded_action"]
+        NT = predicted_action.shape[0] * predicted_action.shape[1]
+
+        cbet_logits = pred["cbet_logits"]
+
+        predicted_action = einops.rearrange(
+            predicted_action, "N T (W A) -> (N T) W A", W=self.config.action_chunk_size
+        )
+
+        action_seq = einops.rearrange(action_seq, "N T W A -> (N T) W A")
+        # Figure out the loss for the actions.
+        # First, we need to find the closest cluster center for each ground truth action.
+        with torch.no_grad():
+            state_vq, action_bins = self.vqvae_model.get_code(action_seq)  # action_bins: NT, G
+
+        # Now we can compute the loss.
+
+        # offset loss is L1 distance between the predicted action and ground truth action
+        offset_loss = F.l1_loss(action_seq, predicted_action)
+
+        # calculate primary code prediction loss
+        cbet_loss1 = self._focal_loss_fn(
+            cbet_logits[:, 0, :],
+            action_bins[:, 0],
+        )
+        # calculate secondary code prediction loss
+        cbet_loss2 = self._focal_loss_fn(
+            cbet_logits[:, 1, :],
+            action_bins[:, 1],
+        )
+        # add all the prediction loss
+        cbet_loss = (
+            cbet_loss1 * self.config.primary_code_loss_weight
+            + cbet_loss2 * self.config.secondary_code_loss_weight
+        )
+
+        equal_primary_code_rate = torch.sum((action_bins[:, 0] == sampled_centers[:, 0]).int()) / (NT)
+        equal_secondary_code_rate = torch.sum((action_bins[:, 1] == sampled_centers[:, 1]).int()) / (NT)
+
+        action_mse_error = torch.mean((action_seq - predicted_action) ** 2)
+        vq_action_error = torch.mean(torch.abs(action_seq - decoded_action))
+        offset_action_error = torch.mean(torch.abs(action_seq - predicted_action))
+        action_error_max = torch.max(torch.abs(action_seq - predicted_action))
+
+        loss = cbet_loss + self.config.offset_loss_weight * offset_loss
+
+        loss_dict = {
+            "loss": loss,
+            "classification_loss": cbet_loss.detach().cpu().item(),
+            "offset_loss": offset_loss.detach().cpu().item(),
+            "equal_primary_code_rate": equal_primary_code_rate.detach().cpu().item(),
+            "equal_secondary_code_rate": equal_secondary_code_rate.detach().cpu().item(),
+            "vq_action_error": vq_action_error.detach().cpu().item(),
+            "offset_action_error": offset_action_error.detach().cpu().item(),
+            "action_error_max": action_error_max.detach().cpu().item(),
+            "action_mse_error": action_mse_error.detach().cpu().item(),
+        }
+        return loss_dict
+
+
+class VQBeTRgbEncoder(nn.Module):
+    """Encode an RGB image into a 1D feature vector.
+
+    Includes the ability to normalize and crop the image first.
+
+    Same with DiffusionRgbEncoder from modeling_diffusion.py
+    """
+
+    def __init__(self, config: VQBeTConfig):
+        super().__init__()
+        # Set up optional preprocessing.
+        if config.crop_shape is not None:
+            self.do_crop = True
+            # Always use center crop for eval
+            self.center_crop = torchvision.transforms.CenterCrop(config.crop_shape)
+            if config.crop_is_random:
+                self.maybe_random_crop = torchvision.transforms.RandomCrop(config.crop_shape)
+            else:
+                self.maybe_random_crop = self.center_crop
+        else:
+            self.do_crop = False
+
+        # Set up backbone.
+        backbone_model = getattr(torchvision.models, config.vision_backbone)(
+            weights=config.pretrained_backbone_weights
+        )
+        # Note: This assumes that the layer4 feature map is children()[-3]
+        # TODO(alexander-soare): Use a safer alternative.
+        self.backbone = nn.Sequential(*(list(backbone_model.children())[:-2]))
+        if config.use_group_norm:
+            if config.pretrained_backbone_weights:
+                raise ValueError(
+                    "You can't replace BatchNorm in a pretrained model without ruining the weights!"
+                )
+            self.backbone = _replace_submodules(
+                root_module=self.backbone,
+                predicate=lambda x: isinstance(x, nn.BatchNorm2d),
+                func=lambda x: nn.GroupNorm(num_groups=x.num_features // 16, num_channels=x.num_features),
+            )
+
+        # Set up pooling and final layers.
+        # Use a dry run to get the feature map shape.
+        # The dummy input should take the number of image channels from `config.image_features` and it should
+        # use the height and width from `config.crop_shape` if it is provided, otherwise it should use the
+        # height and width from `config.image_features`.
+
+        images_shape = next(iter(config.image_features.values())).shape
+        dummy_shape_h_w = config.crop_shape if config.crop_shape is not None else images_shape[1:]
+        dummy_shape = (1, images_shape[0], *dummy_shape_h_w)
+        feature_map_shape = get_output_shape(self.backbone, dummy_shape)[1:]
+
+        self.pool = SpatialSoftmax(feature_map_shape, num_kp=config.spatial_softmax_num_keypoints)
+        self.feature_dim = config.spatial_softmax_num_keypoints * 2
+        self.out = nn.Linear(config.spatial_softmax_num_keypoints * 2, self.feature_dim)
+        self.relu = nn.ReLU()
+
+    def forward(self, x: Tensor) -> Tensor:
+        """
+        Args:
+            x: (B, C, H, W) image tensor with pixel values in [0, 1].
+        Returns:
+            (B, D) image feature.
+        """
+        # Preprocess: maybe crop (if it was set up in the __init__).
+        if self.do_crop:
+            if self.training:  # noqa: SIM108
+                x = self.maybe_random_crop(x)
+            else:
+                # Always use center crop for eval.
+                x = self.center_crop(x)
+        # Extract backbone feature.
+        x = torch.flatten(self.pool(self.backbone(x)), start_dim=1)
+        # Final linear layer with non-linearity.
+        x = self.relu(self.out(x))
+        return x
+
+
+def _replace_submodules(
+    root_module: nn.Module, predicate: Callable[[nn.Module], bool], func: Callable[[nn.Module], nn.Module]
+) -> nn.Module:
+    """
+    Args:
+        root_module: The module for which the submodules need to be replaced
+        predicate: Takes a module as an argument and must return True if the that module is to be replaced.
+        func: Takes a module as an argument and returns a new module to replace it with.
+    Returns:
+        The root module with its submodules replaced.
+    """
+    if predicate(root_module):
+        return func(root_module)
+
+    replace_list = [k.split(".") for k, m in root_module.named_modules(remove_duplicate=True) if predicate(m)]
+    for *parents, k in replace_list:
+        parent_module = root_module
+        if len(parents) > 0:
+            parent_module = root_module.get_submodule(".".join(parents))
+        if isinstance(parent_module, nn.Sequential):
+            src_module = parent_module[int(k)]
+        else:
+            src_module = getattr(parent_module, k)
+        tgt_module = func(src_module)
+        if isinstance(parent_module, nn.Sequential):
+            parent_module[int(k)] = tgt_module
+        else:
+            setattr(parent_module, k, tgt_module)
+    # verify that all BN are replaced
+    assert not any(predicate(m) for _, m in root_module.named_modules(remove_duplicate=True))
+    return root_module
+
+
+class VqVae(nn.Module):
+    def __init__(
+        self,
+        config: VQBeTConfig,
+    ):
+        """
+        VQ-VAE is composed of three parts: encoder, vq_layer, and decoder.
+        Encoder and decoder are MLPs consisting of an input, output layer, and hidden layer, respectively.
+        The vq_layer uses residual VQs.
+
+        This class contains functions for training the encoder and decoder along with the residual VQ layer (for training phase 1),
+        as well as functions to help BeT training part in training phase 2.
+        """
+
+        super().__init__()
+        self.config = config
+        # 'discretized' indicates whether the Residual VQ part is trained or not. (After finishing the training, we set discretized=True)
+        self.register_buffer("discretized", torch.tensor(False))
+        self.optimized_steps = 0
+        # we use the fixed number of layers for Residual VQ across all environments.
+        self.vqvae_num_layers = 2
+
+        self.vq_layer = ResidualVQ(
+            dim=config.vqvae_embedding_dim,
+            num_quantizers=self.vqvae_num_layers,
+            codebook_size=config.vqvae_n_embed,
+        )
+
+        self.encoder = MLP(
+            in_channels=self.config.action_feature.shape[0] * self.config.action_chunk_size,
+            hidden_channels=[
+                config.vqvae_enc_hidden_dim,
+                config.vqvae_enc_hidden_dim,
+                config.vqvae_embedding_dim,
+            ],
+        )
+        self.decoder = MLP(
+            in_channels=config.vqvae_embedding_dim,
+            hidden_channels=[
+                config.vqvae_enc_hidden_dim,
+                config.vqvae_enc_hidden_dim,
+                self.config.action_feature.shape[0] * self.config.action_chunk_size,
+            ],
+        )
+
+    def get_embeddings_from_code(self, encoding_indices):
+        # This function gets code indices as inputs, and outputs embedding vectors corresponding to the code indices.
+        with torch.no_grad():
+            z_embed = self.vq_layer.get_codebook_vector_from_indices(encoding_indices)
+            # since the RVQ has multiple layers, it adds the vectors in the axis of layers to provide a vector for that code combination.
+            z_embed = z_embed.sum(dim=0)
+        return z_embed
+
+    def get_action_from_latent(self, latent):
+        # given latent vector, this function outputs the decoded action.
+        output = self.decoder(latent)
+        if self.config.action_chunk_size == 1:
+            return einops.rearrange(output, "N (T A) -> N T A", A=self.config.action_feature.shape[0])
+        else:
+            return einops.rearrange(output, "N (T A) -> N T A", A=self.config.action_feature.shape[0])
+
+    def get_code(self, state):
+        # in phase 2 of VQ-BeT training, we need a `ground truth labels of action data` to calculate the Focal loss for code prediction head. (please refer to section 3.3 in the paper https://arxiv.org/pdf/2403.03181)
+        # this function outputs the `GT code` of given action using frozen encoder and quantization layers. (please refer to Figure 2. in the paper https://arxiv.org/pdf/2403.03181)
+        state = einops.rearrange(state, "N T A -> N (T A)")
+        with torch.no_grad():
+            state_rep = self.encoder(state)
+            state_rep_shape = state_rep.shape[:-1]
+            state_rep_flat = state_rep.view(state_rep.size(0), -1, state_rep.size(1))
+            state_rep_flat, vq_code, vq_loss_state = self.vq_layer(state_rep_flat)
+            state_vq = state_rep_flat.view(*state_rep_shape, -1)
+            vq_code = vq_code.view(*state_rep_shape, -1)
+            vq_loss_state = torch.sum(vq_loss_state)
+            return state_vq, vq_code
+
+    def vqvae_forward(self, state):
+        # This function passes the given data through Residual VQ with Encoder and Decoder. Please refer to section 3.2 in the paper https://arxiv.org/pdf/2403.03181).
+        state = einops.rearrange(state, "N T A -> N (T A)")
+        # We start with passing action (or action chunk) at:t+n through the encoder ϕ.
+        state_rep = self.encoder(state)
+        state_rep_shape = state_rep.shape[:-1]
+        state_rep_flat = state_rep.view(state_rep.size(0), -1, state_rep.size(1))
+        # The resulting latent embedding vector x = ϕ(at:t+n) is then mapped to an embedding vector in the codebook of the RVQ layers by the nearest neighbor look-up.
+        state_rep_flat, vq_code, vq_loss_state = self.vq_layer(state_rep_flat)
+        state_vq = state_rep_flat.view(*state_rep_shape, -1)
+        vq_code = vq_code.view(*state_rep_shape, -1)
+        # since the RVQ has multiple layers, it adds the vectors in the axis of layers to provide a vector for that code combination.
+        vq_loss_state = torch.sum(vq_loss_state)
+        # Then, the discretized vector zq(x) is reconstructed as ψ(zq(x)) by passing through the decoder ψ.
+        dec_out = self.decoder(state_vq)
+        # Calculate L1 reconstruction loss
+        encoder_loss = (state - dec_out).abs().mean()
+        # add encoder reconstruction loss and commitment loss
+        rep_loss = encoder_loss + vq_loss_state * 5
+
+        metric = (
+            encoder_loss.clone().detach(),
+            vq_loss_state.clone().detach(),
+            vq_code,
+            rep_loss.item(),
+        )
+        return rep_loss, metric
+
+
+class FocalLoss(nn.Module):
+    """
+    From https://github.com/notmahi/miniBET/blob/main/behavior_transformer/bet.py
+    """
+
+    def __init__(self, gamma: float = 0, size_average: bool = True):
+        super().__init__()
+        self.gamma = gamma
+        self.size_average = size_average
+
+    def forward(self, input, target):
+        if len(input.shape) == 3:
+            N, T, _ = input.shape
+            logpt = F.log_softmax(input, dim=-1)
+            logpt = logpt.gather(-1, target.view(N, T, 1)).view(N, T)
+        elif len(input.shape) == 2:
+            logpt = F.log_softmax(input, dim=-1)
+            logpt = logpt.gather(-1, target.view(-1, 1)).view(-1)
+        pt = logpt.exp()
+
+        loss = -1 * (1 - pt) ** self.gamma * logpt
+        if self.size_average:
+            return loss.mean()
+        else:
+            return loss.sum()
+
+
+class MLP(torch.nn.Sequential):
+    def __init__(
+        self,
+        in_channels: int,
+        hidden_channels: List[int],
+    ):
+        layers = []
+        in_dim = in_channels
+        for hidden_dim in hidden_channels[:-1]:
+            layers.append(torch.nn.Linear(in_dim, hidden_dim))
+            layers.append(torch.nn.ReLU())
+            in_dim = hidden_dim
+
+        layers.append(torch.nn.Linear(in_dim, hidden_channels[-1]))
+
+        super().__init__(*layers)

lerobot/common/robot_devices/cameras/configs.py

@@ -0,0 +1,114 @@
+# 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 abc
+from dataclasses import dataclass
+
+import draccus
+
+
+@dataclass
+class CameraConfig(draccus.ChoiceRegistry, abc.ABC):
+    @property
+    def type(self) -> str:
+        return self.get_choice_name(self.__class__)
+
+
+@CameraConfig.register_subclass("opencv")
+@dataclass
+class OpenCVCameraConfig(CameraConfig):
+    """
+    Example of tested options for Intel Real Sense D405:
+
+    ```python
+    OpenCVCameraConfig(0, 30, 640, 480)
+    OpenCVCameraConfig(0, 60, 640, 480)
+    OpenCVCameraConfig(0, 90, 640, 480)
+    OpenCVCameraConfig(0, 30, 1280, 720)
+    ```
+    """
+
+    camera_index: int
+    fps: int | None = None
+    width: int | None = None
+    height: int | None = None
+    color_mode: str = "rgb"
+    channels: int | None = None
+    rotation: int | None = None
+    mock: bool = False
+
+    def __post_init__(self):
+        if self.color_mode not in ["rgb", "bgr"]:
+            raise ValueError(
+                f"`color_mode` is expected to be 'rgb' or 'bgr', but {self.color_mode} is provided."
+            )
+
+        self.channels = 3
+
+        if self.rotation not in [-90, None, 90, 180]:
+            raise ValueError(f"`rotation` must be in [-90, None, 90, 180] (got {self.rotation})")
+
+
+@CameraConfig.register_subclass("intelrealsense")
+@dataclass
+class IntelRealSenseCameraConfig(CameraConfig):
+    """
+    Example of tested options for Intel Real Sense D405:
+
+    ```python
+    IntelRealSenseCameraConfig(128422271347, 30, 640, 480)
+    IntelRealSenseCameraConfig(128422271347, 60, 640, 480)
+    IntelRealSenseCameraConfig(128422271347, 90, 640, 480)
+    IntelRealSenseCameraConfig(128422271347, 30, 1280, 720)
+    IntelRealSenseCameraConfig(128422271347, 30, 640, 480, use_depth=True)
+    IntelRealSenseCameraConfig(128422271347, 30, 640, 480, rotation=90)
+    ```
+    """
+
+    name: str | None = None
+    serial_number: int | None = None
+    fps: int | None = None
+    width: int | None = None
+    height: int | None = None
+    color_mode: str = "rgb"
+    channels: int | None = None
+    use_depth: bool = False
+    force_hardware_reset: bool = True
+    rotation: int | None = None
+    mock: bool = False
+
+    def __post_init__(self):
+        # bool is stronger than is None, since it works with empty strings
+        if bool(self.name) and bool(self.serial_number):
+            raise ValueError(
+                f"One of them must be set: name or serial_number, but {self.name=} and {self.serial_number=} provided."
+            )
+
+        if self.color_mode not in ["rgb", "bgr"]:
+            raise ValueError(
+                f"`color_mode` is expected to be 'rgb' or 'bgr', but {self.color_mode} is provided."
+            )
+
+        self.channels = 3
+
+        at_least_one_is_not_none = self.fps is not None or self.width is not None or self.height is not None
+        at_least_one_is_none = self.fps is None or self.width is None or self.height is None
+        if at_least_one_is_not_none and at_least_one_is_none:
+            raise ValueError(
+                "For `fps`, `width` and `height`, either all of them need to be set, or none of them, "
+                f"but {self.fps=}, {self.width=}, {self.height=} were provided."
+            )
+
+        if self.rotation not in [-90, None, 90, 180]:
+            raise ValueError(f"`rotation` must be in [-90, None, 90, 180] (got {self.rotation})")

lerobot/common/robot_devices/cameras/intelrealsense.py

@@ -0,0 +1,538 @@
+# 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.
+
+"""
+This file contains utilities for recording frames from Intel Realsense cameras.
+"""
+
+import argparse
+import concurrent.futures
+import logging
+import math
+import shutil
+import threading
+import time
+import traceback
+from collections import Counter
+from pathlib import Path
+from threading import Thread
+
+import numpy as np
+from PIL import Image
+
+from lerobot.common.robot_devices.cameras.configs import IntelRealSenseCameraConfig
+from lerobot.common.robot_devices.utils import (
+    RobotDeviceAlreadyConnectedError,
+    RobotDeviceNotConnectedError,
+    busy_wait,
+)
+from lerobot.common.utils.utils import capture_timestamp_utc
+
+SERIAL_NUMBER_INDEX = 1
+
+
+def find_cameras(raise_when_empty=True, mock=False) -> list[dict]:
+    """
+    Find the names and the serial numbers of the Intel RealSense cameras
+    connected to the computer.
+    """
+    if mock:
+        import tests.cameras.mock_pyrealsense2 as rs
+    else:
+        import pyrealsense2 as rs
+
+    cameras = []
+    for device in rs.context().query_devices():
+        serial_number = int(device.get_info(rs.camera_info(SERIAL_NUMBER_INDEX)))
+        name = device.get_info(rs.camera_info.name)
+        cameras.append(
+            {
+                "serial_number": serial_number,
+                "name": name,
+            }
+        )
+
+    if raise_when_empty and len(cameras) == 0:
+        raise OSError(
+            "Not a single camera was detected. Try re-plugging, or re-installing `librealsense` and its python wrapper `pyrealsense2`, or updating the firmware."
+        )
+
+    return cameras
+
+
+def save_image(img_array, serial_number, frame_index, images_dir):
+    try:
+        img = Image.fromarray(img_array)
+        path = images_dir / f"camera_{serial_number}_frame_{frame_index:06d}.png"
+        path.parent.mkdir(parents=True, exist_ok=True)
+        img.save(str(path), quality=100)
+        logging.info(f"Saved image: {path}")
+    except Exception as e:
+        logging.error(f"Failed to save image for camera {serial_number} frame {frame_index}: {e}")
+
+
+def save_images_from_cameras(
+    images_dir: Path,
+    serial_numbers: list[int] | None = None,
+    fps=None,
+    width=None,
+    height=None,
+    record_time_s=2,
+    mock=False,
+):
+    """
+    Initializes all the cameras and saves images to the directory. Useful to visually identify the camera
+    associated to a given serial number.
+    """
+    if serial_numbers is None or len(serial_numbers) == 0:
+        camera_infos = find_cameras(mock=mock)
+        serial_numbers = [cam["serial_number"] for cam in camera_infos]
+
+    if mock:
+        import tests.cameras.mock_cv2 as cv2
+    else:
+        import cv2
+
+    print("Connecting cameras")
+    cameras = []
+    for cam_sn in serial_numbers:
+        print(f"{cam_sn=}")
+        config = IntelRealSenseCameraConfig(
+            serial_number=cam_sn, fps=fps, width=width, height=height, mock=mock
+        )
+        camera = IntelRealSenseCamera(config)
+        camera.connect()
+        print(
+            f"IntelRealSenseCamera({camera.serial_number}, fps={camera.fps}, width={camera.capture_width}, height={camera.capture_height}, color_mode={camera.color_mode})"
+        )
+        cameras.append(camera)
+
+    images_dir = Path(images_dir)
+    if images_dir.exists():
+        shutil.rmtree(
+            images_dir,
+        )
+    images_dir.mkdir(parents=True, exist_ok=True)
+
+    print(f"Saving images to {images_dir}")
+    frame_index = 0
+    start_time = time.perf_counter()
+    try:
+        with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
+            while True:
+                now = time.perf_counter()
+
+                for camera in cameras:
+                    # If we use async_read when fps is None, the loop will go full speed, and we will end up
+                    # saving the same images from the cameras multiple times until the RAM/disk is full.
+                    image = camera.read() if fps is None else camera.async_read()
+                    if image is None:
+                        print("No Frame")
+
+                    bgr_converted_image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+
+                    executor.submit(
+                        save_image,
+                        bgr_converted_image,
+                        camera.serial_number,
+                        frame_index,
+                        images_dir,
+                    )
+
+                if fps is not None:
+                    dt_s = time.perf_counter() - now
+                    busy_wait(1 / fps - dt_s)
+
+                if time.perf_counter() - start_time > record_time_s:
+                    break
+
+                print(f"Frame: {frame_index:04d}\tLatency (ms): {(time.perf_counter() - now) * 1000:.2f}")
+
+                frame_index += 1
+    finally:
+        print(f"Images have been saved to {images_dir}")
+        for camera in cameras:
+            camera.disconnect()
+
+
+class IntelRealSenseCamera:
+    """
+    The IntelRealSenseCamera class is similar to OpenCVCamera class but adds additional features for Intel Real Sense cameras:
+    - is instantiated with the serial number of the camera - won't randomly change as it can be the case of OpenCVCamera for Linux,
+    - can also be instantiated with the camera's name — if it's unique — using IntelRealSenseCamera.init_from_name(),
+    - depth map can be returned.
+
+    To find the camera indices of your cameras, you can run our utility script that will save a few frames for each camera:
+    ```bash
+    python lerobot/common/robot_devices/cameras/intelrealsense.py --images-dir outputs/images_from_intelrealsense_cameras
+    ```
+
+    When an IntelRealSenseCamera is instantiated, if no specific config is provided, the default fps, width, height and color_mode
+    of the given camera will be used.
+
+    Example of instantiating with a serial number:
+    ```python
+    from lerobot.common.robot_devices.cameras.configs import IntelRealSenseCameraConfig
+
+    config = IntelRealSenseCameraConfig(serial_number=128422271347)
+    camera = IntelRealSenseCamera(config)
+    camera.connect()
+    color_image = camera.read()
+    # when done using the camera, consider disconnecting
+    camera.disconnect()
+    ```
+
+    Example of instantiating with a name if it's unique:
+    ```
+    config = IntelRealSenseCameraConfig(name="Intel RealSense D405")
+    ```
+
+    Example of changing default fps, width, height and color_mode:
+    ```python
+    config = IntelRealSenseCameraConfig(serial_number=128422271347, fps=30, width=1280, height=720)
+    config = IntelRealSenseCameraConfig(serial_number=128422271347, fps=90, width=640, height=480)
+    config = IntelRealSenseCameraConfig(serial_number=128422271347, fps=90, width=640, height=480, color_mode="bgr")
+    # Note: might error out upon `camera.connect()` if these settings are not compatible with the camera
+    ```
+
+    Example of returning depth:
+    ```python
+    config = IntelRealSenseCameraConfig(serial_number=128422271347, use_depth=True)
+    camera = IntelRealSenseCamera(config)
+    camera.connect()
+    color_image, depth_map = camera.read()
+    ```
+    """
+
+    def __init__(
+        self,
+        config: IntelRealSenseCameraConfig,
+    ):
+        self.config = config
+        if config.name is not None:
+            self.serial_number = self.find_serial_number_from_name(config.name)
+        else:
+            self.serial_number = config.serial_number
+
+        # Store the raw (capture) resolution from the config.
+        self.capture_width = config.width
+        self.capture_height = config.height
+
+        # If rotated by ±90, swap width and height.
+        if config.rotation in [-90, 90]:
+            self.width = config.height
+            self.height = config.width
+        else:
+            self.width = config.width
+            self.height = config.height
+
+        self.fps = config.fps
+        self.channels = config.channels
+        self.color_mode = config.color_mode
+        self.use_depth = config.use_depth
+        self.force_hardware_reset = config.force_hardware_reset
+        self.mock = config.mock
+
+        self.camera = None
+        self.is_connected = False
+        self.thread = None
+        self.stop_event = None
+        self.color_image = None
+        self.depth_map = None
+        self.logs = {}
+
+        if self.mock:
+            import tests.cameras.mock_cv2 as cv2
+        else:
+            import cv2
+
+        self.rotation = None
+        if config.rotation == -90:
+            self.rotation = cv2.ROTATE_90_COUNTERCLOCKWISE
+        elif config.rotation == 90:
+            self.rotation = cv2.ROTATE_90_CLOCKWISE
+        elif config.rotation == 180:
+            self.rotation = cv2.ROTATE_180
+
+    def find_serial_number_from_name(self, name):
+        camera_infos = find_cameras()
+        camera_names = [cam["name"] for cam in camera_infos]
+        this_name_count = Counter(camera_names)[name]
+        if this_name_count > 1:
+            # TODO(aliberts): Test this with multiple identical cameras (Aloha)
+            raise ValueError(
+                f"Multiple {name} cameras have been detected. Please use their serial number to instantiate them."
+            )
+
+        name_to_serial_dict = {cam["name"]: cam["serial_number"] for cam in camera_infos}
+        cam_sn = name_to_serial_dict[name]
+
+        return cam_sn
+
+    def connect(self):
+        if self.is_connected:
+            raise RobotDeviceAlreadyConnectedError(
+                f"IntelRealSenseCamera({self.serial_number}) is already connected."
+            )
+
+        if self.mock:
+            import tests.cameras.mock_pyrealsense2 as rs
+        else:
+            import pyrealsense2 as rs
+
+        config = rs.config()
+        config.enable_device(str(self.serial_number))
+
+        if self.fps and self.capture_width and self.capture_height:
+            # TODO(rcadene): can we set rgb8 directly?
+            config.enable_stream(
+                rs.stream.color, self.capture_width, self.capture_height, rs.format.rgb8, self.fps
+            )
+        else:
+            config.enable_stream(rs.stream.color)
+
+        if self.use_depth:
+            if self.fps and self.capture_width and self.capture_height:
+                config.enable_stream(
+                    rs.stream.depth, self.capture_width, self.capture_height, rs.format.z16, self.fps
+                )
+            else:
+                config.enable_stream(rs.stream.depth)
+
+        self.camera = rs.pipeline()
+        try:
+            profile = self.camera.start(config)
+            is_camera_open = True
+        except RuntimeError:
+            is_camera_open = False
+            traceback.print_exc()
+
+        # If the camera doesn't work, display the camera indices corresponding to
+        # valid cameras.
+        if not is_camera_open:
+            # Verify that the provided `serial_number` is valid before printing the traceback
+            camera_infos = find_cameras()
+            serial_numbers = [cam["serial_number"] for cam in camera_infos]
+            if self.serial_number not in serial_numbers:
+                raise ValueError(
+                    f"`serial_number` is expected to be one of these available cameras {serial_numbers}, but {self.serial_number} is provided instead. "
+                    "To find the serial number you should use, run `python lerobot/common/robot_devices/cameras/intelrealsense.py`."
+                )
+
+            raise OSError(f"Can't access IntelRealSenseCamera({self.serial_number}).")
+
+        color_stream = profile.get_stream(rs.stream.color)
+        color_profile = color_stream.as_video_stream_profile()
+        actual_fps = color_profile.fps()
+        actual_width = color_profile.width()
+        actual_height = color_profile.height()
+
+        # Using `math.isclose` since actual fps can be a float (e.g. 29.9 instead of 30)
+        if self.fps is not None and not math.isclose(self.fps, actual_fps, rel_tol=1e-3):
+            # Using `OSError` since it's a broad that encompasses issues related to device communication
+            raise OSError(
+                f"Can't set {self.fps=} for IntelRealSenseCamera({self.serial_number}). Actual value is {actual_fps}."
+            )
+        if self.capture_width is not None and self.capture_width != actual_width:
+            raise OSError(
+                f"Can't set {self.capture_width=} for IntelRealSenseCamera({self.serial_number}). Actual value is {actual_width}."
+            )
+        if self.capture_height is not None and self.capture_height != actual_height:
+            raise OSError(
+                f"Can't set {self.capture_height=} for IntelRealSenseCamera({self.serial_number}). Actual value is {actual_height}."
+            )
+
+        self.fps = round(actual_fps)
+        self.capture_width = round(actual_width)
+        self.capture_height = round(actual_height)
+
+        self.is_connected = True
+
+    def read(self, temporary_color: str | None = None) -> np.ndarray | tuple[np.ndarray, np.ndarray]:
+        """Read a frame from the camera returned in the format height x width x channels (e.g. 480 x 640 x 3)
+        of type `np.uint8`, contrarily to the pytorch format which is float channel first.
+
+        When `use_depth=True`, returns a tuple `(color_image, depth_map)` with a depth map in the format
+        height x width (e.g. 480 x 640) of type np.uint16.
+
+        Note: Reading a frame is done every `camera.fps` times per second, and it is blocking.
+        If you are reading data from other sensors, we advise to use `camera.async_read()` which is non blocking version of `camera.read()`.
+        """
+        if not self.is_connected:
+            raise RobotDeviceNotConnectedError(
+                f"IntelRealSenseCamera({self.serial_number}) is not connected. Try running `camera.connect()` first."
+            )
+
+        if self.mock:
+            import tests.cameras.mock_cv2 as cv2
+        else:
+            import cv2
+
+        start_time = time.perf_counter()
+
+        frame = self.camera.wait_for_frames(timeout_ms=5000)
+
+        color_frame = frame.get_color_frame()
+
+        if not color_frame:
+            raise OSError(f"Can't capture color image from IntelRealSenseCamera({self.serial_number}).")
+
+        color_image = np.asanyarray(color_frame.get_data())
+
+        requested_color_mode = self.color_mode if temporary_color is None else temporary_color
+        if requested_color_mode not in ["rgb", "bgr"]:
+            raise ValueError(
+                f"Expected color values are 'rgb' or 'bgr', but {requested_color_mode} is provided."
+            )
+
+        # IntelRealSense uses RGB format as default (red, green, blue).
+        if requested_color_mode == "bgr":
+            color_image = cv2.cvtColor(color_image, cv2.COLOR_RGB2BGR)
+
+        h, w, _ = color_image.shape
+        if h != self.capture_height or w != self.capture_width:
+            raise OSError(
+                f"Can't capture color image with expected height and width ({self.height} x {self.width}). ({h} x {w}) returned instead."
+            )
+
+        if self.rotation is not None:
+            color_image = cv2.rotate(color_image, self.rotation)
+
+        # log the number of seconds it took to read the image
+        self.logs["delta_timestamp_s"] = time.perf_counter() - start_time
+
+        # log the utc time at which the image was received
+        self.logs["timestamp_utc"] = capture_timestamp_utc()
+
+        if self.use_depth:
+            depth_frame = frame.get_depth_frame()
+            if not depth_frame:
+                raise OSError(f"Can't capture depth image from IntelRealSenseCamera({self.serial_number}).")
+
+            depth_map = np.asanyarray(depth_frame.get_data())
+
+            h, w = depth_map.shape
+            if h != self.capture_height or w != self.capture_width:
+                raise OSError(
+                    f"Can't capture depth map with expected height and width ({self.height} x {self.width}). ({h} x {w}) returned instead."
+                )
+
+            if self.rotation is not None:
+                depth_map = cv2.rotate(depth_map, self.rotation)
+
+            return color_image, depth_map
+        else:
+            return color_image
+
+    def read_loop(self):
+        while not self.stop_event.is_set():
+            if self.use_depth:
+                self.color_image, self.depth_map = self.read()
+            else:
+                self.color_image = self.read()
+
+    def async_read(self):
+        """Access the latest color image"""
+        if not self.is_connected:
+            raise RobotDeviceNotConnectedError(
+                f"IntelRealSenseCamera({self.serial_number}) is not connected. Try running `camera.connect()` first."
+            )
+
+        if self.thread is None:
+            self.stop_event = threading.Event()
+            self.thread = Thread(target=self.read_loop, args=())
+            self.thread.daemon = True
+            self.thread.start()
+
+        num_tries = 0
+        while self.color_image is None:
+            # TODO(rcadene, aliberts): intelrealsense has diverged compared to opencv over here
+            num_tries += 1
+            time.sleep(1 / self.fps)
+            if num_tries > self.fps and (self.thread.ident is None or not self.thread.is_alive()):
+                raise Exception(
+                    "The thread responsible for `self.async_read()` took too much time to start. There might be an issue. Verify that `self.thread.start()` has been called."
+                )
+
+        if self.use_depth:
+            return self.color_image, self.depth_map
+        else:
+            return self.color_image
+
+    def disconnect(self):
+        if not self.is_connected:
+            raise RobotDeviceNotConnectedError(
+                f"IntelRealSenseCamera({self.serial_number}) is not connected. Try running `camera.connect()` first."
+            )
+
+        if self.thread is not None and self.thread.is_alive():
+            # wait for the thread to finish
+            self.stop_event.set()
+            self.thread.join()
+            self.thread = None
+            self.stop_event = None
+
+        self.camera.stop()
+        self.camera = None
+
+        self.is_connected = False
+
+    def __del__(self):
+        if getattr(self, "is_connected", False):
+            self.disconnect()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Save a few frames using `IntelRealSenseCamera` for all cameras connected to the computer, or a selected subset."
+    )
+    parser.add_argument(
+        "--serial-numbers",
+        type=int,
+        nargs="*",
+        default=None,
+        help="List of serial numbers used to instantiate the `IntelRealSenseCamera`. If not provided, find and use all available camera indices.",
+    )
+    parser.add_argument(
+        "--fps",
+        type=int,
+        default=30,
+        help="Set the number of frames recorded per seconds for all cameras. If not provided, use the default fps of each camera.",
+    )
+    parser.add_argument(
+        "--width",
+        type=str,
+        default=640,
+        help="Set the width for all cameras. If not provided, use the default width of each camera.",
+    )
+    parser.add_argument(
+        "--height",
+        type=str,
+        default=480,
+        help="Set the height for all cameras. If not provided, use the default height of each camera.",
+    )
+    parser.add_argument(
+        "--images-dir",
+        type=Path,
+        default="outputs/images_from_intelrealsense_cameras",
+        help="Set directory to save a few frames for each camera.",
+    )
+    parser.add_argument(
+        "--record-time-s",
+        type=float,
+        default=2.0,
+        help="Set the number of seconds used to record the frames. By default, 2 seconds.",
+    )
+    args = parser.parse_args()
+    save_images_from_cameras(**vars(args))

lerobot/common/robot_devices/cameras/opencv.py

@@ -0,0 +1,518 @@
+# 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.
+
+"""
+This file contains utilities for recording frames from cameras. For more info look at `OpenCVCamera` docstring.
+"""
+
+import argparse
+import concurrent.futures
+import math
+import platform
+import shutil
+import threading
+import time
+from pathlib import Path
+from threading import Thread
+
+import numpy as np
+from PIL import Image
+
+from lerobot.common.robot_devices.cameras.configs import OpenCVCameraConfig
+from lerobot.common.robot_devices.utils import (
+    RobotDeviceAlreadyConnectedError,
+    RobotDeviceNotConnectedError,
+    busy_wait,
+)
+from lerobot.common.utils.utils import capture_timestamp_utc
+
+# The maximum opencv device index depends on your operating system. For instance,
+# if you have 3 cameras, they should be associated to index 0, 1, and 2. This is the case
+# on MacOS. However, on Ubuntu, the indices are different like 6, 16, 23.
+# When you change the USB port or reboot the computer, the operating system might
+# treat the same cameras as new devices. Thus we select a higher bound to search indices.
+MAX_OPENCV_INDEX = 60
+
+
+def find_cameras(raise_when_empty=False, max_index_search_range=MAX_OPENCV_INDEX, mock=False) -> list[dict]:
+    cameras = []
+    if platform.system() == "Linux":
+        print("Linux detected. Finding available camera indices through scanning '/dev/video*' ports")
+        possible_ports = [str(port) for port in Path("/dev").glob("video*")]
+        ports = _find_cameras(possible_ports, mock=mock)
+        for port in ports:
+            cameras.append(
+                {
+                    "port": port,
+                    "index": int(port.removeprefix("/dev/video")),
+                }
+            )
+    else:
+        print(
+            "Mac or Windows detected. Finding available camera indices through "
+            f"scanning all indices from 0 to {MAX_OPENCV_INDEX}"
+        )
+        possible_indices = range(max_index_search_range)
+        indices = _find_cameras(possible_indices, mock=mock)
+        for index in indices:
+            cameras.append(
+                {
+                    "port": None,
+                    "index": index,
+                }
+            )
+
+    return cameras
+
+
+def _find_cameras(
+    possible_camera_ids: list[int | str], raise_when_empty=False, mock=False
+) -> list[int | str]:
+    if mock:
+        import tests.cameras.mock_cv2 as cv2
+    else:
+        import cv2
+
+    camera_ids = []
+    for camera_idx in possible_camera_ids:
+        camera = cv2.VideoCapture(camera_idx)
+        is_open = camera.isOpened()
+        camera.release()
+
+        if is_open:
+            print(f"Camera found at index {camera_idx}")
+            camera_ids.append(camera_idx)
+
+    if raise_when_empty and len(camera_ids) == 0:
+        raise OSError(
+            "Not a single camera was detected. Try re-plugging, or re-installing `opencv2`, "
+            "or your camera driver, or make sure your camera is compatible with opencv2."
+        )
+
+    return camera_ids
+
+
+def is_valid_unix_path(path: str) -> bool:
+    """Note: if 'path' points to a symlink, this will return True only if the target exists"""
+    p = Path(path)
+    return p.is_absolute() and p.exists()
+
+
+def get_camera_index_from_unix_port(port: Path) -> int:
+    return int(str(port.resolve()).removeprefix("/dev/video"))
+
+
+def save_image(img_array, camera_index, frame_index, images_dir):
+    img = Image.fromarray(img_array)
+    path = images_dir / f"camera_{camera_index:02d}_frame_{frame_index:06d}.png"
+    path.parent.mkdir(parents=True, exist_ok=True)
+    img.save(str(path), quality=100)
+
+
+def save_images_from_cameras(
+    images_dir: Path,
+    camera_ids: list | None = None,
+    fps=None,
+    width=None,
+    height=None,
+    record_time_s=2,
+    mock=False,
+):
+    """
+    Initializes all the cameras and saves images to the directory. Useful to visually identify the camera
+    associated to a given camera index.
+    """
+    if camera_ids is None or len(camera_ids) == 0:
+        camera_infos = find_cameras(mock=mock)
+        camera_ids = [cam["index"] for cam in camera_infos]
+
+    print("Connecting cameras")
+    cameras = []
+    for cam_idx in camera_ids:
+        config = OpenCVCameraConfig(camera_index=cam_idx, fps=fps, width=width, height=height, mock=mock)
+        camera = OpenCVCamera(config)
+        camera.connect()
+        print(
+            f"OpenCVCamera({camera.camera_index}, fps={camera.fps}, width={camera.capture_width}, "
+            f"height={camera.capture_height}, color_mode={camera.color_mode})"
+        )
+        cameras.append(camera)
+
+    images_dir = Path(images_dir)
+    if images_dir.exists():
+        shutil.rmtree(
+            images_dir,
+        )
+    images_dir.mkdir(parents=True, exist_ok=True)
+
+    print(f"Saving images to {images_dir}")
+    frame_index = 0
+    start_time = time.perf_counter()
+    with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
+        while True:
+            now = time.perf_counter()
+
+            for camera in cameras:
+                # If we use async_read when fps is None, the loop will go full speed, and we will endup
+                # saving the same images from the cameras multiple times until the RAM/disk is full.
+                image = camera.read() if fps is None else camera.async_read()
+
+                executor.submit(
+                    save_image,
+                    image,
+                    camera.camera_index,
+                    frame_index,
+                    images_dir,
+                )
+
+            if fps is not None:
+                dt_s = time.perf_counter() - now
+                busy_wait(1 / fps - dt_s)
+
+            print(f"Frame: {frame_index:04d}\tLatency (ms): {(time.perf_counter() - now) * 1000:.2f}")
+
+            if time.perf_counter() - start_time > record_time_s:
+                break
+
+            frame_index += 1
+
+    print(f"Images have been saved to {images_dir}")
+
+
+class OpenCVCamera:
+    """
+    The OpenCVCamera class allows to efficiently record images from cameras. It relies on opencv2 to communicate
+    with the cameras. Most cameras are compatible. For more info, see the [Video I/O with OpenCV Overview](https://docs.opencv.org/4.x/d0/da7/videoio_overview.html).
+
+    An OpenCVCamera instance requires a camera index (e.g. `OpenCVCamera(camera_index=0)`). When you only have one camera
+    like a webcam of a laptop, the camera index is expected to be 0, but it might also be very different, and the camera index
+    might change if you reboot your computer or re-plug your camera. This behavior depends on your operation system.
+
+    To find the camera indices of your cameras, you can run our utility script that will be save a few frames for each camera:
+    ```bash
+    python lerobot/common/robot_devices/cameras/opencv.py --images-dir outputs/images_from_opencv_cameras
+    ```
+
+    When an OpenCVCamera is instantiated, if no specific config is provided, the default fps, width, height and color_mode
+    of the given camera will be used.
+
+    Example of usage:
+    ```python
+    from lerobot.common.robot_devices.cameras.configs import OpenCVCameraConfig
+
+    config = OpenCVCameraConfig(camera_index=0)
+    camera = OpenCVCamera(config)
+    camera.connect()
+    color_image = camera.read()
+    # when done using the camera, consider disconnecting
+    camera.disconnect()
+    ```
+
+    Example of changing default fps, width, height and color_mode:
+    ```python
+    config = OpenCVCameraConfig(camera_index=0, fps=30, width=1280, height=720)
+    config = OpenCVCameraConfig(camera_index=0, fps=90, width=640, height=480)
+    config = OpenCVCameraConfig(camera_index=0, fps=90, width=640, height=480, color_mode="bgr")
+    # Note: might error out open `camera.connect()` if these settings are not compatible with the camera
+    ```
+    """
+
+    def __init__(self, config: OpenCVCameraConfig):
+        self.config = config
+        self.camera_index = config.camera_index
+        self.port = None
+
+        # Linux uses ports for connecting to cameras
+        if platform.system() == "Linux":
+            if isinstance(self.camera_index, int):
+                self.port = Path(f"/dev/video{self.camera_index}")
+            elif isinstance(self.camera_index, str) and is_valid_unix_path(self.camera_index):
+                self.port = Path(self.camera_index)
+                # Retrieve the camera index from a potentially symlinked path
+                self.camera_index = get_camera_index_from_unix_port(self.port)
+            else:
+                raise ValueError(f"Please check the provided camera_index: {self.camera_index}")
+
+        # Store the raw (capture) resolution from the config.
+        self.capture_width = config.width
+        self.capture_height = config.height
+
+        # If rotated by ±90, swap width and height.
+        if config.rotation in [-90, 90]:
+            self.width = config.height
+            self.height = config.width
+        else:
+            self.width = config.width
+            self.height = config.height
+
+        self.fps = config.fps
+        self.channels = config.channels
+        self.color_mode = config.color_mode
+        self.mock = config.mock
+
+        self.camera = None
+        self.is_connected = False
+        self.thread = None
+        self.stop_event = None
+        self.color_image = None
+        self.logs = {}
+
+        if self.mock:
+            import tests.cameras.mock_cv2 as cv2
+        else:
+            import cv2
+
+        self.rotation = None
+        if config.rotation == -90:
+            self.rotation = cv2.ROTATE_90_COUNTERCLOCKWISE
+        elif config.rotation == 90:
+            self.rotation = cv2.ROTATE_90_CLOCKWISE
+        elif config.rotation == 180:
+            self.rotation = cv2.ROTATE_180
+
+    def connect(self):
+        if self.is_connected:
+            raise RobotDeviceAlreadyConnectedError(f"OpenCVCamera({self.camera_index}) is already connected.")
+
+        if self.mock:
+            import tests.cameras.mock_cv2 as cv2
+        else:
+            import cv2
+
+            # Use 1 thread to avoid blocking the main thread. Especially useful during data collection
+            # when other threads are used to save the images.
+            cv2.setNumThreads(1)
+
+        backend = (
+            cv2.CAP_V4L2
+            if platform.system() == "Linux"
+            else cv2.CAP_DSHOW
+            if platform.system() == "Windows"
+            else cv2.CAP_AVFOUNDATION
+            if platform.system() == "Darwin"
+            else cv2.CAP_ANY
+        )
+
+        camera_idx = f"/dev/video{self.camera_index}" if platform.system() == "Linux" else self.camera_index
+        # First create a temporary camera trying to access `camera_index`,
+        # and verify it is a valid camera by calling `isOpened`.
+        tmp_camera = cv2.VideoCapture(camera_idx, backend)
+        is_camera_open = tmp_camera.isOpened()
+        # Release camera to make it accessible for `find_camera_indices`
+        tmp_camera.release()
+        del tmp_camera
+
+        # If the camera doesn't work, display the camera indices corresponding to
+        # valid cameras.
+        if not is_camera_open:
+            # Verify that the provided `camera_index` is valid before printing the traceback
+            cameras_info = find_cameras()
+            available_cam_ids = [cam["index"] for cam in cameras_info]
+            if self.camera_index not in available_cam_ids:
+                raise ValueError(
+                    f"`camera_index` is expected to be one of these available cameras {available_cam_ids}, but {self.camera_index} is provided instead. "
+                    "To find the camera index you should use, run `python lerobot/common/robot_devices/cameras/opencv.py`."
+                )
+
+            raise OSError(f"Can't access OpenCVCamera({camera_idx}).")
+
+        # Secondly, create the camera that will be used downstream.
+        # Note: For some unknown reason, calling `isOpened` blocks the camera which then
+        # needs to be re-created.
+        self.camera = cv2.VideoCapture(camera_idx, backend)
+
+        if self.fps is not None:
+            self.camera.set(cv2.CAP_PROP_FPS, self.fps)
+        if self.capture_width is not None:
+            self.camera.set(cv2.CAP_PROP_FRAME_WIDTH, self.capture_width)
+        if self.capture_height is not None:
+            self.camera.set(cv2.CAP_PROP_FRAME_HEIGHT, self.capture_height)
+
+        actual_fps = self.camera.get(cv2.CAP_PROP_FPS)
+        actual_width = self.camera.get(cv2.CAP_PROP_FRAME_WIDTH)
+        actual_height = self.camera.get(cv2.CAP_PROP_FRAME_HEIGHT)
+
+        # Using `math.isclose` since actual fps can be a float (e.g. 29.9 instead of 30)
+        if self.fps is not None and not math.isclose(self.fps, actual_fps, rel_tol=1e-3):
+            # Using `OSError` since it's a broad that encompasses issues related to device communication
+            raise OSError(
+                f"Can't set {self.fps=} for OpenCVCamera({self.camera_index}). Actual value is {actual_fps}."
+            )
+        if self.capture_width is not None and not math.isclose(
+            self.capture_width, actual_width, rel_tol=1e-3
+        ):
+            raise OSError(
+                f"Can't set {self.capture_width=} for OpenCVCamera({self.camera_index}). Actual value is {actual_width}."
+            )
+        if self.capture_height is not None and not math.isclose(
+            self.capture_height, actual_height, rel_tol=1e-3
+        ):
+            raise OSError(
+                f"Can't set {self.capture_height=} for OpenCVCamera({self.camera_index}). Actual value is {actual_height}."
+            )
+
+        self.fps = round(actual_fps)
+        self.capture_width = round(actual_width)
+        self.capture_height = round(actual_height)
+        self.is_connected = True
+
+    def read(self, temporary_color_mode: str | None = None) -> np.ndarray:
+        """Read a frame from the camera returned in the format (height, width, channels)
+        (e.g. 480 x 640 x 3), contrarily to the pytorch format which is channel first.
+
+        Note: Reading a frame is done every `camera.fps` times per second, and it is blocking.
+        If you are reading data from other sensors, we advise to use `camera.async_read()` which is non blocking version of `camera.read()`.
+        """
+        if not self.is_connected:
+            raise RobotDeviceNotConnectedError(
+                f"OpenCVCamera({self.camera_index}) is not connected. Try running `camera.connect()` first."
+            )
+
+        start_time = time.perf_counter()
+
+        ret, color_image = self.camera.read()
+
+        if not ret:
+            raise OSError(f"Can't capture color image from camera {self.camera_index}.")
+
+        requested_color_mode = self.color_mode if temporary_color_mode is None else temporary_color_mode
+
+        if requested_color_mode not in ["rgb", "bgr"]:
+            raise ValueError(
+                f"Expected color values are 'rgb' or 'bgr', but {requested_color_mode} is provided."
+            )
+
+        # OpenCV uses BGR format as default (blue, green, red) for all operations, including displaying images.
+        # However, Deep Learning framework such as LeRobot uses RGB format as default to train neural networks,
+        # so we convert the image color from BGR to RGB.
+        if requested_color_mode == "rgb":
+            if self.mock:
+                import tests.cameras.mock_cv2 as cv2
+            else:
+                import cv2
+
+            color_image = cv2.cvtColor(color_image, cv2.COLOR_BGR2RGB)
+
+        h, w, _ = color_image.shape
+        if h != self.capture_height or w != self.capture_width:
+            raise OSError(
+                f"Can't capture color image with expected height and width ({self.height} x {self.width}). ({h} x {w}) returned instead."
+            )
+
+        if self.rotation is not None:
+            color_image = cv2.rotate(color_image, self.rotation)
+
+        # log the number of seconds it took to read the image
+        self.logs["delta_timestamp_s"] = time.perf_counter() - start_time
+
+        # log the utc time at which the image was received
+        self.logs["timestamp_utc"] = capture_timestamp_utc()
+
+        self.color_image = color_image
+
+        return color_image
+
+    def read_loop(self):
+        while not self.stop_event.is_set():
+            try:
+                self.color_image = self.read()
+            except Exception as e:
+                print(f"Error reading in thread: {e}")
+
+    def async_read(self):
+        if not self.is_connected:
+            raise RobotDeviceNotConnectedError(
+                f"OpenCVCamera({self.camera_index}) is not connected. Try running `camera.connect()` first."
+            )
+
+        if self.thread is None:
+            self.stop_event = threading.Event()
+            self.thread = Thread(target=self.read_loop, args=())
+            self.thread.daemon = True
+            self.thread.start()
+
+        num_tries = 0
+        while True:
+            if self.color_image is not None:
+                return self.color_image
+
+            time.sleep(1 / self.fps)
+            num_tries += 1
+            if num_tries > self.fps * 2:
+                raise TimeoutError("Timed out waiting for async_read() to start.")
+
+    def disconnect(self):
+        if not self.is_connected:
+            raise RobotDeviceNotConnectedError(
+                f"OpenCVCamera({self.camera_index}) is not connected. Try running `camera.connect()` first."
+            )
+
+        if self.thread is not None:
+            self.stop_event.set()
+            self.thread.join()  # wait for the thread to finish
+            self.thread = None
+            self.stop_event = None
+
+        self.camera.release()
+        self.camera = None
+        self.is_connected = False
+
+    def __del__(self):
+        if getattr(self, "is_connected", False):
+            self.disconnect()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Save a few frames using `OpenCVCamera` for all cameras connected to the computer, or a selected subset."
+    )
+    parser.add_argument(
+        "--camera-ids",
+        type=int,
+        nargs="*",
+        default=None,
+        help="List of camera indices used to instantiate the `OpenCVCamera`. If not provided, find and use all available camera indices.",
+    )
+    parser.add_argument(
+        "--fps",
+        type=int,
+        default=None,
+        help="Set the number of frames recorded per seconds for all cameras. If not provided, use the default fps of each camera.",
+    )
+    parser.add_argument(
+        "--width",
+        type=str,
+        default=None,
+        help="Set the width for all cameras. If not provided, use the default width of each camera.",
+    )
+    parser.add_argument(
+        "--height",
+        type=str,
+        default=None,
+        help="Set the height for all cameras. If not provided, use the default height of each camera.",
+    )
+    parser.add_argument(
+        "--images-dir",
+        type=Path,
+        default="outputs/images_from_opencv_cameras",
+        help="Set directory to save a few frames for each camera.",
+    )
+    parser.add_argument(
+        "--record-time-s",
+        type=float,
+        default=4.0,
+        help="Set the number of seconds used to record the frames. By default, 2 seconds.",
+    )
+    args = parser.parse_args()
+    save_images_from_cameras(**vars(args))

lerobot/common/robot_devices/cameras/utils.py

@@ -0,0 +1,67 @@
+# 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 typing import Protocol
+
+import numpy as np
+
+from lerobot.common.robot_devices.cameras.configs import (
+    CameraConfig,
+    IntelRealSenseCameraConfig,
+    OpenCVCameraConfig,
+)
+
+
+# Defines a camera type
+class Camera(Protocol):
+    def connect(self): ...
+    def read(self, temporary_color: str | None = None) -> np.ndarray: ...
+    def async_read(self) -> np.ndarray: ...
+    def disconnect(self): ...
+
+
+def make_cameras_from_configs(camera_configs: dict[str, CameraConfig]) -> list[Camera]:
+    cameras = {}
+
+    for key, cfg in camera_configs.items():
+        if cfg.type == "opencv":
+            from lerobot.common.robot_devices.cameras.opencv import OpenCVCamera
+
+            cameras[key] = OpenCVCamera(cfg)
+
+        elif cfg.type == "intelrealsense":
+            from lerobot.common.robot_devices.cameras.intelrealsense import IntelRealSenseCamera
+
+            cameras[key] = IntelRealSenseCamera(cfg)
+        else:
+            raise ValueError(f"The camera type '{cfg.type}' is not valid.")
+
+    return cameras
+
+
+def make_camera(camera_type, **kwargs) -> Camera:
+    if camera_type == "opencv":
+        from lerobot.common.robot_devices.cameras.opencv import OpenCVCamera
+
+        config = OpenCVCameraConfig(**kwargs)
+        return OpenCVCamera(config)
+
+    elif camera_type == "intelrealsense":
+        from lerobot.common.robot_devices.cameras.intelrealsense import IntelRealSenseCamera
+
+        config = IntelRealSenseCameraConfig(**kwargs)
+        return IntelRealSenseCamera(config)
+
+    else:
+        raise ValueError(f"The camera type '{camera_type}' is not valid.")

lerobot/common/robot_devices/control_configs.py

@@ -0,0 +1,129 @@
+# 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 dataclasses import dataclass
+from pathlib import Path
+
+import draccus
+
+from lerobot.common.robot_devices.robots.configs import RobotConfig
+from lerobot.configs import parser
+from lerobot.configs.policies import PreTrainedConfig
+
+
+@dataclass
+class ControlConfig(draccus.ChoiceRegistry):
+    pass
+
+
+@ControlConfig.register_subclass("calibrate")
+@dataclass
+class CalibrateControlConfig(ControlConfig):
+    # List of arms to calibrate (e.g. `--arms='["left_follower","right_follower"]' left_leader`)
+    arms: list[str] | None = None
+
+
+@ControlConfig.register_subclass("teleoperate")
+@dataclass
+class TeleoperateControlConfig(ControlConfig):
+    # Limit the maximum frames per second. By default, no limit.
+    fps: int | None = None
+    teleop_time_s: float | None = None
+    # Display all cameras on screen
+    display_cameras: bool = True
+
+
+@ControlConfig.register_subclass("record")
+@dataclass
+class RecordControlConfig(ControlConfig):
+    # Dataset identifier. By convention it should match '{hf_username}/{dataset_name}' (e.g. `lerobot/test`).
+    repo_id: str
+    # A short but accurate description of the task performed during the recording (e.g. "Pick the Lego block and drop it in the box on the right.")
+    single_task: str
+    # Root directory where the dataset will be stored (e.g. 'dataset/path').
+    root: str | Path | None = None
+    policy: PreTrainedConfig | None = None
+    # Limit the frames per second. By default, uses the policy fps.
+    fps: int | None = None
+    # Number of seconds before starting data collection. It allows the robot devices to warmup and synchronize.
+    warmup_time_s: int | float = 10
+    # Number of seconds for data recording for each episode.
+    episode_time_s: int | float = 60
+    # Number of seconds for resetting the environment after each episode.
+    reset_time_s: int | float = 60
+    # Number of episodes to record.
+    num_episodes: int = 50
+    # Encode frames in the dataset into video
+    video: bool = True
+    # Upload dataset to Hugging Face hub.
+    push_to_hub: bool = True
+    # Upload on private repository on the Hugging Face hub.
+    private: bool = False
+    # Add tags to your dataset on the hub.
+    tags: list[str] | None = None
+    # Number of subprocesses handling the saving of frames as PNG. Set to 0 to use threads only;
+    # set to ≥1 to use subprocesses, each using threads to write images. The best number of processes
+    # and threads depends on your system. We recommend 4 threads per camera with 0 processes.
+    # If fps is unstable, adjust the thread count. If still unstable, try using 1 or more subprocesses.
+    num_image_writer_processes: int = 0
+    # Number of threads writing the frames as png images on disk, per camera.
+    # Too many threads might cause unstable teleoperation fps due to main thread being blocked.
+    # Not enough threads might cause low camera fps.
+    num_image_writer_threads_per_camera: int = 4
+    # Display all cameras on screen
+    display_cameras: bool = True
+    # Use vocal synthesis to read events.
+    play_sounds: bool = True
+    # Resume recording on an existing dataset.
+    resume: bool = False
+
+    def __post_init__(self):
+        # HACK: We parse again the cli args here to get the pretrained path if there was one.
+        policy_path = parser.get_path_arg("control.policy")
+        if policy_path:
+            cli_overrides = parser.get_cli_overrides("control.policy")
+            self.policy = PreTrainedConfig.from_pretrained(policy_path, cli_overrides=cli_overrides)
+            self.policy.pretrained_path = policy_path
+
+
+@ControlConfig.register_subclass("replay")
+@dataclass
+class ReplayControlConfig(ControlConfig):
+    # Dataset identifier. By convention it should match '{hf_username}/{dataset_name}' (e.g. `lerobot/test`).
+    repo_id: str
+    # Index of the episode to replay.
+    episode: int
+    # Root directory where the dataset will be stored (e.g. 'dataset/path').
+    root: str | Path | None = None
+    # Limit the frames per second. By default, uses the dataset fps.
+    fps: int | None = None
+    # Use vocal synthesis to read events.
+    play_sounds: bool = True
+
+
+@ControlConfig.register_subclass("remote_robot")
+@dataclass
+class RemoteRobotConfig(ControlConfig):
+    log_interval: int = 100
+
+
+@dataclass
+class ControlPipelineConfig:
+    robot: RobotConfig
+    control: ControlConfig
+
+    @classmethod
+    def __get_path_fields__(cls) -> list[str]:
+        """This enables the parser to load config from the policy using `--policy.path=local/dir`"""
+        return ["control.policy"]

lerobot/common/robot_devices/control_utils.py

@@ -0,0 +1,347 @@
+# 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.
+
+########################################################################################
+# Utilities
+########################################################################################
+
+
+import logging
+import time
+import traceback
+from contextlib import nullcontext
+from copy import copy
+from functools import cache
+
+import cv2
+import torch
+from deepdiff import DeepDiff
+from termcolor import colored
+
+from lerobot.common.datasets.image_writer import safe_stop_image_writer
+from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.common.datasets.utils import get_features_from_robot
+from lerobot.common.policies.pretrained import PreTrainedPolicy
+from lerobot.common.robot_devices.robots.utils import Robot
+from lerobot.common.robot_devices.utils import busy_wait
+from lerobot.common.utils.utils import get_safe_torch_device, has_method
+
+
+def log_control_info(robot: Robot, dt_s, episode_index=None, frame_index=None, fps=None):
+    log_items = []
+    if episode_index is not None:
+        log_items.append(f"ep:{episode_index}")
+    if frame_index is not None:
+        log_items.append(f"frame:{frame_index}")
+
+    def log_dt(shortname, dt_val_s):
+        nonlocal log_items, fps
+        info_str = f"{shortname}:{dt_val_s * 1000:5.2f} ({1 / dt_val_s:3.1f}hz)"
+        if fps is not None:
+            actual_fps = 1 / dt_val_s
+            if actual_fps < fps - 1:
+                info_str = colored(info_str, "yellow")
+        log_items.append(info_str)
+
+    # total step time displayed in milliseconds and its frequency
+    log_dt("dt", dt_s)
+
+    # TODO(aliberts): move robot-specific logs logic in robot.print_logs()
+    if not robot.robot_type.startswith("stretch"):
+        for name in robot.leader_arms:
+            key = f"read_leader_{name}_pos_dt_s"
+            if key in robot.logs:
+                log_dt("dtRlead", robot.logs[key])
+
+        for name in robot.follower_arms:
+            key = f"write_follower_{name}_goal_pos_dt_s"
+            if key in robot.logs:
+                log_dt("dtWfoll", robot.logs[key])
+
+            key = f"read_follower_{name}_pos_dt_s"
+            if key in robot.logs:
+                log_dt("dtRfoll", robot.logs[key])
+
+        for name in robot.cameras:
+            key = f"read_camera_{name}_dt_s"
+            if key in robot.logs:
+                log_dt(f"dtR{name}", robot.logs[key])
+
+    info_str = " ".join(log_items)
+    logging.info(info_str)
+
+
+@cache
+def is_headless():
+    """Detects if python is running without a monitor."""
+    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, policy, device, use_amp):
+    observation = copy(observation)
+    with (
+        torch.inference_mode(),
+        torch.autocast(device_type=device.type) if device.type == "cuda" and use_amp else nullcontext(),
+    ):
+        # Convert to pytorch format: channel first and float32 in [0,1] with batch dimension
+        for name in observation:
+            if "image" in name:
+                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)
+
+        # Compute the next action with the policy
+        # based on the current observation
+        action = policy.select_action(observation)
+
+        # Remove batch dimension
+        action = action.squeeze(0)
+
+        # Move to cpu, if not already the case
+        action = action.to("cpu")
+
+    return action
+
+
+def init_keyboard_listener():
+    # 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 warmup_record(
+    robot,
+    events,
+    enable_teleoperation,
+    warmup_time_s,
+    display_cameras,
+    fps,
+):
+    control_loop(
+        robot=robot,
+        control_time_s=warmup_time_s,
+        display_cameras=display_cameras,
+        events=events,
+        fps=fps,
+        teleoperate=enable_teleoperation,
+    )
+
+
+def record_episode(
+    robot,
+    dataset,
+    events,
+    episode_time_s,
+    display_cameras,
+    policy,
+    fps,
+    single_task,
+):
+    control_loop(
+        robot=robot,
+        control_time_s=episode_time_s,
+        display_cameras=display_cameras,
+        dataset=dataset,
+        events=events,
+        policy=policy,
+        fps=fps,
+        teleoperate=policy is None,
+        single_task=single_task,
+    )
+
+
+@safe_stop_image_writer
+def control_loop(
+    robot,
+    control_time_s=None,
+    teleoperate=False,
+    display_cameras=False,
+    dataset: LeRobotDataset | None = None,
+    events=None,
+    policy: PreTrainedPolicy = None,
+    fps: int | None = None,
+    single_task: str | None = None,
+):
+    # TODO(rcadene): Add option to record logs
+    if not robot.is_connected:
+        robot.connect()
+
+    if events is None:
+        events = {"exit_early": False}
+
+    if control_time_s is None:
+        control_time_s = float("inf")
+
+    if teleoperate and policy is not None:
+        raise ValueError("When `teleoperate` is True, `policy` should be None.")
+
+    if dataset is not None and single_task is None:
+        raise ValueError("You need to provide a task as argument in `single_task`.")
+
+    if dataset is not None and fps is not None and dataset.fps != fps:
+        raise ValueError(f"The dataset fps should be equal to requested fps ({dataset['fps']} != {fps}).")
+
+    timestamp = 0
+    start_episode_t = time.perf_counter()
+    while timestamp < control_time_s:
+        start_loop_t = time.perf_counter()
+
+        if teleoperate:
+            observation, action = robot.teleop_step(record_data=True)
+        else:
+            observation = robot.capture_observation()
+
+            if policy is not None:
+                pred_action = predict_action(
+                    observation, policy, get_safe_torch_device(policy.config.device), policy.config.use_amp
+                )
+                # Action can eventually be clipped using `max_relative_target`,
+                # so action actually sent is saved in the dataset.
+                action = robot.send_action(pred_action)
+                action = {"action": action}
+
+        if dataset is not None:
+            frame = {**observation, **action, "task": single_task}
+            dataset.add_frame(frame)
+
+        if display_cameras and not is_headless():
+            image_keys = [key for key in observation if "image" in key]
+            for key in image_keys:
+                cv2.imshow(key, cv2.cvtColor(observation[key].numpy(), cv2.COLOR_RGB2BGR))
+            cv2.waitKey(1)
+
+        if fps is not None:
+            dt_s = time.perf_counter() - start_loop_t
+            busy_wait(1 / fps - dt_s)
+
+        dt_s = time.perf_counter() - start_loop_t
+        log_control_info(robot, dt_s, fps=fps)
+
+        timestamp = time.perf_counter() - start_episode_t
+        if events["exit_early"]:
+            events["exit_early"] = False
+            break
+
+
+def reset_environment(robot, events, reset_time_s, fps):
+    # TODO(rcadene): refactor warmup_record and reset_environment
+    if has_method(robot, "teleop_safety_stop"):
+        robot.teleop_safety_stop()
+
+    control_loop(
+        robot=robot,
+        control_time_s=reset_time_s,
+        events=events,
+        fps=fps,
+        teleoperate=True,
+    )
+
+
+def stop_recording(robot, listener, display_cameras):
+    robot.disconnect()
+
+    if not is_headless():
+        if listener is not None:
+            listener.stop()
+
+        if display_cameras:
+            cv2.destroyAllWindows()
+
+
+def sanity_check_dataset_name(repo_id, policy_cfg):
+    _, dataset_name = repo_id.split("/")
+    # either repo_id doesnt start with "eval_" and there is no policy
+    # or repo_id starts with "eval_" and there is a policy
+
+    # Check if dataset_name starts with "eval_" but policy is missing
+    if dataset_name.startswith("eval_") and policy_cfg is None:
+        raise ValueError(
+            f"Your dataset name begins with 'eval_' ({dataset_name}), but no policy is provided ({policy_cfg.type})."
+        )
+
+    # Check if dataset_name does not start with "eval_" but policy is provided
+    if not dataset_name.startswith("eval_") and policy_cfg is not None:
+        raise ValueError(
+            f"Your dataset name does not begin with 'eval_' ({dataset_name}), but a policy is provided ({policy_cfg.type})."
+        )
+
+
+def sanity_check_dataset_robot_compatibility(
+    dataset: LeRobotDataset, robot: Robot, fps: int, use_videos: bool
+) -> None:
+    fields = [
+        ("robot_type", dataset.meta.robot_type, robot.robot_type),
+        ("fps", dataset.fps, fps),
+        ("features", dataset.features, get_features_from_robot(robot, use_videos)),
+    ]
+
+    mismatches = []
+    for field, dataset_value, present_value in fields:
+        diff = DeepDiff(dataset_value, present_value, exclude_regex_paths=[r".*\['info'\]$"])
+        if diff:
+            mismatches.append(f"{field}: expected {present_value}, got {dataset_value}")
+
+    if mismatches:
+        raise ValueError(
+            "Dataset metadata compatibility check failed with mismatches:\n" + "\n".join(mismatches)
+        )

lerobot/common/robot_devices/motors/configs.py

@@ -0,0 +1,41 @@
+# 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 abc
+from dataclasses import dataclass
+
+import draccus
+
+
+@dataclass
+class MotorsBusConfig(draccus.ChoiceRegistry, abc.ABC):
+    @property
+    def type(self) -> str:
+        return self.get_choice_name(self.__class__)
+
+
+@MotorsBusConfig.register_subclass("dynamixel")
+@dataclass
+class DynamixelMotorsBusConfig(MotorsBusConfig):
+    port: str
+    motors: dict[str, tuple[int, str]]
+    mock: bool = False
+
+
+@MotorsBusConfig.register_subclass("feetech")
+@dataclass
+class FeetechMotorsBusConfig(MotorsBusConfig):
+    port: str
+    motors: dict[str, tuple[int, str]]
+    mock: bool = False

lerobot/common/robot_devices/motors/dynamixel.py

@@ -0,0 +1,873 @@
+# 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 enum
+import logging
+import math
+import time
+import traceback
+from copy import deepcopy
+
+import numpy as np
+import tqdm
+
+from lerobot.common.robot_devices.motors.configs import DynamixelMotorsBusConfig
+from lerobot.common.robot_devices.utils import RobotDeviceAlreadyConnectedError, RobotDeviceNotConnectedError
+from lerobot.common.utils.utils import capture_timestamp_utc
+
+PROTOCOL_VERSION = 2.0
+BAUDRATE = 1_000_000
+TIMEOUT_MS = 1000
+
+MAX_ID_RANGE = 252
+
+# The following bounds define the lower and upper joints range (after calibration).
+# For joints in degree (i.e. revolute joints), their nominal range is [-180, 180] degrees
+# which corresponds to a half rotation on the left and half rotation on the right.
+# Some joints might require higher range, so we allow up to [-270, 270] degrees until
+# an error is raised.
+LOWER_BOUND_DEGREE = -270
+UPPER_BOUND_DEGREE = 270
+# For joints in percentage (i.e. joints that move linearly like the prismatic joint of a gripper),
+# their nominal range is [0, 100] %. For instance, for Aloha gripper, 0% is fully
+# closed, and 100% is fully open. To account for slight calibration issue, we allow up to
+# [-10, 110] until an error is raised.
+LOWER_BOUND_LINEAR = -10
+UPPER_BOUND_LINEAR = 110
+
+HALF_TURN_DEGREE = 180
+
+# https://emanual.robotis.com/docs/en/dxl/x/xl330-m077
+# https://emanual.robotis.com/docs/en/dxl/x/xl330-m288
+# https://emanual.robotis.com/docs/en/dxl/x/xl430-w250
+# https://emanual.robotis.com/docs/en/dxl/x/xm430-w350
+# https://emanual.robotis.com/docs/en/dxl/x/xm540-w270
+# https://emanual.robotis.com/docs/en/dxl/x/xc430-w150
+
+# data_name: (address, size_byte)
+X_SERIES_CONTROL_TABLE = {
+    "Model_Number": (0, 2),
+    "Model_Information": (2, 4),
+    "Firmware_Version": (6, 1),
+    "ID": (7, 1),
+    "Baud_Rate": (8, 1),
+    "Return_Delay_Time": (9, 1),
+    "Drive_Mode": (10, 1),
+    "Operating_Mode": (11, 1),
+    "Secondary_ID": (12, 1),
+    "Protocol_Type": (13, 1),
+    "Homing_Offset": (20, 4),
+    "Moving_Threshold": (24, 4),
+    "Temperature_Limit": (31, 1),
+    "Max_Voltage_Limit": (32, 2),
+    "Min_Voltage_Limit": (34, 2),
+    "PWM_Limit": (36, 2),
+    "Current_Limit": (38, 2),
+    "Acceleration_Limit": (40, 4),
+    "Velocity_Limit": (44, 4),
+    "Max_Position_Limit": (48, 4),
+    "Min_Position_Limit": (52, 4),
+    "Shutdown": (63, 1),
+    "Torque_Enable": (64, 1),
+    "LED": (65, 1),
+    "Status_Return_Level": (68, 1),
+    "Registered_Instruction": (69, 1),
+    "Hardware_Error_Status": (70, 1),
+    "Velocity_I_Gain": (76, 2),
+    "Velocity_P_Gain": (78, 2),
+    "Position_D_Gain": (80, 2),
+    "Position_I_Gain": (82, 2),
+    "Position_P_Gain": (84, 2),
+    "Feedforward_2nd_Gain": (88, 2),
+    "Feedforward_1st_Gain": (90, 2),
+    "Bus_Watchdog": (98, 1),
+    "Goal_PWM": (100, 2),
+    "Goal_Current": (102, 2),
+    "Goal_Velocity": (104, 4),
+    "Profile_Acceleration": (108, 4),
+    "Profile_Velocity": (112, 4),
+    "Goal_Position": (116, 4),
+    "Realtime_Tick": (120, 2),
+    "Moving": (122, 1),
+    "Moving_Status": (123, 1),
+    "Present_PWM": (124, 2),
+    "Present_Current": (126, 2),
+    "Present_Velocity": (128, 4),
+    "Present_Position": (132, 4),
+    "Velocity_Trajectory": (136, 4),
+    "Position_Trajectory": (140, 4),
+    "Present_Input_Voltage": (144, 2),
+    "Present_Temperature": (146, 1),
+}
+
+X_SERIES_BAUDRATE_TABLE = {
+    0: 9_600,
+    1: 57_600,
+    2: 115_200,
+    3: 1_000_000,
+    4: 2_000_000,
+    5: 3_000_000,
+    6: 4_000_000,
+}
+
+CALIBRATION_REQUIRED = ["Goal_Position", "Present_Position"]
+CONVERT_UINT32_TO_INT32_REQUIRED = ["Goal_Position", "Present_Position"]
+
+MODEL_CONTROL_TABLE = {
+    "x_series": X_SERIES_CONTROL_TABLE,
+    "xl330-m077": X_SERIES_CONTROL_TABLE,
+    "xl330-m288": X_SERIES_CONTROL_TABLE,
+    "xl430-w250": X_SERIES_CONTROL_TABLE,
+    "xm430-w350": X_SERIES_CONTROL_TABLE,
+    "xm540-w270": X_SERIES_CONTROL_TABLE,
+    "xc430-w150": X_SERIES_CONTROL_TABLE,
+}
+
+MODEL_RESOLUTION = {
+    "x_series": 4096,
+    "xl330-m077": 4096,
+    "xl330-m288": 4096,
+    "xl430-w250": 4096,
+    "xm430-w350": 4096,
+    "xm540-w270": 4096,
+    "xc430-w150": 4096,
+}
+
+MODEL_BAUDRATE_TABLE = {
+    "x_series": X_SERIES_BAUDRATE_TABLE,
+    "xl330-m077": X_SERIES_BAUDRATE_TABLE,
+    "xl330-m288": X_SERIES_BAUDRATE_TABLE,
+    "xl430-w250": X_SERIES_BAUDRATE_TABLE,
+    "xm430-w350": X_SERIES_BAUDRATE_TABLE,
+    "xm540-w270": X_SERIES_BAUDRATE_TABLE,
+    "xc430-w150": X_SERIES_BAUDRATE_TABLE,
+}
+
+NUM_READ_RETRY = 10
+NUM_WRITE_RETRY = 10
+
+
+def convert_degrees_to_steps(degrees: float | np.ndarray, models: str | list[str]) -> np.ndarray:
+    """This function converts the degree range to the step range for indicating motors rotation.
+    It assumes a motor achieves a full rotation by going from -180 degree position to +180.
+    The motor resolution (e.g. 4096) corresponds to the number of steps needed to achieve a full rotation.
+    """
+    resolutions = [MODEL_RESOLUTION[model] for model in models]
+    steps = degrees / 180 * np.array(resolutions) / 2
+    steps = steps.astype(int)
+    return steps
+
+
+def convert_to_bytes(value, bytes, mock=False):
+    if mock:
+        return value
+
+    import dynamixel_sdk as dxl
+
+    # Note: No need to convert back into unsigned int, since this byte preprocessing
+    # already handles it for us.
+    if bytes == 1:
+        data = [
+            dxl.DXL_LOBYTE(dxl.DXL_LOWORD(value)),
+        ]
+    elif bytes == 2:
+        data = [
+            dxl.DXL_LOBYTE(dxl.DXL_LOWORD(value)),
+            dxl.DXL_HIBYTE(dxl.DXL_LOWORD(value)),
+        ]
+    elif bytes == 4:
+        data = [
+            dxl.DXL_LOBYTE(dxl.DXL_LOWORD(value)),
+            dxl.DXL_HIBYTE(dxl.DXL_LOWORD(value)),
+            dxl.DXL_LOBYTE(dxl.DXL_HIWORD(value)),
+            dxl.DXL_HIBYTE(dxl.DXL_HIWORD(value)),
+        ]
+    else:
+        raise NotImplementedError(
+            f"Value of the number of bytes to be sent is expected to be in [1, 2, 4], but "
+            f"{bytes} is provided instead."
+        )
+    return data
+
+
+def get_group_sync_key(data_name, motor_names):
+    group_key = f"{data_name}_" + "_".join(motor_names)
+    return group_key
+
+
+def get_result_name(fn_name, data_name, motor_names):
+    group_key = get_group_sync_key(data_name, motor_names)
+    rslt_name = f"{fn_name}_{group_key}"
+    return rslt_name
+
+
+def get_queue_name(fn_name, data_name, motor_names):
+    group_key = get_group_sync_key(data_name, motor_names)
+    queue_name = f"{fn_name}_{group_key}"
+    return queue_name
+
+
+def get_log_name(var_name, fn_name, data_name, motor_names):
+    group_key = get_group_sync_key(data_name, motor_names)
+    log_name = f"{var_name}_{fn_name}_{group_key}"
+    return log_name
+
+
+def assert_same_address(model_ctrl_table, motor_models, data_name):
+    all_addr = []
+    all_bytes = []
+    for model in motor_models:
+        addr, bytes = model_ctrl_table[model][data_name]
+        all_addr.append(addr)
+        all_bytes.append(bytes)
+
+    if len(set(all_addr)) != 1:
+        raise NotImplementedError(
+            f"At least two motor models use a different address for `data_name`='{data_name}' ({list(zip(motor_models, all_addr, strict=False))}). Contact a LeRobot maintainer."
+        )
+
+    if len(set(all_bytes)) != 1:
+        raise NotImplementedError(
+            f"At least two motor models use a different bytes representation for `data_name`='{data_name}' ({list(zip(motor_models, all_bytes, strict=False))}). Contact a LeRobot maintainer."
+        )
+
+
+class TorqueMode(enum.Enum):
+    ENABLED = 1
+    DISABLED = 0
+
+
+class DriveMode(enum.Enum):
+    NON_INVERTED = 0
+    INVERTED = 1
+
+
+class CalibrationMode(enum.Enum):
+    # Joints with rotational motions are expressed in degrees in nominal range of [-180, 180]
+    DEGREE = 0
+    # Joints with linear motions (like gripper of Aloha) are expressed in nominal range of [0, 100]
+    LINEAR = 1
+
+
+class JointOutOfRangeError(Exception):
+    def __init__(self, message="Joint is out of range"):
+        self.message = message
+        super().__init__(self.message)
+
+
+class DynamixelMotorsBus:
+    """
+    The DynamixelMotorsBus class allows to efficiently read and write to the attached motors. It relies on
+    the python dynamixel sdk to communicate with the motors. For more info, see the [Dynamixel SDK Documentation](https://emanual.robotis.com/docs/en/software/dynamixel/dynamixel_sdk/sample_code/python_read_write_protocol_2_0/#python-read-write-protocol-20).
+
+    A DynamixelMotorsBus instance requires a port (e.g. `DynamixelMotorsBus(port="/dev/tty.usbmodem575E0031751"`)).
+    To find the port, you can run our utility script:
+    ```bash
+    python lerobot/scripts/find_motors_bus_port.py
+    >>> Finding all available ports for the MotorBus.
+    >>> ['/dev/tty.usbmodem575E0032081', '/dev/tty.usbmodem575E0031751']
+    >>> Remove the usb cable from your DynamixelMotorsBus and press Enter when done.
+    >>> The port of this DynamixelMotorsBus is /dev/tty.usbmodem575E0031751.
+    >>> Reconnect the usb cable.
+    ```
+
+    Example of usage for 1 motor connected to the bus:
+    ```python
+    motor_name = "gripper"
+    motor_index = 6
+    motor_model = "xl330-m288"
+
+    config = DynamixelMotorsBusConfig(
+        port="/dev/tty.usbmodem575E0031751",
+        motors={motor_name: (motor_index, motor_model)},
+    )
+    motors_bus = DynamixelMotorsBus(config)
+    motors_bus.connect()
+
+    position = motors_bus.read("Present_Position")
+
+    # move from a few motor steps as an example
+    few_steps = 30
+    motors_bus.write("Goal_Position", position + few_steps)
+
+    # when done, consider disconnecting
+    motors_bus.disconnect()
+    ```
+    """
+
+    def __init__(
+        self,
+        config: DynamixelMotorsBusConfig,
+    ):
+        self.port = config.port
+        self.motors = config.motors
+        self.mock = config.mock
+
+        self.model_ctrl_table = deepcopy(MODEL_CONTROL_TABLE)
+        self.model_resolution = deepcopy(MODEL_RESOLUTION)
+
+        self.port_handler = None
+        self.packet_handler = None
+        self.calibration = None
+        self.is_connected = False
+        self.group_readers = {}
+        self.group_writers = {}
+        self.logs = {}
+
+    def connect(self):
+        if self.is_connected:
+            raise RobotDeviceAlreadyConnectedError(
+                f"DynamixelMotorsBus({self.port}) is already connected. Do not call `motors_bus.connect()` twice."
+            )
+
+        if self.mock:
+            import tests.motors.mock_dynamixel_sdk as dxl
+        else:
+            import dynamixel_sdk as dxl
+
+        self.port_handler = dxl.PortHandler(self.port)
+        self.packet_handler = dxl.PacketHandler(PROTOCOL_VERSION)
+
+        try:
+            if not self.port_handler.openPort():
+                raise OSError(f"Failed to open port '{self.port}'.")
+        except Exception:
+            traceback.print_exc()
+            print(
+                "\nTry running `python lerobot/scripts/find_motors_bus_port.py` to make sure you are using the correct port.\n"
+            )
+            raise
+
+        # Allow to read and write
+        self.is_connected = True
+
+        self.port_handler.setPacketTimeoutMillis(TIMEOUT_MS)
+
+    def reconnect(self):
+        if self.mock:
+            import tests.motors.mock_dynamixel_sdk as dxl
+        else:
+            import dynamixel_sdk as dxl
+
+        self.port_handler = dxl.PortHandler(self.port)
+        self.packet_handler = dxl.PacketHandler(PROTOCOL_VERSION)
+
+        if not self.port_handler.openPort():
+            raise OSError(f"Failed to open port '{self.port}'.")
+
+        self.is_connected = True
+
+    def are_motors_configured(self):
+        # Only check the motor indices and not baudrate, since if the motor baudrates are incorrect,
+        # a ConnectionError will be raised anyway.
+        try:
+            return (self.motor_indices == self.read("ID")).all()
+        except ConnectionError as e:
+            print(e)
+            return False
+
+    def find_motor_indices(self, possible_ids=None, num_retry=2):
+        if possible_ids is None:
+            possible_ids = range(MAX_ID_RANGE)
+
+        indices = []
+        for idx in tqdm.tqdm(possible_ids):
+            try:
+                present_idx = self.read_with_motor_ids(self.motor_models, [idx], "ID", num_retry=num_retry)[0]
+            except ConnectionError:
+                continue
+
+            if idx != present_idx:
+                # sanity check
+                raise OSError(
+                    "Motor index used to communicate through the bus is not the same as the one present in the motor memory. The motor memory might be damaged."
+                )
+            indices.append(idx)
+
+        return indices
+
+    def set_bus_baudrate(self, baudrate):
+        present_bus_baudrate = self.port_handler.getBaudRate()
+        if present_bus_baudrate != baudrate:
+            print(f"Setting bus baud rate to {baudrate}. Previously {present_bus_baudrate}.")
+            self.port_handler.setBaudRate(baudrate)
+
+            if self.port_handler.getBaudRate() != baudrate:
+                raise OSError("Failed to write bus baud rate.")
+
+    @property
+    def motor_names(self) -> list[str]:
+        return list(self.motors.keys())
+
+    @property
+    def motor_models(self) -> list[str]:
+        return [model for _, model in self.motors.values()]
+
+    @property
+    def motor_indices(self) -> list[int]:
+        return [idx for idx, _ in self.motors.values()]
+
+    def set_calibration(self, calibration: dict[str, list]):
+        self.calibration = calibration
+
+    def apply_calibration_autocorrect(self, values: np.ndarray | list, motor_names: list[str] | None):
+        """This function applies the calibration, automatically detects out of range errors for motors values and attempts to correct.
+
+        For more info, see docstring of `apply_calibration` and `autocorrect_calibration`.
+        """
+        try:
+            values = self.apply_calibration(values, motor_names)
+        except JointOutOfRangeError as e:
+            print(e)
+            self.autocorrect_calibration(values, motor_names)
+            values = self.apply_calibration(values, motor_names)
+        return values
+
+    def apply_calibration(self, values: np.ndarray | list, motor_names: list[str] | None):
+        """Convert from unsigned int32 joint position range [0, 2**32[ to the universal float32 nominal degree range ]-180.0, 180.0[ with
+        a "zero position" at 0 degree.
+
+        Note: We say "nominal degree range" since the motors can take values outside this range. For instance, 190 degrees, if the motor
+        rotate more than a half a turn from the zero position. However, most motors can't rotate more than 180 degrees and will stay in this range.
+
+        Joints values are original in [0, 2**32[ (unsigned int32). Each motor are expected to complete a full rotation
+        when given a goal position that is + or - their resolution. For instance, dynamixel xl330-m077 have a resolution of 4096, and
+        at any position in their original range, let's say the position 56734, they complete a full rotation clockwise by moving to 60830,
+        or anticlockwise by moving to 52638. The position in the original range is arbitrary and might change a lot between each motor.
+        To harmonize between motors of the same model, different robots, or even models of different brands, we propose to work
+        in the centered nominal degree range ]-180, 180[.
+        """
+        if motor_names is None:
+            motor_names = self.motor_names
+
+        # Convert from unsigned int32 original range [0, 2**32] to signed float32 range
+        values = values.astype(np.float32)
+
+        for i, name in enumerate(motor_names):
+            calib_idx = self.calibration["motor_names"].index(name)
+            calib_mode = self.calibration["calib_mode"][calib_idx]
+
+            if CalibrationMode[calib_mode] == CalibrationMode.DEGREE:
+                drive_mode = self.calibration["drive_mode"][calib_idx]
+                homing_offset = self.calibration["homing_offset"][calib_idx]
+                _, model = self.motors[name]
+                resolution = self.model_resolution[model]
+
+                # Update direction of rotation of the motor to match between leader and follower.
+                # In fact, the motor of the leader for a given joint can be assembled in an
+                # opposite direction in term of rotation than the motor of the follower on the same joint.
+                if drive_mode:
+                    values[i] *= -1
+
+                # Convert from range [-2**31, 2**31] to
+                # nominal range [-resolution//2, resolution//2] (e.g. [-2048, 2048])
+                values[i] += homing_offset
+
+                # Convert from range [-resolution//2, resolution//2] to
+                # universal float32 centered degree range [-180, 180]
+                # (e.g. 2048 / (4096 // 2) * 180 = 180)
+                values[i] = values[i] / (resolution // 2) * HALF_TURN_DEGREE
+
+                if (values[i] < LOWER_BOUND_DEGREE) or (values[i] > UPPER_BOUND_DEGREE):
+                    raise JointOutOfRangeError(
+                        f"Wrong motor position range detected for {name}. "
+                        f"Expected to be in nominal range of [-{HALF_TURN_DEGREE}, {HALF_TURN_DEGREE}] degrees (a full rotation), "
+                        f"with a maximum range of [{LOWER_BOUND_DEGREE}, {UPPER_BOUND_DEGREE}] degrees to account for joints that can rotate a bit more, "
+                        f"but present value is {values[i]} degree. "
+                        "This might be due to a cable connection issue creating an artificial 360 degrees jump in motor values. "
+                        "You need to recalibrate by running: `python lerobot/scripts/control_robot.py calibrate`"
+                    )
+
+            elif CalibrationMode[calib_mode] == CalibrationMode.LINEAR:
+                start_pos = self.calibration["start_pos"][calib_idx]
+                end_pos = self.calibration["end_pos"][calib_idx]
+
+                # Rescale the present position to a nominal range [0, 100] %,
+                # useful for joints with linear motions like Aloha gripper
+                values[i] = (values[i] - start_pos) / (end_pos - start_pos) * 100
+
+                if (values[i] < LOWER_BOUND_LINEAR) or (values[i] > UPPER_BOUND_LINEAR):
+                    raise JointOutOfRangeError(
+                        f"Wrong motor position range detected for {name}. "
+                        f"Expected to be in nominal range of [0, 100] % (a full linear translation), "
+                        f"with a maximum range of [{LOWER_BOUND_LINEAR}, {UPPER_BOUND_LINEAR}] % to account for some imprecision during calibration, "
+                        f"but present value is {values[i]} %. "
+                        "This might be due to a cable connection issue creating an artificial jump in motor values. "
+                        "You need to recalibrate by running: `python lerobot/scripts/control_robot.py calibrate`"
+                    )
+
+        return values
+
+    def autocorrect_calibration(self, values: np.ndarray | list, motor_names: list[str] | None):
+        """This function automatically detects issues with values of motors after calibration, and correct for these issues.
+
+        Some motors might have values outside of expected maximum bounds after calibration.
+        For instance, for a joint in degree, its value can be outside [-270, 270] degrees, which is totally unexpected given
+        a nominal range of [-180, 180] degrees, which represents half a turn to the left or right starting from zero position.
+
+        Known issues:
+        #1: Motor value randomly shifts of a full turn, caused by hardware/connection errors.
+        #2: Motor internal homing offset is shifted by a full turn, caused by using default calibration (e.g Aloha).
+        #3: motor internal homing offset is shifted by less or more than a full turn, caused by using default calibration
+            or by human error during manual calibration.
+
+        Issues #1 and #2 can be solved by shifting the calibration homing offset by a full turn.
+        Issue #3 will be visually detected by user and potentially captured by the safety feature `max_relative_target`,
+        that will slow down the motor, raise an error asking to recalibrate. Manual recalibrating will solve the issue.
+
+        Note: A full turn corresponds to 360 degrees but also to 4096 steps for a motor resolution of 4096.
+        """
+        if motor_names is None:
+            motor_names = self.motor_names
+
+        # Convert from unsigned int32 original range [0, 2**32] to signed float32 range
+        values = values.astype(np.float32)
+
+        for i, name in enumerate(motor_names):
+            calib_idx = self.calibration["motor_names"].index(name)
+            calib_mode = self.calibration["calib_mode"][calib_idx]
+
+            if CalibrationMode[calib_mode] == CalibrationMode.DEGREE:
+                drive_mode = self.calibration["drive_mode"][calib_idx]
+                homing_offset = self.calibration["homing_offset"][calib_idx]
+                _, model = self.motors[name]
+                resolution = self.model_resolution[model]
+
+                # Update direction of rotation of the motor to match between leader and follower.
+                # In fact, the motor of the leader for a given joint can be assembled in an
+                # opposite direction in term of rotation than the motor of the follower on the same joint.
+                if drive_mode:
+                    values[i] *= -1
+
+                # Convert from initial range to range [-180, 180] degrees
+                calib_val = (values[i] + homing_offset) / (resolution // 2) * HALF_TURN_DEGREE
+                in_range = (calib_val > LOWER_BOUND_DEGREE) and (calib_val < UPPER_BOUND_DEGREE)
+
+                # Solve this inequality to find the factor to shift the range into [-180, 180] degrees
+                # values[i] = (values[i] + homing_offset + resolution * factor) / (resolution // 2) * HALF_TURN_DEGREE
+                # - HALF_TURN_DEGREE <= (values[i] + homing_offset + resolution * factor) / (resolution // 2) * HALF_TURN_DEGREE <= HALF_TURN_DEGREE
+                # (- (resolution // 2) - values[i] - homing_offset) / resolution <= factor <= ((resolution // 2) - values[i] - homing_offset) / resolution
+                low_factor = (-(resolution // 2) - values[i] - homing_offset) / resolution
+                upp_factor = ((resolution // 2) - values[i] - homing_offset) / resolution
+
+            elif CalibrationMode[calib_mode] == CalibrationMode.LINEAR:
+                start_pos = self.calibration["start_pos"][calib_idx]
+                end_pos = self.calibration["end_pos"][calib_idx]
+
+                # Convert from initial range to range [0, 100] in %
+                calib_val = (values[i] - start_pos) / (end_pos - start_pos) * 100
+                in_range = (calib_val > LOWER_BOUND_LINEAR) and (calib_val < UPPER_BOUND_LINEAR)
+
+                # Solve this inequality to find the factor to shift the range into [0, 100] %
+                # values[i] = (values[i] - start_pos + resolution * factor) / (end_pos + resolution * factor - start_pos - resolution * factor) * 100
+                # values[i] = (values[i] - start_pos + resolution * factor) / (end_pos - start_pos) * 100
+                # 0 <= (values[i] - start_pos + resolution * factor) / (end_pos - start_pos) * 100 <= 100
+                # (start_pos - values[i]) / resolution <= factor <= (end_pos - values[i]) / resolution
+                low_factor = (start_pos - values[i]) / resolution
+                upp_factor = (end_pos - values[i]) / resolution
+
+            if not in_range:
+                # Get first integer between the two bounds
+                if low_factor < upp_factor:
+                    factor = math.ceil(low_factor)
+
+                    if factor > upp_factor:
+                        raise ValueError(f"No integer found between bounds [{low_factor=}, {upp_factor=}]")
+                else:
+                    factor = math.ceil(upp_factor)
+
+                    if factor > low_factor:
+                        raise ValueError(f"No integer found between bounds [{low_factor=}, {upp_factor=}]")
+
+                if CalibrationMode[calib_mode] == CalibrationMode.DEGREE:
+                    out_of_range_str = f"{LOWER_BOUND_DEGREE} < {calib_val} < {UPPER_BOUND_DEGREE} degrees"
+                    in_range_str = f"{LOWER_BOUND_DEGREE} < {calib_val} < {UPPER_BOUND_DEGREE} degrees"
+                elif CalibrationMode[calib_mode] == CalibrationMode.LINEAR:
+                    out_of_range_str = f"{LOWER_BOUND_LINEAR} < {calib_val} < {UPPER_BOUND_LINEAR} %"
+                    in_range_str = f"{LOWER_BOUND_LINEAR} < {calib_val} < {UPPER_BOUND_LINEAR} %"
+
+                logging.warning(
+                    f"Auto-correct calibration of motor '{name}' by shifting value by {abs(factor)} full turns, "
+                    f"from '{out_of_range_str}' to '{in_range_str}'."
+                )
+
+                # A full turn corresponds to 360 degrees but also to 4096 steps for a motor resolution of 4096.
+                self.calibration["homing_offset"][calib_idx] += resolution * factor
+
+    def revert_calibration(self, values: np.ndarray | list, motor_names: list[str] | None):
+        """Inverse of `apply_calibration`."""
+        if motor_names is None:
+            motor_names = self.motor_names
+
+        for i, name in enumerate(motor_names):
+            calib_idx = self.calibration["motor_names"].index(name)
+            calib_mode = self.calibration["calib_mode"][calib_idx]
+
+            if CalibrationMode[calib_mode] == CalibrationMode.DEGREE:
+                drive_mode = self.calibration["drive_mode"][calib_idx]
+                homing_offset = self.calibration["homing_offset"][calib_idx]
+                _, model = self.motors[name]
+                resolution = self.model_resolution[model]
+
+                # Convert from nominal 0-centered degree range [-180, 180] to
+                # 0-centered resolution range (e.g. [-2048, 2048] for resolution=4096)
+                values[i] = values[i] / HALF_TURN_DEGREE * (resolution // 2)
+
+                # Subtract the homing offsets to come back to actual motor range of values
+                # which can be arbitrary.
+                values[i] -= homing_offset
+
+                # Remove drive mode, which is the rotation direction of the motor, to come back to
+                # actual motor rotation direction which can be arbitrary.
+                if drive_mode:
+                    values[i] *= -1
+
+            elif CalibrationMode[calib_mode] == CalibrationMode.LINEAR:
+                start_pos = self.calibration["start_pos"][calib_idx]
+                end_pos = self.calibration["end_pos"][calib_idx]
+
+                # Convert from nominal lnear range of [0, 100] % to
+                # actual motor range of values which can be arbitrary.
+                values[i] = values[i] / 100 * (end_pos - start_pos) + start_pos
+
+        values = np.round(values).astype(np.int32)
+        return values
+
+    def read_with_motor_ids(self, motor_models, motor_ids, data_name, num_retry=NUM_READ_RETRY):
+        if self.mock:
+            import tests.motors.mock_dynamixel_sdk as dxl
+        else:
+            import dynamixel_sdk as dxl
+
+        return_list = True
+        if not isinstance(motor_ids, list):
+            return_list = False
+            motor_ids = [motor_ids]
+
+        assert_same_address(self.model_ctrl_table, self.motor_models, data_name)
+        addr, bytes = self.model_ctrl_table[motor_models[0]][data_name]
+        group = dxl.GroupSyncRead(self.port_handler, self.packet_handler, addr, bytes)
+        for idx in motor_ids:
+            group.addParam(idx)
+
+        for _ in range(num_retry):
+            comm = group.txRxPacket()
+            if comm == dxl.COMM_SUCCESS:
+                break
+
+        if comm != dxl.COMM_SUCCESS:
+            raise ConnectionError(
+                f"Read failed due to communication error on port {self.port_handler.port_name} for indices {motor_ids}: "
+                f"{self.packet_handler.getTxRxResult(comm)}"
+            )
+
+        values = []
+        for idx in motor_ids:
+            value = group.getData(idx, addr, bytes)
+            values.append(value)
+
+        if return_list:
+            return values
+        else:
+            return values[0]
+
+    def read(self, data_name, motor_names: str | list[str] | None = None):
+        if not self.is_connected:
+            raise RobotDeviceNotConnectedError(
+                f"DynamixelMotorsBus({self.port}) is not connected. You need to run `motors_bus.connect()`."
+            )
+
+        start_time = time.perf_counter()
+
+        if self.mock:
+            import tests.motors.mock_dynamixel_sdk as dxl
+        else:
+            import dynamixel_sdk as dxl
+
+        if motor_names is None:
+            motor_names = self.motor_names
+
+        if isinstance(motor_names, str):
+            motor_names = [motor_names]
+
+        motor_ids = []
+        models = []
+        for name in motor_names:
+            motor_idx, model = self.motors[name]
+            motor_ids.append(motor_idx)
+            models.append(model)
+
+        assert_same_address(self.model_ctrl_table, models, data_name)
+        addr, bytes = self.model_ctrl_table[model][data_name]
+        group_key = get_group_sync_key(data_name, motor_names)
+
+        if data_name not in self.group_readers:
+            # create new group reader
+            self.group_readers[group_key] = dxl.GroupSyncRead(
+                self.port_handler, self.packet_handler, addr, bytes
+            )
+            for idx in motor_ids:
+                self.group_readers[group_key].addParam(idx)
+
+        for _ in range(NUM_READ_RETRY):
+            comm = self.group_readers[group_key].txRxPacket()
+            if comm == dxl.COMM_SUCCESS:
+                break
+
+        if comm != dxl.COMM_SUCCESS:
+            raise ConnectionError(
+                f"Read failed due to communication error on port {self.port} for group_key {group_key}: "
+                f"{self.packet_handler.getTxRxResult(comm)}"
+            )
+
+        values = []
+        for idx in motor_ids:
+            value = self.group_readers[group_key].getData(idx, addr, bytes)
+            values.append(value)
+
+        values = np.array(values)
+
+        # Convert to signed int to use range [-2048, 2048] for our motor positions.
+        if data_name in CONVERT_UINT32_TO_INT32_REQUIRED:
+            values = values.astype(np.int32)
+
+        if data_name in CALIBRATION_REQUIRED and self.calibration is not None:
+            values = self.apply_calibration_autocorrect(values, motor_names)
+
+        # log the number of seconds it took to read the data from the motors
+        delta_ts_name = get_log_name("delta_timestamp_s", "read", data_name, motor_names)
+        self.logs[delta_ts_name] = time.perf_counter() - start_time
+
+        # log the utc time at which the data was received
+        ts_utc_name = get_log_name("timestamp_utc", "read", data_name, motor_names)
+        self.logs[ts_utc_name] = capture_timestamp_utc()
+
+        return values
+
+    def write_with_motor_ids(self, motor_models, motor_ids, data_name, values, num_retry=NUM_WRITE_RETRY):
+        if self.mock:
+            import tests.motors.mock_dynamixel_sdk as dxl
+        else:
+            import dynamixel_sdk as dxl
+
+        if not isinstance(motor_ids, list):
+            motor_ids = [motor_ids]
+        if not isinstance(values, list):
+            values = [values]
+
+        assert_same_address(self.model_ctrl_table, motor_models, data_name)
+        addr, bytes = self.model_ctrl_table[motor_models[0]][data_name]
+        group = dxl.GroupSyncWrite(self.port_handler, self.packet_handler, addr, bytes)
+        for idx, value in zip(motor_ids, values, strict=True):
+            data = convert_to_bytes(value, bytes, self.mock)
+            group.addParam(idx, data)
+
+        for _ in range(num_retry):
+            comm = group.txPacket()
+            if comm == dxl.COMM_SUCCESS:
+                break
+
+        if comm != dxl.COMM_SUCCESS:
+            raise ConnectionError(
+                f"Write failed due to communication error on port {self.port_handler.port_name} for indices {motor_ids}: "
+                f"{self.packet_handler.getTxRxResult(comm)}"
+            )
+
+    def write(self, data_name, values: int | float | np.ndarray, motor_names: str | list[str] | None = None):
+        if not self.is_connected:
+            raise RobotDeviceNotConnectedError(
+                f"DynamixelMotorsBus({self.port}) is not connected. You need to run `motors_bus.connect()`."
+            )
+
+        start_time = time.perf_counter()
+
+        if self.mock:
+            import tests.motors.mock_dynamixel_sdk as dxl
+        else:
+            import dynamixel_sdk as dxl
+
+        if motor_names is None:
+            motor_names = self.motor_names
+
+        if isinstance(motor_names, str):
+            motor_names = [motor_names]
+
+        if isinstance(values, (int, float, np.integer)):
+            values = [int(values)] * len(motor_names)
+
+        values = np.array(values)
+
+        motor_ids = []
+        models = []
+        for name in motor_names:
+            motor_idx, model = self.motors[name]
+            motor_ids.append(motor_idx)
+            models.append(model)
+
+        if data_name in CALIBRATION_REQUIRED and self.calibration is not None:
+            values = self.revert_calibration(values, motor_names)
+
+        values = values.tolist()
+
+        assert_same_address(self.model_ctrl_table, models, data_name)
+        addr, bytes = self.model_ctrl_table[model][data_name]
+        group_key = get_group_sync_key(data_name, motor_names)
+
+        init_group = data_name not in self.group_readers
+        if init_group:
+            self.group_writers[group_key] = dxl.GroupSyncWrite(
+                self.port_handler, self.packet_handler, addr, bytes
+            )
+
+        for idx, value in zip(motor_ids, values, strict=True):
+            data = convert_to_bytes(value, bytes, self.mock)
+            if init_group:
+                self.group_writers[group_key].addParam(idx, data)
+            else:
+                self.group_writers[group_key].changeParam(idx, data)
+
+        comm = self.group_writers[group_key].txPacket()
+        if comm != dxl.COMM_SUCCESS:
+            raise ConnectionError(
+                f"Write failed due to communication error on port {self.port} for group_key {group_key}: "
+                f"{self.packet_handler.getTxRxResult(comm)}"
+            )
+
+        # log the number of seconds it took to write the data to the motors
+        delta_ts_name = get_log_name("delta_timestamp_s", "write", data_name, motor_names)
+        self.logs[delta_ts_name] = time.perf_counter() - start_time
+
+        # TODO(rcadene): should we log the time before sending the write command?
+        # log the utc time when the write has been completed
+        ts_utc_name = get_log_name("timestamp_utc", "write", data_name, motor_names)
+        self.logs[ts_utc_name] = capture_timestamp_utc()
+
+    def disconnect(self):
+        if not self.is_connected:
+            raise RobotDeviceNotConnectedError(
+                f"DynamixelMotorsBus({self.port}) is not connected. Try running `motors_bus.connect()` first."
+            )
+
+        if self.port_handler is not None:
+            self.port_handler.closePort()
+            self.port_handler = None
+
+        self.packet_handler = None
+        self.group_readers = {}
+        self.group_writers = {}
+        self.is_connected = False
+
+    def __del__(self):
+        if getattr(self, "is_connected", False):
+            self.disconnect()

lerobot/common/robot_devices/motors/feetech.py

@@ -0,0 +1,898 @@
+# 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 enum
+import logging
+import math
+import time
+import traceback
+from copy import deepcopy
+
+import numpy as np
+import tqdm
+
+from lerobot.common.robot_devices.motors.configs import FeetechMotorsBusConfig
+from lerobot.common.robot_devices.utils import RobotDeviceAlreadyConnectedError, RobotDeviceNotConnectedError
+from lerobot.common.utils.utils import capture_timestamp_utc
+
+PROTOCOL_VERSION = 0
+BAUDRATE = 1_000_000
+TIMEOUT_MS = 1000
+
+MAX_ID_RANGE = 252
+
+# The following bounds define the lower and upper joints range (after calibration).
+# For joints in degree (i.e. revolute joints), their nominal range is [-180, 180] degrees
+# which corresponds to a half rotation on the left and half rotation on the right.
+# Some joints might require higher range, so we allow up to [-270, 270] degrees until
+# an error is raised.
+LOWER_BOUND_DEGREE = -270
+UPPER_BOUND_DEGREE = 270
+# For joints in percentage (i.e. joints that move linearly like the prismatic joint of a gripper),
+# their nominal range is [0, 100] %. For instance, for Aloha gripper, 0% is fully
+# closed, and 100% is fully open. To account for slight calibration issue, we allow up to
+# [-10, 110] until an error is raised.
+LOWER_BOUND_LINEAR = -10
+UPPER_BOUND_LINEAR = 110
+
+HALF_TURN_DEGREE = 180
+
+
+# See this link for STS3215 Memory Table:
+# https://docs.google.com/spreadsheets/d/1GVs7W1VS1PqdhA1nW-abeyAHhTUxKUdR/edit?usp=sharing&ouid=116566590112741600240&rtpof=true&sd=true
+# data_name: (address, size_byte)
+SCS_SERIES_CONTROL_TABLE = {
+    "Model": (3, 2),
+    "ID": (5, 1),
+    "Baud_Rate": (6, 1),
+    "Return_Delay": (7, 1),
+    "Response_Status_Level": (8, 1),
+    "Min_Angle_Limit": (9, 2),
+    "Max_Angle_Limit": (11, 2),
+    "Max_Temperature_Limit": (13, 1),
+    "Max_Voltage_Limit": (14, 1),
+    "Min_Voltage_Limit": (15, 1),
+    "Max_Torque_Limit": (16, 2),
+    "Phase": (18, 1),
+    "Unloading_Condition": (19, 1),
+    "LED_Alarm_Condition": (20, 1),
+    "P_Coefficient": (21, 1),
+    "D_Coefficient": (22, 1),
+    "I_Coefficient": (23, 1),
+    "Minimum_Startup_Force": (24, 2),
+    "CW_Dead_Zone": (26, 1),
+    "CCW_Dead_Zone": (27, 1),
+    "Protection_Current": (28, 2),
+    "Angular_Resolution": (30, 1),
+    "Offset": (31, 2),
+    "Mode": (33, 1),
+    "Protective_Torque": (34, 1),
+    "Protection_Time": (35, 1),
+    "Overload_Torque": (36, 1),
+    "Speed_closed_loop_P_proportional_coefficient": (37, 1),
+    "Over_Current_Protection_Time": (38, 1),
+    "Velocity_closed_loop_I_integral_coefficient": (39, 1),
+    "Torque_Enable": (40, 1),
+    "Acceleration": (41, 1),
+    "Goal_Position": (42, 2),
+    "Goal_Time": (44, 2),
+    "Goal_Speed": (46, 2),
+    "Torque_Limit": (48, 2),
+    "Lock": (55, 1),
+    "Present_Position": (56, 2),
+    "Present_Speed": (58, 2),
+    "Present_Load": (60, 2),
+    "Present_Voltage": (62, 1),
+    "Present_Temperature": (63, 1),
+    "Status": (65, 1),
+    "Moving": (66, 1),
+    "Present_Current": (69, 2),
+    # Not in the Memory Table
+    "Maximum_Acceleration": (85, 2),
+}
+
+SCS_SERIES_BAUDRATE_TABLE = {
+    0: 1_000_000,
+    1: 500_000,
+    2: 250_000,
+    3: 128_000,
+    4: 115_200,
+    5: 57_600,
+    6: 38_400,
+    7: 19_200,
+}
+
+CALIBRATION_REQUIRED = ["Goal_Position", "Present_Position"]
+CONVERT_UINT32_TO_INT32_REQUIRED = ["Goal_Position", "Present_Position"]
+
+
+MODEL_CONTROL_TABLE = {
+    "scs_series": SCS_SERIES_CONTROL_TABLE,
+    "sts3215": SCS_SERIES_CONTROL_TABLE,
+}
+
+MODEL_RESOLUTION = {
+    "scs_series": 4096,
+    "sts3215": 4096,
+}
+
+MODEL_BAUDRATE_TABLE = {
+    "scs_series": SCS_SERIES_BAUDRATE_TABLE,
+    "sts3215": SCS_SERIES_BAUDRATE_TABLE,
+}
+
+# High number of retries is needed for feetech compared to dynamixel motors.
+NUM_READ_RETRY = 20
+NUM_WRITE_RETRY = 20
+
+
+def convert_degrees_to_steps(degrees: float | np.ndarray, models: str | list[str]) -> np.ndarray:
+    """This function converts the degree range to the step range for indicating motors rotation.
+    It assumes a motor achieves a full rotation by going from -180 degree position to +180.
+    The motor resolution (e.g. 4096) corresponds to the number of steps needed to achieve a full rotation.
+    """
+    resolutions = [MODEL_RESOLUTION[model] for model in models]
+    steps = degrees / 180 * np.array(resolutions) / 2
+    steps = steps.astype(int)
+    return steps
+
+
+def convert_to_bytes(value, bytes, mock=False):
+    if mock:
+        return value
+
+    import scservo_sdk as scs
+
+    # Note: No need to convert back into unsigned int, since this byte preprocessing
+    # already handles it for us.
+    if bytes == 1:
+        data = [
+            scs.SCS_LOBYTE(scs.SCS_LOWORD(value)),
+        ]
+    elif bytes == 2:
+        data = [
+            scs.SCS_LOBYTE(scs.SCS_LOWORD(value)),
+            scs.SCS_HIBYTE(scs.SCS_LOWORD(value)),
+        ]
+    elif bytes == 4:
+        data = [
+            scs.SCS_LOBYTE(scs.SCS_LOWORD(value)),
+            scs.SCS_HIBYTE(scs.SCS_LOWORD(value)),
+            scs.SCS_LOBYTE(scs.SCS_HIWORD(value)),
+            scs.SCS_HIBYTE(scs.SCS_HIWORD(value)),
+        ]
+    else:
+        raise NotImplementedError(
+            f"Value of the number of bytes to be sent is expected to be in [1, 2, 4], but "
+            f"{bytes} is provided instead."
+        )
+    return data
+
+
+def get_group_sync_key(data_name, motor_names):
+    group_key = f"{data_name}_" + "_".join(motor_names)
+    return group_key
+
+
+def get_result_name(fn_name, data_name, motor_names):
+    group_key = get_group_sync_key(data_name, motor_names)
+    rslt_name = f"{fn_name}_{group_key}"
+    return rslt_name
+
+
+def get_queue_name(fn_name, data_name, motor_names):
+    group_key = get_group_sync_key(data_name, motor_names)
+    queue_name = f"{fn_name}_{group_key}"
+    return queue_name
+
+
+def get_log_name(var_name, fn_name, data_name, motor_names):
+    group_key = get_group_sync_key(data_name, motor_names)
+    log_name = f"{var_name}_{fn_name}_{group_key}"
+    return log_name
+
+
+def assert_same_address(model_ctrl_table, motor_models, data_name):
+    all_addr = []
+    all_bytes = []
+    for model in motor_models:
+        addr, bytes = model_ctrl_table[model][data_name]
+        all_addr.append(addr)
+        all_bytes.append(bytes)
+
+    if len(set(all_addr)) != 1:
+        raise NotImplementedError(
+            f"At least two motor models use a different address for `data_name`='{data_name}' ({list(zip(motor_models, all_addr, strict=False))}). Contact a LeRobot maintainer."
+        )
+
+    if len(set(all_bytes)) != 1:
+        raise NotImplementedError(
+            f"At least two motor models use a different bytes representation for `data_name`='{data_name}' ({list(zip(motor_models, all_bytes, strict=False))}). Contact a LeRobot maintainer."
+        )
+
+
+class TorqueMode(enum.Enum):
+    ENABLED = 1
+    DISABLED = 0
+
+
+class DriveMode(enum.Enum):
+    NON_INVERTED = 0
+    INVERTED = 1
+
+
+class CalibrationMode(enum.Enum):
+    # Joints with rotational motions are expressed in degrees in nominal range of [-180, 180]
+    DEGREE = 0
+    # Joints with linear motions (like gripper of Aloha) are expressed in nominal range of [0, 100]
+    LINEAR = 1
+
+
+class JointOutOfRangeError(Exception):
+    def __init__(self, message="Joint is out of range"):
+        self.message = message
+        super().__init__(self.message)
+
+
+class FeetechMotorsBus:
+    """
+    The FeetechMotorsBus class allows to efficiently read and write to the attached motors. It relies on
+    the python feetech sdk to communicate with the motors. For more info, see the [feetech SDK Documentation](https://emanual.robotis.com/docs/en/software/feetech/feetech_sdk/sample_code/python_read_write_protocol_2_0/#python-read-write-protocol-20).
+
+    A FeetechMotorsBus instance requires a port (e.g. `FeetechMotorsBus(port="/dev/tty.usbmodem575E0031751"`)).
+    To find the port, you can run our utility script:
+    ```bash
+    python lerobot/scripts/find_motors_bus_port.py
+    >>> Finding all available ports for the MotorsBus.
+    >>> ['/dev/tty.usbmodem575E0032081', '/dev/tty.usbmodem575E0031751']
+    >>> Remove the usb cable from your FeetechMotorsBus and press Enter when done.
+    >>> The port of this FeetechMotorsBus is /dev/tty.usbmodem575E0031751.
+    >>> Reconnect the usb cable.
+    ```
+
+    Example of usage for 1 motor connected to the bus:
+    ```python
+    motor_name = "gripper"
+    motor_index = 6
+    motor_model = "sts3215"
+
+    config = FeetechMotorsBusConfig(
+        port="/dev/tty.usbmodem575E0031751",
+        motors={motor_name: (motor_index, motor_model)},
+    )
+    motors_bus = FeetechMotorsBus(config)
+    motors_bus.connect()
+
+    position = motors_bus.read("Present_Position")
+
+    # move from a few motor steps as an example
+    few_steps = 30
+    motors_bus.write("Goal_Position", position + few_steps)
+
+    # when done, consider disconnecting
+    motors_bus.disconnect()
+    ```
+    """
+
+    def __init__(
+        self,
+        config: FeetechMotorsBusConfig,
+    ):
+        self.port = config.port
+        self.motors = config.motors
+        self.mock = config.mock
+
+        self.model_ctrl_table = deepcopy(MODEL_CONTROL_TABLE)
+        self.model_resolution = deepcopy(MODEL_RESOLUTION)
+
+        self.port_handler = None
+        self.packet_handler = None
+        self.calibration = None
+        self.is_connected = False
+        self.group_readers = {}
+        self.group_writers = {}
+        self.logs = {}
+
+        self.track_positions = {}
+
+    def connect(self):
+        if self.is_connected:
+            raise RobotDeviceAlreadyConnectedError(
+                f"FeetechMotorsBus({self.port}) is already connected. Do not call `motors_bus.connect()` twice."
+            )
+
+        if self.mock:
+            import tests.motors.mock_scservo_sdk as scs
+        else:
+            import scservo_sdk as scs
+
+        self.port_handler = scs.PortHandler(self.port)
+        self.packet_handler = scs.PacketHandler(PROTOCOL_VERSION)
+
+        try:
+            if not self.port_handler.openPort():
+                raise OSError(f"Failed to open port '{self.port}'.")
+        except Exception:
+            traceback.print_exc()
+            print(
+                "\nTry running `python lerobot/scripts/find_motors_bus_port.py` to make sure you are using the correct port.\n"
+            )
+            raise
+
+        # Allow to read and write
+        self.is_connected = True
+
+        self.port_handler.setPacketTimeoutMillis(TIMEOUT_MS)
+
+    def reconnect(self):
+        if self.mock:
+            import tests.motors.mock_scservo_sdk as scs
+        else:
+            import scservo_sdk as scs
+
+        self.port_handler = scs.PortHandler(self.port)
+        self.packet_handler = scs.PacketHandler(PROTOCOL_VERSION)
+
+        if not self.port_handler.openPort():
+            raise OSError(f"Failed to open port '{self.port}'.")
+
+        self.is_connected = True
+
+    def are_motors_configured(self):
+        # Only check the motor indices and not baudrate, since if the motor baudrates are incorrect,
+        # a ConnectionError will be raised anyway.
+        try:
+            return (self.motor_indices == self.read("ID")).all()
+        except ConnectionError as e:
+            print(e)
+            return False
+
+    def find_motor_indices(self, possible_ids=None, num_retry=2):
+        if possible_ids is None:
+            possible_ids = range(MAX_ID_RANGE)
+
+        indices = []
+        for idx in tqdm.tqdm(possible_ids):
+            try:
+                present_idx = self.read_with_motor_ids(self.motor_models, [idx], "ID", num_retry=num_retry)[0]
+            except ConnectionError:
+                continue
+
+            if idx != present_idx:
+                # sanity check
+                raise OSError(
+                    "Motor index used to communicate through the bus is not the same as the one present in the motor memory. The motor memory might be damaged."
+                )
+            indices.append(idx)
+
+        return indices
+
+    def set_bus_baudrate(self, baudrate):
+        present_bus_baudrate = self.port_handler.getBaudRate()
+        if present_bus_baudrate != baudrate:
+            print(f"Setting bus baud rate to {baudrate}. Previously {present_bus_baudrate}.")
+            self.port_handler.setBaudRate(baudrate)
+
+            if self.port_handler.getBaudRate() != baudrate:
+                raise OSError("Failed to write bus baud rate.")
+
+    @property
+    def motor_names(self) -> list[str]:
+        return list(self.motors.keys())
+
+    @property
+    def motor_models(self) -> list[str]:
+        return [model for _, model in self.motors.values()]
+
+    @property
+    def motor_indices(self) -> list[int]:
+        return [idx for idx, _ in self.motors.values()]
+
+    def set_calibration(self, calibration: dict[str, list]):
+        self.calibration = calibration
+
+    def apply_calibration_autocorrect(self, values: np.ndarray | list, motor_names: list[str] | None):
+        """This function apply the calibration, automatically detects out of range errors for motors values and attempt to correct.
+
+        For more info, see docstring of `apply_calibration` and `autocorrect_calibration`.
+        """
+        try:
+            values = self.apply_calibration(values, motor_names)
+        except JointOutOfRangeError as e:
+            print(e)
+            self.autocorrect_calibration(values, motor_names)
+            values = self.apply_calibration(values, motor_names)
+        return values
+
+    def apply_calibration(self, values: np.ndarray | list, motor_names: list[str] | None):
+        """Convert from unsigned int32 joint position range [0, 2**32[ to the universal float32 nominal degree range ]-180.0, 180.0[ with
+        a "zero position" at 0 degree.
+
+        Note: We say "nominal degree range" since the motors can take values outside this range. For instance, 190 degrees, if the motor
+        rotate more than a half a turn from the zero position. However, most motors can't rotate more than 180 degrees and will stay in this range.
+
+        Joints values are original in [0, 2**32[ (unsigned int32). Each motor are expected to complete a full rotation
+        when given a goal position that is + or - their resolution. For instance, feetech xl330-m077 have a resolution of 4096, and
+        at any position in their original range, let's say the position 56734, they complete a full rotation clockwise by moving to 60830,
+        or anticlockwise by moving to 52638. The position in the original range is arbitrary and might change a lot between each motor.
+        To harmonize between motors of the same model, different robots, or even models of different brands, we propose to work
+        in the centered nominal degree range ]-180, 180[.
+        """
+        if motor_names is None:
+            motor_names = self.motor_names
+
+        # Convert from unsigned int32 original range [0, 2**32] to signed float32 range
+        values = values.astype(np.float32)
+
+        for i, name in enumerate(motor_names):
+            calib_idx = self.calibration["motor_names"].index(name)
+            calib_mode = self.calibration["calib_mode"][calib_idx]
+
+            if CalibrationMode[calib_mode] == CalibrationMode.DEGREE:
+                drive_mode = self.calibration["drive_mode"][calib_idx]
+                homing_offset = self.calibration["homing_offset"][calib_idx]
+                _, model = self.motors[name]
+                resolution = self.model_resolution[model]
+
+                # Update direction of rotation of the motor to match between leader and follower.
+                # In fact, the motor of the leader for a given joint can be assembled in an
+                # opposite direction in term of rotation than the motor of the follower on the same joint.
+                if drive_mode:
+                    values[i] *= -1
+
+                # Convert from range [-2**31, 2**31[ to
+                # nominal range ]-resolution, resolution[ (e.g. ]-2048, 2048[)
+                values[i] += homing_offset
+
+                # Convert from range ]-resolution, resolution[ to
+                # universal float32 centered degree range ]-180, 180[
+                values[i] = values[i] / (resolution // 2) * HALF_TURN_DEGREE
+
+                if (values[i] < LOWER_BOUND_DEGREE) or (values[i] > UPPER_BOUND_DEGREE):
+                    raise JointOutOfRangeError(
+                        f"Wrong motor position range detected for {name}. "
+                        f"Expected to be in nominal range of [-{HALF_TURN_DEGREE}, {HALF_TURN_DEGREE}] degrees (a full rotation), "
+                        f"with a maximum range of [{LOWER_BOUND_DEGREE}, {UPPER_BOUND_DEGREE}] degrees to account for joints that can rotate a bit more, "
+                        f"but present value is {values[i]} degree. "
+                        "This might be due to a cable connection issue creating an artificial 360 degrees jump in motor values. "
+                        "You need to recalibrate by running: `python lerobot/scripts/control_robot.py calibrate`"
+                    )
+
+            elif CalibrationMode[calib_mode] == CalibrationMode.LINEAR:
+                start_pos = self.calibration["start_pos"][calib_idx]
+                end_pos = self.calibration["end_pos"][calib_idx]
+
+                # Rescale the present position to a nominal range [0, 100] %,
+                # useful for joints with linear motions like Aloha gripper
+                values[i] = (values[i] - start_pos) / (end_pos - start_pos) * 100
+
+                if (values[i] < LOWER_BOUND_LINEAR) or (values[i] > UPPER_BOUND_LINEAR):
+                    raise JointOutOfRangeError(
+                        f"Wrong motor position range detected for {name}. "
+                        f"Expected to be in nominal range of [0, 100] % (a full linear translation), "
+                        f"with a maximum range of [{LOWER_BOUND_LINEAR}, {UPPER_BOUND_LINEAR}] % to account for some imprecision during calibration, "
+                        f"but present value is {values[i]} %. "
+                        "This might be due to a cable connection issue creating an artificial jump in motor values. "
+                        "You need to recalibrate by running: `python lerobot/scripts/control_robot.py calibrate`"
+                    )
+
+        return values
+
+    def autocorrect_calibration(self, values: np.ndarray | list, motor_names: list[str] | None):
+        """This function automatically detects issues with values of motors after calibration, and correct for these issues.
+
+        Some motors might have values outside of expected maximum bounds after calibration.
+        For instance, for a joint in degree, its value can be outside [-270, 270] degrees, which is totally unexpected given
+        a nominal range of [-180, 180] degrees, which represents half a turn to the left or right starting from zero position.
+
+        Known issues:
+        #1: Motor value randomly shifts of a full turn, caused by hardware/connection errors.
+        #2: Motor internal homing offset is shifted of a full turn, caused by using default calibration (e.g Aloha).
+        #3: motor internal homing offset is shifted of less or more than a full turn, caused by using default calibration
+            or by human error during manual calibration.
+
+        Issues #1 and #2 can be solved by shifting the calibration homing offset by a full turn.
+        Issue #3 will be visually detected by user and potentially captured by the safety feature `max_relative_target`,
+        that will slow down the motor, raise an error asking to recalibrate. Manual recalibrating will solve the issue.
+
+        Note: A full turn corresponds to 360 degrees but also to 4096 steps for a motor resolution of 4096.
+        """
+        if motor_names is None:
+            motor_names = self.motor_names
+
+        # Convert from unsigned int32 original range [0, 2**32] to signed float32 range
+        values = values.astype(np.float32)
+
+        for i, name in enumerate(motor_names):
+            calib_idx = self.calibration["motor_names"].index(name)
+            calib_mode = self.calibration["calib_mode"][calib_idx]
+
+            if CalibrationMode[calib_mode] == CalibrationMode.DEGREE:
+                drive_mode = self.calibration["drive_mode"][calib_idx]
+                homing_offset = self.calibration["homing_offset"][calib_idx]
+                _, model = self.motors[name]
+                resolution = self.model_resolution[model]
+
+                if drive_mode:
+                    values[i] *= -1
+
+                # Convert from initial range to range [-180, 180] degrees
+                calib_val = (values[i] + homing_offset) / (resolution // 2) * HALF_TURN_DEGREE
+                in_range = (calib_val > LOWER_BOUND_DEGREE) and (calib_val < UPPER_BOUND_DEGREE)
+
+                # Solve this inequality to find the factor to shift the range into [-180, 180] degrees
+                # values[i] = (values[i] + homing_offset + resolution * factor) / (resolution // 2) * HALF_TURN_DEGREE
+                # - HALF_TURN_DEGREE <= (values[i] + homing_offset + resolution * factor) / (resolution // 2) * HALF_TURN_DEGREE <= HALF_TURN_DEGREE
+                # (- HALF_TURN_DEGREE / HALF_TURN_DEGREE * (resolution // 2) - values[i] - homing_offset) / resolution <= factor <= (HALF_TURN_DEGREE / 180 * (resolution // 2) - values[i] - homing_offset) / resolution
+                low_factor = (
+                    -HALF_TURN_DEGREE / HALF_TURN_DEGREE * (resolution // 2) - values[i] - homing_offset
+                ) / resolution
+                upp_factor = (
+                    HALF_TURN_DEGREE / HALF_TURN_DEGREE * (resolution // 2) - values[i] - homing_offset
+                ) / resolution
+
+            elif CalibrationMode[calib_mode] == CalibrationMode.LINEAR:
+                start_pos = self.calibration["start_pos"][calib_idx]
+                end_pos = self.calibration["end_pos"][calib_idx]
+
+                # Convert from initial range to range [0, 100] in %
+                calib_val = (values[i] - start_pos) / (end_pos - start_pos) * 100
+                in_range = (calib_val > LOWER_BOUND_LINEAR) and (calib_val < UPPER_BOUND_LINEAR)
+
+                # Solve this inequality to find the factor to shift the range into [0, 100] %
+                # values[i] = (values[i] - start_pos + resolution * factor) / (end_pos + resolution * factor - start_pos - resolution * factor) * 100
+                # values[i] = (values[i] - start_pos + resolution * factor) / (end_pos - start_pos) * 100
+                # 0 <= (values[i] - start_pos + resolution * factor) / (end_pos - start_pos) * 100 <= 100
+                # (start_pos - values[i]) / resolution <= factor <= (end_pos - values[i]) / resolution
+                low_factor = (start_pos - values[i]) / resolution
+                upp_factor = (end_pos - values[i]) / resolution
+
+            if not in_range:
+                # Get first integer between the two bounds
+                if low_factor < upp_factor:
+                    factor = math.ceil(low_factor)
+
+                    if factor > upp_factor:
+                        raise ValueError(f"No integer found between bounds [{low_factor=}, {upp_factor=}]")
+                else:
+                    factor = math.ceil(upp_factor)
+
+                    if factor > low_factor:
+                        raise ValueError(f"No integer found between bounds [{low_factor=}, {upp_factor=}]")
+
+                if CalibrationMode[calib_mode] == CalibrationMode.DEGREE:
+                    out_of_range_str = f"{LOWER_BOUND_DEGREE} < {calib_val} < {UPPER_BOUND_DEGREE} degrees"
+                    in_range_str = f"{LOWER_BOUND_DEGREE} < {calib_val} < {UPPER_BOUND_DEGREE} degrees"
+                elif CalibrationMode[calib_mode] == CalibrationMode.LINEAR:
+                    out_of_range_str = f"{LOWER_BOUND_LINEAR} < {calib_val} < {UPPER_BOUND_LINEAR} %"
+                    in_range_str = f"{LOWER_BOUND_LINEAR} < {calib_val} < {UPPER_BOUND_LINEAR} %"
+
+                logging.warning(
+                    f"Auto-correct calibration of motor '{name}' by shifting value by {abs(factor)} full turns, "
+                    f"from '{out_of_range_str}' to '{in_range_str}'."
+                )
+
+                # A full turn corresponds to 360 degrees but also to 4096 steps for a motor resolution of 4096.
+                self.calibration["homing_offset"][calib_idx] += resolution * factor
+
+    def revert_calibration(self, values: np.ndarray | list, motor_names: list[str] | None):
+        """Inverse of `apply_calibration`."""
+        if motor_names is None:
+            motor_names = self.motor_names
+
+        for i, name in enumerate(motor_names):
+            calib_idx = self.calibration["motor_names"].index(name)
+            calib_mode = self.calibration["calib_mode"][calib_idx]
+
+            if CalibrationMode[calib_mode] == CalibrationMode.DEGREE:
+                drive_mode = self.calibration["drive_mode"][calib_idx]
+                homing_offset = self.calibration["homing_offset"][calib_idx]
+                _, model = self.motors[name]
+                resolution = self.model_resolution[model]
+
+                # Convert from nominal 0-centered degree range [-180, 180] to
+                # 0-centered resolution range (e.g. [-2048, 2048] for resolution=4096)
+                values[i] = values[i] / HALF_TURN_DEGREE * (resolution // 2)
+
+                # Subtract the homing offsets to come back to actual motor range of values
+                # which can be arbitrary.
+                values[i] -= homing_offset
+
+                # Remove drive mode, which is the rotation direction of the motor, to come back to
+                # actual motor rotation direction which can be arbitrary.
+                if drive_mode:
+                    values[i] *= -1
+
+            elif CalibrationMode[calib_mode] == CalibrationMode.LINEAR:
+                start_pos = self.calibration["start_pos"][calib_idx]
+                end_pos = self.calibration["end_pos"][calib_idx]
+
+                # Convert from nominal lnear range of [0, 100] % to
+                # actual motor range of values which can be arbitrary.
+                values[i] = values[i] / 100 * (end_pos - start_pos) + start_pos
+
+        values = np.round(values).astype(np.int32)
+        return values
+
+    def avoid_rotation_reset(self, values, motor_names, data_name):
+        if data_name not in self.track_positions:
+            self.track_positions[data_name] = {
+                "prev": [None] * len(self.motor_names),
+                # Assume False at initialization
+                "below_zero": [False] * len(self.motor_names),
+                "above_max": [False] * len(self.motor_names),
+            }
+
+        track = self.track_positions[data_name]
+
+        if motor_names is None:
+            motor_names = self.motor_names
+
+        for i, name in enumerate(motor_names):
+            idx = self.motor_names.index(name)
+
+            if track["prev"][idx] is None:
+                track["prev"][idx] = values[i]
+                continue
+
+            # Detect a full rotation occurred
+            if abs(track["prev"][idx] - values[i]) > 2048:
+                # Position went below 0 and got reset to 4095
+                if track["prev"][idx] < values[i]:
+                    # So we set negative value by adding a full rotation
+                    values[i] -= 4096
+
+                # Position went above 4095 and got reset to 0
+                elif track["prev"][idx] > values[i]:
+                    # So we add a full rotation
+                    values[i] += 4096
+
+            track["prev"][idx] = values[i]
+
+        return values
+
+    def read_with_motor_ids(self, motor_models, motor_ids, data_name, num_retry=NUM_READ_RETRY):
+        if self.mock:
+            import tests.motors.mock_scservo_sdk as scs
+        else:
+            import scservo_sdk as scs
+
+        return_list = True
+        if not isinstance(motor_ids, list):
+            return_list = False
+            motor_ids = [motor_ids]
+
+        assert_same_address(self.model_ctrl_table, self.motor_models, data_name)
+        addr, bytes = self.model_ctrl_table[motor_models[0]][data_name]
+        group = scs.GroupSyncRead(self.port_handler, self.packet_handler, addr, bytes)
+        for idx in motor_ids:
+            group.addParam(idx)
+
+        for _ in range(num_retry):
+            comm = group.txRxPacket()
+            if comm == scs.COMM_SUCCESS:
+                break
+
+        if comm != scs.COMM_SUCCESS:
+            raise ConnectionError(
+                f"Read failed due to communication error on port {self.port_handler.port_name} for indices {motor_ids}: "
+                f"{self.packet_handler.getTxRxResult(comm)}"
+            )
+
+        values = []
+        for idx in motor_ids:
+            value = group.getData(idx, addr, bytes)
+            values.append(value)
+
+        if return_list:
+            return values
+        else:
+            return values[0]
+
+    def read(self, data_name, motor_names: str | list[str] | None = None):
+        if self.mock:
+            import tests.motors.mock_scservo_sdk as scs
+        else:
+            import scservo_sdk as scs
+
+        if not self.is_connected:
+            raise RobotDeviceNotConnectedError(
+                f"FeetechMotorsBus({self.port}) is not connected. You need to run `motors_bus.connect()`."
+            )
+
+        start_time = time.perf_counter()
+
+        if motor_names is None:
+            motor_names = self.motor_names
+
+        if isinstance(motor_names, str):
+            motor_names = [motor_names]
+
+        motor_ids = []
+        models = []
+        for name in motor_names:
+            motor_idx, model = self.motors[name]
+            motor_ids.append(motor_idx)
+            models.append(model)
+
+        assert_same_address(self.model_ctrl_table, models, data_name)
+        addr, bytes = self.model_ctrl_table[model][data_name]
+        group_key = get_group_sync_key(data_name, motor_names)
+
+        if data_name not in self.group_readers:
+            # Very Important to flush the buffer!
+            self.port_handler.ser.reset_output_buffer()
+            self.port_handler.ser.reset_input_buffer()
+
+            # create new group reader
+            self.group_readers[group_key] = scs.GroupSyncRead(
+                self.port_handler, self.packet_handler, addr, bytes
+            )
+            for idx in motor_ids:
+                self.group_readers[group_key].addParam(idx)
+
+        for _ in range(NUM_READ_RETRY):
+            comm = self.group_readers[group_key].txRxPacket()
+            if comm == scs.COMM_SUCCESS:
+                break
+
+        if comm != scs.COMM_SUCCESS:
+            raise ConnectionError(
+                f"Read failed due to communication error on port {self.port} for group_key {group_key}: "
+                f"{self.packet_handler.getTxRxResult(comm)}"
+            )
+
+        values = []
+        for idx in motor_ids:
+            value = self.group_readers[group_key].getData(idx, addr, bytes)
+            values.append(value)
+
+        values = np.array(values)
+
+        # Convert to signed int to use range [-2048, 2048] for our motor positions.
+        if data_name in CONVERT_UINT32_TO_INT32_REQUIRED:
+            values = values.astype(np.int32)
+
+        if data_name in CALIBRATION_REQUIRED:
+            values = self.avoid_rotation_reset(values, motor_names, data_name)
+
+        if data_name in CALIBRATION_REQUIRED and self.calibration is not None:
+            values = self.apply_calibration_autocorrect(values, motor_names)
+
+        # log the number of seconds it took to read the data from the motors
+        delta_ts_name = get_log_name("delta_timestamp_s", "read", data_name, motor_names)
+        self.logs[delta_ts_name] = time.perf_counter() - start_time
+
+        # log the utc time at which the data was received
+        ts_utc_name = get_log_name("timestamp_utc", "read", data_name, motor_names)
+        self.logs[ts_utc_name] = capture_timestamp_utc()
+
+        return values
+
+    def write_with_motor_ids(self, motor_models, motor_ids, data_name, values, num_retry=NUM_WRITE_RETRY):
+        if self.mock:
+            import tests.motors.mock_scservo_sdk as scs
+        else:
+            import scservo_sdk as scs
+
+        if not isinstance(motor_ids, list):
+            motor_ids = [motor_ids]
+        if not isinstance(values, list):
+            values = [values]
+
+        assert_same_address(self.model_ctrl_table, motor_models, data_name)
+        addr, bytes = self.model_ctrl_table[motor_models[0]][data_name]
+        group = scs.GroupSyncWrite(self.port_handler, self.packet_handler, addr, bytes)
+        for idx, value in zip(motor_ids, values, strict=True):
+            data = convert_to_bytes(value, bytes, self.mock)
+            group.addParam(idx, data)
+
+        for _ in range(num_retry):
+            comm = group.txPacket()
+            if comm == scs.COMM_SUCCESS:
+                break
+
+        if comm != scs.COMM_SUCCESS:
+            raise ConnectionError(
+                f"Write failed due to communication error on port {self.port_handler.port_name} for indices {motor_ids}: "
+                f"{self.packet_handler.getTxRxResult(comm)}"
+            )
+
+    def write(self, data_name, values: int | float | np.ndarray, motor_names: str | list[str] | None = None):
+        if not self.is_connected:
+            raise RobotDeviceNotConnectedError(
+                f"FeetechMotorsBus({self.port}) is not connected. You need to run `motors_bus.connect()`."
+            )
+
+        start_time = time.perf_counter()
+
+        if self.mock:
+            import tests.motors.mock_scservo_sdk as scs
+        else:
+            import scservo_sdk as scs
+
+        if motor_names is None:
+            motor_names = self.motor_names
+
+        if isinstance(motor_names, str):
+            motor_names = [motor_names]
+
+        if isinstance(values, (int, float, np.integer)):
+            values = [int(values)] * len(motor_names)
+
+        values = np.array(values)
+
+        motor_ids = []
+        models = []
+        for name in motor_names:
+            motor_idx, model = self.motors[name]
+            motor_ids.append(motor_idx)
+            models.append(model)
+
+        if data_name in CALIBRATION_REQUIRED and self.calibration is not None:
+            values = self.revert_calibration(values, motor_names)
+
+        values = values.tolist()
+
+        assert_same_address(self.model_ctrl_table, models, data_name)
+        addr, bytes = self.model_ctrl_table[model][data_name]
+        group_key = get_group_sync_key(data_name, motor_names)
+
+        init_group = data_name not in self.group_readers
+        if init_group:
+            self.group_writers[group_key] = scs.GroupSyncWrite(
+                self.port_handler, self.packet_handler, addr, bytes
+            )
+
+        for idx, value in zip(motor_ids, values, strict=True):
+            data = convert_to_bytes(value, bytes, self.mock)
+            if init_group:
+                self.group_writers[group_key].addParam(idx, data)
+            else:
+                self.group_writers[group_key].changeParam(idx, data)
+
+        comm = self.group_writers[group_key].txPacket()
+        if comm != scs.COMM_SUCCESS:
+            raise ConnectionError(
+                f"Write failed due to communication error on port {self.port} for group_key {group_key}: "
+                f"{self.packet_handler.getTxRxResult(comm)}"
+            )
+
+        # log the number of seconds it took to write the data to the motors
+        delta_ts_name = get_log_name("delta_timestamp_s", "write", data_name, motor_names)
+        self.logs[delta_ts_name] = time.perf_counter() - start_time
+
+        # TODO(rcadene): should we log the time before sending the write command?
+        # log the utc time when the write has been completed
+        ts_utc_name = get_log_name("timestamp_utc", "write", data_name, motor_names)
+        self.logs[ts_utc_name] = capture_timestamp_utc()
+
+    def disconnect(self):
+        if not self.is_connected:
+            raise RobotDeviceNotConnectedError(
+                f"FeetechMotorsBus({self.port}) is not connected. Try running `motors_bus.connect()` first."
+            )
+
+        if self.port_handler is not None:
+            self.port_handler.closePort()
+            self.port_handler = None
+
+        self.packet_handler = None
+        self.group_readers = {}
+        self.group_writers = {}
+        self.is_connected = False
+
+    def __del__(self):
+        if getattr(self, "is_connected", False):
+            self.disconnect()

lerobot/common/robot_devices/motors/utils.py

@@ -0,0 +1,67 @@
+# 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 typing import Protocol
+
+from lerobot.common.robot_devices.motors.configs import (
+    DynamixelMotorsBusConfig,
+    FeetechMotorsBusConfig,
+    MotorsBusConfig,
+)
+
+
+class MotorsBus(Protocol):
+    def motor_names(self): ...
+    def set_calibration(self): ...
+    def apply_calibration(self): ...
+    def revert_calibration(self): ...
+    def read(self): ...
+    def write(self): ...
+
+
+def make_motors_buses_from_configs(motors_bus_configs: dict[str, MotorsBusConfig]) -> list[MotorsBus]:
+    motors_buses = {}
+
+    for key, cfg in motors_bus_configs.items():
+        if cfg.type == "dynamixel":
+            from lerobot.common.robot_devices.motors.dynamixel import DynamixelMotorsBus
+
+            motors_buses[key] = DynamixelMotorsBus(cfg)
+
+        elif cfg.type == "feetech":
+            from lerobot.common.robot_devices.motors.feetech import FeetechMotorsBus
+
+            motors_buses[key] = FeetechMotorsBus(cfg)
+
+        else:
+            raise ValueError(f"The motor type '{cfg.type}' is not valid.")
+
+    return motors_buses
+
+
+def make_motors_bus(motor_type: str, **kwargs) -> MotorsBus:
+    if motor_type == "dynamixel":
+        from lerobot.common.robot_devices.motors.dynamixel import DynamixelMotorsBus
+
+        config = DynamixelMotorsBusConfig(**kwargs)
+        return DynamixelMotorsBus(config)
+
+    elif motor_type == "feetech":
+        from lerobot.common.robot_devices.motors.feetech import FeetechMotorsBus
+
+        config = FeetechMotorsBusConfig(**kwargs)
+        return FeetechMotorsBus(config)
+
+    else:
+        raise ValueError(f"The motor type '{motor_type}' is not valid.")

lerobot/common/robot_devices/robots/configs.py

@@ -0,0 +1,613 @@
+# 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 abc
+from dataclasses import dataclass, field
+from typing import Sequence
+
+import draccus
+
+from lerobot.common.robot_devices.cameras.configs import (
+    CameraConfig,
+    IntelRealSenseCameraConfig,
+    OpenCVCameraConfig,
+)
+from lerobot.common.robot_devices.motors.configs import (
+    DynamixelMotorsBusConfig,
+    FeetechMotorsBusConfig,
+    MotorsBusConfig,
+)
+
+
+@dataclass
+class RobotConfig(draccus.ChoiceRegistry, abc.ABC):
+    @property
+    def type(self) -> str:
+        return self.get_choice_name(self.__class__)
+
+
+# TODO(rcadene, aliberts): remove ManipulatorRobotConfig abstraction
+@dataclass
+class ManipulatorRobotConfig(RobotConfig):
+    leader_arms: dict[str, MotorsBusConfig] = field(default_factory=lambda: {})
+    follower_arms: dict[str, MotorsBusConfig] = field(default_factory=lambda: {})
+    cameras: dict[str, CameraConfig] = field(default_factory=lambda: {})
+
+    # Optionally limit the magnitude of the relative positional target vector for safety purposes.
+    # Set this to a positive scalar to have the same value for all motors, or a list that is the same length
+    # as the number of motors in your follower arms (assumes all follower arms have the same number of
+    # motors).
+    max_relative_target: list[float] | float | None = None
+
+    # Optionally set the leader arm in torque mode with the gripper motor set to this angle. This makes it
+    # possible to squeeze the gripper and have it spring back to an open position on its own. If None, the
+    # gripper is not put in torque mode.
+    gripper_open_degree: float | None = None
+
+    mock: bool = False
+
+    def __post_init__(self):
+        if self.mock:
+            for arm in self.leader_arms.values():
+                if not arm.mock:
+                    arm.mock = True
+            for arm in self.follower_arms.values():
+                if not arm.mock:
+                    arm.mock = True
+            for cam in self.cameras.values():
+                if not cam.mock:
+                    cam.mock = True
+
+        if self.max_relative_target is not None and isinstance(self.max_relative_target, Sequence):
+            for name in self.follower_arms:
+                if len(self.follower_arms[name].motors) != len(self.max_relative_target):
+                    raise ValueError(
+                        f"len(max_relative_target)={len(self.max_relative_target)} but the follower arm with name {name} has "
+                        f"{len(self.follower_arms[name].motors)} motors. Please make sure that the "
+                        f"`max_relative_target` list has as many parameters as there are motors per arm. "
+                        "Note: This feature does not yet work with robots where different follower arms have "
+                        "different numbers of motors."
+                    )
+
+
+@RobotConfig.register_subclass("aloha")
+@dataclass
+class AlohaRobotConfig(ManipulatorRobotConfig):
+    # Specific to Aloha, LeRobot comes with default calibration files. Assuming the motors have been
+    # properly assembled, no manual calibration step is expected. If you need to run manual calibration,
+    # simply update this path to ".cache/calibration/aloha"
+    calibration_dir: str = ".cache/calibration/aloha_default"
+
+    # /!\ FOR SAFETY, READ THIS /!\
+    # `max_relative_target` limits the magnitude of the relative positional target vector for safety purposes.
+    # Set this to a positive scalar to have the same value for all motors, or a list that is the same length as
+    # the number of motors in your follower arms.
+    # For Aloha, for every goal position request, motor rotations are capped at 5 degrees by default.
+    # When you feel more confident with teleoperation or running the policy, you can extend
+    # this safety limit and even removing it by setting it to `null`.
+    # Also, everything is expected to work safely out-of-the-box, but we highly advise to
+    # first try to teleoperate the grippers only (by commenting out the rest of the motors in this yaml),
+    # then to gradually add more motors (by uncommenting), until you can teleoperate both arms fully
+    max_relative_target: int | None = 5
+
+    leader_arms: dict[str, MotorsBusConfig] = field(
+        default_factory=lambda: {
+            "left": DynamixelMotorsBusConfig(
+                # window_x
+                port="/dev/ttyDXL_leader_left",
+                motors={
+                    # name: (index, model)
+                    "waist": [1, "xm430-w350"],
+                    "shoulder": [2, "xm430-w350"],
+                    "shoulder_shadow": [3, "xm430-w350"],
+                    "elbow": [4, "xm430-w350"],
+                    "elbow_shadow": [5, "xm430-w350"],
+                    "forearm_roll": [6, "xm430-w350"],
+                    "wrist_angle": [7, "xm430-w350"],
+                    "wrist_rotate": [8, "xl430-w250"],
+                    "gripper": [9, "xc430-w150"],
+                },
+            ),
+            "right": DynamixelMotorsBusConfig(
+                # window_x
+                port="/dev/ttyDXL_leader_right",
+                motors={
+                    # name: (index, model)
+                    "waist": [1, "xm430-w350"],
+                    "shoulder": [2, "xm430-w350"],
+                    "shoulder_shadow": [3, "xm430-w350"],
+                    "elbow": [4, "xm430-w350"],
+                    "elbow_shadow": [5, "xm430-w350"],
+                    "forearm_roll": [6, "xm430-w350"],
+                    "wrist_angle": [7, "xm430-w350"],
+                    "wrist_rotate": [8, "xl430-w250"],
+                    "gripper": [9, "xc430-w150"],
+                },
+            ),
+        }
+    )
+
+    follower_arms: dict[str, MotorsBusConfig] = field(
+        default_factory=lambda: {
+            "left": DynamixelMotorsBusConfig(
+                port="/dev/ttyDXL_follower_left",
+                motors={
+                    # name: (index, model)
+                    "waist": [1, "xm540-w270"],
+                    "shoulder": [2, "xm540-w270"],
+                    "shoulder_shadow": [3, "xm540-w270"],
+                    "elbow": [4, "xm540-w270"],
+                    "elbow_shadow": [5, "xm540-w270"],
+                    "forearm_roll": [6, "xm540-w270"],
+                    "wrist_angle": [7, "xm540-w270"],
+                    "wrist_rotate": [8, "xm430-w350"],
+                    "gripper": [9, "xm430-w350"],
+                },
+            ),
+            "right": DynamixelMotorsBusConfig(
+                port="/dev/ttyDXL_follower_right",
+                motors={
+                    # name: (index, model)
+                    "waist": [1, "xm540-w270"],
+                    "shoulder": [2, "xm540-w270"],
+                    "shoulder_shadow": [3, "xm540-w270"],
+                    "elbow": [4, "xm540-w270"],
+                    "elbow_shadow": [5, "xm540-w270"],
+                    "forearm_roll": [6, "xm540-w270"],
+                    "wrist_angle": [7, "xm540-w270"],
+                    "wrist_rotate": [8, "xm430-w350"],
+                    "gripper": [9, "xm430-w350"],
+                },
+            ),
+        }
+    )
+
+    # Troubleshooting: If one of your IntelRealSense cameras freeze during
+    # data recording due to bandwidth limit, you might need to plug the camera
+    # on another USB hub or PCIe card.
+    cameras: dict[str, CameraConfig] = field(
+        default_factory=lambda: {
+            "cam_high": IntelRealSenseCameraConfig(
+                serial_number=128422271347,
+                fps=30,
+                width=640,
+                height=480,
+            ),
+            "cam_low": IntelRealSenseCameraConfig(
+                serial_number=130322270656,
+                fps=30,
+                width=640,
+                height=480,
+            ),
+            "cam_left_wrist": IntelRealSenseCameraConfig(
+                serial_number=218622272670,
+                fps=30,
+                width=640,
+                height=480,
+            ),
+            "cam_right_wrist": IntelRealSenseCameraConfig(
+                serial_number=130322272300,
+                fps=30,
+                width=640,
+                height=480,
+            ),
+        }
+    )
+
+    mock: bool = False
+
+
+@RobotConfig.register_subclass("koch")
+@dataclass
+class KochRobotConfig(ManipulatorRobotConfig):
+    calibration_dir: str = ".cache/calibration/koch"
+    # `max_relative_target` limits the magnitude of the relative positional target vector for safety purposes.
+    # Set this to a positive scalar to have the same value for all motors, or a list that is the same length as
+    # the number of motors in your follower arms.
+    max_relative_target: int | None = None
+
+    leader_arms: dict[str, MotorsBusConfig] = field(
+        default_factory=lambda: {
+            "main": DynamixelMotorsBusConfig(
+                port="/dev/tty.usbmodem585A0085511",
+                motors={
+                    # name: (index, model)
+                    "shoulder_pan": [1, "xl330-m077"],
+                    "shoulder_lift": [2, "xl330-m077"],
+                    "elbow_flex": [3, "xl330-m077"],
+                    "wrist_flex": [4, "xl330-m077"],
+                    "wrist_roll": [5, "xl330-m077"],
+                    "gripper": [6, "xl330-m077"],
+                },
+            ),
+        }
+    )
+
+    follower_arms: dict[str, MotorsBusConfig] = field(
+        default_factory=lambda: {
+            "main": DynamixelMotorsBusConfig(
+                port="/dev/tty.usbmodem585A0076891",
+                motors={
+                    # name: (index, model)
+                    "shoulder_pan": [1, "xl430-w250"],
+                    "shoulder_lift": [2, "xl430-w250"],
+                    "elbow_flex": [3, "xl330-m288"],
+                    "wrist_flex": [4, "xl330-m288"],
+                    "wrist_roll": [5, "xl330-m288"],
+                    "gripper": [6, "xl330-m288"],
+                },
+            ),
+        }
+    )
+
+    cameras: dict[str, CameraConfig] = field(
+        default_factory=lambda: {
+            "laptop": OpenCVCameraConfig(
+                camera_index=0,
+                fps=30,
+                width=640,
+                height=480,
+            ),
+            "phone": OpenCVCameraConfig(
+                camera_index=1,
+                fps=30,
+                width=640,
+                height=480,
+            ),
+        }
+    )
+
+    # ~ Koch specific settings ~
+    # Sets the leader arm in torque mode with the gripper motor set to this angle. This makes it possible
+    # to squeeze the gripper and have it spring back to an open position on its own.
+    gripper_open_degree: float = 35.156
+
+    mock: bool = False
+
+
+@RobotConfig.register_subclass("koch_bimanual")
+@dataclass
+class KochBimanualRobotConfig(ManipulatorRobotConfig):
+    calibration_dir: str = ".cache/calibration/koch_bimanual"
+    # `max_relative_target` limits the magnitude of the relative positional target vector for safety purposes.
+    # Set this to a positive scalar to have the same value for all motors, or a list that is the same length as
+    # the number of motors in your follower arms.
+    max_relative_target: int | None = None
+
+    leader_arms: dict[str, MotorsBusConfig] = field(
+        default_factory=lambda: {
+            "left": DynamixelMotorsBusConfig(
+                port="/dev/tty.usbmodem585A0085511",
+                motors={
+                    # name: (index, model)
+                    "shoulder_pan": [1, "xl330-m077"],
+                    "shoulder_lift": [2, "xl330-m077"],
+                    "elbow_flex": [3, "xl330-m077"],
+                    "wrist_flex": [4, "xl330-m077"],
+                    "wrist_roll": [5, "xl330-m077"],
+                    "gripper": [6, "xl330-m077"],
+                },
+            ),
+            "right": DynamixelMotorsBusConfig(
+                port="/dev/tty.usbmodem575E0031751",
+                motors={
+                    # name: (index, model)
+                    "shoulder_pan": [1, "xl330-m077"],
+                    "shoulder_lift": [2, "xl330-m077"],
+                    "elbow_flex": [3, "xl330-m077"],
+                    "wrist_flex": [4, "xl330-m077"],
+                    "wrist_roll": [5, "xl330-m077"],
+                    "gripper": [6, "xl330-m077"],
+                },
+            ),
+        }
+    )
+
+    follower_arms: dict[str, MotorsBusConfig] = field(
+        default_factory=lambda: {
+            "left": DynamixelMotorsBusConfig(
+                port="/dev/tty.usbmodem585A0076891",
+                motors={
+                    # name: (index, model)
+                    "shoulder_pan": [1, "xl430-w250"],
+                    "shoulder_lift": [2, "xl430-w250"],
+                    "elbow_flex": [3, "xl330-m288"],
+                    "wrist_flex": [4, "xl330-m288"],
+                    "wrist_roll": [5, "xl330-m288"],
+                    "gripper": [6, "xl330-m288"],
+                },
+            ),
+            "right": DynamixelMotorsBusConfig(
+                port="/dev/tty.usbmodem575E0032081",
+                motors={
+                    # name: (index, model)
+                    "shoulder_pan": [1, "xl430-w250"],
+                    "shoulder_lift": [2, "xl430-w250"],
+                    "elbow_flex": [3, "xl330-m288"],
+                    "wrist_flex": [4, "xl330-m288"],
+                    "wrist_roll": [5, "xl330-m288"],
+                    "gripper": [6, "xl330-m288"],
+                },
+            ),
+        }
+    )
+
+    cameras: dict[str, CameraConfig] = field(
+        default_factory=lambda: {
+            "laptop": OpenCVCameraConfig(
+                camera_index=0,
+                fps=30,
+                width=640,
+                height=480,
+            ),
+            "phone": OpenCVCameraConfig(
+                camera_index=1,
+                fps=30,
+                width=640,
+                height=480,
+            ),
+        }
+    )
+
+    # ~ Koch specific settings ~
+    # Sets the leader arm in torque mode with the gripper motor set to this angle. This makes it possible
+    # to squeeze the gripper and have it spring back to an open position on its own.
+    gripper_open_degree: float = 35.156
+
+    mock: bool = False
+
+
+@RobotConfig.register_subclass("moss")
+@dataclass
+class MossRobotConfig(ManipulatorRobotConfig):
+    calibration_dir: str = ".cache/calibration/moss"
+    # `max_relative_target` limits the magnitude of the relative positional target vector for safety purposes.
+    # Set this to a positive scalar to have the same value for all motors, or a list that is the same length as
+    # the number of motors in your follower arms.
+    max_relative_target: int | None = None
+
+    leader_arms: dict[str, MotorsBusConfig] = field(
+        default_factory=lambda: {
+            "main": FeetechMotorsBusConfig(
+                port="/dev/tty.usbmodem58760431091",
+                motors={
+                    # name: (index, model)
+                    "shoulder_pan": [1, "sts3215"],
+                    "shoulder_lift": [2, "sts3215"],
+                    "elbow_flex": [3, "sts3215"],
+                    "wrist_flex": [4, "sts3215"],
+                    "wrist_roll": [5, "sts3215"],
+                    "gripper": [6, "sts3215"],
+                },
+            ),
+        }
+    )
+
+    follower_arms: dict[str, MotorsBusConfig] = field(
+        default_factory=lambda: {
+            "main": FeetechMotorsBusConfig(
+                port="/dev/tty.usbmodem585A0076891",
+                motors={
+                    # name: (index, model)
+                    "shoulder_pan": [1, "sts3215"],
+                    "shoulder_lift": [2, "sts3215"],
+                    "elbow_flex": [3, "sts3215"],
+                    "wrist_flex": [4, "sts3215"],
+                    "wrist_roll": [5, "sts3215"],
+                    "gripper": [6, "sts3215"],
+                },
+            ),
+        }
+    )
+
+    cameras: dict[str, CameraConfig] = field(
+        default_factory=lambda: {
+            "laptop": OpenCVCameraConfig(
+                camera_index=0,
+                fps=30,
+                width=640,
+                height=480,
+            ),
+            "phone": OpenCVCameraConfig(
+                camera_index=1,
+                fps=30,
+                width=640,
+                height=480,
+            ),
+        }
+    )
+
+    mock: bool = False
+
+
+@RobotConfig.register_subclass("so100")
+@dataclass
+class So100RobotConfig(ManipulatorRobotConfig):
+    calibration_dir: str = ".cache/calibration/so100"
+    # `max_relative_target` limits the magnitude of the relative positional target vector for safety purposes.
+    # Set this to a positive scalar to have the same value for all motors, or a list that is the same length as
+    # the number of motors in your follower arms.
+    max_relative_target: int | None = None
+
+    leader_arms: dict[str, MotorsBusConfig] = field(
+        default_factory=lambda: {
+            "main": FeetechMotorsBusConfig(
+                port="/dev/tty.usbmodem58760431091",
+                motors={
+                    # name: (index, model)
+                    "shoulder_pan": [1, "sts3215"],
+                    "shoulder_lift": [2, "sts3215"],
+                    "elbow_flex": [3, "sts3215"],
+                    "wrist_flex": [4, "sts3215"],
+                    "wrist_roll": [5, "sts3215"],
+                    "gripper": [6, "sts3215"],
+                },
+            ),
+        }
+    )
+
+    follower_arms: dict[str, MotorsBusConfig] = field(
+        default_factory=lambda: {
+            "main": FeetechMotorsBusConfig(
+                port="/dev/tty.usbmodem585A0076891",
+                motors={
+                    # name: (index, model)
+                    "shoulder_pan": [1, "sts3215"],
+                    "shoulder_lift": [2, "sts3215"],
+                    "elbow_flex": [3, "sts3215"],
+                    "wrist_flex": [4, "sts3215"],
+                    "wrist_roll": [5, "sts3215"],
+                    "gripper": [6, "sts3215"],
+                },
+            ),
+        }
+    )
+
+    cameras: dict[str, CameraConfig] = field(
+        default_factory=lambda: {
+            "laptop": OpenCVCameraConfig(
+                camera_index=0,
+                fps=30,
+                width=640,
+                height=480,
+            ),
+            "phone": OpenCVCameraConfig(
+                camera_index=1,
+                fps=30,
+                width=640,
+                height=480,
+            ),
+        }
+    )
+
+    mock: bool = False
+
+
+@RobotConfig.register_subclass("stretch")
+@dataclass
+class StretchRobotConfig(RobotConfig):
+    # `max_relative_target` limits the magnitude of the relative positional target vector for safety purposes.
+    # Set this to a positive scalar to have the same value for all motors, or a list that is the same length as
+    # the number of motors in your follower arms.
+    max_relative_target: int | None = None
+
+    cameras: dict[str, CameraConfig] = field(
+        default_factory=lambda: {
+            "navigation": OpenCVCameraConfig(
+                camera_index="/dev/hello-nav-head-camera",
+                fps=10,
+                width=1280,
+                height=720,
+                rotation=-90,
+            ),
+            "head": IntelRealSenseCameraConfig(
+                name="Intel RealSense D435I",
+                fps=30,
+                width=640,
+                height=480,
+                rotation=90,
+            ),
+            "wrist": IntelRealSenseCameraConfig(
+                name="Intel RealSense D405",
+                fps=30,
+                width=640,
+                height=480,
+            ),
+        }
+    )
+
+    mock: bool = False
+
+
+@RobotConfig.register_subclass("lekiwi")
+@dataclass
+class LeKiwiRobotConfig(RobotConfig):
+    # `max_relative_target` limits the magnitude of the relative positional target vector for safety purposes.
+    # Set this to a positive scalar to have the same value for all motors, or a list that is the same length as
+    # the number of motors in your follower arms.
+    max_relative_target: int | None = None
+
+    # Network Configuration
+    ip: str = "192.168.0.193"
+    port: int = 5555
+    video_port: int = 5556
+
+    cameras: dict[str, CameraConfig] = field(
+        default_factory=lambda: {
+            "front": OpenCVCameraConfig(
+                camera_index="/dev/video0", fps=30, width=640, height=480, rotation=90
+            ),
+            "wrist": OpenCVCameraConfig(
+                camera_index="/dev/video2", fps=30, width=640, height=480, rotation=180
+            ),
+        }
+    )
+
+    calibration_dir: str = ".cache/calibration/lekiwi"
+
+    leader_arms: dict[str, MotorsBusConfig] = field(
+        default_factory=lambda: {
+            "main": FeetechMotorsBusConfig(
+                port="/dev/tty.usbmodem585A0077581",
+                motors={
+                    # name: (index, model)
+                    "shoulder_pan": [1, "sts3215"],
+                    "shoulder_lift": [2, "sts3215"],
+                    "elbow_flex": [3, "sts3215"],
+                    "wrist_flex": [4, "sts3215"],
+                    "wrist_roll": [5, "sts3215"],
+                    "gripper": [6, "sts3215"],
+                },
+            ),
+        }
+    )
+
+    follower_arms: dict[str, MotorsBusConfig] = field(
+        default_factory=lambda: {
+            "main": FeetechMotorsBusConfig(
+                port="/dev/ttyACM0",
+                motors={
+                    # name: (index, model)
+                    "shoulder_pan": [1, "sts3215"],
+                    "shoulder_lift": [2, "sts3215"],
+                    "elbow_flex": [3, "sts3215"],
+                    "wrist_flex": [4, "sts3215"],
+                    "wrist_roll": [5, "sts3215"],
+                    "gripper": [6, "sts3215"],
+                    "left_wheel": (7, "sts3215"),
+                    "back_wheel": (8, "sts3215"),
+                    "right_wheel": (9, "sts3215"),
+                },
+            ),
+        }
+    )
+
+    teleop_keys: dict[str, str] = field(
+        default_factory=lambda: {
+            # Movement
+            "forward": "w",
+            "backward": "s",
+            "left": "a",
+            "right": "d",
+            "rotate_left": "z",
+            "rotate_right": "x",
+            # Speed control
+            "speed_up": "r",
+            "speed_down": "f",
+            # quit teleop
+            "quit": "q",
+        }
+    )
+
+    mock: bool = False

lerobot/common/robot_devices/robots/dynamixel_calibration.py

@@ -0,0 +1,144 @@
+# 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.
+
+"""Logic to calibrate a robot arm built with dynamixel motors"""
+# TODO(rcadene, aliberts): move this logic into the robot code when refactoring
+
+import numpy as np
+
+from lerobot.common.robot_devices.motors.dynamixel import (
+    CalibrationMode,
+    TorqueMode,
+    convert_degrees_to_steps,
+)
+from lerobot.common.robot_devices.motors.utils import MotorsBus
+
+URL_TEMPLATE = (
+    "https://raw.githubusercontent.com/huggingface/lerobot/main/media/{robot}/{arm}_{position}.webp"
+)
+
+# The following positions are provided in nominal degree range ]-180, +180[
+# For more info on these constants, see comments in the code where they get used.
+ZERO_POSITION_DEGREE = 0
+ROTATED_POSITION_DEGREE = 90
+
+
+def assert_drive_mode(drive_mode):
+    # `drive_mode` is in [0,1] with 0 means original rotation direction for the motor, and 1 means inverted.
+    if not np.all(np.isin(drive_mode, [0, 1])):
+        raise ValueError(f"`drive_mode` contains values other than 0 or 1: ({drive_mode})")
+
+
+def apply_drive_mode(position, drive_mode):
+    assert_drive_mode(drive_mode)
+    # Convert `drive_mode` from [0, 1] with 0 indicates original rotation direction and 1 inverted,
+    # to [-1, 1] with 1 indicates original rotation direction and -1 inverted.
+    signed_drive_mode = -(drive_mode * 2 - 1)
+    position *= signed_drive_mode
+    return position
+
+
+def compute_nearest_rounded_position(position, models):
+    delta_turn = convert_degrees_to_steps(ROTATED_POSITION_DEGREE, models)
+    nearest_pos = np.round(position.astype(float) / delta_turn) * delta_turn
+    return nearest_pos.astype(position.dtype)
+
+
+def run_arm_calibration(arm: MotorsBus, robot_type: str, arm_name: str, arm_type: str):
+    """This function ensures that a neural network trained on data collected on a given robot
+    can work on another robot. For instance before calibration, setting a same goal position
+    for each motor of two different robots will get two very different positions. But after calibration,
+    the two robots will move to the same position.To this end, this function computes the homing offset
+    and the drive mode for each motor of a given robot.
+
+    Homing offset is used to shift the motor position to a ]-2048, +2048[ nominal range (when the motor uses 2048 steps
+    to complete a half a turn). This range is set around an arbitrary "zero position" corresponding to all motor positions
+    being 0. During the calibration process, you will need to manually move the robot to this "zero position".
+
+    Drive mode is used to invert the rotation direction of the motor. This is useful when some motors have been assembled
+    in the opposite orientation for some robots. During the calibration process, you will need to manually move the robot
+    to the "rotated position".
+
+    After calibration, the homing offsets and drive modes are stored in a cache.
+
+    Example of usage:
+    ```python
+    run_arm_calibration(arm, "koch", "left", "follower")
+    ```
+    """
+    if (arm.read("Torque_Enable") != TorqueMode.DISABLED.value).any():
+        raise ValueError("To run calibration, the torque must be disabled on all motors.")
+
+    print(f"\nRunning calibration of {robot_type} {arm_name} {arm_type}...")
+
+    print("\nMove arm to zero position")
+    print("See: " + URL_TEMPLATE.format(robot=robot_type, arm=arm_type, position="zero"))
+    input("Press Enter to continue...")
+
+    # We arbitrarily chose our zero target position to be a straight horizontal position with gripper upwards and closed.
+    # It is easy to identify and all motors are in a "quarter turn" position. Once calibration is done, this position will
+    # correspond to every motor angle being 0. If you set all 0 as Goal Position, the arm will move in this position.
+    zero_target_pos = convert_degrees_to_steps(ZERO_POSITION_DEGREE, arm.motor_models)
+
+    # Compute homing offset so that `present_position + homing_offset ~= target_position`.
+    zero_pos = arm.read("Present_Position")
+    zero_nearest_pos = compute_nearest_rounded_position(zero_pos, arm.motor_models)
+    homing_offset = zero_target_pos - zero_nearest_pos
+
+    # The rotated target position corresponds to a rotation of a quarter turn from the zero position.
+    # This allows to identify the rotation direction of each motor.
+    # For instance, if the motor rotates 90 degree, and its value is -90 after applying the homing offset, then we know its rotation direction
+    # is inverted. However, for the calibration being successful, we need everyone to follow the same target position.
+    # Sometimes, there is only one possible rotation direction. For instance, if the gripper is closed, there is only one direction which
+    # corresponds to opening the gripper. When the rotation direction is ambiguous, we arbitrarily rotate clockwise from the point of view
+    # of the previous motor in the kinetic chain.
+    print("\nMove arm to rotated target position")
+    print("See: " + URL_TEMPLATE.format(robot=robot_type, arm=arm_type, position="rotated"))
+    input("Press Enter to continue...")
+
+    rotated_target_pos = convert_degrees_to_steps(ROTATED_POSITION_DEGREE, arm.motor_models)
+
+    # Find drive mode by rotating each motor by a quarter of a turn.
+    # Drive mode indicates if the motor rotation direction should be inverted (=1) or not (=0).
+    rotated_pos = arm.read("Present_Position")
+    drive_mode = (rotated_pos < zero_pos).astype(np.int32)
+
+    # Re-compute homing offset to take into account drive mode
+    rotated_drived_pos = apply_drive_mode(rotated_pos, drive_mode)
+    rotated_nearest_pos = compute_nearest_rounded_position(rotated_drived_pos, arm.motor_models)
+    homing_offset = rotated_target_pos - rotated_nearest_pos
+
+    print("\nMove arm to rest position")
+    print("See: " + URL_TEMPLATE.format(robot=robot_type, arm=arm_type, position="rest"))
+    input("Press Enter to continue...")
+    print()
+
+    # Joints with rotational motions are expressed in degrees in nominal range of [-180, 180]
+    calib_mode = [CalibrationMode.DEGREE.name] * len(arm.motor_names)
+
+    # TODO(rcadene): make type of joints (DEGREE or LINEAR) configurable from yaml?
+    if robot_type in ["aloha"] and "gripper" in arm.motor_names:
+        # Joints with linear motions (like gripper of Aloha) are expressed in nominal range of [0, 100]
+        calib_idx = arm.motor_names.index("gripper")
+        calib_mode[calib_idx] = CalibrationMode.LINEAR.name
+
+    calib_data = {
+        "homing_offset": homing_offset.tolist(),
+        "drive_mode": drive_mode.tolist(),
+        "start_pos": zero_pos.tolist(),
+        "end_pos": rotated_pos.tolist(),
+        "calib_mode": calib_mode,
+        "motor_names": arm.motor_names,
+    }
+    return calib_data

lerobot/common/robot_devices/robots/feetech_calibration.py

@@ -0,0 +1,498 @@
+# 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.
+
+"""Logic to calibrate a robot arm built with feetech motors"""
+# TODO(rcadene, aliberts): move this logic into the robot code when refactoring
+
+import time
+
+import numpy as np
+
+from lerobot.common.robot_devices.motors.feetech import (
+    CalibrationMode,
+    TorqueMode,
+    convert_degrees_to_steps,
+)
+from lerobot.common.robot_devices.motors.utils import MotorsBus
+
+URL_TEMPLATE = (
+    "https://raw.githubusercontent.com/huggingface/lerobot/main/media/{robot}/{arm}_{position}.webp"
+)
+
+# The following positions are provided in nominal degree range ]-180, +180[
+# For more info on these constants, see comments in the code where they get used.
+ZERO_POSITION_DEGREE = 0
+ROTATED_POSITION_DEGREE = 90
+
+
+def assert_drive_mode(drive_mode):
+    # `drive_mode` is in [0,1] with 0 means original rotation direction for the motor, and 1 means inverted.
+    if not np.all(np.isin(drive_mode, [0, 1])):
+        raise ValueError(f"`drive_mode` contains values other than 0 or 1: ({drive_mode})")
+
+
+def apply_drive_mode(position, drive_mode):
+    assert_drive_mode(drive_mode)
+    # Convert `drive_mode` from [0, 1] with 0 indicates original rotation direction and 1 inverted,
+    # to [-1, 1] with 1 indicates original rotation direction and -1 inverted.
+    signed_drive_mode = -(drive_mode * 2 - 1)
+    position *= signed_drive_mode
+    return position
+
+
+def move_until_block(arm, motor_name, positive_direction=True, while_move_hook=None):
+    count = 0
+    while True:
+        present_pos = arm.read("Present_Position", motor_name)
+        if positive_direction:
+            # Move +100 steps every time. Lower the steps to lower the speed at which the arm moves.
+            arm.write("Goal_Position", present_pos + 100, motor_name)
+        else:
+            arm.write("Goal_Position", present_pos - 100, motor_name)
+
+        if while_move_hook is not None:
+            while_move_hook()
+
+        present_pos = arm.read("Present_Position", motor_name).item()
+        present_speed = arm.read("Present_Speed", motor_name).item()
+        present_current = arm.read("Present_Current", motor_name).item()
+        # present_load = arm.read("Present_Load", motor_name).item()
+        # present_voltage = arm.read("Present_Voltage", motor_name).item()
+        # present_temperature = arm.read("Present_Temperature", motor_name).item()
+
+        # print(f"{present_pos=}")
+        # print(f"{present_speed=}")
+        # print(f"{present_current=}")
+        # print(f"{present_load=}")
+        # print(f"{present_voltage=}")
+        # print(f"{present_temperature=}")
+
+        if present_speed == 0 and present_current > 40:
+            count += 1
+            if count > 100 or present_current > 300:
+                return present_pos
+        else:
+            count = 0
+
+
+def move_to_calibrate(
+    arm,
+    motor_name,
+    invert_drive_mode=False,
+    positive_first=True,
+    in_between_move_hook=None,
+    while_move_hook=None,
+):
+    initial_pos = arm.read("Present_Position", motor_name)
+
+    if positive_first:
+        p_present_pos = move_until_block(
+            arm, motor_name, positive_direction=True, while_move_hook=while_move_hook
+        )
+    else:
+        n_present_pos = move_until_block(
+            arm, motor_name, positive_direction=False, while_move_hook=while_move_hook
+        )
+
+    if in_between_move_hook is not None:
+        in_between_move_hook()
+
+    if positive_first:
+        n_present_pos = move_until_block(
+            arm, motor_name, positive_direction=False, while_move_hook=while_move_hook
+        )
+    else:
+        p_present_pos = move_until_block(
+            arm, motor_name, positive_direction=True, while_move_hook=while_move_hook
+        )
+
+    zero_pos = (n_present_pos + p_present_pos) / 2
+
+    calib_data = {
+        "initial_pos": initial_pos,
+        "homing_offset": zero_pos if invert_drive_mode else -zero_pos,
+        "invert_drive_mode": invert_drive_mode,
+        "drive_mode": -1 if invert_drive_mode else 0,
+        "zero_pos": zero_pos,
+        "start_pos": n_present_pos if invert_drive_mode else p_present_pos,
+        "end_pos": p_present_pos if invert_drive_mode else n_present_pos,
+    }
+    return calib_data
+
+
+def apply_offset(calib, offset):
+    calib["zero_pos"] += offset
+    if calib["drive_mode"]:
+        calib["homing_offset"] += offset
+    else:
+        calib["homing_offset"] -= offset
+    return calib
+
+
+def run_arm_auto_calibration(arm: MotorsBus, robot_type: str, arm_name: str, arm_type: str):
+    if robot_type == "so100":
+        return run_arm_auto_calibration_so100(arm, robot_type, arm_name, arm_type)
+    elif robot_type == "moss":
+        return run_arm_auto_calibration_moss(arm, robot_type, arm_name, arm_type)
+    else:
+        raise ValueError(robot_type)
+
+
+def run_arm_auto_calibration_so100(arm: MotorsBus, robot_type: str, arm_name: str, arm_type: str):
+    """All the offsets and magic numbers are hand tuned, and are unique to SO-100 follower arms"""
+    if (arm.read("Torque_Enable") != TorqueMode.DISABLED.value).any():
+        raise ValueError("To run calibration, the torque must be disabled on all motors.")
+
+    if not (robot_type == "so100" and arm_type == "follower"):
+        raise NotImplementedError("Auto calibration only supports the follower of so100 arms for now.")
+
+    print(f"\nRunning calibration of {robot_type} {arm_name} {arm_type}...")
+
+    print("\nMove arm to initial position")
+    print("See: " + URL_TEMPLATE.format(robot=robot_type, arm=arm_type, position="initial"))
+    input("Press Enter to continue...")
+
+    # Lower the acceleration of the motors (in [0,254])
+    initial_acceleration = arm.read("Acceleration")
+    arm.write("Lock", 0)
+    arm.write("Acceleration", 10)
+    time.sleep(1)
+
+    arm.write("Torque_Enable", TorqueMode.ENABLED.value)
+
+    print(f'{arm.read("Present_Position", "elbow_flex")=}')
+
+    calib = {}
+
+    init_wf_pos = arm.read("Present_Position", "wrist_flex")
+    init_sl_pos = arm.read("Present_Position", "shoulder_lift")
+    init_ef_pos = arm.read("Present_Position", "elbow_flex")
+    arm.write("Goal_Position", init_wf_pos - 800, "wrist_flex")
+    arm.write("Goal_Position", init_sl_pos + 150 + 1024, "shoulder_lift")
+    arm.write("Goal_Position", init_ef_pos - 2048, "elbow_flex")
+    time.sleep(2)
+
+    print("Calibrate shoulder_pan")
+    calib["shoulder_pan"] = move_to_calibrate(arm, "shoulder_pan")
+    arm.write("Goal_Position", calib["shoulder_pan"]["zero_pos"], "shoulder_pan")
+    time.sleep(1)
+
+    print("Calibrate gripper")
+    calib["gripper"] = move_to_calibrate(arm, "gripper", invert_drive_mode=True)
+    time.sleep(1)
+
+    print("Calibrate wrist_flex")
+    calib["wrist_flex"] = move_to_calibrate(arm, "wrist_flex")
+    calib["wrist_flex"] = apply_offset(calib["wrist_flex"], offset=80)
+
+    def in_between_move_hook():
+        nonlocal arm, calib
+        time.sleep(2)
+        ef_pos = arm.read("Present_Position", "elbow_flex")
+        sl_pos = arm.read("Present_Position", "shoulder_lift")
+        arm.write("Goal_Position", ef_pos + 1024, "elbow_flex")
+        arm.write("Goal_Position", sl_pos - 1024, "shoulder_lift")
+        time.sleep(2)
+
+    print("Calibrate elbow_flex")
+    calib["elbow_flex"] = move_to_calibrate(
+        arm, "elbow_flex", positive_first=False, in_between_move_hook=in_between_move_hook
+    )
+    calib["elbow_flex"] = apply_offset(calib["elbow_flex"], offset=80 - 1024)
+
+    arm.write("Goal_Position", calib["elbow_flex"]["zero_pos"] + 1024 + 512, "elbow_flex")
+    time.sleep(1)
+
+    def in_between_move_hook():
+        nonlocal arm, calib
+        arm.write("Goal_Position", calib["elbow_flex"]["zero_pos"], "elbow_flex")
+
+    print("Calibrate shoulder_lift")
+    calib["shoulder_lift"] = move_to_calibrate(
+        arm,
+        "shoulder_lift",
+        invert_drive_mode=True,
+        positive_first=False,
+        in_between_move_hook=in_between_move_hook,
+    )
+    # add an 30 steps as offset to align with body
+    calib["shoulder_lift"] = apply_offset(calib["shoulder_lift"], offset=1024 - 50)
+
+    def while_move_hook():
+        nonlocal arm, calib
+        positions = {
+            "shoulder_lift": round(calib["shoulder_lift"]["zero_pos"] - 1600),
+            "elbow_flex": round(calib["elbow_flex"]["zero_pos"] + 1700),
+            "wrist_flex": round(calib["wrist_flex"]["zero_pos"] + 800),
+            "gripper": round(calib["gripper"]["end_pos"]),
+        }
+        arm.write("Goal_Position", list(positions.values()), list(positions.keys()))
+
+    arm.write("Goal_Position", round(calib["shoulder_lift"]["zero_pos"] - 1600), "shoulder_lift")
+    time.sleep(2)
+    arm.write("Goal_Position", round(calib["elbow_flex"]["zero_pos"] + 1700), "elbow_flex")
+    time.sleep(2)
+    arm.write("Goal_Position", round(calib["wrist_flex"]["zero_pos"] + 800), "wrist_flex")
+    time.sleep(2)
+    arm.write("Goal_Position", round(calib["gripper"]["end_pos"]), "gripper")
+    time.sleep(2)
+
+    print("Calibrate wrist_roll")
+    calib["wrist_roll"] = move_to_calibrate(
+        arm, "wrist_roll", invert_drive_mode=True, positive_first=False, while_move_hook=while_move_hook
+    )
+
+    arm.write("Goal_Position", calib["wrist_roll"]["zero_pos"], "wrist_roll")
+    time.sleep(1)
+    arm.write("Goal_Position", calib["gripper"]["start_pos"], "gripper")
+    time.sleep(1)
+    arm.write("Goal_Position", calib["wrist_flex"]["zero_pos"], "wrist_flex")
+    time.sleep(1)
+    arm.write("Goal_Position", calib["elbow_flex"]["zero_pos"] + 2048, "elbow_flex")
+    arm.write("Goal_Position", calib["shoulder_lift"]["zero_pos"] - 2048, "shoulder_lift")
+    time.sleep(1)
+    arm.write("Goal_Position", calib["shoulder_pan"]["zero_pos"], "shoulder_pan")
+    time.sleep(1)
+
+    calib_modes = []
+    for name in arm.motor_names:
+        if name == "gripper":
+            calib_modes.append(CalibrationMode.LINEAR.name)
+        else:
+            calib_modes.append(CalibrationMode.DEGREE.name)
+
+    calib_dict = {
+        "homing_offset": [calib[name]["homing_offset"] for name in arm.motor_names],
+        "drive_mode": [calib[name]["drive_mode"] for name in arm.motor_names],
+        "start_pos": [calib[name]["start_pos"] for name in arm.motor_names],
+        "end_pos": [calib[name]["end_pos"] for name in arm.motor_names],
+        "calib_mode": calib_modes,
+        "motor_names": arm.motor_names,
+    }
+
+    # Re-enable original accerlation
+    arm.write("Lock", 0)
+    arm.write("Acceleration", initial_acceleration)
+    time.sleep(1)
+
+    return calib_dict
+
+
+def run_arm_auto_calibration_moss(arm: MotorsBus, robot_type: str, arm_name: str, arm_type: str):
+    """All the offsets and magic numbers are hand tuned, and are unique to SO-100 follower arms"""
+    if (arm.read("Torque_Enable") != TorqueMode.DISABLED.value).any():
+        raise ValueError("To run calibration, the torque must be disabled on all motors.")
+
+    if not (robot_type == "moss" and arm_type == "follower"):
+        raise NotImplementedError("Auto calibration only supports the follower of moss arms for now.")
+
+    print(f"\nRunning calibration of {robot_type} {arm_name} {arm_type}...")
+
+    print("\nMove arm to initial position")
+    print("See: " + URL_TEMPLATE.format(robot=robot_type, arm=arm_type, position="initial"))
+    input("Press Enter to continue...")
+
+    # Lower the acceleration of the motors (in [0,254])
+    initial_acceleration = arm.read("Acceleration")
+    arm.write("Lock", 0)
+    arm.write("Acceleration", 10)
+    time.sleep(1)
+
+    arm.write("Torque_Enable", TorqueMode.ENABLED.value)
+
+    sl_pos = arm.read("Present_Position", "shoulder_lift")
+    arm.write("Goal_Position", sl_pos - 1024 - 450, "shoulder_lift")
+    ef_pos = arm.read("Present_Position", "elbow_flex")
+    arm.write("Goal_Position", ef_pos + 1024 + 450, "elbow_flex")
+    time.sleep(2)
+
+    calib = {}
+
+    print("Calibrate shoulder_pan")
+    calib["shoulder_pan"] = move_to_calibrate(arm, "shoulder_pan")
+    arm.write("Goal_Position", calib["shoulder_pan"]["zero_pos"], "shoulder_pan")
+    time.sleep(1)
+
+    print("Calibrate gripper")
+    calib["gripper"] = move_to_calibrate(arm, "gripper", invert_drive_mode=True)
+    time.sleep(1)
+
+    print("Calibrate wrist_flex")
+    calib["wrist_flex"] = move_to_calibrate(arm, "wrist_flex", invert_drive_mode=True)
+    calib["wrist_flex"] = apply_offset(calib["wrist_flex"], offset=-210 + 1024)
+
+    wr_pos = arm.read("Present_Position", "wrist_roll")
+    arm.write("Goal_Position", calib["wrist_flex"]["zero_pos"] - 1024, "wrist_flex")
+    time.sleep(1)
+    arm.write("Goal_Position", wr_pos - 1024, "wrist_roll")
+    time.sleep(1)
+    arm.write("Goal_Position", calib["wrist_flex"]["zero_pos"] - 2048, "wrist_flex")
+    time.sleep(1)
+    arm.write("Goal_Position", calib["gripper"]["end_pos"], "gripper")
+    time.sleep(1)
+
+    print("Calibrate wrist_roll")
+    calib["wrist_roll"] = move_to_calibrate(arm, "wrist_roll", invert_drive_mode=True)
+    calib["wrist_roll"] = apply_offset(calib["wrist_roll"], offset=790)
+
+    arm.write("Goal_Position", calib["wrist_roll"]["zero_pos"] - 1024, "wrist_roll")
+    arm.write("Goal_Position", calib["gripper"]["start_pos"], "gripper")
+    arm.write("Goal_Position", calib["wrist_flex"]["zero_pos"] - 1024, "wrist_flex")
+    time.sleep(1)
+    arm.write("Goal_Position", calib["wrist_roll"]["zero_pos"], "wrist_roll")
+    arm.write("Goal_Position", calib["wrist_flex"]["zero_pos"] - 2048, "wrist_flex")
+
+    def in_between_move_elbow_flex_hook():
+        nonlocal arm, calib
+        arm.write("Goal_Position", calib["wrist_flex"]["zero_pos"], "wrist_flex")
+
+    print("Calibrate elbow_flex")
+    calib["elbow_flex"] = move_to_calibrate(
+        arm,
+        "elbow_flex",
+        invert_drive_mode=True,
+        in_between_move_hook=in_between_move_elbow_flex_hook,
+    )
+    arm.write("Goal_Position", calib["wrist_flex"]["zero_pos"] - 1024, "wrist_flex")
+
+    def in_between_move_shoulder_lift_hook():
+        nonlocal arm, calib
+        sl = arm.read("Present_Position", "shoulder_lift")
+        arm.write("Goal_Position", sl - 1500, "shoulder_lift")
+        time.sleep(1)
+        arm.write("Goal_Position", calib["elbow_flex"]["zero_pos"] + 1536, "elbow_flex")
+        time.sleep(1)
+        arm.write("Goal_Position", calib["wrist_flex"]["start_pos"], "wrist_flex")
+        time.sleep(1)
+
+    print("Calibrate shoulder_lift")
+    calib["shoulder_lift"] = move_to_calibrate(
+        arm, "shoulder_lift", in_between_move_hook=in_between_move_shoulder_lift_hook
+    )
+    calib["shoulder_lift"] = apply_offset(calib["shoulder_lift"], offset=-1024)
+
+    arm.write("Goal_Position", calib["wrist_flex"]["zero_pos"] - 1024, "wrist_flex")
+    time.sleep(1)
+    arm.write("Goal_Position", calib["shoulder_lift"]["zero_pos"] + 2048, "shoulder_lift")
+    arm.write("Goal_Position", calib["elbow_flex"]["zero_pos"] - 1024 - 400, "elbow_flex")
+    time.sleep(2)
+
+    calib_modes = []
+    for name in arm.motor_names:
+        if name == "gripper":
+            calib_modes.append(CalibrationMode.LINEAR.name)
+        else:
+            calib_modes.append(CalibrationMode.DEGREE.name)
+
+    calib_dict = {
+        "homing_offset": [calib[name]["homing_offset"] for name in arm.motor_names],
+        "drive_mode": [calib[name]["drive_mode"] for name in arm.motor_names],
+        "start_pos": [calib[name]["start_pos"] for name in arm.motor_names],
+        "end_pos": [calib[name]["end_pos"] for name in arm.motor_names],
+        "calib_mode": calib_modes,
+        "motor_names": arm.motor_names,
+    }
+
+    # Re-enable original accerlation
+    arm.write("Lock", 0)
+    arm.write("Acceleration", initial_acceleration)
+    time.sleep(1)
+
+    return calib_dict
+
+
+def run_arm_manual_calibration(arm: MotorsBus, robot_type: str, arm_name: str, arm_type: str):
+    """This function ensures that a neural network trained on data collected on a given robot
+    can work on another robot. For instance before calibration, setting a same goal position
+    for each motor of two different robots will get two very different positions. But after calibration,
+    the two robots will move to the same position.To this end, this function computes the homing offset
+    and the drive mode for each motor of a given robot.
+
+    Homing offset is used to shift the motor position to a ]-2048, +2048[ nominal range (when the motor uses 2048 steps
+    to complete a half a turn). This range is set around an arbitrary "zero position" corresponding to all motor positions
+    being 0. During the calibration process, you will need to manually move the robot to this "zero position".
+
+    Drive mode is used to invert the rotation direction of the motor. This is useful when some motors have been assembled
+    in the opposite orientation for some robots. During the calibration process, you will need to manually move the robot
+    to the "rotated position".
+
+    After calibration, the homing offsets and drive modes are stored in a cache.
+
+    Example of usage:
+    ```python
+    run_arm_calibration(arm, "so100", "left", "follower")
+    ```
+    """
+    if (arm.read("Torque_Enable") != TorqueMode.DISABLED.value).any():
+        raise ValueError("To run calibration, the torque must be disabled on all motors.")
+
+    print(f"\nRunning calibration of {robot_type} {arm_name} {arm_type}...")
+
+    print("\nMove arm to zero position")
+    print("See: " + URL_TEMPLATE.format(robot=robot_type, arm=arm_type, position="zero"))
+    input("Press Enter to continue...")
+
+    # We arbitrarily chose our zero target position to be a straight horizontal position with gripper upwards and closed.
+    # It is easy to identify and all motors are in a "quarter turn" position. Once calibration is done, this position will
+    # correspond to every motor angle being 0. If you set all 0 as Goal Position, the arm will move in this position.
+    zero_target_pos = convert_degrees_to_steps(ZERO_POSITION_DEGREE, arm.motor_models)
+
+    # Compute homing offset so that `present_position + homing_offset ~= target_position`.
+    zero_pos = arm.read("Present_Position")
+    homing_offset = zero_target_pos - zero_pos
+
+    # The rotated target position corresponds to a rotation of a quarter turn from the zero position.
+    # This allows to identify the rotation direction of each motor.
+    # For instance, if the motor rotates 90 degree, and its value is -90 after applying the homing offset, then we know its rotation direction
+    # is inverted. However, for the calibration being successful, we need everyone to follow the same target position.
+    # Sometimes, there is only one possible rotation direction. For instance, if the gripper is closed, there is only one direction which
+    # corresponds to opening the gripper. When the rotation direction is ambiguous, we arbitrarily rotate clockwise from the point of view
+    # of the previous motor in the kinetic chain.
+    print("\nMove arm to rotated target position")
+    print("See: " + URL_TEMPLATE.format(robot=robot_type, arm=arm_type, position="rotated"))
+    input("Press Enter to continue...")
+
+    rotated_target_pos = convert_degrees_to_steps(ROTATED_POSITION_DEGREE, arm.motor_models)
+
+    # Find drive mode by rotating each motor by a quarter of a turn.
+    # Drive mode indicates if the motor rotation direction should be inverted (=1) or not (=0).
+    rotated_pos = arm.read("Present_Position")
+    drive_mode = (rotated_pos < zero_pos).astype(np.int32)
+
+    # Re-compute homing offset to take into account drive mode
+    rotated_drived_pos = apply_drive_mode(rotated_pos, drive_mode)
+    homing_offset = rotated_target_pos - rotated_drived_pos
+
+    print("\nMove arm to rest position")
+    print("See: " + URL_TEMPLATE.format(robot=robot_type, arm=arm_type, position="rest"))
+    input("Press Enter to continue...")
+    print()
+
+    # Joints with rotational motions are expressed in degrees in nominal range of [-180, 180]
+    calib_modes = []
+    for name in arm.motor_names:
+        if name == "gripper":
+            calib_modes.append(CalibrationMode.LINEAR.name)
+        else:
+            calib_modes.append(CalibrationMode.DEGREE.name)
+
+    calib_dict = {
+        "homing_offset": homing_offset.tolist(),
+        "drive_mode": drive_mode.tolist(),
+        "start_pos": zero_pos.tolist(),
+        "end_pos": rotated_pos.tolist(),
+        "calib_mode": calib_modes,
+        "motor_names": arm.motor_names,
+    }
+    return calib_dict

lerobot/common/robot_devices/robots/lekiwi_remote.py

@@ -0,0 +1,224 @@
+# 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 base64
+import json
+import threading
+import time
+from pathlib import Path
+
+import cv2
+import zmq
+
+from lerobot.common.robot_devices.robots.mobile_manipulator import LeKiwi
+
+
+def setup_zmq_sockets(config):
+    context = zmq.Context()
+    cmd_socket = context.socket(zmq.PULL)
+    cmd_socket.setsockopt(zmq.CONFLATE, 1)
+    cmd_socket.bind(f"tcp://*:{config.port}")
+
+    video_socket = context.socket(zmq.PUSH)
+    video_socket.setsockopt(zmq.CONFLATE, 1)
+    video_socket.bind(f"tcp://*:{config.video_port}")
+
+    return context, cmd_socket, video_socket
+
+
+def run_camera_capture(cameras, images_lock, latest_images_dict, stop_event):
+    while not stop_event.is_set():
+        local_dict = {}
+        for name, cam in cameras.items():
+            frame = cam.async_read()
+            ret, buffer = cv2.imencode(".jpg", frame, [int(cv2.IMWRITE_JPEG_QUALITY), 90])
+            if ret:
+                local_dict[name] = base64.b64encode(buffer).decode("utf-8")
+            else:
+                local_dict[name] = ""
+        with images_lock:
+            latest_images_dict.update(local_dict)
+        time.sleep(0.01)
+
+
+def calibrate_follower_arm(motors_bus, calib_dir_str):
+    """
+    Calibrates the follower arm. Attempts to load an existing calibration file;
+    if not found, runs manual calibration and saves the result.
+    """
+    calib_dir = Path(calib_dir_str)
+    calib_dir.mkdir(parents=True, exist_ok=True)
+    calib_file = calib_dir / "main_follower.json"
+    try:
+        from lerobot.common.robot_devices.robots.feetech_calibration import run_arm_manual_calibration
+    except ImportError:
+        print("[WARNING] Calibration function not available. Skipping calibration.")
+        return
+
+    if calib_file.exists():
+        with open(calib_file) as f:
+            calibration = json.load(f)
+        print(f"[INFO] Loaded calibration from {calib_file}")
+    else:
+        print("[INFO] Calibration file not found. Running manual calibration...")
+        calibration = run_arm_manual_calibration(motors_bus, "lekiwi", "follower_arm", "follower")
+        print(f"[INFO] Calibration complete. Saving to {calib_file}")
+        with open(calib_file, "w") as f:
+            json.dump(calibration, f)
+    try:
+        motors_bus.set_calibration(calibration)
+        print("[INFO] Applied calibration for follower arm.")
+    except Exception as e:
+        print(f"[WARNING] Could not apply calibration: {e}")
+
+
+def run_lekiwi(robot_config):
+    """
+    Runs the LeKiwi robot:
+      - Sets up cameras and connects them.
+      - Initializes the follower arm motors.
+      - Calibrates the follower arm if necessary.
+      - Creates ZeroMQ sockets for receiving commands and streaming observations.
+      - Processes incoming commands (arm and wheel commands) and sends back sensor and camera data.
+    """
+    # Import helper functions and classes
+    from lerobot.common.robot_devices.cameras.utils import make_cameras_from_configs
+    from lerobot.common.robot_devices.motors.feetech import FeetechMotorsBus, TorqueMode
+
+    # Initialize cameras from the robot configuration.
+    cameras = make_cameras_from_configs(robot_config.cameras)
+    for cam in cameras.values():
+        cam.connect()
+
+    # Initialize the motors bus using the follower arm configuration.
+    motor_config = robot_config.follower_arms.get("main")
+    if motor_config is None:
+        print("[ERROR] Follower arm 'main' configuration not found.")
+        return
+    motors_bus = FeetechMotorsBus(motor_config)
+    motors_bus.connect()
+
+    # Calibrate the follower arm.
+    calibrate_follower_arm(motors_bus, robot_config.calibration_dir)
+
+    # Create the LeKiwi robot instance.
+    robot = LeKiwi(motors_bus)
+
+    # Define the expected arm motor IDs.
+    arm_motor_ids = ["shoulder_pan", "shoulder_lift", "elbow_flex", "wrist_flex", "wrist_roll", "gripper"]
+
+    # Disable torque for each arm motor.
+    for motor in arm_motor_ids:
+        motors_bus.write("Torque_Enable", TorqueMode.DISABLED.value, motor)
+
+    # Set up ZeroMQ sockets.
+    context, cmd_socket, video_socket = setup_zmq_sockets(robot_config)
+
+    # Start the camera capture thread.
+    latest_images_dict = {}
+    images_lock = threading.Lock()
+    stop_event = threading.Event()
+    cam_thread = threading.Thread(
+        target=run_camera_capture, args=(cameras, images_lock, latest_images_dict, stop_event), daemon=True
+    )
+    cam_thread.start()
+
+    last_cmd_time = time.time()
+    print("LeKiwi robot server started. Waiting for commands...")
+
+    try:
+        while True:
+            loop_start_time = time.time()
+
+            # Process incoming commands (non-blocking).
+            while True:
+                try:
+                    msg = cmd_socket.recv_string(zmq.NOBLOCK)
+                except zmq.Again:
+                    break
+                try:
+                    data = json.loads(msg)
+                    # Process arm position commands.
+                    if "arm_positions" in data:
+                        arm_positions = data["arm_positions"]
+                        if not isinstance(arm_positions, list):
+                            print(f"[ERROR] Invalid arm_positions: {arm_positions}")
+                        elif len(arm_positions) < len(arm_motor_ids):
+                            print(
+                                f"[WARNING] Received {len(arm_positions)} arm positions, expected {len(arm_motor_ids)}"
+                            )
+                        else:
+                            for motor, pos in zip(arm_motor_ids, arm_positions, strict=False):
+                                motors_bus.write("Goal_Position", pos, motor)
+                    # Process wheel (base) commands.
+                    if "raw_velocity" in data:
+                        raw_command = data["raw_velocity"]
+                        # Expect keys: "left_wheel", "back_wheel", "right_wheel".
+                        command_speeds = [
+                            int(raw_command.get("left_wheel", 0)),
+                            int(raw_command.get("back_wheel", 0)),
+                            int(raw_command.get("right_wheel", 0)),
+                        ]
+                        robot.set_velocity(command_speeds)
+                        last_cmd_time = time.time()
+                except Exception as e:
+                    print(f"[ERROR] Parsing message failed: {e}")
+
+            # Watchdog: stop the robot if no command is received for over 0.5 seconds.
+            now = time.time()
+            if now - last_cmd_time > 0.5:
+                robot.stop()
+                last_cmd_time = now
+
+            # Read current wheel speeds from the robot.
+            current_velocity = robot.read_velocity()
+
+            # Read the follower arm state from the motors bus.
+            follower_arm_state = []
+            for motor in arm_motor_ids:
+                try:
+                    pos = motors_bus.read("Present_Position", motor)
+                    # Convert the position to a float (or use as is if already numeric).
+                    follower_arm_state.append(float(pos) if not isinstance(pos, (int, float)) else pos)
+                except Exception as e:
+                    print(f"[ERROR] Reading motor {motor} failed: {e}")
+
+            # Get the latest camera images.
+            with images_lock:
+                images_dict_copy = dict(latest_images_dict)
+
+            # Build the observation dictionary.
+            observation = {
+                "images": images_dict_copy,
+                "present_speed": current_velocity,
+                "follower_arm_state": follower_arm_state,
+            }
+            # Send the observation over the video socket.
+            video_socket.send_string(json.dumps(observation))
+
+            # Ensure a short sleep to avoid overloading the CPU.
+            elapsed = time.time() - loop_start_time
+            time.sleep(
+                max(0.033 - elapsed, 0)
+            )  # If robot jitters increase the sleep and monitor cpu load with `top` in cmd
+    except KeyboardInterrupt:
+        print("Shutting down LeKiwi server.")
+    finally:
+        stop_event.set()
+        cam_thread.join()
+        robot.stop()
+        motors_bus.disconnect()
+        cmd_socket.close()
+        video_socket.close()
+        context.term()

lerobot/common/robot_devices/robots/manipulator.py

@@ -0,0 +1,627 @@
+# 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.
+
+"""Contains logic to instantiate a robot, read information from its motors and cameras,
+and send orders to its motors.
+"""
+# TODO(rcadene, aliberts): reorganize the codebase into one file per robot, with the associated
+# calibration procedure, to make it easy for people to add their own robot.
+
+import json
+import logging
+import time
+import warnings
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from lerobot.common.robot_devices.cameras.utils import make_cameras_from_configs
+from lerobot.common.robot_devices.motors.utils import MotorsBus, make_motors_buses_from_configs
+from lerobot.common.robot_devices.robots.configs import ManipulatorRobotConfig
+from lerobot.common.robot_devices.robots.utils import get_arm_id
+from lerobot.common.robot_devices.utils import RobotDeviceAlreadyConnectedError, RobotDeviceNotConnectedError
+
+
+def ensure_safe_goal_position(
+    goal_pos: torch.Tensor, present_pos: torch.Tensor, max_relative_target: float | list[float]
+):
+    # Cap relative action target magnitude for safety.
+    diff = goal_pos - present_pos
+    max_relative_target = torch.tensor(max_relative_target)
+    safe_diff = torch.minimum(diff, max_relative_target)
+    safe_diff = torch.maximum(safe_diff, -max_relative_target)
+    safe_goal_pos = present_pos + safe_diff
+
+    if not torch.allclose(goal_pos, safe_goal_pos):
+        logging.warning(
+            "Relative goal position magnitude had to be clamped to be safe.\n"
+            f"  requested relative goal position target: {diff}\n"
+            f"    clamped relative goal position target: {safe_diff}"
+        )
+
+    return safe_goal_pos
+
+
+class ManipulatorRobot:
+    # TODO(rcadene): Implement force feedback
+    """This class allows to control any manipulator robot of various number of motors.
+
+    Non exhaustive list of robots:
+    - [Koch v1.0](https://github.com/AlexanderKoch-Koch/low_cost_robot), with and without the wrist-to-elbow expansion, developed
+    by Alexander Koch from [Tau Robotics](https://tau-robotics.com)
+    - [Koch v1.1](https://github.com/jess-moss/koch-v1-1) developed by Jess Moss
+    - [Aloha](https://www.trossenrobotics.com/aloha-kits) developed by Trossen Robotics
+
+    Example of instantiation, a pre-defined robot config is required:
+    ```python
+    robot = ManipulatorRobot(KochRobotConfig())
+    ```
+
+    Example of overwriting motors during instantiation:
+    ```python
+    # Defines how to communicate with the motors of the leader and follower arms
+    leader_arms = {
+        "main": DynamixelMotorsBusConfig(
+            port="/dev/tty.usbmodem575E0031751",
+            motors={
+                # name: (index, model)
+                "shoulder_pan": (1, "xl330-m077"),
+                "shoulder_lift": (2, "xl330-m077"),
+                "elbow_flex": (3, "xl330-m077"),
+                "wrist_flex": (4, "xl330-m077"),
+                "wrist_roll": (5, "xl330-m077"),
+                "gripper": (6, "xl330-m077"),
+            },
+        ),
+    }
+    follower_arms = {
+        "main": DynamixelMotorsBusConfig(
+            port="/dev/tty.usbmodem575E0032081",
+            motors={
+                # name: (index, model)
+                "shoulder_pan": (1, "xl430-w250"),
+                "shoulder_lift": (2, "xl430-w250"),
+                "elbow_flex": (3, "xl330-m288"),
+                "wrist_flex": (4, "xl330-m288"),
+                "wrist_roll": (5, "xl330-m288"),
+                "gripper": (6, "xl330-m288"),
+            },
+        ),
+    }
+    robot_config = KochRobotConfig(leader_arms=leader_arms, follower_arms=follower_arms)
+    robot = ManipulatorRobot(robot_config)
+    ```
+
+    Example of overwriting cameras during instantiation:
+    ```python
+    # Defines how to communicate with 2 cameras connected to the computer.
+    # Here, the webcam of the laptop and the phone (connected in USB to the laptop)
+    # can be reached respectively using the camera indices 0 and 1. These indices can be
+    # arbitrary. See the documentation of `OpenCVCamera` to find your own camera indices.
+    cameras = {
+        "laptop": OpenCVCamera(camera_index=0, fps=30, width=640, height=480),
+        "phone": OpenCVCamera(camera_index=1, fps=30, width=640, height=480),
+    }
+    robot = ManipulatorRobot(KochRobotConfig(cameras=cameras))
+    ```
+
+    Once the robot is instantiated, connect motors buses and cameras if any (Required):
+    ```python
+    robot.connect()
+    ```
+
+    Example of highest frequency teleoperation, which doesn't require cameras:
+    ```python
+    while True:
+        robot.teleop_step()
+    ```
+
+    Example of highest frequency data collection from motors and cameras (if any):
+    ```python
+    while True:
+        observation, action = robot.teleop_step(record_data=True)
+    ```
+
+    Example of controlling the robot with a policy:
+    ```python
+    while True:
+        # Uses the follower arms and cameras to capture an observation
+        observation = robot.capture_observation()
+
+        # Assumes a policy has been instantiated
+        with torch.inference_mode():
+            action = policy.select_action(observation)
+
+        # Orders the robot to move
+        robot.send_action(action)
+    ```
+
+    Example of disconnecting which is not mandatory since we disconnect when the object is deleted:
+    ```python
+    robot.disconnect()
+    ```
+    """
+
+    def __init__(
+        self,
+        config: ManipulatorRobotConfig,
+    ):
+        self.config = config
+        self.robot_type = self.config.type
+        self.calibration_dir = Path(self.config.calibration_dir)
+        self.leader_arms = make_motors_buses_from_configs(self.config.leader_arms)
+        self.follower_arms = make_motors_buses_from_configs(self.config.follower_arms)
+        self.cameras = make_cameras_from_configs(self.config.cameras)
+        self.is_connected = False
+        self.logs = {}
+
+    def get_motor_names(self, arm: dict[str, MotorsBus]) -> list:
+        return [f"{arm}_{motor}" for arm, bus in arm.items() for motor in bus.motors]
+
+    @property
+    def camera_features(self) -> dict:
+        cam_ft = {}
+        for cam_key, cam in self.cameras.items():
+            key = f"observation.images.{cam_key}"
+            cam_ft[key] = {
+                "shape": (cam.height, cam.width, cam.channels),
+                "names": ["height", "width", "channels"],
+                "info": None,
+            }
+        return cam_ft
+
+    @property
+    def motor_features(self) -> dict:
+        action_names = self.get_motor_names(self.leader_arms)
+        state_names = self.get_motor_names(self.leader_arms)
+        return {
+            "action": {
+                "dtype": "float32",
+                "shape": (len(action_names),),
+                "names": action_names,
+            },
+            "observation.state": {
+                "dtype": "float32",
+                "shape": (len(state_names),),
+                "names": state_names,
+            },
+        }
+
+    @property
+    def features(self):
+        return {**self.motor_features, **self.camera_features}
+
+    @property
+    def has_camera(self):
+        return len(self.cameras) > 0
+
+    @property
+    def num_cameras(self):
+        return len(self.cameras)
+
+    @property
+    def available_arms(self):
+        available_arms = []
+        for name in self.follower_arms:
+            arm_id = get_arm_id(name, "follower")
+            available_arms.append(arm_id)
+        for name in self.leader_arms:
+            arm_id = get_arm_id(name, "leader")
+            available_arms.append(arm_id)
+        return available_arms
+
+    def connect(self):
+        if self.is_connected:
+            raise RobotDeviceAlreadyConnectedError(
+                "ManipulatorRobot is already connected. Do not run `robot.connect()` twice."
+            )
+
+        if not self.leader_arms and not self.follower_arms and not self.cameras:
+            raise ValueError(
+                "ManipulatorRobot doesn't have any device to connect. See example of usage in docstring of the class."
+            )
+
+        # Connect the arms
+        for name in self.follower_arms:
+            print(f"Connecting {name} follower arm.")
+            self.follower_arms[name].connect()
+        for name in self.leader_arms:
+            print(f"Connecting {name} leader arm.")
+            self.leader_arms[name].connect()
+
+        if self.robot_type in ["koch", "koch_bimanual", "aloha"]:
+            from lerobot.common.robot_devices.motors.dynamixel import TorqueMode
+        elif self.robot_type in ["so100", "moss", "lekiwi"]:
+            from lerobot.common.robot_devices.motors.feetech import TorqueMode
+
+        # We assume that at connection time, arms are in a rest position, and torque can
+        # be safely disabled to run calibration and/or set robot preset configurations.
+        for name in self.follower_arms:
+            self.follower_arms[name].write("Torque_Enable", TorqueMode.DISABLED.value)
+        for name in self.leader_arms:
+            self.leader_arms[name].write("Torque_Enable", TorqueMode.DISABLED.value)
+
+        self.activate_calibration()
+
+        # Set robot preset (e.g. torque in leader gripper for Koch v1.1)
+        if self.robot_type in ["koch", "koch_bimanual"]:
+            self.set_koch_robot_preset()
+        elif self.robot_type == "aloha":
+            self.set_aloha_robot_preset()
+        elif self.robot_type in ["so100", "moss", "lekiwi"]:
+            self.set_so100_robot_preset()
+
+        # Enable torque on all motors of the follower arms
+        for name in self.follower_arms:
+            print(f"Activating torque on {name} follower arm.")
+            self.follower_arms[name].write("Torque_Enable", 1)
+
+        if self.config.gripper_open_degree is not None:
+            if self.robot_type not in ["koch", "koch_bimanual"]:
+                raise NotImplementedError(
+                    f"{self.robot_type} does not support position AND current control in the handle, which is require to set the gripper open."
+                )
+            # Set the leader arm in torque mode with the gripper motor set to an angle. This makes it possible
+            # to squeeze the gripper and have it spring back to an open position on its own.
+            for name in self.leader_arms:
+                self.leader_arms[name].write("Torque_Enable", 1, "gripper")
+                self.leader_arms[name].write("Goal_Position", self.config.gripper_open_degree, "gripper")
+
+        # Check both arms can be read
+        for name in self.follower_arms:
+            self.follower_arms[name].read("Present_Position")
+        for name in self.leader_arms:
+            self.leader_arms[name].read("Present_Position")
+
+        # Connect the cameras
+        for name in self.cameras:
+            self.cameras[name].connect()
+
+        self.is_connected = True
+
+    def activate_calibration(self):
+        """After calibration all motors function in human interpretable ranges.
+        Rotations are expressed in degrees in nominal range of [-180, 180],
+        and linear motions (like gripper of Aloha) in nominal range of [0, 100].
+        """
+
+        def load_or_run_calibration_(name, arm, arm_type):
+            arm_id = get_arm_id(name, arm_type)
+            arm_calib_path = self.calibration_dir / f"{arm_id}.json"
+
+            if arm_calib_path.exists():
+                with open(arm_calib_path) as f:
+                    calibration = json.load(f)
+            else:
+                # TODO(rcadene): display a warning in __init__ if calibration file not available
+                print(f"Missing calibration file '{arm_calib_path}'")
+
+                if self.robot_type in ["koch", "koch_bimanual", "aloha"]:
+                    from lerobot.common.robot_devices.robots.dynamixel_calibration import run_arm_calibration
+
+                    calibration = run_arm_calibration(arm, self.robot_type, name, arm_type)
+
+                elif self.robot_type in ["so100", "moss", "lekiwi"]:
+                    from lerobot.common.robot_devices.robots.feetech_calibration import (
+                        run_arm_manual_calibration,
+                    )
+
+                    calibration = run_arm_manual_calibration(arm, self.robot_type, name, arm_type)
+
+                print(f"Calibration is done! Saving calibration file '{arm_calib_path}'")
+                arm_calib_path.parent.mkdir(parents=True, exist_ok=True)
+                with open(arm_calib_path, "w") as f:
+                    json.dump(calibration, f)
+
+            return calibration
+
+        for name, arm in self.follower_arms.items():
+            calibration = load_or_run_calibration_(name, arm, "follower")
+            arm.set_calibration(calibration)
+        for name, arm in self.leader_arms.items():
+            calibration = load_or_run_calibration_(name, arm, "leader")
+            arm.set_calibration(calibration)
+
+    def set_koch_robot_preset(self):
+        def set_operating_mode_(arm):
+            from lerobot.common.robot_devices.motors.dynamixel import TorqueMode
+
+            if (arm.read("Torque_Enable") != TorqueMode.DISABLED.value).any():
+                raise ValueError("To run set robot preset, the torque must be disabled on all motors.")
+
+            # Use 'extended position mode' for all motors except gripper, because in joint mode the servos can't
+            # rotate more than 360 degrees (from 0 to 4095) And some mistake can happen while assembling the arm,
+            # you could end up with a servo with a position 0 or 4095 at a crucial point See [
+            # https://emanual.robotis.com/docs/en/dxl/x/x_series/#operating-mode11]
+            all_motors_except_gripper = [name for name in arm.motor_names if name != "gripper"]
+            if len(all_motors_except_gripper) > 0:
+                # 4 corresponds to Extended Position on Koch motors
+                arm.write("Operating_Mode", 4, all_motors_except_gripper)
+
+            # Use 'position control current based' for gripper to be limited by the limit of the current.
+            # For the follower gripper, it means it can grasp an object without forcing too much even tho,
+            # it's goal position is a complete grasp (both gripper fingers are ordered to join and reach a touch).
+            # For the leader gripper, it means we can use it as a physical trigger, since we can force with our finger
+            # to make it move, and it will move back to its original target position when we release the force.
+            # 5 corresponds to Current Controlled Position on Koch gripper motors "xl330-m077, xl330-m288"
+            arm.write("Operating_Mode", 5, "gripper")
+
+        for name in self.follower_arms:
+            set_operating_mode_(self.follower_arms[name])
+
+            # Set better PID values to close the gap between recorded states and actions
+            # TODO(rcadene): Implement an automatic procedure to set optimal PID values for each motor
+            self.follower_arms[name].write("Position_P_Gain", 1500, "elbow_flex")
+            self.follower_arms[name].write("Position_I_Gain", 0, "elbow_flex")
+            self.follower_arms[name].write("Position_D_Gain", 600, "elbow_flex")
+
+        if self.config.gripper_open_degree is not None:
+            for name in self.leader_arms:
+                set_operating_mode_(self.leader_arms[name])
+
+                # Enable torque on the gripper of the leader arms, and move it to 45 degrees,
+                # so that we can use it as a trigger to close the gripper of the follower arms.
+                self.leader_arms[name].write("Torque_Enable", 1, "gripper")
+                self.leader_arms[name].write("Goal_Position", self.config.gripper_open_degree, "gripper")
+
+    def set_aloha_robot_preset(self):
+        def set_shadow_(arm):
+            # Set secondary/shadow ID for shoulder and elbow. These joints have two motors.
+            # As a result, if only one of them is required to move to a certain position,
+            # the other will follow. This is to avoid breaking the motors.
+            if "shoulder_shadow" in arm.motor_names:
+                shoulder_idx = arm.read("ID", "shoulder")
+                arm.write("Secondary_ID", shoulder_idx, "shoulder_shadow")
+
+            if "elbow_shadow" in arm.motor_names:
+                elbow_idx = arm.read("ID", "elbow")
+                arm.write("Secondary_ID", elbow_idx, "elbow_shadow")
+
+        for name in self.follower_arms:
+            set_shadow_(self.follower_arms[name])
+
+        for name in self.leader_arms:
+            set_shadow_(self.leader_arms[name])
+
+        for name in self.follower_arms:
+            # Set a velocity limit of 131 as advised by Trossen Robotics
+            self.follower_arms[name].write("Velocity_Limit", 131)
+
+            # Use 'extended position mode' for all motors except gripper, because in joint mode the servos can't
+            # rotate more than 360 degrees (from 0 to 4095) And some mistake can happen while assembling the arm,
+            # you could end up with a servo with a position 0 or 4095 at a crucial point See [
+            # https://emanual.robotis.com/docs/en/dxl/x/x_series/#operating-mode11]
+            all_motors_except_gripper = [
+                name for name in self.follower_arms[name].motor_names if name != "gripper"
+            ]
+            if len(all_motors_except_gripper) > 0:
+                # 4 corresponds to Extended Position on Aloha motors
+                self.follower_arms[name].write("Operating_Mode", 4, all_motors_except_gripper)
+
+            # Use 'position control current based' for follower gripper to be limited by the limit of the current.
+            # It can grasp an object without forcing too much even tho,
+            # it's goal position is a complete grasp (both gripper fingers are ordered to join and reach a touch).
+            # 5 corresponds to Current Controlled Position on Aloha gripper follower "xm430-w350"
+            self.follower_arms[name].write("Operating_Mode", 5, "gripper")
+
+            # Note: We can't enable torque on the leader gripper since "xc430-w150" doesn't have
+            # a Current Controlled Position mode.
+
+        if self.config.gripper_open_degree is not None:
+            warnings.warn(
+                f"`gripper_open_degree` is set to {self.config.gripper_open_degree}, but None is expected for Aloha instead",
+                stacklevel=1,
+            )
+
+    def set_so100_robot_preset(self):
+        for name in self.follower_arms:
+            # Mode=0 for Position Control
+            self.follower_arms[name].write("Mode", 0)
+            # Set P_Coefficient to lower value to avoid shakiness (Default is 32)
+            self.follower_arms[name].write("P_Coefficient", 16)
+            # Set I_Coefficient and D_Coefficient to default value 0 and 32
+            self.follower_arms[name].write("I_Coefficient", 0)
+            self.follower_arms[name].write("D_Coefficient", 32)
+            # Close the write lock so that Maximum_Acceleration gets written to EPROM address,
+            # which is mandatory for Maximum_Acceleration to take effect after rebooting.
+            self.follower_arms[name].write("Lock", 0)
+            # Set Maximum_Acceleration to 254 to speedup acceleration and deceleration of
+            # the motors. Note: this configuration is not in the official STS3215 Memory Table
+            self.follower_arms[name].write("Maximum_Acceleration", 254)
+            self.follower_arms[name].write("Acceleration", 254)
+
+    def teleop_step(
+        self, record_data=False
+    ) -> None | tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]:
+        if not self.is_connected:
+            raise RobotDeviceNotConnectedError(
+                "ManipulatorRobot is not connected. You need to run `robot.connect()`."
+            )
+
+        # Prepare to assign the position of the leader to the follower
+        leader_pos = {}
+        for name in self.leader_arms:
+            before_lread_t = time.perf_counter()
+            leader_pos[name] = self.leader_arms[name].read("Present_Position")
+            leader_pos[name] = torch.from_numpy(leader_pos[name])
+            self.logs[f"read_leader_{name}_pos_dt_s"] = time.perf_counter() - before_lread_t
+
+        # Send goal position to the follower
+        follower_goal_pos = {}
+        for name in self.follower_arms:
+            before_fwrite_t = time.perf_counter()
+            goal_pos = leader_pos[name]
+
+            # Cap goal position when too far away from present position.
+            # Slower fps expected due to reading from the follower.
+            if self.config.max_relative_target is not None:
+                present_pos = self.follower_arms[name].read("Present_Position")
+                present_pos = torch.from_numpy(present_pos)
+                goal_pos = ensure_safe_goal_position(goal_pos, present_pos, self.config.max_relative_target)
+
+            # Used when record_data=True
+            follower_goal_pos[name] = goal_pos
+
+            goal_pos = goal_pos.numpy().astype(np.float32)
+            self.follower_arms[name].write("Goal_Position", goal_pos)
+            self.logs[f"write_follower_{name}_goal_pos_dt_s"] = time.perf_counter() - before_fwrite_t
+
+        # Early exit when recording data is not requested
+        if not record_data:
+            return
+
+        # TODO(rcadene): Add velocity and other info
+        # Read follower position
+        follower_pos = {}
+        for name in self.follower_arms:
+            before_fread_t = time.perf_counter()
+            follower_pos[name] = self.follower_arms[name].read("Present_Position")
+            follower_pos[name] = torch.from_numpy(follower_pos[name])
+            self.logs[f"read_follower_{name}_pos_dt_s"] = time.perf_counter() - before_fread_t
+
+        # Create state by concatenating follower current position
+        state = []
+        for name in self.follower_arms:
+            if name in follower_pos:
+                state.append(follower_pos[name])
+        state = torch.cat(state)
+
+        # Create action by concatenating follower goal position
+        action = []
+        for name in self.follower_arms:
+            if name in follower_goal_pos:
+                action.append(follower_goal_pos[name])
+        action = torch.cat(action)
+
+        # Capture images from cameras
+        images = {}
+        for name in self.cameras:
+            before_camread_t = time.perf_counter()
+            images[name] = self.cameras[name].async_read()
+            images[name] = torch.from_numpy(images[name])
+            self.logs[f"read_camera_{name}_dt_s"] = self.cameras[name].logs["delta_timestamp_s"]
+            self.logs[f"async_read_camera_{name}_dt_s"] = time.perf_counter() - before_camread_t
+
+        # Populate output dictionaries
+        obs_dict, action_dict = {}, {}
+        obs_dict["observation.state"] = state
+        action_dict["action"] = action
+        for name in self.cameras:
+            obs_dict[f"observation.images.{name}"] = images[name]
+
+        return obs_dict, action_dict
+
+    def capture_observation(self):
+        """The returned observations do not have a batch dimension."""
+        if not self.is_connected:
+            raise RobotDeviceNotConnectedError(
+                "ManipulatorRobot is not connected. You need to run `robot.connect()`."
+            )
+
+        # Read follower position
+        follower_pos = {}
+        for name in self.follower_arms:
+            before_fread_t = time.perf_counter()
+            follower_pos[name] = self.follower_arms[name].read("Present_Position")
+            follower_pos[name] = torch.from_numpy(follower_pos[name])
+            self.logs[f"read_follower_{name}_pos_dt_s"] = time.perf_counter() - before_fread_t
+
+        # Create state by concatenating follower current position
+        state = []
+        for name in self.follower_arms:
+            if name in follower_pos:
+                state.append(follower_pos[name])
+        state = torch.cat(state)
+
+        # Capture images from cameras
+        images = {}
+        for name in self.cameras:
+            before_camread_t = time.perf_counter()
+            images[name] = self.cameras[name].async_read()
+            images[name] = torch.from_numpy(images[name])
+            self.logs[f"read_camera_{name}_dt_s"] = self.cameras[name].logs["delta_timestamp_s"]
+            self.logs[f"async_read_camera_{name}_dt_s"] = time.perf_counter() - before_camread_t
+
+        # Populate output dictionaries and format to pytorch
+        obs_dict = {}
+        obs_dict["observation.state"] = state
+        for name in self.cameras:
+            obs_dict[f"observation.images.{name}"] = images[name]
+        return obs_dict
+
+    def send_action(self, action: torch.Tensor) -> torch.Tensor:
+        """Command the follower arms to move to a target joint configuration.
+
+        The relative action magnitude may be clipped depending on the configuration parameter
+        `max_relative_target`. In this case, the action sent differs from original action.
+        Thus, this function always returns the action actually sent.
+
+        Args:
+            action: tensor containing the concatenated goal positions for the follower arms.
+        """
+        if not self.is_connected:
+            raise RobotDeviceNotConnectedError(
+                "ManipulatorRobot is not connected. You need to run `robot.connect()`."
+            )
+
+        from_idx = 0
+        to_idx = 0
+        action_sent = []
+        for name in self.follower_arms:
+            # Get goal position of each follower arm by splitting the action vector
+            to_idx += len(self.follower_arms[name].motor_names)
+            goal_pos = action[from_idx:to_idx]
+            from_idx = to_idx
+
+            # Cap goal position when too far away from present position.
+            # Slower fps expected due to reading from the follower.
+            if self.config.max_relative_target is not None:
+                present_pos = self.follower_arms[name].read("Present_Position")
+                present_pos = torch.from_numpy(present_pos)
+                goal_pos = ensure_safe_goal_position(goal_pos, present_pos, self.config.max_relative_target)
+
+            # Save tensor to concat and return
+            action_sent.append(goal_pos)
+
+            # Send goal position to each follower
+            goal_pos = goal_pos.numpy().astype(np.float32)
+            self.follower_arms[name].write("Goal_Position", goal_pos)
+
+        return torch.cat(action_sent)
+
+    def print_logs(self):
+        pass
+        # TODO(aliberts): move robot-specific logs logic here
+
+    def disconnect(self):
+        if not self.is_connected:
+            raise RobotDeviceNotConnectedError(
+                "ManipulatorRobot is not connected. You need to run `robot.connect()` before disconnecting."
+            )
+
+        for name in self.follower_arms:
+            self.follower_arms[name].disconnect()
+
+        for name in self.leader_arms:
+            self.leader_arms[name].disconnect()
+
+        for name in self.cameras:
+            self.cameras[name].disconnect()
+
+        self.is_connected = False
+
+    def __del__(self):
+        if getattr(self, "is_connected", False):
+            self.disconnect()

lerobot/common/robot_devices/robots/mobile_manipulator.py

@@ -0,0 +1,703 @@
+# 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 base64
+import json
+import os
+import sys
+from pathlib import Path
+
+import cv2
+import numpy as np
+import torch
+import zmq
+
+from lerobot.common.robot_devices.cameras.utils import make_cameras_from_configs
+from lerobot.common.robot_devices.motors.feetech import TorqueMode
+from lerobot.common.robot_devices.motors.utils import MotorsBus, make_motors_buses_from_configs
+from lerobot.common.robot_devices.robots.configs import LeKiwiRobotConfig
+from lerobot.common.robot_devices.robots.feetech_calibration import run_arm_manual_calibration
+from lerobot.common.robot_devices.robots.utils import get_arm_id
+from lerobot.common.robot_devices.utils import RobotDeviceNotConnectedError
+
+PYNPUT_AVAILABLE = True
+try:
+    # Only import if there's a valid X server or if we're not on a Pi
+    if ("DISPLAY" not in os.environ) and ("linux" in sys.platform):
+        print("No DISPLAY set. Skipping pynput import.")
+        raise ImportError("pynput blocked intentionally due to no display.")
+
+    from pynput import keyboard
+except ImportError:
+    keyboard = None
+    PYNPUT_AVAILABLE = False
+except Exception as e:
+    keyboard = None
+    PYNPUT_AVAILABLE = False
+    print(f"Could not import pynput: {e}")
+
+
+class MobileManipulator:
+    """
+    MobileManipulator is a class for connecting to and controlling a remote mobile manipulator robot.
+    The robot includes a three omniwheel mobile base and a remote follower arm.
+    The leader arm is connected locally (on the laptop) and its joint positions are recorded and then
+    forwarded to the remote follower arm (after applying a safety clamp).
+    In parallel, keyboard teleoperation is used to generate raw velocity commands for the wheels.
+    """
+
+    def __init__(self, config: LeKiwiRobotConfig):
+        """
+        Expected keys in config:
+          - ip, port, video_port for the remote connection.
+          - calibration_dir, leader_arms, follower_arms, max_relative_target, etc.
+        """
+        self.robot_type = config.type
+        self.config = config
+        self.remote_ip = config.ip
+        self.remote_port = config.port
+        self.remote_port_video = config.video_port
+        self.calibration_dir = Path(self.config.calibration_dir)
+        self.logs = {}
+
+        self.teleop_keys = self.config.teleop_keys
+
+        # For teleoperation, the leader arm (local) is used to record the desired arm pose.
+        self.leader_arms = make_motors_buses_from_configs(self.config.leader_arms)
+
+        self.follower_arms = make_motors_buses_from_configs(self.config.follower_arms)
+
+        self.cameras = make_cameras_from_configs(self.config.cameras)
+
+        self.is_connected = False
+
+        self.last_frames = {}
+        self.last_present_speed = {}
+        self.last_remote_arm_state = torch.zeros(6, dtype=torch.float32)
+
+        # Define three speed levels and a current index
+        self.speed_levels = [
+            {"xy": 0.1, "theta": 30},  # slow
+            {"xy": 0.2, "theta": 60},  # medium
+            {"xy": 0.3, "theta": 90},  # fast
+        ]
+        self.speed_index = 0  # Start at slow
+
+        # ZeroMQ context and sockets.
+        self.context = None
+        self.cmd_socket = None
+        self.video_socket = None
+
+        # Keyboard state for base teleoperation.
+        self.running = True
+        self.pressed_keys = {
+            "forward": False,
+            "backward": False,
+            "left": False,
+            "right": False,
+            "rotate_left": False,
+            "rotate_right": False,
+        }
+
+        if PYNPUT_AVAILABLE:
+            print("pynput is available - enabling local keyboard listener.")
+            self.listener = keyboard.Listener(
+                on_press=self.on_press,
+                on_release=self.on_release,
+            )
+            self.listener.start()
+        else:
+            print("pynput not available - skipping local keyboard listener.")
+            self.listener = None
+
+    def get_motor_names(self, arms: dict[str, MotorsBus]) -> list:
+        return [f"{arm}_{motor}" for arm, bus in arms.items() for motor in bus.motors]
+
+    @property
+    def camera_features(self) -> dict:
+        cam_ft = {}
+        for cam_key, cam in self.cameras.items():
+            key = f"observation.images.{cam_key}"
+            cam_ft[key] = {
+                "shape": (cam.height, cam.width, cam.channels),
+                "names": ["height", "width", "channels"],
+                "info": None,
+            }
+        return cam_ft
+
+    @property
+    def motor_features(self) -> dict:
+        follower_arm_names = [
+            "shoulder_pan",
+            "shoulder_lift",
+            "elbow_flex",
+            "wrist_flex",
+            "wrist_roll",
+            "gripper",
+        ]
+        observations = ["x_mm", "y_mm", "theta"]
+        combined_names = follower_arm_names + observations
+        return {
+            "action": {
+                "dtype": "float32",
+                "shape": (len(combined_names),),
+                "names": combined_names,
+            },
+            "observation.state": {
+                "dtype": "float32",
+                "shape": (len(combined_names),),
+                "names": combined_names,
+            },
+        }
+
+    @property
+    def features(self):
+        return {**self.motor_features, **self.camera_features}
+
+    @property
+    def has_camera(self):
+        return len(self.cameras) > 0
+
+    @property
+    def num_cameras(self):
+        return len(self.cameras)
+
+    @property
+    def available_arms(self):
+        available = []
+        for name in self.leader_arms:
+            available.append(get_arm_id(name, "leader"))
+        for name in self.follower_arms:
+            available.append(get_arm_id(name, "follower"))
+        return available
+
+    def on_press(self, key):
+        try:
+            # Movement
+            if key.char == self.teleop_keys["forward"]:
+                self.pressed_keys["forward"] = True
+            elif key.char == self.teleop_keys["backward"]:
+                self.pressed_keys["backward"] = True
+            elif key.char == self.teleop_keys["left"]:
+                self.pressed_keys["left"] = True
+            elif key.char == self.teleop_keys["right"]:
+                self.pressed_keys["right"] = True
+            elif key.char == self.teleop_keys["rotate_left"]:
+                self.pressed_keys["rotate_left"] = True
+            elif key.char == self.teleop_keys["rotate_right"]:
+                self.pressed_keys["rotate_right"] = True
+
+            # Quit teleoperation
+            elif key.char == self.teleop_keys["quit"]:
+                self.running = False
+                return False
+
+            # Speed control
+            elif key.char == self.teleop_keys["speed_up"]:
+                self.speed_index = min(self.speed_index + 1, 2)
+                print(f"Speed index increased to {self.speed_index}")
+            elif key.char == self.teleop_keys["speed_down"]:
+                self.speed_index = max(self.speed_index - 1, 0)
+                print(f"Speed index decreased to {self.speed_index}")
+
+        except AttributeError:
+            # e.g., if key is special like Key.esc
+            if key == keyboard.Key.esc:
+                self.running = False
+                return False
+
+    def on_release(self, key):
+        try:
+            if hasattr(key, "char"):
+                if key.char == self.teleop_keys["forward"]:
+                    self.pressed_keys["forward"] = False
+                elif key.char == self.teleop_keys["backward"]:
+                    self.pressed_keys["backward"] = False
+                elif key.char == self.teleop_keys["left"]:
+                    self.pressed_keys["left"] = False
+                elif key.char == self.teleop_keys["right"]:
+                    self.pressed_keys["right"] = False
+                elif key.char == self.teleop_keys["rotate_left"]:
+                    self.pressed_keys["rotate_left"] = False
+                elif key.char == self.teleop_keys["rotate_right"]:
+                    self.pressed_keys["rotate_right"] = False
+        except AttributeError:
+            pass
+
+    def connect(self):
+        if not self.leader_arms:
+            raise ValueError("MobileManipulator has no leader arm to connect.")
+        for name in self.leader_arms:
+            print(f"Connecting {name} leader arm.")
+            self.calibrate_leader()
+
+        # Set up ZeroMQ sockets to communicate with the remote mobile robot.
+        self.context = zmq.Context()
+        self.cmd_socket = self.context.socket(zmq.PUSH)
+        connection_string = f"tcp://{self.remote_ip}:{self.remote_port}"
+        self.cmd_socket.connect(connection_string)
+        self.cmd_socket.setsockopt(zmq.CONFLATE, 1)
+        self.video_socket = self.context.socket(zmq.PULL)
+        video_connection = f"tcp://{self.remote_ip}:{self.remote_port_video}"
+        self.video_socket.connect(video_connection)
+        self.video_socket.setsockopt(zmq.CONFLATE, 1)
+        print(
+            f"[INFO] Connected to remote robot at {connection_string} and video stream at {video_connection}."
+        )
+        self.is_connected = True
+
+    def load_or_run_calibration_(self, name, arm, arm_type):
+        arm_id = get_arm_id(name, arm_type)
+        arm_calib_path = self.calibration_dir / f"{arm_id}.json"
+
+        if arm_calib_path.exists():
+            with open(arm_calib_path) as f:
+                calibration = json.load(f)
+        else:
+            print(f"Missing calibration file '{arm_calib_path}'")
+            calibration = run_arm_manual_calibration(arm, self.robot_type, name, arm_type)
+            print(f"Calibration is done! Saving calibration file '{arm_calib_path}'")
+            arm_calib_path.parent.mkdir(parents=True, exist_ok=True)
+            with open(arm_calib_path, "w") as f:
+                json.dump(calibration, f)
+
+        return calibration
+
+    def calibrate_leader(self):
+        for name, arm in self.leader_arms.items():
+            # Connect the bus
+            arm.connect()
+
+            # Disable torque on all motors
+            for motor_id in arm.motors:
+                arm.write("Torque_Enable", TorqueMode.DISABLED.value, motor_id)
+
+            # Now run calibration
+            calibration = self.load_or_run_calibration_(name, arm, "leader")
+            arm.set_calibration(calibration)
+
+    def calibrate_follower(self):
+        for name, bus in self.follower_arms.items():
+            bus.connect()
+
+            # Disable torque on all motors
+            for motor_id in bus.motors:
+                bus.write("Torque_Enable", 0, motor_id)
+
+            # Then filter out wheels
+            arm_only_dict = {k: v for k, v in bus.motors.items() if not k.startswith("wheel_")}
+            if not arm_only_dict:
+                continue
+
+            original_motors = bus.motors
+            bus.motors = arm_only_dict
+
+            calibration = self.load_or_run_calibration_(name, bus, "follower")
+            bus.set_calibration(calibration)
+
+            bus.motors = original_motors
+
+    def _get_data(self):
+        """
+        Polls the video socket for up to 15 ms. If data arrives, decode only
+        the *latest* message, returning frames, speed, and arm state. If
+        nothing arrives for any field, use the last known values.
+        """
+        frames = {}
+        present_speed = {}
+        remote_arm_state_tensor = torch.zeros(6, dtype=torch.float32)
+
+        # Poll up to 15 ms
+        poller = zmq.Poller()
+        poller.register(self.video_socket, zmq.POLLIN)
+        socks = dict(poller.poll(15))
+        if self.video_socket not in socks or socks[self.video_socket] != zmq.POLLIN:
+            # No new data arrived → reuse ALL old data
+            return (self.last_frames, self.last_present_speed, self.last_remote_arm_state)
+
+        # Drain all messages, keep only the last
+        last_msg = None
+        while True:
+            try:
+                obs_string = self.video_socket.recv_string(zmq.NOBLOCK)
+                last_msg = obs_string
+            except zmq.Again:
+                break
+
+        if not last_msg:
+            # No new message → also reuse old
+            return (self.last_frames, self.last_present_speed, self.last_remote_arm_state)
+
+        # Decode only the final message
+        try:
+            observation = json.loads(last_msg)
+
+            images_dict = observation.get("images", {})
+            new_speed = observation.get("present_speed", {})
+            new_arm_state = observation.get("follower_arm_state", None)
+
+            # Convert images
+            for cam_name, image_b64 in images_dict.items():
+                if image_b64:
+                    jpg_data = base64.b64decode(image_b64)
+                    np_arr = np.frombuffer(jpg_data, dtype=np.uint8)
+                    frame_candidate = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
+                    if frame_candidate is not None:
+                        frames[cam_name] = frame_candidate
+
+            # If remote_arm_state is None and frames is None there is no message then use the previous message
+            if new_arm_state is not None and frames is not None:
+                self.last_frames = frames
+
+                remote_arm_state_tensor = torch.tensor(new_arm_state, dtype=torch.float32)
+                self.last_remote_arm_state = remote_arm_state_tensor
+
+                present_speed = new_speed
+                self.last_present_speed = new_speed
+            else:
+                frames = self.last_frames
+
+                remote_arm_state_tensor = self.last_remote_arm_state
+
+                present_speed = self.last_present_speed
+
+        except Exception as e:
+            print(f"[DEBUG] Error decoding video message: {e}")
+            # If decode fails, fall back to old data
+            return (self.last_frames, self.last_present_speed, self.last_remote_arm_state)
+
+        return frames, present_speed, remote_arm_state_tensor
+
+    def _process_present_speed(self, present_speed: dict) -> torch.Tensor:
+        state_tensor = torch.zeros(3, dtype=torch.int32)
+        if present_speed:
+            decoded = {key: MobileManipulator.raw_to_degps(value) for key, value in present_speed.items()}
+            if "1" in decoded:
+                state_tensor[0] = decoded["1"]
+            if "2" in decoded:
+                state_tensor[1] = decoded["2"]
+            if "3" in decoded:
+                state_tensor[2] = decoded["3"]
+        return state_tensor
+
+    def teleop_step(
+        self, record_data: bool = False
+    ) -> None | tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]:
+        if not self.is_connected:
+            raise RobotDeviceNotConnectedError("MobileManipulator is not connected. Run `connect()` first.")
+
+        speed_setting = self.speed_levels[self.speed_index]
+        xy_speed = speed_setting["xy"]  # e.g. 0.1, 0.25, or 0.4
+        theta_speed = speed_setting["theta"]  # e.g. 30, 60, or 90
+
+        # Prepare to assign the position of the leader to the follower
+        arm_positions = []
+        for name in self.leader_arms:
+            pos = self.leader_arms[name].read("Present_Position")
+            pos_tensor = torch.from_numpy(pos).float()
+            arm_positions.extend(pos_tensor.tolist())
+
+        y_cmd = 0.0  # m/s forward/backward
+        x_cmd = 0.0  # m/s lateral
+        theta_cmd = 0.0  # deg/s rotation
+        if self.pressed_keys["forward"]:
+            y_cmd += xy_speed
+        if self.pressed_keys["backward"]:
+            y_cmd -= xy_speed
+        if self.pressed_keys["left"]:
+            x_cmd += xy_speed
+        if self.pressed_keys["right"]:
+            x_cmd -= xy_speed
+        if self.pressed_keys["rotate_left"]:
+            theta_cmd += theta_speed
+        if self.pressed_keys["rotate_right"]:
+            theta_cmd -= theta_speed
+
+        wheel_commands = self.body_to_wheel_raw(x_cmd, y_cmd, theta_cmd)
+
+        message = {"raw_velocity": wheel_commands, "arm_positions": arm_positions}
+        self.cmd_socket.send_string(json.dumps(message))
+
+        if not record_data:
+            return
+
+        obs_dict = self.capture_observation()
+
+        arm_state_tensor = torch.tensor(arm_positions, dtype=torch.float32)
+
+        wheel_velocity_tuple = self.wheel_raw_to_body(wheel_commands)
+        wheel_velocity_mm = (
+            wheel_velocity_tuple[0] * 1000.0,
+            wheel_velocity_tuple[1] * 1000.0,
+            wheel_velocity_tuple[2],
+        )
+        wheel_tensor = torch.tensor(wheel_velocity_mm, dtype=torch.float32)
+        action_tensor = torch.cat([arm_state_tensor, wheel_tensor])
+        action_dict = {"action": action_tensor}
+
+        return obs_dict, action_dict
+
+    def capture_observation(self) -> dict:
+        """
+        Capture observations from the remote robot: current follower arm positions,
+        present wheel speeds (converted to body-frame velocities: x, y, theta),
+        and a camera frame.
+        """
+        if not self.is_connected:
+            raise RobotDeviceNotConnectedError("Not connected. Run `connect()` first.")
+
+        frames, present_speed, remote_arm_state_tensor = self._get_data()
+
+        body_state = self.wheel_raw_to_body(present_speed)
+
+        body_state_mm = (body_state[0] * 1000.0, body_state[1] * 1000.0, body_state[2])  # Convert x,y to mm/s
+        wheel_state_tensor = torch.tensor(body_state_mm, dtype=torch.float32)
+        combined_state_tensor = torch.cat((remote_arm_state_tensor, wheel_state_tensor), dim=0)
+
+        obs_dict = {"observation.state": combined_state_tensor}
+
+        # Loop over each configured camera
+        for cam_name, cam in self.cameras.items():
+            frame = frames.get(cam_name, None)
+            if frame is None:
+                # Create a black image using the camera's configured width, height, and channels
+                frame = np.zeros((cam.height, cam.width, cam.channels), dtype=np.uint8)
+            obs_dict[f"observation.images.{cam_name}"] = torch.from_numpy(frame)
+
+        return obs_dict
+
+    def send_action(self, action: torch.Tensor) -> torch.Tensor:
+        if not self.is_connected:
+            raise RobotDeviceNotConnectedError("Not connected. Run `connect()` first.")
+
+        # Ensure the action tensor has at least 9 elements:
+        #   - First 6: arm positions.
+        #   - Last 3: base commands.
+        if action.numel() < 9:
+            # Pad with zeros if there are not enough elements.
+            padded = torch.zeros(9, dtype=action.dtype)
+            padded[: action.numel()] = action
+            action = padded
+
+        # Extract arm and base actions.
+        arm_actions = action[:6].flatten()
+        base_actions = action[6:].flatten()
+
+        x_cmd_mm = base_actions[0].item()  # mm/s
+        y_cmd_mm = base_actions[1].item()  # mm/s
+        theta_cmd = base_actions[2].item()  # deg/s
+
+        # Convert mm/s to m/s for the kinematics calculations.
+        x_cmd = x_cmd_mm / 1000.0  # m/s
+        y_cmd = y_cmd_mm / 1000.0  # m/s
+
+        # Compute wheel commands from body commands.
+        wheel_commands = self.body_to_wheel_raw(x_cmd, y_cmd, theta_cmd)
+
+        arm_positions_list = arm_actions.tolist()
+
+        message = {"raw_velocity": wheel_commands, "arm_positions": arm_positions_list}
+        self.cmd_socket.send_string(json.dumps(message))
+
+        return action
+
+    def print_logs(self):
+        pass
+
+    def disconnect(self):
+        if not self.is_connected:
+            raise RobotDeviceNotConnectedError("Not connected.")
+        if self.cmd_socket:
+            stop_cmd = {
+                "raw_velocity": {"left_wheel": 0, "back_wheel": 0, "right_wheel": 0},
+                "arm_positions": {},
+            }
+            self.cmd_socket.send_string(json.dumps(stop_cmd))
+            self.cmd_socket.close()
+        if self.video_socket:
+            self.video_socket.close()
+        if self.context:
+            self.context.term()
+        if PYNPUT_AVAILABLE:
+            self.listener.stop()
+        self.is_connected = False
+        print("[INFO] Disconnected from remote robot.")
+
+    def __del__(self):
+        if getattr(self, "is_connected", False):
+            self.disconnect()
+        if PYNPUT_AVAILABLE:
+            self.listener.stop()
+
+    @staticmethod
+    def degps_to_raw(degps: float) -> int:
+        steps_per_deg = 4096.0 / 360.0
+        speed_in_steps = abs(degps) * steps_per_deg
+        speed_int = int(round(speed_in_steps))
+        if speed_int > 0x7FFF:
+            speed_int = 0x7FFF
+        if degps < 0:
+            return speed_int | 0x8000
+        else:
+            return speed_int & 0x7FFF
+
+    @staticmethod
+    def raw_to_degps(raw_speed: int) -> float:
+        steps_per_deg = 4096.0 / 360.0
+        magnitude = raw_speed & 0x7FFF
+        degps = magnitude / steps_per_deg
+        if raw_speed & 0x8000:
+            degps = -degps
+        return degps
+
+    def body_to_wheel_raw(
+        self,
+        x_cmd: float,
+        y_cmd: float,
+        theta_cmd: float,
+        wheel_radius: float = 0.05,
+        base_radius: float = 0.125,
+        max_raw: int = 3000,
+    ) -> dict:
+        """
+        Convert desired body-frame velocities into wheel raw commands.
+
+        Parameters:
+          x_cmd      : Linear velocity in x (m/s).
+          y_cmd      : Linear velocity in y (m/s).
+          theta_cmd  : Rotational velocity (deg/s).
+          wheel_radius: Radius of each wheel (meters).
+          base_radius : Distance from the center of rotation to each wheel (meters).
+          max_raw    : Maximum allowed raw command (ticks) per wheel.
+
+        Returns:
+          A dictionary with wheel raw commands:
+             {"left_wheel": value, "back_wheel": value, "right_wheel": value}.
+
+        Notes:
+          - Internally, the method converts theta_cmd to rad/s for the kinematics.
+          - The raw command is computed from the wheels angular speed in deg/s
+            using degps_to_raw(). If any command exceeds max_raw, all commands
+            are scaled down proportionally.
+        """
+        # Convert rotational velocity from deg/s to rad/s.
+        theta_rad = theta_cmd * (np.pi / 180.0)
+        # Create the body velocity vector [x, y, theta_rad].
+        velocity_vector = np.array([x_cmd, y_cmd, theta_rad])
+
+        # Define the wheel mounting angles (defined from y axis cw)
+        angles = np.radians(np.array([300, 180, 60]))
+        # Build the kinematic matrix: each row maps body velocities to a wheel’s linear speed.
+        # The third column (base_radius) accounts for the effect of rotation.
+        m = np.array([[np.cos(a), np.sin(a), base_radius] for a in angles])
+
+        # Compute each wheel’s linear speed (m/s) and then its angular speed (rad/s).
+        wheel_linear_speeds = m.dot(velocity_vector)
+        wheel_angular_speeds = wheel_linear_speeds / wheel_radius
+
+        # Convert wheel angular speeds from rad/s to deg/s.
+        wheel_degps = wheel_angular_speeds * (180.0 / np.pi)
+
+        # Scaling
+        steps_per_deg = 4096.0 / 360.0
+        raw_floats = [abs(degps) * steps_per_deg for degps in wheel_degps]
+        max_raw_computed = max(raw_floats)
+        if max_raw_computed > max_raw:
+            scale = max_raw / max_raw_computed
+            wheel_degps = wheel_degps * scale
+
+        # Convert each wheel’s angular speed (deg/s) to a raw integer.
+        wheel_raw = [MobileManipulator.degps_to_raw(deg) for deg in wheel_degps]
+
+        return {"left_wheel": wheel_raw[0], "back_wheel": wheel_raw[1], "right_wheel": wheel_raw[2]}
+
+    def wheel_raw_to_body(
+        self, wheel_raw: dict, wheel_radius: float = 0.05, base_radius: float = 0.125
+    ) -> tuple:
+        """
+        Convert wheel raw command feedback back into body-frame velocities.
+
+        Parameters:
+          wheel_raw   : Dictionary with raw wheel commands (keys: "left_wheel", "back_wheel", "right_wheel").
+          wheel_radius: Radius of each wheel (meters).
+          base_radius : Distance from the robot center to each wheel (meters).
+
+        Returns:
+          A tuple (x_cmd, y_cmd, theta_cmd) where:
+             x_cmd      : Linear velocity in x (m/s).
+             y_cmd      : Linear velocity in y (m/s).
+             theta_cmd  : Rotational velocity in deg/s.
+        """
+        # Extract the raw values in order.
+        raw_list = [
+            int(wheel_raw.get("left_wheel", 0)),
+            int(wheel_raw.get("back_wheel", 0)),
+            int(wheel_raw.get("right_wheel", 0)),
+        ]
+
+        # Convert each raw command back to an angular speed in deg/s.
+        wheel_degps = np.array([MobileManipulator.raw_to_degps(r) for r in raw_list])
+        # Convert from deg/s to rad/s.
+        wheel_radps = wheel_degps * (np.pi / 180.0)
+        # Compute each wheel’s linear speed (m/s) from its angular speed.
+        wheel_linear_speeds = wheel_radps * wheel_radius
+
+        # Define the wheel mounting angles (defined from y axis cw)
+        angles = np.radians(np.array([300, 180, 60]))
+        m = np.array([[np.cos(a), np.sin(a), base_radius] for a in angles])
+
+        # Solve the inverse kinematics: body_velocity = M⁻¹ · wheel_linear_speeds.
+        m_inv = np.linalg.inv(m)
+        velocity_vector = m_inv.dot(wheel_linear_speeds)
+        x_cmd, y_cmd, theta_rad = velocity_vector
+        theta_cmd = theta_rad * (180.0 / np.pi)
+        return (x_cmd, y_cmd, theta_cmd)
+
+
+class LeKiwi:
+    def __init__(self, motor_bus):
+        """
+        Initializes the LeKiwi with Feetech motors bus.
+        """
+        self.motor_bus = motor_bus
+        self.motor_ids = ["left_wheel", "back_wheel", "right_wheel"]
+
+        # Initialize motors in velocity mode.
+        self.motor_bus.write("Lock", 0)
+        self.motor_bus.write("Mode", [1, 1, 1], self.motor_ids)
+        self.motor_bus.write("Lock", 1)
+        print("Motors set to velocity mode.")
+
+    def read_velocity(self):
+        """
+        Reads the raw speeds for all wheels. Returns a dictionary with motor names:
+        """
+        raw_speeds = self.motor_bus.read("Present_Speed", self.motor_ids)
+        return {
+            "left_wheel": int(raw_speeds[0]),
+            "back_wheel": int(raw_speeds[1]),
+            "right_wheel": int(raw_speeds[2]),
+        }
+
+    def set_velocity(self, command_speeds):
+        """
+        Sends raw velocity commands (16-bit encoded values) directly to the motor bus.
+        The order of speeds must correspond to self.motor_ids.
+        """
+        self.motor_bus.write("Goal_Speed", command_speeds, self.motor_ids)
+
+    def stop(self):
+        """Stops the robot by setting all motor speeds to zero."""
+        self.motor_bus.write("Goal_Speed", [0, 0, 0], self.motor_ids)
+        print("Motors stopped.")

lerobot/common/robot_devices/robots/stretch.py

@@ -0,0 +1,208 @@
+#!/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 time
+from dataclasses import replace
+
+import torch
+from stretch_body.gamepad_teleop import GamePadTeleop
+from stretch_body.robot import Robot as StretchAPI
+from stretch_body.robot_params import RobotParams
+
+from lerobot.common.robot_devices.robots.configs import StretchRobotConfig
+
+
+class StretchRobot(StretchAPI):
+    """Wrapper of stretch_body.robot.Robot"""
+
+    def __init__(self, config: StretchRobotConfig | None = None, **kwargs):
+        super().__init__()
+        if config is None:
+            self.config = StretchRobotConfig(**kwargs)
+        else:
+            # Overwrite config arguments using kwargs
+            self.config = replace(config, **kwargs)
+
+        self.robot_type = self.config.type
+        self.cameras = self.config.cameras
+        self.is_connected = False
+        self.teleop = None
+        self.logs = {}
+
+        # TODO(aliberts): test this
+        RobotParams.set_logging_level("WARNING")
+        RobotParams.set_logging_formatter("brief_console_formatter")
+
+        self.state_keys = None
+        self.action_keys = None
+
+    def connect(self) -> None:
+        self.is_connected = self.startup()
+        if not self.is_connected:
+            print("Another process is already using Stretch. Try running 'stretch_free_robot_process.py'")
+            raise ConnectionError()
+
+        for name in self.cameras:
+            self.cameras[name].connect()
+            self.is_connected = self.is_connected and self.cameras[name].is_connected
+
+        if not self.is_connected:
+            print("Could not connect to the cameras, check that all cameras are plugged-in.")
+            raise ConnectionError()
+
+        self.run_calibration()
+
+    def run_calibration(self) -> None:
+        if not self.is_homed():
+            self.home()
+
+    def teleop_step(
+        self, record_data=False
+    ) -> None | tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]:
+        # TODO(aliberts): return ndarrays instead of torch.Tensors
+        if not self.is_connected:
+            raise ConnectionError()
+
+        if self.teleop is None:
+            self.teleop = GamePadTeleop(robot_instance=False)
+            self.teleop.startup(robot=self)
+
+        before_read_t = time.perf_counter()
+        state = self.get_state()
+        action = self.teleop.gamepad_controller.get_state()
+        self.logs["read_pos_dt_s"] = time.perf_counter() - before_read_t
+
+        before_write_t = time.perf_counter()
+        self.teleop.do_motion(robot=self)
+        self.push_command()
+        self.logs["write_pos_dt_s"] = time.perf_counter() - before_write_t
+
+        if self.state_keys is None:
+            self.state_keys = list(state)
+
+        if not record_data:
+            return
+
+        state = torch.as_tensor(list(state.values()))
+        action = torch.as_tensor(list(action.values()))
+
+        # Capture images from cameras
+        images = {}
+        for name in self.cameras:
+            before_camread_t = time.perf_counter()
+            images[name] = self.cameras[name].async_read()
+            images[name] = torch.from_numpy(images[name])
+            self.logs[f"read_camera_{name}_dt_s"] = self.cameras[name].logs["delta_timestamp_s"]
+            self.logs[f"async_read_camera_{name}_dt_s"] = time.perf_counter() - before_camread_t
+
+        # Populate output dictionaries
+        obs_dict, action_dict = {}, {}
+        obs_dict["observation.state"] = state
+        action_dict["action"] = action
+        for name in self.cameras:
+            obs_dict[f"observation.images.{name}"] = images[name]
+
+        return obs_dict, action_dict
+
+    def get_state(self) -> dict:
+        status = self.get_status()
+        return {
+            "head_pan.pos": status["head"]["head_pan"]["pos"],
+            "head_tilt.pos": status["head"]["head_tilt"]["pos"],
+            "lift.pos": status["lift"]["pos"],
+            "arm.pos": status["arm"]["pos"],
+            "wrist_pitch.pos": status["end_of_arm"]["wrist_pitch"]["pos"],
+            "wrist_roll.pos": status["end_of_arm"]["wrist_roll"]["pos"],
+            "wrist_yaw.pos": status["end_of_arm"]["wrist_yaw"]["pos"],
+            "gripper.pos": status["end_of_arm"]["stretch_gripper"]["pos"],
+            "base_x.vel": status["base"]["x_vel"],
+            "base_y.vel": status["base"]["y_vel"],
+            "base_theta.vel": status["base"]["theta_vel"],
+        }
+
+    def capture_observation(self) -> dict:
+        # TODO(aliberts): return ndarrays instead of torch.Tensors
+        before_read_t = time.perf_counter()
+        state = self.get_state()
+        self.logs["read_pos_dt_s"] = time.perf_counter() - before_read_t
+
+        if self.state_keys is None:
+            self.state_keys = list(state)
+
+        state = torch.as_tensor(list(state.values()))
+
+        # Capture images from cameras
+        images = {}
+        for name in self.cameras:
+            before_camread_t = time.perf_counter()
+            images[name] = self.cameras[name].async_read()
+            images[name] = torch.from_numpy(images[name])
+            self.logs[f"read_camera_{name}_dt_s"] = self.cameras[name].logs["delta_timestamp_s"]
+            self.logs[f"async_read_camera_{name}_dt_s"] = time.perf_counter() - before_camread_t
+
+        # Populate output dictionaries
+        obs_dict = {}
+        obs_dict["observation.state"] = state
+        for name in self.cameras:
+            obs_dict[f"observation.images.{name}"] = images[name]
+
+        return obs_dict
+
+    def send_action(self, action: torch.Tensor) -> torch.Tensor:
+        # TODO(aliberts): return ndarrays instead of torch.Tensors
+        if not self.is_connected:
+            raise ConnectionError()
+
+        if self.teleop is None:
+            self.teleop = GamePadTeleop(robot_instance=False)
+            self.teleop.startup(robot=self)
+
+        if self.action_keys is None:
+            dummy_action = self.teleop.gamepad_controller.get_state()
+            self.action_keys = list(dummy_action.keys())
+
+        action_dict = dict(zip(self.action_keys, action.tolist(), strict=True))
+
+        before_write_t = time.perf_counter()
+        self.teleop.do_motion(state=action_dict, robot=self)
+        self.push_command()
+        self.logs["write_pos_dt_s"] = time.perf_counter() - before_write_t
+
+        # TODO(aliberts): return action_sent when motion is limited
+        return action
+
+    def print_logs(self) -> None:
+        pass
+        # TODO(aliberts): move robot-specific logs logic here
+
+    def teleop_safety_stop(self) -> None:
+        if self.teleop is not None:
+            self.teleop._safety_stop(robot=self)
+
+    def disconnect(self) -> None:
+        self.stop()
+        if self.teleop is not None:
+            self.teleop.gamepad_controller.stop()
+            self.teleop.stop()
+
+        if len(self.cameras) > 0:
+            for cam in self.cameras.values():
+                cam.disconnect()
+
+        self.is_connected = False
+
+    def __del__(self):
+        self.disconnect()

lerobot/common/robot_devices/robots/utils.py

@@ -0,0 +1,86 @@
+# 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 typing import Protocol
+
+from lerobot.common.robot_devices.robots.configs import (
+    AlohaRobotConfig,
+    KochBimanualRobotConfig,
+    KochRobotConfig,
+    LeKiwiRobotConfig,
+    ManipulatorRobotConfig,
+    MossRobotConfig,
+    RobotConfig,
+    So100RobotConfig,
+    StretchRobotConfig,
+)
+
+
+def get_arm_id(name, arm_type):
+    """Returns the string identifier of a robot arm. For instance, for a bimanual manipulator
+    like Aloha, it could be left_follower, right_follower, left_leader, or right_leader.
+    """
+    return f"{name}_{arm_type}"
+
+
+class Robot(Protocol):
+    # TODO(rcadene, aliberts): Add unit test checking the protocol is implemented in the corresponding classes
+    robot_type: str
+    features: dict
+
+    def connect(self): ...
+    def run_calibration(self): ...
+    def teleop_step(self, record_data=False): ...
+    def capture_observation(self): ...
+    def send_action(self, action): ...
+    def disconnect(self): ...
+
+
+def make_robot_config(robot_type: str, **kwargs) -> RobotConfig:
+    if robot_type == "aloha":
+        return AlohaRobotConfig(**kwargs)
+    elif robot_type == "koch":
+        return KochRobotConfig(**kwargs)
+    elif robot_type == "koch_bimanual":
+        return KochBimanualRobotConfig(**kwargs)
+    elif robot_type == "moss":
+        return MossRobotConfig(**kwargs)
+    elif robot_type == "so100":
+        return So100RobotConfig(**kwargs)
+    elif robot_type == "stretch":
+        return StretchRobotConfig(**kwargs)
+    elif robot_type == "lekiwi":
+        return LeKiwiRobotConfig(**kwargs)
+    else:
+        raise ValueError(f"Robot type '{robot_type}' is not available.")
+
+
+def make_robot_from_config(config: RobotConfig):
+    if isinstance(config, ManipulatorRobotConfig):
+        from lerobot.common.robot_devices.robots.manipulator import ManipulatorRobot
+
+        return ManipulatorRobot(config)
+    elif isinstance(config, LeKiwiRobotConfig):
+        from lerobot.common.robot_devices.robots.mobile_manipulator import MobileManipulator
+
+        return MobileManipulator(config)
+    else:
+        from lerobot.common.robot_devices.robots.stretch import StretchRobot
+
+        return StretchRobot(config)
+
+
+def make_robot(robot_type: str, **kwargs) -> Robot:
+    config = make_robot_config(robot_type, **kwargs)
+    return make_robot_from_config(config)

lerobot/common/robot_devices/utils.py

@@ -0,0 +1,65 @@
+# 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 platform
+import time
+
+
+def busy_wait(seconds):
+    if platform.system() == "Darwin":
+        # On Mac, `time.sleep` is not accurate and we need to use this while loop trick,
+        # but it consumes CPU cycles.
+        # TODO(rcadene): find an alternative: from python 11, time.sleep is precise
+        end_time = time.perf_counter() + seconds
+        while time.perf_counter() < end_time:
+            pass
+    else:
+        # On Linux time.sleep is accurate
+        if seconds > 0:
+            time.sleep(seconds)
+
+
+def safe_disconnect(func):
+    # TODO(aliberts): Allow to pass custom exceptions
+    # (e.g. ThreadServiceExit, KeyboardInterrupt, SystemExit, UnpluggedError, DynamixelCommError)
+    def wrapper(robot, *args, **kwargs):
+        try:
+            return func(robot, *args, **kwargs)
+        except Exception as e:
+            if robot.is_connected:
+                robot.disconnect()
+            raise e
+
+    return wrapper
+
+
+class RobotDeviceNotConnectedError(Exception):
+    """Exception raised when the robot device is not connected."""
+
+    def __init__(
+        self, message="This robot device is not connected. Try calling `robot_device.connect()` first."
+    ):
+        self.message = message
+        super().__init__(self.message)
+
+
+class RobotDeviceAlreadyConnectedError(Exception):
+    """Exception raised when the robot device is already connected."""
+
+    def __init__(
+        self,
+        message="This robot device is already connected. Try not calling `robot_device.connect()` twice.",
+    ):
+        self.message = message
+        super().__init__(self.message)

lerobot/common/utils/benchmark.py

@@ -0,0 +1,92 @@
+#!/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 threading
+import time
+from contextlib import ContextDecorator
+
+
+class TimeBenchmark(ContextDecorator):
+    """
+    Measures execution time using a context manager or decorator.
+
+    This class supports both context manager and decorator usage, and is thread-safe for multithreaded
+    environments.
+
+    Args:
+        print: If True, prints the elapsed time upon exiting the context or completing the function. Defaults
+        to False.
+
+    Examples:
+
+        Using as a context manager:
+
+        >>> benchmark = TimeBenchmark()
+        >>> with benchmark:
+        ...     time.sleep(1)
+        >>> print(f"Block took {benchmark.result:.4f} seconds")
+        Block took approximately 1.0000 seconds
+
+        Using with multithreading:
+
+        ```python
+        import threading
+
+        benchmark = TimeBenchmark()
+
+        def context_manager_example():
+            with benchmark:
+                time.sleep(0.01)
+            print(f"Block took {benchmark.result_ms:.2f} milliseconds")
+
+        threads = []
+        for _ in range(3):
+            t1 = threading.Thread(target=context_manager_example)
+            threads.append(t1)
+
+        for t in threads:
+            t.start()
+
+        for t in threads:
+            t.join()
+        ```
+        Expected output:
+        Block took approximately 10.00 milliseconds
+        Block took approximately 10.00 milliseconds
+        Block took approximately 10.00 milliseconds
+    """
+
+    def __init__(self, print=False):
+        self.local = threading.local()
+        self.print_time = print
+
+    def __enter__(self):
+        self.local.start_time = time.perf_counter()
+        return self
+
+    def __exit__(self, *exc):
+        self.local.end_time = time.perf_counter()
+        self.local.elapsed_time = self.local.end_time - self.local.start_time
+        if self.print_time:
+            print(f"Elapsed time: {self.local.elapsed_time:.4f} seconds")
+        return False
+
+    @property
+    def result(self):
+        return getattr(self.local, "elapsed_time", None)
+
+    @property
+    def result_ms(self):
+        return self.result * 1e3

lerobot/common/utils/hub.py

@@ -0,0 +1,202 @@
+# 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 pathlib import Path
+from tempfile import TemporaryDirectory
+from typing import Any, Type, TypeVar
+
+from huggingface_hub import HfApi
+from huggingface_hub.utils import validate_hf_hub_args
+
+T = TypeVar("T", bound="HubMixin")
+
+
+class HubMixin:
+    """
+    A Mixin containing the functionality to push an object to the hub.
+
+    This is similar to huggingface_hub.ModelHubMixin but is lighter and makes less assumptions about its
+    subclasses (in particular, the fact that it's not necessarily a model).
+
+    The inheriting classes must implement '_save_pretrained' and 'from_pretrained'.
+    """
+
+    def save_pretrained(
+        self,
+        save_directory: str | Path,
+        *,
+        repo_id: str | None = None,
+        push_to_hub: bool = False,
+        card_kwargs: dict[str, Any] | None = None,
+        **push_to_hub_kwargs,
+    ) -> str | None:
+        """
+        Save object in local directory.
+
+        Args:
+            save_directory (`str` or `Path`):
+                Path to directory in which the object will be saved.
+            push_to_hub (`bool`, *optional*, defaults to `False`):
+                Whether or not to push your object to the Huggingface Hub after saving it.
+            repo_id (`str`, *optional*):
+                ID of your repository on the Hub. Used only if `push_to_hub=True`. Will default to the folder name if
+                not provided.
+            card_kwargs (`Dict[str, Any]`, *optional*):
+                Additional arguments passed to the card template to customize the card.
+            push_to_hub_kwargs:
+                Additional key word arguments passed along to the [`~HubMixin.push_to_hub`] method.
+        Returns:
+            `str` or `None`: url of the commit on the Hub if `push_to_hub=True`, `None` otherwise.
+        """
+        save_directory = Path(save_directory)
+        save_directory.mkdir(parents=True, exist_ok=True)
+
+        # save object (weights, files, etc.)
+        self._save_pretrained(save_directory)
+
+        # push to the Hub if required
+        if push_to_hub:
+            if repo_id is None:
+                repo_id = save_directory.name  # Defaults to `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) -> None:
+        """
+        Overwrite this method in subclass to define how to save your object.
+
+        Args:
+            save_directory (`str` or `Path`):
+                Path to directory in which the object files will be saved.
+        """
+        raise NotImplementedError
+
+    @classmethod
+    @validate_hf_hub_args
+    def from_pretrained(
+        cls: Type[T],
+        pretrained_name_or_path: str | Path,
+        *,
+        force_download: bool = False,
+        resume_download: bool | None = None,
+        proxies: dict | None = None,
+        token: str | bool | None = None,
+        cache_dir: str | Path | None = None,
+        local_files_only: bool = False,
+        revision: str | None = None,
+        **kwargs,
+    ) -> T:
+        """
+        Download the object from the Huggingface Hub and instantiate it.
+
+        Args:
+            pretrained_name_or_path (`str`, `Path`):
+                - Either the `repo_id` (string) of the object hosted on the Hub, e.g. `lerobot/diffusion_pusht`.
+                - Or a path to a `directory` containing the object files saved using `.save_pretrained`,
+                    e.g., `../path/to/my_model_directory/`.
+            revision (`str`, *optional*):
+                Revision on the Hub. Can be a branch name, a git tag or any commit id.
+                Defaults to the latest commit on `main` branch.
+            force_download (`bool`, *optional*, defaults to `False`):
+                Whether to force (re-)downloading the files from the Hub, overriding the existing cache.
+            proxies (`Dict[str, str]`, *optional*):
+                A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',
+                'http://hostname': 'foo.bar:4012'}`. The proxies are used on every request.
+            token (`str` or `bool`, *optional*):
+                The token to use as HTTP bearer authorization for remote files. By default, it will use the token
+                cached when running `huggingface-cli login`.
+            cache_dir (`str`, `Path`, *optional*):
+                Path to the folder where cached files are stored.
+            local_files_only (`bool`, *optional*, defaults to `False`):
+                If `True`, avoid downloading the file and return the path to the local cached file if it exists.
+            kwargs (`Dict`, *optional*):
+                Additional kwargs to pass to the object during initialization.
+        """
+        raise NotImplementedError
+
+    @validate_hf_hub_args
+    def push_to_hub(
+        self,
+        repo_id: str,
+        *,
+        commit_message: str | None = None,
+        private: bool | None = None,
+        token: str | None = None,
+        branch: str | None = None,
+        create_pr: bool | None = None,
+        allow_patterns: list[str] | str | None = None,
+        ignore_patterns: list[str] | str | None = None,
+        delete_patterns: list[str] | str | None = None,
+        card_kwargs: dict[str, Any] | None = None,
+    ) -> str:
+        """
+        Upload model checkpoint to the Hub.
+
+        Use `allow_patterns` and `ignore_patterns` to precisely filter which files should be pushed to the hub. Use
+        `delete_patterns` to delete existing remote files in the same commit. See [`upload_folder`] reference for more
+        details.
+
+        Args:
+            repo_id (`str`):
+                ID of the repository to push to (example: `"username/my-model"`).
+            commit_message (`str`, *optional*):
+                Message to commit while pushing.
+            private (`bool`, *optional*):
+                Whether the repository created should be private.
+                If `None` (default), the repo will be public unless the organization's default is private.
+            token (`str`, *optional*):
+                The token to use as HTTP bearer authorization for remote files. By default, it will use the token
+                cached when running `huggingface-cli login`.
+            branch (`str`, *optional*):
+                The git branch on which to push the model. This defaults to `"main"`.
+            create_pr (`boolean`, *optional*):
+                Whether or not to create a Pull Request from `branch` with that commit. Defaults to `False`.
+            allow_patterns (`List[str]` or `str`, *optional*):
+                If provided, only files matching at least one pattern are pushed.
+            ignore_patterns (`List[str]` or `str`, *optional*):
+                If provided, files matching any of the patterns are not pushed.
+            delete_patterns (`List[str]` or `str`, *optional*):
+                If provided, remote files matching any of the patterns will be deleted from the repo.
+            card_kwargs (`Dict[str, Any]`, *optional*):
+                Additional arguments passed to the card template to customize the card.
+
+        Returns:
+            The url of the commit of your object in the given repository.
+        """
+        api = HfApi(token=token)
+        repo_id = api.create_repo(repo_id=repo_id, private=private, exist_ok=True).repo_id
+
+        if commit_message is None:
+            if "Policy" in self.__class__.__name__:
+                commit_message = "Upload policy"
+            elif "Config" in self.__class__.__name__:
+                commit_message = "Upload config"
+            else:
+                commit_message = f"Upload {self.__class__.__name__}"
+
+        # Push the files to the repo in a single commit
+        with TemporaryDirectory(ignore_cleanup_errors=True) as tmp:
+            saved_path = Path(tmp) / repo_id
+            self.save_pretrained(saved_path, card_kwargs=card_kwargs)
+            return api.upload_folder(
+                repo_id=repo_id,
+                repo_type="model",
+                folder_path=saved_path,
+                commit_message=commit_message,
+                revision=branch,
+                create_pr=create_pr,
+                allow_patterns=allow_patterns,
+                ignore_patterns=ignore_patterns,
+                delete_patterns=delete_patterns,
+            )

lerobot/common/utils/import_utils.py

@@ -0,0 +1,59 @@
+#!/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 importlib
+import logging
+
+
+def is_package_available(pkg_name: str, return_version: bool = False) -> tuple[bool, str] | bool:
+    """Copied from https://github.com/huggingface/transformers/blob/main/src/transformers/utils/import_utils.py
+    Check if the package spec exists and grab its version to avoid importing a local directory.
+    **Note:** this doesn't work for all packages.
+    """
+    package_exists = importlib.util.find_spec(pkg_name) is not None
+    package_version = "N/A"
+    if package_exists:
+        try:
+            # Primary method to get the package version
+            package_version = importlib.metadata.version(pkg_name)
+        except importlib.metadata.PackageNotFoundError:
+            # Fallback method: Only for "torch" and versions containing "dev"
+            if pkg_name == "torch":
+                try:
+                    package = importlib.import_module(pkg_name)
+                    temp_version = getattr(package, "__version__", "N/A")
+                    # Check if the version contains "dev"
+                    if "dev" in temp_version:
+                        package_version = temp_version
+                        package_exists = True
+                    else:
+                        package_exists = False
+                except ImportError:
+                    # If the package can't be imported, it's not available
+                    package_exists = False
+            else:
+                # For packages other than "torch", don't attempt the fallback and set as not available
+                package_exists = False
+        logging.debug(f"Detected {pkg_name} version: {package_version}")
+    if return_version:
+        return package_exists, package_version
+    else:
+        return package_exists
+
+
+_torch_available, _torch_version = is_package_available("torch", return_version=True)
+_gym_xarm_available = is_package_available("gym_xarm")
+_gym_aloha_available = is_package_available("gym_aloha")
+_gym_pusht_available = is_package_available("gym_pusht")

lerobot/common/utils/io_utils.py

@@ -0,0 +1,111 @@
+#!/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 warnings
+from pathlib import Path
+from typing import TypeVar
+
+import imageio
+
+JsonLike = str | int | float | bool | None | list["JsonLike"] | dict[str, "JsonLike"] | tuple["JsonLike", ...]
+T = TypeVar("T", bound=JsonLike)
+
+
+def write_video(video_path, stacked_frames, fps):
+    # Filter out DeprecationWarnings raised from pkg_resources
+    with warnings.catch_warnings():
+        warnings.filterwarnings(
+            "ignore", "pkg_resources is deprecated as an API", category=DeprecationWarning
+        )
+        imageio.mimsave(video_path, stacked_frames, fps=fps)
+
+
+def deserialize_json_into_object(fpath: Path, obj: T) -> T:
+    """
+    Loads the JSON data from `fpath` and recursively fills `obj` with the
+    corresponding values (strictly matching structure and types).
+    Tuples in `obj` are expected to be lists in the JSON data, which will be
+    converted back into tuples.
+    """
+    with open(fpath, encoding="utf-8") as f:
+        data = json.load(f)
+
+    def _deserialize(target, source):
+        """
+        Recursively overwrite the structure in `target` with data from `source`,
+        performing strict checks on structure and type.
+        Returns the updated version of `target` (especially important for tuples).
+        """
+
+        # If the target is a dictionary, source must be a dictionary as well.
+        if isinstance(target, dict):
+            if not isinstance(source, dict):
+                raise TypeError(f"Type mismatch: expected dict, got {type(source)}")
+
+            # Check that they have exactly the same set of keys.
+            if target.keys() != source.keys():
+                raise ValueError(
+                    f"Dictionary keys do not match.\nExpected: {target.keys()}, got: {source.keys()}"
+                )
+
+            # Recursively update each key.
+            for k in target:
+                target[k] = _deserialize(target[k], source[k])
+
+            return target
+
+        # If the target is a list, source must be a list as well.
+        elif isinstance(target, list):
+            if not isinstance(source, list):
+                raise TypeError(f"Type mismatch: expected list, got {type(source)}")
+
+            # Check length
+            if len(target) != len(source):
+                raise ValueError(f"List length mismatch: expected {len(target)}, got {len(source)}")
+
+            # Recursively update each element.
+            for i in range(len(target)):
+                target[i] = _deserialize(target[i], source[i])
+
+            return target
+
+        # If the target is a tuple, the source must be a list in JSON,
+        # which we'll convert back to a tuple.
+        elif isinstance(target, tuple):
+            if not isinstance(source, list):
+                raise TypeError(f"Type mismatch: expected list (for tuple), got {type(source)}")
+
+            if len(target) != len(source):
+                raise ValueError(f"Tuple length mismatch: expected {len(target)}, got {len(source)}")
+
+            # Convert each element, forming a new tuple.
+            converted_items = []
+            for t_item, s_item in zip(target, source, strict=False):
+                converted_items.append(_deserialize(t_item, s_item))
+
+            # Return a brand new tuple (tuples are immutable in Python).
+            return tuple(converted_items)
+
+        # Otherwise, we're dealing with a "primitive" (int, float, str, bool, None).
+        else:
+            # Check the exact type.  If these must match 1:1, do:
+            if type(target) is not type(source):
+                raise TypeError(f"Type mismatch: expected {type(target)}, got {type(source)}")
+            return source
+
+    # Perform the in-place/recursive deserialization
+    updated_obj = _deserialize(obj, data)
+    return updated_obj

lerobot/common/utils/logging_utils.py

@@ -0,0 +1,163 @@
+#!/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.
+from typing import Any
+
+from lerobot.common.utils.utils import format_big_number
+
+
+class AverageMeter:
+    """
+    Computes and stores the average and current value
+    Adapted from https://github.com/pytorch/examples/blob/main/imagenet/main.py
+    """
+
+    def __init__(self, name: str, fmt: str = ":f"):
+        self.name = name
+        self.fmt = fmt
+        self.reset()
+
+    def reset(self) -> None:
+        self.val = 0.0
+        self.avg = 0.0
+        self.sum = 0.0
+        self.count = 0.0
+
+    def update(self, val: float, n: int = 1) -> None:
+        self.val = val
+        self.sum += val * n
+        self.count += n
+        self.avg = self.sum / self.count
+
+    def __str__(self):
+        fmtstr = "{name}:{avg" + self.fmt + "}"
+        return fmtstr.format(**self.__dict__)
+
+
+class MetricsTracker:
+    """
+    A helper class to track and log metrics over time.
+
+    Usage pattern:
+
+    ```python
+    # initialize, potentially with non-zero initial step (e.g. if resuming run)
+    metrics = {"loss": AverageMeter("loss", ":.3f")}
+    train_metrics = MetricsTracker(cfg, dataset, metrics, initial_step=step)
+
+    # update metrics derived from step (samples, episodes, epochs) at each training step
+    train_metrics.step()
+
+    # update various metrics
+    loss = policy.forward(batch)
+    train_metrics.loss = loss
+
+    # display current metrics
+    logging.info(train_metrics)
+
+    # export for wandb
+    wandb.log(train_metrics.to_dict())
+
+    # reset averages after logging
+    train_metrics.reset_averages()
+    ```
+    """
+
+    __keys__ = [
+        "_batch_size",
+        "_num_frames",
+        "_avg_samples_per_ep",
+        "metrics",
+        "steps",
+        "samples",
+        "episodes",
+        "epochs",
+    ]
+
+    def __init__(
+        self,
+        batch_size: int,
+        num_frames: int,
+        num_episodes: int,
+        metrics: dict[str, AverageMeter],
+        initial_step: int = 0,
+    ):
+        self.__dict__.update({k: None for k in self.__keys__})
+        self._batch_size = batch_size
+        self._num_frames = num_frames
+        self._avg_samples_per_ep = num_frames / num_episodes
+        self.metrics = metrics
+
+        self.steps = initial_step
+        # A sample is an (observation,action) pair, where observation and action
+        # can be on multiple timestamps. In a batch, we have `batch_size` number of samples.
+        self.samples = self.steps * self._batch_size
+        self.episodes = self.samples / self._avg_samples_per_ep
+        self.epochs = self.samples / self._num_frames
+
+    def __getattr__(self, name: str) -> int | dict[str, AverageMeter] | AverageMeter | Any:
+        if name in self.__dict__:
+            return self.__dict__[name]
+        elif name in self.metrics:
+            return self.metrics[name]
+        else:
+            raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{name}'")
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        if name in self.__dict__:
+            super().__setattr__(name, value)
+        elif name in self.metrics:
+            self.metrics[name].update(value)
+        else:
+            raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{name}'")
+
+    def step(self) -> None:
+        """
+        Updates metrics that depend on 'step' for one step.
+        """
+        self.steps += 1
+        self.samples += self._batch_size
+        self.episodes = self.samples / self._avg_samples_per_ep
+        self.epochs = self.samples / self._num_frames
+
+    def __str__(self) -> str:
+        display_list = [
+            f"step:{format_big_number(self.steps)}",
+            # number of samples seen during training
+            f"smpl:{format_big_number(self.samples)}",
+            # number of episodes seen during training
+            f"ep:{format_big_number(self.episodes)}",
+            # number of time all unique samples are seen
+            f"epch:{self.epochs:.2f}",
+            *[str(m) for m in self.metrics.values()],
+        ]
+        return " ".join(display_list)
+
+    def to_dict(self, use_avg: bool = True) -> dict[str, int | float]:
+        """
+        Returns the current metric values (or averages if `use_avg=True`) as a dict.
+        """
+        return {
+            "steps": self.steps,
+            "samples": self.samples,
+            "episodes": self.episodes,
+            "epochs": self.epochs,
+            **{k: m.avg if use_avg else m.val for k, m in self.metrics.items()},
+        }
+
+    def reset_averages(self) -> None:
+        """Resets average meters."""
+        for m in self.metrics.values():
+            m.reset()

lerobot/common/utils/random_utils.py

@@ -0,0 +1,191 @@
+#!/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 random
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Any, Generator
+
+import numpy as np
+import torch
+from safetensors.torch import load_file, save_file
+
+from lerobot.common.constants import RNG_STATE
+from lerobot.common.datasets.utils import flatten_dict, unflatten_dict
+
+
+def serialize_python_rng_state() -> dict[str, torch.Tensor]:
+    """
+    Returns the rng state for `random` in the form of a flat dict[str, torch.Tensor] to be saved using
+    `safetensors.save_file()` or `torch.save()`.
+    """
+    py_state = random.getstate()
+    return {
+        "py_rng_version": torch.tensor([py_state[0]], dtype=torch.int64),
+        "py_rng_state": torch.tensor(py_state[1], dtype=torch.int64),
+    }
+
+
+def deserialize_python_rng_state(rng_state_dict: dict[str, torch.Tensor]) -> None:
+    """
+    Restores the rng state for `random` from a dictionary produced by `serialize_python_rng_state()`.
+    """
+    py_state = (rng_state_dict["py_rng_version"].item(), tuple(rng_state_dict["py_rng_state"].tolist()), None)
+    random.setstate(py_state)
+
+
+def serialize_numpy_rng_state() -> dict[str, torch.Tensor]:
+    """
+    Returns the rng state for `numpy` in the form of a flat dict[str, torch.Tensor] to be saved using
+    `safetensors.save_file()` or `torch.save()`.
+    """
+    np_state = np.random.get_state()
+    # Ensure no breaking changes from numpy
+    assert np_state[0] == "MT19937"
+    return {
+        "np_rng_state_values": torch.tensor(np_state[1], dtype=torch.int64),
+        "np_rng_state_index": torch.tensor([np_state[2]], dtype=torch.int64),
+        "np_rng_has_gauss": torch.tensor([np_state[3]], dtype=torch.int64),
+        "np_rng_cached_gaussian": torch.tensor([np_state[4]], dtype=torch.float32),
+    }
+
+
+def deserialize_numpy_rng_state(rng_state_dict: dict[str, torch.Tensor]) -> None:
+    """
+    Restores the rng state for `numpy` from a dictionary produced by `serialize_numpy_rng_state()`.
+    """
+    np_state = (
+        "MT19937",
+        rng_state_dict["np_rng_state_values"].numpy(),
+        rng_state_dict["np_rng_state_index"].item(),
+        rng_state_dict["np_rng_has_gauss"].item(),
+        rng_state_dict["np_rng_cached_gaussian"].item(),
+    )
+    np.random.set_state(np_state)
+
+
+def serialize_torch_rng_state() -> dict[str, torch.Tensor]:
+    """
+    Returns the rng state for `torch` in the form of a flat dict[str, torch.Tensor] to be saved using
+    `safetensors.save_file()` or `torch.save()`.
+    """
+    torch_rng_state_dict = {"torch_rng_state": torch.get_rng_state()}
+    if torch.cuda.is_available():
+        torch_rng_state_dict["torch_cuda_rng_state"] = torch.cuda.get_rng_state()
+    return torch_rng_state_dict
+
+
+def deserialize_torch_rng_state(rng_state_dict: dict[str, torch.Tensor]) -> None:
+    """
+    Restores the rng state for `torch` from a dictionary produced by `serialize_torch_rng_state()`.
+    """
+    torch.set_rng_state(rng_state_dict["torch_rng_state"])
+    if torch.cuda.is_available() and "torch_cuda_rng_state" in rng_state_dict:
+        torch.cuda.set_rng_state(rng_state_dict["torch_cuda_rng_state"])
+
+
+def serialize_rng_state() -> dict[str, torch.Tensor]:
+    """
+    Returns the rng state for `random`, `numpy`, and `torch`, in the form of a flat
+    dict[str, torch.Tensor] to be saved using `safetensors.save_file()` `torch.save()`.
+    """
+    py_rng_state_dict = serialize_python_rng_state()
+    np_rng_state_dict = serialize_numpy_rng_state()
+    torch_rng_state_dict = serialize_torch_rng_state()
+
+    return {
+        **py_rng_state_dict,
+        **np_rng_state_dict,
+        **torch_rng_state_dict,
+    }
+
+
+def deserialize_rng_state(rng_state_dict: dict[str, torch.Tensor]) -> None:
+    """
+    Restores the rng state for `random`, `numpy`, and `torch` from a dictionary produced by
+    `serialize_rng_state()`.
+    """
+    py_rng_state_dict = {k: v for k, v in rng_state_dict.items() if k.startswith("py")}
+    np_rng_state_dict = {k: v for k, v in rng_state_dict.items() if k.startswith("np")}
+    torch_rng_state_dict = {k: v for k, v in rng_state_dict.items() if k.startswith("torch")}
+
+    deserialize_python_rng_state(py_rng_state_dict)
+    deserialize_numpy_rng_state(np_rng_state_dict)
+    deserialize_torch_rng_state(torch_rng_state_dict)
+
+
+def save_rng_state(save_dir: Path) -> None:
+    rng_state_dict = serialize_rng_state()
+    flat_rng_state_dict = flatten_dict(rng_state_dict)
+    save_file(flat_rng_state_dict, save_dir / RNG_STATE)
+
+
+def load_rng_state(save_dir: Path) -> None:
+    flat_rng_state_dict = load_file(save_dir / RNG_STATE)
+    rng_state_dict = unflatten_dict(flat_rng_state_dict)
+    deserialize_rng_state(rng_state_dict)
+
+
+def get_rng_state() -> dict[str, Any]:
+    """Get the random state for `random`, `numpy`, and `torch`."""
+    random_state_dict = {
+        "random_state": random.getstate(),
+        "numpy_random_state": np.random.get_state(),
+        "torch_random_state": torch.random.get_rng_state(),
+    }
+    if torch.cuda.is_available():
+        random_state_dict["torch_cuda_random_state"] = torch.cuda.random.get_rng_state()
+    return random_state_dict
+
+
+def set_rng_state(random_state_dict: dict[str, Any]):
+    """Set the random state for `random`, `numpy`, and `torch`.
+
+    Args:
+        random_state_dict: A dictionary of the form returned by `get_rng_state`.
+    """
+    random.setstate(random_state_dict["random_state"])
+    np.random.set_state(random_state_dict["numpy_random_state"])
+    torch.random.set_rng_state(random_state_dict["torch_random_state"])
+    if torch.cuda.is_available():
+        torch.cuda.random.set_rng_state(random_state_dict["torch_cuda_random_state"])
+
+
+def set_seed(seed) -> None:
+    """Set seed for reproducibility."""
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+@contextmanager
+def seeded_context(seed: int) -> Generator[None, None, None]:
+    """Set the seed when entering a context, and restore the prior random state at exit.
+
+    Example usage:
+
+    ```
+    a = random.random()  # produces some random number
+    with seeded_context(1337):
+        b = random.random()  # produces some other random number
+    c = random.random()  # produces yet another random number, but the same it would have if we never made `b`
+    ```
+    """
+    random_state_dict = get_rng_state()
+    set_seed(seed)
+    yield None
+    set_rng_state(random_state_dict)

lerobot/common/utils/train_utils.py

@@ -0,0 +1,161 @@
+#!/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 logging
+from pathlib import Path
+
+from termcolor import colored
+from torch.optim import Optimizer
+from torch.optim.lr_scheduler import LRScheduler
+
+from lerobot.common.constants import (
+    CHECKPOINTS_DIR,
+    LAST_CHECKPOINT_LINK,
+    PRETRAINED_MODEL_DIR,
+    TRAINING_STATE_DIR,
+    TRAINING_STEP,
+)
+from lerobot.common.datasets.utils import load_json, write_json
+from lerobot.common.optim.optimizers import load_optimizer_state, save_optimizer_state
+from lerobot.common.optim.schedulers import load_scheduler_state, save_scheduler_state
+from lerobot.common.policies.pretrained import PreTrainedPolicy
+from lerobot.common.utils.random_utils import load_rng_state, save_rng_state
+from lerobot.configs.train import TrainPipelineConfig
+
+
+def log_output_dir(out_dir):
+    logging.info(colored("Output dir:", "yellow", attrs=["bold"]) + f" {out_dir}")
+
+
+def get_step_identifier(step: int, total_steps: int) -> str:
+    num_digits = max(6, len(str(total_steps)))
+    return f"{step:0{num_digits}d}"
+
+
+def get_step_checkpoint_dir(output_dir: Path, total_steps: int, step: int) -> Path:
+    """Returns the checkpoint sub-directory corresponding to the step number."""
+    step_identifier = get_step_identifier(step, total_steps)
+    return output_dir / CHECKPOINTS_DIR / step_identifier
+
+
+def save_training_step(step: int, save_dir: Path) -> None:
+    write_json({"step": step}, save_dir / TRAINING_STEP)
+
+
+def load_training_step(save_dir: Path) -> int:
+    training_step = load_json(save_dir / TRAINING_STEP)
+    return training_step["step"]
+
+
+def update_last_checkpoint(checkpoint_dir: Path) -> Path:
+    last_checkpoint_dir = checkpoint_dir.parent / LAST_CHECKPOINT_LINK
+    if last_checkpoint_dir.is_symlink():
+        last_checkpoint_dir.unlink()
+    relative_target = checkpoint_dir.relative_to(checkpoint_dir.parent)
+    last_checkpoint_dir.symlink_to(relative_target)
+
+
+def save_checkpoint(
+    checkpoint_dir: Path,
+    step: int,
+    cfg: TrainPipelineConfig,
+    policy: PreTrainedPolicy,
+    optimizer: Optimizer,
+    scheduler: LRScheduler | None = None,
+) -> None:
+    """This function creates the following directory structure:
+
+    005000/  #  training step at checkpoint
+    ├── pretrained_model/
+    │   ├── config.json  # policy config
+    │   ├── model.safetensors  # policy weights
+    │   └── train_config.json  # train config
+    └── training_state/
+        ├── optimizer_param_groups.json  #  optimizer param groups
+        ├── optimizer_state.safetensors  # optimizer state
+        ├── rng_state.safetensors  # rng states
+        ├── scheduler_state.json  # scheduler state
+        └── training_step.json  # training step
+
+    Args:
+        cfg (TrainPipelineConfig): The training config used for this run.
+        step (int): The training step at that checkpoint.
+        policy (PreTrainedPolicy): The policy to save.
+        optimizer (Optimizer | None, optional): The optimizer to save the state from. Defaults to None.
+        scheduler (LRScheduler | None, optional): The scheduler to save the state from. Defaults to None.
+    """
+    pretrained_dir = checkpoint_dir / PRETRAINED_MODEL_DIR
+    policy.save_pretrained(pretrained_dir)
+    cfg.save_pretrained(pretrained_dir)
+    save_training_state(checkpoint_dir, step, optimizer, scheduler)
+
+
+def save_training_state(
+    checkpoint_dir: Path,
+    train_step: int,
+    optimizer: Optimizer | None = None,
+    scheduler: LRScheduler | None = None,
+) -> None:
+    """
+    Saves the training step, optimizer state, scheduler state, and rng state.
+
+    Args:
+        save_dir (Path): The directory to save artifacts to.
+        train_step (int): Current training step.
+        optimizer (Optimizer | None, optional): The optimizer from which to save the state_dict.
+            Defaults to None.
+        scheduler (LRScheduler | None, optional): The scheduler from which to save the state_dict.
+            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)
+    save_rng_state(save_dir)
+    if optimizer is not None:
+        save_optimizer_state(optimizer, save_dir)
+    if scheduler is not None:
+        save_scheduler_state(scheduler, save_dir)
+
+
+def load_training_state(
+    checkpoint_dir: Path, optimizer: Optimizer, scheduler: LRScheduler | None
+) -> tuple[int, Optimizer, LRScheduler | None]:
+    """
+    Loads the training step, optimizer state, scheduler state, and rng state.
+    This is used to resume a training run.
+
+    Args:
+        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).
+
+    Raises:
+        NotADirectoryError: If 'checkpoint_dir' doesn't contain a 'training_state' dir
+
+    Returns:
+        tuple[int, Optimizer, LRScheduler | None]: training step, optimizer and scheduler with their
+            state_dict loaded.
+    """
+    training_state_dir = checkpoint_dir / TRAINING_STATE_DIR
+    if not training_state_dir.is_dir():
+        raise NotADirectoryError(training_state_dir)
+
+    load_rng_state(training_state_dir)
+    step = load_training_step(training_state_dir)
+    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

lerobot/common/utils/utils.py

@@ -0,0 +1,230 @@
+#!/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 logging
+import os
+import os.path as osp
+import platform
+import subprocess
+from copy import copy
+from datetime import datetime, timezone
+from pathlib import Path
+
+import numpy as np
+import torch
+
+
+def none_or_int(value):
+    if value == "None":
+        return None
+    return int(value)
+
+
+def inside_slurm():
+    """Check whether the python process was launched through slurm"""
+    # TODO(rcadene): return False for interactive mode `--pty bash`
+    return "SLURM_JOB_ID" in os.environ
+
+
+def auto_select_torch_device() -> torch.device:
+    """Tries to select automatically a torch device."""
+    if torch.cuda.is_available():
+        logging.info("Cuda backend detected, using cuda.")
+        return torch.device("cuda")
+    elif torch.backends.mps.is_available():
+        logging.info("Metal backend detected, using cuda.")
+        return torch.device("mps")
+    else:
+        logging.warning("No accelerated backend detected. Using default cpu, this will be slow.")
+        return torch.device("cpu")
+
+
+# TODO(Steven): Remove log. log shouldn't be an argument, this should be handled by the logger level
+def get_safe_torch_device(try_device: str, log: bool = False) -> torch.device:
+    """Given a string, return a torch.device with checks on whether the device is available."""
+    try_device = str(try_device)
+    match try_device:
+        case "cuda":
+            assert torch.cuda.is_available()
+            device = torch.device("cuda")
+        case "mps":
+            assert torch.backends.mps.is_available()
+            device = torch.device("mps")
+        case "cpu":
+            device = torch.device("cpu")
+            if log:
+                logging.warning("Using CPU, this will be slow.")
+        case _:
+            device = torch.device(try_device)
+            if log:
+                logging.warning(f"Using custom {try_device} device.")
+
+    return device
+
+
+def get_safe_dtype(dtype: torch.dtype, device: str | torch.device):
+    """
+    mps is currently not compatible with float64
+    """
+    if isinstance(device, torch.device):
+        device = device.type
+    if device == "mps" and dtype == torch.float64:
+        return torch.float32
+    else:
+        return dtype
+
+
+def is_torch_device_available(try_device: str) -> bool:
+    try_device = str(try_device)  # Ensure try_device is a string
+    if try_device == "cuda":
+        return torch.cuda.is_available()
+    elif try_device == "mps":
+        return torch.backends.mps.is_available()
+    elif try_device == "cpu":
+        return True
+    else:
+        raise ValueError(f"Unknown device {try_device}. Supported devices are: cuda, mps or cpu.")
+
+
+def is_amp_available(device: str):
+    if device in ["cuda", "cpu"]:
+        return True
+    elif device == "mps":
+        return False
+    else:
+        raise ValueError(f"Unknown device '{device}.")
+
+
+def init_logging():
+    def custom_format(record):
+        dt = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+        fnameline = f"{record.pathname}:{record.lineno}"
+        message = f"{record.levelname} {dt} {fnameline[-15:]:>15} {record.msg}"
+        return message
+
+    logging.basicConfig(level=logging.INFO)
+
+    for handler in logging.root.handlers[:]:
+        logging.root.removeHandler(handler)
+
+    formatter = logging.Formatter()
+    formatter.format = custom_format
+    console_handler = logging.StreamHandler()
+    console_handler.setFormatter(formatter)
+    logging.getLogger().addHandler(console_handler)
+
+
+def format_big_number(num, precision=0):
+    suffixes = ["", "K", "M", "B", "T", "Q"]
+    divisor = 1000.0
+
+    for suffix in suffixes:
+        if abs(num) < divisor:
+            return f"{num:.{precision}f}{suffix}"
+        num /= divisor
+
+    return num
+
+
+def _relative_path_between(path1: Path, path2: Path) -> Path:
+    """Returns path1 relative to path2."""
+    path1 = path1.absolute()
+    path2 = path2.absolute()
+    try:
+        return path1.relative_to(path2)
+    except ValueError:  # most likely because path1 is not a subpath of path2
+        common_parts = Path(osp.commonpath([path1, path2])).parts
+        return Path(
+            "/".join([".."] * (len(path2.parts) - len(common_parts)) + list(path1.parts[len(common_parts) :]))
+        )
+
+
+def print_cuda_memory_usage():
+    """Use this function to locate and debug memory leak."""
+    import gc
+
+    gc.collect()
+    # Also clear the cache if you want to fully release the memory
+    torch.cuda.empty_cache()
+    print("Current GPU Memory Allocated: {:.2f} MB".format(torch.cuda.memory_allocated(0) / 1024**2))
+    print("Maximum GPU Memory Allocated: {:.2f} MB".format(torch.cuda.max_memory_allocated(0) / 1024**2))
+    print("Current GPU Memory Reserved: {:.2f} MB".format(torch.cuda.memory_reserved(0) / 1024**2))
+    print("Maximum GPU Memory Reserved: {:.2f} MB".format(torch.cuda.max_memory_reserved(0) / 1024**2))
+
+
+def capture_timestamp_utc():
+    return datetime.now(timezone.utc)
+
+
+def say(text, blocking=False):
+    system = platform.system()
+
+    if system == "Darwin":
+        cmd = ["say", text]
+
+    elif system == "Linux":
+        cmd = ["spd-say", text]
+        if blocking:
+            cmd.append("--wait")
+
+    elif system == "Windows":
+        cmd = [
+            "PowerShell",
+            "-Command",
+            "Add-Type -AssemblyName System.Speech; "
+            f"(New-Object System.Speech.Synthesis.SpeechSynthesizer).Speak('{text}')",
+        ]
+
+    else:
+        raise RuntimeError("Unsupported operating system for text-to-speech.")
+
+    if blocking:
+        subprocess.run(cmd, check=True)
+    else:
+        subprocess.Popen(cmd, creationflags=subprocess.CREATE_NO_WINDOW if system == "Windows" else 0)
+
+
+def log_say(text, play_sounds, blocking=False):
+    logging.info(text)
+
+    if play_sounds:
+        say(text, blocking)
+
+
+def get_channel_first_image_shape(image_shape: tuple) -> tuple:
+    shape = copy(image_shape)
+    if shape[2] < shape[0] and shape[2] < shape[1]:  # (h, w, c) -> (c, h, w)
+        shape = (shape[2], shape[0], shape[1])
+    elif not (shape[0] < shape[1] and shape[0] < shape[2]):
+        raise ValueError(image_shape)
+
+    return shape
+
+
+def has_method(cls: object, method_name: str) -> bool:
+    return hasattr(cls, method_name) and callable(getattr(cls, method_name))
+
+
+def is_valid_numpy_dtype_string(dtype_str: str) -> bool:
+    """
+    Return True if a given string can be converted to a numpy dtype.
+    """
+    try:
+        # Attempt to convert the string to a numpy dtype
+        np.dtype(dtype_str)
+        return True
+    except TypeError:
+        # If a TypeError is raised, the string is not a valid dtype
+        return False

lerobot/common/utils/wandb_utils.py

@@ -0,0 +1,127 @@
+#!/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 logging
+import os
+import re
+from glob import glob
+from pathlib import Path
+
+from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE
+from termcolor import colored
+
+from lerobot.common.constants import PRETRAINED_MODEL_DIR
+from lerobot.configs.train import TrainPipelineConfig
+
+
+def cfg_to_group(cfg: TrainPipelineConfig, return_list: bool = False) -> list[str] | str:
+    """Return a group name for logging. Optionally returns group name as list."""
+    lst = [
+        f"policy:{cfg.policy.type}",
+        f"dataset:{cfg.dataset.repo_id}",
+        f"seed:{cfg.seed}",
+    ]
+    if cfg.env is not None:
+        lst.append(f"env:{cfg.env.type}")
+    return lst if return_list else "-".join(lst)
+
+
+def get_wandb_run_id_from_filesystem(log_dir: Path) -> str:
+    # Get the WandB run ID.
+    paths = glob(str(log_dir / "wandb/latest-run/run-*"))
+    if len(paths) != 1:
+        raise RuntimeError("Couldn't get the previous WandB run ID for run resumption.")
+    match = re.search(r"run-([^\.]+).wandb", paths[0].split("/")[-1])
+    if match is None:
+        raise RuntimeError("Couldn't get the previous WandB run ID for run resumption.")
+    wandb_run_id = match.groups(0)[0]
+    return wandb_run_id
+
+
+def get_safe_wandb_artifact_name(name: str):
+    """WandB artifacts don't accept ":" or "/" in their name."""
+    return name.replace(":", "_").replace("/", "_")
+
+
+class WandBLogger:
+    """A helper class to log object using wandb."""
+
+    def __init__(self, cfg: TrainPipelineConfig):
+        self.cfg = cfg.wandb
+        self.log_dir = cfg.output_dir
+        self.job_name = cfg.job_name
+        self.env_fps = cfg.env.fps if cfg.env else None
+        self._group = cfg_to_group(cfg)
+
+        # Set up WandB.
+        os.environ["WANDB_SILENT"] = "True"
+        import wandb
+
+        wandb_run_id = (
+            cfg.wandb.run_id
+            if cfg.wandb.run_id
+            else get_wandb_run_id_from_filesystem(self.log_dir)
+            if cfg.resume
+            else None
+        )
+        wandb.init(
+            id=wandb_run_id,
+            project=self.cfg.project,
+            entity=self.cfg.entity,
+            name=self.job_name,
+            notes=self.cfg.notes,
+            tags=cfg_to_group(cfg, return_list=True),
+            dir=self.log_dir,
+            config=cfg.to_dict(),
+            # TODO(rcadene): try set to True
+            save_code=False,
+            # TODO(rcadene): split train and eval, and run async eval with job_type="eval"
+            job_type="train_eval",
+            resume="must" if cfg.resume else None,
+        )
+        print(colored("Logs will be synced with wandb.", "blue", attrs=["bold"]))
+        logging.info(f"Track this run --> {colored(wandb.run.get_url(), 'yellow', attrs=['bold'])}")
+        self._wandb = wandb
+
+    def log_policy(self, checkpoint_dir: Path):
+        """Checkpoints the policy to wandb."""
+        if self.cfg.disable_artifact:
+            return
+
+        step_id = checkpoint_dir.name
+        artifact_name = f"{self._group}-{step_id}"
+        artifact_name = get_safe_wandb_artifact_name(artifact_name)
+        artifact = self._wandb.Artifact(artifact_name, type="model")
+        artifact.add_file(checkpoint_dir / PRETRAINED_MODEL_DIR / SAFETENSORS_SINGLE_FILE)
+        self._wandb.log_artifact(artifact)
+
+    def log_dict(self, d: dict, step: int, mode: str = "train"):
+        if mode not in {"train", "eval"}:
+            raise ValueError(mode)
+
+        for k, v in d.items():
+            if not isinstance(v, (int, float, str)):
+                logging.warning(
+                    f'WandB logging of key "{k}" was ignored as its type is not handled by this wrapper.'
+                )
+                continue
+            self._wandb.log({f"{mode}/{k}": v}, step=step)
+
+    def log_video(self, video_path: str, step: int, mode: str = "train"):
+        if mode not in {"train", "eval"}:
+            raise ValueError(mode)
+
+        wandb_video = self._wandb.Video(video_path, fps=self.env_fps, format="mp4")
+        self._wandb.log({f"{mode}/video": wandb_video}, step=step)

lerobot/configs/default.py

@@ -0,0 +1,70 @@
+#!/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.
+
+from dataclasses import dataclass, field
+
+from lerobot.common import (
+    policies,  # noqa: F401
+)
+from lerobot.common.datasets.transforms import ImageTransformsConfig
+from lerobot.common.datasets.video_utils import get_safe_default_codec
+
+
+@dataclass
+class DatasetConfig:
+    # You may provide a list of datasets here. `train.py` creates them all and concatenates them. Note: only data
+    # keys common between the datasets are kept. Each dataset gets and additional transform that inserts the
+    # "dataset_index" into the returned item. The index mapping is made according to the order in which the
+    # datasets are provided.
+    repo_id: str
+    # Root directory where the dataset will be stored (e.g. 'dataset/path').
+    root: str | None = None
+    episodes: list[int] | None = None
+    image_transforms: ImageTransformsConfig = field(default_factory=ImageTransformsConfig)
+    revision: str | None = None
+    use_imagenet_stats: bool = True
+    video_backend: str = field(default_factory=get_safe_default_codec)
+
+
+@dataclass
+class WandBConfig:
+    enable: bool = False
+    # Set to true to disable saving an artifact despite training.save_checkpoint=True
+    disable_artifact: bool = False
+    project: str = "lerobot"
+    entity: str | None = None
+    notes: str | None = None
+    run_id: str | None = None
+
+
+@dataclass
+class EvalConfig:
+    n_episodes: int = 50
+    # `batch_size` specifies the number of environments to use in a gym.vector.VectorEnv.
+    batch_size: int = 50
+    # `use_async_envs` specifies whether to use asynchronous environments (multiprocessing).
+    use_async_envs: bool = False
+
+    def __post_init__(self):
+        if self.batch_size > self.n_episodes:
+            raise ValueError(
+                "The eval batch size is greater than the number of eval episodes "
+                f"({self.batch_size} > {self.n_episodes}). As a result, {self.batch_size} "
+                f"eval environments will be instantiated, but only {self.n_episodes} will be used. "
+                "This might significantly slow down evaluation. To fix this, you should update your command "
+                f"to increase the number of episodes to match the batch size (e.g. `eval.n_episodes={self.batch_size}`), "
+                f"or lower the batch size (e.g. `eval.batch_size={self.n_episodes}`)."
+            )

lerobot/configs/eval.py

@@ -0,0 +1,65 @@
+# 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 datetime as dt
+import logging
+from dataclasses import dataclass, field
+from pathlib import Path
+
+from lerobot.common import envs, policies  # noqa: F401
+from lerobot.configs import parser
+from lerobot.configs.default import EvalConfig
+from lerobot.configs.policies import PreTrainedConfig
+
+
+@dataclass
+class EvalPipelineConfig:
+    # 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
+    # (useful for debugging). This argument is mutually exclusive with `--config`.
+    env: envs.EnvConfig
+    eval: EvalConfig = field(default_factory=EvalConfig)
+    policy: PreTrainedConfig | None = None
+    output_dir: Path | None = None
+    job_name: str | None = None
+    seed: int | None = 1000
+
+    def __post_init__(self):
+        # HACK: We parse again the cli args here to get the pretrained path if there was one.
+        policy_path = parser.get_path_arg("policy")
+        if policy_path:
+            cli_overrides = parser.get_cli_overrides("policy")
+            self.policy = PreTrainedConfig.from_pretrained(policy_path, cli_overrides=cli_overrides)
+            self.policy.pretrained_path = policy_path
+
+        else:
+            logging.warning(
+                "No pretrained path was provided, evaluated policy will be built from scratch (random weights)."
+            )
+
+        if not self.job_name:
+            if self.env is None:
+                self.job_name = f"{self.policy.type}"
+            else:
+                self.job_name = f"{self.env.type}_{self.policy.type}"
+
+        if not self.output_dir:
+            now = dt.datetime.now()
+            eval_dir = f"{now:%Y-%m-%d}/{now:%H-%M-%S}_{self.job_name}"
+            self.output_dir = Path("outputs/eval") / eval_dir
+
+    @classmethod
+    def __get_path_fields__(cls) -> list[str]:
+        """This enables the parser to load config from the policy using `--policy.path=local/dir`"""
+        return ["policy"]

lerobot/configs/parser.py

@@ -0,0 +1,232 @@
+# 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 importlib
+import inspect
+import pkgutil
+import sys
+from argparse import ArgumentError
+from functools import wraps
+from pathlib import Path
+from typing import Sequence
+
+import draccus
+
+from lerobot.common.utils.utils import has_method
+
+PATH_KEY = "path"
+PLUGIN_DISCOVERY_SUFFIX = "discover_packages_path"
+draccus.set_config_type("json")
+
+
+def get_cli_overrides(field_name: str, args: Sequence[str] | None = None) -> list[str] | None:
+    """Parses arguments from cli at a given nested attribute level.
+
+    For example, supposing the main script was called with:
+    python myscript.py --arg1=1 --arg2.subarg1=abc --arg2.subarg2=some/path
+
+    If called during execution of myscript.py, get_cli_overrides("arg2") will return:
+    ["--subarg1=abc" "--subarg2=some/path"]
+    """
+    if args is None:
+        args = sys.argv[1:]
+    attr_level_args = []
+    detect_string = f"--{field_name}."
+    exclude_strings = (f"--{field_name}.{draccus.CHOICE_TYPE_KEY}=", f"--{field_name}.{PATH_KEY}=")
+    for arg in args:
+        if arg.startswith(detect_string) and not arg.startswith(exclude_strings):
+            denested_arg = f"--{arg.removeprefix(detect_string)}"
+            attr_level_args.append(denested_arg)
+
+    return attr_level_args
+
+
+def parse_arg(arg_name: str, args: Sequence[str] | None = None) -> str | None:
+    if args is None:
+        args = sys.argv[1:]
+    prefix = f"--{arg_name}="
+    for arg in args:
+        if arg.startswith(prefix):
+            return arg[len(prefix) :]
+    return None
+
+
+def parse_plugin_args(plugin_arg_suffix: str, args: Sequence[str]) -> dict:
+    """Parse plugin-related arguments from command-line arguments.
+
+    This function extracts arguments from command-line arguments that match a specified suffix pattern.
+    It processes arguments in the format '--key=value' and returns them as a dictionary.
+
+    Args:
+        plugin_arg_suffix (str): The suffix to identify plugin-related arguments.
+        cli_args (Sequence[str]): A sequence of command-line arguments to parse.
+
+    Returns:
+        dict: A dictionary containing the parsed plugin arguments where:
+            - Keys are the argument names (with '--' prefix removed if present)
+            - Values are the corresponding argument values
+
+    Example:
+        >>> args = ['--env.discover_packages_path=my_package',
+        ...         '--other_arg=value']
+        >>> parse_plugin_args('discover_packages_path', args)
+        {'env.discover_packages_path': 'my_package'}
+    """
+    plugin_args = {}
+    for arg in args:
+        if "=" in arg and plugin_arg_suffix in arg:
+            key, value = arg.split("=", 1)
+            # Remove leading '--' if present
+            if key.startswith("--"):
+                key = key[2:]
+            plugin_args[key] = value
+    return plugin_args
+
+
+class PluginLoadError(Exception):
+    """Raised when a plugin fails to load."""
+
+
+def load_plugin(plugin_path: str) -> None:
+    """Load and initialize a plugin from a given Python package path.
+
+    This function attempts to load a plugin by importing its package and any submodules.
+    Plugin registration is expected to happen during package initialization, i.e. when
+    the package is imported the gym environment should be registered and the config classes
+    registered with their parents using the `register_subclass` decorator.
+
+    Args:
+        plugin_path (str): The Python package path to the plugin (e.g. "mypackage.plugins.myplugin")
+
+    Raises:
+        PluginLoadError: If the plugin cannot be loaded due to import errors or if the package path is invalid.
+
+    Examples:
+        >>> load_plugin("external_plugin.core")       # Loads plugin from external package
+
+    Notes:
+        - The plugin package should handle its own registration during import
+        - All submodules in the plugin package will be imported
+        - Implementation follows the plugin discovery pattern from Python packaging guidelines
+
+    See Also:
+        https://packaging.python.org/en/latest/guides/creating-and-discovering-plugins/
+    """
+    try:
+        package_module = importlib.import_module(plugin_path, __package__)
+    except (ImportError, ModuleNotFoundError) as e:
+        raise PluginLoadError(
+            f"Failed to load plugin '{plugin_path}'. Verify the path and installation: {str(e)}"
+        ) from e
+
+    def iter_namespace(ns_pkg):
+        return pkgutil.iter_modules(ns_pkg.__path__, ns_pkg.__name__ + ".")
+
+    try:
+        for _finder, pkg_name, _ispkg in iter_namespace(package_module):
+            importlib.import_module(pkg_name)
+    except ImportError as e:
+        raise PluginLoadError(
+            f"Failed to load plugin '{plugin_path}'. Verify the path and installation: {str(e)}"
+        ) from e
+
+
+def get_path_arg(field_name: str, args: Sequence[str] | None = None) -> str | None:
+    return parse_arg(f"{field_name}.{PATH_KEY}", args)
+
+
+def get_type_arg(field_name: str, args: Sequence[str] | None = None) -> str | None:
+    return parse_arg(f"{field_name}.{draccus.CHOICE_TYPE_KEY}", args)
+
+
+def filter_arg(field_to_filter: str, args: Sequence[str] | None = None) -> list[str]:
+    return [arg for arg in args if not arg.startswith(f"--{field_to_filter}=")]
+
+
+def filter_path_args(fields_to_filter: str | list[str], args: Sequence[str] | None = None) -> list[str]:
+    """
+    Filters command-line arguments related to fields with specific path arguments.
+
+    Args:
+        fields_to_filter (str | list[str]): A single str or a list of str whose arguments need to be filtered.
+        args (Sequence[str] | None): The sequence of command-line arguments to be filtered.
+            Defaults to None.
+
+    Returns:
+        list[str]: A filtered list of arguments, with arguments related to the specified
+        fields removed.
+
+    Raises:
+        ArgumentError: If both a path argument (e.g., `--field_name.path`) and a type
+            argument (e.g., `--field_name.type`) are specified for the same field.
+    """
+    if isinstance(fields_to_filter, str):
+        fields_to_filter = [fields_to_filter]
+
+    filtered_args = args
+    for field in fields_to_filter:
+        if get_path_arg(field, args):
+            if get_type_arg(field, args):
+                raise ArgumentError(
+                    argument=None,
+                    message=f"Cannot specify both --{field}.{PATH_KEY} and --{field}.{draccus.CHOICE_TYPE_KEY}",
+                )
+            filtered_args = [arg for arg in filtered_args if not arg.startswith(f"--{field}.")]
+
+    return filtered_args
+
+
+def wrap(config_path: Path | None = None):
+    """
+    HACK: Similar to draccus.wrap but does three additional things:
+        - Will remove '.path' arguments from CLI in order to process them later on.
+        - If a 'config_path' is passed and the main config class has a 'from_pretrained' method, will
+          initialize it from there to allow to fetch configs from the hub directly
+        - Will load plugins specified in the CLI arguments. These plugins will typically register
+            their own subclasses of config classes, so that draccus can find the right class to instantiate
+            from the CLI '.type' arguments
+    """
+
+    def wrapper_outer(fn):
+        @wraps(fn)
+        def wrapper_inner(*args, **kwargs):
+            argspec = inspect.getfullargspec(fn)
+            argtype = argspec.annotations[argspec.args[0]]
+            if len(args) > 0 and type(args[0]) is argtype:
+                cfg = args[0]
+                args = args[1:]
+            else:
+                cli_args = sys.argv[1:]
+                plugin_args = parse_plugin_args(PLUGIN_DISCOVERY_SUFFIX, cli_args)
+                for plugin_cli_arg, plugin_path in plugin_args.items():
+                    try:
+                        load_plugin(plugin_path)
+                    except PluginLoadError as e:
+                        # add the relevant CLI arg to the error message
+                        raise PluginLoadError(f"{e}\nFailed plugin CLI Arg: {plugin_cli_arg}") from e
+                    cli_args = filter_arg(plugin_cli_arg, cli_args)
+                config_path_cli = parse_arg("config_path", cli_args)
+                if has_method(argtype, "__get_path_fields__"):
+                    path_fields = argtype.__get_path_fields__()
+                    cli_args = filter_path_args(path_fields, cli_args)
+                if has_method(argtype, "from_pretrained") and config_path_cli:
+                    cli_args = filter_arg("config_path", cli_args)
+                    cfg = argtype.from_pretrained(config_path_cli, cli_args=cli_args)
+                else:
+                    cfg = draccus.parse(config_class=argtype, config_path=config_path, args=cli_args)
+            response = fn(cfg, *args, **kwargs)
+            return response
+
+        return wrapper_inner
+
+    return wrapper_outer

lerobot/configs/policies.py

@@ -0,0 +1,176 @@
+# 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 abc
+import logging
+import os
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Type, TypeVar
+
+import draccus
+from huggingface_hub import hf_hub_download
+from huggingface_hub.constants import CONFIG_NAME
+from huggingface_hub.errors import HfHubHTTPError
+
+from lerobot.common.optim.optimizers import OptimizerConfig
+from lerobot.common.optim.schedulers import LRSchedulerConfig
+from lerobot.common.utils.hub import HubMixin
+from lerobot.common.utils.utils import auto_select_torch_device, is_amp_available, is_torch_device_available
+from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
+
+# Generic variable that is either PreTrainedConfig or a subclass thereof
+T = TypeVar("T", bound="PreTrainedConfig")
+
+
+@dataclass
+class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
+    """
+    Base configuration class for policy models.
+
+    Args:
+        n_obs_steps: Number of environment steps worth of observations to pass to the policy (takes the
+            current step and additional steps going back).
+        input_shapes: A dictionary defining the shapes of the input data for the policy.
+        output_shapes: A dictionary defining the shapes of the output data for the policy.
+        input_normalization_modes: A dictionary with key representing the modality and the value specifies the
+            normalization mode to apply.
+        output_normalization_modes: Similar dictionary as `input_normalization_modes`, but to unnormalize to
+            the original scale.
+    """
+
+    n_obs_steps: int = 1
+    normalization_mapping: dict[str, NormalizationMode] = field(default_factory=dict)
+
+    input_features: dict[str, PolicyFeature] = field(default_factory=dict)
+    output_features: dict[str, PolicyFeature] = field(default_factory=dict)
+
+    device: str | None = None  # cuda | cpu | mp
+    # `use_amp` determines whether to use Automatic Mixed Precision (AMP) for training and evaluation. With AMP,
+    # automatic gradient scaling is used.
+    use_amp: bool = False
+
+    def __post_init__(self):
+        self.pretrained_path = None
+        if not self.device or not is_torch_device_available(self.device):
+            auto_device = auto_select_torch_device()
+            logging.warning(f"Device '{self.device}' is not available. Switching to '{auto_device}'.")
+            self.device = auto_device.type
+
+        # Automatically deactivate AMP if necessary
+        if self.use_amp and not is_amp_available(self.device):
+            logging.warning(
+                f"Automatic Mixed Precision (amp) is not available on device '{self.device}'. Deactivating AMP."
+            )
+            self.use_amp = False
+
+    @property
+    def type(self) -> str:
+        return self.get_choice_name(self.__class__)
+
+    @abc.abstractproperty
+    def observation_delta_indices(self) -> list | None:
+        raise NotImplementedError
+
+    @abc.abstractproperty
+    def action_delta_indices(self) -> list | None:
+        raise NotImplementedError
+
+    @abc.abstractproperty
+    def reward_delta_indices(self) -> list | None:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def get_optimizer_preset(self) -> OptimizerConfig:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def get_scheduler_preset(self) -> LRSchedulerConfig | None:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def validate_features(self) -> None:
+        raise NotImplementedError
+
+    @property
+    def robot_state_feature(self) -> PolicyFeature | None:
+        for _, ft in self.input_features.items():
+            if ft.type is FeatureType.STATE:
+                return ft
+        return None
+
+    @property
+    def env_state_feature(self) -> PolicyFeature | None:
+        for _, ft in self.input_features.items():
+            if ft.type is FeatureType.ENV:
+                return ft
+        return None
+
+    @property
+    def image_features(self) -> dict[str, PolicyFeature]:
+        return {key: ft for key, ft in self.input_features.items() if ft.type is FeatureType.VISUAL}
+
+    @property
+    def action_feature(self) -> PolicyFeature | None:
+        for _, ft in self.output_features.items():
+            if ft.type is FeatureType.ACTION:
+                return ft
+        return None
+
+    def _save_pretrained(self, save_directory: Path) -> None:
+        with open(save_directory / CONFIG_NAME, "w") as f, draccus.config_type("json"):
+            draccus.dump(self, f, indent=4)
+
+    @classmethod
+    def from_pretrained(
+        cls: Type[T],
+        pretrained_name_or_path: str | Path,
+        *,
+        force_download: bool = False,
+        resume_download: bool = None,
+        proxies: dict | None = None,
+        token: str | bool | None = None,
+        cache_dir: str | Path | None = None,
+        local_files_only: bool = False,
+        revision: str | None = None,
+        **policy_kwargs,
+    ) -> T:
+        model_id = str(pretrained_name_or_path)
+        config_file: str | None = None
+        if Path(model_id).is_dir():
+            if CONFIG_NAME in os.listdir(model_id):
+                config_file = os.path.join(model_id, CONFIG_NAME)
+            else:
+                print(f"{CONFIG_NAME} not found in {Path(model_id).resolve()}")
+        else:
+            try:
+                config_file = hf_hub_download(
+                    repo_id=model_id,
+                    filename=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,
+                )
+            except HfHubHTTPError as e:
+                raise FileNotFoundError(
+                    f"{CONFIG_NAME} not found on the HuggingFace Hub in {model_id}"
+                ) from e
+
+        # HACK: this is very ugly, ideally we'd like to be able to do that natively with draccus
+        # something like --policy.path (in addition to --policy.type)
+        cli_overrides = policy_kwargs.pop("cli_overrides", [])
+        return draccus.parse(cls, config_file, args=cli_overrides)

lerobot/configs/train.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.
+import datetime as dt
+import os
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Type
+
+import draccus
+from huggingface_hub import hf_hub_download
+from huggingface_hub.errors import HfHubHTTPError
+
+from lerobot.common import envs
+from lerobot.common.optim import OptimizerConfig
+from lerobot.common.optim.schedulers import LRSchedulerConfig
+from lerobot.common.utils.hub import HubMixin
+from lerobot.configs import parser
+from lerobot.configs.default import DatasetConfig, EvalConfig, WandBConfig
+from lerobot.configs.policies import PreTrainedConfig
+
+TRAIN_CONFIG_NAME = "train_config.json"
+
+
+@dataclass
+class TrainPipelineConfig(HubMixin):
+    dataset: DatasetConfig
+    env: envs.EnvConfig | None = None
+    policy: PreTrainedConfig | None = None
+    # Set `dir` to where you would like to save all of the run outputs. If you run another training session
+    # 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.
+    resume: bool = False
+    # `seed` is used for training (eg: model initialization, dataset shuffling)
+    # AND for the evaluation environments.
+    seed: int | None = 1000
+    # Number of workers for the dataloader.
+    num_workers: int = 4
+    batch_size: int = 8
+    steps: int = 100_000
+    eval_freq: int = 20_000
+    log_freq: int = 200
+    save_checkpoint: bool = True
+    # Checkpoint is saved every `save_freq` training iterations and after the last training step.
+    save_freq: int = 20_000
+    use_policy_training_preset: bool = True
+    optimizer: OptimizerConfig | None = None
+    scheduler: LRSchedulerConfig | None = None
+    eval: EvalConfig = field(default_factory=EvalConfig)
+    wandb: WandBConfig = field(default_factory=WandBConfig)
+
+    def __post_init__(self):
+        self.checkpoint_path = None
+
+    def validate(self):
+        # HACK: We parse again the cli args here to get the pretrained paths if there was some.
+        policy_path = parser.get_path_arg("policy")
+        if policy_path:
+            # Only load the policy config
+            cli_overrides = parser.get_cli_overrides("policy")
+            self.policy = PreTrainedConfig.from_pretrained(policy_path, cli_overrides=cli_overrides)
+            self.policy.pretrained_path = policy_path
+        elif self.resume:
+            # The entire train config is already loaded, we just need to get the checkpoint dir
+            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 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."
+                )
+            policy_path = Path(config_path).parent
+            self.policy.pretrained_path = policy_path
+            self.checkpoint_path = policy_path.parent
+
+        if not self.job_name:
+            if self.env is None:
+                self.job_name = f"{self.policy.type}"
+            else:
+                self.job_name = f"{self.env.type}_{self.policy.type}"
+
+        if not self.resume and isinstance(self.output_dir, Path) and self.output_dir.is_dir():
+            raise FileExistsError(
+                f"Output directory {self.output_dir} already exists and resume is {self.resume}. "
+                f"Please change your output directory so that {self.output_dir} is not overwritten."
+            )
+        elif not self.output_dir:
+            now = dt.datetime.now()
+            train_dir = f"{now:%Y-%m-%d}/{now:%H-%M-%S}_{self.job_name}"
+            self.output_dir = Path("outputs/train") / train_dir
+
+        if isinstance(self.dataset.repo_id, list):
+            raise NotImplementedError("LeRobotMultiDataset is not currently implemented.")
+
+        if not self.use_policy_training_preset and (self.optimizer is None or self.scheduler is None):
+            raise ValueError("Optimizer and Scheduler must be set when the policy presets are not used.")
+        elif self.use_policy_training_preset and not self.resume:
+            self.optimizer = self.policy.get_optimizer_preset()
+            self.scheduler = self.policy.get_scheduler_preset()
+
+    @classmethod
+    def __get_path_fields__(cls) -> list[str]:
+        """This enables the parser to load config from the policy using `--policy.path=local/dir`"""
+        return ["policy"]
+
+    def to_dict(self) -> dict:
+        return draccus.encode(self)
+
+    def _save_pretrained(self, save_directory: Path) -> None:
+        with open(save_directory / TRAIN_CONFIG_NAME, "w") as f, draccus.config_type("json"):
+            draccus.dump(self, f, indent=4)
+
+    @classmethod
+    def from_pretrained(
+        cls: Type["TrainPipelineConfig"],
+        pretrained_name_or_path: str | Path,
+        *,
+        force_download: bool = False,
+        resume_download: bool = None,
+        proxies: dict | None = None,
+        token: str | bool | None = None,
+        cache_dir: str | Path | None = None,
+        local_files_only: bool = False,
+        revision: str | None = None,
+        **kwargs,
+    ) -> "TrainPipelineConfig":
+        model_id = str(pretrained_name_or_path)
+        config_file: str | None = None
+        if Path(model_id).is_dir():
+            if TRAIN_CONFIG_NAME in os.listdir(model_id):
+                config_file = os.path.join(model_id, TRAIN_CONFIG_NAME)
+            else:
+                print(f"{TRAIN_CONFIG_NAME} not found in {Path(model_id).resolve()}")
+        elif Path(model_id).is_file():
+            config_file = model_id
+        else:
+            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,
+                )
+            except HfHubHTTPError as e:
+                raise FileNotFoundError(
+                    f"{TRAIN_CONFIG_NAME} not found on the HuggingFace Hub in {model_id}"
+                ) from e
+
+        cli_args = kwargs.pop("cli_args", [])
+        cfg = draccus.parse(cls, config_file, args=cli_args)
+
+        return cfg

lerobot/configs/types.py

@@ -0,0 +1,41 @@
+# 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.
+# Note: We subclass str so that serialization is straightforward
+# https://stackoverflow.com/questions/24481852/serialising-an-enum-member-to-json
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any, Protocol
+
+
+class FeatureType(str, Enum):
+    STATE = "STATE"
+    VISUAL = "VISUAL"
+    ENV = "ENV"
+    ACTION = "ACTION"
+
+
+class NormalizationMode(str, Enum):
+    MIN_MAX = "MIN_MAX"
+    MEAN_STD = "MEAN_STD"
+    IDENTITY = "IDENTITY"
+
+
+class DictLike(Protocol):
+    def __getitem__(self, key: Any) -> Any: ...
+
+
+@dataclass
+class PolicyFeature:
+    type: FeatureType
+    shape: tuple

lerobot/scripts/configure_motor.py

@@ -0,0 +1,176 @@
+# 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.
+"""
+This script configure a single motor at a time to a given ID and baudrate.
+
+Example of usage:
+```bash
+python lerobot/scripts/configure_motor.py \
+  --port /dev/tty.usbmodem585A0080521 \
+  --brand feetech \
+  --model sts3215 \
+  --baudrate 1000000 \
+  --ID 1
+```
+"""
+
+import argparse
+import time
+
+
+def get_motor_bus_cls(brand: str) -> tuple:
+    if brand == "feetech":
+        from lerobot.common.robot_devices.motors.configs import FeetechMotorsBusConfig
+        from lerobot.common.robot_devices.motors.feetech import (
+            MODEL_BAUDRATE_TABLE,
+            SCS_SERIES_BAUDRATE_TABLE,
+            FeetechMotorsBus,
+        )
+
+        return FeetechMotorsBusConfig, FeetechMotorsBus, MODEL_BAUDRATE_TABLE, SCS_SERIES_BAUDRATE_TABLE
+
+    elif brand == "dynamixel":
+        from lerobot.common.robot_devices.motors.configs import DynamixelMotorsBusConfig
+        from lerobot.common.robot_devices.motors.dynamixel import (
+            MODEL_BAUDRATE_TABLE,
+            X_SERIES_BAUDRATE_TABLE,
+            DynamixelMotorsBus,
+        )
+
+        return DynamixelMotorsBusConfig, DynamixelMotorsBus, MODEL_BAUDRATE_TABLE, X_SERIES_BAUDRATE_TABLE
+
+    else:
+        raise ValueError(
+            f"Currently we do not support this motor brand: {brand}. We currently support feetech and dynamixel motors."
+        )
+
+
+def configure_motor(port, brand, model, motor_idx_des, baudrate_des):
+    motor_bus_config_cls, motor_bus_cls, model_baudrate_table, series_baudrate_table = get_motor_bus_cls(
+        brand
+    )
+
+    # Check if the provided model exists in the model_baud_rate_table
+    if model not in model_baudrate_table:
+        raise ValueError(
+            f"Invalid model '{model}' for brand '{brand}'. Supported models: {list(model_baudrate_table.keys())}"
+        )
+
+    # Setup motor names, indices, and models
+    motor_name = "motor"
+    motor_index_arbitrary = motor_idx_des  # Use the motor ID passed via argument
+    motor_model = model  # Use the motor model passed via argument
+
+    config = motor_bus_config_cls(port=port, motors={motor_name: (motor_index_arbitrary, motor_model)})
+
+    # Initialize the MotorBus with the correct port and motor configurations
+    motor_bus = motor_bus_cls(config=config)
+
+    # Try to connect to the motor bus and handle any connection-specific errors
+    try:
+        motor_bus.connect()
+        print(f"Connected on port {motor_bus.port}")
+    except OSError as e:
+        print(f"Error occurred when connecting to the motor bus: {e}")
+        return
+
+    # Motor bus is connected, proceed with the rest of the operations
+    try:
+        print("Scanning all baudrates and motor indices")
+        all_baudrates = set(series_baudrate_table.values())
+        motor_index = -1  # Set the motor index to an out-of-range value.
+
+        for baudrate in all_baudrates:
+            motor_bus.set_bus_baudrate(baudrate)
+            present_ids = motor_bus.find_motor_indices(list(range(1, 10)))
+            if len(present_ids) > 1:
+                raise ValueError(
+                    "Error: More than one motor ID detected. This script is designed to only handle one motor at a time. Please disconnect all but one motor."
+                )
+
+            if len(present_ids) == 1:
+                if motor_index != -1:
+                    raise ValueError(
+                        "Error: More than one motor ID detected. This script is designed to only handle one motor at a time. Please disconnect all but one motor."
+                    )
+                motor_index = present_ids[0]
+                break
+
+        if motor_index == -1:
+            raise ValueError("No motors detected. Please ensure you have one motor connected.")
+
+        print(f"Motor index found at: {motor_index}")
+
+        if brand == "feetech":
+            # Allows ID and BAUDRATE to be written in memory
+            motor_bus.write_with_motor_ids(motor_bus.motor_models, motor_index, "Lock", 0)
+
+        if baudrate != baudrate_des:
+            print(f"Setting its baudrate to {baudrate_des}")
+            baudrate_idx = list(series_baudrate_table.values()).index(baudrate_des)
+
+            # The write can fail, so we allow retries
+            motor_bus.write_with_motor_ids(motor_bus.motor_models, motor_index, "Baud_Rate", baudrate_idx)
+            time.sleep(0.5)
+            motor_bus.set_bus_baudrate(baudrate_des)
+            present_baudrate_idx = motor_bus.read_with_motor_ids(
+                motor_bus.motor_models, motor_index, "Baud_Rate", num_retry=2
+            )
+
+            if present_baudrate_idx != baudrate_idx:
+                raise OSError("Failed to write baudrate.")
+
+        print(f"Setting its index to desired index {motor_idx_des}")
+        if brand == "feetech":
+            motor_bus.write_with_motor_ids(motor_bus.motor_models, motor_index, "Lock", 0)
+        motor_bus.write_with_motor_ids(motor_bus.motor_models, motor_index, "ID", motor_idx_des)
+
+        present_idx = motor_bus.read_with_motor_ids(motor_bus.motor_models, motor_idx_des, "ID", num_retry=2)
+        if present_idx != motor_idx_des:
+            raise OSError("Failed to write index.")
+
+        if brand == "feetech":
+            # Set Maximum_Acceleration to 254 to speedup acceleration and deceleration of
+            # the motors. Note: this configuration is not in the official STS3215 Memory Table
+            motor_bus.write("Lock", 0)
+            motor_bus.write("Maximum_Acceleration", 254)
+
+            motor_bus.write("Goal_Position", 2048)
+            time.sleep(4)
+            print("Present Position", motor_bus.read("Present_Position"))
+
+            motor_bus.write("Offset", 0)
+            time.sleep(4)
+            print("Offset", motor_bus.read("Offset"))
+
+    except Exception as e:
+        print(f"Error occurred during motor configuration: {e}")
+
+    finally:
+        motor_bus.disconnect()
+        print("Disconnected from motor bus.")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--port", type=str, required=True, help="Motors bus port (e.g. dynamixel,feetech)")
+    parser.add_argument("--brand", type=str, required=True, help="Motor brand (e.g. dynamixel,feetech)")
+    parser.add_argument("--model", type=str, required=True, help="Motor model (e.g. xl330-m077,sts3215)")
+    parser.add_argument("--ID", type=int, required=True, help="Desired ID of the current motor (e.g. 1,2,3)")
+    parser.add_argument(
+        "--baudrate", type=int, default=1000000, help="Desired baudrate for the motor (default: 1000000)"
+    )
+    args = parser.parse_args()
+
+    configure_motor(args.port, args.brand, args.model, args.ID, args.baudrate)

lerobot/scripts/control_robot.py

@@ -0,0 +1,393 @@
+# 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.
+"""
+Utilities to control a robot.
+
+Useful to record a dataset, replay a recorded episode, run the policy on your robot
+and record an evaluation dataset, and to recalibrate your robot if needed.
+
+Examples of usage:
+
+- Recalibrate your robot:
+```bash
+python lerobot/scripts/control_robot.py \
+    --robot.type=so100 \
+    --control.type=calibrate
+```
+
+- Unlimited teleoperation at highest frequency (~200 Hz is expected), to exit with CTRL+C:
+```bash
+python lerobot/scripts/control_robot.py \
+    --robot.type=so100 \
+    --robot.cameras='{}' \
+    --control.type=teleoperate
+
+# Add the cameras from the robot definition to visualize them:
+python lerobot/scripts/control_robot.py \
+    --robot.type=so100 \
+    --control.type=teleoperate
+```
+
+- Unlimited teleoperation at a limited frequency of 30 Hz, to simulate data recording frequency:
+```bash
+python lerobot/scripts/control_robot.py \
+    --robot.type=so100 \
+    --control.type=teleoperate \
+    --control.fps=30
+```
+
+- Record one episode in order to test replay:
+```bash
+python lerobot/scripts/control_robot.py \
+    --robot.type=so100 \
+    --control.type=record \
+    --control.fps=30 \
+    --control.single_task="Grasp a lego block and put it in the bin." \
+    --control.repo_id=$USER/koch_test \
+    --control.num_episodes=1 \
+    --control.push_to_hub=True
+```
+
+- Visualize dataset:
+```bash
+python lerobot/scripts/visualize_dataset.py \
+    --repo-id $USER/koch_test \
+    --episode-index 0
+```
+
+- Replay this test episode:
+```bash
+python lerobot/scripts/control_robot.py replay \
+    --robot.type=so100 \
+    --control.type=replay \
+    --control.fps=30 \
+    --control.repo_id=$USER/koch_test \
+    --control.episode=0
+```
+
+- Record a full dataset in order to train a policy, with 2 seconds of warmup,
+30 seconds of recording for each episode, and 10 seconds to reset the environment in between episodes:
+```bash
+python lerobot/scripts/control_robot.py record \
+    --robot.type=so100 \
+    --control.type=record \
+    --control.fps 30 \
+    --control.repo_id=$USER/koch_pick_place_lego \
+    --control.num_episodes=50 \
+    --control.warmup_time_s=2 \
+    --control.episode_time_s=30 \
+    --control.reset_time_s=10
+```
+
+- For remote controlled robots like LeKiwi, run this script on the robot edge device (e.g. RaspBerryPi):
+```bash
+python lerobot/scripts/control_robot.py \
+  --robot.type=lekiwi \
+  --control.type=remote_robot
+```
+
+**NOTE**: You can use your keyboard to control data recording flow.
+- Tap right arrow key '->' to early exit while recording an episode and go to resseting the environment.
+- Tap right arrow key '->' to early exit while resetting the environment and got to recording the next episode.
+- Tap left arrow key '<-' to early exit and re-record the current episode.
+- Tap escape key 'esc' to stop the data recording.
+This might require a sudo permission to allow your terminal to monitor keyboard events.
+
+**NOTE**: You can resume/continue data recording by running the same data recording command and adding `--control.resume=true`.
+
+- Train on this dataset with the ACT policy:
+```bash
+python lerobot/scripts/train.py \
+  --dataset.repo_id=${HF_USER}/koch_pick_place_lego \
+  --policy.type=act \
+  --output_dir=outputs/train/act_koch_pick_place_lego \
+  --job_name=act_koch_pick_place_lego \
+  --device=cuda \
+  --wandb.enable=true
+```
+
+- Run the pretrained policy on the robot:
+```bash
+python lerobot/scripts/control_robot.py \
+    --robot.type=so100 \
+    --control.type=record \
+    --control.fps=30 \
+    --control.single_task="Grasp a lego block and put it in the bin." \
+    --control.repo_id=$USER/eval_act_koch_pick_place_lego \
+    --control.num_episodes=10 \
+    --control.warmup_time_s=2 \
+    --control.episode_time_s=30 \
+    --control.reset_time_s=10 \
+    --control.push_to_hub=true \
+    --control.policy.path=outputs/train/act_koch_pick_place_lego/checkpoints/080000/pretrained_model
+```
+"""
+
+import logging
+import time
+from dataclasses import asdict
+from pprint import pformat
+
+# from safetensors.torch import load_file, save_file
+from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.common.policies.factory import make_policy
+from lerobot.common.robot_devices.control_configs import (
+    CalibrateControlConfig,
+    ControlPipelineConfig,
+    RecordControlConfig,
+    RemoteRobotConfig,
+    ReplayControlConfig,
+    TeleoperateControlConfig,
+)
+from lerobot.common.robot_devices.control_utils import (
+    control_loop,
+    init_keyboard_listener,
+    log_control_info,
+    record_episode,
+    reset_environment,
+    sanity_check_dataset_name,
+    sanity_check_dataset_robot_compatibility,
+    stop_recording,
+    warmup_record,
+)
+from lerobot.common.robot_devices.robots.utils import Robot, make_robot_from_config
+from lerobot.common.robot_devices.utils import busy_wait, safe_disconnect
+from lerobot.common.utils.utils import has_method, init_logging, log_say
+from lerobot.configs import parser
+
+########################################################################################
+# Control modes
+########################################################################################
+
+
+@safe_disconnect
+def calibrate(robot: Robot, cfg: CalibrateControlConfig):
+    # TODO(aliberts): move this code in robots' classes
+    if robot.robot_type.startswith("stretch"):
+        if not robot.is_connected:
+            robot.connect()
+        if not robot.is_homed():
+            robot.home()
+        return
+
+    arms = robot.available_arms if cfg.arms is None else cfg.arms
+    unknown_arms = [arm_id for arm_id in arms if arm_id not in robot.available_arms]
+    available_arms_str = " ".join(robot.available_arms)
+    unknown_arms_str = " ".join(unknown_arms)
+
+    if arms is None or len(arms) == 0:
+        raise ValueError(
+            "No arm provided. Use `--arms` as argument with one or more available arms.\n"
+            f"For instance, to recalibrate all arms add: `--arms {available_arms_str}`"
+        )
+
+    if len(unknown_arms) > 0:
+        raise ValueError(
+            f"Unknown arms provided ('{unknown_arms_str}'). Available arms are `{available_arms_str}`."
+        )
+
+    for arm_id in arms:
+        arm_calib_path = robot.calibration_dir / f"{arm_id}.json"
+        if arm_calib_path.exists():
+            print(f"Removing '{arm_calib_path}'")
+            arm_calib_path.unlink()
+        else:
+            print(f"Calibration file not found '{arm_calib_path}'")
+
+    if robot.is_connected:
+        robot.disconnect()
+
+    if robot.robot_type.startswith("lekiwi") and "main_follower" in arms:
+        print("Calibrating only the lekiwi follower arm 'main_follower'...")
+        robot.calibrate_follower()
+        return
+
+    if robot.robot_type.startswith("lekiwi") and "main_leader" in arms:
+        print("Calibrating only the lekiwi leader arm 'main_leader'...")
+        robot.calibrate_leader()
+        return
+
+    # Calling `connect` automatically runs calibration
+    # when the calibration file is missing
+    robot.connect()
+    robot.disconnect()
+    print("Calibration is done! You can now teleoperate and record datasets!")
+
+
+@safe_disconnect
+def teleoperate(robot: Robot, cfg: TeleoperateControlConfig):
+    control_loop(
+        robot,
+        control_time_s=cfg.teleop_time_s,
+        fps=cfg.fps,
+        teleoperate=True,
+        display_cameras=cfg.display_cameras,
+    )
+
+
+@safe_disconnect
+def record(
+    robot: Robot,
+    cfg: RecordControlConfig,
+) -> LeRobotDataset:
+    # TODO(rcadene): Add option to record logs
+    if cfg.resume:
+        dataset = LeRobotDataset(
+            cfg.repo_id,
+            root=cfg.root,
+        )
+        if len(robot.cameras) > 0:
+            dataset.start_image_writer(
+                num_processes=cfg.num_image_writer_processes,
+                num_threads=cfg.num_image_writer_threads_per_camera * len(robot.cameras),
+            )
+        sanity_check_dataset_robot_compatibility(dataset, robot, cfg.fps, cfg.video)
+    else:
+        # Create empty dataset or load existing saved episodes
+        sanity_check_dataset_name(cfg.repo_id, cfg.policy)
+        dataset = LeRobotDataset.create(
+            cfg.repo_id,
+            cfg.fps,
+            root=cfg.root,
+            robot=robot,
+            use_videos=cfg.video,
+            image_writer_processes=cfg.num_image_writer_processes,
+            image_writer_threads=cfg.num_image_writer_threads_per_camera * len(robot.cameras),
+        )
+
+    # Load pretrained policy
+    policy = None if cfg.policy is None else make_policy(cfg.policy, ds_meta=dataset.meta)
+
+    if not robot.is_connected:
+        robot.connect()
+
+    listener, events = init_keyboard_listener()
+
+    # Execute a few seconds without recording to:
+    # 1. teleoperate the robot to move it in starting position if no policy provided,
+    # 2. give times to the robot devices to connect and start synchronizing,
+    # 3. place the cameras windows on screen
+    enable_teleoperation = policy is None
+    log_say("Warmup record", cfg.play_sounds)
+    warmup_record(robot, events, enable_teleoperation, cfg.warmup_time_s, cfg.display_cameras, cfg.fps)
+
+    if has_method(robot, "teleop_safety_stop"):
+        robot.teleop_safety_stop()
+
+    recorded_episodes = 0
+    while True:
+        if recorded_episodes >= cfg.num_episodes:
+            break
+
+        log_say(f"Recording episode {dataset.num_episodes}", cfg.play_sounds)
+        record_episode(
+            robot=robot,
+            dataset=dataset,
+            events=events,
+            episode_time_s=cfg.episode_time_s,
+            display_cameras=cfg.display_cameras,
+            policy=policy,
+            fps=cfg.fps,
+            single_task=cfg.single_task,
+        )
+
+        # Execute a few seconds without recording to give time to manually reset the environment
+        # Current code logic doesn't allow to teleoperate during this time.
+        # TODO(rcadene): add an option to enable teleoperation during reset
+        # Skip reset for the last episode to be recorded
+        if not events["stop_recording"] and (
+            (recorded_episodes < cfg.num_episodes - 1) or events["rerecord_episode"]
+        ):
+            log_say("Reset the environment", cfg.play_sounds)
+            reset_environment(robot, events, cfg.reset_time_s, cfg.fps)
+
+        if events["rerecord_episode"]:
+            log_say("Re-record episode", cfg.play_sounds)
+            events["rerecord_episode"] = False
+            events["exit_early"] = False
+            dataset.clear_episode_buffer()
+            continue
+
+        dataset.save_episode()
+        recorded_episodes += 1
+
+        if events["stop_recording"]:
+            break
+
+    log_say("Stop recording", cfg.play_sounds, blocking=True)
+    stop_recording(robot, listener, cfg.display_cameras)
+
+    if cfg.push_to_hub:
+        dataset.push_to_hub(tags=cfg.tags, private=cfg.private)
+
+    log_say("Exiting", cfg.play_sounds)
+    return dataset
+
+
+@safe_disconnect
+def replay(
+    robot: Robot,
+    cfg: ReplayControlConfig,
+):
+    # TODO(rcadene, aliberts): refactor with control_loop, once `dataset` is an instance of LeRobotDataset
+    # TODO(rcadene): Add option to record logs
+
+    dataset = LeRobotDataset(cfg.repo_id, root=cfg.root, episodes=[cfg.episode])
+    actions = dataset.hf_dataset.select_columns("action")
+
+    if not robot.is_connected:
+        robot.connect()
+
+    log_say("Replaying episode", cfg.play_sounds, blocking=True)
+    for idx in range(dataset.num_frames):
+        start_episode_t = time.perf_counter()
+
+        action = actions[idx]["action"]
+        robot.send_action(action)
+
+        dt_s = time.perf_counter() - start_episode_t
+        busy_wait(1 / cfg.fps - dt_s)
+
+        dt_s = time.perf_counter() - start_episode_t
+        log_control_info(robot, dt_s, fps=cfg.fps)
+
+
+@parser.wrap()
+def control_robot(cfg: ControlPipelineConfig):
+    init_logging()
+    logging.info(pformat(asdict(cfg)))
+
+    robot = make_robot_from_config(cfg.robot)
+
+    if isinstance(cfg.control, CalibrateControlConfig):
+        calibrate(robot, cfg.control)
+    elif isinstance(cfg.control, TeleoperateControlConfig):
+        teleoperate(robot, cfg.control)
+    elif isinstance(cfg.control, RecordControlConfig):
+        record(robot, cfg.control)
+    elif isinstance(cfg.control, ReplayControlConfig):
+        replay(robot, cfg.control)
+    elif isinstance(cfg.control, RemoteRobotConfig):
+        from lerobot.common.robot_devices.robots.lekiwi_remote import run_lekiwi
+
+        run_lekiwi(cfg.robot)
+
+    if robot.is_connected:
+        # Disconnect manually to avoid a "Core dump" during process
+        # termination due to camera threads not properly exiting.
+        robot.disconnect()
+
+
+if __name__ == "__main__":
+    control_robot()

lerobot/scripts/control_sim_robot.py

@@ -0,0 +1,561 @@
+# 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.
+"""
+Utilities to control a robot in simulation.
+
+Useful to record a dataset, replay a recorded episode and record an evaluation dataset.
+
+Examples of usage:
+
+
+- Unlimited teleoperation at a limited frequency of 30 Hz, to simulate data recording frequency.
+  You can modify this value depending on how fast your simulation can run:
+```bash
+python lerobot/scripts/control_robot.py teleoperate \
+    --fps 30 \
+    --robot-path lerobot/configs/robot/your_robot_config.yaml \
+    --sim-config lerobot/configs/env/your_sim_config.yaml
+```
+
+- Record one episode in order to test replay:
+```bash
+python lerobot/scripts/control_sim_robot.py record \
+    --robot-path lerobot/configs/robot/your_robot_config.yaml \
+    --sim-config lerobot/configs/env/your_sim_config.yaml \
+    --fps 30 \
+    --repo-id $USER/robot_sim_test \
+    --num-episodes 1 \
+    --run-compute-stats 0
+```
+
+Enable the --push-to-hub 1 to push the recorded dataset to the huggingface hub.
+
+- Visualize dataset:
+```bash
+python lerobot/scripts/visualize_dataset.py \
+    --repo-id $USER/robot_sim_test \
+    --episode-index 0
+```
+
+- Replay a sequence of test episodes:
+```bash
+python lerobot/scripts/control_sim_robot.py replay \
+    --robot-path lerobot/configs/robot/your_robot_config.yaml \
+    --sim-config lerobot/configs/env/your_sim_config.yaml \
+    --fps 30 \
+    --repo-id $USER/robot_sim_test \
+    --episode 0
+```
+Note: The seed is saved, therefore, during replay we can load the same environment state as the one during collection.
+
+- Record a full dataset in order to train a policy,
+30 seconds of recording for each episode, and 10 seconds to reset the environment in between episodes:
+```bash
+python lerobot/scripts/control_sim_robot.py record \
+    --robot-path lerobot/configs/robot/your_robot_config.yaml \
+    --sim-config lerobot/configs/env/your_sim_config.yaml \
+    --fps 30 \
+    --repo-id $USER/robot_sim_test \
+    --num-episodes 50 \
+    --episode-time-s 30 \
+```
+
+**NOTE**: You can use your keyboard to control data recording flow.
+- Tap right arrow key '->' to early exit while recording an episode and go to resetting the environment.
+- Tap right arrow key '->' to early exit while resetting the environment and got to recording the next episode.
+- Tap left arrow key '<-' to early exit and re-record the current episode.
+- Tap escape key 'esc' to stop the data recording.
+This might require a sudo permission to allow your terminal to monitor keyboard events.
+
+**NOTE**: You can resume/continue data recording by running the same data recording command twice.
+"""
+
+import argparse
+import importlib
+import logging
+import time
+from pathlib import Path
+
+import cv2
+import gymnasium as gym
+import numpy as np
+import torch
+
+from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.common.robot_devices.control_utils import (
+    init_keyboard_listener,
+    init_policy,
+    is_headless,
+    log_control_info,
+    predict_action,
+    sanity_check_dataset_name,
+    sanity_check_dataset_robot_compatibility,
+    stop_recording,
+)
+from lerobot.common.robot_devices.robots.utils import Robot, make_robot
+from lerobot.common.robot_devices.utils import busy_wait
+from lerobot.common.utils.utils import init_hydra_config, init_logging, log_say
+
+raise NotImplementedError("This script is currently deactivated")
+
+DEFAULT_FEATURES = {
+    "next.reward": {
+        "dtype": "float32",
+        "shape": (1,),
+        "names": None,
+    },
+    "next.success": {
+        "dtype": "bool",
+        "shape": (1,),
+        "names": None,
+    },
+    "seed": {
+        "dtype": "int64",
+        "shape": (1,),
+        "names": None,
+    },
+    "timestamp": {
+        "dtype": "float32",
+        "shape": (1,),
+        "names": None,
+    },
+}
+
+
+########################################################################################
+# Utilities
+########################################################################################
+def none_or_int(value):
+    if value == "None":
+        return None
+    return int(value)
+
+
+def init_sim_calibration(robot, cfg):
+    # Constants necessary for transforming the joint pos of the real robot to the sim
+    # depending on the robot description used in that sim.
+    start_pos = np.array(robot.leader_arms.main.calibration["start_pos"])
+    axis_directions = np.array(cfg.get("axis_directions", [1]))
+    offsets = np.array(cfg.get("offsets", [0])) * np.pi
+
+    return {"start_pos": start_pos, "axis_directions": axis_directions, "offsets": offsets}
+
+
+def real_positions_to_sim(real_positions, axis_directions, start_pos, offsets):
+    """Counts - starting position -> radians -> align axes -> offset"""
+    return axis_directions * (real_positions - start_pos) * 2.0 * np.pi / 4096 + offsets
+
+
+########################################################################################
+# Control modes
+########################################################################################
+
+
+def teleoperate(env, robot: Robot, process_action_fn, teleop_time_s=None):
+    env = env()
+    env.reset()
+    start_teleop_t = time.perf_counter()
+    while True:
+        leader_pos = robot.leader_arms.main.read("Present_Position")
+        action = process_action_fn(leader_pos)
+        env.step(np.expand_dims(action, 0))
+        if teleop_time_s is not None and time.perf_counter() - start_teleop_t > teleop_time_s:
+            print("Teleoperation processes finished.")
+            break
+
+
+def record(
+    env,
+    robot: Robot,
+    process_action_from_leader,
+    root: Path,
+    repo_id: str,
+    task: str,
+    fps: int | None = None,
+    tags: list[str] | None = None,
+    pretrained_policy_name_or_path: str = None,
+    policy_overrides: bool | None = None,
+    episode_time_s: int = 30,
+    num_episodes: int = 50,
+    video: bool = True,
+    push_to_hub: bool = True,
+    num_image_writer_processes: int = 0,
+    num_image_writer_threads_per_camera: int = 4,
+    display_cameras: bool = False,
+    play_sounds: bool = True,
+    resume: bool = False,
+    local_files_only: bool = False,
+    run_compute_stats: bool = True,
+) -> LeRobotDataset:
+    # Load pretrained policy
+    policy = None
+    if pretrained_policy_name_or_path is not None:
+        policy, policy_fps, device, use_amp = init_policy(pretrained_policy_name_or_path, policy_overrides)
+
+        if fps is None:
+            fps = policy_fps
+            logging.warning(f"No fps provided, so using the fps from policy config ({policy_fps}).")
+
+    if policy is None and process_action_from_leader is None:
+        raise ValueError("Either policy or process_action_fn has to be set to enable control in sim.")
+
+    # initialize listener before sim env
+    listener, events = init_keyboard_listener()
+
+    # create sim env
+    env = env()
+
+    # Create empty dataset or load existing saved episodes
+    num_cameras = sum([1 if "image" in key else 0 for key in env.observation_space])
+
+    # get image keys
+    image_keys = [key for key in env.observation_space if "image" in key]
+    state_keys_dict = env_cfg.state_keys
+
+    if resume:
+        dataset = LeRobotDataset(
+            repo_id,
+            root=root,
+            local_files_only=local_files_only,
+        )
+        dataset.start_image_writer(
+            num_processes=num_image_writer_processes,
+            num_threads=num_image_writer_threads_per_camera * num_cameras,
+        )
+        sanity_check_dataset_robot_compatibility(dataset, robot, fps, video)
+    else:
+        features = DEFAULT_FEATURES
+        # add image keys to features
+        for key in image_keys:
+            shape = env.observation_space[key].shape
+            if not key.startswith("observation.image."):
+                key = "observation.image." + key
+            features[key] = {"dtype": "video", "names": ["channels", "height", "width"], "shape": shape}
+
+        for key, obs_key in state_keys_dict.items():
+            features[key] = {
+                "dtype": "float32",
+                "names": None,
+                "shape": env.observation_space[obs_key].shape,
+            }
+
+        features["action"] = {"dtype": "float32", "shape": env.action_space.shape, "names": None}
+
+        # Create empty dataset or load existing saved episodes
+        sanity_check_dataset_name(repo_id, policy)
+        dataset = LeRobotDataset.create(
+            repo_id,
+            fps,
+            root=root,
+            features=features,
+            use_videos=video,
+            image_writer_processes=num_image_writer_processes,
+            image_writer_threads=num_image_writer_threads_per_camera * num_cameras,
+        )
+
+    recorded_episodes = 0
+    while True:
+        log_say(f"Recording episode {dataset.num_episodes}", play_sounds)
+
+        if events is None:
+            events = {"exit_early": False}
+
+        if episode_time_s is None:
+            episode_time_s = float("inf")
+
+        timestamp = 0
+        start_episode_t = time.perf_counter()
+
+        seed = np.random.randint(0, 1e5)
+        observation, info = env.reset(seed=seed)
+
+        while timestamp < episode_time_s:
+            start_loop_t = time.perf_counter()
+
+            if policy is not None:
+                action = predict_action(observation, policy, device, use_amp)
+            else:
+                leader_pos = robot.leader_arms.main.read("Present_Position")
+                action = process_action_from_leader(leader_pos)
+
+            observation, reward, terminated, _, info = env.step(action)
+
+            success = info.get("is_success", False)
+            env_timestamp = info.get("timestamp", dataset.episode_buffer["size"] / fps)
+
+            frame = {
+                "action": torch.from_numpy(action),
+                "next.reward": reward,
+                "next.success": success,
+                "seed": seed,
+                "timestamp": env_timestamp,
+            }
+
+            for key in image_keys:
+                if not key.startswith("observation.image"):
+                    frame["observation.image." + key] = observation[key]
+                else:
+                    frame[key] = observation[key]
+
+            for key, obs_key in state_keys_dict.items():
+                frame[key] = torch.from_numpy(observation[obs_key])
+
+            dataset.add_frame(frame)
+
+            if display_cameras and not is_headless():
+                for key in image_keys:
+                    cv2.imshow(key, cv2.cvtColor(observation[key], cv2.COLOR_RGB2BGR))
+                cv2.waitKey(1)
+
+            if fps is not None:
+                dt_s = time.perf_counter() - start_loop_t
+                busy_wait(1 / fps - dt_s)
+
+            dt_s = time.perf_counter() - start_loop_t
+            log_control_info(robot, dt_s, fps=fps)
+
+            timestamp = time.perf_counter() - start_episode_t
+            if events["exit_early"] or terminated:
+                events["exit_early"] = False
+                break
+
+        if events["rerecord_episode"]:
+            log_say("Re-record episode", play_sounds)
+            events["rerecord_episode"] = False
+            events["exit_early"] = False
+            dataset.clear_episode_buffer()
+            continue
+
+        dataset.save_episode(task=task)
+        recorded_episodes += 1
+
+        if events["stop_recording"] or recorded_episodes >= num_episodes:
+            break
+        else:
+            logging.info("Waiting for a few seconds before starting next episode recording...")
+            busy_wait(3)
+
+    log_say("Stop recording", play_sounds, blocking=True)
+    stop_recording(robot, listener, display_cameras)
+
+    if run_compute_stats:
+        logging.info("Computing dataset statistics")
+    dataset.consolidate(run_compute_stats)
+
+    if push_to_hub:
+        dataset.push_to_hub(tags=tags)
+
+    log_say("Exiting", play_sounds)
+    return dataset
+
+
+def replay(
+    env, root: Path, repo_id: str, episode: int, fps: int | None = None, local_files_only: bool = True
+):
+    env = env()
+
+    local_dir = Path(root) / repo_id
+    if not local_dir.exists():
+        raise ValueError(local_dir)
+
+    dataset = LeRobotDataset(repo_id, root=root, local_files_only=local_files_only)
+    items = dataset.hf_dataset.select_columns("action")
+    seeds = dataset.hf_dataset.select_columns("seed")["seed"]
+
+    from_idx = dataset.episode_data_index["from"][episode].item()
+    to_idx = dataset.episode_data_index["to"][episode].item()
+    env.reset(seed=seeds[from_idx].item())
+    logging.info("Replaying episode")
+    log_say("Replaying episode", play_sounds=True)
+    for idx in range(from_idx, to_idx):
+        start_episode_t = time.perf_counter()
+        action = items[idx]["action"]
+        env.step(action.unsqueeze(0).numpy())
+        dt_s = time.perf_counter() - start_episode_t
+        busy_wait(1 / fps - dt_s)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    subparsers = parser.add_subparsers(dest="mode", required=True)
+
+    # Set common options for all the subparsers
+    base_parser = argparse.ArgumentParser(add_help=False)
+    base_parser.add_argument(
+        "--robot-path",
+        type=str,
+        default="lerobot/configs/robot/koch.yaml",
+        help="Path to robot yaml file used to instantiate the robot using `make_robot` factory function.",
+    )
+
+    base_parser.add_argument(
+        "--sim-config",
+        help="Path to a yaml config you want to use for initializing a sim environment based on gym ",
+    )
+
+    parser_record = subparsers.add_parser("teleoperate", parents=[base_parser])
+
+    parser_record = subparsers.add_parser("record", parents=[base_parser])
+    parser_record.add_argument(
+        "--fps", type=none_or_int, default=None, help="Frames per second (set to None to disable)"
+    )
+    parser_record.add_argument(
+        "--root",
+        type=Path,
+        default=None,
+        help="Root directory where the dataset will be stored locally at '{root}/{repo_id}' (e.g. 'data/hf_username/dataset_name').",
+    )
+    parser_record.add_argument(
+        "--repo-id",
+        type=str,
+        default="lerobot/test",
+        help="Dataset identifier. By convention it should match '{hf_username}/{dataset_name}' (e.g. `lerobot/test`).",
+    )
+    parser_record.add_argument(
+        "--episode-time-s",
+        type=int,
+        default=60,
+        help="Number of seconds for data recording for each episode.",
+    )
+    parser_record.add_argument(
+        "--task",
+        type=str,
+        required=True,
+        help="A description of the task preformed during recording that can be used as a language instruction.",
+    )
+    parser_record.add_argument("--num-episodes", type=int, default=50, help="Number of episodes to record.")
+    parser_record.add_argument(
+        "--run-compute-stats",
+        type=int,
+        default=1,
+        help="By default, run the computation of the data statistics at the end of data collection. Compute intensive and not required to just replay an episode.",
+    )
+    parser_record.add_argument(
+        "--push-to-hub",
+        type=int,
+        default=1,
+        help="Upload dataset to Hugging Face hub.",
+    )
+    parser_record.add_argument(
+        "--tags",
+        type=str,
+        nargs="*",
+        help="Add tags to your dataset on the hub.",
+    )
+    parser_record.add_argument(
+        "--num-image-writer-processes",
+        type=int,
+        default=0,
+        help=(
+            "Number of subprocesses handling the saving of frames as PNG. Set to 0 to use threads only; "
+            "set to ≥1 to use subprocesses, each using threads to write images. The best number of processes "
+            "and threads depends on your system. We recommend 4 threads per camera with 0 processes. "
+            "If fps is unstable, adjust the thread count. If still unstable, try using 1 or more subprocesses."
+        ),
+    )
+    parser_record.add_argument(
+        "--num-image-writer-threads-per-camera",
+        type=int,
+        default=4,
+        help=(
+            "Number of threads writing the frames as png images on disk, per camera. "
+            "Too much threads might cause unstable teleoperation fps due to main thread being blocked. "
+            "Not enough threads might cause low camera fps."
+        ),
+    )
+    parser_record.add_argument(
+        "--display-cameras",
+        type=int,
+        default=0,
+        help="Visualize image observations with opencv.",
+    )
+    parser_record.add_argument(
+        "--resume",
+        type=int,
+        default=0,
+        help="Resume recording on an existing dataset.",
+    )
+    parser_replay = subparsers.add_parser("replay", parents=[base_parser])
+    parser_replay.add_argument(
+        "--fps", type=none_or_int, default=None, help="Frames per second (set to None to disable)"
+    )
+    parser_replay.add_argument(
+        "--root",
+        type=Path,
+        default=None,
+        help="Root directory where the dataset will be stored locally (e.g. 'data/hf_username/dataset_name'). By default, stored in cache folder.",
+    )
+    parser_replay.add_argument(
+        "--repo-id",
+        type=str,
+        default="lerobot/test",
+        help="Dataset identifier. By convention it should match '{hf_username}/{dataset_name}' (e.g. `lerobot/test`).",
+    )
+    parser_replay.add_argument("--episode", type=int, default=0, help="Index of the episodes to replay.")
+
+    args = parser.parse_args()
+
+    init_logging()
+
+    control_mode = args.mode
+    robot_path = args.robot_path
+    env_config_path = args.sim_config
+    kwargs = vars(args)
+    del kwargs["mode"]
+    del kwargs["robot_path"]
+    del kwargs["sim_config"]
+
+    # make gym env
+    env_cfg = init_hydra_config(env_config_path)
+    importlib.import_module(f"gym_{env_cfg.env.type}")
+
+    def env_constructor():
+        return gym.make(env_cfg.env.handle, disable_env_checker=True, **env_cfg.env.gym)
+
+    robot = None
+    process_leader_actions_fn = None
+
+    if control_mode in ["teleoperate", "record"]:
+        # make robot
+        robot_overrides = ["~cameras", "~follower_arms"]
+        # TODO(rcadene): remove
+        robot_cfg = init_hydra_config(robot_path, robot_overrides)
+        robot = make_robot(robot_cfg)
+        robot.connect()
+
+        calib_kwgs = init_sim_calibration(robot, env_cfg.calibration)
+
+        def process_leader_actions_fn(action):
+            return real_positions_to_sim(action, **calib_kwgs)
+
+        robot.leader_arms.main.calibration = None
+
+    if control_mode == "teleoperate":
+        teleoperate(env_constructor, robot, process_leader_actions_fn)
+
+    elif control_mode == "record":
+        record(env_constructor, robot, process_leader_actions_fn, **kwargs)
+
+    elif control_mode == "replay":
+        replay(env_constructor, **kwargs)
+
+    else:
+        raise ValueError(
+            f"Invalid control mode: '{control_mode}', only valid modes are teleoperate, record and replay."
+        )
+
+    if robot and robot.is_connected:
+        # Disconnect manually to avoid a "Core dump" during process
+        # termination due to camera threads not properly exiting.
+        robot.disconnect()

lerobot/scripts/display_sys_info.py

@@ -0,0 +1,90 @@
+#!/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.
+
+"""Use this script to get a quick summary of your system config.
+It should be able to run without any of LeRobot's dependencies or LeRobot itself installed.
+"""
+
+import platform
+
+HAS_HF_HUB = True
+HAS_HF_DATASETS = True
+HAS_NP = True
+HAS_TORCH = True
+HAS_LEROBOT = True
+
+try:
+    import huggingface_hub
+except ImportError:
+    HAS_HF_HUB = False
+
+try:
+    import datasets
+except ImportError:
+    HAS_HF_DATASETS = False
+
+try:
+    import numpy as np
+except ImportError:
+    HAS_NP = False
+
+try:
+    import torch
+except ImportError:
+    HAS_TORCH = False
+
+try:
+    import lerobot
+except ImportError:
+    HAS_LEROBOT = False
+
+
+lerobot_version = lerobot.__version__ if HAS_LEROBOT else "N/A"
+hf_hub_version = huggingface_hub.__version__ if HAS_HF_HUB else "N/A"
+hf_datasets_version = datasets.__version__ if HAS_HF_DATASETS else "N/A"
+np_version = np.__version__ if HAS_NP else "N/A"
+
+torch_version = torch.__version__ if HAS_TORCH else "N/A"
+torch_cuda_available = torch.cuda.is_available() if HAS_TORCH else "N/A"
+cuda_version = torch._C._cuda_getCompiledVersion() if HAS_TORCH and torch.version.cuda is not None else "N/A"
+
+
+# TODO(aliberts): refactor into an actual command `lerobot env`
+def display_sys_info() -> dict:
+    """Run this to get basic system info to help for tracking issues & bugs."""
+    info = {
+        "`lerobot` version": lerobot_version,
+        "Platform": platform.platform(),
+        "Python version": platform.python_version(),
+        "Huggingface_hub version": hf_hub_version,
+        "Dataset version": hf_datasets_version,
+        "Numpy version": np_version,
+        "PyTorch version (GPU?)": f"{torch_version} ({torch_cuda_available})",
+        "Cuda version": cuda_version,
+        "Using GPU in script?": "<fill in>",
+        # "Using distributed or parallel set-up in script?": "<fill in>",
+    }
+    print("\nCopy-and-paste the text below in your GitHub issue and FILL OUT the last point.\n")
+    print(format_dict(info))
+    return info
+
+
+def format_dict(d: dict) -> str:
+    return "\n".join([f"- {prop}: {val}" for prop, val in d.items()]) + "\n"
+
+
+if __name__ == "__main__":
+    display_sys_info()

lerobot/scripts/eval.py

@@ -0,0 +1,502 @@
+#!/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.
+"""Evaluate a policy on an environment by running rollouts and computing metrics.
+
+Usage examples:
+
+You want to evaluate a model from the hub (eg: https://huggingface.co/lerobot/diffusion_pusht)
+for 10 episodes.
+
+```
+python lerobot/scripts/eval.py \
+    --policy.path=lerobot/diffusion_pusht \
+    --env.type=pusht \
+    --eval.batch_size=10 \
+    --eval.n_episodes=10 \
+    --use_amp=false \
+    --device=cuda
+```
+
+OR, you want to evaluate a model checkpoint from the LeRobot training script for 10 episodes.
+```
+python lerobot/scripts/eval.py \
+    --policy.path=outputs/train/diffusion_pusht/checkpoints/005000/pretrained_model \
+    --env.type=pusht \
+    --eval.batch_size=10 \
+    --eval.n_episodes=10 \
+    --use_amp=false \
+    --device=cuda
+```
+
+Note that in both examples, the repo/folder should contain at least `config.json` and `model.safetensors` files.
+
+You can learn about the CLI options for this script in the `EvalPipelineConfig` in lerobot/configs/eval.py
+"""
+
+import json
+import logging
+import threading
+import time
+from contextlib import nullcontext
+from copy import deepcopy
+from dataclasses import asdict
+from pathlib import Path
+from pprint import pformat
+from typing import Callable
+
+import einops
+import gymnasium as gym
+import numpy as np
+import torch
+from termcolor import colored
+from torch import Tensor, nn
+from tqdm import trange
+
+from lerobot.common.envs.factory import make_env
+from lerobot.common.envs.utils import preprocess_observation
+from lerobot.common.policies.factory import make_policy
+from lerobot.common.policies.pretrained import PreTrainedPolicy
+from lerobot.common.policies.utils import get_device_from_parameters
+from lerobot.common.utils.io_utils import write_video
+from lerobot.common.utils.random_utils import set_seed
+from lerobot.common.utils.utils import (
+    get_safe_torch_device,
+    init_logging,
+    inside_slurm,
+)
+from lerobot.configs import parser
+from lerobot.configs.eval import EvalPipelineConfig
+
+
+def rollout(
+    env: gym.vector.VectorEnv,
+    policy: PreTrainedPolicy,
+    seeds: list[int] | None = None,
+    return_observations: bool = False,
+    render_callback: Callable[[gym.vector.VectorEnv], None] | None = None,
+) -> dict:
+    """Run a batched policy rollout once through a batch of environments.
+
+    Note that all environments in the batch are run until the last environment is done. This means some
+    data will probably need to be discarded (for environments that aren't the first one to be done).
+
+    The return dictionary contains:
+        (optional) "observation": A a dictionary of (batch, sequence + 1, *) tensors mapped to observation
+            keys. NOTE the that this has an extra sequence element relative to the other keys in the
+            dictionary. This is because an extra observation is included for after the environment is
+            terminated or truncated.
+        "action": A (batch, sequence, action_dim) tensor of actions applied based on the observations (not
+            including the last observations).
+        "reward": A (batch, sequence) tensor of rewards received for applying the actions.
+        "success": A (batch, sequence) tensor of success conditions (the only time this can be True is upon
+            environment termination/truncation).
+        "done": A (batch, sequence) tensor of **cumulative** done conditions. For any given batch element,
+            the first True is followed by True's all the way till the end. This can be used for masking
+            extraneous elements from the sequences above.
+
+    Args:
+        env: The batch of environments.
+        policy: The policy. Must be a PyTorch nn module.
+        seeds: The environments are seeded once at the start of the rollout. If provided, this argument
+            specifies the seeds for each of the environments.
+        return_observations: Whether to include all observations in the returned rollout data. Observations
+            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.
+    Returns:
+        The dictionary described above.
+    """
+    assert isinstance(policy, nn.Module), "Policy must be a PyTorch nn module."
+    device = get_device_from_parameters(policy)
+
+    # Reset the policy and environments.
+    policy.reset()
+
+    observation, info = env.reset(seed=seeds)
+    if render_callback is not None:
+        render_callback(env)
+
+    all_observations = []
+    all_actions = []
+    all_rewards = []
+    all_successes = []
+    all_dones = []
+
+    step = 0
+    # Keep track of which environments are done.
+    done = np.array([False] * env.num_envs)
+    max_steps = env.call("_max_episode_steps")[0]
+    progbar = trange(
+        max_steps,
+        desc=f"Running rollout with at most {max_steps} steps",
+        disable=inside_slurm(),  # we dont want progress bar when we use slurm, since it clutters the logs
+        leave=False,
+    )
+    while not np.all(done):
+        # Numpy array to tensor and changing dictionary keys to LeRobot policy format.
+        observation = preprocess_observation(observation)
+        if return_observations:
+            all_observations.append(deepcopy(observation))
+
+        observation = {
+            key: observation[key].to(device, non_blocking=device.type == "cuda") for key in observation
+        }
+
+        with torch.inference_mode():
+            action = policy.select_action(observation)
+
+        # Convert to CPU / numpy.
+        action = action.to("cpu").numpy()
+        assert action.ndim == 2, "Action dimensions should be (batch, action_dim)"
+
+        # Apply the next action.
+        observation, reward, terminated, truncated, info = env.step(action)
+        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 of none of the envs finished.
+        if "final_info" in info:
+            successes = [info["is_success"] if info is not None else False for info in info["final_info"]]
+        else:
+            successes = [False] * env.num_envs
+
+        # Keep track of which environments are done so far.
+        done = terminated | truncated | done
+
+        all_actions.append(torch.from_numpy(action))
+        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()
+
+    # Track the final observation.
+    if return_observations:
+        observation = preprocess_observation(observation)
+        all_observations.append(deepcopy(observation))
+
+    # Stack the sequence along the first dimension so that we have (batch, sequence, *) tensors.
+    ret = {
+        "action": torch.stack(all_actions, dim=1),
+        "reward": torch.stack(all_rewards, dim=1),
+        "success": torch.stack(all_successes, dim=1),
+        "done": torch.stack(all_dones, dim=1),
+    }
+    if return_observations:
+        stacked_observations = {}
+        for key in all_observations[0]:
+            stacked_observations[key] = torch.stack([obs[key] for obs in all_observations], dim=1)
+        ret["observation"] = stacked_observations
+
+    if hasattr(policy, "use_original_modules"):
+        policy.use_original_modules()
+
+    return ret
+
+
+def eval_policy(
+    env: gym.vector.VectorEnv,
+    policy: PreTrainedPolicy,
+    n_episodes: int,
+    max_episodes_rendered: int = 0,
+    videos_dir: Path | None = None,
+    return_episode_data: bool = False,
+    start_seed: int | None = None,
+) -> dict:
+    """
+    Args:
+        env: The batch of environments.
+        policy: The policy.
+        n_episodes: The number of episodes to evaluate.
+        max_episodes_rendered: Maximum number of episodes to render into videos.
+        videos_dir: Where to save rendered videos.
+        return_episode_data: Whether to return episode data for online training. Incorporates the data into
+            the "episodes" key of the returned dictionary.
+        start_seed: The first seed to use for the first individual rollout. For all subsequent rollouts the
+            seed is incremented by 1. If not provided, the environments are not manually seeded.
+    Returns:
+        Dictionary with metrics and data regarding the rollouts.
+    """
+    if max_episodes_rendered > 0 and not videos_dir:
+        raise ValueError("If max_episodes_rendered > 0, videos_dir must be provided.")
+
+    if not isinstance(policy, PreTrainedPolicy):
+        raise ValueError(
+            f"Policy of type 'PreTrainedPolicy' is expected, but type '{type(policy)}' was provided."
+        )
+
+    start = time.time()
+    policy.eval()
+
+    # Determine how many batched rollouts we need to get n_episodes. Note that if n_episodes is not evenly
+    # divisible by env.num_envs we end up discarding some data in the last batch.
+    n_batches = n_episodes // env.num_envs + int((n_episodes % env.num_envs) != 0)
+
+    # Keep track of some metrics.
+    sum_rewards = []
+    max_rewards = []
+    all_successes = []
+    all_seeds = []
+    threads = []  # for video saving threads
+    n_episodes_rendered = 0  # for saving the correct number of videos
+
+    # Callback for visualization.
+    def render_frame(env: gym.vector.VectorEnv):
+        # noqa: B023
+        if n_episodes_rendered >= max_episodes_rendered:
+            return
+        n_to_render_now = min(max_episodes_rendered - n_episodes_rendered, env.num_envs)
+        if isinstance(env, gym.vector.SyncVectorEnv):
+            ep_frames.append(np.stack([env.envs[i].render() for i in range(n_to_render_now)]))  # noqa: B023
+        elif isinstance(env, gym.vector.AsyncVectorEnv):
+            # Here we must render all frames and discard any we don't need.
+            ep_frames.append(np.stack(env.call("render")[:n_to_render_now]))
+
+    if max_episodes_rendered > 0:
+        video_paths: list[str] = []
+
+    if return_episode_data:
+        episode_data: dict | None = None
+
+    # we dont want progress bar when we use slurm, since it clutters the logs
+    progbar = trange(n_batches, desc="Stepping through eval batches", disable=inside_slurm())
+    for batch_ix in progbar:
+        # Cache frames for rendering videos. Each item will be (b, h, w, c), and the list indexes the rollout
+        # step.
+        if max_episodes_rendered > 0:
+            ep_frames: list[np.ndarray] = []
+
+        if start_seed is None:
+            seeds = None
+        else:
+            seeds = range(
+                start_seed + (batch_ix * env.num_envs), start_seed + ((batch_ix + 1) * env.num_envs)
+            )
+        rollout_data = rollout(
+            env,
+            policy,
+            seeds=list(seeds) if seeds else None,
+            return_observations=return_episode_data,
+            render_callback=render_frame if max_episodes_rendered > 0 else None,
+        )
+
+        # Figure out where in each rollout sequence the first done condition was encountered (results after
+        # this won't be included).
+        n_steps = rollout_data["done"].shape[1]
+        # Note: this relies on a property of argmax: that it returns the first occurrence as a tiebreaker.
+        done_indices = torch.argmax(rollout_data["done"].to(int), dim=1)
+
+        # Make a mask with shape (batch, n_steps) to mask out rollout data after the first done
+        # (batch-element-wise). Note the `done_indices + 1` to make sure to keep the data from the done step.
+        mask = (torch.arange(n_steps) <= einops.repeat(done_indices + 1, "b -> b s", s=n_steps)).int()
+        # Extend metrics.
+        batch_sum_rewards = einops.reduce((rollout_data["reward"] * mask), "b n -> b", "sum")
+        sum_rewards.extend(batch_sum_rewards.tolist())
+        batch_max_rewards = einops.reduce((rollout_data["reward"] * mask), "b n -> b", "max")
+        max_rewards.extend(batch_max_rewards.tolist())
+        batch_successes = einops.reduce((rollout_data["success"] * mask), "b n -> b", "any")
+        all_successes.extend(batch_successes.tolist())
+        if seeds:
+            all_seeds.extend(seeds)
+        else:
+            all_seeds.append(None)
+
+        # FIXME: episode_data is either None or it doesn't exist
+        if return_episode_data:
+            this_episode_data = _compile_episode_data(
+                rollout_data,
+                done_indices,
+                start_episode_index=batch_ix * env.num_envs,
+                start_data_index=(0 if episode_data is None else (episode_data["index"][-1].item() + 1)),
+                fps=env.unwrapped.metadata["render_fps"],
+            )
+            if episode_data is None:
+                episode_data = this_episode_data
+            else:
+                # Some sanity checks to make sure we are correctly compiling the data.
+                assert episode_data["episode_index"][-1] + 1 == this_episode_data["episode_index"][0]
+                assert episode_data["index"][-1] + 1 == this_episode_data["index"][0]
+                # Concatenate the episode data.
+                episode_data = {k: torch.cat([episode_data[k], this_episode_data[k]]) for k in episode_data}
+
+        # Maybe render video for visualization.
+        if max_episodes_rendered > 0 and len(ep_frames) > 0:
+            batch_stacked_frames = np.stack(ep_frames, axis=1)  # (b, t, *)
+            for stacked_frames, done_index in zip(
+                batch_stacked_frames, done_indices.flatten().tolist(), strict=False
+            ):
+                if n_episodes_rendered >= max_episodes_rendered:
+                    break
+
+                videos_dir.mkdir(parents=True, exist_ok=True)
+                video_path = videos_dir / f"eval_episode_{n_episodes_rendered}.mp4"
+                video_paths.append(str(video_path))
+                thread = threading.Thread(
+                    target=write_video,
+                    args=(
+                        str(video_path),
+                        stacked_frames[: done_index + 1],  # + 1 to capture the last observation
+                        env.unwrapped.metadata["render_fps"],
+                    ),
+                )
+                thread.start()
+                threads.append(thread)
+                n_episodes_rendered += 1
+
+        progbar.set_postfix(
+            {"running_success_rate": f"{np.mean(all_successes[:n_episodes]).item() * 100:.1f}%"}
+        )
+
+    # Wait till all video rendering threads are done.
+    for thread in threads:
+        thread.join()
+
+    # Compile eval info.
+    info = {
+        "per_episode": [
+            {
+                "episode_ix": i,
+                "sum_reward": sum_reward,
+                "max_reward": max_reward,
+                "success": success,
+                "seed": seed,
+            }
+            for i, (sum_reward, max_reward, success, seed) in enumerate(
+                zip(
+                    sum_rewards[:n_episodes],
+                    max_rewards[:n_episodes],
+                    all_successes[:n_episodes],
+                    all_seeds[:n_episodes],
+                    strict=True,
+                )
+            )
+        ],
+        "aggregated": {
+            "avg_sum_reward": float(np.nanmean(sum_rewards[:n_episodes])),
+            "avg_max_reward": float(np.nanmean(max_rewards[:n_episodes])),
+            "pc_success": float(np.nanmean(all_successes[:n_episodes]) * 100),
+            "eval_s": time.time() - start,
+            "eval_ep_s": (time.time() - start) / n_episodes,
+        },
+    }
+
+    if return_episode_data:
+        info["episodes"] = episode_data
+
+    if max_episodes_rendered > 0:
+        info["video_paths"] = video_paths
+
+    return info
+
+
+def _compile_episode_data(
+    rollout_data: dict, done_indices: Tensor, start_episode_index: int, start_data_index: int, fps: float
+) -> dict:
+    """Convenience function for `eval_policy(return_episode_data=True)`
+
+    Compiles all the rollout data into a Hugging Face dataset.
+
+    Similar logic is implemented when datasets are pushed to hub (see: `push_to_hub`).
+    """
+    ep_dicts = []
+    total_frames = 0
+    for ep_ix in range(rollout_data["action"].shape[0]):
+        # + 2 to include the first done frame and the last observation frame.
+        num_frames = done_indices[ep_ix].item() + 2
+        total_frames += num_frames
+
+        # Here we do `num_frames - 1` as we don't want to include the last observation frame just yet.
+        ep_dict = {
+            "action": rollout_data["action"][ep_ix, : num_frames - 1],
+            "episode_index": torch.tensor([start_episode_index + ep_ix] * (num_frames - 1)),
+            "frame_index": torch.arange(0, num_frames - 1, 1),
+            "timestamp": torch.arange(0, num_frames - 1, 1) / fps,
+            "next.done": rollout_data["done"][ep_ix, : num_frames - 1],
+            "next.success": rollout_data["success"][ep_ix, : num_frames - 1],
+            "next.reward": rollout_data["reward"][ep_ix, : num_frames - 1].type(torch.float32),
+        }
+
+        # For the last observation frame, all other keys will just be copy padded.
+        for k in ep_dict:
+            ep_dict[k] = torch.cat([ep_dict[k], ep_dict[k][-1:]])
+
+        for key in rollout_data["observation"]:
+            ep_dict[key] = rollout_data["observation"][key][ep_ix, :num_frames]
+
+        ep_dicts.append(ep_dict)
+
+    data_dict = {}
+    for key in ep_dicts[0]:
+        data_dict[key] = torch.cat([x[key] for x in ep_dicts])
+
+    data_dict["index"] = torch.arange(start_data_index, start_data_index + total_frames, 1)
+
+    return data_dict
+
+
+@parser.wrap()
+def eval_main(cfg: EvalPipelineConfig):
+    logging.info(pformat(asdict(cfg)))
+
+    # Check device is available
+    device = get_safe_torch_device(cfg.policy.device, log=True)
+
+    torch.backends.cudnn.benchmark = True
+    torch.backends.cuda.matmul.allow_tf32 = True
+    set_seed(cfg.seed)
+
+    logging.info(colored("Output dir:", "yellow", attrs=["bold"]) + f" {cfg.output_dir}")
+
+    logging.info("Making environment.")
+    env = make_env(cfg.env, n_envs=cfg.eval.batch_size, use_async_envs=cfg.eval.use_async_envs)
+
+    logging.info("Making policy.")
+
+    policy = make_policy(
+        cfg=cfg.policy,
+        env_cfg=cfg.env,
+    )
+    policy.eval()
+
+    with torch.no_grad(), torch.autocast(device_type=device.type) if cfg.policy.use_amp else nullcontext():
+        info = eval_policy(
+            env,
+            policy,
+            cfg.eval.n_episodes,
+            max_episodes_rendered=10,
+            videos_dir=Path(cfg.output_dir) / "videos",
+            start_seed=cfg.seed,
+        )
+    print(info["aggregated"])
+
+    # Save info
+    with open(Path(cfg.output_dir) / "eval_info.json", "w") as f:
+        json.dump(info, f, indent=2)
+
+    env.close()
+
+    logging.info("End of eval")
+
+
+if __name__ == "__main__":
+    init_logging()
+    eval_main()

lerobot/scripts/find_motors_bus_port.py

@@ -0,0 +1,55 @@
+# 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 os
+import time
+from pathlib import Path
+
+from serial.tools import list_ports  # Part of pyserial library
+
+
+def find_available_ports():
+    if os.name == "nt":  # Windows
+        # List COM ports using pyserial
+        ports = [port.device for port in list_ports.comports()]
+    else:  # Linux/macOS
+        # List /dev/tty* ports for Unix-based systems
+        ports = [str(path) for path in Path("/dev").glob("tty*")]
+    return ports
+
+
+def find_port():
+    print("Finding all available ports for the MotorsBus.")
+    ports_before = find_available_ports()
+    print("Ports before disconnecting:", ports_before)
+
+    print("Remove the USB cable from your MotorsBus and press Enter when done.")
+    input()  # Wait for user to disconnect the device
+
+    time.sleep(0.5)  # Allow some time for port to be released
+    ports_after = find_available_ports()
+    ports_diff = list(set(ports_before) - set(ports_after))
+
+    if len(ports_diff) == 1:
+        port = ports_diff[0]
+        print(f"The port of this MotorsBus is '{port}'")
+        print("Reconnect the USB cable.")
+    elif len(ports_diff) == 0:
+        raise OSError(f"Could not detect the port. No difference was found ({ports_diff}).")
+    else:
+        raise OSError(f"Could not detect the port. More than one port was found ({ports_diff}).")
+
+
+if __name__ == "__main__":
+    # Helper to find the USB port associated with your MotorsBus.
+    find_port()

lerobot/scripts/push_dataset_to_hub.py

@@ -0,0 +1,364 @@
+#!/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.
+"""
+Use this script to convert your dataset into LeRobot dataset format and upload it to the Hugging Face hub,
+or store it locally. LeRobot dataset format is lightweight, fast to load from, and does not require any
+installation of neural net specific packages like pytorch, tensorflow, jax.
+
+Example of how to download raw datasets, convert them into LeRobotDataset format, and push them to the hub:
+```
+python lerobot/scripts/push_dataset_to_hub.py \
+--raw-dir data/pusht_raw \
+--raw-format pusht_zarr \
+--repo-id lerobot/pusht
+
+python lerobot/scripts/push_dataset_to_hub.py \
+--raw-dir data/xarm_lift_medium_raw \
+--raw-format xarm_pkl \
+--repo-id lerobot/xarm_lift_medium
+
+python lerobot/scripts/push_dataset_to_hub.py \
+--raw-dir data/aloha_sim_insertion_scripted_raw \
+--raw-format aloha_hdf5 \
+--repo-id lerobot/aloha_sim_insertion_scripted
+
+python lerobot/scripts/push_dataset_to_hub.py \
+--raw-dir data/umi_cup_in_the_wild_raw \
+--raw-format umi_zarr \
+--repo-id lerobot/umi_cup_in_the_wild
+```
+"""
+
+import argparse
+import json
+import shutil
+import warnings
+from pathlib import Path
+from typing import Any
+
+import torch
+from huggingface_hub import HfApi
+from safetensors.torch import save_file
+
+from lerobot.common.datasets.compute_stats import compute_stats
+from lerobot.common.datasets.lerobot_dataset import CODEBASE_VERSION, LeRobotDataset
+from lerobot.common.datasets.push_dataset_to_hub.utils import check_repo_id
+from lerobot.common.datasets.utils import create_branch, create_lerobot_dataset_card, flatten_dict
+
+
+def get_from_raw_to_lerobot_format_fn(raw_format: str):
+    if raw_format == "pusht_zarr":
+        from lerobot.common.datasets.push_dataset_to_hub.pusht_zarr_format import from_raw_to_lerobot_format
+    elif raw_format == "umi_zarr":
+        from lerobot.common.datasets.push_dataset_to_hub.umi_zarr_format import from_raw_to_lerobot_format
+    elif raw_format == "aloha_hdf5":
+        from lerobot.common.datasets.push_dataset_to_hub.aloha_hdf5_format import from_raw_to_lerobot_format
+    elif raw_format in ["rlds", "openx"]:
+        from lerobot.common.datasets.push_dataset_to_hub.openx_rlds_format import from_raw_to_lerobot_format
+    elif raw_format == "dora_parquet":
+        from lerobot.common.datasets.push_dataset_to_hub.dora_parquet_format import from_raw_to_lerobot_format
+    elif raw_format == "xarm_pkl":
+        from lerobot.common.datasets.push_dataset_to_hub.xarm_pkl_format import from_raw_to_lerobot_format
+    elif raw_format == "cam_png":
+        from lerobot.common.datasets.push_dataset_to_hub.cam_png_format import from_raw_to_lerobot_format
+    else:
+        raise ValueError(
+            f"The selected {raw_format} can't be found. Did you add it to `lerobot/scripts/push_dataset_to_hub.py::get_from_raw_to_lerobot_format_fn`?"
+        )
+
+    return from_raw_to_lerobot_format
+
+
+def save_meta_data(
+    info: dict[str, Any], stats: dict, episode_data_index: dict[str, list], meta_data_dir: Path
+):
+    meta_data_dir.mkdir(parents=True, exist_ok=True)
+
+    # save info
+    info_path = meta_data_dir / "info.json"
+    with open(str(info_path), "w") as f:
+        json.dump(info, f, indent=4)
+
+    # save stats
+    stats_path = meta_data_dir / "stats.safetensors"
+    save_file(flatten_dict(stats), stats_path)
+
+    # save episode_data_index
+    episode_data_index = {key: torch.tensor(episode_data_index[key]) for key in episode_data_index}
+    ep_data_idx_path = meta_data_dir / "episode_data_index.safetensors"
+    save_file(episode_data_index, ep_data_idx_path)
+
+
+def push_meta_data_to_hub(repo_id: str, meta_data_dir: str | Path, revision: str | None):
+    """Expect all meta data files to be all stored in a single "meta_data" directory.
+    On the hugging face repositery, they will be uploaded in a "meta_data" directory at the root.
+    """
+    api = HfApi()
+    api.upload_folder(
+        folder_path=meta_data_dir,
+        path_in_repo="meta_data",
+        repo_id=repo_id,
+        revision=revision,
+        repo_type="dataset",
+    )
+
+
+def push_dataset_card_to_hub(
+    repo_id: str,
+    revision: str | None,
+    tags: list | None = None,
+    license: str = "apache-2.0",
+    **card_kwargs,
+):
+    """Creates and pushes a LeRobotDataset Card with appropriate tags to easily find it on the hub."""
+    card = create_lerobot_dataset_card(tags=tags, license=license, **card_kwargs)
+    card.push_to_hub(repo_id=repo_id, repo_type="dataset", revision=revision)
+
+
+def push_videos_to_hub(repo_id: str, videos_dir: str | Path, revision: str | None):
+    """Expect mp4 files to be all stored in a single "videos" directory.
+    On the hugging face repositery, they will be uploaded in a "videos" directory at the root.
+    """
+    api = HfApi()
+    api.upload_folder(
+        folder_path=videos_dir,
+        path_in_repo="videos",
+        repo_id=repo_id,
+        revision=revision,
+        repo_type="dataset",
+        allow_patterns="*.mp4",
+    )
+
+
+def push_dataset_to_hub(
+    raw_dir: Path,
+    raw_format: str,
+    repo_id: str,
+    push_to_hub: bool = True,
+    local_dir: Path | None = None,
+    fps: int | None = None,
+    video: bool = True,
+    batch_size: int = 32,
+    num_workers: int = 8,
+    episodes: list[int] | None = None,
+    force_override: bool = False,
+    resume: bool = False,
+    cache_dir: Path = Path("/tmp"),
+    tests_data_dir: Path | None = None,
+    encoding: dict | None = None,
+):
+    check_repo_id(repo_id)
+    user_id, dataset_id = repo_id.split("/")
+
+    # Robustify when `raw_dir` is str instead of Path
+    raw_dir = Path(raw_dir)
+    if not raw_dir.exists():
+        raise NotADirectoryError(
+            f"{raw_dir} does not exists. Check your paths or run this command to download an existing raw dataset on the hub: "
+            f"`python lerobot/common/datasets/push_dataset_to_hub/_download_raw.py --raw-dir your/raw/dir --repo-id your/repo/id_raw`"
+        )
+
+    if local_dir:
+        # Robustify when `local_dir` is str instead of Path
+        local_dir = Path(local_dir)
+
+        # Send warning if local_dir isn't well formatted
+        if local_dir.parts[-2] != user_id or local_dir.parts[-1] != dataset_id:
+            warnings.warn(
+                f"`local_dir` ({local_dir}) doesn't contain a community or user id `/` the name of the dataset that match the `repo_id` (e.g. 'data/lerobot/pusht'). Following this naming convention is advised, but not mandatory.",
+                stacklevel=1,
+            )
+
+        # Check we don't override an existing `local_dir` by mistake
+        if local_dir.exists():
+            if force_override:
+                shutil.rmtree(local_dir)
+            elif not resume:
+                raise ValueError(f"`local_dir` already exists ({local_dir}). Use `--force-override 1`.")
+
+        meta_data_dir = local_dir / "meta_data"
+        videos_dir = local_dir / "videos"
+    else:
+        # Temporary directory used to store images, videos, meta_data
+        meta_data_dir = Path(cache_dir) / "meta_data"
+        videos_dir = Path(cache_dir) / "videos"
+
+    if raw_format is None:
+        # TODO(rcadene, adilzouitine): implement auto_find_raw_format
+        raise NotImplementedError()
+        # raw_format = auto_find_raw_format(raw_dir)
+
+    # convert dataset from original raw format to LeRobot format
+    from_raw_to_lerobot_format = get_from_raw_to_lerobot_format_fn(raw_format)
+
+    hf_dataset, episode_data_index, info = from_raw_to_lerobot_format(
+        raw_dir,
+        videos_dir,
+        fps,
+        video,
+        episodes,
+        encoding,
+    )
+
+    lerobot_dataset = LeRobotDataset.from_preloaded(
+        repo_id=repo_id,
+        hf_dataset=hf_dataset,
+        episode_data_index=episode_data_index,
+        info=info,
+        videos_dir=videos_dir,
+    )
+    stats = compute_stats(lerobot_dataset, batch_size, num_workers)
+
+    if local_dir:
+        hf_dataset = hf_dataset.with_format(None)  # to remove transforms that cant be saved
+        hf_dataset.save_to_disk(str(local_dir / "train"))
+
+    if push_to_hub or local_dir:
+        # mandatory for upload
+        save_meta_data(info, stats, episode_data_index, meta_data_dir)
+
+    if push_to_hub:
+        hf_dataset.push_to_hub(repo_id, revision="main")
+        push_meta_data_to_hub(repo_id, meta_data_dir, revision="main")
+        push_dataset_card_to_hub(repo_id, revision="main")
+        if video:
+            push_videos_to_hub(repo_id, videos_dir, revision="main")
+        create_branch(repo_id, repo_type="dataset", branch=CODEBASE_VERSION)
+
+    if tests_data_dir:
+        # get the first episode
+        num_items_first_ep = episode_data_index["to"][0] - episode_data_index["from"][0]
+        test_hf_dataset = hf_dataset.select(range(num_items_first_ep))
+        episode_data_index = {k: v[:1] for k, v in episode_data_index.items()}
+
+        test_hf_dataset = test_hf_dataset.with_format(None)
+        test_hf_dataset.save_to_disk(str(tests_data_dir / repo_id / "train"))
+
+        tests_meta_data = tests_data_dir / repo_id / "meta_data"
+        save_meta_data(info, stats, episode_data_index, tests_meta_data)
+
+        # copy videos of first episode to tests directory
+        episode_index = 0
+        tests_videos_dir = tests_data_dir / repo_id / "videos"
+        tests_videos_dir.mkdir(parents=True, exist_ok=True)
+        for key in lerobot_dataset.camera_keys:
+            fname = f"{key}_episode_{episode_index:06d}.mp4"
+            shutil.copy(videos_dir / fname, tests_videos_dir / fname)
+
+    if local_dir is None:
+        # clear cache
+        shutil.rmtree(meta_data_dir)
+        shutil.rmtree(videos_dir)
+
+    return lerobot_dataset
+
+
+def main():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--raw-dir",
+        type=Path,
+        required=True,
+        help="Directory containing input raw datasets (e.g. `data/aloha_mobile_chair_raw` or `data/pusht_raw).",
+    )
+    # TODO(rcadene): add automatic detection of the format
+    parser.add_argument(
+        "--raw-format",
+        type=str,
+        required=True,
+        help="Dataset type (e.g. `pusht_zarr`, `umi_zarr`, `aloha_hdf5`, `xarm_pkl`, `dora_parquet`, `rlds`, `openx`).",
+    )
+    parser.add_argument(
+        "--repo-id",
+        type=str,
+        required=True,
+        help="Repositery identifier on Hugging Face: a community or a user name `/` the name of the dataset (e.g. `lerobot/pusht`, `cadene/aloha_sim_insertion_human`).",
+    )
+    parser.add_argument(
+        "--local-dir",
+        type=Path,
+        help="When provided, writes the dataset converted to LeRobotDataset format in this directory  (e.g. `data/lerobot/aloha_mobile_chair`).",
+    )
+    parser.add_argument(
+        "--push-to-hub",
+        type=int,
+        default=1,
+        help="Upload to hub.",
+    )
+    parser.add_argument(
+        "--fps",
+        type=int,
+        help="Frame rate used to collect videos. If not provided, use the default one specified in the code.",
+    )
+    parser.add_argument(
+        "--video",
+        type=int,
+        default=1,
+        help="Convert each episode of the raw dataset to an mp4 video. This option allows 60 times lower disk space consumption and 25 faster loading time during training.",
+    )
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=32,
+        help="Batch size loaded by DataLoader for computing the dataset statistics.",
+    )
+    parser.add_argument(
+        "--num-workers",
+        type=int,
+        default=8,
+        help="Number of processes of Dataloader for computing the dataset statistics.",
+    )
+    parser.add_argument(
+        "--episodes",
+        type=int,
+        nargs="*",
+        help="When provided, only converts the provided episodes (e.g `--episodes 2 3 4`). Useful to test the code on 1 episode.",
+    )
+    parser.add_argument(
+        "--force-override",
+        type=int,
+        default=0,
+        help="When set to 1, removes provided output directory if it already exists. By default, raises a ValueError exception.",
+    )
+    parser.add_argument(
+        "--resume",
+        type=int,
+        default=0,
+        help="When set to 1, resumes a previous run.",
+    )
+    parser.add_argument(
+        "--cache-dir",
+        type=Path,
+        required=False,
+        default="/tmp",
+        help="Directory to store the temporary videos and images generated while creating the dataset.",
+    )
+    parser.add_argument(
+        "--tests-data-dir",
+        type=Path,
+        help=(
+            "When provided, save tests artifacts into the given directory "
+            "(e.g. `--tests-data-dir tests/data` will save to tests/data/{--repo-id})."
+        ),
+    )
+
+    args = parser.parse_args()
+    push_dataset_to_hub(**vars(args))
+
+
+if __name__ == "__main__":
+    main()

lerobot/scripts/push_pretrained.py

@@ -0,0 +1,71 @@
+#!/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.
+"""
+Once you have trained a policy with our training script (lerobot/scripts/train.py), use this script to push it
+to the hub.
+
+Example:
+
+```bash
+python lerobot/scripts/push_pretrained.py \
+    --pretrained_path=outputs/train/act_aloha_sim_transfer_cube_human/checkpoints/last/pretrained_model \
+    --repo_id=lerobot/act_aloha_sim_transfer_cube_human
+```
+"""
+
+from dataclasses import dataclass
+from pathlib import Path
+
+import draccus
+from huggingface_hub import HfApi
+
+
+@dataclass
+class PushPreTrainedConfig:
+    pretrained_path: Path
+    repo_id: str
+    branch: str | None = None
+    private: bool = False
+    exist_ok: bool = False
+
+
+@draccus.wrap()
+def main(cfg: PushPreTrainedConfig):
+    hub_api = HfApi()
+    hub_api.create_repo(
+        repo_id=cfg.repo_id,
+        private=cfg.private,
+        repo_type="model",
+        exist_ok=cfg.exist_ok,
+    )
+    if cfg.branch:
+        hub_api.create_branch(
+            repo_id=cfg.repo_id,
+            branch=cfg.branch,
+            repo_type="model",
+            exist_ok=cfg.exist_ok,
+        )
+
+    hub_api.upload_folder(
+        repo_id=cfg.repo_id,
+        folder_path=cfg.pretrained_path,
+        repo_type="model",
+        revision=cfg.branch,
+    )
+
+
+if __name__ == "__main__":
+    main()

lerobot/scripts/train.py

@@ -0,0 +1,288 @@
+#!/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 logging
+import time
+from contextlib import nullcontext
+from pprint import pformat
+from typing import Any
+
+import torch
+from termcolor import colored
+from torch.amp import GradScaler
+from torch.optim import Optimizer
+
+from lerobot.common.datasets.factory import make_dataset
+from lerobot.common.datasets.sampler import EpisodeAwareSampler
+from lerobot.common.datasets.utils import cycle
+from lerobot.common.envs.factory import make_env
+from lerobot.common.optim.factory import make_optimizer_and_scheduler
+from lerobot.common.policies.factory import make_policy
+from lerobot.common.policies.pretrained import PreTrainedPolicy
+from lerobot.common.policies.utils import get_device_from_parameters
+from lerobot.common.utils.logging_utils import AverageMeter, MetricsTracker
+from lerobot.common.utils.random_utils import set_seed
+from lerobot.common.utils.train_utils import (
+    get_step_checkpoint_dir,
+    get_step_identifier,
+    load_training_state,
+    save_checkpoint,
+    update_last_checkpoint,
+)
+from lerobot.common.utils.utils import (
+    format_big_number,
+    get_safe_torch_device,
+    has_method,
+    init_logging,
+)
+from lerobot.common.utils.wandb_utils import WandBLogger
+from lerobot.configs import parser
+from lerobot.configs.train import TrainPipelineConfig
+from lerobot.scripts.eval import eval_policy
+
+
+def update_policy(
+    train_metrics: MetricsTracker,
+    policy: PreTrainedPolicy,
+    batch: Any,
+    optimizer: Optimizer,
+    grad_clip_norm: float,
+    grad_scaler: GradScaler,
+    lr_scheduler=None,
+    use_amp: bool = False,
+    lock=None,
+) -> tuple[MetricsTracker, dict]:
+    start_time = time.perf_counter()
+    device = get_device_from_parameters(policy)
+    policy.train()
+    with torch.autocast(device_type=device.type) if use_amp else nullcontext():
+        loss, output_dict = policy.forward(batch)
+        # TODO(rcadene): policy.unnormalize_outputs(out_dict)
+    grad_scaler.scale(loss).backward()
+
+    # Unscale the gradient of the optimizer's assigned params in-place **prior to gradient clipping**.
+    grad_scaler.unscale_(optimizer)
+
+    grad_norm = torch.nn.utils.clip_grad_norm_(
+        policy.parameters(),
+        grad_clip_norm,
+        error_if_nonfinite=False,
+    )
+
+    # Optimizer's gradients are already unscaled, so scaler.step does not unscale them,
+    # although it still skips optimizer.step() if the gradients contain infs or NaNs.
+    with lock if lock is not None else nullcontext():
+        grad_scaler.step(optimizer)
+    # Updates the scale for next iteration.
+    grad_scaler.update()
+
+    optimizer.zero_grad()
+
+    # Step through pytorch scheduler at every batch instead of epoch
+    if lr_scheduler is not None:
+        lr_scheduler.step()
+
+    if has_method(policy, "update"):
+        # To possibly update an internal buffer (for instance an Exponential Moving Average like in TDMPC).
+        policy.update()
+
+    train_metrics.loss = loss.item()
+    train_metrics.grad_norm = grad_norm.item()
+    train_metrics.lr = optimizer.param_groups[0]["lr"]
+    train_metrics.update_s = time.perf_counter() - start_time
+    return train_metrics, output_dict
+
+
+@parser.wrap()
+def train(cfg: TrainPipelineConfig):
+    cfg.validate()
+    logging.info(pformat(cfg.to_dict()))
+
+    if cfg.wandb.enable and cfg.wandb.project:
+        wandb_logger = WandBLogger(cfg)
+    else:
+        wandb_logger = None
+        logging.info(colored("Logs will be saved locally.", "yellow", attrs=["bold"]))
+
+    if cfg.seed is not None:
+        set_seed(cfg.seed)
+
+    # Check device is available
+    device = get_safe_torch_device(cfg.policy.device, log=True)
+    torch.backends.cudnn.benchmark = True
+    torch.backends.cuda.matmul.allow_tf32 = True
+
+    logging.info("Creating dataset")
+    dataset = make_dataset(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:
+        logging.info("Creating env")
+        eval_env = make_env(cfg.env, n_envs=cfg.eval.batch_size)
+
+    logging.info("Creating policy")
+    policy = make_policy(
+        cfg=cfg.policy,
+        ds_meta=dataset.meta,
+    )
+
+    logging.info("Creating optimizer and scheduler")
+    optimizer, lr_scheduler = make_optimizer_and_scheduler(cfg, policy)
+    grad_scaler = GradScaler(device.type, enabled=cfg.policy.use_amp)
+
+    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)
+
+    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())
+
+    logging.info(colored("Output dir:", "yellow", attrs=["bold"]) + f" {cfg.output_dir}")
+    if cfg.env is not None:
+        logging.info(f"{cfg.env.task=}")
+    logging.info(f"{cfg.steps=} ({format_big_number(cfg.steps)})")
+    logging.info(f"{dataset.num_frames=} ({format_big_number(dataset.num_frames)})")
+    logging.info(f"{dataset.num_episodes=}")
+    logging.info(f"{num_learnable_params=} ({format_big_number(num_learnable_params)})")
+    logging.info(f"{num_total_params=} ({format_big_number(num_total_params)})")
+
+    # create dataloader for offline training
+    if hasattr(cfg.policy, "drop_n_last_frames"):
+        shuffle = False
+        sampler = EpisodeAwareSampler(
+            dataset.episode_data_index,
+            drop_n_last_frames=cfg.policy.drop_n_last_frames,
+            shuffle=True,
+        )
+    else:
+        shuffle = True
+        sampler = None
+
+    dataloader = torch.utils.data.DataLoader(
+        dataset,
+        num_workers=cfg.num_workers,
+        batch_size=cfg.batch_size,
+        shuffle=shuffle,
+        sampler=sampler,
+        pin_memory=device.type != "cpu",
+        drop_last=False,
+    )
+    dl_iter = cycle(dataloader)
+
+    policy.train()
+
+    train_metrics = {
+        "loss": AverageMeter("loss", ":.3f"),
+        "grad_norm": AverageMeter("grdn", ":.3f"),
+        "lr": AverageMeter("lr", ":0.1e"),
+        "update_s": AverageMeter("updt_s", ":.3f"),
+        "dataloading_s": AverageMeter("data_s", ":.3f"),
+    }
+
+    train_tracker = MetricsTracker(
+        cfg.batch_size, dataset.num_frames, dataset.num_episodes, train_metrics, initial_step=step
+    )
+
+    logging.info("Start offline training on a fixed dataset")
+    for _ in range(step, cfg.steps):
+        start_time = time.perf_counter()
+        batch = next(dl_iter)
+        train_tracker.dataloading_s = time.perf_counter() - start_time
+
+        for key in batch:
+            if isinstance(batch[key], torch.Tensor):
+                batch[key] = batch[key].to(device, non_blocking=True)
+
+        train_tracker, output_dict = update_policy(
+            train_tracker,
+            policy,
+            batch,
+            optimizer,
+            cfg.optimizer.grad_clip_norm,
+            grad_scaler=grad_scaler,
+            lr_scheduler=lr_scheduler,
+            use_amp=cfg.policy.use_amp,
+        )
+
+        # Note: eval and checkpoint happens *after* the `step`th training update has completed, so we
+        # increment `step` here.
+        step += 1
+        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
+
+        if is_log_step:
+            logging.info(train_tracker)
+            if wandb_logger:
+                wandb_log_dict = train_tracker.to_dict()
+                if output_dict:
+                    wandb_log_dict.update(output_dict)
+                wandb_logger.log_dict(wandb_log_dict, step)
+            train_tracker.reset_averages()
+
+        if cfg.save_checkpoint and is_saving_step:
+            logging.info(f"Checkpoint policy after step {step}")
+            checkpoint_dir = get_step_checkpoint_dir(cfg.output_dir, cfg.steps, step)
+            save_checkpoint(checkpoint_dir, step, cfg, policy, optimizer, lr_scheduler)
+            update_last_checkpoint(checkpoint_dir)
+            if wandb_logger:
+                wandb_logger.log_policy(checkpoint_dir)
+
+        if cfg.env and is_eval_step:
+            step_id = get_step_identifier(step, cfg.steps)
+            logging.info(f"Eval policy at step {step}")
+            with (
+                torch.no_grad(),
+                torch.autocast(device_type=device.type) if cfg.policy.use_amp else nullcontext(),
+            ):
+                eval_info = eval_policy(
+                    eval_env,
+                    policy,
+                    cfg.eval.n_episodes,
+                    videos_dir=cfg.output_dir / "eval" / f"videos_step_{step_id}",
+                    max_episodes_rendered=4,
+                    start_seed=cfg.seed,
+                )
+
+            eval_metrics = {
+                "avg_sum_reward": AverageMeter("∑rwrd", ":.3f"),
+                "pc_success": AverageMeter("success", ":.1f"),
+                "eval_s": AverageMeter("eval_s", ":.3f"),
+            }
+            eval_tracker = MetricsTracker(
+                cfg.batch_size, dataset.num_frames, dataset.num_episodes, eval_metrics, initial_step=step
+            )
+            eval_tracker.eval_s = eval_info["aggregated"].pop("eval_s")
+            eval_tracker.avg_sum_reward = eval_info["aggregated"].pop("avg_sum_reward")
+            eval_tracker.pc_success = eval_info["aggregated"].pop("pc_success")
+            logging.info(eval_tracker)
+            if wandb_logger:
+                wandb_log_dict = {**eval_tracker.to_dict(), **eval_info}
+                wandb_logger.log_dict(wandb_log_dict, step, mode="eval")
+                wandb_logger.log_video(eval_info["video_paths"][0], step, mode="eval")
+
+    if eval_env:
+        eval_env.close()
+    logging.info("End of training")
+
+
+if __name__ == "__main__":
+    init_logging()
+    train()

lerobot/scripts/visualize_dataset.py

@@ -0,0 +1,292 @@
+#!/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.
+""" Visualize data of **all** frames of any episode of a dataset of type LeRobotDataset.
+
+Note: The last frame of the episode doesn't always correspond to a final state.
+That's because our datasets are composed of transition from state to state up to
+the antepenultimate state associated to the ultimate action to arrive in the final state.
+However, there might not be a transition from a final state to another state.
+
+Note: This script aims to visualize the data used to train the neural networks.
+~What you see is what you get~. When visualizing image modality, it is often expected to observe
+lossy compression artifacts since these images have been decoded from compressed mp4 videos to
+save disk space. The compression factor applied has been tuned to not affect success rate.
+
+Examples:
+
+- Visualize data stored on a local machine:
+```
+local$ python lerobot/scripts/visualize_dataset.py \
+    --repo-id lerobot/pusht \
+    --episode-index 0
+```
+
+- Visualize data stored on a distant machine with a local viewer:
+```
+distant$ python lerobot/scripts/visualize_dataset.py \
+    --repo-id lerobot/pusht \
+    --episode-index 0 \
+    --save 1 \
+    --output-dir path/to/directory
+
+local$ scp distant:path/to/directory/lerobot_pusht_episode_0.rrd .
+local$ rerun lerobot_pusht_episode_0.rrd
+```
+
+- Visualize data stored on a distant machine through streaming:
+(You need to forward the websocket port to the distant machine, with
+`ssh -L 9087:localhost:9087 username@remote-host`)
+```
+distant$ python lerobot/scripts/visualize_dataset.py \
+    --repo-id lerobot/pusht \
+    --episode-index 0 \
+    --mode distant \
+    --ws-port 9087
+
+local$ rerun ws://localhost:9087
+```
+
+"""
+
+import argparse
+import gc
+import logging
+import time
+from pathlib import Path
+from typing import Iterator
+
+import numpy as np
+import rerun as rr
+import torch
+import torch.utils.data
+import tqdm
+
+from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
+
+
+class EpisodeSampler(torch.utils.data.Sampler):
+    def __init__(self, dataset: LeRobotDataset, episode_index: int):
+        from_idx = dataset.episode_data_index["from"][episode_index].item()
+        to_idx = dataset.episode_data_index["to"][episode_index].item()
+        self.frame_ids = range(from_idx, to_idx)
+
+    def __iter__(self) -> Iterator:
+        return iter(self.frame_ids)
+
+    def __len__(self) -> int:
+        return len(self.frame_ids)
+
+
+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}"
+    hwc_uint8_numpy = (chw_float32_torch * 255).type(torch.uint8).permute(1, 2, 0).numpy()
+    return hwc_uint8_numpy
+
+
+def visualize_dataset(
+    dataset: LeRobotDataset,
+    episode_index: int,
+    batch_size: int = 32,
+    num_workers: int = 0,
+    mode: str = "local",
+    web_port: int = 9090,
+    ws_port: int = 9087,
+    save: bool = False,
+    output_dir: Path | None = None,
+) -> Path | None:
+    if save:
+        assert output_dir is not None, (
+            "Set an output directory where to write .rrd files with `--output-dir path/to/directory`."
+        )
+
+    repo_id = dataset.repo_id
+
+    logging.info("Loading dataloader")
+    episode_sampler = EpisodeSampler(dataset, episode_index)
+    dataloader = torch.utils.data.DataLoader(
+        dataset,
+        num_workers=num_workers,
+        batch_size=batch_size,
+        sampler=episode_sampler,
+    )
+
+    logging.info("Starting Rerun")
+
+    if mode not in ["local", "distant"]:
+        raise ValueError(mode)
+
+    spawn_local_viewer = mode == "local" and not save
+    rr.init(f"{repo_id}/episode_{episode_index}", spawn=spawn_local_viewer)
+
+    # Manually call python garbage collector after `rr.init` to avoid hanging in a blocking flush
+    # when iterating on a dataloader with `num_workers` > 0
+    # TODO(rcadene): remove `gc.collect` when rerun version 0.16 is out, which includes a fix
+    gc.collect()
+
+    if mode == "distant":
+        rr.serve(open_browser=False, web_port=web_port, ws_port=ws_port)
+
+    logging.info("Logging to Rerun")
+
+    for batch in tqdm.tqdm(dataloader, total=len(dataloader)):
+        # iterate over the batch
+        for i in range(len(batch["index"])):
+            rr.set_time_sequence("frame_index", batch["frame_index"][i].item())
+            rr.set_time_seconds("timestamp", batch["timestamp"][i].item())
+
+            # display each camera image
+            for key in dataset.meta.camera_keys:
+                # TODO(rcadene): add `.compress()`? is it lossless?
+                rr.log(key, rr.Image(to_hwc_uint8_numpy(batch[key][i])))
+
+            # display each dimension of 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.Scalar(val.item()))
+
+            # display each dimension of observed state space (e.g. agent position in joint space)
+            if "observation.state" in batch:
+                for dim_idx, val in enumerate(batch["observation.state"][i]):
+                    rr.log(f"state/{dim_idx}", rr.Scalar(val.item()))
+
+            if "next.done" in batch:
+                rr.log("next.done", rr.Scalar(batch["next.done"][i].item()))
+
+            if "next.reward" in batch:
+                rr.log("next.reward", rr.Scalar(batch["next.reward"][i].item()))
+
+            if "next.success" in batch:
+                rr.log("next.success", rr.Scalar(batch["next.success"][i].item()))
+
+    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"
+        rr.save(rrd_path)
+        return rrd_path
+
+    elif mode == "distant":
+        # stop the process from exiting since it is serving the websocket connection
+        try:
+            while True:
+                time.sleep(1)
+        except KeyboardInterrupt:
+            print("Ctrl-C received. Exiting.")
+
+
+def main():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--repo-id",
+        type=str,
+        required=True,
+        help="Name of hugging face repository containing a LeRobotDataset dataset (e.g. `lerobot/pusht`).",
+    )
+    parser.add_argument(
+        "--episode-index",
+        type=int,
+        required=True,
+        help="Episode to visualize.",
+    )
+    parser.add_argument(
+        "--root",
+        type=Path,
+        default=None,
+        help="Root directory for the dataset stored locally (e.g. `--root data`). By default, the dataset will be loaded from hugging face cache folder, or downloaded from the hub if available.",
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=Path,
+        default=None,
+        help="Directory path to write a .rrd file when `--save 1` is set.",
+    )
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=32,
+        help="Batch size loaded by DataLoader.",
+    )
+    parser.add_argument(
+        "--num-workers",
+        type=int,
+        default=4,
+        help="Number of processes of Dataloader for loading the data.",
+    )
+    parser.add_argument(
+        "--mode",
+        type=str,
+        default="local",
+        help=(
+            "Mode of viewing between 'local' or 'distant'. "
+            "'local' requires data to be on a local machine. It spawns a viewer to visualize the data locally. "
+            "'distant' creates a server on the distant machine where the data is stored. "
+            "Visualize the data by connecting to the server with `rerun ws://localhost:PORT` on the local machine."
+        ),
+    )
+    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,
+        default=9087,
+        help="Web socket port for rerun.io when `--mode distant` is set.",
+    )
+    parser.add_argument(
+        "--save",
+        type=int,
+        default=0,
+        help=(
+            "Save a .rrd file in the directory provided by `--output-dir`. "
+            "It also deactivates the spawning of a viewer. "
+            "Visualize the data by running `rerun path/to/file.rrd` on your local machine."
+        ),
+    )
+
+    parser.add_argument(
+        "--tolerance-s",
+        type=float,
+        default=1e-4,
+        help=(
+            "Tolerance in seconds used to ensure data timestamps respect the dataset fps value"
+            "This is argument passed to the constructor of LeRobotDataset and maps to its tolerance_s constructor argument"
+            "If not given, defaults to 1e-4."
+        ),
+    )
+
+    args = parser.parse_args()
+    kwargs = vars(args)
+    repo_id = kwargs.pop("repo_id")
+    root = kwargs.pop("root")
+    tolerance_s = kwargs.pop("tolerance_s")
+
+    logging.info("Loading dataset")
+    dataset = LeRobotDataset(repo_id, root=root, tolerance_s=tolerance_s)
+
+    visualize_dataset(dataset, **vars(args))
+
+
+if __name__ == "__main__":
+    main()

lerobot/scripts/visualize_dataset_html.py

@@ -0,0 +1,479 @@
+#!/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.
+""" Visualize data of **all** frames of any episode of a dataset of type LeRobotDataset.
+
+Note: The last frame of the episode doesnt always correspond to a final state.
+That's because our datasets are composed of transition from state to state up to
+the antepenultimate state associated to the ultimate action to arrive in the final state.
+However, there might not be a transition from a final state to another state.
+
+Note: This script aims to visualize the data used to train the neural networks.
+~What you see is what you get~. When visualizing image modality, it is often expected to observe
+lossly compression artifacts since these images have been decoded from compressed mp4 videos to
+save disk space. The compression factor applied has been tuned to not affect success rate.
+
+Example of usage:
+
+- Visualize data stored on a local machine:
+```bash
+local$ python lerobot/scripts/visualize_dataset_html.py \
+    --repo-id lerobot/pusht
+
+local$ open http://localhost:9090
+```
+
+- Visualize data stored on a distant machine with a local viewer:
+```bash
+distant$ python lerobot/scripts/visualize_dataset_html.py \
+    --repo-id lerobot/pusht
+
+local$ ssh -L 9090:localhost:9090 distant  # create a ssh tunnel
+local$ open http://localhost:9090
+```
+
+- Select episodes to visualize:
+```bash
+python lerobot/scripts/visualize_dataset_html.py \
+    --repo-id lerobot/pusht \
+    --episodes 7 3 5 1 4
+```
+"""
+
+import argparse
+import csv
+import json
+import logging
+import re
+import shutil
+import tempfile
+from io import StringIO
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import requests
+from flask import Flask, redirect, render_template, request, url_for
+
+from lerobot import available_datasets
+from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.common.datasets.utils import IterableNamespace
+from lerobot.common.utils.utils import init_logging
+
+
+def run_server(
+    dataset: LeRobotDataset | IterableNamespace | None,
+    episodes: list[int] | None,
+    host: str,
+    port: str,
+    static_folder: Path,
+    template_folder: Path,
+):
+    app = Flask(__name__, static_folder=static_folder.resolve(), template_folder=template_folder.resolve())
+    app.config["SEND_FILE_MAX_AGE_DEFAULT"] = 0  # specifying not to cache
+
+    @app.route("/")
+    def hommepage(dataset=dataset):
+        if dataset:
+            dataset_namespace, dataset_name = dataset.repo_id.split("/")
+            return redirect(
+                url_for(
+                    "show_episode",
+                    dataset_namespace=dataset_namespace,
+                    dataset_name=dataset_name,
+                    episode_id=0,
+                )
+            )
+
+        dataset_param, episode_param = None, None
+        all_params = request.args
+        if "dataset" in all_params:
+            dataset_param = all_params["dataset"]
+        if "episode" in all_params:
+            episode_param = int(all_params["episode"])
+
+        if dataset_param:
+            dataset_namespace, dataset_name = dataset_param.split("/")
+            return redirect(
+                url_for(
+                    "show_episode",
+                    dataset_namespace=dataset_namespace,
+                    dataset_name=dataset_name,
+                    episode_id=episode_param if episode_param is not None else 0,
+                )
+            )
+
+        featured_datasets = [
+            "lerobot/aloha_static_cups_open",
+            "lerobot/columbia_cairlab_pusht_real",
+            "lerobot/taco_play",
+        ]
+        return render_template(
+            "visualize_dataset_homepage.html",
+            featured_datasets=featured_datasets,
+            lerobot_datasets=available_datasets,
+        )
+
+    @app.route("/<string:dataset_namespace>/<string:dataset_name>")
+    def show_first_episode(dataset_namespace, dataset_name):
+        first_episode_id = 0
+        return redirect(
+            url_for(
+                "show_episode",
+                dataset_namespace=dataset_namespace,
+                dataset_name=dataset_name,
+                episode_id=first_episode_id,
+            )
+        )
+
+    @app.route("/<string:dataset_namespace>/<string:dataset_name>/episode_<int:episode_id>")
+    def show_episode(dataset_namespace, dataset_name, episode_id, dataset=dataset, episodes=episodes):
+        repo_id = f"{dataset_namespace}/{dataset_name}"
+        try:
+            if dataset is None:
+                dataset = get_dataset_info(repo_id)
+        except FileNotFoundError:
+            return (
+                "Make sure to convert your LeRobotDataset to v2 & above. See how to convert your dataset at https://github.com/huggingface/lerobot/pull/461",
+                400,
+            )
+        dataset_version = (
+            str(dataset.meta._version) if isinstance(dataset, LeRobotDataset) else dataset.codebase_version
+        )
+        match = re.search(r"v(\d+)\.", dataset_version)
+        if match:
+            major_version = int(match.group(1))
+            if major_version < 2:
+                return "Make sure to convert your LeRobotDataset to v2 & above."
+
+        episode_data_csv_str, columns, ignored_columns = get_episode_data(dataset, episode_id)
+        dataset_info = {
+            "repo_id": f"{dataset_namespace}/{dataset_name}",
+            "num_samples": dataset.num_frames
+            if isinstance(dataset, LeRobotDataset)
+            else dataset.total_frames,
+            "num_episodes": dataset.num_episodes
+            if isinstance(dataset, LeRobotDataset)
+            else dataset.total_episodes,
+            "fps": dataset.fps,
+        }
+        if isinstance(dataset, LeRobotDataset):
+            video_paths = [
+                dataset.meta.get_video_file_path(episode_id, key) for key in dataset.meta.video_keys
+            ]
+            videos_info = [
+                {"url": url_for("static", filename=video_path), "filename": video_path.parent.name}
+                for video_path in video_paths
+            ]
+            tasks = dataset.meta.episodes[episode_id]["tasks"]
+        else:
+            video_keys = [key for key, ft in dataset.features.items() if ft["dtype"] == "video"]
+            videos_info = [
+                {
+                    "url": f"https://huggingface.co/datasets/{repo_id}/resolve/main/"
+                    + dataset.video_path.format(
+                        episode_chunk=int(episode_id) // dataset.chunks_size,
+                        video_key=video_key,
+                        episode_index=episode_id,
+                    ),
+                    "filename": video_key,
+                }
+                for video_key in video_keys
+            ]
+
+            response = requests.get(
+                f"https://huggingface.co/datasets/{repo_id}/resolve/main/meta/episodes.jsonl", timeout=5
+            )
+            response.raise_for_status()
+            # Split into lines and parse each line as JSON
+            tasks_jsonl = [json.loads(line) for line in response.text.splitlines() if line.strip()]
+
+            filtered_tasks_jsonl = [row for row in tasks_jsonl if row["episode_index"] == episode_id]
+            tasks = filtered_tasks_jsonl[0]["tasks"]
+
+        videos_info[0]["language_instruction"] = tasks
+
+        if episodes is None:
+            episodes = list(
+                range(dataset.num_episodes if isinstance(dataset, LeRobotDataset) else dataset.total_episodes)
+            )
+
+        return render_template(
+            "visualize_dataset_template.html",
+            episode_id=episode_id,
+            episodes=episodes,
+            dataset_info=dataset_info,
+            videos_info=videos_info,
+            episode_data_csv_str=episode_data_csv_str,
+            columns=columns,
+            ignored_columns=ignored_columns,
+        )
+
+    app.run(host=host, port=port)
+
+
+def get_ep_csv_fname(episode_id: int):
+    ep_csv_fname = f"episode_{episode_id}.csv"
+    return ep_csv_fname
+
+
+def get_episode_data(dataset: LeRobotDataset | IterableNamespace, episode_index):
+    """Get a csv str containing timeseries data of an episode (e.g. state and action).
+    This file will be loaded by Dygraph javascript to plot data in real time."""
+    columns = []
+
+    selected_columns = [col for col, ft in dataset.features.items() if ft["dtype"] in ["float32", "int32"]]
+    selected_columns.remove("timestamp")
+
+    ignored_columns = []
+    for column_name in selected_columns:
+        shape = dataset.features[column_name]["shape"]
+        shape_dim = len(shape)
+        if shape_dim > 1:
+            selected_columns.remove(column_name)
+            ignored_columns.append(column_name)
+
+    # init header of csv with state and action names
+    header = ["timestamp"]
+
+    for column_name in selected_columns:
+        dim_state = (
+            dataset.meta.shapes[column_name][0]
+            if isinstance(dataset, LeRobotDataset)
+            else dataset.features[column_name].shape[0]
+        )
+
+        if "names" in dataset.features[column_name] and dataset.features[column_name]["names"]:
+            column_names = dataset.features[column_name]["names"]
+            while not isinstance(column_names, list):
+                column_names = list(column_names.values())[0]
+        else:
+            column_names = [f"{column_name}_{i}" for i in range(dim_state)]
+        columns.append({"key": column_name, "value": column_names})
+
+        header += column_names
+
+    selected_columns.insert(0, "timestamp")
+
+    if isinstance(dataset, LeRobotDataset):
+        from_idx = dataset.episode_data_index["from"][episode_index]
+        to_idx = dataset.episode_data_index["to"][episode_index]
+        data = (
+            dataset.hf_dataset.select(range(from_idx, to_idx))
+            .select_columns(selected_columns)
+            .with_format("pandas")
+        )
+    else:
+        repo_id = dataset.repo_id
+
+        url = f"https://huggingface.co/datasets/{repo_id}/resolve/main/" + dataset.data_path.format(
+            episode_chunk=int(episode_index) // dataset.chunks_size, episode_index=episode_index
+        )
+        df = pd.read_parquet(url)
+        data = df[selected_columns]  # Select specific columns
+
+    rows = np.hstack(
+        (
+            np.expand_dims(data["timestamp"], axis=1),
+            *[np.vstack(data[col]) for col in selected_columns[1:]],
+        )
+    ).tolist()
+
+    # Convert data to CSV string
+    csv_buffer = StringIO()
+    csv_writer = csv.writer(csv_buffer)
+    # Write header
+    csv_writer.writerow(header)
+    # Write data rows
+    csv_writer.writerows(rows)
+    csv_string = csv_buffer.getvalue()
+
+    return csv_string, columns, ignored_columns
+
+
+def get_episode_video_paths(dataset: LeRobotDataset, ep_index: int) -> list[str]:
+    # get first frame of episode (hack to get video_path of the episode)
+    first_frame_idx = dataset.episode_data_index["from"][ep_index].item()
+    return [
+        dataset.hf_dataset.select_columns(key)[first_frame_idx][key]["path"]
+        for key in dataset.meta.video_keys
+    ]
+
+
+def get_episode_language_instruction(dataset: LeRobotDataset, ep_index: int) -> list[str]:
+    # check if the dataset has language instructions
+    if "language_instruction" not in dataset.features:
+        return None
+
+    # get first frame index
+    first_frame_idx = dataset.episode_data_index["from"][ep_index].item()
+
+    language_instruction = dataset.hf_dataset[first_frame_idx]["language_instruction"]
+    # TODO (michel-aractingi) hack to get the sentence, some strings in openx are badly stored
+    # with the tf.tensor appearing in the string
+    return language_instruction.removeprefix("tf.Tensor(b'").removesuffix("', shape=(), dtype=string)")
+
+
+def get_dataset_info(repo_id: str) -> IterableNamespace:
+    response = requests.get(
+        f"https://huggingface.co/datasets/{repo_id}/resolve/main/meta/info.json", timeout=5
+    )
+    response.raise_for_status()  # Raises an HTTPError for bad responses
+    dataset_info = response.json()
+    dataset_info["repo_id"] = repo_id
+    return IterableNamespace(dataset_info)
+
+
+def visualize_dataset_html(
+    dataset: LeRobotDataset | None,
+    episodes: list[int] | None = None,
+    output_dir: Path | None = None,
+    serve: bool = True,
+    host: str = "127.0.0.1",
+    port: int = 9090,
+    force_override: bool = False,
+) -> Path | None:
+    init_logging()
+
+    template_dir = Path(__file__).resolve().parent.parent / "templates"
+
+    if output_dir is None:
+        # Create a temporary directory that will be automatically cleaned up
+        output_dir = tempfile.mkdtemp(prefix="lerobot_visualize_dataset_")
+
+    output_dir = Path(output_dir)
+    if output_dir.exists():
+        if force_override:
+            shutil.rmtree(output_dir)
+        else:
+            logging.info(f"Output directory already exists. Loading from it: '{output_dir}'")
+
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    static_dir = output_dir / "static"
+    static_dir.mkdir(parents=True, exist_ok=True)
+
+    if dataset is None:
+        if serve:
+            run_server(
+                dataset=None,
+                episodes=None,
+                host=host,
+                port=port,
+                static_folder=static_dir,
+                template_folder=template_dir,
+            )
+    else:
+        # Create a simlink from the dataset video folder containing mp4 files to the output directory
+        # so that the http server can get access to the mp4 files.
+        if isinstance(dataset, LeRobotDataset):
+            ln_videos_dir = static_dir / "videos"
+            if not ln_videos_dir.exists():
+                ln_videos_dir.symlink_to((dataset.root / "videos").resolve())
+
+        if serve:
+            run_server(dataset, episodes, host, port, static_dir, template_dir)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--repo-id",
+        type=str,
+        default=None,
+        help="Name of hugging face repositery containing a LeRobotDataset dataset (e.g. `lerobot/pusht` for https://huggingface.co/datasets/lerobot/pusht).",
+    )
+    parser.add_argument(
+        "--root",
+        type=Path,
+        default=None,
+        help="Root directory for a dataset stored locally (e.g. `--root data`). By default, the dataset will be loaded from hugging face cache folder, or downloaded from the hub if available.",
+    )
+    parser.add_argument(
+        "--load-from-hf-hub",
+        type=int,
+        default=0,
+        help="Load videos and parquet files from HF Hub rather than local system.",
+    )
+    parser.add_argument(
+        "--episodes",
+        type=int,
+        nargs="*",
+        default=None,
+        help="Episode indices to visualize (e.g. `0 1 5 6` to load episodes of index 0, 1, 5 and 6). By default loads all episodes.",
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=Path,
+        default=None,
+        help="Directory path to write html files and kickoff a web server. By default write them to 'outputs/visualize_dataset/REPO_ID'.",
+    )
+    parser.add_argument(
+        "--serve",
+        type=int,
+        default=1,
+        help="Launch web server.",
+    )
+    parser.add_argument(
+        "--host",
+        type=str,
+        default="127.0.0.1",
+        help="Web host used by the http server.",
+    )
+    parser.add_argument(
+        "--port",
+        type=int,
+        default=9090,
+        help="Web port used by the http server.",
+    )
+    parser.add_argument(
+        "--force-override",
+        type=int,
+        default=0,
+        help="Delete the output directory if it exists already.",
+    )
+
+    parser.add_argument(
+        "--tolerance-s",
+        type=float,
+        default=1e-4,
+        help=(
+            "Tolerance in seconds used to ensure data timestamps respect the dataset fps value"
+            "This is argument passed to the constructor of LeRobotDataset and maps to its tolerance_s constructor argument"
+            "If not given, defaults to 1e-4."
+        ),
+    )
+
+    args = parser.parse_args()
+    kwargs = vars(args)
+    repo_id = kwargs.pop("repo_id")
+    load_from_hf_hub = kwargs.pop("load_from_hf_hub")
+    root = kwargs.pop("root")
+    tolerance_s = kwargs.pop("tolerance_s")
+
+    dataset = None
+    if repo_id:
+        dataset = (
+            LeRobotDataset(repo_id, root=root, tolerance_s=tolerance_s)
+            if not load_from_hf_hub
+            else get_dataset_info(repo_id)
+        )
+
+    visualize_dataset_html(dataset, **vars(args))
+
+
+if __name__ == "__main__":
+    main()

lerobot/scripts/visualize_image_transforms.py

@@ -0,0 +1,130 @@
+#!/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.
+""" Visualize effects of image transforms for a given configuration.
+
+This script will generate examples of transformed images as they are output by LeRobot dataset.
+Additionally, each individual transform can be visualized separately as well as examples of combined transforms
+
+Example:
+```bash
+python lerobot/scripts/visualize_image_transforms.py \
+    --repo_id=lerobot/pusht \
+    --episodes='[0]' \
+    --image_transforms.enable=True
+```
+"""
+
+import logging
+from copy import deepcopy
+from dataclasses import replace
+from pathlib import Path
+
+import draccus
+from torchvision.transforms import ToPILImage
+
+from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.common.datasets.transforms import (
+    ImageTransforms,
+    ImageTransformsConfig,
+    make_transform_from_config,
+)
+from lerobot.configs.default import DatasetConfig
+
+OUTPUT_DIR = Path("outputs/image_transforms")
+to_pil = ToPILImage()
+
+
+def save_all_transforms(cfg: ImageTransformsConfig, original_frame, output_dir, n_examples):
+    output_dir_all = output_dir / "all"
+    output_dir_all.mkdir(parents=True, exist_ok=True)
+
+    tfs = ImageTransforms(cfg)
+    for i in range(1, n_examples + 1):
+        transformed_frame = tfs(original_frame)
+        to_pil(transformed_frame).save(output_dir_all / f"{i}.png", quality=100)
+
+    print("Combined transforms examples saved to:")
+    print(f"    {output_dir_all}")
+
+
+def save_each_transform(cfg: ImageTransformsConfig, original_frame, output_dir, n_examples):
+    if not cfg.enable:
+        logging.warning(
+            "No single transforms will be saved, because `image_transforms.enable=False`. To enable, set `enable` to True in `ImageTransformsConfig` or in the command line with `--image_transforms.enable=True`."
+        )
+        return
+
+    print("Individual transforms examples saved to:")
+    for tf_name, tf_cfg in cfg.tfs.items():
+        # Apply a few transformation with random value in min_max range
+        output_dir_single = output_dir / tf_name
+        output_dir_single.mkdir(parents=True, exist_ok=True)
+
+        tf = make_transform_from_config(tf_cfg)
+        for i in range(1, n_examples + 1):
+            transformed_frame = tf(original_frame)
+            to_pil(transformed_frame).save(output_dir_single / f"{i}.png", quality=100)
+
+        # Apply min, max, average transformations
+        tf_cfg_kwgs_min = deepcopy(tf_cfg.kwargs)
+        tf_cfg_kwgs_max = deepcopy(tf_cfg.kwargs)
+        tf_cfg_kwgs_avg = deepcopy(tf_cfg.kwargs)
+
+        for key, (min_, max_) in tf_cfg.kwargs.items():
+            avg = (min_ + max_) / 2
+            tf_cfg_kwgs_min[key] = [min_, min_]
+            tf_cfg_kwgs_max[key] = [max_, max_]
+            tf_cfg_kwgs_avg[key] = [avg, avg]
+
+        tf_min = make_transform_from_config(replace(tf_cfg, **{"kwargs": tf_cfg_kwgs_min}))
+        tf_max = make_transform_from_config(replace(tf_cfg, **{"kwargs": tf_cfg_kwgs_max}))
+        tf_avg = make_transform_from_config(replace(tf_cfg, **{"kwargs": tf_cfg_kwgs_avg}))
+
+        tf_frame_min = tf_min(original_frame)
+        tf_frame_max = tf_max(original_frame)
+        tf_frame_avg = tf_avg(original_frame)
+
+        to_pil(tf_frame_min).save(output_dir_single / "min.png", quality=100)
+        to_pil(tf_frame_max).save(output_dir_single / "max.png", quality=100)
+        to_pil(tf_frame_avg).save(output_dir_single / "mean.png", quality=100)
+
+        print(f"    {output_dir_single}")
+
+
+@draccus.wrap()
+def visualize_image_transforms(cfg: DatasetConfig, output_dir: Path = OUTPUT_DIR, n_examples: int = 5):
+    dataset = LeRobotDataset(
+        repo_id=cfg.repo_id,
+        episodes=cfg.episodes,
+        revision=cfg.revision,
+        video_backend=cfg.video_backend,
+    )
+
+    output_dir = output_dir / cfg.repo_id.split("/")[-1]
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    # Get 1st frame from 1st camera of 1st episode
+    original_frame = dataset[0][dataset.meta.camera_keys[0]]
+    to_pil(original_frame).save(output_dir / "original_frame.png", quality=100)
+    print("\nOriginal frame saved to:")
+    print(f"    {output_dir / 'original_frame.png'}.")
+
+    save_all_transforms(cfg.image_transforms, original_frame, output_dir, n_examples)
+    save_each_transform(cfg.image_transforms, original_frame, output_dir, n_examples)
+
+
+if __name__ == "__main__":
+    visualize_image_transforms()

lerobot/templates/visualize_dataset_homepage.html

@@ -0,0 +1,68 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>Interactive Video Background Page</title>
+    <script src="https://cdn.tailwindcss.com"></script>
+    <script defer src="https://cdn.jsdelivr.net/npm/alpinejs@3.x.x/dist/cdn.min.js"></script>
+</head>
+<body class="h-screen overflow-hidden font-mono text-white" x-data="{
+    inputValue: '',
+    navigateToDataset() {
+        const trimmedValue = this.inputValue.trim();
+        if (trimmedValue) {
+            window.location.href = `/${trimmedValue}`;
+        }
+    }
+}">
+    <div class="fixed inset-0 w-full h-full overflow-hidden">
+        <video class="absolute min-w-full min-h-full w-auto h-auto top-1/2 left-1/2 transform -translate-x-1/2 -translate-y-1/2" autoplay muted loop>
+            <source src="https://huggingface.co/datasets/cadene/koch_bimanual_folding/resolve/v1.6/videos/observation.images.phone_episode_000037.mp4" type="video/mp4">
+            Your browser does not support HTML5 video.
+        </video>
+    </div>
+    <div class="fixed inset-0 bg-black bg-opacity-80"></div>
+    <div class="relative z-10 flex flex-col items-center justify-center h-screen">
+        <div class="text-center mb-8">
+            <h1 class="text-4xl font-bold mb-4">LeRobot Dataset Visualizer</h1>
+
+            <a href="https://x.com/RemiCadene/status/1825455895561859185" target="_blank" rel="noopener noreferrer" class="underline">create & train your own robots</a>
+
+            <p class="text-xl mb-4"></p>
+            <div class="text-left inline-block">
+                <h3 class="font-semibold mb-2 mt-4">Example Datasets:</h3>
+                <ul class="list-disc list-inside">
+                    {% for dataset in featured_datasets %}
+                        <li><a href="/{{ dataset }}" class="text-blue-300 hover:text-blue-100 hover:underline">{{ dataset }}</a></li>
+                    {% endfor %}
+                </ul>
+            </div>
+        </div>
+        <div class="flex w-full max-w-lg px-4 mb-4">
+            <input
+                type="text"
+                x-model="inputValue"
+                @keyup.enter="navigateToDataset"
+                placeholder="enter dataset id (ex: lerobot/droid_100)"
+                class="flex-grow px-4 py-2 rounded-l bg-white bg-opacity-20 text-white placeholder-gray-300 focus:outline-none focus:ring-2 focus:ring-blue-300"
+            >
+            <button
+                @click="navigateToDataset"
+                class="px-4 py-2 bg-blue-500 text-white rounded-r hover:bg-blue-600 focus:outline-none focus:ring-2 focus:ring-blue-300"
+            >
+                Go
+            </button>
+        </div>
+
+        <details class="mt-4 max-w-full px-4">
+            <summary>More example datasets</summary>
+            <ul class="list-disc list-inside max-h-28 overflow-y-auto break-all">
+                {% for dataset in lerobot_datasets %}
+                    <li><a href="/{{ dataset }}" class="text-blue-300 hover:text-blue-100 hover:underline">{{ dataset }}</a></li>
+                {% endfor %}
+            </ul>
+        </details>
+    </div>
+</body>
+</html>