mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-26 11:16:00 +00:00
Pepijn
121 lines
4.5 KiB
Python
121 lines
4.5 KiB
Python
# Copyright 2025 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.
|
|
|
|
"""Train directly from episode-scoped Parquet and MP4 streams.
|
|
|
|
For normal training, prefer ``lerobot-train --dataset.streaming=true`` so distributed sharding,
|
|
checkpoint resume, and device placement are configured by the training pipeline. This lower-level
|
|
example shows the underlying Python API.
|
|
"""
|
|
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
from lerobot.configs import FeatureType
|
|
from lerobot.datasets import LeRobotDatasetMetadata, StreamingLeRobotDataset
|
|
from lerobot.policies import make_pre_post_processors
|
|
from lerobot.policies.act import ACTConfig, ACTPolicy
|
|
from lerobot.utils.constants import ACTION
|
|
from lerobot.utils.feature_utils import dataset_to_policy_features
|
|
|
|
|
|
def main():
|
|
# Create a directory to store the training checkpoint.
|
|
output_directory = Path("outputs/train/example_streaming_dataset")
|
|
output_directory.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Selects the "best" device available
|
|
device = (
|
|
torch.device("cuda")
|
|
if torch.cuda.is_available()
|
|
else torch.device("mps")
|
|
if torch.backends.mps.is_available()
|
|
else torch.device("cpu")
|
|
)
|
|
print(f"Using device: {device}")
|
|
|
|
training_steps = 10
|
|
log_freq = 1
|
|
|
|
dataset_id = "lerobot/droid_1.0.1" # 26M frames! Would require 4TB of disk space if installed locally (:
|
|
dataset_metadata = LeRobotDatasetMetadata(dataset_id)
|
|
features = dataset_to_policy_features(dataset_metadata.features)
|
|
output_features = {key: ft for key, ft in features.items() if ft.type is FeatureType.ACTION}
|
|
input_features = {key: ft for key, ft in features.items() if key not in output_features}
|
|
|
|
# We can now instantiate our policy with this config and the dataset stats.
|
|
cfg = ACTConfig(input_features=input_features, output_features=output_features)
|
|
policy = ACTPolicy(cfg)
|
|
policy.train()
|
|
policy.to(device)
|
|
preprocessor, postprocessor = make_pre_post_processors(cfg, dataset_stats=dataset_metadata.stats)
|
|
|
|
# Delta timestamps are used to (1) augment frames used during training and (2) supervise the policy.
|
|
# Here, we use delta-timestamps to only provide ground truth actions for supervision
|
|
delta_timestamps = {
|
|
ACTION: [t / dataset_metadata.fps for t in range(cfg.n_action_steps)],
|
|
}
|
|
|
|
# The first run resolves or locally builds a revision-safe MP4 index sidecar. It is never
|
|
# uploaded implicitly. Episode rows and video byte ranges are then prefetched together.
|
|
dataset = StreamingLeRobotDataset(
|
|
dataset_id,
|
|
delta_timestamps=delta_timestamps,
|
|
tolerance_s=1e-3,
|
|
episode_pool_size=16,
|
|
prefetch_episodes=4,
|
|
byte_budget_gb=4,
|
|
repeat=True,
|
|
)
|
|
|
|
optimizer = torch.optim.Adam(policy.parameters(), lr=1e-4)
|
|
dataloader = torch.utils.data.DataLoader(
|
|
dataset,
|
|
# One worker owns the rank-level pool. Internal fetch concurrency is configured through
|
|
# StreamingLeRobotDataset.max_num_shards (the lerobot-train CLI derives it from num_workers).
|
|
num_workers=1,
|
|
batch_size=16,
|
|
pin_memory=device.type != "cpu",
|
|
prefetch_factor=2, # bounded decoded-batch queue while the policy trains
|
|
persistent_workers=True,
|
|
)
|
|
|
|
# Run training loop.
|
|
step = 0
|
|
done = False
|
|
while not done:
|
|
for batch in dataloader:
|
|
batch = preprocessor(batch)
|
|
loss, _ = policy.forward(batch)
|
|
loss.backward()
|
|
optimizer.step()
|
|
optimizer.zero_grad()
|
|
|
|
if step % log_freq == 0:
|
|
print(f"step: {step} loss: {loss.item():.3f}")
|
|
step += 1
|
|
if step >= training_steps:
|
|
done = True
|
|
break
|
|
|
|
# Save a policy checkpoint.
|
|
policy.save_pretrained(output_directory)
|
|
preprocessor.save_pretrained(output_directory)
|
|
postprocessor.save_pretrained(output_directory)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|