Delete examples
Browse files- examples/1_load_lerobot_dataset.py +0 -148
- examples/2_evaluate_pretrained_policy.py +0 -139
- examples/3_train_policy.py +0 -120
- examples/4_train_policy_with_script.md +0 -274
- examples/advanced/1_add_image_transforms.py +0 -67
- examples/advanced/2_calculate_validation_loss.py +0 -104
- examples/backward_compatibility/replay.py +0 -105
- examples/lekiwi/evaluate.py +0 -32
- examples/lekiwi/record.py +0 -67
- examples/lekiwi/replay.py +0 -25
- examples/lekiwi/teleoperate.py +0 -32
examples/1_load_lerobot_dataset.py
DELETED
|
@@ -1,148 +0,0 @@
|
|
| 1 |
-
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
|
| 2 |
-
#
|
| 3 |
-
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
-
# you may not use this file except in compliance with the License.
|
| 5 |
-
# You may obtain a copy of the License at
|
| 6 |
-
#
|
| 7 |
-
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
-
#
|
| 9 |
-
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
-
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
-
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
-
# See the License for the specific language governing permissions and
|
| 13 |
-
# limitations under the License.
|
| 14 |
-
|
| 15 |
-
"""
|
| 16 |
-
This script demonstrates the use of `LeRobotDataset` class for handling and processing robotic datasets from Hugging Face.
|
| 17 |
-
It illustrates how to load datasets, manipulate them, and apply transformations suitable for machine learning tasks in PyTorch.
|
| 18 |
-
|
| 19 |
-
Features included in this script:
|
| 20 |
-
- Viewing a dataset's metadata and exploring its properties.
|
| 21 |
-
- Loading an existing dataset from the hub or a subset of it.
|
| 22 |
-
- Accessing frames by episode number.
|
| 23 |
-
- Using advanced dataset features like timestamp-based frame selection.
|
| 24 |
-
- Demonstrating compatibility with PyTorch DataLoader for batch processing.
|
| 25 |
-
|
| 26 |
-
The script ends with examples of how to batch process data using PyTorch's DataLoader.
|
| 27 |
-
"""
|
| 28 |
-
|
| 29 |
-
from pprint import pprint
|
| 30 |
-
|
| 31 |
-
import torch
|
| 32 |
-
from huggingface_hub import HfApi
|
| 33 |
-
|
| 34 |
-
import lerobot
|
| 35 |
-
from lerobot.common.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata
|
| 36 |
-
|
| 37 |
-
# We ported a number of existing datasets ourselves, use this to see the list:
|
| 38 |
-
print("List of available datasets:")
|
| 39 |
-
pprint(lerobot.available_datasets)
|
| 40 |
-
|
| 41 |
-
# You can also browse through the datasets created/ported by the community on the hub using the hub api:
|
| 42 |
-
hub_api = HfApi()
|
| 43 |
-
repo_ids = [info.id for info in hub_api.list_datasets(task_categories="robotics", tags=["LeRobot"])]
|
| 44 |
-
pprint(repo_ids)
|
| 45 |
-
|
| 46 |
-
# Or simply explore them in your web browser directly at:
|
| 47 |
-
# https://huggingface.co/datasets?other=LeRobot
|
| 48 |
-
|
| 49 |
-
# Let's take this one for this example
|
| 50 |
-
repo_id = "lerobot/aloha_mobile_cabinet"
|
| 51 |
-
# We can have a look and fetch its metadata to know more about it:
|
| 52 |
-
ds_meta = LeRobotDatasetMetadata(repo_id)
|
| 53 |
-
|
| 54 |
-
# By instantiating just this class, you can quickly access useful information about the content and the
|
| 55 |
-
# structure of the dataset without downloading the actual data yet (only metadata files — which are
|
| 56 |
-
# lightweight).
|
| 57 |
-
print(f"Total number of episodes: {ds_meta.total_episodes}")
|
| 58 |
-
print(f"Average number of frames per episode: {ds_meta.total_frames / ds_meta.total_episodes:.3f}")
|
| 59 |
-
print(f"Frames per second used during data collection: {ds_meta.fps}")
|
| 60 |
-
print(f"Robot type: {ds_meta.robot_type}")
|
| 61 |
-
print(f"keys to access images from cameras: {ds_meta.camera_keys=}\n")
|
| 62 |
-
|
| 63 |
-
print("Tasks:")
|
| 64 |
-
print(ds_meta.tasks)
|
| 65 |
-
print("Features:")
|
| 66 |
-
pprint(ds_meta.features)
|
| 67 |
-
|
| 68 |
-
# You can also get a short summary by simply printing the object:
|
| 69 |
-
print(ds_meta)
|
| 70 |
-
|
| 71 |
-
# You can then load the actual dataset from the hub.
|
| 72 |
-
# Either load any subset of episodes:
|
| 73 |
-
dataset = LeRobotDataset(repo_id, episodes=[0, 10, 11, 23])
|
| 74 |
-
|
| 75 |
-
# And see how many frames you have:
|
| 76 |
-
print(f"Selected episodes: {dataset.episodes}")
|
| 77 |
-
print(f"Number of episodes selected: {dataset.num_episodes}")
|
| 78 |
-
print(f"Number of frames selected: {dataset.num_frames}")
|
| 79 |
-
|
| 80 |
-
# Or simply load the entire dataset:
|
| 81 |
-
dataset = LeRobotDataset(repo_id)
|
| 82 |
-
print(f"Number of episodes selected: {dataset.num_episodes}")
|
| 83 |
-
print(f"Number of frames selected: {dataset.num_frames}")
|
| 84 |
-
|
| 85 |
-
# The previous metadata class is contained in the 'meta' attribute of the dataset:
|
| 86 |
-
print(dataset.meta)
|
| 87 |
-
|
| 88 |
-
# LeRobotDataset actually wraps an underlying Hugging Face dataset
|
| 89 |
-
# (see https://huggingface.co/docs/datasets for more information).
|
| 90 |
-
print(dataset.hf_dataset)
|
| 91 |
-
|
| 92 |
-
# LeRobot datasets also subclasses PyTorch datasets so you can do everything you know and love from working
|
| 93 |
-
# with the latter, like iterating through the dataset.
|
| 94 |
-
# The __getitem__ iterates over the frames of the dataset. Since our datasets are also structured by
|
| 95 |
-
# episodes, you can access the frame indices of any episode using the episode_data_index. Here, we access
|
| 96 |
-
# frame indices associated to the first episode:
|
| 97 |
-
episode_index = 0
|
| 98 |
-
from_idx = dataset.episode_data_index["from"][episode_index].item()
|
| 99 |
-
to_idx = dataset.episode_data_index["to"][episode_index].item()
|
| 100 |
-
|
| 101 |
-
# Then we grab all the image frames from the first camera:
|
| 102 |
-
camera_key = dataset.meta.camera_keys[0]
|
| 103 |
-
frames = [dataset[idx][camera_key] for idx in range(from_idx, to_idx)]
|
| 104 |
-
|
| 105 |
-
# The objects returned by the dataset are all torch.Tensors
|
| 106 |
-
print(type(frames[0]))
|
| 107 |
-
print(frames[0].shape)
|
| 108 |
-
|
| 109 |
-
# Since we're using pytorch, the shape is in pytorch, channel-first convention (c, h, w).
|
| 110 |
-
# We can compare this shape with the information available for that feature
|
| 111 |
-
pprint(dataset.features[camera_key])
|
| 112 |
-
# In particular:
|
| 113 |
-
print(dataset.features[camera_key]["shape"])
|
| 114 |
-
# The shape is in (h, w, c) which is a more universal format.
|
| 115 |
-
|
| 116 |
-
# For many machine learning applications we need to load the history of past observations or trajectories of
|
| 117 |
-
# future actions. Our datasets can load previous and future frames for each key/modality, using timestamps
|
| 118 |
-
# differences with the current loaded frame. For instance:
|
| 119 |
-
delta_timestamps = {
|
| 120 |
-
# loads 4 images: 1 second before current frame, 500 ms before, 200 ms before, and current frame
|
| 121 |
-
camera_key: [-1, -0.5, -0.20, 0],
|
| 122 |
-
# loads 6 state vectors: 1.5 seconds before, 1 second before, ... 200 ms, 100 ms, and current frame
|
| 123 |
-
"observation.state": [-1.5, -1, -0.5, -0.20, -0.10, 0],
|
| 124 |
-
# loads 64 action vectors: current frame, 1 frame in the future, 2 frames, ... 63 frames in the future
|
| 125 |
-
"action": [t / dataset.fps for t in range(64)],
|
| 126 |
-
}
|
| 127 |
-
# Note that in any case, these delta_timestamps values need to be multiples of (1/fps) so that added to any
|
| 128 |
-
# timestamp, you still get a valid timestamp.
|
| 129 |
-
|
| 130 |
-
dataset = LeRobotDataset(repo_id, delta_timestamps=delta_timestamps)
|
| 131 |
-
print(f"\n{dataset[0][camera_key].shape=}") # (4, c, h, w)
|
| 132 |
-
print(f"{dataset[0]['observation.state'].shape=}") # (6, c)
|
| 133 |
-
print(f"{dataset[0]['action'].shape=}\n") # (64, c)
|
| 134 |
-
|
| 135 |
-
# Finally, our datasets are fully compatible with PyTorch dataloaders and samplers because they are just
|
| 136 |
-
# PyTorch datasets.
|
| 137 |
-
dataloader = torch.utils.data.DataLoader(
|
| 138 |
-
dataset,
|
| 139 |
-
num_workers=0,
|
| 140 |
-
batch_size=32,
|
| 141 |
-
shuffle=True,
|
| 142 |
-
)
|
| 143 |
-
|
| 144 |
-
for batch in dataloader:
|
| 145 |
-
print(f"{batch[camera_key].shape=}") # (32, 4, c, h, w)
|
| 146 |
-
print(f"{batch['observation.state'].shape=}") # (32, 6, c)
|
| 147 |
-
print(f"{batch['action'].shape=}") # (32, 64, c)
|
| 148 |
-
break
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
examples/2_evaluate_pretrained_policy.py
DELETED
|
@@ -1,139 +0,0 @@
|
|
| 1 |
-
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
|
| 2 |
-
#
|
| 3 |
-
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
-
# you may not use this file except in compliance with the License.
|
| 5 |
-
# You may obtain a copy of the License at
|
| 6 |
-
#
|
| 7 |
-
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
-
#
|
| 9 |
-
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
-
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
-
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
-
# See the License for the specific language governing permissions and
|
| 13 |
-
# limitations under the License.
|
| 14 |
-
|
| 15 |
-
"""
|
| 16 |
-
This script demonstrates how to evaluate a pretrained policy from the HuggingFace Hub or from your local
|
| 17 |
-
training outputs directory. In the latter case, you might want to run examples/3_train_policy.py first.
|
| 18 |
-
|
| 19 |
-
It requires the installation of the 'gym_pusht' simulation environment. Install it by running:
|
| 20 |
-
```bash
|
| 21 |
-
pip install -e ".[pusht]"
|
| 22 |
-
```
|
| 23 |
-
"""
|
| 24 |
-
|
| 25 |
-
from pathlib import Path
|
| 26 |
-
|
| 27 |
-
import gym_pusht # noqa: F401
|
| 28 |
-
import gymnasium as gym
|
| 29 |
-
import imageio
|
| 30 |
-
import numpy
|
| 31 |
-
import torch
|
| 32 |
-
|
| 33 |
-
from lerobot.common.policies.diffusion.modeling_diffusion import DiffusionPolicy
|
| 34 |
-
|
| 35 |
-
# Create a directory to store the video of the evaluation
|
| 36 |
-
output_directory = Path("outputs/eval/example_pusht_diffusion")
|
| 37 |
-
output_directory.mkdir(parents=True, exist_ok=True)
|
| 38 |
-
|
| 39 |
-
# Select your device
|
| 40 |
-
device = "cuda"
|
| 41 |
-
|
| 42 |
-
# Provide the [hugging face repo id](https://huggingface.co/lerobot/diffusion_pusht):
|
| 43 |
-
pretrained_policy_path = "lerobot/diffusion_pusht"
|
| 44 |
-
# OR a path to a local outputs/train folder.
|
| 45 |
-
# pretrained_policy_path = Path("outputs/train/example_pusht_diffusion")
|
| 46 |
-
|
| 47 |
-
policy = DiffusionPolicy.from_pretrained(pretrained_policy_path)
|
| 48 |
-
|
| 49 |
-
# Initialize evaluation environment to render two observation types:
|
| 50 |
-
# an image of the scene and state/position of the agent. The environment
|
| 51 |
-
# also automatically stops running after 300 interactions/steps.
|
| 52 |
-
env = gym.make(
|
| 53 |
-
"gym_pusht/PushT-v0",
|
| 54 |
-
obs_type="pixels_agent_pos",
|
| 55 |
-
max_episode_steps=300,
|
| 56 |
-
)
|
| 57 |
-
|
| 58 |
-
# We can verify that the shapes of the features expected by the policy match the ones from the observations
|
| 59 |
-
# produced by the environment
|
| 60 |
-
print(policy.config.input_features)
|
| 61 |
-
print(env.observation_space)
|
| 62 |
-
|
| 63 |
-
# Similarly, we can check that the actions produced by the policy will match the actions expected by the
|
| 64 |
-
# environment
|
| 65 |
-
print(policy.config.output_features)
|
| 66 |
-
print(env.action_space)
|
| 67 |
-
|
| 68 |
-
# Reset the policy and environments to prepare for rollout
|
| 69 |
-
policy.reset()
|
| 70 |
-
numpy_observation, info = env.reset(seed=42)
|
| 71 |
-
|
| 72 |
-
# Prepare to collect every rewards and all the frames of the episode,
|
| 73 |
-
# from initial state to final state.
|
| 74 |
-
rewards = []
|
| 75 |
-
frames = []
|
| 76 |
-
|
| 77 |
-
# Render frame of the initial state
|
| 78 |
-
frames.append(env.render())
|
| 79 |
-
|
| 80 |
-
step = 0
|
| 81 |
-
done = False
|
| 82 |
-
while not done:
|
| 83 |
-
# Prepare observation for the policy running in Pytorch
|
| 84 |
-
state = torch.from_numpy(numpy_observation["agent_pos"])
|
| 85 |
-
image = torch.from_numpy(numpy_observation["pixels"])
|
| 86 |
-
|
| 87 |
-
# Convert to float32 with image from channel first in [0,255]
|
| 88 |
-
# to channel last in [0,1]
|
| 89 |
-
state = state.to(torch.float32)
|
| 90 |
-
image = image.to(torch.float32) / 255
|
| 91 |
-
image = image.permute(2, 0, 1)
|
| 92 |
-
|
| 93 |
-
# Send data tensors from CPU to GPU
|
| 94 |
-
state = state.to(device, non_blocking=True)
|
| 95 |
-
image = image.to(device, non_blocking=True)
|
| 96 |
-
|
| 97 |
-
# Add extra (empty) batch dimension, required to forward the policy
|
| 98 |
-
state = state.unsqueeze(0)
|
| 99 |
-
image = image.unsqueeze(0)
|
| 100 |
-
|
| 101 |
-
# Create the policy input dictionary
|
| 102 |
-
observation = {
|
| 103 |
-
"observation.state": state,
|
| 104 |
-
"observation.image": image,
|
| 105 |
-
}
|
| 106 |
-
|
| 107 |
-
# Predict the next action with respect to the current observation
|
| 108 |
-
with torch.inference_mode():
|
| 109 |
-
action = policy.select_action(observation)
|
| 110 |
-
|
| 111 |
-
# Prepare the action for the environment
|
| 112 |
-
numpy_action = action.squeeze(0).to("cpu").numpy()
|
| 113 |
-
|
| 114 |
-
# Step through the environment and receive a new observation
|
| 115 |
-
numpy_observation, reward, terminated, truncated, info = env.step(numpy_action)
|
| 116 |
-
print(f"{step=} {reward=} {terminated=}")
|
| 117 |
-
|
| 118 |
-
# Keep track of all the rewards and frames
|
| 119 |
-
rewards.append(reward)
|
| 120 |
-
frames.append(env.render())
|
| 121 |
-
|
| 122 |
-
# The rollout is considered done when the success state is reached (i.e. terminated is True),
|
| 123 |
-
# or the maximum number of iterations is reached (i.e. truncated is True)
|
| 124 |
-
done = terminated | truncated | done
|
| 125 |
-
step += 1
|
| 126 |
-
|
| 127 |
-
if terminated:
|
| 128 |
-
print("Success!")
|
| 129 |
-
else:
|
| 130 |
-
print("Failure!")
|
| 131 |
-
|
| 132 |
-
# Get the speed of environment (i.e. its number of frames per second).
|
| 133 |
-
fps = env.metadata["render_fps"]
|
| 134 |
-
|
| 135 |
-
# Encode all frames into a mp4 video.
|
| 136 |
-
video_path = output_directory / "rollout.mp4"
|
| 137 |
-
imageio.mimsave(str(video_path), numpy.stack(frames), fps=fps)
|
| 138 |
-
|
| 139 |
-
print(f"Video of the evaluation is available in '{video_path}'.")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
examples/3_train_policy.py
DELETED
|
@@ -1,120 +0,0 @@
|
|
| 1 |
-
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
|
| 2 |
-
#
|
| 3 |
-
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
-
# you may not use this file except in compliance with the License.
|
| 5 |
-
# You may obtain a copy of the License at
|
| 6 |
-
#
|
| 7 |
-
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
-
#
|
| 9 |
-
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
-
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
-
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
-
# See the License for the specific language governing permissions and
|
| 13 |
-
# limitations under the License.
|
| 14 |
-
|
| 15 |
-
"""This script demonstrates how to train Diffusion Policy on the PushT environment.
|
| 16 |
-
|
| 17 |
-
Once you have trained a model with this script, you can try to evaluate it on
|
| 18 |
-
examples/2_evaluate_pretrained_policy.py
|
| 19 |
-
"""
|
| 20 |
-
|
| 21 |
-
from pathlib import Path
|
| 22 |
-
|
| 23 |
-
import torch
|
| 24 |
-
|
| 25 |
-
from lerobot.common.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata
|
| 26 |
-
from lerobot.common.datasets.utils import dataset_to_policy_features
|
| 27 |
-
from lerobot.common.policies.diffusion.configuration_diffusion import DiffusionConfig
|
| 28 |
-
from lerobot.common.policies.diffusion.modeling_diffusion import DiffusionPolicy
|
| 29 |
-
from lerobot.configs.types import FeatureType
|
| 30 |
-
|
| 31 |
-
|
| 32 |
-
def main():
|
| 33 |
-
# Create a directory to store the training checkpoint.
|
| 34 |
-
output_directory = Path("outputs/train/example_pusht_diffusion")
|
| 35 |
-
output_directory.mkdir(parents=True, exist_ok=True)
|
| 36 |
-
|
| 37 |
-
# # Select your device
|
| 38 |
-
device = torch.device("cuda")
|
| 39 |
-
|
| 40 |
-
# Number of offline training steps (we'll only do offline training for this example.)
|
| 41 |
-
# Adjust as you prefer. 5000 steps are needed to get something worth evaluating.
|
| 42 |
-
training_steps = 5000
|
| 43 |
-
log_freq = 1
|
| 44 |
-
|
| 45 |
-
# When starting from scratch (i.e. not from a pretrained policy), we need to specify 2 things before
|
| 46 |
-
# creating the policy:
|
| 47 |
-
# - input/output shapes: to properly size the policy
|
| 48 |
-
# - dataset stats: for normalization and denormalization of input/outputs
|
| 49 |
-
dataset_metadata = LeRobotDatasetMetadata("lerobot/pusht")
|
| 50 |
-
features = dataset_to_policy_features(dataset_metadata.features)
|
| 51 |
-
output_features = {key: ft for key, ft in features.items() if ft.type is FeatureType.ACTION}
|
| 52 |
-
input_features = {key: ft for key, ft in features.items() if key not in output_features}
|
| 53 |
-
|
| 54 |
-
# Policies are initialized with a configuration class, in this case `DiffusionConfig`. For this example,
|
| 55 |
-
# we'll just use the defaults and so no arguments other than input/output features need to be passed.
|
| 56 |
-
cfg = DiffusionConfig(input_features=input_features, output_features=output_features)
|
| 57 |
-
|
| 58 |
-
# We can now instantiate our policy with this config and the dataset stats.
|
| 59 |
-
policy = DiffusionPolicy(cfg, dataset_stats=dataset_metadata.stats)
|
| 60 |
-
policy.train()
|
| 61 |
-
policy.to(device)
|
| 62 |
-
|
| 63 |
-
# Another policy-dataset interaction is with the delta_timestamps. Each policy expects a given number frames
|
| 64 |
-
# which can differ for inputs, outputs and rewards (if there are some).
|
| 65 |
-
delta_timestamps = {
|
| 66 |
-
"observation.image": [i / dataset_metadata.fps for i in cfg.observation_delta_indices],
|
| 67 |
-
"observation.state": [i / dataset_metadata.fps for i in cfg.observation_delta_indices],
|
| 68 |
-
"action": [i / dataset_metadata.fps for i in cfg.action_delta_indices],
|
| 69 |
-
}
|
| 70 |
-
|
| 71 |
-
# In this case with the standard configuration for Diffusion Policy, it is equivalent to this:
|
| 72 |
-
delta_timestamps = {
|
| 73 |
-
# Load the previous image and state at -0.1 seconds before current frame,
|
| 74 |
-
# then load current image and state corresponding to 0.0 second.
|
| 75 |
-
"observation.image": [-0.1, 0.0],
|
| 76 |
-
"observation.state": [-0.1, 0.0],
|
| 77 |
-
# Load the previous action (-0.1), the next action to be executed (0.0),
|
| 78 |
-
# and 14 future actions with a 0.1 seconds spacing. All these actions will be
|
| 79 |
-
# used to supervise the policy.
|
| 80 |
-
"action": [-0.1, 0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 1.1, 1.2, 1.3, 1.4],
|
| 81 |
-
}
|
| 82 |
-
|
| 83 |
-
# We can then instantiate the dataset with these delta_timestamps configuration.
|
| 84 |
-
dataset = LeRobotDataset("lerobot/pusht", delta_timestamps=delta_timestamps)
|
| 85 |
-
|
| 86 |
-
# Then we create our optimizer and dataloader for offline training.
|
| 87 |
-
optimizer = torch.optim.Adam(policy.parameters(), lr=1e-4)
|
| 88 |
-
dataloader = torch.utils.data.DataLoader(
|
| 89 |
-
dataset,
|
| 90 |
-
num_workers=4,
|
| 91 |
-
batch_size=64,
|
| 92 |
-
shuffle=True,
|
| 93 |
-
pin_memory=device.type != "cpu",
|
| 94 |
-
drop_last=True,
|
| 95 |
-
)
|
| 96 |
-
|
| 97 |
-
# Run training loop.
|
| 98 |
-
step = 0
|
| 99 |
-
done = False
|
| 100 |
-
while not done:
|
| 101 |
-
for batch in dataloader:
|
| 102 |
-
batch = {k: (v.to(device) if isinstance(v, torch.Tensor) else v) for k, v in batch.items()}
|
| 103 |
-
loss, _ = policy.forward(batch)
|
| 104 |
-
loss.backward()
|
| 105 |
-
optimizer.step()
|
| 106 |
-
optimizer.zero_grad()
|
| 107 |
-
|
| 108 |
-
if step % log_freq == 0:
|
| 109 |
-
print(f"step: {step} loss: {loss.item():.3f}")
|
| 110 |
-
step += 1
|
| 111 |
-
if step >= training_steps:
|
| 112 |
-
done = True
|
| 113 |
-
break
|
| 114 |
-
|
| 115 |
-
# Save a policy checkpoint.
|
| 116 |
-
policy.save_pretrained(output_directory)
|
| 117 |
-
|
| 118 |
-
|
| 119 |
-
if __name__ == "__main__":
|
| 120 |
-
main()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
examples/4_train_policy_with_script.md
DELETED
|
@@ -1,274 +0,0 @@
|
|
| 1 |
-
This tutorial will explain the training script, how to use it, and particularly how to configure everything needed for the training run.
|
| 2 |
-
> **Note:** The following assumes you're running these commands on a machine equipped with a cuda GPU. If you don't have one (or if you're using a Mac), you can add `--policy.device=cpu` (`--policy.device=mps` respectively). However, be advised that the code executes much slower on cpu.
|
| 3 |
-
|
| 4 |
-
|
| 5 |
-
## The training script
|
| 6 |
-
|
| 7 |
-
LeRobot offers a training script at [`lerobot/scripts/train.py`](../lerobot/scripts/train.py). At a high level it does the following:
|
| 8 |
-
|
| 9 |
-
- Initialize/load a configuration for the following steps using.
|
| 10 |
-
- Instantiates a dataset.
|
| 11 |
-
- (Optional) Instantiates a simulation environment corresponding to that dataset.
|
| 12 |
-
- Instantiates a policy.
|
| 13 |
-
- Runs a standard training loop with forward pass, backward pass, optimization step, and occasional logging, evaluation (of the policy on the environment), and checkpointing.
|
| 14 |
-
|
| 15 |
-
## Overview of the configuration system
|
| 16 |
-
|
| 17 |
-
In the training script, the main function `train` expects a `TrainPipelineConfig` object:
|
| 18 |
-
```python
|
| 19 |
-
# train.py
|
| 20 |
-
@parser.wrap()
|
| 21 |
-
def train(cfg: TrainPipelineConfig):
|
| 22 |
-
```
|
| 23 |
-
|
| 24 |
-
You can inspect the `TrainPipelineConfig` defined in [`lerobot/configs/train.py`](../lerobot/configs/train.py) (which is heavily commented and meant to be a reference to understand any option)
|
| 25 |
-
|
| 26 |
-
When running the script, inputs for the command line are parsed thanks to the `@parser.wrap()` decorator and an instance of this class is automatically generated. Under the hood, this is done with [Draccus](https://github.com/dlwh/draccus) which is a tool dedicated to this purpose. If you're familiar with Hydra, Draccus can similarly load configurations from config files (.json, .yaml) and also override their values through command line inputs. Unlike Hydra, these configurations are pre-defined in the code through dataclasses rather than being defined entirely in config files. This allows for more rigorous serialization/deserialization, typing, and to manipulate configuration as objects directly in the code and not as dictionaries or namespaces (which enables nice features in an IDE such as autocomplete, jump-to-def, etc.)
|
| 27 |
-
|
| 28 |
-
Let's have a look at a simplified example. Amongst other attributes, the training config has the following attributes:
|
| 29 |
-
```python
|
| 30 |
-
@dataclass
|
| 31 |
-
class TrainPipelineConfig:
|
| 32 |
-
dataset: DatasetConfig
|
| 33 |
-
env: envs.EnvConfig | None = None
|
| 34 |
-
policy: PreTrainedConfig | None = None
|
| 35 |
-
```
|
| 36 |
-
in which `DatasetConfig` for example is defined as such:
|
| 37 |
-
```python
|
| 38 |
-
@dataclass
|
| 39 |
-
class DatasetConfig:
|
| 40 |
-
repo_id: str
|
| 41 |
-
episodes: list[int] | None = None
|
| 42 |
-
video_backend: str = "pyav"
|
| 43 |
-
```
|
| 44 |
-
|
| 45 |
-
This creates a hierarchical relationship where, for example assuming we have a `cfg` instance of `TrainPipelineConfig`, we can access the `repo_id` value with `cfg.dataset.repo_id`.
|
| 46 |
-
From the command line, we can specify this value by using a very similar syntax `--dataset.repo_id=repo/id`.
|
| 47 |
-
|
| 48 |
-
By default, every field takes its default value specified in the dataclass. If a field doesn't have a default value, it needs to be specified either from the command line or from a config file – which path is also given in the command line (more in this below). In the example above, the `dataset` field doesn't have a default value which means it must be specified.
|
| 49 |
-
|
| 50 |
-
|
| 51 |
-
## Specifying values from the CLI
|
| 52 |
-
|
| 53 |
-
Let's say that we want to train [Diffusion Policy](../lerobot/common/policies/diffusion) on the [pusht](https://huggingface.co/datasets/lerobot/pusht) dataset, using the [gym_pusht](https://github.com/huggingface/gym-pusht) environment for evaluation. The command to do so would look like this:
|
| 54 |
-
```bash
|
| 55 |
-
python lerobot/scripts/train.py \
|
| 56 |
-
--dataset.repo_id=lerobot/pusht \
|
| 57 |
-
--policy.type=diffusion \
|
| 58 |
-
--env.type=pusht
|
| 59 |
-
```
|
| 60 |
-
|
| 61 |
-
Let's break this down:
|
| 62 |
-
- To specify the dataset, we just need to specify its `repo_id` on the hub which is the only required argument in the `DatasetConfig`. The rest of the fields have default values and in this case we are fine with those so we can just add the option `--dataset.repo_id=lerobot/pusht`.
|
| 63 |
-
- To specify the policy, we can just select diffusion policy using `--policy` appended with `.type`. Here, `.type` is a special argument which allows us to select config classes inheriting from `draccus.ChoiceRegistry` and that have been decorated with the `register_subclass()` method. To have a better explanation of this feature, have a look at this [Draccus demo](https://github.com/dlwh/draccus?tab=readme-ov-file#more-flexible-configuration-with-choice-types). In our code, we use this mechanism mainly to select policies, environments, robots, and some other components like optimizers. The policies available to select are located in [lerobot/common/policies](../lerobot/common/policies)
|
| 64 |
-
- Similarly, we select the environment with `--env.type=pusht`. The different environment configs are available in [`lerobot/common/envs/configs.py`](../lerobot/common/envs/configs.py)
|
| 65 |
-
|
| 66 |
-
Let's see another example. Let's say you've been training [ACT](../lerobot/common/policies/act) on [lerobot/aloha_sim_insertion_human](https://huggingface.co/datasets/lerobot/aloha_sim_insertion_human) using the [gym-aloha](https://github.com/huggingface/gym-aloha) environment for evaluation with:
|
| 67 |
-
```bash
|
| 68 |
-
python lerobot/scripts/train.py \
|
| 69 |
-
--policy.type=act \
|
| 70 |
-
--dataset.repo_id=lerobot/aloha_sim_insertion_human \
|
| 71 |
-
--env.type=aloha \
|
| 72 |
-
--output_dir=outputs/train/act_aloha_insertion
|
| 73 |
-
```
|
| 74 |
-
> Notice we added `--output_dir` to explicitly tell where to write outputs from this run (checkpoints, training state, configs etc.). This is not mandatory and if you don't specify it, a default directory will be created from the current date and time, env.type and policy.type. This will typically look like `outputs/train/2025-01-24/16-10-05_aloha_act`.
|
| 75 |
-
|
| 76 |
-
We now want to train a different policy for aloha on another task. We'll change the dataset and use [lerobot/aloha_sim_transfer_cube_human](https://huggingface.co/datasets/lerobot/aloha_sim_transfer_cube_human) instead. Of course, we also need to change the task of the environment as well to match this other task.
|
| 77 |
-
Looking at the [`AlohaEnv`](../lerobot/common/envs/configs.py) config, the task is `"AlohaInsertion-v0"` by default, which corresponds to the task we trained on in the command above. The [gym-aloha](https://github.com/huggingface/gym-aloha?tab=readme-ov-file#description) environment also has the `AlohaTransferCube-v0` task which corresponds to this other task we want to train on. Putting this together, we can train this new policy on this different task using:
|
| 78 |
-
```bash
|
| 79 |
-
python lerobot/scripts/train.py \
|
| 80 |
-
--policy.type=act \
|
| 81 |
-
--dataset.repo_id=lerobot/aloha_sim_transfer_cube_human \
|
| 82 |
-
--env.type=aloha \
|
| 83 |
-
--env.task=AlohaTransferCube-v0 \
|
| 84 |
-
--output_dir=outputs/train/act_aloha_transfer
|
| 85 |
-
```
|
| 86 |
-
|
| 87 |
-
## Loading from a config file
|
| 88 |
-
|
| 89 |
-
Now, let's assume that we want to reproduce the run just above. That run has produced a `train_config.json` file in its checkpoints, which serializes the `TrainPipelineConfig` instance it used:
|
| 90 |
-
```json
|
| 91 |
-
{
|
| 92 |
-
"dataset": {
|
| 93 |
-
"repo_id": "lerobot/aloha_sim_transfer_cube_human",
|
| 94 |
-
"episodes": null,
|
| 95 |
-
...
|
| 96 |
-
},
|
| 97 |
-
"env": {
|
| 98 |
-
"type": "aloha",
|
| 99 |
-
"task": "AlohaTransferCube-v0",
|
| 100 |
-
"fps": 50,
|
| 101 |
-
...
|
| 102 |
-
},
|
| 103 |
-
"policy": {
|
| 104 |
-
"type": "act",
|
| 105 |
-
"n_obs_steps": 1,
|
| 106 |
-
...
|
| 107 |
-
},
|
| 108 |
-
...
|
| 109 |
-
}
|
| 110 |
-
```
|
| 111 |
-
|
| 112 |
-
We can then simply load the config values from this file using:
|
| 113 |
-
```bash
|
| 114 |
-
python lerobot/scripts/train.py \
|
| 115 |
-
--config_path=outputs/train/act_aloha_transfer/checkpoints/last/pretrained_model/ \
|
| 116 |
-
--output_dir=outputs/train/act_aloha_transfer_2
|
| 117 |
-
```
|
| 118 |
-
`--config_path` is also a special argument which allows to initialize the config from a local config file. It can point to a directory that contains `train_config.json` or to the config file itself directly.
|
| 119 |
-
|
| 120 |
-
Similarly to Hydra, we can still override some parameters in the CLI if we want to, e.g.:
|
| 121 |
-
```bash
|
| 122 |
-
python lerobot/scripts/train.py \
|
| 123 |
-
--config_path=outputs/train/act_aloha_transfer/checkpoints/last/pretrained_model/ \
|
| 124 |
-
--output_dir=outputs/train/act_aloha_transfer_2
|
| 125 |
-
--policy.n_action_steps=80
|
| 126 |
-
```
|
| 127 |
-
> Note: While `--output_dir` is not required in general, in this case we need to specify it since it will otherwise take the value from the `train_config.json` (which is `outputs/train/act_aloha_transfer`). In order to prevent accidental deletion of previous run checkpoints, we raise an error if you're trying to write in an existing directory. This is not the case when resuming a run, which is what you'll learn next.
|
| 128 |
-
|
| 129 |
-
`--config_path` can also accept the repo_id of a repo on the hub that contains a `train_config.json` file, e.g. running:
|
| 130 |
-
```bash
|
| 131 |
-
python lerobot/scripts/train.py --config_path=lerobot/diffusion_pusht
|
| 132 |
-
```
|
| 133 |
-
will start a training run with the same configuration used for training [lerobot/diffusion_pusht](https://huggingface.co/lerobot/diffusion_pusht)
|
| 134 |
-
|
| 135 |
-
|
| 136 |
-
## Resume training
|
| 137 |
-
|
| 138 |
-
Being able to resume a training run is important in case it crashed or aborted for any reason. We'll demonstrate how to do that here.
|
| 139 |
-
|
| 140 |
-
Let's reuse the command from the previous run and add a few more options:
|
| 141 |
-
```bash
|
| 142 |
-
python lerobot/scripts/train.py \
|
| 143 |
-
--policy.type=act \
|
| 144 |
-
--dataset.repo_id=lerobot/aloha_sim_transfer_cube_human \
|
| 145 |
-
--env.type=aloha \
|
| 146 |
-
--env.task=AlohaTransferCube-v0 \
|
| 147 |
-
--log_freq=25 \
|
| 148 |
-
--save_freq=100 \
|
| 149 |
-
--output_dir=outputs/train/run_resumption
|
| 150 |
-
```
|
| 151 |
-
|
| 152 |
-
Here we've taken care to set up the log frequency and checkpointing frequency to low numbers so we can showcase resumption. You should be able to see some logging and have a first checkpoint within 1 minute (depending on hardware). Wait for the first checkpoint to happen, you should see a line that looks like this in your terminal:
|
| 153 |
-
```
|
| 154 |
-
INFO 2025-01-24 16:10:56 ts/train.py:263 Checkpoint policy after step 100
|
| 155 |
-
```
|
| 156 |
-
Now let's simulate a crash by killing the process (hit `ctrl`+`c`). We can then simply resume this run from the last checkpoint available with:
|
| 157 |
-
```bash
|
| 158 |
-
python lerobot/scripts/train.py \
|
| 159 |
-
--config_path=outputs/train/run_resumption/checkpoints/last/pretrained_model/ \
|
| 160 |
-
--resume=true
|
| 161 |
-
```
|
| 162 |
-
You should see from the logging that your training picks up from where it left off.
|
| 163 |
-
|
| 164 |
-
Another reason for which you might want to resume a run is simply to extend training and add more training steps. The number of training steps is set by the option `--steps`, which is 100 000 by default.
|
| 165 |
-
You could double the number of steps of the previous run with:
|
| 166 |
-
```bash
|
| 167 |
-
python lerobot/scripts/train.py \
|
| 168 |
-
--config_path=outputs/train/run_resumption/checkpoints/last/pretrained_model/ \
|
| 169 |
-
--resume=true \
|
| 170 |
-
--steps=200000
|
| 171 |
-
```
|
| 172 |
-
|
| 173 |
-
## Outputs of a run
|
| 174 |
-
In the output directory, there will be a folder called `checkpoints` with the following structure:
|
| 175 |
-
```bash
|
| 176 |
-
outputs/train/run_resumption/checkpoints
|
| 177 |
-
├── 000100 # checkpoint_dir for training step 100
|
| 178 |
-
│ ├── pretrained_model/
|
| 179 |
-
│ │ ├── config.json # policy config
|
| 180 |
-
│ │ ├── model.safetensors # policy weights
|
| 181 |
-
│ │ └── train_config.json # train config
|
| 182 |
-
│ └── training_state/
|
| 183 |
-
│ ├── optimizer_param_groups.json # optimizer param groups
|
| 184 |
-
│ ├── optimizer_state.safetensors # optimizer state
|
| 185 |
-
│ ├── rng_state.safetensors # rng states
|
| 186 |
-
│ ├── scheduler_state.json # scheduler state
|
| 187 |
-
│ └── training_step.json # training step
|
| 188 |
-
├── 000200
|
| 189 |
-
└── last -> 000200 # symlink to the last available checkpoint
|
| 190 |
-
```
|
| 191 |
-
|
| 192 |
-
## Fine-tuning a pre-trained policy
|
| 193 |
-
|
| 194 |
-
In addition to the features currently in Draccus, we've added a special `.path` argument for the policy, which allows to load a policy as you would with `PreTrainedPolicy.from_pretrained()`. In that case, `path` can be a local directory that contains a checkpoint or a repo_id pointing to a pretrained policy on the hub.
|
| 195 |
-
|
| 196 |
-
For example, we could fine-tune a [policy pre-trained on the aloha transfer task](https://huggingface.co/lerobot/act_aloha_sim_transfer_cube_human) on the aloha insertion task. We can achieve this with:
|
| 197 |
-
```bash
|
| 198 |
-
python lerobot/scripts/train.py \
|
| 199 |
-
--policy.path=lerobot/act_aloha_sim_transfer_cube_human \
|
| 200 |
-
--dataset.repo_id=lerobot/aloha_sim_insertion_human \
|
| 201 |
-
--env.type=aloha \
|
| 202 |
-
--env.task=AlohaInsertion-v0
|
| 203 |
-
```
|
| 204 |
-
|
| 205 |
-
When doing so, keep in mind that the features of the fine-tuning dataset would have to match the input/output features of the pretrained policy.
|
| 206 |
-
|
| 207 |
-
## Typical logs and metrics
|
| 208 |
-
|
| 209 |
-
When you start the training process, you will first see your full configuration being printed in the terminal. You can check it to make sure that you configured your run correctly. The final configuration will also be saved with the checkpoint.
|
| 210 |
-
|
| 211 |
-
After that, you will see training log like this one:
|
| 212 |
-
```
|
| 213 |
-
INFO 2024-08-14 13:35:12 ts/train.py:192 step:0 smpl:64 ep:1 epch:0.00 loss:1.112 grdn:15.387 lr:2.0e-07 updt_s:1.738 data_s:4.774
|
| 214 |
-
```
|
| 215 |
-
or evaluation log:
|
| 216 |
-
```
|
| 217 |
-
INFO 2024-08-14 13:38:45 ts/train.py:226 step:100 smpl:6K ep:52 epch:0.25 ∑rwrd:20.693 success:0.0% eval_s:120.266
|
| 218 |
-
```
|
| 219 |
-
|
| 220 |
-
These logs will also be saved in wandb if `wandb.enable` is set to `true`. Here are the meaning of some abbreviations:
|
| 221 |
-
- `smpl`: number of samples seen during training.
|
| 222 |
-
- `ep`: number of episodes seen during training. An episode contains multiple samples in a complete manipulation task.
|
| 223 |
-
- `epch`: number of time all unique samples are seen (epoch).
|
| 224 |
-
- `grdn`: gradient norm.
|
| 225 |
-
- `∑rwrd`: compute the sum of rewards in every evaluation episode and then take an average of them.
|
| 226 |
-
- `success`: average success rate of eval episodes. Reward and success are usually different except for the sparsing reward setting, where reward=1 only when the task is completed successfully.
|
| 227 |
-
- `eval_s`: time to evaluate the policy in the environment, in second.
|
| 228 |
-
- `updt_s`: time to update the network parameters, in second.
|
| 229 |
-
- `data_s`: time to load a batch of data, in second.
|
| 230 |
-
|
| 231 |
-
Some metrics are useful for initial performance profiling. For example, if you find the current GPU utilization is low via the `nvidia-smi` command and `data_s` sometimes is too high, you may need to modify batch size or number of dataloading workers to accelerate dataloading. We also recommend [pytorch profiler](https://github.com/huggingface/lerobot?tab=readme-ov-file#improve-your-code-with-profiling) for detailed performance probing.
|
| 232 |
-
|
| 233 |
-
## In short
|
| 234 |
-
|
| 235 |
-
We'll summarize here the main use cases to remember from this tutorial.
|
| 236 |
-
|
| 237 |
-
#### Train a policy from scratch – CLI
|
| 238 |
-
```bash
|
| 239 |
-
python lerobot/scripts/train.py \
|
| 240 |
-
--policy.type=act \ # <- select 'act' policy
|
| 241 |
-
--env.type=pusht \ # <- select 'pusht' environment
|
| 242 |
-
--dataset.repo_id=lerobot/pusht # <- train on this dataset
|
| 243 |
-
```
|
| 244 |
-
|
| 245 |
-
#### Train a policy from scratch - config file + CLI
|
| 246 |
-
```bash
|
| 247 |
-
python lerobot/scripts/train.py \
|
| 248 |
-
--config_path=path/to/pretrained_model \ # <- can also be a repo_id
|
| 249 |
-
--policy.n_action_steps=80 # <- you may still override values
|
| 250 |
-
```
|
| 251 |
-
|
| 252 |
-
#### Resume/continue a training run
|
| 253 |
-
```bash
|
| 254 |
-
python lerobot/scripts/train.py \
|
| 255 |
-
--config_path=checkpoint/pretrained_model/ \
|
| 256 |
-
--resume=true \
|
| 257 |
-
--steps=200000 # <- you can change some training parameters
|
| 258 |
-
```
|
| 259 |
-
|
| 260 |
-
#### Fine-tuning
|
| 261 |
-
```bash
|
| 262 |
-
python lerobot/scripts/train.py \
|
| 263 |
-
--policy.path=lerobot/act_aloha_sim_transfer_cube_human \ # <- can also be a local path to a checkpoint
|
| 264 |
-
--dataset.repo_id=lerobot/aloha_sim_insertion_human \
|
| 265 |
-
--env.type=aloha \
|
| 266 |
-
--env.task=AlohaInsertion-v0
|
| 267 |
-
```
|
| 268 |
-
|
| 269 |
-
---
|
| 270 |
-
|
| 271 |
-
Now that you know the basics of how to train a policy, you might want to know how to apply this knowledge to actual robots, or how to record your own datasets and train policies on your specific task?
|
| 272 |
-
If that's the case, head over to the next tutorial [`7_get_started_with_real_robot.md`](./7_get_started_with_real_robot.md).
|
| 273 |
-
|
| 274 |
-
Or in the meantime, happy training! 🤗
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
examples/advanced/1_add_image_transforms.py
DELETED
|
@@ -1,67 +0,0 @@
|
|
| 1 |
-
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
|
| 2 |
-
#
|
| 3 |
-
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
-
# you may not use this file except in compliance with the License.
|
| 5 |
-
# You may obtain a copy of the License at
|
| 6 |
-
#
|
| 7 |
-
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
-
#
|
| 9 |
-
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
-
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
-
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
-
# See the License for the specific language governing permissions and
|
| 13 |
-
# limitations under the License.
|
| 14 |
-
|
| 15 |
-
"""
|
| 16 |
-
This script demonstrates how to use torchvision's image transformation with LeRobotDataset for data
|
| 17 |
-
augmentation purposes. The transformations are passed to the dataset as an argument upon creation, and
|
| 18 |
-
transforms are applied to the observation images before they are returned in the dataset's __getitem__.
|
| 19 |
-
"""
|
| 20 |
-
|
| 21 |
-
from pathlib import Path
|
| 22 |
-
|
| 23 |
-
from torchvision.transforms import ToPILImage, v2
|
| 24 |
-
|
| 25 |
-
from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
|
| 26 |
-
|
| 27 |
-
dataset_repo_id = "lerobot/aloha_static_screw_driver"
|
| 28 |
-
|
| 29 |
-
# Create a LeRobotDataset with no transformations
|
| 30 |
-
dataset = LeRobotDataset(dataset_repo_id, episodes=[0])
|
| 31 |
-
# This is equivalent to `dataset = LeRobotDataset(dataset_repo_id, image_transforms=None)`
|
| 32 |
-
|
| 33 |
-
# Get the index of the first observation in the first episode
|
| 34 |
-
first_idx = dataset.episode_data_index["from"][0].item()
|
| 35 |
-
|
| 36 |
-
# Get the frame corresponding to the first camera
|
| 37 |
-
frame = dataset[first_idx][dataset.meta.camera_keys[0]]
|
| 38 |
-
|
| 39 |
-
|
| 40 |
-
# Define the transformations
|
| 41 |
-
transforms = v2.Compose(
|
| 42 |
-
[
|
| 43 |
-
v2.ColorJitter(brightness=(0.5, 1.5)),
|
| 44 |
-
v2.ColorJitter(contrast=(0.5, 1.5)),
|
| 45 |
-
v2.ColorJitter(hue=(-0.1, 0.1)),
|
| 46 |
-
v2.RandomAdjustSharpness(sharpness_factor=2, p=1),
|
| 47 |
-
]
|
| 48 |
-
)
|
| 49 |
-
|
| 50 |
-
# Create another LeRobotDataset with the defined transformations
|
| 51 |
-
transformed_dataset = LeRobotDataset(dataset_repo_id, episodes=[0], image_transforms=transforms)
|
| 52 |
-
|
| 53 |
-
# Get a frame from the transformed dataset
|
| 54 |
-
transformed_frame = transformed_dataset[first_idx][transformed_dataset.meta.camera_keys[0]]
|
| 55 |
-
|
| 56 |
-
# Create a directory to store output images
|
| 57 |
-
output_dir = Path("outputs/image_transforms")
|
| 58 |
-
output_dir.mkdir(parents=True, exist_ok=True)
|
| 59 |
-
|
| 60 |
-
# Save the original frame
|
| 61 |
-
to_pil = ToPILImage()
|
| 62 |
-
to_pil(frame).save(output_dir / "original_frame.png", quality=100)
|
| 63 |
-
print(f"Original frame saved to {output_dir / 'original_frame.png'}.")
|
| 64 |
-
|
| 65 |
-
# Save the transformed frame
|
| 66 |
-
to_pil(transformed_frame).save(output_dir / "transformed_frame.png", quality=100)
|
| 67 |
-
print(f"Transformed frame saved to {output_dir / 'transformed_frame.png'}.")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
examples/advanced/2_calculate_validation_loss.py
DELETED
|
@@ -1,104 +0,0 @@
|
|
| 1 |
-
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
|
| 2 |
-
#
|
| 3 |
-
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
-
# you may not use this file except in compliance with the License.
|
| 5 |
-
# You may obtain a copy of the License at
|
| 6 |
-
#
|
| 7 |
-
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
-
#
|
| 9 |
-
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
-
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
-
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
-
# See the License for the specific language governing permissions and
|
| 13 |
-
# limitations under the License.
|
| 14 |
-
|
| 15 |
-
"""This script demonstrates how to slice a dataset and calculate the loss on a subset of the data.
|
| 16 |
-
|
| 17 |
-
This technique can be useful for debugging and testing purposes, as well as identifying whether a policy
|
| 18 |
-
is learning effectively.
|
| 19 |
-
|
| 20 |
-
Furthermore, relying on validation loss to evaluate performance is generally not considered a good practice,
|
| 21 |
-
especially in the context of imitation learning. The most reliable approach is to evaluate the policy directly
|
| 22 |
-
on the target environment, whether that be in simulation or the real world.
|
| 23 |
-
"""
|
| 24 |
-
|
| 25 |
-
import math
|
| 26 |
-
|
| 27 |
-
import torch
|
| 28 |
-
|
| 29 |
-
from lerobot.common.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata
|
| 30 |
-
from lerobot.common.policies.diffusion.modeling_diffusion import DiffusionPolicy
|
| 31 |
-
|
| 32 |
-
|
| 33 |
-
def main():
|
| 34 |
-
device = torch.device("cuda")
|
| 35 |
-
|
| 36 |
-
# Download the diffusion policy for pusht environment
|
| 37 |
-
pretrained_policy_path = "lerobot/diffusion_pusht"
|
| 38 |
-
# OR uncomment the following to evaluate a policy from the local outputs/train folder.
|
| 39 |
-
# pretrained_policy_path = Path("outputs/train/example_pusht_diffusion")
|
| 40 |
-
|
| 41 |
-
policy = DiffusionPolicy.from_pretrained(pretrained_policy_path)
|
| 42 |
-
policy.eval()
|
| 43 |
-
policy.to(device)
|
| 44 |
-
|
| 45 |
-
# Set up the dataset.
|
| 46 |
-
delta_timestamps = {
|
| 47 |
-
# Load the previous image and state at -0.1 seconds before current frame,
|
| 48 |
-
# then load current image and state corresponding to 0.0 second.
|
| 49 |
-
"observation.image": [-0.1, 0.0],
|
| 50 |
-
"observation.state": [-0.1, 0.0],
|
| 51 |
-
# Load the previous action (-0.1), the next action to be executed (0.0),
|
| 52 |
-
# and 14 future actions with a 0.1 seconds spacing. All these actions will be
|
| 53 |
-
# used to calculate the loss.
|
| 54 |
-
"action": [-0.1, 0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 1.1, 1.2, 1.3, 1.4],
|
| 55 |
-
}
|
| 56 |
-
|
| 57 |
-
# Load the last 10% of episodes of the dataset as a validation set.
|
| 58 |
-
# - Load dataset metadata
|
| 59 |
-
dataset_metadata = LeRobotDatasetMetadata("lerobot/pusht")
|
| 60 |
-
# - Calculate train and val episodes
|
| 61 |
-
total_episodes = dataset_metadata.total_episodes
|
| 62 |
-
episodes = list(range(dataset_metadata.total_episodes))
|
| 63 |
-
num_train_episodes = math.floor(total_episodes * 90 / 100)
|
| 64 |
-
train_episodes = episodes[:num_train_episodes]
|
| 65 |
-
val_episodes = episodes[num_train_episodes:]
|
| 66 |
-
print(f"Number of episodes in full dataset: {total_episodes}")
|
| 67 |
-
print(f"Number of episodes in training dataset (90% subset): {len(train_episodes)}")
|
| 68 |
-
print(f"Number of episodes in validation dataset (10% subset): {len(val_episodes)}")
|
| 69 |
-
# - Load train and val datasets
|
| 70 |
-
train_dataset = LeRobotDataset(
|
| 71 |
-
"lerobot/pusht", episodes=train_episodes, delta_timestamps=delta_timestamps
|
| 72 |
-
)
|
| 73 |
-
val_dataset = LeRobotDataset("lerobot/pusht", episodes=val_episodes, delta_timestamps=delta_timestamps)
|
| 74 |
-
print(f"Number of frames in training dataset (90% subset): {len(train_dataset)}")
|
| 75 |
-
print(f"Number of frames in validation dataset (10% subset): {len(val_dataset)}")
|
| 76 |
-
|
| 77 |
-
# Create dataloader for evaluation.
|
| 78 |
-
val_dataloader = torch.utils.data.DataLoader(
|
| 79 |
-
val_dataset,
|
| 80 |
-
num_workers=4,
|
| 81 |
-
batch_size=64,
|
| 82 |
-
shuffle=False,
|
| 83 |
-
pin_memory=device != torch.device("cpu"),
|
| 84 |
-
drop_last=False,
|
| 85 |
-
)
|
| 86 |
-
|
| 87 |
-
# Run validation loop.
|
| 88 |
-
loss_cumsum = 0
|
| 89 |
-
n_examples_evaluated = 0
|
| 90 |
-
for batch in val_dataloader:
|
| 91 |
-
batch = {k: v.to(device, non_blocking=True) for k, v in batch.items()}
|
| 92 |
-
loss, _ = policy.forward(batch)
|
| 93 |
-
|
| 94 |
-
loss_cumsum += loss.item()
|
| 95 |
-
n_examples_evaluated += batch["index"].shape[0]
|
| 96 |
-
|
| 97 |
-
# Calculate the average loss over the validation set.
|
| 98 |
-
average_loss = loss_cumsum / n_examples_evaluated
|
| 99 |
-
|
| 100 |
-
print(f"Average loss on validation set: {average_loss:.4f}")
|
| 101 |
-
|
| 102 |
-
|
| 103 |
-
if __name__ == "__main__":
|
| 104 |
-
main()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
examples/backward_compatibility/replay.py
DELETED
|
@@ -1,105 +0,0 @@
|
|
| 1 |
-
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
|
| 2 |
-
#
|
| 3 |
-
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
-
# you may not use this file except in compliance with the License.
|
| 5 |
-
# You may obtain a copy of the License at
|
| 6 |
-
#
|
| 7 |
-
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
-
#
|
| 9 |
-
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
-
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
-
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
-
# See the License for the specific language governing permissions and
|
| 13 |
-
# limitations under the License.
|
| 14 |
-
|
| 15 |
-
"""
|
| 16 |
-
Replays the actions of an episode from a dataset on a robot.
|
| 17 |
-
|
| 18 |
-
Example:
|
| 19 |
-
|
| 20 |
-
```shell
|
| 21 |
-
python -m lerobot.replay \
|
| 22 |
-
--robot.type=so100_follower \
|
| 23 |
-
--robot.port=/dev/tty.usbmodem58760431541 \
|
| 24 |
-
--robot.id=black \
|
| 25 |
-
--dataset.repo_id=aliberts/record-test \
|
| 26 |
-
--dataset.episode=2
|
| 27 |
-
```
|
| 28 |
-
"""
|
| 29 |
-
|
| 30 |
-
import logging
|
| 31 |
-
import time
|
| 32 |
-
from dataclasses import asdict, dataclass
|
| 33 |
-
from pathlib import Path
|
| 34 |
-
from pprint import pformat
|
| 35 |
-
|
| 36 |
-
import draccus
|
| 37 |
-
|
| 38 |
-
from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
|
| 39 |
-
from lerobot.common.robots import ( # noqa: F401
|
| 40 |
-
Robot,
|
| 41 |
-
RobotConfig,
|
| 42 |
-
koch_follower,
|
| 43 |
-
make_robot_from_config,
|
| 44 |
-
so100_follower,
|
| 45 |
-
so101_follower,
|
| 46 |
-
)
|
| 47 |
-
from lerobot.common.utils.robot_utils import busy_wait
|
| 48 |
-
from lerobot.common.utils.utils import (
|
| 49 |
-
init_logging,
|
| 50 |
-
log_say,
|
| 51 |
-
)
|
| 52 |
-
|
| 53 |
-
|
| 54 |
-
@dataclass
|
| 55 |
-
class DatasetReplayConfig:
|
| 56 |
-
# Dataset identifier. By convention it should match '{hf_username}/{dataset_name}' (e.g. `lerobot/test`).
|
| 57 |
-
repo_id: str
|
| 58 |
-
# Episode to replay.
|
| 59 |
-
episode: int
|
| 60 |
-
# Root directory where the dataset will be stored (e.g. 'dataset/path').
|
| 61 |
-
root: str | Path | None = None
|
| 62 |
-
# Limit the frames per second. By default, uses the policy fps.
|
| 63 |
-
fps: int = 30
|
| 64 |
-
|
| 65 |
-
|
| 66 |
-
@dataclass
|
| 67 |
-
class ReplayConfig:
|
| 68 |
-
robot: RobotConfig
|
| 69 |
-
dataset: DatasetReplayConfig
|
| 70 |
-
# Use vocal synthesis to read events.
|
| 71 |
-
play_sounds: bool = True
|
| 72 |
-
|
| 73 |
-
|
| 74 |
-
@draccus.wrap()
|
| 75 |
-
def replay(cfg: ReplayConfig):
|
| 76 |
-
init_logging()
|
| 77 |
-
logging.info(pformat(asdict(cfg)))
|
| 78 |
-
|
| 79 |
-
robot = make_robot_from_config(cfg.robot)
|
| 80 |
-
dataset = LeRobotDataset(cfg.dataset.repo_id, root=cfg.dataset.root, episodes=[cfg.dataset.episode])
|
| 81 |
-
actions = dataset.hf_dataset.select_columns("action")
|
| 82 |
-
robot.connect()
|
| 83 |
-
|
| 84 |
-
log_say("Replaying episode", cfg.play_sounds, blocking=True)
|
| 85 |
-
for idx in range(dataset.num_frames):
|
| 86 |
-
start_episode_t = time.perf_counter()
|
| 87 |
-
|
| 88 |
-
action_array = actions[idx]["action"]
|
| 89 |
-
action = {}
|
| 90 |
-
for i, name in enumerate(dataset.features["action"]["names"]):
|
| 91 |
-
key = f"{name.removeprefix('main_')}.pos"
|
| 92 |
-
action[key] = action_array[i].item()
|
| 93 |
-
|
| 94 |
-
action["shoulder_lift.pos"] = -(action["shoulder_lift.pos"] - 90)
|
| 95 |
-
action["elbow_flex.pos"] -= 90
|
| 96 |
-
robot.send_action(action)
|
| 97 |
-
|
| 98 |
-
dt_s = time.perf_counter() - start_episode_t
|
| 99 |
-
busy_wait(1 / dataset.fps - dt_s)
|
| 100 |
-
|
| 101 |
-
robot.disconnect()
|
| 102 |
-
|
| 103 |
-
|
| 104 |
-
if __name__ == "__main__":
|
| 105 |
-
replay()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
examples/lekiwi/evaluate.py
DELETED
|
@@ -1,32 +0,0 @@
|
|
| 1 |
-
from lerobot.common.datasets.utils import build_dataset_frame, hw_to_dataset_features
|
| 2 |
-
from lerobot.common.policies.act.modeling_act import ACTPolicy
|
| 3 |
-
from lerobot.common.robots.lekiwi import LeKiwiClient, LeKiwiClientConfig
|
| 4 |
-
from lerobot.common.utils.control_utils import predict_action
|
| 5 |
-
from lerobot.common.utils.utils import get_safe_torch_device
|
| 6 |
-
|
| 7 |
-
NB_CYCLES_CLIENT_CONNECTION = 1000
|
| 8 |
-
|
| 9 |
-
robot_config = LeKiwiClientConfig(remote_ip="172.18.134.136", id="lekiwi")
|
| 10 |
-
robot = LeKiwiClient(robot_config)
|
| 11 |
-
|
| 12 |
-
robot.connect()
|
| 13 |
-
|
| 14 |
-
policy = ACTPolicy.from_pretrained("pepijn223/act_lekiwi_circle")
|
| 15 |
-
policy.reset()
|
| 16 |
-
|
| 17 |
-
obs_features = hw_to_dataset_features(robot.observation_features, "observation")
|
| 18 |
-
|
| 19 |
-
print("Running inference")
|
| 20 |
-
i = 0
|
| 21 |
-
while i < NB_CYCLES_CLIENT_CONNECTION:
|
| 22 |
-
obs = robot.get_observation()
|
| 23 |
-
|
| 24 |
-
observation_frame = build_dataset_frame(obs_features, obs, prefix="observation")
|
| 25 |
-
action_values = predict_action(
|
| 26 |
-
observation_frame, policy, get_safe_torch_device(policy.config.device), policy.config.use_amp
|
| 27 |
-
)
|
| 28 |
-
action = {key: action_values[i].item() for i, key in enumerate(robot.action_features)}
|
| 29 |
-
robot.send_action(action)
|
| 30 |
-
i += 1
|
| 31 |
-
|
| 32 |
-
robot.disconnect()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
examples/lekiwi/record.py
DELETED
|
@@ -1,67 +0,0 @@
|
|
| 1 |
-
import time
|
| 2 |
-
|
| 3 |
-
from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
|
| 4 |
-
from lerobot.common.datasets.utils import hw_to_dataset_features
|
| 5 |
-
from lerobot.common.robots.lekiwi.config_lekiwi import LeKiwiClientConfig
|
| 6 |
-
from lerobot.common.robots.lekiwi.lekiwi_client import LeKiwiClient
|
| 7 |
-
from lerobot.common.teleoperators.keyboard import KeyboardTeleop, KeyboardTeleopConfig
|
| 8 |
-
from lerobot.common.teleoperators.so100_leader import SO100Leader, SO100LeaderConfig
|
| 9 |
-
|
| 10 |
-
NB_CYCLES_CLIENT_CONNECTION = 250
|
| 11 |
-
|
| 12 |
-
leader_arm_config = SO100LeaderConfig(port="/dev/tty.usbmodem58760431551")
|
| 13 |
-
leader_arm = SO100Leader(leader_arm_config)
|
| 14 |
-
|
| 15 |
-
keyboard_config = KeyboardTeleopConfig()
|
| 16 |
-
keyboard = KeyboardTeleop(keyboard_config)
|
| 17 |
-
|
| 18 |
-
robot_config = LeKiwiClientConfig(remote_ip="172.18.134.136", id="lekiwi")
|
| 19 |
-
robot = LeKiwiClient(robot_config)
|
| 20 |
-
|
| 21 |
-
action_features = hw_to_dataset_features(robot.action_features, "action")
|
| 22 |
-
obs_features = hw_to_dataset_features(robot.observation_features, "observation")
|
| 23 |
-
dataset_features = {**action_features, **obs_features}
|
| 24 |
-
|
| 25 |
-
dataset = LeRobotDataset.create(
|
| 26 |
-
repo_id="pepijn223/lekiwi" + str(int(time.time())),
|
| 27 |
-
fps=10,
|
| 28 |
-
features=dataset_features,
|
| 29 |
-
robot_type=robot.name,
|
| 30 |
-
)
|
| 31 |
-
|
| 32 |
-
leader_arm.connect()
|
| 33 |
-
keyboard.connect()
|
| 34 |
-
robot.connect()
|
| 35 |
-
|
| 36 |
-
if not robot.is_connected or not leader_arm.is_connected or not keyboard.is_connected:
|
| 37 |
-
exit()
|
| 38 |
-
|
| 39 |
-
print("Starting LeKiwi recording")
|
| 40 |
-
i = 0
|
| 41 |
-
while i < NB_CYCLES_CLIENT_CONNECTION:
|
| 42 |
-
arm_action = leader_arm.get_action()
|
| 43 |
-
arm_action = {f"arm_{k}": v for k, v in arm_action.items()}
|
| 44 |
-
|
| 45 |
-
keyboard_keys = keyboard.get_action()
|
| 46 |
-
|
| 47 |
-
base_action = robot._from_keyboard_to_base_action(keyboard_keys)
|
| 48 |
-
|
| 49 |
-
action = {**arm_action, **base_action} if len(base_action) > 0 else arm_action
|
| 50 |
-
|
| 51 |
-
action_sent = robot.send_action(action)
|
| 52 |
-
observation = robot.get_observation()
|
| 53 |
-
|
| 54 |
-
frame = {**action_sent, **observation}
|
| 55 |
-
task = "Dummy Example Task Dataset"
|
| 56 |
-
|
| 57 |
-
dataset.add_frame(frame, task)
|
| 58 |
-
i += 1
|
| 59 |
-
|
| 60 |
-
print("Disconnecting Teleop Devices and LeKiwi Client")
|
| 61 |
-
robot.disconnect()
|
| 62 |
-
leader_arm.disconnect()
|
| 63 |
-
keyboard.disconnect()
|
| 64 |
-
|
| 65 |
-
print("Uploading dataset to the hub")
|
| 66 |
-
dataset.save_episode()
|
| 67 |
-
dataset.push_to_hub()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
examples/lekiwi/replay.py
DELETED
|
@@ -1,25 +0,0 @@
|
|
| 1 |
-
import time
|
| 2 |
-
|
| 3 |
-
from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
|
| 4 |
-
from lerobot.common.robots.lekiwi.config_lekiwi import LeKiwiClientConfig
|
| 5 |
-
from lerobot.common.robots.lekiwi.lekiwi_client import LeKiwiClient
|
| 6 |
-
from lerobot.common.utils.robot_utils import busy_wait
|
| 7 |
-
|
| 8 |
-
robot_config = LeKiwiClientConfig(remote_ip="172.18.134.136", id="lekiwi")
|
| 9 |
-
robot = LeKiwiClient(robot_config)
|
| 10 |
-
|
| 11 |
-
dataset = LeRobotDataset("pepijn223/lekiwi1749025613", episodes=[0])
|
| 12 |
-
|
| 13 |
-
robot.connect()
|
| 14 |
-
|
| 15 |
-
print("Replaying episode…")
|
| 16 |
-
for _, action_array in enumerate(dataset.hf_dataset["action"]):
|
| 17 |
-
t0 = time.perf_counter()
|
| 18 |
-
|
| 19 |
-
action = {name: float(action_array[i]) for i, name in enumerate(dataset.features["action"]["names"])}
|
| 20 |
-
robot.send_action(action)
|
| 21 |
-
|
| 22 |
-
busy_wait(max(1.0 / dataset.fps - (time.perf_counter() - t0), 0.0))
|
| 23 |
-
|
| 24 |
-
print("Disconnecting LeKiwi Client")
|
| 25 |
-
robot.disconnect()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
examples/lekiwi/teleoperate.py
DELETED
|
@@ -1,32 +0,0 @@
|
|
| 1 |
-
from lerobot.common.robots.lekiwi import LeKiwiClient, LeKiwiClientConfig
|
| 2 |
-
from lerobot.common.teleoperators.keyboard.teleop_keyboard import KeyboardTeleop, KeyboardTeleopConfig
|
| 3 |
-
from lerobot.common.teleoperators.so100_leader import SO100Leader, SO100LeaderConfig
|
| 4 |
-
|
| 5 |
-
robot_config = LeKiwiClientConfig(remote_ip="172.18.134.136", id="my_lekiwi")
|
| 6 |
-
|
| 7 |
-
teleop__arm_config = SO100LeaderConfig(
|
| 8 |
-
port="/dev/tty.usbmodem58760431551",
|
| 9 |
-
id="my_awesome_leader_arm",
|
| 10 |
-
)
|
| 11 |
-
|
| 12 |
-
teleop_keyboard_config = KeyboardTeleopConfig(
|
| 13 |
-
id="my_laptop_keyboard",
|
| 14 |
-
)
|
| 15 |
-
|
| 16 |
-
robot = LeKiwiClient(robot_config)
|
| 17 |
-
teleop_arm = SO100Leader(teleop__arm_config)
|
| 18 |
-
telep_keyboard = KeyboardTeleop(teleop_keyboard_config)
|
| 19 |
-
robot.connect()
|
| 20 |
-
teleop_arm.connect()
|
| 21 |
-
telep_keyboard.connect()
|
| 22 |
-
|
| 23 |
-
while True:
|
| 24 |
-
observation = robot.get_observation()
|
| 25 |
-
|
| 26 |
-
arm_action = teleop_arm.get_action()
|
| 27 |
-
arm_action = {f"arm_{k}": v for k, v in arm_action.items()}
|
| 28 |
-
|
| 29 |
-
keyboard_keys = telep_keyboard.get_action()
|
| 30 |
-
base_action = robot._from_keyboard_to_base_action(keyboard_keys)
|
| 31 |
-
|
| 32 |
-
robot.send_action(arm_action | base_action)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|