From 0da98afd639dc3e4977f10591ee940ab833e1361 Mon Sep 17 00:00:00 2001 From: Khalil Meftah Date: Tue, 30 Jun 2026 17:17:50 +0200 Subject: [PATCH 1/2] Feat(robot): add MIT control mode to ReBot (#3778) * fix(config): update joint limits for RebotB601Follower and RebotArm102Leader * feat(config): add MIT control mode ReBot - Add configurable arm control mode (mit default, pos_vel fallback) with tunable mit_kp / mit_kd - Add optional gripper control mode (force_pos default, mit optional) with gripper_mit_kp / gripper_mit_kd - Update tests for MIT arm routing, gripper mode routing, and revised joint limits * fix(robots): restore joint clipping and wrist_yaw fallback in ReBot B601 send_action * feat(robot): increase gripper velocity and torque for rebot arm --- .../bi_rebot_b601_follower.py | 12 +++++ .../config_rebot_b601_follower.py | 29 ++++++++--- .../rebot_b601_follower.py | 48 ++++++++++++++----- .../config_rebot_102_leader.py | 2 +- tests/robots/test_rebot_b601_follower.py | 23 +++++++-- 5 files changed, 92 insertions(+), 22 deletions(-) diff --git a/src/lerobot/robots/bi_rebot_b601_follower/bi_rebot_b601_follower.py b/src/lerobot/robots/bi_rebot_b601_follower/bi_rebot_b601_follower.py index c320cee8b..7780ac0fe 100644 --- a/src/lerobot/robots/bi_rebot_b601_follower/bi_rebot_b601_follower.py +++ b/src/lerobot/robots/bi_rebot_b601_follower/bi_rebot_b601_follower.py @@ -65,7 +65,13 @@ class BiRebotB601Follower(BimanualMixin, Robot): cameras=left_arm_cameras, motor_can_ids=config.left_arm_config.motor_can_ids, pos_vel_velocity=config.left_arm_config.pos_vel_velocity, + control_mode=config.left_arm_config.control_mode, + mit_kp=config.left_arm_config.mit_kp, + mit_kd=config.left_arm_config.mit_kd, + gripper_control_mode=config.left_arm_config.gripper_control_mode, gripper_torque_ratio=config.left_arm_config.gripper_torque_ratio, + gripper_mit_kp=config.left_arm_config.gripper_mit_kp, + gripper_mit_kd=config.left_arm_config.gripper_mit_kd, joint_limits=config.left_arm_config.joint_limits, ) @@ -80,7 +86,13 @@ class BiRebotB601Follower(BimanualMixin, Robot): cameras=config.right_arm_config.cameras, motor_can_ids=config.right_arm_config.motor_can_ids, pos_vel_velocity=config.right_arm_config.pos_vel_velocity, + control_mode=config.right_arm_config.control_mode, + mit_kp=config.right_arm_config.mit_kp, + mit_kd=config.right_arm_config.mit_kd, + gripper_control_mode=config.right_arm_config.gripper_control_mode, gripper_torque_ratio=config.right_arm_config.gripper_torque_ratio, + gripper_mit_kp=config.right_arm_config.gripper_mit_kp, + gripper_mit_kd=config.right_arm_config.gripper_mit_kd, joint_limits=config.right_arm_config.joint_limits, ) diff --git a/src/lerobot/robots/rebot_b601_follower/config_rebot_b601_follower.py b/src/lerobot/robots/rebot_b601_follower/config_rebot_b601_follower.py index 096548afb..88af07e59 100644 --- a/src/lerobot/robots/rebot_b601_follower/config_rebot_b601_follower.py +++ b/src/lerobot/robots/rebot_b601_follower/config_rebot_b601_follower.py @@ -65,18 +65,33 @@ class RebotB601FollowerConfig: } ) - # Target velocity for joints running in POS_VEL mode, in degrees/s. A scalar is - # applied to every joint; a list provides one value per joint (in motor order). - pos_vel_velocity: float | list[float] = field(default_factory=lambda: [150.0] * 7) + # Max speed (deg/s) per joint for POS_VEL arms and FORCE_POS gripper (motor order). + pos_vel_velocity: float | list[float] = field( + default_factory=lambda: [150.0, 150.0, 150.0, 150.0, 150.0, 150.0, 900.0] + ) - # Torque/current ratio for the gripper's FORCE_POS mode, in range [0, 1]. - gripper_torque_ratio: float = 0.1 + # Arm control: "mit" or "pos_vel". + control_mode: str = "mit" + + # MIT kp/kd per arm joint (motor order). Unused when control_mode="pos_vel". + mit_kp: float | list[float] = field(default_factory=lambda: [45.0, 45.0, 45.0, 8.0, 9.0, 8.0, 8.0]) + mit_kd: float | list[float] = field(default_factory=lambda: [12.0, 12.0, 12.0, 1.0, 1.0, 1.0, 1.0]) + + # Gripper control: "force_pos" or "mit". + gripper_control_mode: str = "force_pos" + + # FORCE_POS only: max grip force, in [0, 1]. + gripper_torque_ratio: float = 0.07 + + # MIT only. + gripper_mit_kp: float = 8.0 + gripper_mit_kd: float = 0.3 # Soft joint limits (degrees). These are clipped against on every action. joint_limits: dict[str, tuple[float, float]] = field( default_factory=lambda: { - "shoulder_pan": (-145.0, 145.0), - "shoulder_lift": (-170.0, 1.0), + "shoulder_pan": (-150.0, 150.0), + "shoulder_lift": (-200.0, 1.0), "elbow_flex": (-200.0, 1.0), "wrist_flex": (-80.0, 90.0), "wrist_yaw": (-90.0, 90.0), diff --git a/src/lerobot/robots/rebot_b601_follower/rebot_b601_follower.py b/src/lerobot/robots/rebot_b601_follower/rebot_b601_follower.py index bf989702b..bff0b86c1 100644 --- a/src/lerobot/robots/rebot_b601_follower/rebot_b601_follower.py +++ b/src/lerobot/robots/rebot_b601_follower/rebot_b601_follower.py @@ -174,11 +174,25 @@ class RebotB601Follower(Robot): print(f"Calibration saved to {self.calibration_fpath}") def configure(self) -> None: + if self.config.control_mode not in ("pos_vel", "mit"): + raise ValueError( + f"Unsupported control_mode '{self.config.control_mode}'. Use 'pos_vel' or 'mit'." + ) + if self.config.gripper_control_mode not in ("force_pos", "mit"): + raise ValueError( + f"Unsupported gripper_control_mode '{self.config.gripper_control_mode}'. " + "Use 'force_pos' or 'mit'." + ) + use_mit = self.config.control_mode == "mit" + gripper_use_mit = self.config.gripper_control_mode == "mit" self.bus.enable_all() for motor_name, motor in self.motors.items(): - target_mode = ( - MotorBridgeMode.FORCE_POS if motor_name == GRIPPER_MOTOR else MotorBridgeMode.POS_VEL - ) + if motor_name == GRIPPER_MOTOR: + target_mode = MotorBridgeMode.MIT if gripper_use_mit else MotorBridgeMode.FORCE_POS + elif use_mit: + target_mode = MotorBridgeMode.MIT + else: + target_mode = MotorBridgeMode.POS_VEL for attempt in range(_ENSURE_MODE_RETRIES + 1): try: motor.ensure_mode(target_mode) @@ -264,22 +278,34 @@ class RebotB601Follower(Robot): goal_present_pos = {key: (g, present_pos.get(key, g)) for key, g in goal_pos.items()} goal_pos = ensure_safe_goal_position(goal_present_pos, self.config.max_relative_target) + use_mit = self.config.control_mode == "mit" for motor_name, position_deg in goal_pos.items(): motor = self.motors.get(motor_name) if motor is None: continue idx = self.motor_names.index(motor_name) - vel_deg_s = ( - self.config.pos_vel_velocity[idx] - if isinstance(self.config.pos_vel_velocity, list) - else self.config.pos_vel_velocity - ) pos_rad = math.radians(position_deg) - vel_rad = math.radians(vel_deg_s) if motor_name == GRIPPER_MOTOR: - motor.send_force_pos(pos_rad, vel_rad, self.config.gripper_torque_ratio) + if self.config.gripper_control_mode == "mit": + motor.send_mit(pos_rad, 0.0, self.config.gripper_mit_kp, self.config.gripper_mit_kd, 0.0) + else: + vel_deg_s = ( + self.config.pos_vel_velocity[idx] + if isinstance(self.config.pos_vel_velocity, list) + else self.config.pos_vel_velocity + ) + motor.send_force_pos(pos_rad, math.radians(vel_deg_s), self.config.gripper_torque_ratio) + elif use_mit: + kp = self.config.mit_kp[idx] if isinstance(self.config.mit_kp, list) else self.config.mit_kp + kd = self.config.mit_kd[idx] if isinstance(self.config.mit_kd, list) else self.config.mit_kd + motor.send_mit(pos_rad, 0.0, kp, kd, 0.0) else: - motor.send_pos_vel(pos_rad, vel_rad) + vel_deg_s = ( + self.config.pos_vel_velocity[idx] + if isinstance(self.config.pos_vel_velocity, list) + else self.config.pos_vel_velocity + ) + motor.send_pos_vel(pos_rad, math.radians(vel_deg_s)) return {f"{motor}.pos": val for motor, val in goal_pos.items()} diff --git a/src/lerobot/teleoperators/rebot_102_leader/config_rebot_102_leader.py b/src/lerobot/teleoperators/rebot_102_leader/config_rebot_102_leader.py index d1beea2ed..81a1c2c39 100644 --- a/src/lerobot/teleoperators/rebot_102_leader/config_rebot_102_leader.py +++ b/src/lerobot/teleoperators/rebot_102_leader/config_rebot_102_leader.py @@ -65,7 +65,7 @@ class RebotArm102LeaderConfig: joint_ranges: dict[str, list[int]] = field( default_factory=lambda: { "shoulder_pan": [-150, 150], - "shoulder_lift": [-170, 1], + "shoulder_lift": [-200, 1], "elbow_flex": [-200, 1], "wrist_flex": [-80, 90], "wrist_yaw": [-90, 90], diff --git a/tests/robots/test_rebot_b601_follower.py b/tests/robots/test_rebot_b601_follower.py index 553675be0..0861827be 100644 --- a/tests/robots/test_rebot_b601_follower.py +++ b/tests/robots/test_rebot_b601_follower.py @@ -91,10 +91,11 @@ def test_get_observation_converts_to_degrees(follower): def test_send_action_clips_to_joint_limits(follower): - # shoulder_pan limit is (-145, 145); request beyond the upper bound. + # shoulder_pan limit is (-150, 150); request beyond the upper bound. returned = follower.send_action({"shoulder_pan.pos": 999.0}) - assert returned["shoulder_pan.pos"] == 145.0 - follower.motors["shoulder_pan"].send_pos_vel.assert_called_once() + assert returned["shoulder_pan.pos"] == 150.0 + # Default control_mode is "mit", so arm joints are driven via send_mit. + follower.motors["shoulder_pan"].send_mit.assert_called_once() def test_send_action_routes_gripper_to_force_pos(follower): @@ -103,6 +104,22 @@ def test_send_action_routes_gripper_to_force_pos(follower): follower.motors["gripper"].send_pos_vel.assert_not_called() +def test_gripper_mit_mode_routes_to_send_mit(): + bus_mock = _make_bus_mock() + with ( + patch(f"{_MODULE}.require_package", lambda *a, **kw: None), + patch(f"{_MODULE}.MotorBridgeController") as controller_cls, + patch(f"{_MODULE}.MotorBridgeMode", MagicMock()), + ): + controller_cls.from_dm_serial.return_value = bus_mock + cfg = RebotB601FollowerRobotConfig(port="/dev/null", gripper_control_mode="mit") + robot = RebotB601Follower(cfg) + robot.connect(calibrate=False) + robot.send_action({"gripper.pos": -10.0}) + robot.motors["gripper"].send_mit.assert_called_once() + robot.motors["gripper"].send_force_pos.assert_not_called() + + def test_bimanual_prefixes_features(): with patch(f"{_MODULE}.require_package", lambda *a, **kw: None): cfg = BiRebotB601FollowerConfig( From 8414188db0b178b947985a7a9a91314708837315 Mon Sep 17 00:00:00 2001 From: Caroline Pascal Date: Tue, 30 Jun 2026 20:21:06 +0200 Subject: [PATCH 2/2] fix(datasets dependency): removing datasets dependency in pretrained.py (#3897) --- src/lerobot/policies/pretrained.py | 56 ++++++++++++++-------------- src/lerobot/scripts/lerobot_train.py | 4 +- 2 files changed, 29 insertions(+), 31 deletions(-) diff --git a/src/lerobot/policies/pretrained.py b/src/lerobot/policies/pretrained.py index aea5f1b08..702569b8c 100644 --- a/src/lerobot/policies/pretrained.py +++ b/src/lerobot/policies/pretrained.py @@ -11,6 +11,8 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations + import abc import builtins import dataclasses @@ -19,7 +21,7 @@ import os from importlib.resources import files from pathlib import Path from tempfile import TemporaryDirectory -from typing import TypedDict, TypeVar, Unpack +from typing import TYPE_CHECKING, TypedDict, TypeVar, Unpack import packaging import safetensors @@ -38,10 +40,13 @@ from .utils import log_model_loading_keys T = TypeVar("T", bound="PreTrainedPolicy") +if TYPE_CHECKING: + from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata + def _build_card_context( cfg: TrainPipelineConfig | None, - dataset_repo_id: str | None, + dataset_meta: LeRobotDatasetMetadata | None, input_features: dict | None, output_features: dict | None, ) -> dict: @@ -72,30 +77,16 @@ def _build_card_context( "lerobot_version": __version__, } - if dataset_repo_id: - dataset_cfg = getattr(cfg, "dataset", None) - try: - from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata - - meta = LeRobotDatasetMetadata( - dataset_repo_id, - root=getattr(dataset_cfg, "root", None), - revision=getattr(dataset_cfg, "revision", None), - ) - context["dataset"] = { - "repo_id": dataset_repo_id, - "episodes": meta.total_episodes, - "frames": meta.total_frames, - "fps": meta.fps, - "tasks": [str(task) for task in meta.tasks.index], - } - context["robot_type"] = meta.robot_type - context["cameras"] = [key.split(".")[-1] for key in meta.camera_keys] - except Exception as e: # noqa: BLE001 — dataset details are optional, never fail the push - logging.warning( - f"Could not load dataset metadata for '{dataset_repo_id}'; those sections will be " - f"omitted from the model card. ({e})" - ) + if dataset_meta is not None: + context["dataset"] = { + "repo_id": dataset_meta.repo_id, + "episodes": dataset_meta.total_episodes, + "frames": dataset_meta.total_frames, + "fps": dataset_meta.fps, + "tasks": [str(task) for task in dataset_meta.tasks.index], + } + context["robot_type"] = dataset_meta.robot_type + context["cameras"] = [key.split(".")[-1] for key in dataset_meta.camera_keys] return context @@ -304,6 +295,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC): cfg: TrainPipelineConfig, peft_model=None, state_dict: dict[str, Tensor] | None = None, + dataset_meta: LeRobotDatasetMetadata | None = None, ): api = HfApi() repo_id = api.create_repo( @@ -325,7 +317,12 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC): self.save_pretrained(saved_path, state_dict=state_dict) card = self.generate_model_card( - cfg.dataset.repo_id, self.config.type, self.config.license, self.config.tags, cfg=cfg + cfg.dataset.repo_id, + self.config.type, + self.config.license, + self.config.tags, + cfg=cfg, + dataset_meta=dataset_meta, ) card.save(str(saved_path / "README.md")) @@ -352,6 +349,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC): license: str | None, tags: list[str] | None, cfg: TrainPipelineConfig | None = None, + dataset_meta: LeRobotDatasetMetadata | None = None, ) -> ModelCard: base_model_mapping = { "smolvla": "lerobot/smolvla_base", @@ -372,7 +370,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC): ) context = _build_card_context( - cfg, dataset_repo_id, self.config.input_features, self.config.output_features + cfg, dataset_meta, self.config.input_features, self.config.output_features ) # Used by the template to pre-fill commands and the "Fine-tuned from" line. context["policy_repo_id"] = getattr(self.config, "repo_id", None) @@ -389,7 +387,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC): self, peft_config=None, peft_cli_overrides: dict | None = None, - ) -> "PreTrainedPolicy": + ) -> PreTrainedPolicy: """ Wrap this policy with PEFT adapters for parameter-efficient fine-tuning. diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index f2a152df9..44c94a1eb 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/src/lerobot/scripts/lerobot_train.py @@ -736,9 +736,9 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): unwrapped_model = accelerator.unwrap_model(policy) # PEFT only applies when training a policy — reward models use the plain path. if not cfg.is_reward_model_training and cfg.policy.use_peft: - unwrapped_model.push_model_to_hub(cfg, peft_model=unwrapped_model) + unwrapped_model.push_model_to_hub(cfg, peft_model=unwrapped_model, dataset_meta=dataset.meta) else: - unwrapped_model.push_model_to_hub(cfg, state_dict=model_state_dict) + unwrapped_model.push_model_to_hub(cfg, state_dict=model_state_dict, dataset_meta=dataset.meta) preprocessor.push_to_hub(active_cfg.repo_id) postprocessor.push_to_hub(active_cfg.repo_id)