mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-29 12:39:41 +00:00
Compare commits
19 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 6906178b39 | |||
| cbc8bfb2e6 | |||
| 0d1be72dc8 | |||
| 96b7c212c4 | |||
| 4303b3c930 | |||
| 63dca86df8 | |||
| 8a0cc3d664 | |||
| 8bb8ed4803 | |||
| 095856b06a | |||
| 563f42bdb1 | |||
| 8fff0fde7c | |||
| 04de496547 | |||
| baf9b50365 | |||
| a0fdbf037a | |||
| c085531b17 | |||
| c7c6205332 | |||
| 4e54be1334 | |||
| fde9d08281 | |||
| 46044fed75 |
@@ -0,0 +1,25 @@
|
|||||||
|
# AI Usage Policy
|
||||||
|
|
||||||
|
The LeRobot project welcomes contributions from everyone, and we have a few guidelines regarding AI usage to ensure high code quality, clear communication, and a healthy open-source ecosystem:
|
||||||
|
|
||||||
|
- **Please disclose significant AI assistance.** If you used AI tools (e.g., Copilot, Claude, Cursor, ChatGPT) to generate a substantial portion of your code or text, let us know in your PR description. Transparency helps us review your changes more effectively.
|
||||||
|
- **Own your code (The Human-in-the-Loop).** You must fully understand all the changes you are proposing. If you cannot explain what your AI-assisted code does or how it interacts with LeRobot's broader architecture, please take the time to learn and test it before submitting.
|
||||||
|
- **Keep issues and discussions focused.** You are welcome to use AI to help draft issues or PR descriptions, but please review and edit them carefully before posting. AI can often be overly verbose; trimming the noise and getting straight to the point helps our maintainers address your needs faster.
|
||||||
|
|
||||||
|
Our core maintainers also use AI tools to aid their workflows, but they do so while bringing deep contextual knowledge of the LeRobot codebase to validate the output. We ask all contributors to apply that same level of rigor.
|
||||||
|
|
||||||
|
## Remember the Human Maintainers
|
||||||
|
|
||||||
|
Please remember that LeRobot is maintained by a dedicated team of humans.
|
||||||
|
|
||||||
|
Every discussion, issue, and pull request is read and reviewed by real people. While AI tools can generate thousands of lines of code in seconds, reviewing that code still takes human time and energy. Submitting unverified or low-effort AI output puts an unfair burden on our maintainers.
|
||||||
|
|
||||||
|
Today, the quality of the AI output still heavily depends on the developer driving the tool. We ask that you respect our maintainers' time by thoroughly vetting, testing, and refining your submissions.
|
||||||
|
|
||||||
|
## AI is Welcome Here
|
||||||
|
|
||||||
|
LeRobot operates at the cutting edge of AI and robotics, and many of our maintainers actively embrace AI coding assistants as valuable productivity tools. We are a pro-AI project!
|
||||||
|
|
||||||
|
Our reason for having an AI policy is not an anti-AI stance. Rather, it exists to ensure that AI is used to enhance human contributions, not replace them with unverified noise. It's about how the tools are used, not the tools themselves.
|
||||||
|
|
||||||
|
We value the unique human insight you bring to the LeRobot community. Let AI empower your workflow, but always let your own judgment take the wheel.
|
||||||
+1
-1
@@ -2,7 +2,7 @@
|
|||||||
|
|
||||||
Everyone is welcome to contribute, and we value everybody's contribution. Code is not the only way to help the community. Answering questions, helping others, reaching out, and improving the documentation are immensely valuable.
|
Everyone is welcome to contribute, and we value everybody's contribution. Code is not the only way to help the community. Answering questions, helping others, reaching out, and improving the documentation are immensely valuable.
|
||||||
|
|
||||||
Whichever way you choose to contribute, please be mindful to respect our [code of conduct](./CODE_OF_CONDUCT.md).
|
Whichever way you choose to contribute, please be mindful to respect our [code of conduct](./CODE_OF_CONDUCT.md) and our [AI policy](./AI_POLICY.md).
|
||||||
|
|
||||||
## Ways to Contribute
|
## Ways to Contribute
|
||||||
|
|
||||||
|
|||||||
@@ -48,7 +48,7 @@ python -m lerobot.async_inference.robot_client \
|
|||||||
--task="dummy" \ # POLICY: The task to run the policy on (`Fold my t-shirt`). Not necessarily defined for all policies, such as `act`
|
--task="dummy" \ # POLICY: The task to run the policy on (`Fold my t-shirt`). Not necessarily defined for all policies, such as `act`
|
||||||
--policy_type=your_policy_type \ # POLICY: the type of policy to run (smolvla, act, etc)
|
--policy_type=your_policy_type \ # POLICY: the type of policy to run (smolvla, act, etc)
|
||||||
--pretrained_name_or_path=user/model \ # POLICY: the model name/path on server to the checkpoint to run (e.g., lerobot/smolvla_base)
|
--pretrained_name_or_path=user/model \ # POLICY: the model name/path on server to the checkpoint to run (e.g., lerobot/smolvla_base)
|
||||||
--policy_device=mps \ # POLICY: the device to run the policy on, on the server
|
--policy_device=mps \ # POLICY: the device to run the policy on, on the server (cuda, mps, xpu, cpu)
|
||||||
--actions_per_chunk=50 \ # POLICY: the number of actions to output at once
|
--actions_per_chunk=50 \ # POLICY: the number of actions to output at once
|
||||||
--chunk_size_threshold=0.5 \ # CLIENT: the threshold for the chunk size before sending a new observation to the server
|
--chunk_size_threshold=0.5 \ # CLIENT: the threshold for the chunk size before sending a new observation to the server
|
||||||
--aggregate_fn_name=weighted_average \ # CLIENT: the function to aggregate actions on overlapping portions
|
--aggregate_fn_name=weighted_average \ # CLIENT: the function to aggregate actions on overlapping portions
|
||||||
|
|||||||
@@ -170,13 +170,13 @@ Once you can drive the robot well, you can start recording data to train AI mode
|
|||||||
We use Hugging Face to store your data online. First, log in with your token from [Hugging Face settings](https://huggingface.co/settings/tokens):
|
We use Hugging Face to store your data online. First, log in with your token from [Hugging Face settings](https://huggingface.co/settings/tokens):
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
huggingface-cli login --token ${HUGGINGFACE_TOKEN} --add-to-git-credential
|
hf auth login --token ${HUGGINGFACE_TOKEN} --add-to-git-credential
|
||||||
```
|
```
|
||||||
|
|
||||||
Store your Hugging Face username:
|
Store your Hugging Face username:
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
HF_USER=$(huggingface-cli whoami | head -n 1)
|
HF_USER=$(hf auth whoami | awk -F': *' 'NR==1 {print $2}')
|
||||||
echo $HF_USER
|
echo $HF_USER
|
||||||
```
|
```
|
||||||
|
|
||||||
|
|||||||
@@ -155,10 +155,10 @@ Upload your repository to Hugging Face:
|
|||||||
pip install huggingface_hub
|
pip install huggingface_hub
|
||||||
|
|
||||||
# Login to Hugging Face
|
# Login to Hugging Face
|
||||||
huggingface-cli login
|
hf auth login
|
||||||
|
|
||||||
# Create a new repository
|
# Create a new repository
|
||||||
huggingface-cli repo create my-custom-env --type space --org my-org
|
hf repo create my-org/my-custom-env
|
||||||
|
|
||||||
# Initialize git and push
|
# Initialize git and push
|
||||||
git init
|
git init
|
||||||
|
|||||||
@@ -159,7 +159,7 @@ We use the Hugging Face hub features for uploading your dataset. If you haven't
|
|||||||
Add your token to the CLI by running this command:
|
Add your token to the CLI by running this command:
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
huggingface-cli login --token ${HUGGINGFACE_TOKEN} --add-to-git-credential
|
hf auth login --token ${HUGGINGFACE_TOKEN} --add-to-git-credential
|
||||||
```
|
```
|
||||||
|
|
||||||
Then store your Hugging Face repository name in a variable:
|
Then store your Hugging Face repository name in a variable:
|
||||||
@@ -327,7 +327,7 @@ You can look for other LeRobot datasets on the hub by searching for `LeRobot` [t
|
|||||||
You can also push your local dataset to the Hub manually, running:
|
You can also push your local dataset to the Hub manually, running:
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
huggingface-cli upload ${HF_USER}/record-test ~/.cache/huggingface/lerobot/{repo-id} --repo-type dataset
|
hf upload ${HF_USER}/record-test ~/.cache/huggingface/lerobot/{repo-id} --repo-type dataset
|
||||||
```
|
```
|
||||||
|
|
||||||
#### Record function
|
#### Record function
|
||||||
@@ -491,7 +491,7 @@ If your local computer doesn't have a powerful GPU you could utilize Google Cola
|
|||||||
Once training is done, upload the latest checkpoint with:
|
Once training is done, upload the latest checkpoint with:
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
huggingface-cli upload ${HF_USER}/act_so101_test \
|
hf upload ${HF_USER}/act_so101_test \
|
||||||
outputs/train/act_so101_test/checkpoints/last/pretrained_model
|
outputs/train/act_so101_test/checkpoints/last/pretrained_model
|
||||||
```
|
```
|
||||||
|
|
||||||
@@ -499,7 +499,7 @@ You can also upload intermediate checkpoints with:
|
|||||||
|
|
||||||
```bash
|
```bash
|
||||||
CKPT=010000
|
CKPT=010000
|
||||||
huggingface-cli upload ${HF_USER}/act_so101_test${CKPT} \
|
hf upload ${HF_USER}/act_so101_test${CKPT} \
|
||||||
outputs/train/act_so101_test/checkpoints/${CKPT}/pretrained_model
|
outputs/train/act_so101_test/checkpoints/${CKPT}/pretrained_model
|
||||||
```
|
```
|
||||||
|
|
||||||
|
|||||||
@@ -279,13 +279,13 @@ We use the Hugging Face hub features for uploading your dataset. If you haven't
|
|||||||
Add your token to the CLI by running this command:
|
Add your token to the CLI by running this command:
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
huggingface-cli login --token ${HUGGINGFACE_TOKEN} --add-to-git-credential
|
hf auth login --token ${HUGGINGFACE_TOKEN} --add-to-git-credential
|
||||||
```
|
```
|
||||||
|
|
||||||
Then store your Hugging Face repository name in a variable:
|
Then store your Hugging Face repository name in a variable:
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
HF_USER=$(huggingface-cli whoami | head -n 1)
|
HF_USER=$(hf auth whoami | awk -F': *' 'NR==1 {print $2}')
|
||||||
echo $HF_USER
|
echo $HF_USER
|
||||||
```
|
```
|
||||||
|
|
||||||
|
|||||||
@@ -57,7 +57,7 @@ class DatasetReplayConfig:
|
|||||||
repo_id: str
|
repo_id: str
|
||||||
# Episode to replay.
|
# Episode to replay.
|
||||||
episode: int
|
episode: int
|
||||||
# Root directory where the dataset will be stored (e.g. 'dataset/path').
|
# Root directory where the dataset will be stored (e.g. 'dataset/path'). If None, defaults to $HF_LEROBOT_HOME/repo_id.
|
||||||
root: str | Path | None = None
|
root: str | Path | None = None
|
||||||
# Limit the frames per second. By default, uses the policy fps.
|
# Limit the frames per second. By default, uses the policy fps.
|
||||||
fps: int = 30
|
fps: int = 30
|
||||||
|
|||||||
+1
-1
@@ -25,7 +25,7 @@ discord = "https://discord.gg/s3KuuzsPFb"
|
|||||||
|
|
||||||
[project]
|
[project]
|
||||||
name = "lerobot"
|
name = "lerobot"
|
||||||
version = "0.4.4"
|
version = "0.4.5"
|
||||||
description = "🤗 LeRobot: State-of-the-art Machine Learning for Real-World Robotics in Pytorch"
|
description = "🤗 LeRobot: State-of-the-art Machine Learning for Real-World Robotics in Pytorch"
|
||||||
dynamic = ["readme"]
|
dynamic = ["readme"]
|
||||||
license = { text = "Apache-2.0" }
|
license = { text = "Apache-2.0" }
|
||||||
|
|||||||
@@ -49,23 +49,18 @@ import torch
|
|||||||
|
|
||||||
from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig # noqa: F401
|
from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig # noqa: F401
|
||||||
from lerobot.cameras.realsense.configuration_realsense import RealSenseCameraConfig # noqa: F401
|
from lerobot.cameras.realsense.configuration_realsense import RealSenseCameraConfig # noqa: F401
|
||||||
from lerobot.robots import ( # noqa: F401
|
from lerobot.robots import (
|
||||||
Robot,
|
RobotConfig, # noqa: F401
|
||||||
RobotConfig,
|
|
||||||
bi_so_follower,
|
|
||||||
koch_follower,
|
|
||||||
make_robot_from_config,
|
make_robot_from_config,
|
||||||
omx_follower,
|
|
||||||
so_follower,
|
|
||||||
)
|
)
|
||||||
from lerobot.transport import (
|
from lerobot.transport import (
|
||||||
services_pb2, # type: ignore
|
services_pb2, # type: ignore
|
||||||
services_pb2_grpc, # type: ignore
|
services_pb2_grpc, # type: ignore
|
||||||
)
|
)
|
||||||
from lerobot.transport.utils import grpc_channel_options, send_bytes_in_chunks
|
from lerobot.transport.utils import grpc_channel_options, send_bytes_in_chunks
|
||||||
|
from lerobot.utils.import_utils import register_third_party_plugins
|
||||||
|
|
||||||
from .configs import RobotClientConfig
|
from .configs import RobotClientConfig
|
||||||
from .constants import SUPPORTED_ROBOTS
|
|
||||||
from .helpers import (
|
from .helpers import (
|
||||||
Action,
|
Action,
|
||||||
FPSTracker,
|
FPSTracker,
|
||||||
@@ -485,8 +480,9 @@ class RobotClient:
|
|||||||
def async_client(cfg: RobotClientConfig):
|
def async_client(cfg: RobotClientConfig):
|
||||||
logging.info(pformat(asdict(cfg)))
|
logging.info(pformat(asdict(cfg)))
|
||||||
|
|
||||||
if cfg.robot.type not in SUPPORTED_ROBOTS:
|
# TODO: Assert if checking robot support is still needed with the plugin system
|
||||||
raise ValueError(f"Robot {cfg.robot.type} not yet supported!")
|
# if cfg.robot.type not in SUPPORTED_ROBOTS:
|
||||||
|
# raise ValueError(f"Robot {cfg.robot.type} not yet supported!")
|
||||||
|
|
||||||
client = RobotClient(cfg)
|
client = RobotClient(cfg)
|
||||||
|
|
||||||
@@ -512,4 +508,5 @@ def async_client(cfg: RobotClientConfig):
|
|||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
|
register_third_party_plugins()
|
||||||
async_client() # run the client
|
async_client() # run the client
|
||||||
|
|||||||
@@ -27,7 +27,7 @@ class DatasetConfig:
|
|||||||
# "dataset_index" into the returned item. The index mapping is made according to the order in which the
|
# "dataset_index" into the returned item. The index mapping is made according to the order in which the
|
||||||
# datasets are provided.
|
# datasets are provided.
|
||||||
repo_id: str
|
repo_id: str
|
||||||
# Root directory where the dataset will be stored (e.g. 'dataset/path').
|
# Root directory where the dataset will be stored (e.g. 'dataset/path'). If None, defaults to $HF_LEROBOT_HOME/repo_id.
|
||||||
root: str | None = None
|
root: str | None = None
|
||||||
episodes: list[int] | None = None
|
episodes: list[int] | None = None
|
||||||
image_transforms: ImageTransformsConfig = field(default_factory=ImageTransformsConfig)
|
image_transforms: ImageTransformsConfig = field(default_factory=ImageTransformsConfig)
|
||||||
|
|||||||
@@ -21,6 +21,8 @@ from pathlib import Path
|
|||||||
|
|
||||||
import datasets
|
import datasets
|
||||||
import pandas as pd
|
import pandas as pd
|
||||||
|
import pyarrow as pa
|
||||||
|
import pyarrow.parquet as pq
|
||||||
import tqdm
|
import tqdm
|
||||||
|
|
||||||
from lerobot.datasets.compute_stats import aggregate_stats
|
from lerobot.datasets.compute_stats import aggregate_stats
|
||||||
@@ -35,7 +37,6 @@ from lerobot.datasets.utils import (
|
|||||||
get_file_size_in_mb,
|
get_file_size_in_mb,
|
||||||
get_hf_features_from_features,
|
get_hf_features_from_features,
|
||||||
get_parquet_file_size_in_mb,
|
get_parquet_file_size_in_mb,
|
||||||
to_parquet_with_hf_images,
|
|
||||||
update_chunk_file_indices,
|
update_chunk_file_indices,
|
||||||
write_info,
|
write_info,
|
||||||
write_stats,
|
write_stats,
|
||||||
@@ -80,28 +81,41 @@ def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]):
|
|||||||
return fps, robot_type, features
|
return fps, robot_type, features
|
||||||
|
|
||||||
|
|
||||||
def update_data_df(df, src_meta, dst_meta):
|
def update_data_table(table: pa.Table, src_meta, dst_meta) -> pa.Table:
|
||||||
"""Updates a data DataFrame with new indices and task mappings for aggregation.
|
"""Updates a pyarrow Table with new indices and task mappings for aggregation.
|
||||||
|
|
||||||
Adjusts episode indices, frame indices, and task indices to account for
|
Adjusts episode indices, frame indices, and task indices to account for
|
||||||
previously aggregated data in the destination dataset.
|
previously aggregated data in the destination dataset.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
df: DataFrame containing the data to be updated.
|
table: pyarrow Table containing the data to be updated.
|
||||||
src_meta: Source dataset metadata.
|
src_meta: Source dataset metadata.
|
||||||
dst_meta: Destination dataset metadata.
|
dst_meta: Destination dataset metadata.
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
pd.DataFrame: Updated DataFrame with adjusted indices.
|
pa.Table: Updated Table with adjusted indices.
|
||||||
"""
|
"""
|
||||||
|
ep_offset = dst_meta.info["total_episodes"]
|
||||||
|
idx_offset = dst_meta.info["total_frames"]
|
||||||
|
|
||||||
df["episode_index"] = df["episode_index"] + dst_meta.info["total_episodes"]
|
ep_col = table.column("episode_index")
|
||||||
df["index"] = df["index"] + dst_meta.info["total_frames"]
|
new_ep = pa.array([v + ep_offset for v in ep_col.to_pylist()], type=ep_col.type)
|
||||||
|
table = table.set_column(table.column_names.index("episode_index"), "episode_index", new_ep)
|
||||||
|
|
||||||
src_task_names = src_meta.tasks.index.take(df["task_index"].to_numpy())
|
idx_col = table.column("index")
|
||||||
df["task_index"] = dst_meta.tasks.loc[src_task_names, "task_index"].to_numpy()
|
new_idx = pa.array([v + idx_offset for v in idx_col.to_pylist()], type=idx_col.type)
|
||||||
|
table = table.set_column(table.column_names.index("index"), "index", new_idx)
|
||||||
|
|
||||||
return df
|
old_task_indices = table.column("task_index").to_pylist()
|
||||||
|
src_task_names = src_meta.tasks.index.take(old_task_indices)
|
||||||
|
new_task_indices = dst_meta.tasks.loc[src_task_names, "task_index"].to_numpy().tolist()
|
||||||
|
table = table.set_column(
|
||||||
|
table.column_names.index("task_index"),
|
||||||
|
"task_index",
|
||||||
|
pa.array(new_task_indices, type=table.schema.field("task_index").type),
|
||||||
|
)
|
||||||
|
|
||||||
|
return table
|
||||||
|
|
||||||
|
|
||||||
def update_meta_data(
|
def update_meta_data(
|
||||||
@@ -289,7 +303,9 @@ def aggregate_datasets(
|
|||||||
|
|
||||||
logging.info("Find all tasks")
|
logging.info("Find all tasks")
|
||||||
unique_tasks = pd.concat([m.tasks for m in all_metadata]).index.unique()
|
unique_tasks = pd.concat([m.tasks for m in all_metadata]).index.unique()
|
||||||
dst_meta.tasks = pd.DataFrame({"task_index": range(len(unique_tasks))}, index=unique_tasks)
|
dst_meta.tasks = pd.DataFrame(
|
||||||
|
{"task_index": range(len(unique_tasks))}, index=pd.Index(unique_tasks, name="task")
|
||||||
|
)
|
||||||
|
|
||||||
meta_idx = {"chunk": 0, "file": 0}
|
meta_idx = {"chunk": 0, "file": 0}
|
||||||
data_idx = {"chunk": 0, "file": 0}
|
data_idx = {"chunk": 0, "file": 0}
|
||||||
@@ -466,18 +482,13 @@ def aggregate_data(src_meta, dst_meta, data_idx, data_files_size_in_mb, chunk_si
|
|||||||
src_path = src_meta.root / DEFAULT_DATA_PATH.format(
|
src_path = src_meta.root / DEFAULT_DATA_PATH.format(
|
||||||
chunk_index=src_chunk_idx, file_index=src_file_idx
|
chunk_index=src_chunk_idx, file_index=src_file_idx
|
||||||
)
|
)
|
||||||
if contains_images:
|
table = pq.read_table(src_path)
|
||||||
# Use HuggingFace datasets to read source data to preserve image format
|
table = update_data_table(table, src_meta, dst_meta)
|
||||||
src_ds = datasets.Dataset.from_parquet(str(src_path))
|
|
||||||
df = src_ds.to_pandas()
|
|
||||||
else:
|
|
||||||
df = pd.read_parquet(src_path)
|
|
||||||
df = update_data_df(df, src_meta, dst_meta)
|
|
||||||
|
|
||||||
# Write data and get the actual destination file it was written to
|
# Write data and get the actual destination file it was written to
|
||||||
# This avoids duplicating the rotation logic here
|
# This avoids duplicating the rotation logic here
|
||||||
data_idx, (dst_chunk, dst_file) = append_or_create_parquet_file(
|
data_idx, (dst_chunk, dst_file) = append_or_create_parquet_file(
|
||||||
df,
|
table,
|
||||||
src_path,
|
src_path,
|
||||||
data_idx,
|
data_idx,
|
||||||
data_files_size_in_mb,
|
data_files_size_in_mb,
|
||||||
@@ -552,8 +563,16 @@ def aggregate_metadata(src_meta, dst_meta, meta_idx, data_idx, videos_idx):
|
|||||||
return meta_idx
|
return meta_idx
|
||||||
|
|
||||||
|
|
||||||
|
def _write_table_with_hf_images(
|
||||||
|
table: pa.Table, path: Path, features: datasets.Features | None = None
|
||||||
|
) -> None:
|
||||||
|
"""Write a pyarrow Table to parquet with proper HF image encoding."""
|
||||||
|
ds = datasets.Dataset.from_dict(table.to_pydict(), features=features)
|
||||||
|
ds.to_parquet(path)
|
||||||
|
|
||||||
|
|
||||||
def append_or_create_parquet_file(
|
def append_or_create_parquet_file(
|
||||||
df: pd.DataFrame,
|
data: pd.DataFrame | pa.Table,
|
||||||
src_path: Path,
|
src_path: Path,
|
||||||
idx: dict[str, int],
|
idx: dict[str, int],
|
||||||
max_mb: float,
|
max_mb: float,
|
||||||
@@ -569,7 +588,7 @@ def append_or_create_parquet_file(
|
|||||||
from becoming too large. Handles both regular parquet files and those containing images.
|
from becoming too large. Handles both regular parquet files and those containing images.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
df: DataFrame to write to the parquet file.
|
data: Data to write, as a pandas DataFrame or pyarrow Table.
|
||||||
src_path: Path to the source file (used for size estimation).
|
src_path: Path to the source file (used for size estimation).
|
||||||
idx: Dictionary containing current 'chunk' and 'file' indices.
|
idx: Dictionary containing current 'chunk' and 'file' indices.
|
||||||
max_mb: Maximum allowed file size in MB before rotation.
|
max_mb: Maximum allowed file size in MB before rotation.
|
||||||
@@ -583,15 +602,17 @@ def append_or_create_parquet_file(
|
|||||||
tuple: (updated_idx, (dst_chunk, dst_file)) where updated_idx is the index dict
|
tuple: (updated_idx, (dst_chunk, dst_file)) where updated_idx is the index dict
|
||||||
and (dst_chunk, dst_file) is the actual destination file the data was written to.
|
and (dst_chunk, dst_file) is the actual destination file the data was written to.
|
||||||
"""
|
"""
|
||||||
|
table = data if isinstance(data, pa.Table) else pa.Table.from_pandas(data)
|
||||||
|
|
||||||
dst_chunk, dst_file = idx["chunk"], idx["file"]
|
dst_chunk, dst_file = idx["chunk"], idx["file"]
|
||||||
dst_path = aggr_root / default_path.format(chunk_index=dst_chunk, file_index=dst_file)
|
dst_path = aggr_root / default_path.format(chunk_index=dst_chunk, file_index=dst_file)
|
||||||
|
|
||||||
if not dst_path.exists():
|
if not dst_path.exists():
|
||||||
dst_path.parent.mkdir(parents=True, exist_ok=True)
|
dst_path.parent.mkdir(parents=True, exist_ok=True)
|
||||||
if contains_images:
|
if contains_images:
|
||||||
to_parquet_with_hf_images(df, dst_path, features=hf_features)
|
_write_table_with_hf_images(table, dst_path, features=hf_features)
|
||||||
else:
|
else:
|
||||||
df.to_parquet(dst_path)
|
pq.write_table(table, dst_path)
|
||||||
return idx, (dst_chunk, dst_file)
|
return idx, (dst_chunk, dst_file)
|
||||||
|
|
||||||
src_size = get_parquet_file_size_in_mb(src_path)
|
src_size = get_parquet_file_size_in_mb(src_path)
|
||||||
@@ -602,22 +623,17 @@ def append_or_create_parquet_file(
|
|||||||
dst_chunk, dst_file = idx["chunk"], idx["file"]
|
dst_chunk, dst_file = idx["chunk"], idx["file"]
|
||||||
new_path = aggr_root / default_path.format(chunk_index=dst_chunk, file_index=dst_file)
|
new_path = aggr_root / default_path.format(chunk_index=dst_chunk, file_index=dst_file)
|
||||||
new_path.parent.mkdir(parents=True, exist_ok=True)
|
new_path.parent.mkdir(parents=True, exist_ok=True)
|
||||||
final_df = df
|
final_table = table
|
||||||
target_path = new_path
|
target_path = new_path
|
||||||
else:
|
else:
|
||||||
if contains_images:
|
existing_table = pq.read_table(dst_path)
|
||||||
# Use HuggingFace datasets to read existing data to preserve image format
|
final_table = pa.concat_tables([existing_table, table], promote_options="permissive")
|
||||||
existing_ds = datasets.Dataset.from_parquet(str(dst_path))
|
|
||||||
existing_df = existing_ds.to_pandas()
|
|
||||||
else:
|
|
||||||
existing_df = pd.read_parquet(dst_path)
|
|
||||||
final_df = pd.concat([existing_df, df], ignore_index=True)
|
|
||||||
target_path = dst_path
|
target_path = dst_path
|
||||||
|
|
||||||
if contains_images:
|
if contains_images:
|
||||||
to_parquet_with_hf_images(final_df, target_path, features=hf_features)
|
_write_table_with_hf_images(final_table, target_path, features=hf_features)
|
||||||
else:
|
else:
|
||||||
final_df.to_parquet(target_path)
|
pq.write_table(final_table, target_path)
|
||||||
|
|
||||||
return idx, (dst_chunk, dst_file)
|
return idx, (dst_chunk, dst_file)
|
||||||
|
|
||||||
|
|||||||
@@ -32,6 +32,9 @@ from pathlib import Path
|
|||||||
import datasets
|
import datasets
|
||||||
import numpy as np
|
import numpy as np
|
||||||
import pandas as pd
|
import pandas as pd
|
||||||
|
import pyarrow as pa
|
||||||
|
import pyarrow.compute as pc
|
||||||
|
import pyarrow.dataset as pa_ds
|
||||||
import pyarrow.parquet as pq
|
import pyarrow.parquet as pq
|
||||||
import torch
|
import torch
|
||||||
from tqdm import tqdm
|
from tqdm import tqdm
|
||||||
@@ -47,7 +50,6 @@ from lerobot.datasets.utils import (
|
|||||||
DEFAULT_EPISODES_PATH,
|
DEFAULT_EPISODES_PATH,
|
||||||
get_parquet_file_size_in_mb,
|
get_parquet_file_size_in_mb,
|
||||||
load_episodes,
|
load_episodes,
|
||||||
load_info,
|
|
||||||
update_chunk_file_indices,
|
update_chunk_file_indices,
|
||||||
write_info,
|
write_info,
|
||||||
write_stats,
|
write_stats,
|
||||||
@@ -90,8 +92,8 @@ def delete_episodes(
|
|||||||
Args:
|
Args:
|
||||||
dataset: The source LeRobotDataset.
|
dataset: The source LeRobotDataset.
|
||||||
episode_indices: List of episode indices to delete.
|
episode_indices: List of episode indices to delete.
|
||||||
output_dir: Directory to save the new dataset. If None, uses default location.
|
output_dir: Root directory where the edited dataset will be stored. If not specified, defaults to $HF_LEROBOT_HOME/repo_id. Equivalent to new_root in EditDatasetConfig.
|
||||||
repo_id: Repository ID for the new dataset. If None, appends "_modified" to original.
|
repo_id: Edited dataset identifier. Equivalent to new_repo_id in EditDatasetConfig.
|
||||||
"""
|
"""
|
||||||
if not episode_indices:
|
if not episode_indices:
|
||||||
raise ValueError("No episodes to delete")
|
raise ValueError("No episodes to delete")
|
||||||
@@ -153,7 +155,7 @@ def split_dataset(
|
|||||||
dataset: The source LeRobotDataset to split.
|
dataset: The source LeRobotDataset to split.
|
||||||
splits: Either a dict mapping split names to episode indices, or a dict mapping
|
splits: Either a dict mapping split names to episode indices, or a dict mapping
|
||||||
split names to fractions (must sum to <= 1.0).
|
split names to fractions (must sum to <= 1.0).
|
||||||
output_dir: Base directory for output datasets. If None, uses default location.
|
output_dir: Root directory where the split datasets will be stored. If not specified, defaults to $HF_LEROBOT_HOME/repo_id.
|
||||||
|
|
||||||
Examples:
|
Examples:
|
||||||
Split by specific episodes
|
Split by specific episodes
|
||||||
@@ -244,8 +246,8 @@ def merge_datasets(
|
|||||||
|
|
||||||
Args:
|
Args:
|
||||||
datasets: List of LeRobotDatasets to merge.
|
datasets: List of LeRobotDatasets to merge.
|
||||||
output_repo_id: Repository ID for the merged dataset.
|
output_repo_id: Merged dataset identifier.
|
||||||
output_dir: Directory to save the merged dataset. If None, uses default location.
|
output_dir: Root directory where the merged dataset will be stored. If not specified, defaults to $HF_LEROBOT_HOME/output_repo_id.
|
||||||
"""
|
"""
|
||||||
if not datasets:
|
if not datasets:
|
||||||
raise ValueError("No datasets to merge")
|
raise ValueError("No datasets to merge")
|
||||||
@@ -289,8 +291,8 @@ def modify_features(
|
|||||||
dataset: The source LeRobotDataset.
|
dataset: The source LeRobotDataset.
|
||||||
add_features: Optional dict mapping feature names to (feature_values, feature_info) tuples.
|
add_features: Optional dict mapping feature names to (feature_values, feature_info) tuples.
|
||||||
remove_features: Optional feature name(s) to remove. Can be a single string or list.
|
remove_features: Optional feature name(s) to remove. Can be a single string or list.
|
||||||
output_dir: Directory to save the new dataset. If None, uses default location.
|
output_dir: Root directory where the edited dataset will be stored. If not specified, defaults to $HF_LEROBOT_HOME/repo_id. Equivalent to new_root in EditDatasetConfig.
|
||||||
repo_id: Repository ID for the new dataset. If None, appends "_modified" to original.
|
repo_id: Edited dataset identifier. Equivalent to new_repo_id in EditDatasetConfig.
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
New dataset with features modified.
|
New dataset with features modified.
|
||||||
@@ -391,8 +393,8 @@ def add_features(
|
|||||||
Args:
|
Args:
|
||||||
dataset: The source LeRobotDataset.
|
dataset: The source LeRobotDataset.
|
||||||
features: Dictionary mapping feature names to (feature_values, feature_info) tuples.
|
features: Dictionary mapping feature names to (feature_values, feature_info) tuples.
|
||||||
output_dir: Directory to save the new dataset. If None, uses default location.
|
output_dir: Root directory where the edited dataset will be stored. If not specified, defaults to $HF_LEROBOT_HOME/repo_id. Equivalent to new_root in EditDatasetConfig.
|
||||||
repo_id: Repository ID for the new dataset. If None, appends "_modified" to original.
|
repo_id: Edited dataset identifier. Equivalent to new_repo_id in EditDatasetConfig.
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
New dataset with all features added.
|
New dataset with all features added.
|
||||||
@@ -428,8 +430,8 @@ def remove_feature(
|
|||||||
Args:
|
Args:
|
||||||
dataset: The source LeRobotDataset.
|
dataset: The source LeRobotDataset.
|
||||||
feature_names: Name(s) of features to remove. Can be a single string or list.
|
feature_names: Name(s) of features to remove. Can be a single string or list.
|
||||||
output_dir: Directory to save the new dataset. If None, uses default location.
|
output_dir: Root directory where the edited dataset will be stored. If not specified, defaults to $HF_LEROBOT_HOME/repo_id. Equivalent to new_root in EditDatasetConfig.
|
||||||
repo_id: Repository ID for the new dataset. If None, appends "_modified" to original.
|
repo_id: Edited dataset identifier. Equivalent to new_repo_id in EditDatasetConfig.
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
New dataset with features removed.
|
New dataset with features removed.
|
||||||
@@ -497,13 +499,16 @@ def _copy_and_reindex_data(
|
|||||||
global_index = 0
|
global_index = 0
|
||||||
episode_data_metadata: dict[int, dict] = {}
|
episode_data_metadata: dict[int, dict] = {}
|
||||||
|
|
||||||
|
episode_keys = list(episode_mapping.keys())
|
||||||
|
ep_filter = pa_ds.field("episode_index").isin(episode_keys)
|
||||||
|
|
||||||
if dst_meta.tasks is None:
|
if dst_meta.tasks is None:
|
||||||
all_task_indices = set()
|
all_task_indices: set[int] = set()
|
||||||
for src_path in file_to_episodes:
|
for src_path in file_to_episodes:
|
||||||
df = pd.read_parquet(src_dataset.root / src_path)
|
table = pq.read_table(
|
||||||
mask = df["episode_index"].isin(list(episode_mapping.keys()))
|
src_dataset.root / src_path, columns=["episode_index", "task_index"], filters=ep_filter
|
||||||
task_series: pd.Series = df[mask]["task_index"]
|
)
|
||||||
all_task_indices.update(task_series.unique().tolist())
|
all_task_indices.update(pc.unique(table.column("task_index")).to_pylist())
|
||||||
tasks = [src_dataset.meta.tasks.iloc[idx].name for idx in all_task_indices]
|
tasks = [src_dataset.meta.tasks.iloc[idx].name for idx in all_task_indices]
|
||||||
dst_meta.save_episode_tasks(list(set(tasks)))
|
dst_meta.save_episode_tasks(list(set(tasks)))
|
||||||
|
|
||||||
@@ -515,30 +520,17 @@ def _copy_and_reindex_data(
|
|||||||
task_mapping[old_task_idx] = new_task_idx
|
task_mapping[old_task_idx] = new_task_idx
|
||||||
|
|
||||||
for src_path in tqdm(sorted(file_to_episodes.keys()), desc="Processing data files"):
|
for src_path in tqdm(sorted(file_to_episodes.keys()), desc="Processing data files"):
|
||||||
df = pd.read_parquet(src_dataset.root / src_path)
|
table = pq.read_table(src_dataset.root / src_path, filters=ep_filter)
|
||||||
|
|
||||||
all_episodes_in_file = set(df["episode_index"].unique())
|
|
||||||
episodes_to_keep = file_to_episodes[src_path]
|
episodes_to_keep = file_to_episodes[src_path]
|
||||||
|
|
||||||
if all_episodes_in_file == episodes_to_keep:
|
if table.num_rows == 0:
|
||||||
df["episode_index"] = df["episode_index"].replace(episode_mapping)
|
|
||||||
df["index"] = range(global_index, global_index + len(df))
|
|
||||||
df["task_index"] = df["task_index"].replace(task_mapping)
|
|
||||||
|
|
||||||
first_ep_old_idx = min(episodes_to_keep)
|
|
||||||
src_ep = src_dataset.meta.episodes[first_ep_old_idx]
|
|
||||||
chunk_idx = src_ep["data/chunk_index"]
|
|
||||||
file_idx = src_ep["data/file_index"]
|
|
||||||
else:
|
|
||||||
mask = df["episode_index"].isin(list(episode_mapping.keys()))
|
|
||||||
df = df[mask].copy().reset_index(drop=True)
|
|
||||||
|
|
||||||
if len(df) == 0:
|
|
||||||
continue
|
continue
|
||||||
|
|
||||||
df["episode_index"] = df["episode_index"].replace(episode_mapping)
|
table = _replace_column_values(table, "episode_index", episode_mapping)
|
||||||
df["index"] = range(global_index, global_index + len(df))
|
col_pos = table.column_names.index("index")
|
||||||
df["task_index"] = df["task_index"].replace(task_mapping)
|
new_indices = pa.array(range(global_index, global_index + table.num_rows), type=pa.int64())
|
||||||
|
table = table.set_column(col_pos, "index", new_indices)
|
||||||
|
table = _replace_column_values(table, "task_index", task_mapping)
|
||||||
|
|
||||||
first_ep_old_idx = min(episodes_to_keep)
|
first_ep_old_idx = min(episodes_to_keep)
|
||||||
src_ep = src_dataset.meta.episodes[first_ep_old_idx]
|
src_ep = src_dataset.meta.episodes[first_ep_old_idx]
|
||||||
@@ -548,19 +540,21 @@ def _copy_and_reindex_data(
|
|||||||
dst_path = dst_meta.root / DEFAULT_DATA_PATH.format(chunk_index=chunk_idx, file_index=file_idx)
|
dst_path = dst_meta.root / DEFAULT_DATA_PATH.format(chunk_index=chunk_idx, file_index=file_idx)
|
||||||
dst_path.parent.mkdir(parents=True, exist_ok=True)
|
dst_path.parent.mkdir(parents=True, exist_ok=True)
|
||||||
|
|
||||||
_write_parquet(df, dst_path, dst_meta)
|
_write_parquet(table, dst_path, dst_meta)
|
||||||
|
|
||||||
|
ep_col = table.column("episode_index").to_pylist()
|
||||||
|
idx_col = table.column("index").to_pylist()
|
||||||
for ep_old_idx in episodes_to_keep:
|
for ep_old_idx in episodes_to_keep:
|
||||||
ep_new_idx = episode_mapping[ep_old_idx]
|
ep_new_idx = episode_mapping[ep_old_idx]
|
||||||
ep_df = df[df["episode_index"] == ep_new_idx]
|
ep_indices = [idx_col[i] for i, e in enumerate(ep_col) if e == ep_new_idx]
|
||||||
episode_data_metadata[ep_new_idx] = {
|
episode_data_metadata[ep_new_idx] = {
|
||||||
"data/chunk_index": chunk_idx,
|
"data/chunk_index": chunk_idx,
|
||||||
"data/file_index": file_idx,
|
"data/file_index": file_idx,
|
||||||
"dataset_from_index": int(ep_df["index"].min()),
|
"dataset_from_index": min(ep_indices),
|
||||||
"dataset_to_index": int(ep_df["index"].max() + 1),
|
"dataset_to_index": max(ep_indices) + 1,
|
||||||
}
|
}
|
||||||
|
|
||||||
global_index += len(df)
|
global_index += table.num_rows
|
||||||
|
|
||||||
return episode_data_metadata
|
return episode_data_metadata
|
||||||
|
|
||||||
@@ -911,15 +905,39 @@ def _copy_and_reindex_episodes_metadata(
|
|||||||
write_stats(filtered_stats, dst_meta.root)
|
write_stats(filtered_stats, dst_meta.root)
|
||||||
|
|
||||||
|
|
||||||
def _write_parquet(df: pd.DataFrame, path: Path, meta: LeRobotDatasetMetadata) -> None:
|
def _replace_column_values(table: pa.Table, column: str, mapping: dict) -> pa.Table:
|
||||||
"""Write DataFrame to parquet
|
"""Replace values in a pyarrow Table column using a mapping dict."""
|
||||||
|
old_values = table.column(column).to_pylist()
|
||||||
|
new_values = [mapping.get(v, v) for v in old_values]
|
||||||
|
col_pos = table.column_names.index(column)
|
||||||
|
return table.set_column(col_pos, column, pa.array(new_values, type=table.schema.field(column).type))
|
||||||
|
|
||||||
|
|
||||||
|
def _write_parquet(
|
||||||
|
data: pd.DataFrame | pa.Table | dict, path: Path, meta: LeRobotDatasetMetadata
|
||||||
|
) -> None:
|
||||||
|
"""Write data to parquet.
|
||||||
|
|
||||||
This ensures images are properly embedded and the file can be loaded correctly by HF datasets.
|
This ensures images are properly embedded and the file can be loaded correctly by HF datasets.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
data: Input data as a pandas DataFrame, pyarrow Table, or dict of lists.
|
||||||
|
path: Destination parquet file path.
|
||||||
|
meta: Dataset metadata for feature schema.
|
||||||
"""
|
"""
|
||||||
from lerobot.datasets.utils import embed_images, get_hf_features_from_features
|
from lerobot.datasets.utils import embed_images, get_hf_features_from_features
|
||||||
|
|
||||||
|
if isinstance(data, pd.DataFrame):
|
||||||
|
data_dict = data.to_dict(orient="list")
|
||||||
|
elif isinstance(data, pa.Table):
|
||||||
|
data_dict = data.to_pydict()
|
||||||
|
elif isinstance(data, dict):
|
||||||
|
data_dict = data
|
||||||
|
else:
|
||||||
|
raise TypeError(f"Unsupported data type: {type(data)}")
|
||||||
|
|
||||||
hf_features = get_hf_features_from_features(meta.features)
|
hf_features = get_hf_features_from_features(meta.features)
|
||||||
ep_dataset = datasets.Dataset.from_dict(df.to_dict(orient="list"), features=hf_features, split="train")
|
ep_dataset = datasets.Dataset.from_dict(data_dict, features=hf_features, split="train")
|
||||||
|
|
||||||
if len(meta.image_keys) > 0:
|
if len(meta.image_keys) > 0:
|
||||||
ep_dataset = embed_images(ep_dataset)
|
ep_dataset = embed_images(ep_dataset)
|
||||||
@@ -1476,7 +1494,9 @@ def modify_tasks(
|
|||||||
|
|
||||||
# Collect all unique tasks and create new task mapping
|
# Collect all unique tasks and create new task mapping
|
||||||
unique_tasks = sorted(set(episode_to_task.values()))
|
unique_tasks = sorted(set(episode_to_task.values()))
|
||||||
new_task_df = pd.DataFrame({"task_index": list(range(len(unique_tasks)))}, index=unique_tasks)
|
new_task_df = pd.DataFrame(
|
||||||
|
{"task_index": list(range(len(unique_tasks)))}, index=pd.Index(unique_tasks, name="task")
|
||||||
|
)
|
||||||
task_to_index = {task: idx for idx, task in enumerate(unique_tasks)}
|
task_to_index = {task: idx for idx, task in enumerate(unique_tasks)}
|
||||||
|
|
||||||
logging.info(f"Modifying tasks in {dataset.repo_id}")
|
logging.info(f"Modifying tasks in {dataset.repo_id}")
|
||||||
@@ -1530,7 +1550,7 @@ def modify_tasks(
|
|||||||
|
|
||||||
def convert_image_to_video_dataset(
|
def convert_image_to_video_dataset(
|
||||||
dataset: LeRobotDataset,
|
dataset: LeRobotDataset,
|
||||||
output_dir: Path,
|
output_dir: Path | None = None,
|
||||||
repo_id: str | None = None,
|
repo_id: str | None = None,
|
||||||
vcodec: str = "libsvtav1",
|
vcodec: str = "libsvtav1",
|
||||||
pix_fmt: str = "yuv420p",
|
pix_fmt: str = "yuv420p",
|
||||||
@@ -1549,8 +1569,8 @@ def convert_image_to_video_dataset(
|
|||||||
|
|
||||||
Args:
|
Args:
|
||||||
dataset: The source LeRobot dataset with images
|
dataset: The source LeRobot dataset with images
|
||||||
output_dir: Directory to save the new video dataset
|
output_dir: Root directory where the edited dataset will be stored. If not specified, defaults to $HF_LEROBOT_HOME/repo_id. Equivalent to new_root in EditDatasetConfig.
|
||||||
repo_id: Repository ID for the new dataset (default: original_id + "_video")
|
repo_id: Edited dataset identifier. Equivalent to new_repo_id in EditDatasetConfig.
|
||||||
vcodec: Video codec (default: libsvtav1)
|
vcodec: Video codec (default: libsvtav1)
|
||||||
pix_fmt: Pixel format (default: yuv420p)
|
pix_fmt: Pixel format (default: yuv420p)
|
||||||
g: Group of pictures size (default: 2)
|
g: Group of pictures size (default: 2)
|
||||||
@@ -1601,6 +1621,7 @@ def convert_image_to_video_dataset(
|
|||||||
# Video info will be updated after episodes are encoded
|
# Video info will be updated after episodes are encoded
|
||||||
|
|
||||||
# Create new metadata for video dataset
|
# Create new metadata for video dataset
|
||||||
|
output_dir = Path(output_dir) if output_dir is not None else HF_LEROBOT_HOME / repo_id
|
||||||
new_meta = LeRobotDatasetMetadata.create(
|
new_meta = LeRobotDatasetMetadata.create(
|
||||||
repo_id=repo_id,
|
repo_id=repo_id,
|
||||||
fps=dataset.meta.fps,
|
fps=dataset.meta.fps,
|
||||||
@@ -1775,296 +1796,3 @@ def convert_image_to_video_dataset(
|
|||||||
|
|
||||||
# Return new dataset
|
# Return new dataset
|
||||||
return LeRobotDataset(repo_id=repo_id, root=output_dir)
|
return LeRobotDataset(repo_id=repo_id, root=output_dir)
|
||||||
|
|
||||||
|
|
||||||
def trim_episodes_by_frames(
|
|
||||||
dataset: LeRobotDataset,
|
|
||||||
episode_frames_to_keep: dict[int, list[int]],
|
|
||||||
output_dir: str | Path | None = None,
|
|
||||||
repo_id: str | None = None,
|
|
||||||
) -> LeRobotDataset:
|
|
||||||
"""Trim multiple episodes to keep only specific frames.
|
|
||||||
|
|
||||||
This function creates a new dataset where the specified episodes contain only
|
|
||||||
the frames at the given indices. All other episodes are copied as-is.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
dataset: The source LeRobotDataset.
|
|
||||||
episode_frames_to_keep: Dict mapping episode indices to lists of global frame indices to keep.
|
|
||||||
output_dir: Directory to save the new dataset. If None, uses default location.
|
|
||||||
repo_id: Repository ID for the new dataset. If None, appends "_trimmed" to original.
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
A new LeRobotDataset with the trimmed episodes.
|
|
||||||
"""
|
|
||||||
if not episode_frames_to_keep:
|
|
||||||
raise ValueError("No episodes to trim")
|
|
||||||
|
|
||||||
for ep_idx in episode_frames_to_keep:
|
|
||||||
if ep_idx >= dataset.meta.total_episodes:
|
|
||||||
raise ValueError(f"Episode {ep_idx} does not exist")
|
|
||||||
if not episode_frames_to_keep[ep_idx]:
|
|
||||||
raise ValueError(f"No frames to keep for episode {ep_idx}")
|
|
||||||
|
|
||||||
if repo_id is None:
|
|
||||||
repo_id = f"{dataset.repo_id}_trimmed"
|
|
||||||
output_dir = Path(output_dir) if output_dir is not None else HF_LEROBOT_HOME / repo_id
|
|
||||||
|
|
||||||
total_trimmed = sum(len(frames) for frames in episode_frames_to_keep.values())
|
|
||||||
logging.info(f"Trimming {len(episode_frames_to_keep)} episodes, keeping {total_trimmed} frames total")
|
|
||||||
|
|
||||||
# Create new metadata
|
|
||||||
new_meta = LeRobotDatasetMetadata.create(
|
|
||||||
repo_id=repo_id,
|
|
||||||
fps=dataset.meta.fps,
|
|
||||||
features=dataset.meta.features,
|
|
||||||
robot_type=dataset.meta.robot_type,
|
|
||||||
root=output_dir,
|
|
||||||
use_videos=len(dataset.meta.video_keys) > 0,
|
|
||||||
)
|
|
||||||
|
|
||||||
# Build set of all frames to keep (for episodes being trimmed)
|
|
||||||
# and compute new frame counts per episode
|
|
||||||
all_keep_frames: set[int] = set()
|
|
||||||
trimmed_frame_counts: dict[int, int] = {}
|
|
||||||
for ep_idx, frames in episode_frames_to_keep.items():
|
|
||||||
all_keep_frames.update(frames)
|
|
||||||
trimmed_frame_counts[ep_idx] = len(frames)
|
|
||||||
|
|
||||||
# Copy and filter data
|
|
||||||
_copy_and_reindex_data_with_multi_frame_filter(
|
|
||||||
dataset, new_meta, episode_frames_to_keep, all_keep_frames
|
|
||||||
)
|
|
||||||
|
|
||||||
# Handle videos if present
|
|
||||||
if dataset.meta.video_keys:
|
|
||||||
_copy_and_reindex_videos_with_multi_frame_filter(
|
|
||||||
dataset, new_meta, episode_frames_to_keep
|
|
||||||
)
|
|
||||||
|
|
||||||
# Copy episode metadata
|
|
||||||
_copy_and_reindex_episodes_metadata_for_multi_trim(
|
|
||||||
dataset, new_meta, trimmed_frame_counts
|
|
||||||
)
|
|
||||||
|
|
||||||
logging.info(f"Created trimmed dataset with {new_meta.total_frames} frames at {output_dir}")
|
|
||||||
|
|
||||||
# Return the metadata instead of trying to load as LeRobotDataset
|
|
||||||
# This avoids Hub validation issues when the repo doesn't exist yet
|
|
||||||
return new_meta
|
|
||||||
|
|
||||||
|
|
||||||
# Keep old function for backward compatibility
|
|
||||||
def trim_episode_by_frames(
|
|
||||||
dataset: LeRobotDataset,
|
|
||||||
episode_index: int,
|
|
||||||
keep_frame_indices: list[int],
|
|
||||||
output_dir: str | Path | None = None,
|
|
||||||
repo_id: str | None = None,
|
|
||||||
) -> LeRobotDataset:
|
|
||||||
"""Trim a single episode. Wrapper around trim_episodes_by_frames."""
|
|
||||||
return trim_episodes_by_frames(
|
|
||||||
dataset,
|
|
||||||
episode_frames_to_keep={episode_index: keep_frame_indices},
|
|
||||||
output_dir=output_dir,
|
|
||||||
repo_id=repo_id,
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
def _copy_and_reindex_data_with_multi_frame_filter(
|
|
||||||
src_dataset: LeRobotDataset,
|
|
||||||
dst_meta: LeRobotDatasetMetadata,
|
|
||||||
episode_frames_to_keep: dict[int, list[int]],
|
|
||||||
all_keep_frames: set[int],
|
|
||||||
) -> None:
|
|
||||||
"""Copy data files with frame-level filtering for multiple episodes."""
|
|
||||||
if src_dataset.meta.episodes is None:
|
|
||||||
src_dataset.meta.episodes = load_episodes(src_dataset.meta.root)
|
|
||||||
|
|
||||||
# Copy tasks
|
|
||||||
if dst_meta.tasks is None and src_dataset.meta.tasks is not None:
|
|
||||||
# Tasks are stored with task string as index
|
|
||||||
dst_meta.save_episode_tasks(list(src_dataset.meta.tasks.index))
|
|
||||||
|
|
||||||
# Get all parquet files
|
|
||||||
data_dir = src_dataset.root / "data"
|
|
||||||
parquet_files = sorted(data_dir.glob("chunk-*/file-*.parquet"))
|
|
||||||
|
|
||||||
trim_episode_set = set(episode_frames_to_keep.keys())
|
|
||||||
global_index = 0
|
|
||||||
|
|
||||||
for parquet_path in tqdm(parquet_files, desc="Processing data files"):
|
|
||||||
df = pd.read_parquet(parquet_path)
|
|
||||||
|
|
||||||
# Filter: keep all frames from non-trimmed episodes,
|
|
||||||
# and only specified frames from trimmed episodes
|
|
||||||
mask = (~df["episode_index"].isin(trim_episode_set)) | (df["index"].isin(all_keep_frames))
|
|
||||||
df = df[mask].copy().reset_index(drop=True)
|
|
||||||
|
|
||||||
if len(df) == 0:
|
|
||||||
continue
|
|
||||||
|
|
||||||
# Reindex
|
|
||||||
df["index"] = range(global_index, global_index + len(df))
|
|
||||||
|
|
||||||
# Recalculate frame_index within each episode
|
|
||||||
for ep_idx in df["episode_index"].unique():
|
|
||||||
ep_mask = df["episode_index"] == ep_idx
|
|
||||||
df.loc[ep_mask, "frame_index"] = range(ep_mask.sum())
|
|
||||||
|
|
||||||
# Recalculate timestamps based on frame_index and fps
|
|
||||||
df["timestamp"] = df["frame_index"] / src_dataset.meta.fps
|
|
||||||
|
|
||||||
# Determine output path (keep same structure)
|
|
||||||
rel_path = parquet_path.relative_to(src_dataset.root)
|
|
||||||
dst_path = dst_meta.root / rel_path
|
|
||||||
dst_path.parent.mkdir(parents=True, exist_ok=True)
|
|
||||||
|
|
||||||
_write_parquet(df, dst_path, dst_meta)
|
|
||||||
global_index += len(df)
|
|
||||||
|
|
||||||
|
|
||||||
def _copy_and_reindex_videos_with_multi_frame_filter(
|
|
||||||
src_dataset: LeRobotDataset,
|
|
||||||
dst_meta: LeRobotDatasetMetadata,
|
|
||||||
episode_frames_to_keep: dict[int, list[int]],
|
|
||||||
) -> None:
|
|
||||||
"""Copy video files for trimmed dataset.
|
|
||||||
|
|
||||||
In v3.0 datasets, multiple episodes are concatenated into single video files.
|
|
||||||
Each episode has from_timestamp/to_timestamp indicating its portion of the video.
|
|
||||||
|
|
||||||
For trimming, we copy the original video files as-is and update the metadata
|
|
||||||
timestamps in _copy_and_reindex_episodes_metadata_for_multi_trim.
|
|
||||||
"""
|
|
||||||
for video_key in src_dataset.meta.video_keys:
|
|
||||||
video_dir = src_dataset.root / "videos" / video_key
|
|
||||||
dst_video_dir = dst_meta.root / "videos" / video_key
|
|
||||||
|
|
||||||
if not video_dir.exists():
|
|
||||||
logging.warning(f"Video directory not found: {video_dir}")
|
|
||||||
continue
|
|
||||||
|
|
||||||
# Copy all video files (they contain concatenated episodes)
|
|
||||||
# The metadata timestamps will handle which portions to use
|
|
||||||
copied_files = set()
|
|
||||||
for chunk_dir in video_dir.glob("chunk-*"):
|
|
||||||
dst_chunk_dir = dst_video_dir / chunk_dir.name
|
|
||||||
dst_chunk_dir.mkdir(parents=True, exist_ok=True)
|
|
||||||
|
|
||||||
for video_file in chunk_dir.glob("*.mp4"):
|
|
||||||
if video_file.name not in copied_files:
|
|
||||||
dst_path = dst_chunk_dir / video_file.name
|
|
||||||
if not dst_path.exists():
|
|
||||||
shutil.copy(video_file, dst_path)
|
|
||||||
copied_files.add(video_file.name)
|
|
||||||
|
|
||||||
logging.info(f"Copied {len(copied_files)} video files for {video_key}")
|
|
||||||
|
|
||||||
|
|
||||||
def _trim_video_frames(
|
|
||||||
src_path: Path,
|
|
||||||
dst_path: Path,
|
|
||||||
keep_frame_indices: list[int],
|
|
||||||
fps: float,
|
|
||||||
episode_start_idx: int,
|
|
||||||
) -> None:
|
|
||||||
"""Trim a video to keep only specific frames using ffmpeg."""
|
|
||||||
import subprocess
|
|
||||||
|
|
||||||
# Convert global indices to local indices within the episode
|
|
||||||
local_indices = sorted([idx - episode_start_idx for idx in keep_frame_indices])
|
|
||||||
|
|
||||||
if not local_indices:
|
|
||||||
logging.warning(f"No frames to keep for video {src_path}")
|
|
||||||
return
|
|
||||||
|
|
||||||
# Calculate start and end times
|
|
||||||
start_frame = local_indices[0]
|
|
||||||
end_frame = local_indices[-1]
|
|
||||||
|
|
||||||
start_time = start_frame / fps
|
|
||||||
duration = (end_frame - start_frame + 1) / fps
|
|
||||||
|
|
||||||
# Use ffmpeg to trim
|
|
||||||
cmd = [
|
|
||||||
"ffmpeg", "-y",
|
|
||||||
"-ss", str(start_time),
|
|
||||||
"-i", str(src_path),
|
|
||||||
"-t", str(duration),
|
|
||||||
"-c", "copy", # Fast copy without re-encoding
|
|
||||||
str(dst_path)
|
|
||||||
]
|
|
||||||
|
|
||||||
try:
|
|
||||||
subprocess.run(cmd, check=True, capture_output=True)
|
|
||||||
except subprocess.CalledProcessError as e:
|
|
||||||
logging.error(f"Failed to trim video: {e.stderr.decode()}")
|
|
||||||
# Fallback: copy the whole video
|
|
||||||
shutil.copy(src_path, dst_path)
|
|
||||||
|
|
||||||
|
|
||||||
def _copy_and_reindex_episodes_metadata_for_multi_trim(
|
|
||||||
src_dataset: LeRobotDataset,
|
|
||||||
dst_meta: LeRobotDatasetMetadata,
|
|
||||||
trimmed_frame_counts: dict[int, int],
|
|
||||||
) -> None:
|
|
||||||
"""Copy and update episode metadata for trimmed dataset."""
|
|
||||||
if src_dataset.meta.episodes is None:
|
|
||||||
src_dataset.meta.episodes = load_episodes(src_dataset.meta.root)
|
|
||||||
|
|
||||||
# Calculate new frame counts and indices
|
|
||||||
episodes_data = []
|
|
||||||
global_idx = 0
|
|
||||||
|
|
||||||
for old_ep_idx in range(src_dataset.meta.total_episodes):
|
|
||||||
src_ep = src_dataset.meta.episodes[old_ep_idx]
|
|
||||||
|
|
||||||
if old_ep_idx in trimmed_frame_counts:
|
|
||||||
ep_length = trimmed_frame_counts[old_ep_idx]
|
|
||||||
else:
|
|
||||||
ep_length = src_ep["length"]
|
|
||||||
|
|
||||||
ep_data = {
|
|
||||||
"episode_index": old_ep_idx,
|
|
||||||
"tasks": src_ep["tasks"],
|
|
||||||
"length": ep_length,
|
|
||||||
"data/chunk_index": src_ep["data/chunk_index"],
|
|
||||||
"data/file_index": src_ep["data/file_index"],
|
|
||||||
"dataset_from_index": global_idx,
|
|
||||||
"dataset_to_index": global_idx + ep_length,
|
|
||||||
}
|
|
||||||
|
|
||||||
# Copy video metadata - preserve timestamps for concatenated videos
|
|
||||||
for video_key in src_dataset.meta.video_keys:
|
|
||||||
ep_data[f"videos/{video_key}/chunk_index"] = src_ep[f"videos/{video_key}/chunk_index"]
|
|
||||||
ep_data[f"videos/{video_key}/file_index"] = src_ep[f"videos/{video_key}/file_index"]
|
|
||||||
|
|
||||||
# Keep original from_timestamp (start position in concatenated video)
|
|
||||||
orig_from_ts = src_ep[f"videos/{video_key}/from_timestamp"]
|
|
||||||
ep_data[f"videos/{video_key}/from_timestamp"] = orig_from_ts
|
|
||||||
|
|
||||||
# For trimmed episodes, update to_timestamp based on new length
|
|
||||||
# For non-trimmed episodes, keep original to_timestamp
|
|
||||||
if old_ep_idx in trimmed_frame_counts:
|
|
||||||
ep_data[f"videos/{video_key}/to_timestamp"] = orig_from_ts + (ep_length / src_dataset.meta.fps)
|
|
||||||
else:
|
|
||||||
ep_data[f"videos/{video_key}/to_timestamp"] = src_ep[f"videos/{video_key}/to_timestamp"]
|
|
||||||
|
|
||||||
ep_data["meta/episodes/chunk_index"] = 0
|
|
||||||
ep_data["meta/episodes/file_index"] = 0
|
|
||||||
|
|
||||||
episodes_data.append(ep_data)
|
|
||||||
global_idx += ep_length
|
|
||||||
|
|
||||||
# Save episodes metadata
|
|
||||||
df = pd.DataFrame(episodes_data)
|
|
||||||
episodes_path = dst_meta.root / DEFAULT_EPISODES_PATH.format(chunk_index=0, file_index=0)
|
|
||||||
episodes_path.parent.mkdir(parents=True, exist_ok=True)
|
|
||||||
df.to_parquet(episodes_path)
|
|
||||||
|
|
||||||
# Update info.json
|
|
||||||
info = load_info(src_dataset.root)
|
|
||||||
info["total_episodes"] = len(episodes_data)
|
|
||||||
info["total_frames"] = global_idx
|
|
||||||
write_info(info, dst_meta.root)
|
|
||||||
|
|||||||
@@ -314,7 +314,7 @@ class LeRobotDatasetMetadata:
|
|||||||
if self.tasks is None:
|
if self.tasks is None:
|
||||||
new_tasks = tasks
|
new_tasks = tasks
|
||||||
task_indices = range(len(tasks))
|
task_indices = range(len(tasks))
|
||||||
self.tasks = pd.DataFrame({"task_index": task_indices}, index=tasks)
|
self.tasks = pd.DataFrame({"task_index": task_indices}, index=pd.Index(tasks, name="task"))
|
||||||
else:
|
else:
|
||||||
new_tasks = [task for task in tasks if task not in self.tasks.index]
|
new_tasks = [task for task in tasks if task not in self.tasks.index]
|
||||||
new_task_indices = range(len(self.tasks), len(self.tasks) + len(new_tasks))
|
new_task_indices = range(len(self.tasks), len(self.tasks) + len(new_tasks))
|
||||||
@@ -664,11 +664,11 @@ class LeRobotDataset(torch.utils.data.Dataset):
|
|||||||
for the README).
|
for the README).
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
repo_id (str): This is the repo id that will be used to fetch the dataset. Locally, the dataset
|
repo_id (str): This is the repo id that will be used to fetch the dataset.
|
||||||
will be stored under root/repo_id.
|
root (Path | None, optional): Local directory where the dataset will be downloaded and
|
||||||
root (Path | None, optional): Local directory to use for downloading/writing files. You can also
|
stored. If set, all dataset files will be stored directly under this path. If not set, the
|
||||||
set the HF_LEROBOT_HOME environment variable to point to a different location. Defaults to
|
dataset files will be stored under $HF_LEROBOT_HOME/repo_id (configurable via the
|
||||||
'~/.cache/huggingface/lerobot'.
|
HF_LEROBOT_HOME environment variable).
|
||||||
episodes (list[int] | None, optional): If specified, this will only load episodes specified by
|
episodes (list[int] | None, optional): If specified, this will only load episodes specified by
|
||||||
their episode_index in this list. Defaults to None.
|
their episode_index in this list. Defaults to None.
|
||||||
image_transforms (Callable | None, optional): You can pass standard v2 image transforms from
|
image_transforms (Callable | None, optional): You can pass standard v2 image transforms from
|
||||||
@@ -1771,6 +1771,7 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
|
|||||||
)
|
)
|
||||||
for repo_id, ds in zip(self.repo_ids, self._datasets, strict=True):
|
for repo_id, ds in zip(self.repo_ids, self._datasets, strict=True):
|
||||||
extra_keys = set(ds.features).difference(intersection_features)
|
extra_keys = set(ds.features).difference(intersection_features)
|
||||||
|
if extra_keys:
|
||||||
logging.warning(
|
logging.warning(
|
||||||
f"keys {extra_keys} of {repo_id} were disabled as they are not contained in all the "
|
f"keys {extra_keys} of {repo_id} were disabled as they are not contained in all the "
|
||||||
"other datasets."
|
"other datasets."
|
||||||
|
|||||||
@@ -341,6 +341,7 @@ def write_tasks(tasks: pandas.DataFrame, local_dir: Path) -> None:
|
|||||||
|
|
||||||
def load_tasks(local_dir: Path) -> pandas.DataFrame:
|
def load_tasks(local_dir: Path) -> pandas.DataFrame:
|
||||||
tasks = pd.read_parquet(local_dir / DEFAULT_TASKS_PATH)
|
tasks = pd.read_parquet(local_dir / DEFAULT_TASKS_PATH)
|
||||||
|
tasks.index.name = "task"
|
||||||
return tasks
|
return tasks
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -36,8 +36,11 @@ Convert a local dataset (works in place):
|
|||||||
```bash
|
```bash
|
||||||
python src/lerobot/datasets/v30/convert_dataset_v21_to_v30.py \
|
python src/lerobot/datasets/v30/convert_dataset_v21_to_v30.py \
|
||||||
--repo-id=lerobot/pusht \
|
--repo-id=lerobot/pusht \
|
||||||
--root=/path/to/local/dataset/directory
|
--root=/path/to/local/dataset/directory \
|
||||||
--push-to-hub=false
|
--push-to-hub=false
|
||||||
|
|
||||||
|
N.B. Path semantics (v2): --root is the exact dataset folder containing
|
||||||
|
meta/, data/, videos/. When omitted, defaults to $HF_LEROBOT_HOME/{repo_id}.
|
||||||
```
|
```
|
||||||
|
|
||||||
"""
|
"""
|
||||||
@@ -105,7 +108,7 @@ episodes.jsonl
|
|||||||
{"episode_index": 1, "tasks": ["Put the blue block in the green bowl"], "length": 266}
|
{"episode_index": 1, "tasks": ["Put the blue block in the green bowl"], "length": 266}
|
||||||
|
|
||||||
NEW
|
NEW
|
||||||
meta/episodes/chunk-000/episodes_000.parquet
|
meta/episodes/chunk-000/file_000.parquet
|
||||||
episode_index | video_chunk_index | video_file_index | data_chunk_index | data_file_index | tasks | length
|
episode_index | video_chunk_index | video_file_index | data_chunk_index | data_file_index | tasks | length
|
||||||
-------------------------
|
-------------------------
|
||||||
OLD
|
OLD
|
||||||
@@ -113,15 +116,16 @@ tasks.jsonl
|
|||||||
{"task_index": 1, "task": "Put the blue block in the green bowl"}
|
{"task_index": 1, "task": "Put the blue block in the green bowl"}
|
||||||
|
|
||||||
NEW
|
NEW
|
||||||
meta/tasks/chunk-000/file_000.parquet
|
meta/tasks.parquet
|
||||||
task_index | task
|
task_index | task
|
||||||
-------------------------
|
-------------------------
|
||||||
OLD
|
OLD
|
||||||
episodes_stats.jsonl
|
episodes_stats.jsonl
|
||||||
|
{"episode_index": 1, "stats": {"feature_name": {"min": ..., "max": ..., "mean": ..., "std": ..., "count": ...}}}
|
||||||
|
|
||||||
NEW
|
NEW
|
||||||
meta/episodes_stats/chunk-000/file_000.parquet
|
meta/episodes/chunk-000/file_000.parquet
|
||||||
episode_index | mean | std | min | max
|
episode_index | feature_name/min | feature_name/max | feature_name/mean | feature_name/std | feature_name/count
|
||||||
-------------------------
|
-------------------------
|
||||||
UPDATE
|
UPDATE
|
||||||
meta/info.json
|
meta/info.json
|
||||||
@@ -170,7 +174,7 @@ def convert_tasks(root, new_root):
|
|||||||
tasks, _ = legacy_load_tasks(root)
|
tasks, _ = legacy_load_tasks(root)
|
||||||
task_indices = tasks.keys()
|
task_indices = tasks.keys()
|
||||||
task_strings = tasks.values()
|
task_strings = tasks.values()
|
||||||
df_tasks = pd.DataFrame({"task_index": task_indices}, index=task_strings)
|
df_tasks = pd.DataFrame({"task_index": task_indices}, index=pd.Index(task_strings, name="task"))
|
||||||
write_tasks(df_tasks, new_root)
|
write_tasks(df_tasks, new_root)
|
||||||
|
|
||||||
|
|
||||||
@@ -201,7 +205,6 @@ def convert_data(root: Path, new_root: Path, data_file_size_in_mb: int):
|
|||||||
|
|
||||||
image_keys = get_image_keys(root)
|
image_keys = get_image_keys(root)
|
||||||
|
|
||||||
ep_idx = 0
|
|
||||||
chunk_idx = 0
|
chunk_idx = 0
|
||||||
file_idx = 0
|
file_idx = 0
|
||||||
size_in_mb = 0
|
size_in_mb = 0
|
||||||
@@ -211,9 +214,24 @@ def convert_data(root: Path, new_root: Path, data_file_size_in_mb: int):
|
|||||||
|
|
||||||
logging.info(f"Converting data files from {len(ep_paths)} episodes")
|
logging.info(f"Converting data files from {len(ep_paths)} episodes")
|
||||||
|
|
||||||
for ep_path in tqdm.tqdm(ep_paths, desc="convert data files"):
|
for ep_idx, ep_path in enumerate(tqdm.tqdm(ep_paths, desc="convert data files")):
|
||||||
ep_size_in_mb = get_parquet_file_size_in_mb(ep_path)
|
ep_size_in_mb = get_parquet_file_size_in_mb(ep_path)
|
||||||
ep_num_frames = get_parquet_num_frames(ep_path)
|
ep_num_frames = get_parquet_num_frames(ep_path)
|
||||||
|
|
||||||
|
# Check if we need to start a new file BEFORE creating metadata
|
||||||
|
if size_in_mb + ep_size_in_mb >= data_file_size_in_mb and len(paths_to_cat) > 0:
|
||||||
|
# Write the accumulated data files
|
||||||
|
concat_data_files(paths_to_cat, new_root, chunk_idx, file_idx, image_keys)
|
||||||
|
|
||||||
|
# Move to next file
|
||||||
|
chunk_idx, file_idx = update_chunk_file_indices(chunk_idx, file_idx, DEFAULT_CHUNK_SIZE)
|
||||||
|
|
||||||
|
# Reset for the next file
|
||||||
|
size_in_mb = 0
|
||||||
|
num_frames += ep_num_frames # Still need to accumulate total frames
|
||||||
|
paths_to_cat = []
|
||||||
|
|
||||||
|
# Now create metadata with correct chunk/file indices
|
||||||
ep_metadata = {
|
ep_metadata = {
|
||||||
"episode_index": ep_idx,
|
"episode_index": ep_idx,
|
||||||
"data/chunk_index": chunk_idx,
|
"data/chunk_index": chunk_idx,
|
||||||
@@ -224,20 +242,7 @@ def convert_data(root: Path, new_root: Path, data_file_size_in_mb: int):
|
|||||||
size_in_mb += ep_size_in_mb
|
size_in_mb += ep_size_in_mb
|
||||||
num_frames += ep_num_frames
|
num_frames += ep_num_frames
|
||||||
episodes_metadata.append(ep_metadata)
|
episodes_metadata.append(ep_metadata)
|
||||||
ep_idx += 1
|
|
||||||
|
|
||||||
if size_in_mb < data_file_size_in_mb:
|
|
||||||
paths_to_cat.append(ep_path)
|
paths_to_cat.append(ep_path)
|
||||||
continue
|
|
||||||
|
|
||||||
if paths_to_cat:
|
|
||||||
concat_data_files(paths_to_cat, new_root, chunk_idx, file_idx, image_keys)
|
|
||||||
|
|
||||||
# Reset for the next file
|
|
||||||
size_in_mb = ep_size_in_mb
|
|
||||||
paths_to_cat = [ep_path]
|
|
||||||
|
|
||||||
chunk_idx, file_idx = update_chunk_file_indices(chunk_idx, file_idx, DEFAULT_CHUNK_SIZE)
|
|
||||||
|
|
||||||
# Write remaining data if any
|
# Write remaining data if any
|
||||||
if paths_to_cat:
|
if paths_to_cat:
|
||||||
@@ -469,7 +474,7 @@ def convert_dataset(
|
|||||||
|
|
||||||
# Set root based on whether local dataset path is provided
|
# Set root based on whether local dataset path is provided
|
||||||
use_local_dataset = False
|
use_local_dataset = False
|
||||||
root = HF_LEROBOT_HOME / repo_id if root is None else Path(root) / repo_id
|
root = HF_LEROBOT_HOME / repo_id if root is None else Path(root)
|
||||||
if root.exists():
|
if root.exists():
|
||||||
validate_local_dataset_version(root)
|
validate_local_dataset_version(root)
|
||||||
use_local_dataset = True
|
use_local_dataset = True
|
||||||
@@ -553,7 +558,7 @@ if __name__ == "__main__":
|
|||||||
"--root",
|
"--root",
|
||||||
type=str,
|
type=str,
|
||||||
default=None,
|
default=None,
|
||||||
help="Local directory to use for downloading/writing the dataset.",
|
help="Local directory to use for downloading/writing the dataset. Defaults to $HF_LEROBOT_HOME/repo_id.",
|
||||||
)
|
)
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
"--push-to-hub",
|
"--push-to-hub",
|
||||||
|
|||||||
@@ -55,10 +55,16 @@ class DiffusionConfig(PreTrainedConfig):
|
|||||||
normalization_mapping: A dictionary that maps from a str value of FeatureType (e.g., "STATE", "VISUAL") to
|
normalization_mapping: A dictionary that maps from a str value of FeatureType (e.g., "STATE", "VISUAL") to
|
||||||
a corresponding NormalizationMode (e.g., NormalizationMode.MIN_MAX)
|
a corresponding NormalizationMode (e.g., NormalizationMode.MIN_MAX)
|
||||||
vision_backbone: Name of the torchvision resnet backbone to use for encoding images.
|
vision_backbone: Name of the torchvision resnet backbone to use for encoding images.
|
||||||
crop_shape: (H, W) shape to crop images to as a preprocessing step for the vision backbone. Must fit
|
resize_shape: (H, W) shape to resize images to as a preprocessing step for the vision
|
||||||
within the image size. If None, no cropping is done.
|
backbone. If None, no resizing is done and the original image resolution is used.
|
||||||
crop_is_random: Whether the crop should be random at training time (it's always a center crop in eval
|
crop_ratio: Ratio in (0, 1] used to derive the crop size from resize_shape
|
||||||
mode).
|
(crop_h = int(resize_shape[0] * crop_ratio), likewise for width).
|
||||||
|
Set to 1.0 to disable cropping. Only takes effect when resize_shape is not None.
|
||||||
|
crop_shape: (H, W) shape to crop images to. When resize_shape is set and crop_ratio < 1.0,
|
||||||
|
this is computed automatically. Can also be set directly for legacy configs that use
|
||||||
|
crop-only (without resize). If None and no derivation applies, no cropping is done.
|
||||||
|
crop_is_random: Whether the crop should be random at training time (it's always a center
|
||||||
|
crop in eval mode).
|
||||||
pretrained_backbone_weights: Pretrained weights from torchvision to initialize the backbone.
|
pretrained_backbone_weights: Pretrained weights from torchvision to initialize the backbone.
|
||||||
`None` means no pretrained weights.
|
`None` means no pretrained weights.
|
||||||
use_group_norm: Whether to replace batch normalization with group normalization in the backbone.
|
use_group_norm: Whether to replace batch normalization with group normalization in the backbone.
|
||||||
@@ -114,7 +120,9 @@ class DiffusionConfig(PreTrainedConfig):
|
|||||||
# Architecture / modeling.
|
# Architecture / modeling.
|
||||||
# Vision backbone.
|
# Vision backbone.
|
||||||
vision_backbone: str = "resnet18"
|
vision_backbone: str = "resnet18"
|
||||||
crop_shape: tuple[int, int] | None = (84, 84)
|
resize_shape: tuple[int, int] | None = None
|
||||||
|
crop_ratio: float = 1.0
|
||||||
|
crop_shape: tuple[int, int] | None = None
|
||||||
crop_is_random: bool = True
|
crop_is_random: bool = True
|
||||||
pretrained_backbone_weights: str | None = None
|
pretrained_backbone_weights: str | None = None
|
||||||
use_group_norm: bool = True
|
use_group_norm: bool = True
|
||||||
@@ -175,6 +183,25 @@ class DiffusionConfig(PreTrainedConfig):
|
|||||||
f"Got {self.noise_scheduler_type}."
|
f"Got {self.noise_scheduler_type}."
|
||||||
)
|
)
|
||||||
|
|
||||||
|
if self.resize_shape is not None and (
|
||||||
|
len(self.resize_shape) != 2 or any(d <= 0 for d in self.resize_shape)
|
||||||
|
):
|
||||||
|
raise ValueError(f"`resize_shape` must be a pair of positive integers. Got {self.resize_shape}.")
|
||||||
|
if not (0 < self.crop_ratio <= 1.0):
|
||||||
|
raise ValueError(f"`crop_ratio` must be in (0, 1]. Got {self.crop_ratio}.")
|
||||||
|
|
||||||
|
if self.resize_shape is not None:
|
||||||
|
if self.crop_ratio < 1.0:
|
||||||
|
self.crop_shape = (
|
||||||
|
int(self.resize_shape[0] * self.crop_ratio),
|
||||||
|
int(self.resize_shape[1] * self.crop_ratio),
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
# Explicitly disable cropping for resize+ratio path when crop_ratio == 1.0.
|
||||||
|
self.crop_shape = None
|
||||||
|
if self.crop_shape is not None and (self.crop_shape[0] <= 0 or self.crop_shape[1] <= 0):
|
||||||
|
raise ValueError(f"`crop_shape` must have positive dimensions. Got {self.crop_shape}.")
|
||||||
|
|
||||||
# Check that the horizon size and U-Net downsampling is compatible.
|
# Check that the horizon size and U-Net downsampling is compatible.
|
||||||
# U-Net downsamples by 2 with each stage.
|
# U-Net downsamples by 2 with each stage.
|
||||||
downsampling_factor = 2 ** len(self.down_dims)
|
downsampling_factor = 2 ** len(self.down_dims)
|
||||||
@@ -202,13 +229,12 @@ class DiffusionConfig(PreTrainedConfig):
|
|||||||
if len(self.image_features) == 0 and self.env_state_feature is None:
|
if len(self.image_features) == 0 and self.env_state_feature is None:
|
||||||
raise ValueError("You must provide at least one image or the environment state among the inputs.")
|
raise ValueError("You must provide at least one image or the environment state among the inputs.")
|
||||||
|
|
||||||
if self.crop_shape is not None:
|
if self.resize_shape is None and self.crop_shape is not None:
|
||||||
for key, image_ft in self.image_features.items():
|
for key, image_ft in self.image_features.items():
|
||||||
if self.crop_shape[0] > image_ft.shape[1] or self.crop_shape[1] > image_ft.shape[2]:
|
if self.crop_shape[0] > image_ft.shape[1] or self.crop_shape[1] > image_ft.shape[2]:
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
f"`crop_shape` should fit within the images shapes. Got {self.crop_shape} "
|
f"`crop_shape` should fit within the image shapes. Got {self.crop_shape} "
|
||||||
f"for `crop_shape` and {image_ft.shape} for "
|
f"for `crop_shape` and {image_ft.shape} for `{key}`."
|
||||||
f"`{key}`."
|
|
||||||
)
|
)
|
||||||
|
|
||||||
# Check that all input images have the same shape.
|
# Check that all input images have the same shape.
|
||||||
|
|||||||
@@ -454,12 +454,18 @@ class DiffusionRgbEncoder(nn.Module):
|
|||||||
def __init__(self, config: DiffusionConfig):
|
def __init__(self, config: DiffusionConfig):
|
||||||
super().__init__()
|
super().__init__()
|
||||||
# Set up optional preprocessing.
|
# Set up optional preprocessing.
|
||||||
if config.crop_shape is not None:
|
if config.resize_shape is not None:
|
||||||
|
self.resize = torchvision.transforms.Resize(config.resize_shape)
|
||||||
|
else:
|
||||||
|
self.resize = None
|
||||||
|
|
||||||
|
crop_shape = config.crop_shape
|
||||||
|
if crop_shape is not None:
|
||||||
self.do_crop = True
|
self.do_crop = True
|
||||||
# Always use center crop for eval
|
# Always use center crop for eval
|
||||||
self.center_crop = torchvision.transforms.CenterCrop(config.crop_shape)
|
self.center_crop = torchvision.transforms.CenterCrop(crop_shape)
|
||||||
if config.crop_is_random:
|
if config.crop_is_random:
|
||||||
self.maybe_random_crop = torchvision.transforms.RandomCrop(config.crop_shape)
|
self.maybe_random_crop = torchvision.transforms.RandomCrop(crop_shape)
|
||||||
else:
|
else:
|
||||||
self.maybe_random_crop = self.center_crop
|
self.maybe_random_crop = self.center_crop
|
||||||
else:
|
else:
|
||||||
@@ -485,13 +491,16 @@ class DiffusionRgbEncoder(nn.Module):
|
|||||||
|
|
||||||
# Set up pooling and final layers.
|
# Set up pooling and final layers.
|
||||||
# Use a dry run to get the feature map shape.
|
# Use a dry run to get the feature map shape.
|
||||||
# The dummy input should take the number of image channels from `config.image_features` and it should
|
# The dummy shape mirrors the runtime preprocessing order: resize -> crop.
|
||||||
# use the height and width from `config.crop_shape` if it is provided, otherwise it should use the
|
|
||||||
# height and width from `config.image_features`.
|
|
||||||
|
|
||||||
# Note: we have a check in the config class to make sure all images have the same shape.
|
# Note: we have a check in the config class to make sure all images have the same shape.
|
||||||
images_shape = next(iter(config.image_features.values())).shape
|
images_shape = next(iter(config.image_features.values())).shape
|
||||||
dummy_shape_h_w = config.crop_shape if config.crop_shape is not None else images_shape[1:]
|
if config.crop_shape is not None:
|
||||||
|
dummy_shape_h_w = config.crop_shape
|
||||||
|
elif config.resize_shape is not None:
|
||||||
|
dummy_shape_h_w = config.resize_shape
|
||||||
|
else:
|
||||||
|
dummy_shape_h_w = images_shape[1:]
|
||||||
dummy_shape = (1, images_shape[0], *dummy_shape_h_w)
|
dummy_shape = (1, images_shape[0], *dummy_shape_h_w)
|
||||||
feature_map_shape = get_output_shape(self.backbone, dummy_shape)[1:]
|
feature_map_shape = get_output_shape(self.backbone, dummy_shape)[1:]
|
||||||
|
|
||||||
@@ -507,7 +516,10 @@ class DiffusionRgbEncoder(nn.Module):
|
|||||||
Returns:
|
Returns:
|
||||||
(B, D) image feature.
|
(B, D) image feature.
|
||||||
"""
|
"""
|
||||||
# Preprocess: maybe crop (if it was set up in the __init__).
|
# Preprocess: resize if configured, then crop if configured.
|
||||||
|
|
||||||
|
if self.resize is not None:
|
||||||
|
x = self.resize(x)
|
||||||
if self.do_crop:
|
if self.do_crop:
|
||||||
if self.training: # noqa: SIM108
|
if self.training: # noqa: SIM108
|
||||||
x = self.maybe_random_crop(x)
|
x = self.maybe_random_crop(x)
|
||||||
|
|||||||
@@ -106,6 +106,9 @@ class SmolVLAConfig(PreTrainedConfig):
|
|||||||
# Real-Time Chunking (RTC) configuration
|
# Real-Time Chunking (RTC) configuration
|
||||||
rtc_config: RTCConfig | None = None
|
rtc_config: RTCConfig | None = None
|
||||||
|
|
||||||
|
compile_model: bool = False # Whether to use torch.compile for model optimization
|
||||||
|
compile_mode: str = "max-autotune" # Torch compile mode
|
||||||
|
|
||||||
def __post_init__(self):
|
def __post_init__(self):
|
||||||
super().__post_init__()
|
super().__post_init__()
|
||||||
|
|
||||||
|
|||||||
@@ -593,6 +593,12 @@ class VLAFlowMatching(nn.Module):
|
|||||||
self.prefix_length = self.config.prefix_length
|
self.prefix_length = self.config.prefix_length
|
||||||
self.rtc_processor = rtc_processor
|
self.rtc_processor = rtc_processor
|
||||||
|
|
||||||
|
# Compile model if requested
|
||||||
|
if config.compile_model:
|
||||||
|
torch.set_float32_matmul_precision("high")
|
||||||
|
self.sample_actions = torch.compile(self.sample_actions, mode=config.compile_mode)
|
||||||
|
self.forward = torch.compile(self.forward, mode=config.compile_mode)
|
||||||
|
|
||||||
def _rtc_enabled(self):
|
def _rtc_enabled(self):
|
||||||
return self.config.rtc_config is not None and self.config.rtc_config.enabled
|
return self.config.rtc_config is not None and self.config.rtc_config.enabled
|
||||||
|
|
||||||
|
|||||||
@@ -77,7 +77,6 @@ class SmolVLMWithExpertModel(nn.Module):
|
|||||||
print(f"Loading {model_id} weights ...")
|
print(f"Loading {model_id} weights ...")
|
||||||
self.vlm = AutoModelForImageTextToText.from_pretrained(
|
self.vlm = AutoModelForImageTextToText.from_pretrained(
|
||||||
model_id,
|
model_id,
|
||||||
device_map=device,
|
|
||||||
torch_dtype="bfloat16",
|
torch_dtype="bfloat16",
|
||||||
low_cpu_mem_usage=True,
|
low_cpu_mem_usage=True,
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -56,6 +56,7 @@ from lerobot.teleoperators import ( # noqa: F401
|
|||||||
make_teleoperator_from_config,
|
make_teleoperator_from_config,
|
||||||
omx_leader,
|
omx_leader,
|
||||||
openarm_leader,
|
openarm_leader,
|
||||||
|
openarm_mini,
|
||||||
so_leader,
|
so_leader,
|
||||||
unitree_g1,
|
unitree_g1,
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -132,10 +132,13 @@ def visualize_dataset(
|
|||||||
|
|
||||||
logging.info("Logging to Rerun")
|
logging.info("Logging to Rerun")
|
||||||
|
|
||||||
|
first_index = None
|
||||||
for batch in tqdm.tqdm(dataloader, total=len(dataloader)):
|
for batch in tqdm.tqdm(dataloader, total=len(dataloader)):
|
||||||
|
if first_index is None:
|
||||||
|
first_index = batch["index"][0].item()
|
||||||
# iterate over the batch
|
# iterate over the batch
|
||||||
for i in range(len(batch["index"])):
|
for i in range(len(batch["index"])):
|
||||||
rr.set_time("frame_index", sequence=batch["frame_index"][i].item())
|
rr.set_time("frame_index", sequence=batch["index"][i].item() - first_index)
|
||||||
rr.set_time("timestamp", timestamp=batch["timestamp"][i].item())
|
rr.set_time("timestamp", timestamp=batch["timestamp"][i].item())
|
||||||
|
|
||||||
# display each camera image
|
# display each camera image
|
||||||
|
|||||||
@@ -21,6 +21,9 @@ This script allows you to delete episodes, split datasets, merge datasets,
|
|||||||
remove features, modify tasks, and convert image datasets to video format.
|
remove features, modify tasks, and convert image datasets to video format.
|
||||||
When new_repo_id is specified, creates a new dataset.
|
When new_repo_id is specified, creates a new dataset.
|
||||||
|
|
||||||
|
Path semantics (v2): --root and --new_root are exact dataset folders containing
|
||||||
|
meta/, data/, videos/. When omitted, defaults to $HF_LEROBOT_HOME/{repo_id}.
|
||||||
|
|
||||||
Usage Examples:
|
Usage Examples:
|
||||||
|
|
||||||
Delete episodes 0, 2, and 5 from a dataset:
|
Delete episodes 0, 2, and 5 from a dataset:
|
||||||
@@ -29,16 +32,31 @@ Delete episodes 0, 2, and 5 from a dataset:
|
|||||||
--operation.type delete_episodes \
|
--operation.type delete_episodes \
|
||||||
--operation.episode_indices "[0, 2, 5]"
|
--operation.episode_indices "[0, 2, 5]"
|
||||||
|
|
||||||
Delete episodes and save to a new dataset:
|
Delete episodes from a local dataset at a specific path:
|
||||||
lerobot-edit-dataset \
|
lerobot-edit-dataset \
|
||||||
--repo_id lerobot/pusht \
|
--repo_id lerobot/pusht \
|
||||||
--new_repo_id lerobot/pusht_filtered \
|
--root /path/to/pusht \
|
||||||
--operation.type delete_episodes \
|
--operation.type delete_episodes \
|
||||||
--operation.episode_indices "[0, 2, 5]"
|
--operation.episode_indices "[0, 2, 5]"
|
||||||
|
|
||||||
Split dataset by fractions:
|
Delete episodes and save to a new dataset at a specific path and with a new repo_id:
|
||||||
lerobot-edit-dataset \
|
lerobot-edit-dataset \
|
||||||
--repo_id lerobot/pusht \
|
--repo_id lerobot/pusht \
|
||||||
|
--new_repo_id lerobot/pusht_filtered \
|
||||||
|
--new_root /path/to/pusht_filtered \
|
||||||
|
--operation.type delete_episodes \
|
||||||
|
--operation.episode_indices "[0, 2, 5]"
|
||||||
|
|
||||||
|
Split dataset by fractions (pusht_train, pusht_val):
|
||||||
|
lerobot-edit-dataset \
|
||||||
|
--repo_id lerobot/pusht \
|
||||||
|
--operation.type split \
|
||||||
|
--operation.splits '{"train": 0.8, "val": 0.2}'
|
||||||
|
|
||||||
|
Split dataset by fractions and save split datasets to a specific folder (base_folder/train, base_folder/val):
|
||||||
|
lerobot-edit-dataset \
|
||||||
|
--repo_id lerobot/pusht \
|
||||||
|
--new_root /path/to/base_folder \
|
||||||
--operation.type split \
|
--operation.type split \
|
||||||
--operation.splits '{"train": 0.8, "val": 0.2}'
|
--operation.splits '{"train": 0.8, "val": 0.2}'
|
||||||
|
|
||||||
@@ -56,15 +74,29 @@ Split into more than two splits:
|
|||||||
|
|
||||||
Merge multiple datasets:
|
Merge multiple datasets:
|
||||||
lerobot-edit-dataset \
|
lerobot-edit-dataset \
|
||||||
--repo_id lerobot/pusht_merged \
|
--new_repo_id lerobot/pusht_merged \
|
||||||
--operation.type merge \
|
--operation.type merge \
|
||||||
--operation.repo_ids "['lerobot/pusht_train', 'lerobot/pusht_val']"
|
--operation.repo_ids "['lerobot/pusht_train', 'lerobot/pusht_val']"
|
||||||
|
|
||||||
|
Merge multiple datasets to a specific output path:
|
||||||
|
lerobot-edit-dataset \
|
||||||
|
--new_repo_id lerobot/pusht_merged \
|
||||||
|
--new_root /path/to/pusht_merged \
|
||||||
|
--operation.type merge \
|
||||||
|
--operation.repo_ids "['lerobot/pusht_train', 'lerobot/pusht_val']"
|
||||||
|
|
||||||
|
Merge multiple datasets from a list of local dataset paths:
|
||||||
|
lerobot-edit-dataset \
|
||||||
|
--new_repo_id lerobot/pusht_merged \
|
||||||
|
--operation.type merge \
|
||||||
|
--operation.repo_ids "['pusht_train', 'pusht_val']" \
|
||||||
|
--operation.roots "['/path/to/pusht_train', '/path/to/pusht_val']"
|
||||||
|
|
||||||
Remove camera feature:
|
Remove camera feature:
|
||||||
lerobot-edit-dataset \
|
lerobot-edit-dataset \
|
||||||
--repo_id lerobot/pusht \
|
--repo_id lerobot/pusht \
|
||||||
--operation.type remove_feature \
|
--operation.type remove_feature \
|
||||||
--operation.feature_names "['observation.images.top']"
|
--operation.feature_names "['observation.image']"
|
||||||
|
|
||||||
Modify tasks - set a single task for all episodes (WARNING: modifies in-place):
|
Modify tasks - set a single task for all episodes (WARNING: modifies in-place):
|
||||||
lerobot-edit-dataset \
|
lerobot-edit-dataset \
|
||||||
@@ -88,8 +120,8 @@ Modify tasks - set default task with overrides for specific episodes (WARNING: m
|
|||||||
Convert image dataset to video format and save locally:
|
Convert image dataset to video format and save locally:
|
||||||
lerobot-edit-dataset \
|
lerobot-edit-dataset \
|
||||||
--repo_id lerobot/pusht_image \
|
--repo_id lerobot/pusht_image \
|
||||||
--operation.type convert_image_to_video \
|
--new_root /path/to/output/pusht_video \
|
||||||
--operation.output_dir /path/to/output/pusht_video
|
--operation.type convert_image_to_video
|
||||||
|
|
||||||
Convert image dataset to video format and save with new repo_id:
|
Convert image dataset to video format and save with new repo_id:
|
||||||
lerobot-edit-dataset \
|
lerobot-edit-dataset \
|
||||||
@@ -104,28 +136,6 @@ Convert image dataset to video format and push to hub:
|
|||||||
--operation.type convert_image_to_video \
|
--operation.type convert_image_to_video \
|
||||||
--push_to_hub true
|
--push_to_hub true
|
||||||
|
|
||||||
Trim single episode to keep only frames within timestamp range:
|
|
||||||
python -m lerobot.scripts.lerobot_edit_dataset \
|
|
||||||
--repo_id lerobot/pusht \
|
|
||||||
--new_repo_id lerobot/pusht_trimmed \
|
|
||||||
--operation.type trim_episode \
|
|
||||||
--operation.episode_index 0 \
|
|
||||||
--operation.start_timestamp 10.0 \
|
|
||||||
--operation.end_timestamp 30.0
|
|
||||||
|
|
||||||
Trim multiple episodes at once (use null for no limit):
|
|
||||||
python -m lerobot.scripts.lerobot_edit_dataset \
|
|
||||||
--repo_id lerobot/pusht \
|
|
||||||
--operation.type trim_episode \
|
|
||||||
--operation.episode_trims '{"0": [10.0, 30.0], "2": [5.0, null], "3": [null, 20.0]}'
|
|
||||||
|
|
||||||
Trim and re-upload to same repo (overwrites original):
|
|
||||||
python -m lerobot.scripts.lerobot_edit_dataset \
|
|
||||||
--repo_id lerobot/pusht \
|
|
||||||
--operation.type trim_episode \
|
|
||||||
--operation.episode_index 0 \
|
|
||||||
--operation.start_timestamp 10.0 \
|
|
||||||
--push_to_hub true
|
|
||||||
Show dataset information:
|
Show dataset information:
|
||||||
lerobot-edit-dataset \
|
lerobot-edit-dataset \
|
||||||
--repo_id lerobot/pusht_image \
|
--repo_id lerobot/pusht_image \
|
||||||
@@ -189,6 +199,7 @@ class SplitConfig(OperationConfig):
|
|||||||
@dataclass
|
@dataclass
|
||||||
class MergeConfig(OperationConfig):
|
class MergeConfig(OperationConfig):
|
||||||
repo_ids: list[str] | None = None
|
repo_ids: list[str] | None = None
|
||||||
|
roots: list[str] | None = None
|
||||||
|
|
||||||
|
|
||||||
@OperationConfig.register_subclass("remove_feature")
|
@OperationConfig.register_subclass("remove_feature")
|
||||||
@@ -222,59 +233,46 @@ class ConvertImageToVideoConfig(OperationConfig):
|
|||||||
@OperationConfig.register_subclass("info")
|
@OperationConfig.register_subclass("info")
|
||||||
@dataclass
|
@dataclass
|
||||||
class InfoConfig(OperationConfig):
|
class InfoConfig(OperationConfig):
|
||||||
type: str = "info"
|
|
||||||
show_features: bool = False
|
show_features: bool = False
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
|
||||||
class TrimEpisodeConfig:
|
|
||||||
"""Trim episodes to keep only frames within timestamp ranges.
|
|
||||||
|
|
||||||
Supports multiple episodes via episode_trims dict:
|
|
||||||
--operation.episode_trims '{"0": [10.0, 30.0], "2": [5.0, 20.0]}'
|
|
||||||
|
|
||||||
Or single episode via legacy parameters:
|
|
||||||
--operation.episode_index 0 --operation.start_timestamp 10.0 --operation.end_timestamp 30.0
|
|
||||||
"""
|
|
||||||
type: str = "trim_episode"
|
|
||||||
# Multi-episode support: dict mapping episode_index -> [start_timestamp, end_timestamp]
|
|
||||||
# Use null for no limit, e.g. {"0": [10.0, null], "2": [null, 30.0]}
|
|
||||||
episode_trims: dict[str, list[float | None]] | None = None
|
|
||||||
# Legacy single-episode parameters (used if episode_trims is None)
|
|
||||||
episode_index: int | None = None
|
|
||||||
start_timestamp: float | None = None # Keep frames from this timestamp (inclusive)
|
|
||||||
end_timestamp: float | None = None # Keep frames until this timestamp (inclusive)
|
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
class EditDatasetConfig:
|
class EditDatasetConfig:
|
||||||
repo_id: str
|
# Operation configuration.
|
||||||
operation: (
|
|
||||||
DeleteEpisodesConfig | SplitConfig | MergeConfig | RemoveFeatureConfig | ConvertImageToVideoConfig | TrimEpisodeConfig
|
|
||||||
)
|
|
||||||
operation: OperationConfig
|
operation: OperationConfig
|
||||||
|
# Input dataset identifier. Always required unless for Merge operation.
|
||||||
|
repo_id: str | None = None
|
||||||
|
# Root directory where the input dataset is stored. If not specified, defaults to $HF_LEROBOT_HOME/repo_id.
|
||||||
root: str | None = None
|
root: str | None = None
|
||||||
|
# Edited dataset identifier. When both new_repo_id (resp. new_root) and repo_id (resp. root) are identical, modifications are applied in-place and a backup of the original dataset is created. Required for Merge operation.
|
||||||
new_repo_id: str | None = None
|
new_repo_id: str | None = None
|
||||||
|
# Root directory where the edited dataset will be stored. If not specified, defaults to $HF_LEROBOT_HOME/new_repo_id. For Split operation, this is the base directory for the split datasets.
|
||||||
|
new_root: str | None = None
|
||||||
|
# Upload dataset to Hugging Face hub.
|
||||||
push_to_hub: bool = False
|
push_to_hub: bool = False
|
||||||
|
|
||||||
|
|
||||||
def get_output_path(repo_id: str, new_repo_id: str | None, root: Path | None) -> tuple[str, Path]:
|
def get_output_path(
|
||||||
if new_repo_id:
|
repo_id: str,
|
||||||
output_repo_id = new_repo_id
|
new_repo_id: str | None,
|
||||||
output_dir = root / new_repo_id if root else HF_LEROBOT_HOME / new_repo_id
|
root: Path | str | None,
|
||||||
else:
|
new_root: Path | str | None,
|
||||||
output_repo_id = repo_id
|
) -> tuple[str, Path]:
|
||||||
dataset_path = root / repo_id if root else HF_LEROBOT_HOME / repo_id
|
input_path = Path(root) if root else HF_LEROBOT_HOME / repo_id
|
||||||
old_path = Path(str(dataset_path) + "_old")
|
|
||||||
|
|
||||||
if dataset_path.exists():
|
output_repo_id = new_repo_id if new_repo_id else repo_id
|
||||||
if old_path.exists():
|
output_path = Path(new_root) if new_root else HF_LEROBOT_HOME / output_repo_id
|
||||||
shutil.rmtree(old_path)
|
|
||||||
shutil.move(str(dataset_path), str(old_path))
|
|
||||||
|
|
||||||
output_dir = dataset_path
|
# In case of in-place modification, create a backup of the original dataset (if it exists)
|
||||||
|
if output_path == input_path:
|
||||||
|
backup_path = input_path.with_name(input_path.name + "_old")
|
||||||
|
|
||||||
return output_repo_id, output_dir
|
if input_path.exists():
|
||||||
|
if backup_path.exists():
|
||||||
|
shutil.rmtree(backup_path)
|
||||||
|
shutil.move(input_path, backup_path)
|
||||||
|
|
||||||
|
return output_repo_id, output_path
|
||||||
|
|
||||||
|
|
||||||
def handle_delete_episodes(cfg: EditDatasetConfig) -> None:
|
def handle_delete_episodes(cfg: EditDatasetConfig) -> None:
|
||||||
@@ -286,11 +284,15 @@ def handle_delete_episodes(cfg: EditDatasetConfig) -> None:
|
|||||||
|
|
||||||
dataset = LeRobotDataset(cfg.repo_id, root=cfg.root)
|
dataset = LeRobotDataset(cfg.repo_id, root=cfg.root)
|
||||||
output_repo_id, output_dir = get_output_path(
|
output_repo_id, output_dir = get_output_path(
|
||||||
cfg.repo_id, cfg.new_repo_id, Path(cfg.root) if cfg.root else None
|
cfg.repo_id,
|
||||||
|
new_repo_id=cfg.new_repo_id,
|
||||||
|
root=cfg.root,
|
||||||
|
new_root=cfg.new_root,
|
||||||
)
|
)
|
||||||
|
|
||||||
if cfg.new_repo_id is None:
|
# In case of in-place modification, make the dataset point to the backup directory
|
||||||
dataset.root = Path(str(dataset.root) + "_old")
|
if output_dir == dataset.root:
|
||||||
|
dataset.root = dataset.root.with_name(dataset.root.name + "_old")
|
||||||
|
|
||||||
logging.info(f"Deleting episodes {cfg.operation.episode_indices} from {cfg.repo_id}")
|
logging.info(f"Deleting episodes {cfg.operation.episode_indices} from {cfg.repo_id}")
|
||||||
new_dataset = delete_episodes(
|
new_dataset = delete_episodes(
|
||||||
@@ -317,19 +319,27 @@ def handle_split(cfg: EditDatasetConfig) -> None:
|
|||||||
"splits dict must be specified with split names as keys and fractions/episode lists as values"
|
"splits dict must be specified with split names as keys and fractions/episode lists as values"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
if cfg.new_repo_id is not None:
|
||||||
|
logging.warning(
|
||||||
|
"split uses the original dataset identifier --repo_id to generate split names. The --new_repo_id parameter is ignored."
|
||||||
|
)
|
||||||
|
|
||||||
dataset = LeRobotDataset(cfg.repo_id, root=cfg.root)
|
dataset = LeRobotDataset(cfg.repo_id, root=cfg.root)
|
||||||
|
|
||||||
logging.info(f"Splitting dataset {cfg.repo_id} with splits: {cfg.operation.splits}")
|
logging.info(f"Splitting dataset {cfg.repo_id} with splits: {cfg.operation.splits}")
|
||||||
split_datasets = split_dataset(dataset, splits=cfg.operation.splits)
|
split_datasets = split_dataset(
|
||||||
|
dataset,
|
||||||
|
splits=cfg.operation.splits,
|
||||||
|
output_dir=cfg.new_root,
|
||||||
|
)
|
||||||
|
|
||||||
for split_name, split_ds in split_datasets.items():
|
for split_name, split_ds in split_datasets.items():
|
||||||
split_repo_id = f"{cfg.repo_id}_{split_name}"
|
|
||||||
logging.info(
|
logging.info(
|
||||||
f"{split_name}: {split_ds.meta.total_episodes} episodes, {split_ds.meta.total_frames} frames"
|
f"{split_name}: {split_ds.meta.total_episodes} episodes, {split_ds.meta.total_frames} frames"
|
||||||
)
|
)
|
||||||
|
|
||||||
if cfg.push_to_hub:
|
if cfg.push_to_hub:
|
||||||
logging.info(f"Pushing {split_name} split to hub as {split_repo_id}")
|
logging.info(f"Pushing {split_name} split to hub as {split_ds.repo_id}")
|
||||||
LeRobotDataset(split_ds.repo_id, root=split_ds.root).push_to_hub()
|
LeRobotDataset(split_ds.repo_id, root=split_ds.root).push_to_hub()
|
||||||
|
|
||||||
|
|
||||||
@@ -340,18 +350,29 @@ def handle_merge(cfg: EditDatasetConfig) -> None:
|
|||||||
if not cfg.operation.repo_ids:
|
if not cfg.operation.repo_ids:
|
||||||
raise ValueError("repo_ids must be specified for merge operation")
|
raise ValueError("repo_ids must be specified for merge operation")
|
||||||
|
|
||||||
if not cfg.repo_id:
|
if cfg.repo_id is not None or cfg.root is not None:
|
||||||
raise ValueError("repo_id must be specified as the output repository for merged dataset")
|
logging.warning(
|
||||||
|
"merge uses --new_repo_id and --new_root for the merged dataset. The --repo_id and --root parameters are ignored."
|
||||||
|
)
|
||||||
|
|
||||||
|
if cfg.operation.roots:
|
||||||
|
if len(cfg.operation.roots) != len(cfg.operation.repo_ids):
|
||||||
|
raise ValueError("repo_ids and roots must have the same length for merge operation")
|
||||||
|
logging.info(f"Loading {len(cfg.operation.roots)} datasets to merge")
|
||||||
|
datasets = [
|
||||||
|
LeRobotDataset(repo_id=repo_id, root=root)
|
||||||
|
for repo_id, root in zip(cfg.operation.repo_ids, cfg.operation.roots, strict=True)
|
||||||
|
]
|
||||||
|
else:
|
||||||
logging.info(f"Loading {len(cfg.operation.repo_ids)} datasets to merge")
|
logging.info(f"Loading {len(cfg.operation.repo_ids)} datasets to merge")
|
||||||
datasets = [LeRobotDataset(repo_id, root=cfg.root) for repo_id in cfg.operation.repo_ids]
|
datasets = [LeRobotDataset(repo_id) for repo_id in cfg.operation.repo_ids]
|
||||||
|
|
||||||
output_dir = Path(cfg.root) / cfg.repo_id if cfg.root else HF_LEROBOT_HOME / cfg.repo_id
|
output_dir = Path(cfg.new_root) if cfg.new_root else HF_LEROBOT_HOME / cfg.new_repo_id
|
||||||
|
|
||||||
logging.info(f"Merging datasets into {cfg.repo_id}")
|
logging.info(f"Merging datasets into {cfg.new_repo_id}")
|
||||||
merged_dataset = merge_datasets(
|
merged_dataset = merge_datasets(
|
||||||
datasets,
|
datasets,
|
||||||
output_repo_id=cfg.repo_id,
|
output_repo_id=cfg.new_repo_id,
|
||||||
output_dir=output_dir,
|
output_dir=output_dir,
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -361,7 +382,7 @@ def handle_merge(cfg: EditDatasetConfig) -> None:
|
|||||||
)
|
)
|
||||||
|
|
||||||
if cfg.push_to_hub:
|
if cfg.push_to_hub:
|
||||||
logging.info(f"Pushing to hub as {cfg.repo_id}")
|
logging.info(f"Pushing to hub as {cfg.new_repo_id}")
|
||||||
LeRobotDataset(merged_dataset.repo_id, root=output_dir).push_to_hub()
|
LeRobotDataset(merged_dataset.repo_id, root=output_dir).push_to_hub()
|
||||||
|
|
||||||
|
|
||||||
@@ -374,11 +395,15 @@ def handle_remove_feature(cfg: EditDatasetConfig) -> None:
|
|||||||
|
|
||||||
dataset = LeRobotDataset(cfg.repo_id, root=cfg.root)
|
dataset = LeRobotDataset(cfg.repo_id, root=cfg.root)
|
||||||
output_repo_id, output_dir = get_output_path(
|
output_repo_id, output_dir = get_output_path(
|
||||||
cfg.repo_id, cfg.new_repo_id, Path(cfg.root) if cfg.root else None
|
cfg.repo_id,
|
||||||
|
new_repo_id=cfg.new_repo_id,
|
||||||
|
root=cfg.root,
|
||||||
|
new_root=cfg.new_root,
|
||||||
)
|
)
|
||||||
|
|
||||||
if cfg.new_repo_id is None:
|
# In case of in-place modification, make the dataset point to the backup directory
|
||||||
dataset.root = Path(str(dataset.root) + "_old")
|
if output_dir == dataset.root:
|
||||||
|
dataset.root = dataset.root.with_name(dataset.root.name + "_old")
|
||||||
|
|
||||||
logging.info(f"Removing features {cfg.operation.feature_names} from {cfg.repo_id}")
|
logging.info(f"Removing features {cfg.operation.feature_names} from {cfg.repo_id}")
|
||||||
new_dataset = remove_feature(
|
new_dataset = remove_feature(
|
||||||
@@ -396,92 +421,6 @@ def handle_remove_feature(cfg: EditDatasetConfig) -> None:
|
|||||||
LeRobotDataset(output_repo_id, root=output_dir).push_to_hub()
|
LeRobotDataset(output_repo_id, root=output_dir).push_to_hub()
|
||||||
|
|
||||||
|
|
||||||
def handle_trim_episode(cfg: EditDatasetConfig) -> None:
|
|
||||||
"""Trim episodes to keep only frames within timestamp ranges."""
|
|
||||||
if not isinstance(cfg.operation, TrimEpisodeConfig):
|
|
||||||
raise ValueError("Operation config must be TrimEpisodeConfig")
|
|
||||||
|
|
||||||
# Parse episode trims - support both multi-episode dict and legacy single episode
|
|
||||||
episode_trims: dict[int, tuple[float | None, float | None]] = {}
|
|
||||||
|
|
||||||
if cfg.operation.episode_trims is not None:
|
|
||||||
# Multi-episode mode
|
|
||||||
for ep_str, ts_range in cfg.operation.episode_trims.items():
|
|
||||||
ep_idx = int(ep_str)
|
|
||||||
start_ts = ts_range[0] if len(ts_range) > 0 else None
|
|
||||||
end_ts = ts_range[1] if len(ts_range) > 1 else None
|
|
||||||
episode_trims[ep_idx] = (start_ts, end_ts)
|
|
||||||
elif cfg.operation.episode_index is not None:
|
|
||||||
# Legacy single-episode mode
|
|
||||||
if cfg.operation.start_timestamp is None and cfg.operation.end_timestamp is None:
|
|
||||||
raise ValueError("At least one of start_timestamp or end_timestamp must be specified")
|
|
||||||
episode_trims[cfg.operation.episode_index] = (
|
|
||||||
cfg.operation.start_timestamp,
|
|
||||||
cfg.operation.end_timestamp,
|
|
||||||
)
|
|
||||||
else:
|
|
||||||
raise ValueError("Either episode_trims or episode_index must be specified")
|
|
||||||
|
|
||||||
dataset = LeRobotDataset(cfg.repo_id, root=cfg.root)
|
|
||||||
output_repo_id, output_dir = get_output_path(
|
|
||||||
cfg.repo_id, cfg.new_repo_id, Path(cfg.root) if cfg.root else None
|
|
||||||
)
|
|
||||||
|
|
||||||
if cfg.new_repo_id is None:
|
|
||||||
dataset.root = Path(str(dataset.root) + "_old")
|
|
||||||
|
|
||||||
logging.info(f"Trimming {len(episode_trims)} episode(s) from {cfg.repo_id}")
|
|
||||||
|
|
||||||
# Get episode boundaries and find frames to keep for each episode
|
|
||||||
episodes_info = dataset.meta.episodes
|
|
||||||
all_frames_to_keep: dict[int, list[int]] = {}
|
|
||||||
|
|
||||||
for ep_idx, (start_ts, end_ts) in episode_trims.items():
|
|
||||||
if ep_idx >= len(episodes_info["episode_index"]):
|
|
||||||
raise ValueError(f"Episode {ep_idx} does not exist (dataset has {len(episodes_info['episode_index'])} episodes)")
|
|
||||||
|
|
||||||
from_frame = episodes_info["dataset_from_index"][ep_idx]
|
|
||||||
to_frame = episodes_info["dataset_to_index"][ep_idx]
|
|
||||||
|
|
||||||
logging.info(f"Episode {ep_idx}: trimming to [{start_ts}, {end_ts}]")
|
|
||||||
logging.info(f" Original frames: {from_frame} to {to_frame} ({to_frame - from_frame} frames)")
|
|
||||||
|
|
||||||
# Find frames within timestamp range
|
|
||||||
frames_to_keep = []
|
|
||||||
for frame_idx in range(from_frame, to_frame):
|
|
||||||
frame = dataset.hf_dataset[frame_idx]
|
|
||||||
ts = frame["timestamp"]
|
|
||||||
|
|
||||||
in_range = True
|
|
||||||
if start_ts is not None and ts < start_ts:
|
|
||||||
in_range = False
|
|
||||||
if end_ts is not None and ts > end_ts:
|
|
||||||
in_range = False
|
|
||||||
|
|
||||||
if in_range:
|
|
||||||
frames_to_keep.append(frame_idx)
|
|
||||||
|
|
||||||
if not frames_to_keep:
|
|
||||||
raise ValueError(f"Episode {ep_idx}: No frames found in timestamp range [{start_ts}, {end_ts}]")
|
|
||||||
|
|
||||||
logging.info(f" Keeping {len(frames_to_keep)} frames (indices {frames_to_keep[0]} to {frames_to_keep[-1]})")
|
|
||||||
all_frames_to_keep[ep_idx] = frames_to_keep
|
|
||||||
|
|
||||||
from lerobot.datasets.dataset_tools import trim_episodes_by_frames
|
|
||||||
|
|
||||||
new_dataset = trim_episodes_by_frames(
|
|
||||||
dataset,
|
|
||||||
episode_frames_to_keep=all_frames_to_keep,
|
|
||||||
output_dir=output_dir,
|
|
||||||
repo_id=output_repo_id,
|
|
||||||
)
|
|
||||||
|
|
||||||
logging.info(f"Dataset saved to {output_dir}")
|
|
||||||
logging.info(f"Episodes: {new_dataset.meta.total_episodes}, Frames: {new_dataset.meta.total_frames}")
|
|
||||||
|
|
||||||
if cfg.push_to_hub:
|
|
||||||
logging.info(f"Pushing to hub as {output_repo_id}")
|
|
||||||
LeRobotDataset(output_repo_id, root=output_dir).push_to_hub()
|
|
||||||
def handle_modify_tasks(cfg: EditDatasetConfig) -> None:
|
def handle_modify_tasks(cfg: EditDatasetConfig) -> None:
|
||||||
if not isinstance(cfg.operation, ModifyTasksConfig):
|
if not isinstance(cfg.operation, ModifyTasksConfig):
|
||||||
raise ValueError("Operation config must be ModifyTasksConfig")
|
raise ValueError("Operation config must be ModifyTasksConfig")
|
||||||
@@ -492,9 +431,10 @@ def handle_modify_tasks(cfg: EditDatasetConfig) -> None:
|
|||||||
if new_task is None and episode_tasks_raw is None:
|
if new_task is None and episode_tasks_raw is None:
|
||||||
raise ValueError("Must specify at least one of new_task or episode_tasks for modify_tasks operation")
|
raise ValueError("Must specify at least one of new_task or episode_tasks for modify_tasks operation")
|
||||||
|
|
||||||
# Warn about in-place modification behavior
|
if cfg.new_repo_id is not None or cfg.new_root is not None:
|
||||||
if cfg.new_repo_id is not None:
|
logging.warning(
|
||||||
logging.warning("modify_tasks modifies datasets in-place. The --new_repo_id parameter is ignored.")
|
"modify_tasks modifies datasets in-place. The --new_repo_id and --new_root parameters are ignored."
|
||||||
|
)
|
||||||
|
|
||||||
dataset = LeRobotDataset(cfg.repo_id, root=cfg.root)
|
dataset = LeRobotDataset(cfg.repo_id, root=cfg.root)
|
||||||
logging.warning(f"Modifying dataset in-place at {dataset.root}. Original data will be overwritten.")
|
logging.warning(f"Modifying dataset in-place at {dataset.root}. Original data will be overwritten.")
|
||||||
@@ -530,32 +470,30 @@ def handle_convert_image_to_video(cfg: EditDatasetConfig) -> None:
|
|||||||
dataset = LeRobotDataset(cfg.repo_id, root=cfg.root)
|
dataset = LeRobotDataset(cfg.repo_id, root=cfg.root)
|
||||||
|
|
||||||
# Determine output directory and repo_id
|
# Determine output directory and repo_id
|
||||||
# Priority: 1) new_repo_id, 2) operation.output_dir, 3) auto-generated name
|
# Priority: 1) new_root, 2) new_repo_id, 3) operation.output_dir, 4) auto-generated name
|
||||||
output_dir_config = getattr(cfg.operation, "output_dir", None)
|
output_dir_config = getattr(cfg.operation, "output_dir", None)
|
||||||
|
if output_dir_config:
|
||||||
|
logging.warning(
|
||||||
|
"--operation.output_dir is deprecated and will be removed in future versions. "
|
||||||
|
"Please use --new_root instead."
|
||||||
|
)
|
||||||
|
|
||||||
if cfg.new_repo_id:
|
if cfg.new_root:
|
||||||
# Use new_repo_id for both local storage and hub push
|
output_dir = Path(cfg.new_root)
|
||||||
|
output_repo_id = cfg.new_repo_id or f"{cfg.repo_id}_video"
|
||||||
|
logging.info(f"Saving to new_root: {output_dir} as {output_repo_id}")
|
||||||
|
elif cfg.new_repo_id:
|
||||||
output_repo_id = cfg.new_repo_id
|
output_repo_id = cfg.new_repo_id
|
||||||
# Place new dataset as a sibling to the original dataset
|
output_dir = HF_LEROBOT_HOME / cfg.new_repo_id
|
||||||
# Get the parent of the actual dataset root (not cfg.root which might be the lerobot cache dir)
|
|
||||||
# Extract just the dataset name (after last slash) for the local directory
|
|
||||||
local_dir_name = cfg.new_repo_id.split("/")[-1]
|
|
||||||
output_dir = dataset.root.parent / local_dir_name
|
|
||||||
logging.info(f"Saving to new dataset: {cfg.new_repo_id} at {output_dir}")
|
logging.info(f"Saving to new dataset: {cfg.new_repo_id} at {output_dir}")
|
||||||
elif output_dir_config:
|
elif output_dir_config:
|
||||||
# Use custom output directory for local-only storage
|
|
||||||
output_dir = Path(output_dir_config)
|
output_dir = Path(output_dir_config)
|
||||||
# Extract repo name from output_dir for the dataset
|
|
||||||
output_repo_id = output_dir.name
|
output_repo_id = output_dir.name
|
||||||
logging.info(f"Saving to local directory: {output_dir}")
|
logging.info(f"Saving to local directory: {output_dir} as {output_repo_id}")
|
||||||
else:
|
else:
|
||||||
# Auto-generate name: append "_video" to original repo_id
|
|
||||||
output_repo_id = f"{cfg.repo_id}_video"
|
output_repo_id = f"{cfg.repo_id}_video"
|
||||||
# Place new dataset as a sibling to the original dataset
|
output_dir = HF_LEROBOT_HOME / output_repo_id
|
||||||
# Extract just the dataset name (after last slash) for the local directory
|
logging.info(f"Saving to auto-generated location: {output_dir} as {output_repo_id}")
|
||||||
local_dir_name = output_repo_id.split("/")[-1]
|
|
||||||
output_dir = dataset.root.parent / local_dir_name
|
|
||||||
logging.info(f"Saving to auto-generated location: {output_dir}")
|
|
||||||
|
|
||||||
logging.info(f"Converting dataset {cfg.repo_id} to video format")
|
logging.info(f"Converting dataset {cfg.repo_id} to video format")
|
||||||
|
|
||||||
@@ -630,8 +568,20 @@ def handle_info(cfg: EditDatasetConfig):
|
|||||||
sys.stdout.write(f"{feature_dump_str}\n")
|
sys.stdout.write(f"{feature_dump_str}\n")
|
||||||
|
|
||||||
|
|
||||||
|
def _validate_config(cfg: EditDatasetConfig) -> None:
|
||||||
|
if isinstance(cfg.operation, MergeConfig):
|
||||||
|
if not cfg.new_repo_id:
|
||||||
|
raise ValueError("--new_repo_id is required for merge operation (the merged dataset identifier)")
|
||||||
|
else:
|
||||||
|
if not cfg.repo_id:
|
||||||
|
raise ValueError(
|
||||||
|
f"--repo_id is required for {cfg.operation.type} operation (the input dataset identifier)"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
@parser.wrap()
|
@parser.wrap()
|
||||||
def edit_dataset(cfg: EditDatasetConfig) -> None:
|
def edit_dataset(cfg: EditDatasetConfig) -> None:
|
||||||
|
_validate_config(cfg)
|
||||||
operation_type = cfg.operation.type
|
operation_type = cfg.operation.type
|
||||||
|
|
||||||
if operation_type == "delete_episodes":
|
if operation_type == "delete_episodes":
|
||||||
@@ -646,8 +596,6 @@ def edit_dataset(cfg: EditDatasetConfig) -> None:
|
|||||||
handle_modify_tasks(cfg)
|
handle_modify_tasks(cfg)
|
||||||
elif operation_type == "convert_image_to_video":
|
elif operation_type == "convert_image_to_video":
|
||||||
handle_convert_image_to_video(cfg)
|
handle_convert_image_to_video(cfg)
|
||||||
elif operation_type == "trim_episode":
|
|
||||||
handle_trim_episode(cfg)
|
|
||||||
elif operation_type == "info":
|
elif operation_type == "info":
|
||||||
handle_info(cfg)
|
handle_info(cfg)
|
||||||
else:
|
else:
|
||||||
|
|||||||
@@ -61,6 +61,7 @@ from lerobot.teleoperators import ( # noqa: F401
|
|||||||
make_teleoperator_from_config,
|
make_teleoperator_from_config,
|
||||||
omx_leader,
|
omx_leader,
|
||||||
openarm_leader,
|
openarm_leader,
|
||||||
|
openarm_mini,
|
||||||
so_leader,
|
so_leader,
|
||||||
)
|
)
|
||||||
from lerobot.utils.robot_utils import precise_sleep
|
from lerobot.utils.robot_utils import precise_sleep
|
||||||
|
|||||||
@@ -125,6 +125,7 @@ from lerobot.teleoperators import ( # noqa: F401
|
|||||||
make_teleoperator_from_config,
|
make_teleoperator_from_config,
|
||||||
omx_leader,
|
omx_leader,
|
||||||
openarm_leader,
|
openarm_leader,
|
||||||
|
openarm_mini,
|
||||||
reachy2_teleoperator,
|
reachy2_teleoperator,
|
||||||
so_leader,
|
so_leader,
|
||||||
unitree_g1,
|
unitree_g1,
|
||||||
@@ -154,7 +155,7 @@ class DatasetRecordConfig:
|
|||||||
repo_id: str
|
repo_id: str
|
||||||
# A short but accurate description of the task performed during the recording (e.g. "Pick the Lego block and drop it in the box on the right.")
|
# A short but accurate description of the task performed during the recording (e.g. "Pick the Lego block and drop it in the box on the right.")
|
||||||
single_task: str
|
single_task: str
|
||||||
# Root directory where the dataset will be stored (e.g. 'dataset/path').
|
# Root directory where the dataset will be stored (e.g. 'dataset/path'). If None, defaults to $HF_LEROBOT_HOME/repo_id.
|
||||||
root: str | Path | None = None
|
root: str | Path | None = None
|
||||||
# Limit the frames per second.
|
# Limit the frames per second.
|
||||||
fps: int = 30
|
fps: int = 30
|
||||||
@@ -333,6 +334,7 @@ def record_loop(
|
|||||||
preprocessor.reset()
|
preprocessor.reset()
|
||||||
postprocessor.reset()
|
postprocessor.reset()
|
||||||
|
|
||||||
|
no_action_count = 0
|
||||||
timestamp = 0
|
timestamp = 0
|
||||||
start_episode_t = time.perf_counter()
|
start_episode_t = time.perf_counter()
|
||||||
while timestamp < control_time_s:
|
while timestamp < control_time_s:
|
||||||
@@ -380,9 +382,11 @@ def record_loop(
|
|||||||
act = {**arm_action, **base_action} if len(base_action) > 0 else arm_action
|
act = {**arm_action, **base_action} if len(base_action) > 0 else arm_action
|
||||||
act_processed_teleop = teleop_action_processor((act, obs))
|
act_processed_teleop = teleop_action_processor((act, obs))
|
||||||
else:
|
else:
|
||||||
logging.info(
|
no_action_count += 1
|
||||||
"No policy or teleoperator provided, skipping action generation."
|
if no_action_count == 1 or no_action_count % 10 == 0:
|
||||||
"This is likely to happen when resetting the environment without a teleop device."
|
logging.warning(
|
||||||
|
"No policy or teleoperator provided, skipping action generation. "
|
||||||
|
"This is likely to happen when resetting the environment without a teleop device. "
|
||||||
"The robot won't be at its rest position at the start of the next episode."
|
"The robot won't be at its rest position at the start of the next episode."
|
||||||
)
|
)
|
||||||
continue
|
continue
|
||||||
|
|||||||
@@ -80,7 +80,7 @@ class DatasetReplayConfig:
|
|||||||
repo_id: str
|
repo_id: str
|
||||||
# Episode to replay.
|
# Episode to replay.
|
||||||
episode: int
|
episode: int
|
||||||
# Root directory where the dataset will be stored (e.g. 'dataset/path').
|
# Root directory where the dataset will be stored (e.g. 'dataset/path'). If None, defaults to $HF_LEROBOT_HOME/repo_id.
|
||||||
root: str | Path | None = None
|
root: str | Path | None = None
|
||||||
# Limit the frames per second. By default, uses the policy fps.
|
# Limit the frames per second. By default, uses the policy fps.
|
||||||
fps: int = 30
|
fps: int = 30
|
||||||
|
|||||||
@@ -94,6 +94,7 @@ from lerobot.teleoperators import ( # noqa: F401
|
|||||||
make_teleoperator_from_config,
|
make_teleoperator_from_config,
|
||||||
omx_leader,
|
omx_leader,
|
||||||
openarm_leader,
|
openarm_leader,
|
||||||
|
openarm_mini,
|
||||||
reachy2_teleoperator,
|
reachy2_teleoperator,
|
||||||
so_leader,
|
so_leader,
|
||||||
unitree_g1,
|
unitree_g1,
|
||||||
|
|||||||
@@ -380,10 +380,10 @@ def train(cfg: TrainPipelineConfig, accelerator: Accelerator | None = None):
|
|||||||
"dataloading_s": AverageMeter("data_s", ":.3f"),
|
"dataloading_s": AverageMeter("data_s", ":.3f"),
|
||||||
}
|
}
|
||||||
|
|
||||||
# Use effective batch size for proper epoch calculation in distributed training
|
# Keep global batch size for logging; MetricsTracker handles world size internally.
|
||||||
effective_batch_size = cfg.batch_size * accelerator.num_processes
|
effective_batch_size = cfg.batch_size * accelerator.num_processes
|
||||||
train_tracker = MetricsTracker(
|
train_tracker = MetricsTracker(
|
||||||
effective_batch_size,
|
cfg.batch_size,
|
||||||
dataset.num_frames,
|
dataset.num_frames,
|
||||||
dataset.num_episodes,
|
dataset.num_episodes,
|
||||||
train_metrics,
|
train_metrics,
|
||||||
|
|||||||
@@ -104,9 +104,10 @@ class MetricsTracker:
|
|||||||
self.metrics = metrics
|
self.metrics = metrics
|
||||||
|
|
||||||
self.steps = initial_step
|
self.steps = initial_step
|
||||||
|
world_size = accelerator.num_processes if accelerator else 1
|
||||||
# A sample is an (observation,action) pair, where observation and action
|
# A sample is an (observation,action) pair, where observation and action
|
||||||
# can be on multiple timestamps. In a batch, we have `batch_size` number of samples.
|
# can be on multiple timestamps. In a batch, we have `batch_size` number of samples.
|
||||||
self.samples = self.steps * self._batch_size
|
self.samples = self.steps * self._batch_size * world_size
|
||||||
self.episodes = self.samples / self._avg_samples_per_ep
|
self.episodes = self.samples / self._avg_samples_per_ep
|
||||||
self.epochs = self.samples / self._num_frames
|
self.epochs = self.samples / self._num_frames
|
||||||
self.accelerator = accelerator
|
self.accelerator = accelerator
|
||||||
@@ -132,7 +133,8 @@ class MetricsTracker:
|
|||||||
Updates metrics that depend on 'step' for one step.
|
Updates metrics that depend on 'step' for one step.
|
||||||
"""
|
"""
|
||||||
self.steps += 1
|
self.steps += 1
|
||||||
self.samples += self._batch_size * (self.accelerator.num_processes if self.accelerator else 1)
|
world_size = self.accelerator.num_processes if self.accelerator else 1
|
||||||
|
self.samples += self._batch_size * world_size
|
||||||
self.episodes = self.samples / self._avg_samples_per_ep
|
self.episodes = self.samples / self._avg_samples_per_ep
|
||||||
self.epochs = self.samples / self._num_frames
|
self.epochs = self.samples / self._num_frames
|
||||||
|
|
||||||
|
|||||||
@@ -1,3 +1,3 @@
|
|||||||
version https://git-lfs.github.com/spec/v1
|
version https://git-lfs.github.com/spec/v1
|
||||||
oid sha256:19eaaa85f66ba4aa6388dbb83819ffad6ea4363247208f871a8dc385689f6fc8
|
oid sha256:54aecbc1af72a4cd5e9261492f5e7601890517516257aacdf2a0ffb3ce281f1b
|
||||||
size 992
|
size 992
|
||||||
|
|||||||
@@ -1,3 +1,3 @@
|
|||||||
version https://git-lfs.github.com/spec/v1
|
version https://git-lfs.github.com/spec/v1
|
||||||
oid sha256:227296eaeeb54acdc3dae2eb8af3d4d08fb87e245337624447140b1e91cfd002
|
oid sha256:88a9c3775a2aa1e90a08850521970070a4fcf0f6b82aab43cd8ccc5cf77e0013
|
||||||
size 47424
|
size 47424
|
||||||
|
|||||||
@@ -1,3 +1,3 @@
|
|||||||
version https://git-lfs.github.com/spec/v1
|
version https://git-lfs.github.com/spec/v1
|
||||||
oid sha256:271b00cb2f0cd5fd26b1d53463638e3d1a6e92692ec625fcffb420ca190869e5
|
oid sha256:91a2635e05a75fe187a5081504c5f35ce3417378813fa2deaf9ca4e8200e1819
|
||||||
size 68
|
size 68
|
||||||
|
|||||||
@@ -1,3 +1,3 @@
|
|||||||
version https://git-lfs.github.com/spec/v1
|
version https://git-lfs.github.com/spec/v1
|
||||||
oid sha256:778fddbbaa64248cee35cb377c02cc2b6076f7ce5855146de677128900617ddf
|
oid sha256:645bff922ac7bea63ad018ebf77c303c0e4cd2c1c0dc5ef3192865281bef3dc6
|
||||||
size 47424
|
size 47424
|
||||||
|
|||||||
Vendored
+1
-1
@@ -222,7 +222,7 @@ def tasks_factory():
|
|||||||
def _create_tasks(total_tasks: int = 3) -> pd.DataFrame:
|
def _create_tasks(total_tasks: int = 3) -> pd.DataFrame:
|
||||||
ids = list(range(total_tasks))
|
ids = list(range(total_tasks))
|
||||||
tasks = [f"Perform action {i}." for i in ids]
|
tasks = [f"Perform action {i}." for i in ids]
|
||||||
df = pd.DataFrame({"task_index": ids}, index=tasks)
|
df = pd.DataFrame({"task_index": ids}, index=pd.Index(tasks, name="task"))
|
||||||
return df
|
return df
|
||||||
|
|
||||||
return _create_tasks
|
return _create_tasks
|
||||||
|
|||||||
@@ -27,6 +27,7 @@ from lerobot.scripts.lerobot_edit_dataset import (
|
|||||||
OperationConfig,
|
OperationConfig,
|
||||||
RemoveFeatureConfig,
|
RemoveFeatureConfig,
|
||||||
SplitConfig,
|
SplitConfig,
|
||||||
|
_validate_config,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
@@ -51,11 +52,23 @@ class TestOperationTypeParsing:
|
|||||||
],
|
],
|
||||||
)
|
)
|
||||||
def test_operation_type_resolves_correct_class(self, type_name, expected_cls):
|
def test_operation_type_resolves_correct_class(self, type_name, expected_cls):
|
||||||
cfg = parse_cfg(["--repo_id", "test/repo", "--operation.type", type_name])
|
cfg = parse_cfg(
|
||||||
|
["--repo_id", "test/repo", "--new_repo_id", "test/merged", "--operation.type", type_name]
|
||||||
|
)
|
||||||
assert isinstance(cfg.operation, expected_cls), (
|
assert isinstance(cfg.operation, expected_cls), (
|
||||||
f"Expected {expected_cls.__name__}, got {type(cfg.operation).__name__}"
|
f"Expected {expected_cls.__name__}, got {type(cfg.operation).__name__}"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
def test_merge_requires_new_repo_id(self):
|
||||||
|
cfg = parse_cfg(["--operation.type", "merge"])
|
||||||
|
with pytest.raises(ValueError, match="--new_repo_id is required for merge"):
|
||||||
|
_validate_config(cfg)
|
||||||
|
|
||||||
|
def test_non_merge_requires_repo_id(self):
|
||||||
|
cfg = parse_cfg(["--operation.type", "delete_episodes"])
|
||||||
|
with pytest.raises(ValueError, match="--repo_id is required for delete_episodes"):
|
||||||
|
_validate_config(cfg)
|
||||||
|
|
||||||
@pytest.mark.parametrize(
|
@pytest.mark.parametrize(
|
||||||
"type_name, expected_cls",
|
"type_name, expected_cls",
|
||||||
[
|
[
|
||||||
@@ -69,6 +82,8 @@ class TestOperationTypeParsing:
|
|||||||
],
|
],
|
||||||
)
|
)
|
||||||
def test_get_choice_name_roundtrips(self, type_name, expected_cls):
|
def test_get_choice_name_roundtrips(self, type_name, expected_cls):
|
||||||
cfg = parse_cfg(["--repo_id", "test/repo", "--operation.type", type_name])
|
cfg = parse_cfg(
|
||||||
|
["--repo_id", "test/repo", "--new_repo_id", "test/merged", "--operation.type", type_name]
|
||||||
|
)
|
||||||
resolved_name = OperationConfig.get_choice_name(type(cfg.operation))
|
resolved_name = OperationConfig.get_choice_name(type(cfg.operation))
|
||||||
assert resolved_name == type_name
|
assert resolved_name == type_name
|
||||||
|
|||||||
@@ -24,6 +24,11 @@ def mock_metrics():
|
|||||||
return {"loss": AverageMeter("loss", ":.3f"), "accuracy": AverageMeter("accuracy", ":.2f")}
|
return {"loss": AverageMeter("loss", ":.3f"), "accuracy": AverageMeter("accuracy", ":.2f")}
|
||||||
|
|
||||||
|
|
||||||
|
class MockAccelerator:
|
||||||
|
def __init__(self, num_processes: int):
|
||||||
|
self.num_processes = num_processes
|
||||||
|
|
||||||
|
|
||||||
def test_average_meter_initialization():
|
def test_average_meter_initialization():
|
||||||
meter = AverageMeter("loss", ":.2f")
|
meter = AverageMeter("loss", ":.2f")
|
||||||
assert meter.name == "loss"
|
assert meter.name == "loss"
|
||||||
@@ -82,6 +87,37 @@ def test_metrics_tracker_step(mock_metrics):
|
|||||||
assert tracker.epochs == tracker.samples / 1000
|
assert tracker.epochs == tracker.samples / 1000
|
||||||
|
|
||||||
|
|
||||||
|
def test_metrics_tracker_initialization_with_accelerator(mock_metrics):
|
||||||
|
tracker = MetricsTracker(
|
||||||
|
batch_size=32,
|
||||||
|
num_frames=1000,
|
||||||
|
num_episodes=50,
|
||||||
|
metrics=mock_metrics,
|
||||||
|
initial_step=10,
|
||||||
|
accelerator=MockAccelerator(num_processes=2),
|
||||||
|
)
|
||||||
|
assert tracker.steps == 10
|
||||||
|
assert tracker.samples == 10 * 32 * 2
|
||||||
|
assert tracker.episodes == tracker.samples / (1000 / 50)
|
||||||
|
assert tracker.epochs == tracker.samples / 1000
|
||||||
|
|
||||||
|
|
||||||
|
def test_metrics_tracker_step_with_accelerator(mock_metrics):
|
||||||
|
tracker = MetricsTracker(
|
||||||
|
batch_size=32,
|
||||||
|
num_frames=1000,
|
||||||
|
num_episodes=50,
|
||||||
|
metrics=mock_metrics,
|
||||||
|
initial_step=5,
|
||||||
|
accelerator=MockAccelerator(num_processes=2),
|
||||||
|
)
|
||||||
|
tracker.step()
|
||||||
|
assert tracker.steps == 6
|
||||||
|
assert tracker.samples == (5 * 32 * 2) + (32 * 2)
|
||||||
|
assert tracker.episodes == tracker.samples / (1000 / 50)
|
||||||
|
assert tracker.epochs == tracker.samples / 1000
|
||||||
|
|
||||||
|
|
||||||
def test_metrics_tracker_getattr(mock_metrics):
|
def test_metrics_tracker_getattr(mock_metrics):
|
||||||
tracker = MetricsTracker(batch_size=32, num_frames=1000, num_episodes=50, metrics=mock_metrics)
|
tracker = MetricsTracker(batch_size=32, num_frames=1000, num_episodes=50, metrics=mock_metrics)
|
||||||
assert tracker.loss == mock_metrics["loss"]
|
assert tracker.loss == mock_metrics["loss"]
|
||||||
|
|||||||
Reference in New Issue
Block a user