e23b41e79a
Co-authored-by: Pepijn <138571049+pkooij@users.noreply.github.com> Co-authored-by: Steven Palma <imstevenpmwork@ieee.org> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Steven Palma <steven.palma@huggingface.co> Co-authored-by: Adil Zouitine <adilzouitinegm@gmail.com> Co-authored-by: Pepijn <pepijn@huggingface.co>
106 lines
2.9 KiB
Python
106 lines
2.9 KiB
Python
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
"""
|
|
Replays the actions of an episode from a dataset on a robot.
|
|
|
|
Example:
|
|
|
|
```shell
|
|
python -m lerobot.replay \
|
|
--robot.type=so100_follower \
|
|
--robot.port=/dev/tty.usbmodem58760431541 \
|
|
--robot.id=black \
|
|
--dataset.repo_id=aliberts/record-test \
|
|
--dataset.episode=2
|
|
```
|
|
"""
|
|
|
|
import logging
|
|
import time
|
|
from dataclasses import asdict, dataclass
|
|
from pathlib import Path
|
|
from pprint import pformat
|
|
|
|
import draccus
|
|
|
|
from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
|
|
from lerobot.common.robots import ( # noqa: F401
|
|
Robot,
|
|
RobotConfig,
|
|
koch_follower,
|
|
make_robot_from_config,
|
|
so100_follower,
|
|
so101_follower,
|
|
)
|
|
from lerobot.common.utils.robot_utils import busy_wait
|
|
from lerobot.common.utils.utils import (
|
|
init_logging,
|
|
log_say,
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class DatasetReplayConfig:
|
|
# Dataset identifier. By convention it should match '{hf_username}/{dataset_name}' (e.g. `lerobot/test`).
|
|
repo_id: str
|
|
# Episode to replay.
|
|
episode: int
|
|
# Root directory where the dataset will be stored (e.g. 'dataset/path').
|
|
root: str | Path | None = None
|
|
# Limit the frames per second. By default, uses the policy fps.
|
|
fps: int = 30
|
|
|
|
|
|
@dataclass
|
|
class ReplayConfig:
|
|
robot: RobotConfig
|
|
dataset: DatasetReplayConfig
|
|
# Use vocal synthesis to read events.
|
|
play_sounds: bool = True
|
|
|
|
|
|
@draccus.wrap()
|
|
def replay(cfg: ReplayConfig):
|
|
init_logging()
|
|
logging.info(pformat(asdict(cfg)))
|
|
|
|
robot = make_robot_from_config(cfg.robot)
|
|
dataset = LeRobotDataset(cfg.dataset.repo_id, root=cfg.dataset.root, episodes=[cfg.dataset.episode])
|
|
actions = dataset.hf_dataset.select_columns("action")
|
|
robot.connect()
|
|
|
|
log_say("Replaying episode", cfg.play_sounds, blocking=True)
|
|
for idx in range(dataset.num_frames):
|
|
start_episode_t = time.perf_counter()
|
|
|
|
action_array = actions[idx]["action"]
|
|
action = {}
|
|
for i, name in enumerate(dataset.features["action"]["names"]):
|
|
key = f"{name.removeprefix('main_')}.pos"
|
|
action[key] = action_array[i].item()
|
|
|
|
action["shoulder_lift.pos"] = -(action["shoulder_lift.pos"] - 90)
|
|
action["elbow_flex.pos"] -= 90
|
|
robot.send_action(action)
|
|
|
|
dt_s = time.perf_counter() - start_episode_t
|
|
busy_wait(1 / dataset.fps - dt_s)
|
|
|
|
robot.disconnect()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
replay()
|