mirror of
https://github.com/huggingface/lerobot.git
synced 2026-06-12 05:59:53 +00:00
Compare commits
5 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 099d2cb34e | |||
| b9c4dd3d12 | |||
| aa0d3e6608 | |||
| 350d01b74d | |||
| 41166b39fb |
@@ -57,11 +57,11 @@ The `lerobot-rollout --strategy.type=dagger` mode requires **teleoperators with
|
||||
|
||||
**Compatible teleoperators:**
|
||||
|
||||
- `openarm_mini` - OpenArm Mini
|
||||
- `bi_openarm_mini` - Bimanual OpenArm Mini
|
||||
- `so_leader` - SO100 / SO101 leader arm
|
||||
|
||||
> [!IMPORTANT]
|
||||
> The provided commands default to `bi_openarm_follower` + `openarm_mini`.
|
||||
> The provided commands default to `bi_openarm_follower` + `bi_openarm_mini`.
|
||||
> `so_follower` + `so_leader` configs are also registered and can be used via CLI flags.
|
||||
|
||||
---
|
||||
@@ -104,9 +104,9 @@ lerobot-rollout --strategy.type=dagger \
|
||||
--robot.right_arm_config.port=can0 \
|
||||
--robot.right_arm_config.side=right \
|
||||
--robot.cameras='{left_wrist: {type: opencv, index_or_path: "/dev/video0", width: 1280, height: 720, fps: 30}, right_wrist: {type: opencv, index_or_path: "/dev/video4", width: 1280, height: 720, fps: 30}, base: {type: opencv, index_or_path: "/dev/video2", width: 640, height: 480, fps: 30}}' \
|
||||
--teleop.type=openarm_mini \
|
||||
--teleop.port_left=/dev/ttyACM0 \
|
||||
--teleop.port_right=/dev/ttyACM1 \
|
||||
--teleop.type=bi_openarm_mini \
|
||||
--teleop.left_arm_config.port=/dev/ttyACM0 \
|
||||
--teleop.right_arm_config.port=/dev/ttyACM1 \
|
||||
--policy.path=outputs/pretrain/checkpoints/last/pretrained_model \
|
||||
--dataset.repo_id=your-username/rollout_hil_dataset \
|
||||
--dataset.single_task="Fold the T-shirt properly" \
|
||||
@@ -131,9 +131,9 @@ lerobot-rollout --strategy.type=dagger \
|
||||
--robot.right_arm_config.port=can0 \
|
||||
--robot.right_arm_config.side=right \
|
||||
--robot.cameras='{left_wrist: {type: opencv, index_or_path: "/dev/video0", width: 1280, height: 720, fps: 30}, right_wrist: {type: opencv, index_or_path: "/dev/video4", width: 1280, height: 720, fps: 30}, base: {type: opencv, index_or_path: "/dev/video2", width: 640, height: 480, fps: 30}}' \
|
||||
--teleop.type=openarm_mini \
|
||||
--teleop.port_left=/dev/ttyACM0 \
|
||||
--teleop.port_right=/dev/ttyACM1 \
|
||||
--teleop.type=bi_openarm_mini \
|
||||
--teleop.left_arm_config.port=/dev/ttyACM0 \
|
||||
--teleop.right_arm_config.port=/dev/ttyACM1 \
|
||||
--policy.path=outputs/pretrain/checkpoints/last/pretrained_model \
|
||||
--dataset.repo_id=your-username/rollout_hil_rtc_dataset \
|
||||
--dataset.single_task="Fold the T-shirt properly" \
|
||||
|
||||
@@ -117,7 +117,7 @@ lerobot-rollout \
|
||||
--strategy.num_episodes=20 \
|
||||
--policy.path=outputs/pretrain/checkpoints/last/pretrained_model \
|
||||
--robot.type=bi_openarm_follower \
|
||||
--teleop.type=openarm_mini \
|
||||
--teleop.type=bi_openarm_mini \
|
||||
--dataset.repo_id=${HF_USER}/rollout_hil_data \
|
||||
--dataset.single_task="Fold the T-shirt"
|
||||
```
|
||||
|
||||
+1
-1
@@ -124,7 +124,7 @@ hardware = [
|
||||
"lerobot[deepdiff-dep]",
|
||||
]
|
||||
viz = [
|
||||
"rerun-sdk>=0.24.0,<0.34.0",
|
||||
"rerun-sdk>=0.24.0,<0.27.0",
|
||||
]
|
||||
# ── User-facing composite extras (map to CLI scripts) ─────
|
||||
# lerobot-record, lerobot-replay, lerobot-calibrate, lerobot-teleoperate, etc.
|
||||
|
||||
@@ -30,6 +30,7 @@ class EpisodeAwareSampler:
|
||||
drop_n_first_frames: int = 0,
|
||||
drop_n_last_frames: int = 0,
|
||||
shuffle: bool = False,
|
||||
generator: torch.Generator | None = None,
|
||||
):
|
||||
"""Sampler that optionally incorporates episode boundary information.
|
||||
|
||||
@@ -41,6 +42,10 @@ class EpisodeAwareSampler:
|
||||
drop_n_first_frames: Number of frames to drop from the start of each episode.
|
||||
drop_n_last_frames: Number of frames to drop from the end of each episode.
|
||||
shuffle: Whether to shuffle the indices.
|
||||
generator: Generator used for shuffling. Exposing this attribute (even when None) lets
|
||||
`accelerate` register it as the synchronized RNG in distributed training, so
|
||||
every rank draws the same permutation and batch shards stay disjoint. When
|
||||
None, shuffling falls back to the global torch RNG.
|
||||
"""
|
||||
if drop_n_first_frames < 0:
|
||||
raise ValueError(f"drop_n_first_frames must be >= 0, got {drop_n_first_frames}")
|
||||
@@ -73,10 +78,11 @@ class EpisodeAwareSampler:
|
||||
|
||||
self.indices = indices
|
||||
self.shuffle = shuffle
|
||||
self.generator = generator
|
||||
|
||||
def __iter__(self) -> Iterator[int]:
|
||||
if self.shuffle:
|
||||
for i in torch.randperm(len(self.indices)):
|
||||
for i in torch.randperm(len(self.indices), generator=self.generator):
|
||||
yield self.indices[i]
|
||||
else:
|
||||
for i in self.indices:
|
||||
|
||||
@@ -18,7 +18,8 @@ import logging
|
||||
from functools import cached_property
|
||||
|
||||
from lerobot.types import RobotAction, RobotObservation
|
||||
from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
|
||||
from lerobot.utils.bimanual import BimanualMixin
|
||||
from lerobot.utils.decorators import check_if_not_connected
|
||||
|
||||
from ..openarm_follower import OpenArmFollower, OpenArmFollowerConfig
|
||||
from ..robot import Robot
|
||||
@@ -27,7 +28,7 @@ from .config_bi_openarm_follower import BiOpenArmFollowerConfig
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BiOpenArmFollower(Robot):
|
||||
class BiOpenArmFollower(BimanualMixin, Robot):
|
||||
"""
|
||||
Bimanual OpenArm Follower Arms
|
||||
"""
|
||||
@@ -39,15 +40,17 @@ class BiOpenArmFollower(Robot):
|
||||
super().__init__(config)
|
||||
self.config = config
|
||||
|
||||
# Top-level cameras are distributed evenly: each arm's OpenArmFollower
|
||||
# will only open the cameras assigned to it. Per-arm cameras are used
|
||||
# as fallback when top-level cameras are empty.
|
||||
if config.cameras:
|
||||
left_cameras = config.cameras
|
||||
right_cameras = {}
|
||||
else:
|
||||
left_cameras = config.left_arm_config.cameras
|
||||
right_cameras = config.right_arm_config.cameras
|
||||
# Top-level cameras are opened by `left_arm` for convenience, but their
|
||||
# keys stay unprefixed in observations (tracked via `_top_level_cam_keys`).
|
||||
self._top_level_cam_keys = set(config.cameras)
|
||||
_collisions = self._top_level_cam_keys & set(
|
||||
config.left_arm_config.cameras
|
||||
) | self._top_level_cam_keys & set(config.right_arm_config.cameras)
|
||||
if _collisions:
|
||||
raise ValueError(
|
||||
f"Top-level camera names collide with per-arm camera names: {sorted(_collisions)}"
|
||||
)
|
||||
left_arm_cameras = {**config.left_arm_config.cameras, **config.cameras}
|
||||
|
||||
left_arm_config = OpenArmFollowerConfig(
|
||||
id=f"{config.id}_left" if config.id else None,
|
||||
@@ -56,7 +59,7 @@ class BiOpenArmFollower(Robot):
|
||||
disable_torque_on_disconnect=config.left_arm_config.disable_torque_on_disconnect,
|
||||
use_velocity_and_torque=config.left_arm_config.use_velocity_and_torque,
|
||||
max_relative_target=config.left_arm_config.max_relative_target,
|
||||
cameras=left_cameras,
|
||||
cameras=left_arm_cameras,
|
||||
side=config.left_arm_config.side,
|
||||
can_interface=config.left_arm_config.can_interface,
|
||||
use_can_fd=config.left_arm_config.use_can_fd,
|
||||
@@ -75,7 +78,7 @@ class BiOpenArmFollower(Robot):
|
||||
disable_torque_on_disconnect=config.right_arm_config.disable_torque_on_disconnect,
|
||||
use_velocity_and_torque=config.right_arm_config.use_velocity_and_torque,
|
||||
max_relative_target=config.right_arm_config.max_relative_target,
|
||||
cameras=right_cameras,
|
||||
cameras=config.right_arm_config.cameras,
|
||||
side=config.right_arm_config.side,
|
||||
can_interface=config.right_arm_config.can_interface,
|
||||
use_can_fd=config.right_arm_config.use_can_fd,
|
||||
@@ -95,22 +98,19 @@ class BiOpenArmFollower(Robot):
|
||||
|
||||
@property
|
||||
def _motors_ft(self) -> dict[str, type]:
|
||||
left_arm_motors_ft = self.left_arm._motors_ft
|
||||
right_arm_motors_ft = self.right_arm._motors_ft
|
||||
|
||||
# Right first, then left — matches the teleoperator (OpenArmMini) ordering
|
||||
# and the dataset feature names recorded during data collection.
|
||||
return {
|
||||
**{f"right_{k}": v for k, v in right_arm_motors_ft.items()},
|
||||
**{f"left_{k}": v for k, v in left_arm_motors_ft.items()},
|
||||
**{f"left_{k}": v for k, v in self.left_arm._motors_ft.items()},
|
||||
**{f"right_{k}": v for k, v in self.right_arm._motors_ft.items()},
|
||||
}
|
||||
|
||||
@property
|
||||
def _cameras_ft(self) -> dict[str, tuple]:
|
||||
# Cameras already have unique user-chosen names (e.g. "left_wrist", "base",
|
||||
# "right_wrist"), so we merge them directly — unlike motors which need the
|
||||
# left_/right_ prefix to disambiguate identical per-arm joint names.
|
||||
return {**self.left_arm._cameras_ft, **self.right_arm._cameras_ft}
|
||||
out: dict[str, tuple] = {}
|
||||
for k, v in self.left_arm._cameras_ft.items():
|
||||
out[k if k in self._top_level_cam_keys else f"left_{k}"] = v
|
||||
for k, v in self.right_arm._cameras_ft.items():
|
||||
out[f"right_{k}"] = v
|
||||
return out
|
||||
|
||||
@cached_property
|
||||
def observation_features(self) -> dict[str, type | tuple]:
|
||||
@@ -120,27 +120,6 @@ class BiOpenArmFollower(Robot):
|
||||
def action_features(self) -> dict[str, type]:
|
||||
return self._motors_ft
|
||||
|
||||
@property
|
||||
def is_connected(self) -> bool:
|
||||
return self.left_arm.is_connected and self.right_arm.is_connected
|
||||
|
||||
@check_if_already_connected
|
||||
def connect(self, calibrate: bool = True) -> None:
|
||||
self.left_arm.connect(calibrate)
|
||||
self.right_arm.connect(calibrate)
|
||||
|
||||
@property
|
||||
def is_calibrated(self) -> bool:
|
||||
return self.left_arm.is_calibrated and self.right_arm.is_calibrated
|
||||
|
||||
def calibrate(self) -> None:
|
||||
self.left_arm.calibrate()
|
||||
self.right_arm.calibrate()
|
||||
|
||||
def configure(self) -> None:
|
||||
self.left_arm.configure()
|
||||
self.right_arm.configure()
|
||||
|
||||
def setup_motors(self) -> None:
|
||||
raise NotImplementedError(
|
||||
"Motor ID configuration is typically done via manufacturer tools for CAN motors."
|
||||
@@ -148,21 +127,15 @@ class BiOpenArmFollower(Robot):
|
||||
|
||||
@check_if_not_connected
|
||||
def get_observation(self) -> RobotObservation:
|
||||
obs_dict = {}
|
||||
obs_dict: RobotObservation = {}
|
||||
|
||||
# Camera keys that should NOT get the arm prefix (they already have unique names)
|
||||
left_cam_keys = set(self.left_arm.cameras.keys())
|
||||
right_cam_keys = set(self.right_arm.cameras.keys())
|
||||
# Add "left_" prefix to per-arm keys; keep top-level camera keys unprefixed.
|
||||
for key, value in self.left_arm.get_observation().items():
|
||||
obs_dict[key if key in self._top_level_cam_keys else f"left_{key}"] = value
|
||||
|
||||
# Right first, then left — matches the teleoperator (OpenArmMini) ordering
|
||||
# and the dataset feature names recorded during data collection.
|
||||
right_obs = self.right_arm.get_observation()
|
||||
for key, value in right_obs.items():
|
||||
obs_dict[key if key in right_cam_keys else f"right_{key}"] = value
|
||||
|
||||
left_obs = self.left_arm.get_observation()
|
||||
for key, value in left_obs.items():
|
||||
obs_dict[key if key in left_cam_keys else f"left_{key}"] = value
|
||||
# Add "right_" prefix
|
||||
for key, value in self.right_arm.get_observation().items():
|
||||
obs_dict[f"right_{key}"] = value
|
||||
|
||||
return obs_dict
|
||||
|
||||
@@ -189,9 +162,4 @@ class BiOpenArmFollower(Robot):
|
||||
prefixed_sent_action_left = {f"left_{key}": value for key, value in sent_action_left.items()}
|
||||
prefixed_sent_action_right = {f"right_{key}": value for key, value in sent_action_right.items()}
|
||||
|
||||
return {**prefixed_sent_action_right, **prefixed_sent_action_left}
|
||||
|
||||
@check_if_not_connected
|
||||
def disconnect(self):
|
||||
self.left_arm.disconnect()
|
||||
self.right_arm.disconnect()
|
||||
return {**prefixed_sent_action_left, **prefixed_sent_action_right}
|
||||
|
||||
@@ -32,5 +32,7 @@ class BiOpenArmFollowerConfig(RobotConfig):
|
||||
left_arm_config: OpenArmFollowerConfigBase
|
||||
right_arm_config: OpenArmFollowerConfigBase
|
||||
|
||||
# Top-level cameras shared across both arms.
|
||||
# Top-level cameras not attached to a specific side. Keys are kept as-is in
|
||||
# observations (no `left_`/`right_` prefix). Per-arm cameras (declared on
|
||||
# `{left,right}_arm_config.cameras`) are prefixed.
|
||||
cameras: dict[str, CameraConfig] = field(default_factory=dict)
|
||||
|
||||
@@ -18,7 +18,8 @@ import logging
|
||||
from functools import cached_property
|
||||
|
||||
from lerobot.types import RobotAction, RobotObservation
|
||||
from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
|
||||
from lerobot.utils.bimanual import BimanualMixin
|
||||
from lerobot.utils.decorators import check_if_not_connected
|
||||
|
||||
from ..rebot_b601_follower import RebotB601Follower, RebotB601FollowerRobotConfig
|
||||
from ..robot import Robot
|
||||
@@ -27,7 +28,7 @@ from .config_bi_rebot_b601_follower import BiRebotB601FollowerConfig
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BiRebotB601Follower(Robot):
|
||||
class BiRebotB601Follower(BimanualMixin, Robot):
|
||||
"""Bimanual Seeed Studio reBot B601-DM follower.
|
||||
|
||||
Composes two single-arm :class:`RebotB601Follower` instances. Observation and
|
||||
@@ -41,6 +42,18 @@ class BiRebotB601Follower(Robot):
|
||||
super().__init__(config)
|
||||
self.config = config
|
||||
|
||||
# Top-level cameras are opened by `left_arm` for convenience, but their
|
||||
# keys stay unprefixed in observations (tracked via `_top_level_cam_keys`).
|
||||
self._top_level_cam_keys = set(config.cameras)
|
||||
_collisions = self._top_level_cam_keys & set(
|
||||
config.left_arm_config.cameras
|
||||
) | self._top_level_cam_keys & set(config.right_arm_config.cameras)
|
||||
if _collisions:
|
||||
raise ValueError(
|
||||
f"Top-level camera names collide with per-arm camera names: {sorted(_collisions)}"
|
||||
)
|
||||
left_arm_cameras = {**config.left_arm_config.cameras, **config.cameras}
|
||||
|
||||
left_arm_config = RebotB601FollowerRobotConfig(
|
||||
id=f"{config.id}_left" if config.id else None,
|
||||
calibration_dir=config.calibration_dir,
|
||||
@@ -49,7 +62,7 @@ class BiRebotB601Follower(Robot):
|
||||
dm_serial_baud=config.left_arm_config.dm_serial_baud,
|
||||
disable_torque_on_disconnect=config.left_arm_config.disable_torque_on_disconnect,
|
||||
max_relative_target=config.left_arm_config.max_relative_target,
|
||||
cameras=config.left_arm_config.cameras,
|
||||
cameras=left_arm_cameras,
|
||||
motor_can_ids=config.left_arm_config.motor_can_ids,
|
||||
pos_vel_velocity=config.left_arm_config.pos_vel_velocity,
|
||||
gripper_torque_ratio=config.left_arm_config.gripper_torque_ratio,
|
||||
@@ -86,10 +99,12 @@ class BiRebotB601Follower(Robot):
|
||||
|
||||
@property
|
||||
def _cameras_ft(self) -> dict[str, tuple]:
|
||||
return {
|
||||
**{f"left_{k}": v for k, v in self.left_arm._cameras_ft.items()},
|
||||
**{f"right_{k}": v for k, v in self.right_arm._cameras_ft.items()},
|
||||
}
|
||||
out: dict[str, tuple] = {}
|
||||
for k, v in self.left_arm._cameras_ft.items():
|
||||
out[k if k in self._top_level_cam_keys else f"left_{k}"] = v
|
||||
for k, v in self.right_arm._cameras_ft.items():
|
||||
out[f"right_{k}"] = v
|
||||
return out
|
||||
|
||||
@cached_property
|
||||
def observation_features(self) -> dict[str, type | tuple]:
|
||||
@@ -99,32 +114,13 @@ class BiRebotB601Follower(Robot):
|
||||
def action_features(self) -> dict[str, type]:
|
||||
return self._motors_ft
|
||||
|
||||
@property
|
||||
def is_connected(self) -> bool:
|
||||
return self.left_arm.is_connected and self.right_arm.is_connected
|
||||
|
||||
@check_if_already_connected
|
||||
def connect(self, calibrate: bool = True) -> None:
|
||||
self.left_arm.connect(calibrate)
|
||||
self.right_arm.connect(calibrate)
|
||||
|
||||
@property
|
||||
def is_calibrated(self) -> bool:
|
||||
return self.left_arm.is_calibrated and self.right_arm.is_calibrated
|
||||
|
||||
def calibrate(self) -> None:
|
||||
self.left_arm.calibrate()
|
||||
self.right_arm.calibrate()
|
||||
|
||||
def configure(self) -> None:
|
||||
self.left_arm.configure()
|
||||
self.right_arm.configure()
|
||||
|
||||
@check_if_not_connected
|
||||
def get_observation(self) -> RobotObservation:
|
||||
obs_dict = {}
|
||||
obs_dict.update({f"left_{k}": v for k, v in self.left_arm.get_observation().items()})
|
||||
obs_dict.update({f"right_{k}": v for k, v in self.right_arm.get_observation().items()})
|
||||
obs_dict: RobotObservation = {}
|
||||
for k, v in self.left_arm.get_observation().items():
|
||||
obs_dict[k if k in self._top_level_cam_keys else f"left_{k}"] = v
|
||||
for k, v in self.right_arm.get_observation().items():
|
||||
obs_dict[f"right_{k}"] = v
|
||||
return obs_dict
|
||||
|
||||
@check_if_not_connected
|
||||
@@ -143,8 +139,3 @@ class BiRebotB601Follower(Robot):
|
||||
**{f"left_{k}": v for k, v in sent_action_left.items()},
|
||||
**{f"right_{k}": v for k, v in sent_action_right.items()},
|
||||
}
|
||||
|
||||
@check_if_not_connected
|
||||
def disconnect(self) -> None:
|
||||
self.left_arm.disconnect()
|
||||
self.right_arm.disconnect()
|
||||
|
||||
@@ -14,7 +14,9 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from lerobot.cameras import CameraConfig
|
||||
|
||||
from ..config import RobotConfig
|
||||
from ..rebot_b601_follower import RebotB601FollowerConfig
|
||||
@@ -27,3 +29,8 @@ class BiRebotB601FollowerConfig(RobotConfig):
|
||||
|
||||
left_arm_config: RebotB601FollowerConfig
|
||||
right_arm_config: RebotB601FollowerConfig
|
||||
|
||||
# Top-level cameras not attached to a specific side. Keys are kept as-is in
|
||||
# observations (no `left_`/`right_` prefix). Per-arm cameras (declared on
|
||||
# `{left,right}_arm_config.cameras`) are prefixed.
|
||||
cameras: dict[str, CameraConfig] = field(default_factory=dict)
|
||||
|
||||
@@ -18,7 +18,8 @@ import logging
|
||||
from functools import cached_property
|
||||
|
||||
from lerobot.types import RobotAction, RobotObservation
|
||||
from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
|
||||
from lerobot.utils.bimanual import BimanualMixin
|
||||
from lerobot.utils.decorators import check_if_not_connected
|
||||
|
||||
from ..robot import Robot
|
||||
from ..so_follower import SOFollower, SOFollowerRobotConfig
|
||||
@@ -27,7 +28,7 @@ from .config_bi_so_follower import BiSOFollowerConfig
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BiSOFollower(Robot):
|
||||
class BiSOFollower(BimanualMixin, Robot):
|
||||
"""
|
||||
[Bimanual SO Follower Arms](https://github.com/TheRobotStudio/SO-ARM100) designed by TheRobotStudio
|
||||
"""
|
||||
@@ -39,6 +40,18 @@ class BiSOFollower(Robot):
|
||||
super().__init__(config)
|
||||
self.config = config
|
||||
|
||||
# Top-level cameras are opened by `left_arm` for convenience, but their
|
||||
# keys stay unprefixed in observations (tracked via `_top_level_cam_keys`).
|
||||
self._top_level_cam_keys = set(config.cameras)
|
||||
_collisions = self._top_level_cam_keys & set(
|
||||
config.left_arm_config.cameras
|
||||
) | self._top_level_cam_keys & set(config.right_arm_config.cameras)
|
||||
if _collisions:
|
||||
raise ValueError(
|
||||
f"Top-level camera names collide with per-arm camera names: {sorted(_collisions)}"
|
||||
)
|
||||
left_arm_cameras = {**config.left_arm_config.cameras, **config.cameras}
|
||||
|
||||
left_arm_config = SOFollowerRobotConfig(
|
||||
id=f"{config.id}_left" if config.id else None,
|
||||
calibration_dir=config.calibration_dir,
|
||||
@@ -46,7 +59,7 @@ class BiSOFollower(Robot):
|
||||
disable_torque_on_disconnect=config.left_arm_config.disable_torque_on_disconnect,
|
||||
max_relative_target=config.left_arm_config.max_relative_target,
|
||||
use_degrees=config.left_arm_config.use_degrees,
|
||||
cameras=config.left_arm_config.cameras,
|
||||
cameras=left_arm_cameras,
|
||||
)
|
||||
|
||||
right_arm_config = SOFollowerRobotConfig(
|
||||
@@ -77,13 +90,12 @@ class BiSOFollower(Robot):
|
||||
|
||||
@property
|
||||
def _cameras_ft(self) -> dict[str, tuple]:
|
||||
left_arm_cameras_ft = self.left_arm._cameras_ft
|
||||
right_arm_cameras_ft = self.right_arm._cameras_ft
|
||||
|
||||
return {
|
||||
**{f"left_{k}": v for k, v in left_arm_cameras_ft.items()},
|
||||
**{f"right_{k}": v for k, v in right_arm_cameras_ft.items()},
|
||||
}
|
||||
out: dict[str, tuple] = {}
|
||||
for k, v in self.left_arm._cameras_ft.items():
|
||||
out[k if k in self._top_level_cam_keys else f"left_{k}"] = v
|
||||
for k, v in self.right_arm._cameras_ft.items():
|
||||
out[f"right_{k}"] = v
|
||||
return out
|
||||
|
||||
@cached_property
|
||||
def observation_features(self) -> dict[str, type | tuple]:
|
||||
@@ -93,42 +105,21 @@ class BiSOFollower(Robot):
|
||||
def action_features(self) -> dict[str, type]:
|
||||
return self._motors_ft
|
||||
|
||||
@property
|
||||
def is_connected(self) -> bool:
|
||||
return self.left_arm.is_connected and self.right_arm.is_connected
|
||||
|
||||
@check_if_already_connected
|
||||
def connect(self, calibrate: bool = True) -> None:
|
||||
self.left_arm.connect(calibrate)
|
||||
self.right_arm.connect(calibrate)
|
||||
|
||||
@property
|
||||
def is_calibrated(self) -> bool:
|
||||
return self.left_arm.is_calibrated and self.right_arm.is_calibrated
|
||||
|
||||
def calibrate(self) -> None:
|
||||
self.left_arm.calibrate()
|
||||
self.right_arm.calibrate()
|
||||
|
||||
def configure(self) -> None:
|
||||
self.left_arm.configure()
|
||||
self.right_arm.configure()
|
||||
|
||||
def setup_motors(self) -> None:
|
||||
self.left_arm.setup_motors()
|
||||
self.right_arm.setup_motors()
|
||||
|
||||
@check_if_not_connected
|
||||
def get_observation(self) -> RobotObservation:
|
||||
obs_dict = {}
|
||||
obs_dict: RobotObservation = {}
|
||||
|
||||
# Add "left_" prefix
|
||||
left_obs = self.left_arm.get_observation()
|
||||
obs_dict.update({f"left_{key}": value for key, value in left_obs.items()})
|
||||
# Add "left_" prefix to per-arm keys; keep top-level camera keys unprefixed.
|
||||
for key, value in self.left_arm.get_observation().items():
|
||||
obs_dict[key if key in self._top_level_cam_keys else f"left_{key}"] = value
|
||||
|
||||
# Add "right_" prefix
|
||||
right_obs = self.right_arm.get_observation()
|
||||
obs_dict.update({f"right_{key}": value for key, value in right_obs.items()})
|
||||
for key, value in self.right_arm.get_observation().items():
|
||||
obs_dict[f"right_{key}"] = value
|
||||
|
||||
return obs_dict
|
||||
|
||||
@@ -151,8 +142,3 @@ class BiSOFollower(Robot):
|
||||
prefixed_sent_action_right = {f"right_{key}": value for key, value in sent_action_right.items()}
|
||||
|
||||
return {**prefixed_sent_action_left, **prefixed_sent_action_right}
|
||||
|
||||
@check_if_not_connected
|
||||
def disconnect(self):
|
||||
self.left_arm.disconnect()
|
||||
self.right_arm.disconnect()
|
||||
|
||||
@@ -14,7 +14,9 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from lerobot.cameras import CameraConfig
|
||||
|
||||
from ..config import RobotConfig
|
||||
from ..so_follower import SOFollowerConfig
|
||||
@@ -27,3 +29,8 @@ class BiSOFollowerConfig(RobotConfig):
|
||||
|
||||
left_arm_config: SOFollowerConfig
|
||||
right_arm_config: SOFollowerConfig
|
||||
|
||||
# Top-level cameras not attached to a specific side. Keys are kept as-is in
|
||||
# observations (no `left_`/`right_` prefix). Per-arm cameras (declared on
|
||||
# `{left,right}_arm_config.cameras`) are prefixed.
|
||||
cameras: dict[str, CameraConfig] = field(default_factory=dict)
|
||||
|
||||
@@ -54,6 +54,7 @@ from lerobot.teleoperators import ( # noqa: F401
|
||||
Teleoperator,
|
||||
TeleoperatorConfig,
|
||||
bi_openarm_leader,
|
||||
bi_openarm_mini,
|
||||
bi_rebot_102_leader,
|
||||
bi_so_leader,
|
||||
homunculus,
|
||||
|
||||
@@ -77,21 +77,6 @@ from lerobot.utils.constants import ACTION, DONE, OBS_STATE, REWARD
|
||||
from lerobot.utils.utils import init_logging
|
||||
|
||||
|
||||
def get_feature_names(dataset: LeRobotDataset, key: str) -> list[str]:
|
||||
"""Return per-dimension names for a feature from the dataset metadata.
|
||||
|
||||
Only flat-list ``names`` metadata is used. Dict-style ``names`` and missing names fall back to ``{key}_{i}`` indices.
|
||||
"""
|
||||
feature = dataset.features[key]
|
||||
dim = feature["shape"][-1]
|
||||
|
||||
names = feature.get("names")
|
||||
if isinstance(names, list) and len(names) == dim:
|
||||
return [str(name) for name in names]
|
||||
|
||||
return [f"{key}_{d}" for d in range(dim)]
|
||||
|
||||
|
||||
def to_hwc_uint8_numpy(chw_float32_torch: torch.Tensor) -> np.ndarray:
|
||||
assert chw_float32_torch.dtype == torch.float32
|
||||
assert chw_float32_torch.ndim == 3
|
||||
@@ -101,31 +86,6 @@ def to_hwc_uint8_numpy(chw_float32_torch: torch.Tensor) -> np.ndarray:
|
||||
return hwc_uint8_numpy
|
||||
|
||||
|
||||
def build_blueprint_from_dataset(dataset: LeRobotDataset):
|
||||
"""Build a Rerun blueprint laying out camera images and time series for the given dataset.
|
||||
|
||||
Camera images and scalar signals (action, state, reward, done, success) are arranged in a grid.
|
||||
The per-dimension series names for ``action`` and ``state`` are applied directly
|
||||
via blueprint overrides.
|
||||
"""
|
||||
import rerun as rr
|
||||
import rerun.blueprint as rrb
|
||||
|
||||
views = [rrb.Spatial2DView(origin=key, name=key) for key in dataset.meta.camera_keys]
|
||||
|
||||
# Style multi-dimensional signals (action, state) with per-dimension names.
|
||||
for origin, key in ((ACTION, ACTION), ("state", OBS_STATE)):
|
||||
if key in dataset.features:
|
||||
names = get_feature_names(dataset, key)
|
||||
styling = rr.SeriesLines(names=names)
|
||||
views.append(rrb.TimeSeriesView(origin=origin, name=origin, overrides={origin: styling}))
|
||||
for key in (DONE, REWARD, "next.success"):
|
||||
if key in dataset.features:
|
||||
views.append(rrb.TimeSeriesView(origin=key, name=key))
|
||||
|
||||
return rrb.Blueprint(rrb.Grid(*views))
|
||||
|
||||
|
||||
def visualize_dataset(
|
||||
dataset: LeRobotDataset,
|
||||
episode_index: int,
|
||||
@@ -164,8 +124,7 @@ def visualize_dataset(
|
||||
import rerun as rr
|
||||
|
||||
spawn_local_viewer = mode == "local" and not save
|
||||
blueprint = build_blueprint_from_dataset(dataset)
|
||||
rr.init(f"{repo_id}/episode_{episode_index}", spawn=spawn_local_viewer, default_blueprint=blueprint)
|
||||
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
|
||||
@@ -183,21 +142,26 @@ def visualize_dataset(
|
||||
for batch in tqdm.tqdm(dataloader, total=len(dataloader)):
|
||||
if first_index is None:
|
||||
first_index = batch["index"][0].item()
|
||||
|
||||
# iterate over the batch
|
||||
for i in range(len(batch["index"])):
|
||||
rr.set_time("frame_index", sequence=batch["index"][i].item() - first_index)
|
||||
rr.set_time("timestamp", timestamp=batch["timestamp"][i].item())
|
||||
|
||||
# display each camera image
|
||||
for key in dataset.meta.camera_keys:
|
||||
img = to_hwc_uint8_numpy(batch[key][i])
|
||||
img_entity = rr.Image(img).compress() if display_compressed_images else rr.Image(img)
|
||||
rr.log(key, entity=img_entity)
|
||||
|
||||
# display each dimension of action space (e.g. actuators command)
|
||||
if ACTION in batch:
|
||||
rr.log(ACTION, rr.Scalars(batch[ACTION][i].numpy()))
|
||||
for dim_idx, val in enumerate(batch[ACTION][i]):
|
||||
rr.log(f"{ACTION}/{dim_idx}", rr.Scalars(val.item()))
|
||||
|
||||
# display each dimension of observed state space (e.g. agent position in joint space)
|
||||
if OBS_STATE in batch:
|
||||
rr.log("state", rr.Scalars(batch[OBS_STATE][i].numpy()))
|
||||
for dim_idx, val in enumerate(batch[OBS_STATE][i]):
|
||||
rr.log(f"state/{dim_idx}", rr.Scalars(val.item()))
|
||||
|
||||
if DONE in batch:
|
||||
rr.log(DONE, rr.Scalars(batch[DONE][i].item()))
|
||||
@@ -209,6 +173,8 @@ def visualize_dataset(
|
||||
rr.log("next.success", rr.Scalars(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"
|
||||
@@ -216,7 +182,7 @@ def visualize_dataset(
|
||||
return rrd_path
|
||||
|
||||
elif mode == "distant":
|
||||
# Keep the process alive while it serves the gRPC/web connection.
|
||||
# stop the process from exiting since it is serving the websocket connection
|
||||
try:
|
||||
while True:
|
||||
time.sleep(1)
|
||||
@@ -331,14 +297,12 @@ def main():
|
||||
)
|
||||
logging.warning("Setting grpc_port to ws_port value.")
|
||||
kwargs["grpc_port"] = kwargs.pop("ws_port")
|
||||
else:
|
||||
kwargs.pop("ws_port") # Always remove ws_port from kwargs
|
||||
|
||||
init_logging()
|
||||
logging.info("Loading dataset")
|
||||
dataset = LeRobotDataset(repo_id, episodes=[args.episode_index], root=root, tolerance_s=tolerance_s)
|
||||
|
||||
visualize_dataset(dataset, **kwargs)
|
||||
visualize_dataset(dataset, **vars(args))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -57,6 +57,7 @@ from lerobot.robots import ( # noqa: F401
|
||||
from lerobot.teleoperators import ( # noqa: F401
|
||||
TeleoperatorConfig,
|
||||
bi_openarm_leader,
|
||||
bi_openarm_mini,
|
||||
bi_rebot_102_leader,
|
||||
bi_so_leader,
|
||||
gamepad,
|
||||
|
||||
@@ -137,6 +137,7 @@ from lerobot.teleoperators import ( # noqa: F401
|
||||
Teleoperator,
|
||||
TeleoperatorConfig,
|
||||
bi_openarm_leader,
|
||||
bi_openarm_mini,
|
||||
bi_rebot_102_leader,
|
||||
bi_so_leader,
|
||||
homunculus,
|
||||
|
||||
@@ -174,6 +174,7 @@ from lerobot.teleoperators import ( # noqa: F401
|
||||
Teleoperator,
|
||||
TeleoperatorConfig,
|
||||
bi_openarm_leader,
|
||||
bi_openarm_mini,
|
||||
bi_rebot_102_leader,
|
||||
bi_so_leader,
|
||||
homunculus,
|
||||
|
||||
@@ -41,6 +41,7 @@ from lerobot.robots import ( # noqa: F401
|
||||
)
|
||||
from lerobot.teleoperators import ( # noqa: F401
|
||||
TeleoperatorConfig,
|
||||
bi_openarm_mini,
|
||||
bi_rebot_102_leader,
|
||||
bi_so_leader,
|
||||
koch_leader,
|
||||
|
||||
@@ -89,6 +89,7 @@ from lerobot.teleoperators import ( # noqa: F401
|
||||
Teleoperator,
|
||||
TeleoperatorConfig,
|
||||
bi_openarm_leader,
|
||||
bi_openarm_mini,
|
||||
bi_rebot_102_leader,
|
||||
bi_so_leader,
|
||||
gamepad,
|
||||
|
||||
@@ -232,15 +232,18 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
|
||||
torch.backends.cudnn.benchmark = True
|
||||
torch.backends.cuda.matmul.allow_tf32 = True
|
||||
|
||||
# Dataset loading synchronization: main process downloads first to avoid race conditions
|
||||
if is_main_process:
|
||||
logging.info("Creating dataset")
|
||||
# Dataset loading synchronization: each node's local main process downloads first to avoid
|
||||
# race conditions (the global main process only exists on node 0, so gating on it would let
|
||||
# all ranks of the other nodes download and build the Arrow cache concurrently).
|
||||
if accelerator.is_local_main_process:
|
||||
if is_main_process:
|
||||
logging.info("Creating dataset")
|
||||
dataset = make_dataset(cfg)
|
||||
|
||||
accelerator.wait_for_everyone()
|
||||
|
||||
# Now all other processes can safely load the dataset
|
||||
if not is_main_process:
|
||||
# Now all other processes can safely load the dataset from the local cache
|
||||
if not accelerator.is_local_main_process:
|
||||
dataset = make_dataset(cfg)
|
||||
|
||||
# Create environment used for evaluating checkpoints during training on simulation data.
|
||||
@@ -386,12 +389,19 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
|
||||
# create dataloader for offline training
|
||||
if hasattr(active_cfg, "drop_n_last_frames"):
|
||||
shuffle = False
|
||||
# A dedicated generator (rather than the global torch RNG) lets accelerator.prepare
|
||||
# synchronize the shuffle permutation across ranks, keeping batch shards disjoint even
|
||||
# when ranks consume the global RNG asymmetrically (e.g. eval on the main process only).
|
||||
sampler_generator = torch.Generator()
|
||||
if cfg.seed is not None:
|
||||
sampler_generator.manual_seed(cfg.seed)
|
||||
sampler = EpisodeAwareSampler(
|
||||
dataset.meta.episodes["dataset_from_index"],
|
||||
dataset.meta.episodes["dataset_to_index"],
|
||||
episode_indices_to_use=dataset.episodes,
|
||||
drop_n_last_frames=active_cfg.drop_n_last_frames,
|
||||
shuffle=True,
|
||||
generator=sampler_generator,
|
||||
)
|
||||
else:
|
||||
shuffle = True
|
||||
|
||||
@@ -18,7 +18,8 @@ import logging
|
||||
from functools import cached_property
|
||||
|
||||
from lerobot.types import RobotAction
|
||||
from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
|
||||
from lerobot.utils.bimanual import BimanualMixin
|
||||
from lerobot.utils.decorators import check_if_not_connected
|
||||
|
||||
from ..openarm_leader import OpenArmLeader, OpenArmLeaderConfig
|
||||
from ..teleoperator import Teleoperator
|
||||
@@ -27,7 +28,7 @@ from .config_bi_openarm_leader import BiOpenArmLeaderConfig
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BiOpenArmLeader(Teleoperator):
|
||||
class BiOpenArmLeader(BimanualMixin, Teleoperator):
|
||||
"""
|
||||
Bimanual OpenArm Leader Arms
|
||||
"""
|
||||
@@ -86,27 +87,6 @@ class BiOpenArmLeader(Teleoperator):
|
||||
def feedback_features(self) -> dict[str, type]:
|
||||
return {}
|
||||
|
||||
@property
|
||||
def is_connected(self) -> bool:
|
||||
return self.left_arm.is_connected and self.right_arm.is_connected
|
||||
|
||||
@check_if_already_connected
|
||||
def connect(self, calibrate: bool = True) -> None:
|
||||
self.left_arm.connect(calibrate)
|
||||
self.right_arm.connect(calibrate)
|
||||
|
||||
@property
|
||||
def is_calibrated(self) -> bool:
|
||||
return self.left_arm.is_calibrated and self.right_arm.is_calibrated
|
||||
|
||||
def calibrate(self) -> None:
|
||||
self.left_arm.calibrate()
|
||||
self.right_arm.calibrate()
|
||||
|
||||
def configure(self) -> None:
|
||||
self.left_arm.configure()
|
||||
self.right_arm.configure()
|
||||
|
||||
def setup_motors(self) -> None:
|
||||
raise NotImplementedError(
|
||||
"Motor ID configuration is typically done via manufacturer tools for CAN motors."
|
||||
@@ -129,8 +109,3 @@ class BiOpenArmLeader(Teleoperator):
|
||||
def send_feedback(self, feedback: dict[str, float]) -> None:
|
||||
# TODO: Implement force feedback
|
||||
raise NotImplementedError
|
||||
|
||||
@check_if_not_connected
|
||||
def disconnect(self) -> None:
|
||||
self.left_arm.disconnect()
|
||||
self.right_arm.disconnect()
|
||||
|
||||
@@ -23,7 +23,7 @@ from ..openarm_leader import OpenArmLeaderConfigBase
|
||||
@TeleoperatorConfig.register_subclass("bi_openarm_leader")
|
||||
@dataclass
|
||||
class BiOpenArmLeaderConfig(TeleoperatorConfig):
|
||||
"""Configuration class for Bi OpenArm Follower robots."""
|
||||
"""Configuration class for Bi OpenArm Leader teleoperators."""
|
||||
|
||||
left_arm_config: OpenArmLeaderConfigBase
|
||||
right_arm_config: OpenArmLeaderConfigBase
|
||||
|
||||
@@ -0,0 +1,20 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from .bi_openarm_mini import BiOpenArmMini
|
||||
from .config_bi_openarm_mini import BiOpenArmMiniConfig
|
||||
|
||||
__all__ = ["BiOpenArmMini", "BiOpenArmMiniConfig"]
|
||||
@@ -0,0 +1,101 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import logging
|
||||
from functools import cached_property
|
||||
|
||||
from lerobot.types import RobotAction
|
||||
from lerobot.utils.bimanual import BimanualMixin
|
||||
from lerobot.utils.decorators import check_if_not_connected
|
||||
|
||||
from ..openarm_mini import OpenArmMini, OpenArmMiniConfig
|
||||
from ..teleoperator import Teleoperator
|
||||
from .config_bi_openarm_mini import BiOpenArmMiniConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BiOpenArmMini(BimanualMixin, Teleoperator):
|
||||
"""Bimanual OpenArm Mini teleoperator.
|
||||
|
||||
Composes two single-arm :class:`OpenArmMini` instances. Action and feedback
|
||||
keys of each arm are namespaced with a ``left_`` / ``right_`` prefix, so a
|
||||
bimanual leader can teleoperate a bimanual OpenArm follower.
|
||||
"""
|
||||
|
||||
config_class = BiOpenArmMiniConfig
|
||||
name = "bi_openarm_mini"
|
||||
|
||||
def __init__(self, config: BiOpenArmMiniConfig):
|
||||
super().__init__(config)
|
||||
self.config = config
|
||||
|
||||
# `side` is forced to match left/right regardless of what the user passed
|
||||
# on the per-arm base config — the bimanual wrapper owns the side semantics.
|
||||
left_arm_config = OpenArmMiniConfig(
|
||||
id=f"{config.id}_left" if config.id else None,
|
||||
calibration_dir=config.calibration_dir,
|
||||
port=config.left_arm_config.port,
|
||||
side="left",
|
||||
use_degrees=config.left_arm_config.use_degrees,
|
||||
)
|
||||
|
||||
right_arm_config = OpenArmMiniConfig(
|
||||
id=f"{config.id}_right" if config.id else None,
|
||||
calibration_dir=config.calibration_dir,
|
||||
port=config.right_arm_config.port,
|
||||
side="right",
|
||||
use_degrees=config.right_arm_config.use_degrees,
|
||||
)
|
||||
|
||||
self.left_arm = OpenArmMini(left_arm_config)
|
||||
self.right_arm = OpenArmMini(right_arm_config)
|
||||
|
||||
@cached_property
|
||||
def action_features(self) -> dict[str, type]:
|
||||
return {
|
||||
**{f"left_{k}": v for k, v in self.left_arm.action_features.items()},
|
||||
**{f"right_{k}": v for k, v in self.right_arm.action_features.items()},
|
||||
}
|
||||
|
||||
@cached_property
|
||||
def feedback_features(self) -> dict[str, type]:
|
||||
return {
|
||||
**{f"left_{k}": v for k, v in self.left_arm.feedback_features.items()},
|
||||
**{f"right_{k}": v for k, v in self.right_arm.feedback_features.items()},
|
||||
}
|
||||
|
||||
def setup_motors(self) -> None:
|
||||
self.left_arm.setup_motors()
|
||||
self.right_arm.setup_motors()
|
||||
|
||||
@check_if_not_connected
|
||||
def get_action(self) -> RobotAction:
|
||||
action: RobotAction = {}
|
||||
for k, v in self.left_arm.get_action().items():
|
||||
action[f"left_{k}"] = v
|
||||
for k, v in self.right_arm.get_action().items():
|
||||
action[f"right_{k}"] = v
|
||||
return action
|
||||
|
||||
@check_if_not_connected
|
||||
def send_feedback(self, feedback: dict[str, float]) -> None:
|
||||
left_fb = {k.removeprefix("left_"): v for k, v in feedback.items() if k.startswith("left_")}
|
||||
right_fb = {k.removeprefix("right_"): v for k, v in feedback.items() if k.startswith("right_")}
|
||||
if left_fb:
|
||||
self.left_arm.send_feedback(left_fb)
|
||||
if right_fb:
|
||||
self.right_arm.send_feedback(right_fb)
|
||||
@@ -0,0 +1,29 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
from ..config import TeleoperatorConfig
|
||||
from ..openarm_mini import OpenArmMiniConfigBase
|
||||
|
||||
|
||||
@TeleoperatorConfig.register_subclass("bi_openarm_mini")
|
||||
@dataclass
|
||||
class BiOpenArmMiniConfig(TeleoperatorConfig):
|
||||
"""Configuration class for Bi OpenArm Mini teleoperators."""
|
||||
|
||||
left_arm_config: OpenArmMiniConfigBase
|
||||
right_arm_config: OpenArmMiniConfigBase
|
||||
@@ -14,7 +14,7 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from .bi_rebot_102_leader import BiRebotArm102Leader
|
||||
from .config_bi_rebot_102_leader import BiRebotArm102LeaderConfig
|
||||
from .bi_rebot_102_leader import BiRebot102Leader
|
||||
from .config_bi_rebot_102_leader import BiRebot102LeaderConfig
|
||||
|
||||
__all__ = ["BiRebotArm102Leader", "BiRebotArm102LeaderConfig"]
|
||||
__all__ = ["BiRebot102Leader", "BiRebot102LeaderConfig"]
|
||||
|
||||
@@ -18,16 +18,17 @@ import logging
|
||||
from functools import cached_property
|
||||
|
||||
from lerobot.types import RobotAction
|
||||
from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
|
||||
from lerobot.utils.bimanual import BimanualMixin
|
||||
from lerobot.utils.decorators import check_if_not_connected
|
||||
|
||||
from ..rebot_102_leader import RebotArm102Leader, RebotArm102LeaderTeleopConfig
|
||||
from ..teleoperator import Teleoperator
|
||||
from .config_bi_rebot_102_leader import BiRebotArm102LeaderConfig
|
||||
from .config_bi_rebot_102_leader import BiRebot102LeaderConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BiRebotArm102Leader(Teleoperator):
|
||||
class BiRebot102Leader(BimanualMixin, Teleoperator):
|
||||
"""Bimanual Seeed Studio StarArm102 / reBot Arm 102 leader.
|
||||
|
||||
Composes two single-arm :class:`RebotArm102Leader` instances. Action keys of
|
||||
@@ -35,10 +36,10 @@ class BiRebotArm102Leader(Teleoperator):
|
||||
leader can teleoperate a bimanual reBot B601 follower.
|
||||
"""
|
||||
|
||||
config_class = BiRebotArm102LeaderConfig
|
||||
config_class = BiRebot102LeaderConfig
|
||||
name = "bi_rebot_102_leader"
|
||||
|
||||
def __init__(self, config: BiRebotArm102LeaderConfig):
|
||||
def __init__(self, config: BiRebot102LeaderConfig):
|
||||
super().__init__(config)
|
||||
self.config = config
|
||||
|
||||
@@ -76,27 +77,6 @@ class BiRebotArm102Leader(Teleoperator):
|
||||
def feedback_features(self) -> dict[str, type]:
|
||||
return {}
|
||||
|
||||
@property
|
||||
def is_connected(self) -> bool:
|
||||
return self.left_arm.is_connected and self.right_arm.is_connected
|
||||
|
||||
@check_if_already_connected
|
||||
def connect(self, calibrate: bool = True) -> None:
|
||||
self.left_arm.connect(calibrate)
|
||||
self.right_arm.connect(calibrate)
|
||||
|
||||
@property
|
||||
def is_calibrated(self) -> bool:
|
||||
return self.left_arm.is_calibrated and self.right_arm.is_calibrated
|
||||
|
||||
def calibrate(self) -> None:
|
||||
self.left_arm.calibrate()
|
||||
self.right_arm.calibrate()
|
||||
|
||||
def configure(self) -> None:
|
||||
self.left_arm.configure()
|
||||
self.right_arm.configure()
|
||||
|
||||
@check_if_not_connected
|
||||
def get_action(self) -> RobotAction:
|
||||
action_dict = {}
|
||||
@@ -106,8 +86,3 @@ class BiRebotArm102Leader(Teleoperator):
|
||||
|
||||
def send_feedback(self, feedback: dict[str, float]) -> None:
|
||||
raise NotImplementedError("Feedback is not implemented for the reBot Arm 102 leader.")
|
||||
|
||||
@check_if_not_connected
|
||||
def disconnect(self) -> None:
|
||||
self.left_arm.disconnect()
|
||||
self.right_arm.disconnect()
|
||||
|
||||
@@ -22,7 +22,7 @@ from ..rebot_102_leader import RebotArm102LeaderConfig
|
||||
|
||||
@TeleoperatorConfig.register_subclass("bi_rebot_102_leader")
|
||||
@dataclass
|
||||
class BiRebotArm102LeaderConfig(TeleoperatorConfig):
|
||||
class BiRebot102LeaderConfig(TeleoperatorConfig):
|
||||
"""Configuration class for the bimanual reBot Arm 102 leader teleoperator."""
|
||||
|
||||
left_arm_config: RebotArm102LeaderConfig
|
||||
|
||||
@@ -17,7 +17,9 @@
|
||||
import logging
|
||||
from functools import cached_property
|
||||
|
||||
from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
|
||||
from lerobot.types import RobotAction
|
||||
from lerobot.utils.bimanual import BimanualMixin
|
||||
from lerobot.utils.decorators import check_if_not_connected
|
||||
|
||||
from ..so_leader import SOLeader, SOLeaderTeleopConfig
|
||||
from ..teleoperator import Teleoperator
|
||||
@@ -26,7 +28,7 @@ from .config_bi_so_leader import BiSOLeaderConfig
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BiSOLeader(Teleoperator):
|
||||
class BiSOLeader(BimanualMixin, Teleoperator):
|
||||
"""
|
||||
[Bimanual SO Leader Arms](https://github.com/TheRobotStudio/SO-ARM100) designed by TheRobotStudio
|
||||
"""
|
||||
@@ -67,33 +69,12 @@ class BiSOLeader(Teleoperator):
|
||||
def feedback_features(self) -> dict[str, type]:
|
||||
return {}
|
||||
|
||||
@property
|
||||
def is_connected(self) -> bool:
|
||||
return self.left_arm.is_connected and self.right_arm.is_connected
|
||||
|
||||
@check_if_already_connected
|
||||
def connect(self, calibrate: bool = True) -> None:
|
||||
self.left_arm.connect(calibrate)
|
||||
self.right_arm.connect(calibrate)
|
||||
|
||||
@property
|
||||
def is_calibrated(self) -> bool:
|
||||
return self.left_arm.is_calibrated and self.right_arm.is_calibrated
|
||||
|
||||
def calibrate(self) -> None:
|
||||
self.left_arm.calibrate()
|
||||
self.right_arm.calibrate()
|
||||
|
||||
def configure(self) -> None:
|
||||
self.left_arm.configure()
|
||||
self.right_arm.configure()
|
||||
|
||||
def setup_motors(self) -> None:
|
||||
self.left_arm.setup_motors()
|
||||
self.right_arm.setup_motors()
|
||||
|
||||
@check_if_not_connected
|
||||
def get_action(self) -> dict[str, float]:
|
||||
def get_action(self) -> RobotAction:
|
||||
action_dict = {}
|
||||
|
||||
# Add "left_" prefix
|
||||
@@ -109,8 +90,3 @@ class BiSOLeader(Teleoperator):
|
||||
def send_feedback(self, feedback: dict[str, float]) -> None:
|
||||
# TODO: Implement force feedback
|
||||
raise NotImplementedError
|
||||
|
||||
@check_if_not_connected
|
||||
def disconnect(self) -> None:
|
||||
self.left_arm.disconnect()
|
||||
self.right_arm.disconnect()
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
|
||||
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
@@ -14,7 +14,7 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from .config_openarm_mini import OpenArmMiniConfig
|
||||
from .config_openarm_mini import OpenArmMiniConfig, OpenArmMiniConfigBase
|
||||
from .openarm_mini import OpenArmMini
|
||||
|
||||
__all__ = ["OpenArmMini", "OpenArmMiniConfig"]
|
||||
__all__ = ["OpenArmMini", "OpenArmMiniConfig", "OpenArmMiniConfigBase"]
|
||||
|
||||
@@ -19,12 +19,21 @@ from dataclasses import dataclass
|
||||
from ..config import TeleoperatorConfig
|
||||
|
||||
|
||||
@TeleoperatorConfig.register_subclass("openarm_mini")
|
||||
@dataclass
|
||||
class OpenArmMiniConfig(TeleoperatorConfig):
|
||||
"""Configuration for OpenArm Mini teleoperator with Feetech motors (dual arms)."""
|
||||
class OpenArmMiniConfigBase:
|
||||
"""Base configuration for the OpenArm Mini teleoperator (Feetech STS3215, 7DOF + gripper)."""
|
||||
|
||||
port_right: str = "/dev/ttyUSB0"
|
||||
port_left: str = "/dev/ttyUSB1"
|
||||
# Serial port for the Feetech bus (e.g., "/dev/ttyUSB0").
|
||||
port: str
|
||||
|
||||
# Side of the arm: "left" or "right". Controls per-joint direction flips applied
|
||||
# during readout. If `None`, no flipping is applied.
|
||||
side: str | None = None
|
||||
|
||||
use_degrees: bool = True
|
||||
|
||||
|
||||
@TeleoperatorConfig.register_subclass("openarm_mini")
|
||||
@dataclass
|
||||
class OpenArmMiniConfig(TeleoperatorConfig, OpenArmMiniConfigBase):
|
||||
pass
|
||||
|
||||
@@ -31,22 +31,22 @@ from .config_openarm_mini import OpenArmMiniConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Motors whose direction is inverted during readout
|
||||
RIGHT_MOTORS_TO_FLIP = ["joint_1", "joint_2", "joint_3", "joint_4", "joint_5", "joint_7"]
|
||||
LEFT_MOTORS_TO_FLIP = ["joint_1", "joint_3", "joint_4", "joint_5", "joint_6", "joint_7"]
|
||||
# Per-side motor direction flips applied during readout.
|
||||
SIDE_MOTORS_TO_FLIP: dict[str, list[str]] = {
|
||||
"left": ["joint_1", "joint_3", "joint_4", "joint_5", "joint_6", "joint_7"],
|
||||
"right": ["joint_1", "joint_2", "joint_3", "joint_4", "joint_5", "joint_7"],
|
||||
}
|
||||
|
||||
# Leader joint 6 maps to follower joint 7 and vice versa
|
||||
# Leader joint 6 ↔ follower joint 7 (symmetric — its own inverse).
|
||||
JOINT_REMAP = {"joint_6": "joint_7", "joint_7": "joint_6"}
|
||||
JOINT_REMAP_REVERSE = {"joint_7": "joint_6", "joint_6": "joint_7"}
|
||||
|
||||
GRIPPER_TELEOP_TO_DEGREES = -0.65
|
||||
|
||||
|
||||
class OpenArmMini(Teleoperator):
|
||||
"""
|
||||
OpenArm Mini Teleoperator with dual Feetech-based arms (8 motors per arm).
|
||||
"""OpenArm Mini single-arm teleoperator (Feetech STS3215, 7DOF + gripper).
|
||||
|
||||
Each arm has 7 joints plus a gripper, using Feetech STS3215 servos.
|
||||
For the bimanual setup, see :class:`BiOpenArmMini` which composes two of these.
|
||||
"""
|
||||
|
||||
config_class = OpenArmMiniConfig
|
||||
@@ -56,9 +56,12 @@ class OpenArmMini(Teleoperator):
|
||||
super().__init__(config)
|
||||
self.config = config
|
||||
|
||||
if config.side is not None and config.side not in SIDE_MOTORS_TO_FLIP:
|
||||
raise ValueError(f"Invalid side '{config.side}'; expected 'left', 'right', or None.")
|
||||
self._motors_to_flip: list[str] = SIDE_MOTORS_TO_FLIP.get(config.side, []) if config.side else []
|
||||
|
||||
norm_mode_body = MotorNormMode.DEGREES
|
||||
|
||||
motors_right = {
|
||||
motors = {
|
||||
"joint_1": Motor(1, "sts3215", norm_mode_body),
|
||||
"joint_2": Motor(2, "sts3215", norm_mode_body),
|
||||
"joint_3": Motor(3, "sts3215", norm_mode_body),
|
||||
@@ -69,46 +72,15 @@ class OpenArmMini(Teleoperator):
|
||||
"gripper": Motor(8, "sts3215", MotorNormMode.RANGE_0_100),
|
||||
}
|
||||
|
||||
motors_left = {
|
||||
"joint_1": Motor(1, "sts3215", norm_mode_body),
|
||||
"joint_2": Motor(2, "sts3215", norm_mode_body),
|
||||
"joint_3": Motor(3, "sts3215", norm_mode_body),
|
||||
"joint_4": Motor(4, "sts3215", norm_mode_body),
|
||||
"joint_5": Motor(5, "sts3215", norm_mode_body),
|
||||
"joint_6": Motor(6, "sts3215", norm_mode_body),
|
||||
"joint_7": Motor(7, "sts3215", norm_mode_body),
|
||||
"gripper": Motor(8, "sts3215", MotorNormMode.RANGE_0_100),
|
||||
}
|
||||
|
||||
cal_right = {
|
||||
k.replace("right_", ""): v for k, v in (self.calibration or {}).items() if k.startswith("right_")
|
||||
}
|
||||
cal_left = {
|
||||
k.replace("left_", ""): v for k, v in (self.calibration or {}).items() if k.startswith("left_")
|
||||
}
|
||||
|
||||
self.bus_right = FeetechMotorsBus(
|
||||
port=self.config.port_right,
|
||||
motors=motors_right,
|
||||
calibration=cal_right,
|
||||
)
|
||||
|
||||
self.bus_left = FeetechMotorsBus(
|
||||
port=self.config.port_left,
|
||||
motors=motors_left,
|
||||
calibration=cal_left,
|
||||
self.bus = FeetechMotorsBus(
|
||||
port=self.config.port,
|
||||
motors=motors,
|
||||
calibration=self.calibration,
|
||||
)
|
||||
|
||||
@property
|
||||
def action_features(self) -> dict[str, type]:
|
||||
# Right first, then left — matches the robot (BiOpenArmFollower) ordering
|
||||
# and the dataset feature names recorded during data collection.
|
||||
features: dict[str, type] = {}
|
||||
for motor in self.bus_right.motors:
|
||||
features[f"right_{motor}.pos"] = float
|
||||
for motor in self.bus_left.motors:
|
||||
features[f"left_{motor}.pos"] = float
|
||||
return features
|
||||
return {f"{motor}.pos": float for motor in self.bus.motors}
|
||||
|
||||
@property
|
||||
def feedback_features(self) -> dict[str, type]:
|
||||
@@ -116,14 +88,12 @@ class OpenArmMini(Teleoperator):
|
||||
|
||||
@property
|
||||
def is_connected(self) -> bool:
|
||||
return self.bus_right.is_connected and self.bus_left.is_connected
|
||||
return self.bus.is_connected
|
||||
|
||||
@check_if_already_connected
|
||||
def connect(self, calibrate: bool = True) -> None:
|
||||
logger.info(f"Connecting right arm on {self.config.port_right}...")
|
||||
self.bus_right.connect()
|
||||
logger.info(f"Connecting left arm on {self.config.port_left}...")
|
||||
self.bus_left.connect()
|
||||
logger.info(f"Connecting arm on {self.config.port}...")
|
||||
self.bus.connect()
|
||||
|
||||
if calibrate:
|
||||
self.calibrate()
|
||||
@@ -133,14 +103,14 @@ class OpenArmMini(Teleoperator):
|
||||
|
||||
@property
|
||||
def is_calibrated(self) -> bool:
|
||||
return self.bus_right.is_calibrated and self.bus_left.is_calibrated
|
||||
return self.bus.is_calibrated
|
||||
|
||||
def calibrate(self) -> None:
|
||||
"""
|
||||
Run calibration procedure for OpenArm Mini.
|
||||
Run calibration procedure for a single OpenArm Mini arm.
|
||||
|
||||
1. Disable torque
|
||||
2. Ask user to position arms in hanging position with grippers closed
|
||||
2. Ask user to position arm in hanging position with gripper closed
|
||||
3. Set this as zero position via half-turn homing
|
||||
4. Interactive gripper calibration (open/close positions)
|
||||
5. Save calibration
|
||||
@@ -152,70 +122,51 @@ class OpenArmMini(Teleoperator):
|
||||
)
|
||||
if user_input.strip().lower() != "c":
|
||||
logger.info(f"Using existing calibration for {self.id}")
|
||||
cal_right = {
|
||||
k.replace("right_", ""): v for k, v in self.calibration.items() if k.startswith("right_")
|
||||
}
|
||||
cal_left = {
|
||||
k.replace("left_", ""): v for k, v in self.calibration.items() if k.startswith("left_")
|
||||
}
|
||||
self.bus_right.write_calibration(cal_right)
|
||||
self.bus_left.write_calibration(cal_left)
|
||||
self.bus.write_calibration(self.calibration)
|
||||
return
|
||||
|
||||
logger.info(f"\nRunning calibration for {self}")
|
||||
|
||||
self._calibrate_arm("right", self.bus_right)
|
||||
self._calibrate_arm("left", self.bus_left)
|
||||
self.bus.disable_torque()
|
||||
|
||||
self._save_calibration()
|
||||
print(f"\nCalibration complete and saved to {self.calibration_fpath}")
|
||||
logger.info("Setting Phase to 12 for all motors...")
|
||||
for motor in self.bus.motors:
|
||||
self.bus.write("Phase", motor, 12)
|
||||
|
||||
def _calibrate_arm(self, arm_name: str, bus: FeetechMotorsBus) -> None:
|
||||
"""Calibrate a single arm with Feetech motors."""
|
||||
logger.info(f"\n=== Calibrating {arm_name.upper()} arm ===")
|
||||
|
||||
bus.disable_torque()
|
||||
|
||||
logger.info(f"Setting Phase to 12 for all motors in {arm_name.upper()} arm...")
|
||||
for motor in bus.motors:
|
||||
bus.write("Phase", motor, 12)
|
||||
|
||||
for motor in bus.motors:
|
||||
bus.write("Operating_Mode", motor, OperatingMode.POSITION.value)
|
||||
for motor in self.bus.motors:
|
||||
self.bus.write("Operating_Mode", motor, OperatingMode.POSITION.value)
|
||||
|
||||
input(
|
||||
f"\nCalibration: Zero Position ({arm_name.upper()} arm)\n"
|
||||
"\nCalibration: Zero Position\n"
|
||||
"Position the arm in the following configuration:\n"
|
||||
" - Arm hanging straight down\n"
|
||||
" - Gripper closed\n"
|
||||
"Press ENTER when ready..."
|
||||
)
|
||||
|
||||
homing_offsets = bus.set_half_turn_homings()
|
||||
logger.info(f"{arm_name.capitalize()} arm zero position set.")
|
||||
homing_offsets = self.bus.set_half_turn_homings()
|
||||
logger.info("Arm zero position set.")
|
||||
|
||||
print(f"\nSetting motor ranges for {arm_name.upper()} arm\n")
|
||||
print("\nSetting motor ranges\n")
|
||||
|
||||
if self.calibration is None:
|
||||
self.calibration = {}
|
||||
|
||||
motor_resolution = bus.model_resolution_table[list(bus.motors.values())[0].model]
|
||||
motor_resolution = self.bus.model_resolution_table[list(self.bus.motors.values())[0].model]
|
||||
max_res = motor_resolution - 1
|
||||
|
||||
for motor_name, motor in bus.motors.items():
|
||||
prefixed_name = f"{arm_name}_{motor_name}"
|
||||
|
||||
for motor_name, motor in self.bus.motors.items():
|
||||
if motor_name == "gripper":
|
||||
input(
|
||||
f"\nGripper Calibration ({arm_name.upper()} arm)\n"
|
||||
f"Step 1: CLOSE the gripper fully\n"
|
||||
f"Press ENTER when gripper is closed..."
|
||||
"\nGripper Calibration\n"
|
||||
"Step 1: CLOSE the gripper fully\n"
|
||||
"Press ENTER when gripper is closed..."
|
||||
)
|
||||
closed_pos = bus.read("Present_Position", motor_name, normalize=False)
|
||||
closed_pos = self.bus.read("Present_Position", motor_name, normalize=False)
|
||||
logger.info(f" Gripper closed position recorded: {closed_pos}")
|
||||
|
||||
input("\nStep 2: OPEN the gripper fully\nPress ENTER when gripper is fully open...")
|
||||
open_pos = bus.read("Present_Position", motor_name, normalize=False)
|
||||
open_pos = self.bus.read("Present_Position", motor_name, normalize=False)
|
||||
logger.info(f" Gripper open position recorded: {open_pos}")
|
||||
|
||||
if closed_pos < open_pos:
|
||||
@@ -228,16 +179,16 @@ class OpenArmMini(Teleoperator):
|
||||
drive_mode = 1
|
||||
|
||||
logger.info(
|
||||
f" {prefixed_name}: range set to [{range_min}, {range_max}] "
|
||||
f" {motor_name}: range set to [{range_min}, {range_max}] "
|
||||
f"(0=closed, 100=open, drive_mode={drive_mode})"
|
||||
)
|
||||
else:
|
||||
range_min = 0
|
||||
range_max = max_res
|
||||
drive_mode = 0
|
||||
logger.info(f" {prefixed_name}: range set to [0, {max_res}] (full motor range)")
|
||||
logger.info(f" {motor_name}: range set to [0, {max_res}] (full motor range)")
|
||||
|
||||
self.calibration[prefixed_name] = MotorCalibration(
|
||||
self.calibration[motor_name] = MotorCalibration(
|
||||
id=motor.id,
|
||||
drive_mode=drive_mode,
|
||||
homing_offset=homing_offsets[motor_name],
|
||||
@@ -245,108 +196,68 @@ class OpenArmMini(Teleoperator):
|
||||
range_max=range_max,
|
||||
)
|
||||
|
||||
cal_for_bus = {
|
||||
k.replace(f"{arm_name}_", ""): v
|
||||
for k, v in self.calibration.items()
|
||||
if k.startswith(f"{arm_name}_")
|
||||
}
|
||||
bus.write_calibration(cal_for_bus)
|
||||
self.bus.write_calibration(self.calibration)
|
||||
self._save_calibration()
|
||||
print(f"\nCalibration complete and saved to {self.calibration_fpath}")
|
||||
|
||||
def configure(self) -> None:
|
||||
self.bus_right.disable_torque()
|
||||
self.bus_right.configure_motors()
|
||||
for motor in self.bus_right.motors:
|
||||
self.bus_right.write("Operating_Mode", motor, OperatingMode.POSITION.value)
|
||||
|
||||
self.bus_left.disable_torque()
|
||||
self.bus_left.configure_motors()
|
||||
for motor in self.bus_left.motors:
|
||||
self.bus_left.write("Operating_Mode", motor, OperatingMode.POSITION.value)
|
||||
self.bus.disable_torque()
|
||||
self.bus.configure_motors()
|
||||
for motor in self.bus.motors:
|
||||
self.bus.write("Operating_Mode", motor, OperatingMode.POSITION.value)
|
||||
|
||||
def setup_motors(self) -> None:
|
||||
print("\nSetting up RIGHT arm motors...")
|
||||
for motor in reversed(self.bus_right.motors):
|
||||
input(f"Connect the controller board to the RIGHT '{motor}' motor only and press enter.")
|
||||
self.bus_right.setup_motor(motor)
|
||||
print(f"RIGHT '{motor}' motor id set to {self.bus_right.motors[motor].id}")
|
||||
|
||||
print("\nSetting up LEFT arm motors...")
|
||||
for motor in reversed(self.bus_left.motors):
|
||||
input(f"Connect the controller board to the LEFT '{motor}' motor only and press enter.")
|
||||
self.bus_left.setup_motor(motor)
|
||||
print(f"LEFT '{motor}' motor id set to {self.bus_left.motors[motor].id}")
|
||||
for motor in reversed(self.bus.motors):
|
||||
input(f"Connect the controller board to the '{motor}' motor only and press enter.")
|
||||
self.bus.setup_motor(motor)
|
||||
print(f"'{motor}' motor id set to {self.bus.motors[motor].id}")
|
||||
|
||||
@check_if_not_connected
|
||||
def get_action(self) -> RobotAction:
|
||||
"""Get current action from both arms (read positions from all motors)."""
|
||||
"""Get current action (read positions from all motors)."""
|
||||
start = time.perf_counter()
|
||||
|
||||
right_positions = self.bus_right.sync_read("Present_Position")
|
||||
left_positions = self.bus_left.sync_read("Present_Position")
|
||||
positions = self.bus.sync_read("Present_Position")
|
||||
|
||||
# Right first, then left — matches the robot (BiOpenArmFollower) ordering
|
||||
# and the dataset feature names recorded during data collection.
|
||||
# Joint 6↔7 remap: leader joint_6 → follower joint_7 and vice versa.
|
||||
# Per-side direction flip is applied based on the configured `side`.
|
||||
action: dict[str, Any] = {}
|
||||
for motor, val in right_positions.items():
|
||||
for motor, val in positions.items():
|
||||
target = JOINT_REMAP.get(motor, motor)
|
||||
if motor == "gripper":
|
||||
# Convert gripper from teleop 0-100 to openarms degrees: 0→0°, 100→-65°
|
||||
action[f"right_{target}.pos"] = val * GRIPPER_TELEOP_TO_DEGREES
|
||||
action[f"{target}.pos"] = val * GRIPPER_TELEOP_TO_DEGREES
|
||||
else:
|
||||
action[f"right_{target}.pos"] = -val if motor in RIGHT_MOTORS_TO_FLIP else val
|
||||
for motor, val in left_positions.items():
|
||||
target = JOINT_REMAP.get(motor, motor)
|
||||
if motor == "gripper":
|
||||
action[f"left_{target}.pos"] = val * GRIPPER_TELEOP_TO_DEGREES
|
||||
else:
|
||||
action[f"left_{target}.pos"] = -val if motor in LEFT_MOTORS_TO_FLIP else val
|
||||
action[f"{target}.pos"] = -val if motor in self._motors_to_flip else val
|
||||
|
||||
dt_ms = (time.perf_counter() - start) * 1e3
|
||||
logger.debug(f"{self} read action: {dt_ms:.1f}ms")
|
||||
return action
|
||||
|
||||
def enable_torque(self) -> None:
|
||||
"""Enable torque on both arms for position control."""
|
||||
self.bus_right.enable_torque()
|
||||
self.bus_left.enable_torque()
|
||||
self.bus.enable_torque()
|
||||
|
||||
def disable_torque(self) -> None:
|
||||
"""Disable torque on both arms for free movement."""
|
||||
self.bus_right.disable_torque()
|
||||
self.bus_left.disable_torque()
|
||||
self.bus.disable_torque()
|
||||
|
||||
def write_goal_positions(self, positions: dict[str, float]) -> None:
|
||||
"""Write goal positions to motors (inverse of get_action flip/gripper/remap logic)."""
|
||||
right_goals: dict[str, float] = {}
|
||||
left_goals: dict[str, float] = {}
|
||||
|
||||
goals: dict[str, float] = {}
|
||||
for key, val in positions.items():
|
||||
if not key.endswith(".pos"):
|
||||
continue
|
||||
motor_name = key.removesuffix(".pos")
|
||||
if motor_name.startswith("right_"):
|
||||
base = motor_name.removeprefix("right_")
|
||||
# Reverse remap: follower joint_7 → leader joint_6 and vice versa
|
||||
target = JOINT_REMAP_REVERSE.get(base, base)
|
||||
if base == "gripper":
|
||||
# Convert robot degrees to teleop 0-100: 0°→0, -65°→100
|
||||
right_goals[target] = val / GRIPPER_TELEOP_TO_DEGREES
|
||||
else:
|
||||
# Un-flip using the ORIGINAL motor name (target = leader motor)
|
||||
right_goals[target] = -val if target in RIGHT_MOTORS_TO_FLIP else val
|
||||
elif motor_name.startswith("left_"):
|
||||
base = motor_name.removeprefix("left_")
|
||||
target = JOINT_REMAP_REVERSE.get(base, base)
|
||||
if base == "gripper":
|
||||
left_goals[target] = val / GRIPPER_TELEOP_TO_DEGREES
|
||||
else:
|
||||
left_goals[target] = -val if target in LEFT_MOTORS_TO_FLIP else val
|
||||
base = key.removesuffix(".pos")
|
||||
# JOINT_REMAP is symmetric (its own inverse).
|
||||
target = JOINT_REMAP.get(base, base)
|
||||
if base == "gripper":
|
||||
# Convert robot degrees to teleop 0-100: 0°→0, -65°→100
|
||||
goals[target] = val / GRIPPER_TELEOP_TO_DEGREES
|
||||
else:
|
||||
# Un-flip using the ORIGINAL motor name (target = leader motor)
|
||||
goals[target] = -val if target in self._motors_to_flip else val
|
||||
|
||||
if right_goals:
|
||||
self.bus_right.sync_write("Goal_Position", right_goals)
|
||||
if left_goals:
|
||||
self.bus_left.sync_write("Goal_Position", left_goals)
|
||||
if goals:
|
||||
self.bus.sync_write("Goal_Position", goals)
|
||||
|
||||
@check_if_not_connected
|
||||
def send_feedback(self, feedback: dict[str, float]) -> None:
|
||||
@@ -354,6 +265,5 @@ class OpenArmMini(Teleoperator):
|
||||
|
||||
@check_if_not_connected
|
||||
def disconnect(self) -> None:
|
||||
self.bus_right.disconnect()
|
||||
self.bus_left.disconnect()
|
||||
self.bus.disconnect()
|
||||
logger.info(f"{self} disconnected.")
|
||||
|
||||
@@ -99,14 +99,18 @@ def make_teleoperator_from_config(config: TeleoperatorConfig) -> "Teleoperator":
|
||||
from .openarm_mini import OpenArmMini
|
||||
|
||||
return OpenArmMini(config)
|
||||
elif config.type == "bi_openarm_mini":
|
||||
from .bi_openarm_mini import BiOpenArmMini
|
||||
|
||||
return BiOpenArmMini(config)
|
||||
elif config.type == "rebot_102_leader":
|
||||
from .rebot_102_leader import RebotArm102Leader
|
||||
|
||||
return RebotArm102Leader(config)
|
||||
elif config.type == "bi_rebot_102_leader":
|
||||
from .bi_rebot_102_leader import BiRebotArm102Leader
|
||||
from .bi_rebot_102_leader import BiRebot102Leader
|
||||
|
||||
return BiRebotArm102Leader(config)
|
||||
return BiRebot102Leader(config)
|
||||
else:
|
||||
try:
|
||||
return cast("Teleoperator", make_device_from_device_class(config))
|
||||
|
||||
@@ -0,0 +1,63 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from typing import Any
|
||||
|
||||
from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
|
||||
|
||||
|
||||
class BimanualMixin:
|
||||
"""Lifecycle delegation for bimanual robots and teleoperators.
|
||||
|
||||
Concrete subclasses must populate ``self.left_arm`` and ``self.right_arm`` in
|
||||
their own ``__init__``. They retain ownership of feature dicts and the
|
||||
data-routing methods (``get_action`` / ``send_action`` / ``get_observation`` /
|
||||
``send_feedback``), which vary per-embodiment.
|
||||
|
||||
Inherit before the ``Robot`` / ``Teleoperator`` base so the mixin's methods
|
||||
take precedence in the MRO::
|
||||
|
||||
class BiFooFollower(BimanualMixin, Robot): ...
|
||||
"""
|
||||
|
||||
left_arm: Any
|
||||
right_arm: Any
|
||||
|
||||
@property
|
||||
def is_connected(self) -> bool:
|
||||
return self.left_arm.is_connected and self.right_arm.is_connected
|
||||
|
||||
@property
|
||||
def is_calibrated(self) -> bool:
|
||||
return self.left_arm.is_calibrated and self.right_arm.is_calibrated
|
||||
|
||||
@check_if_already_connected
|
||||
def connect(self, calibrate: bool = True) -> None:
|
||||
self.left_arm.connect(calibrate)
|
||||
self.right_arm.connect(calibrate)
|
||||
|
||||
def calibrate(self) -> None:
|
||||
self.left_arm.calibrate()
|
||||
self.right_arm.calibrate()
|
||||
|
||||
def configure(self) -> None:
|
||||
self.left_arm.configure()
|
||||
self.right_arm.configure()
|
||||
|
||||
@check_if_not_connected
|
||||
def disconnect(self) -> None:
|
||||
self.left_arm.disconnect()
|
||||
self.right_arm.disconnect()
|
||||
@@ -38,8 +38,6 @@ def init_rerun(
|
||||
require_package("rerun-sdk", extra="viz", import_name="rerun")
|
||||
import rerun as rr
|
||||
|
||||
log_rerun_data.blueprint = None # Reset blueprint cache for new session
|
||||
|
||||
batch_size = os.getenv("RERUN_FLUSH_NUM_BYTES", "8000")
|
||||
os.environ["RERUN_FLUSH_NUM_BYTES"] = batch_size
|
||||
rr.init(session_name)
|
||||
@@ -65,38 +63,6 @@ def _is_scalar(x):
|
||||
)
|
||||
|
||||
|
||||
def _build_blueprint(observation_paths: set[str], action_paths: set[str], image_paths: set[str]):
|
||||
"""Build a Rerun blueprint laying out camera images, observation and action scalars in separate views.
|
||||
|
||||
Camera images, observation and action scalars are arranged in a grid.
|
||||
"""
|
||||
|
||||
# Safe + zero-overhead: `log_rerun_data` already ran the `require_package` guard and imported rerun.
|
||||
import rerun.blueprint as rrb
|
||||
|
||||
views = [rrb.Spatial2DView(origin=path, name=path) for path in sorted(image_paths)]
|
||||
|
||||
if observation_paths:
|
||||
views.append(rrb.TimeSeriesView(name="observation", contents=sorted(observation_paths)))
|
||||
if action_paths:
|
||||
views.append(rrb.TimeSeriesView(name="action", contents=sorted(action_paths)))
|
||||
|
||||
return rrb.Blueprint(rrb.Grid(*views))
|
||||
|
||||
|
||||
def _ensure_blueprint(observation_paths: set[str], action_paths: set[str], image_paths: set[str]) -> None:
|
||||
"""Build and send the blueprint once, from the first observation and action data."""
|
||||
if getattr(log_rerun_data, "blueprint", None) is not None:
|
||||
return
|
||||
|
||||
# Safe + zero-overhead: `log_rerun_data` already ran the `require_package` guard and imported rerun.
|
||||
import rerun as rr
|
||||
|
||||
blueprint = _build_blueprint(observation_paths, action_paths, image_paths)
|
||||
log_rerun_data.blueprint = blueprint
|
||||
rr.send_blueprint(blueprint)
|
||||
|
||||
|
||||
def log_rerun_data(
|
||||
observation: RobotObservation | None = None,
|
||||
action: RobotAction | None = None,
|
||||
@@ -110,15 +76,11 @@ def log_rerun_data(
|
||||
- Scalars values (floats, ints) are logged as `rr.Scalars`.
|
||||
- 3D NumPy arrays that resemble images (e.g., with 1, 3, or 4 channels first) are transposed
|
||||
from CHW to HWC format, (optionally) compressed to JPEG and logged as `rr.Image` or `rr.EncodedImage`.
|
||||
- 1D NumPy arrays are logged as a single `rr.Scalars` batch under one entity path, so that every
|
||||
dimension shares the same view instead of being split across one view per element.
|
||||
- Multi-dimensional **action** arrays are flattened and logged as a single `rr.Scalars` batch.
|
||||
- 1D NumPy arrays are logged as a series of individual scalars, with each element indexed.
|
||||
- Other multi-dimensional arrays are flattened and logged as individual scalars.
|
||||
|
||||
Keys are automatically namespaced with "observation." or "action." if not already present.
|
||||
|
||||
On the first call, a blueprint is built and sent so observation and action scalars get separate
|
||||
time-series views and each image gets its own spatial view.
|
||||
|
||||
Args:
|
||||
observation: An optional dictionary containing observation data to log.
|
||||
action: An optional dictionary containing action data to log.
|
||||
@@ -128,10 +90,6 @@ def log_rerun_data(
|
||||
require_package("rerun-sdk", extra="viz", import_name="rerun")
|
||||
import rerun as rr
|
||||
|
||||
observation_paths: set[str] = set()
|
||||
action_paths: set[str] = set()
|
||||
image_paths: set[str] = set()
|
||||
|
||||
if observation:
|
||||
for k, v in observation.items():
|
||||
if v is None:
|
||||
@@ -140,19 +98,17 @@ def log_rerun_data(
|
||||
|
||||
if _is_scalar(v):
|
||||
rr.log(key, rr.Scalars(float(v)))
|
||||
observation_paths.add(key)
|
||||
elif isinstance(v, np.ndarray):
|
||||
arr = v
|
||||
# Convert CHW -> HWC when needed
|
||||
if arr.ndim == 3 and arr.shape[0] in (1, 3, 4) and arr.shape[-1] not in (1, 3, 4):
|
||||
arr = np.transpose(arr, (1, 2, 0))
|
||||
if arr.ndim == 1:
|
||||
rr.log(key, rr.Scalars(arr.astype(float)))
|
||||
observation_paths.add(key)
|
||||
for i, vi in enumerate(arr):
|
||||
rr.log(f"{key}_{i}", rr.Scalars(float(vi)))
|
||||
else:
|
||||
img_entity = rr.Image(arr).compress() if compress_images else rr.Image(arr)
|
||||
rr.log(key, entity=img_entity, static=True)
|
||||
image_paths.add(key)
|
||||
|
||||
if action:
|
||||
for k, v in action.items():
|
||||
@@ -162,9 +118,12 @@ def log_rerun_data(
|
||||
|
||||
if _is_scalar(v):
|
||||
rr.log(key, rr.Scalars(float(v)))
|
||||
action_paths.add(key)
|
||||
elif isinstance(v, np.ndarray):
|
||||
rr.log(key, rr.Scalars(v.reshape(-1).astype(float)))
|
||||
action_paths.add(key)
|
||||
|
||||
_ensure_blueprint(observation_paths, action_paths, image_paths)
|
||||
if v.ndim == 1:
|
||||
for i, vi in enumerate(v):
|
||||
rr.log(f"{key}_{i}", rr.Scalars(float(vi)))
|
||||
else:
|
||||
# Fall back to flattening higher-dimensional arrays
|
||||
flat = v.flatten()
|
||||
for i, vi in enumerate(flat):
|
||||
rr.log(f"{key}_{i}", rr.Scalars(float(vi)))
|
||||
|
||||
@@ -114,6 +114,30 @@ def test_shuffle():
|
||||
assert set(sampler) == {0, 1, 2, 3, 4, 5}
|
||||
|
||||
|
||||
def test_shuffle_with_generator_is_deterministic():
|
||||
# Two samplers shuffling with same-seed generators must yield identical permutations.
|
||||
# This is what keeps batch shards disjoint across ranks in distributed training, where
|
||||
# accelerate synchronizes the sampler's generator state instead of the global torch RNG.
|
||||
sampler_a = EpisodeAwareSampler([0], [6], shuffle=True, generator=torch.Generator().manual_seed(42))
|
||||
sampler_b = EpisodeAwareSampler([0], [6], shuffle=True, generator=torch.Generator().manual_seed(42))
|
||||
assert list(sampler_a) == list(sampler_b)
|
||||
|
||||
# Desyncing the global RNG must not affect the permutation.
|
||||
sampler_c = EpisodeAwareSampler([0], [6], shuffle=True, generator=torch.Generator().manual_seed(42))
|
||||
order_before = list(sampler_c)
|
||||
sampler_c.generator.manual_seed(42)
|
||||
torch.randperm(1000) # consume global RNG, as rank-asymmetric code (e.g. eval) would
|
||||
assert list(sampler_c) == order_before
|
||||
|
||||
|
||||
def test_generator_attribute_defaults_to_none():
|
||||
# accelerate detects synchronizable samplers via `hasattr(sampler, "generator")`,
|
||||
# so the attribute must exist even when no generator is passed.
|
||||
sampler = EpisodeAwareSampler([0], [6], shuffle=True)
|
||||
assert sampler.generator is None
|
||||
assert set(sampler) == {0, 1, 2, 3, 4, 5}
|
||||
|
||||
|
||||
def test_negative_drop_first_frames_raises():
|
||||
with pytest.raises(ValueError, match="drop_n_first_frames must be >= 0"):
|
||||
EpisodeAwareSampler([0], [10], drop_n_first_frames=-1)
|
||||
|
||||
@@ -18,7 +18,7 @@ from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from lerobot.teleoperators.bi_rebot_102_leader import BiRebotArm102Leader, BiRebotArm102LeaderConfig
|
||||
from lerobot.teleoperators.bi_rebot_102_leader import BiRebot102Leader, BiRebot102LeaderConfig
|
||||
from lerobot.teleoperators.rebot_102_leader import (
|
||||
RebotArm102Leader,
|
||||
RebotArm102LeaderConfig,
|
||||
@@ -91,11 +91,11 @@ def test_send_feedback_not_implemented(leader):
|
||||
|
||||
def test_bimanual_prefixes_features():
|
||||
with patch(f"{_MODULE}.require_package", lambda *a, **kw: None):
|
||||
cfg = BiRebotArm102LeaderConfig(
|
||||
cfg = BiRebot102LeaderConfig(
|
||||
left_arm_config=RebotArm102LeaderConfig(port="/dev/null0"),
|
||||
right_arm_config=RebotArm102LeaderConfig(port="/dev/null1"),
|
||||
)
|
||||
teleop = BiRebotArm102Leader(cfg)
|
||||
teleop = BiRebot102Leader(cfg)
|
||||
assert any(k.startswith("left_") for k in teleop.action_features)
|
||||
assert any(k.startswith("right_") for k in teleop.action_features)
|
||||
assert "left_gripper.pos" in teleop.action_features
|
||||
|
||||
@@ -30,46 +30,25 @@ from lerobot.utils.constants import OBS_STATE
|
||||
@pytest.fixture
|
||||
def mock_rerun(monkeypatch):
|
||||
"""
|
||||
Provide a mock `rerun` module (and `rerun.blueprint` submodule) so tests don't
|
||||
depend on the real library. Also reload the module-under-test so it binds to
|
||||
this mock `rr`.
|
||||
Provide a mock `rerun` module so tests don't depend on the real library.
|
||||
Also reload the module-under-test so it binds to this mock `rr`.
|
||||
"""
|
||||
calls = []
|
||||
blueprints = []
|
||||
|
||||
class DummyScalar:
|
||||
def __init__(self, value):
|
||||
# Scalars may be built from a single float or from a 1D array batch.
|
||||
self.value = value
|
||||
self.value = float(value)
|
||||
|
||||
class DummyImage:
|
||||
def __init__(self, arr):
|
||||
self.arr = arr
|
||||
|
||||
def compress(self, *a, **k):
|
||||
return self
|
||||
|
||||
def dummy_log(key, obj=None, **kwargs):
|
||||
# Accept either positional `obj` or keyword `entity` and record remaining kwargs.
|
||||
if obj is None and "entity" in kwargs:
|
||||
obj = kwargs.pop("entity")
|
||||
calls.append((key, obj, kwargs))
|
||||
|
||||
def dummy_send_blueprint(blueprint, *a, **k):
|
||||
blueprints.append(blueprint)
|
||||
|
||||
# Mock the `rerun.blueprint` submodule used to build the layout.
|
||||
dummy_rrb = SimpleNamespace(
|
||||
Spatial2DView=lambda origin=None, name=None: SimpleNamespace(
|
||||
kind="Spatial2DView", origin=origin, name=name
|
||||
),
|
||||
TimeSeriesView=lambda name=None, contents=None: SimpleNamespace(
|
||||
kind="TimeSeriesView", name=name, contents=contents
|
||||
),
|
||||
Grid=lambda *views: SimpleNamespace(kind="Grid", views=list(views)),
|
||||
Blueprint=lambda root: SimpleNamespace(kind="Blueprint", root=root),
|
||||
)
|
||||
|
||||
dummy_rr = SimpleNamespace(
|
||||
__name__="rerun",
|
||||
__package__="rerun",
|
||||
@@ -77,23 +56,20 @@ def mock_rerun(monkeypatch):
|
||||
Scalars=DummyScalar,
|
||||
Image=DummyImage,
|
||||
log=dummy_log,
|
||||
send_blueprint=dummy_send_blueprint,
|
||||
init=lambda *a, **k: None,
|
||||
spawn=lambda *a, **k: None,
|
||||
blueprint=dummy_rrb,
|
||||
)
|
||||
|
||||
# Inject fake modules into sys.modules (both `rerun` and `rerun.blueprint`).
|
||||
# Inject fake module into sys.modules
|
||||
monkeypatch.setitem(sys.modules, "rerun", dummy_rr)
|
||||
monkeypatch.setitem(sys.modules, "rerun.blueprint", dummy_rrb)
|
||||
|
||||
# Now import and reload the module under test, to bind to our rerun mock
|
||||
import lerobot.utils.visualization_utils as vu
|
||||
|
||||
importlib.reload(vu)
|
||||
|
||||
# Expose the reloaded module, the call recorder and the captured blueprints
|
||||
yield vu, calls, blueprints
|
||||
# Expose both the reloaded module and the call recorder
|
||||
yield vu, calls
|
||||
|
||||
|
||||
def _keys(calls):
|
||||
@@ -116,13 +92,8 @@ def _kwargs_for(calls, key):
|
||||
raise KeyError(f"Key {key} not found in calls: {calls}")
|
||||
|
||||
|
||||
def _views_by_kind(blueprint, kind):
|
||||
"""Return the views of a given kind from the (single) blueprint's grid."""
|
||||
return [v for v in blueprint.root.views if v.kind == kind]
|
||||
|
||||
|
||||
def test_log_rerun_data_envtransition_scalars_and_image(mock_rerun):
|
||||
vu, calls, blueprints = mock_rerun
|
||||
vu, calls = mock_rerun
|
||||
|
||||
# Build EnvTransition dict
|
||||
obs = {
|
||||
@@ -132,7 +103,7 @@ def test_log_rerun_data_envtransition_scalars_and_image(mock_rerun):
|
||||
}
|
||||
act = {
|
||||
"action.throttle": 0.7,
|
||||
# 1D array should be logged as a single Scalars batch under one entity path
|
||||
# 1D array should log individual Scalars with suffix _i
|
||||
"action.vector": np.array([1.0, 2.0], dtype=np.float32),
|
||||
}
|
||||
transition = {
|
||||
@@ -149,28 +120,31 @@ def test_log_rerun_data_envtransition_scalars_and_image(mock_rerun):
|
||||
# - observation.state.temperature -> Scalars
|
||||
# - observation.camera -> Image (HWC) with static=True
|
||||
# - action.throttle -> Scalars
|
||||
# - action.vector -> single Scalars batch (no per-element suffix)
|
||||
# - action.vector_0, action.vector_1 -> Scalars
|
||||
expected_keys = {
|
||||
f"{OBS_STATE}.temperature",
|
||||
"observation.camera",
|
||||
"action.throttle",
|
||||
"action.vector",
|
||||
"action.vector_0",
|
||||
"action.vector_1",
|
||||
}
|
||||
assert set(_keys(calls)) == expected_keys
|
||||
|
||||
# Check scalar types and values
|
||||
temp_obj = _obj_for(calls, f"{OBS_STATE}.temperature")
|
||||
assert type(temp_obj).__name__ == "DummyScalar"
|
||||
assert float(temp_obj.value) == pytest.approx(25.0)
|
||||
assert temp_obj.value == pytest.approx(25.0)
|
||||
|
||||
throttle_obj = _obj_for(calls, "action.throttle")
|
||||
assert type(throttle_obj).__name__ == "DummyScalar"
|
||||
assert float(throttle_obj.value) == pytest.approx(0.7)
|
||||
assert throttle_obj.value == pytest.approx(0.7)
|
||||
|
||||
# 1D vector logged as a single batched Scalars under one entity path
|
||||
vec = _obj_for(calls, "action.vector")
|
||||
assert type(vec).__name__ == "DummyScalar"
|
||||
np.testing.assert_allclose(np.asarray(vec.value), [1.0, 2.0])
|
||||
v0 = _obj_for(calls, "action.vector_0")
|
||||
v1 = _obj_for(calls, "action.vector_1")
|
||||
assert type(v0).__name__ == "DummyScalar"
|
||||
assert type(v1).__name__ == "DummyScalar"
|
||||
assert v0.value == pytest.approx(1.0)
|
||||
assert v1.value == pytest.approx(2.0)
|
||||
|
||||
# Check image handling: CHW -> HWC
|
||||
img_obj = _obj_for(calls, "observation.camera")
|
||||
@@ -178,24 +152,9 @@ def test_log_rerun_data_envtransition_scalars_and_image(mock_rerun):
|
||||
assert img_obj.arr.shape == (10, 20, 3) # transposed
|
||||
assert _kwargs_for(calls, "observation.camera").get("static", False) is True # static=True for images
|
||||
|
||||
# A blueprint should have been built and sent exactly once, and cached on the function.
|
||||
assert len(blueprints) == 1
|
||||
assert vu.log_rerun_data.blueprint is blueprints[0]
|
||||
|
||||
bp = blueprints[0]
|
||||
# One spatial view per image path
|
||||
spatial_views = _views_by_kind(bp, "Spatial2DView")
|
||||
assert {v.origin for v in spatial_views} == {"observation.camera"}
|
||||
|
||||
# One time-series view each for observation and action scalars
|
||||
ts_views = {v.name: v for v in _views_by_kind(bp, "TimeSeriesView")}
|
||||
assert set(ts_views) == {"observation", "action"}
|
||||
assert ts_views["observation"].contents == [f"{OBS_STATE}.temperature"]
|
||||
assert ts_views["action"].contents == ["action.throttle", "action.vector"]
|
||||
|
||||
|
||||
def test_log_rerun_data_plain_list_ordering_and_prefixes(mock_rerun):
|
||||
vu, calls, blueprints = mock_rerun
|
||||
vu, calls = mock_rerun
|
||||
|
||||
# First dict without prefixes treated as observation
|
||||
# Second dict without prefixes treated as action
|
||||
@@ -214,12 +173,14 @@ def test_log_rerun_data_plain_list_ordering_and_prefixes(mock_rerun):
|
||||
# First dict was treated as observation, second as action
|
||||
vu.log_rerun_data(observation=obs_plain, action=act_plain)
|
||||
|
||||
# Expected keys with auto-prefixes. The 1D vector is a single batched Scalars.
|
||||
# Expected keys with auto-prefixes
|
||||
expected = {
|
||||
"observation.temp",
|
||||
"observation.img",
|
||||
"action.throttle",
|
||||
"action.vec",
|
||||
"action.vec_0",
|
||||
"action.vec_1",
|
||||
"action.vec_2",
|
||||
}
|
||||
logged = set(_keys(calls))
|
||||
assert logged == expected
|
||||
@@ -227,11 +188,11 @@ def test_log_rerun_data_plain_list_ordering_and_prefixes(mock_rerun):
|
||||
# Scalars
|
||||
t = _obj_for(calls, "observation.temp")
|
||||
assert type(t).__name__ == "DummyScalar"
|
||||
assert float(t.value) == pytest.approx(1.5)
|
||||
assert t.value == pytest.approx(1.5)
|
||||
|
||||
throttle = _obj_for(calls, "action.throttle")
|
||||
assert type(throttle).__name__ == "DummyScalar"
|
||||
assert float(throttle.value) == pytest.approx(0.3)
|
||||
assert throttle.value == pytest.approx(0.3)
|
||||
|
||||
# Image stays HWC
|
||||
img = _obj_for(calls, "observation.img")
|
||||
@@ -239,23 +200,15 @@ def test_log_rerun_data_plain_list_ordering_and_prefixes(mock_rerun):
|
||||
assert img.arr.shape == (5, 6, 3)
|
||||
assert _kwargs_for(calls, "observation.img").get("static", False) is True
|
||||
|
||||
# Vector logged as a single batched Scalars under one entity path
|
||||
vec = _obj_for(calls, "action.vec")
|
||||
assert type(vec).__name__ == "DummyScalar"
|
||||
np.testing.assert_allclose(np.asarray(vec.value), [9, 8, 7])
|
||||
|
||||
# Blueprint sent once with the expected view layout
|
||||
assert len(blueprints) == 1
|
||||
bp = blueprints[0]
|
||||
spatial_views = _views_by_kind(bp, "Spatial2DView")
|
||||
assert {v.origin for v in spatial_views} == {"observation.img"}
|
||||
ts_views = {v.name: v for v in _views_by_kind(bp, "TimeSeriesView")}
|
||||
assert ts_views["observation"].contents == ["observation.temp"]
|
||||
assert ts_views["action"].contents == ["action.throttle", "action.vec"]
|
||||
# Vectors
|
||||
for i, val in enumerate([9, 8, 7]):
|
||||
o = _obj_for(calls, f"action.vec_{i}")
|
||||
assert type(o).__name__ == "DummyScalar"
|
||||
assert o.value == pytest.approx(val)
|
||||
|
||||
|
||||
def test_log_rerun_data_kwargs_only(mock_rerun):
|
||||
vu, calls, blueprints = mock_rerun
|
||||
vu, calls = mock_rerun
|
||||
|
||||
vu.log_rerun_data(
|
||||
observation={"observation.temp": 10.0, "observation.gray": np.zeros((8, 8, 1), dtype=np.uint8)},
|
||||
@@ -269,7 +222,7 @@ def test_log_rerun_data_kwargs_only(mock_rerun):
|
||||
|
||||
temp = _obj_for(calls, "observation.temp")
|
||||
assert type(temp).__name__ == "DummyScalar"
|
||||
assert float(temp.value) == pytest.approx(10.0)
|
||||
assert temp.value == pytest.approx(10.0)
|
||||
|
||||
img = _obj_for(calls, "observation.gray")
|
||||
assert type(img).__name__ == "DummyImage"
|
||||
@@ -278,26 +231,4 @@ def test_log_rerun_data_kwargs_only(mock_rerun):
|
||||
|
||||
a = _obj_for(calls, "action.a")
|
||||
assert type(a).__name__ == "DummyScalar"
|
||||
assert float(a.value) == pytest.approx(1.0)
|
||||
|
||||
# Blueprint sent once, with a spatial view for the image and time-series views for scalars
|
||||
assert len(blueprints) == 1
|
||||
bp = blueprints[0]
|
||||
assert {v.origin for v in _views_by_kind(bp, "Spatial2DView")} == {"observation.gray"}
|
||||
ts_views = {v.name: v for v in _views_by_kind(bp, "TimeSeriesView")}
|
||||
assert ts_views["observation"].contents == ["observation.temp"]
|
||||
assert ts_views["action"].contents == ["action.a"]
|
||||
|
||||
|
||||
def test_log_rerun_data_blueprint_sent_only_once(mock_rerun):
|
||||
"""The blueprint is built from the first call and not resent on subsequent calls."""
|
||||
vu, calls, blueprints = mock_rerun
|
||||
|
||||
vu.log_rerun_data(observation={"temp": 1.0}, action={"a": 2.0})
|
||||
assert len(blueprints) == 1
|
||||
first_blueprint = vu.log_rerun_data.blueprint
|
||||
|
||||
vu.log_rerun_data(observation={"temp": 3.0}, action={"a": 4.0})
|
||||
# Still only one blueprint, and the cached one is unchanged.
|
||||
assert len(blueprints) == 1
|
||||
assert vu.log_rerun_data.blueprint is first_blueprint
|
||||
assert a.value == pytest.approx(1.0)
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
version = 1
|
||||
revision = 2
|
||||
revision = 3
|
||||
requires-python = ">=3.12"
|
||||
resolution-markers = [
|
||||
"(python_full_version >= '3.15' and platform_machine == 'AMD64' and sys_platform == 'linux') or (python_full_version >= '3.15' and platform_machine == 'x86_64' and sys_platform == 'linux')",
|
||||
@@ -3257,7 +3257,7 @@ requires-dist = [
|
||||
{ name = "qwen-vl-utils", marker = "extra == 'qwen-vl-utils-dep'", specifier = ">=0.0.11,<0.1.0" },
|
||||
{ name = "reachy2-sdk", marker = "extra == 'reachy2'", specifier = ">=1.0.15,<1.1.0" },
|
||||
{ name = "requests", specifier = ">=2.32.0,<3.0.0" },
|
||||
{ name = "rerun-sdk", marker = "extra == 'viz'", specifier = ">=0.24.0,<0.34.0" },
|
||||
{ name = "rerun-sdk", marker = "extra == 'viz'", specifier = ">=0.24.0,<0.27.0" },
|
||||
{ name = "ruff", marker = "extra == 'dev'", specifier = ">=0.14.1" },
|
||||
{ name = "safetensors", specifier = ">=0.4.3,<1.0.0" },
|
||||
{ name = "scikit-image", marker = "extra == 'video-benchmark'", specifier = ">=0.23.2,<0.26.0" },
|
||||
@@ -5636,21 +5636,21 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "rerun-sdk"
|
||||
version = "0.33.0"
|
||||
version = "0.26.2"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "attrs" },
|
||||
{ name = "numpy" },
|
||||
{ name = "pillow" },
|
||||
{ name = "psutil" },
|
||||
{ name = "pyarrow" },
|
||||
{ name = "typing-extensions" },
|
||||
]
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/31/17/5a521e86ac0064bd0f452e3e98e2422433511b54110423c0217d2cc1234f/rerun_sdk-0.33.0-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:97f123e3ef6aa69b60194bc566e5435c7d4040757ed4f58297ea46c8ef320c5c", size = 125707606, upload-time = "2026-05-29T09:42:53.584Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/34/2f/2ca2599aca03b69fbcac7c8391ef50376968edd7c58b96de53a4b7f20624/rerun_sdk-0.33.0-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:8f734cf59419dcfbc46915bea6cec030224f16e96c3a597f0ccf7cb7b058dd43", size = 135271020, upload-time = "2026-05-29T09:43:00.106Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/2e/ba/d70997b43e6db4f58c4326c29c6a6a384ddc6c2fe125f231c885ad9b3b1f/rerun_sdk-0.33.0-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:53d95609f8b330026bcd041bf6d11b46ee1c18b6fbde155135f291fe86328eeb", size = 139552018, upload-time = "2026-05-29T09:43:06.275Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/14/a5/0cac294d16aff6c9a2f183f838428a0380b4d2fd9e053bb37b3041999ad5/rerun_sdk-0.33.0-cp310-abi3-win_amd64.whl", hash = "sha256:b152992a72ec240062c8c285bd30ab681b464a25efbe1464c66fdac82320de1f", size = 120418186, upload-time = "2026-05-29T09:43:13.733Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/4b/4a/767c20e1529d74d9be5b5e55c6c26b63a6918ef3c1709fc422d08a460114/rerun_sdk-0.26.2-cp39-abi3-macosx_10_12_x86_64.whl", hash = "sha256:3d4151c9a3484e112b53d1df90c8fa07397dc7b8bfbb420f09e011eff20f1ef2", size = 93349439, upload-time = "2025-10-27T11:34:10.745Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/2b/3d/d8dd0af9c287a85d51ec99d69406cc4b94a9feb1d6f192d3bbcaac9f0b81/rerun_sdk-0.26.2-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:03977d2aba4966d9a70b682eca196123fda11408fecd733441ede9916c6341e2", size = 86323042, upload-time = "2025-10-27T11:34:17.995Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/13/29/53d8d98799ab32418fd4ba6834d6a5749c31f56160d3c87f52a7219887e9/rerun_sdk-0.26.2-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:b6128c3c4f014cae5be18e4d37657c5932d1bcdb2ce5e9d4b488a6eed47f7437", size = 92677274, upload-time = "2025-10-27T11:34:22.601Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/f5/86/0b9c8f56398b4fc85f8e99279907c258413a297e5603f8f2537fe5806e51/rerun_sdk-0.26.2-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:a6f97b60aaa7d4e8c6124a3f6b97ce9dbd09520050955f0e0bdacb72b0eb106a", size = 98768129, upload-time = "2025-10-27T11:34:27.36Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/be/e7/99fc91c0f99f69d7d43e1db0a6f6cb8273ffc02111539bfc1fee43749bad/rerun_sdk-0.26.2-cp39-abi3-win_amd64.whl", hash = "sha256:a493ad6c8357022cba2ca6f8954a81d0faf984b0b22154eb1d976bfc7649df63", size = 84267089, upload-time = "2025-10-27T11:34:32.023Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
||||
Reference in New Issue
Block a user