diff --git a/.github/workflows/benchmark_tests.yml b/.github/workflows/benchmark_tests.yml index b82c59a8b..3493e5048 100644 --- a/.github/workflows/benchmark_tests.yml +++ b/.github/workflows/benchmark_tests.yml @@ -167,9 +167,9 @@ jobs: # ── LIBERO TRAIN+EVAL SMOKE ────────────────────────────────────────────── # Train SmolVLA for 1 step (batch_size=1, dataset episode 0 only) then - # immediately runs eval inside the training loop (eval_freq=1, 1 episode). + # immediately runs eval inside the training loop (env_eval_freq=1, 1 episode). # Tests the full train→eval-within-training pipeline end-to-end. - - name: Run Libero train+eval smoke (1 step, eval_freq=1) + - name: Run Libero train+eval smoke (1 step, env_eval_freq=1) if: env.HF_USER_TOKEN != '' run: | docker run --name libero-train-smoke --gpus all \ @@ -196,7 +196,7 @@ jobs: --output_dir=/tmp/train-smoke \ --steps=1 \ --batch_size=1 \ - --eval_freq=1 \ + --env_eval_freq=1 \ --eval.n_episodes=1 \ --eval.batch_size=1 \ --eval.use_async_envs=false \ diff --git a/AGENT_GUIDE.md b/AGENT_GUIDE.md index 57a33fdba..03b270dce 100644 --- a/AGENT_GUIDE.md +++ b/AGENT_GUIDE.md @@ -138,7 +138,7 @@ lerobot-replay --robot.type=so101_follower --robot.port= --robot. --dataset.repo_id=${HF_USER}/my_task --dataset.episode=0 ``` -**4.9 Train** (default: ACT — fastest, lowest memory). Apple silicon: `--policy.device=mps`. See §6/§7 for policy and duration. +**4.9 Train** (default: ACT — fastest, lowest memory). Apple silicon: `--policy.device=mps`. No local GPU? Add `--job.target=` (e.g. `a10g-small`, list them with `hf jobs hardware`) to run on Hugging Face Jobs instead. See §6/§7 for policy and duration. ```bash lerobot-train \ diff --git a/Makefile b/Makefile index d3987101f..ea3b6e261 100644 --- a/Makefile +++ b/Makefile @@ -58,7 +58,7 @@ test-act-ete-train: --dataset.episodes="[0]" \ --batch_size=2 \ --steps=4 \ - --eval_freq=2 \ + --env_eval_freq=2 \ --eval.n_episodes=1 \ --eval.batch_size=1 \ --save_freq=2 \ @@ -96,7 +96,7 @@ test-diffusion-ete-train: --dataset.episodes="[0]" \ --batch_size=2 \ --steps=2 \ - --eval_freq=2 \ + --env_eval_freq=2 \ --eval.n_episodes=1 \ --eval.batch_size=1 \ --save_checkpoint=true \ @@ -126,7 +126,7 @@ test-tdmpc-ete-train: --dataset.episodes="[0]" \ --batch_size=2 \ --steps=2 \ - --eval_freq=2 \ + --env_eval_freq=2 \ --eval.n_episodes=1 \ --eval.batch_size=1 \ --save_checkpoint=true \ @@ -161,7 +161,7 @@ test-smolvla-ete-train: --dataset.episodes="[0]" \ --batch_size=2 \ --steps=4 \ - --eval_freq=2 \ + --env_eval_freq=2 \ --eval.n_episodes=1 \ --eval.batch_size=1 \ --save_freq=2 \ diff --git a/README.md b/README.md index 2a330d823..f72952102 100644 --- a/README.md +++ b/README.md @@ -97,7 +97,7 @@ Training a policy is as simple as running a script configuration: ```bash lerobot-train \ - --policy=act \ + --policy.type=act \ --dataset.repo_id=lerobot/aloha_mobile_cabinet ``` diff --git a/docs/source/cameras.mdx b/docs/source/cameras.mdx index 2dc2859dd..02714d591 100644 --- a/docs/source/cameras.mdx +++ b/docs/source/cameras.mdx @@ -157,6 +157,14 @@ finally: +### Working with depth + +The Intel RealSense and Reachy 2 cameras can capture both color and depth in lockstep. Calling `read()` returns the **color** frame as `(H, W, 3)` `uint8`. Calling `read_depth()` returns the **depth map** as `(H, W, 1)` `uint16`, where each pixel value is the distance from the sensor expressed in **millimetres**. A pixel value of `0` typically means "no measurement available" (out-of-range, occluded, or low-confidence). + +During recording, the control loop peeks the freshest buffered frames non-blockingly via `read_latest()` (color) and `read_latest_depth()` (depth), adding the depth map as a sibling feature (e.g. `front_depth` next to `front`). + +For how depth streams are stored and encoded when recording a dataset, see the [Depth streams](./video_encoding_parameters#depth-streams) section of the video encoding guide. + ## Use your phone's camera diff --git a/docs/source/cheat-sheet.mdx b/docs/source/cheat-sheet.mdx index a6afa14c2..0531c95bf 100644 --- a/docs/source/cheat-sheet.mdx +++ b/docs/source/cheat-sheet.mdx @@ -89,6 +89,36 @@ Control the data recording flow using keyboard shortcuts: - Press **Left Arrow (`←`)**: Delete current episode and retry. - Press **Escape (`ESC`)**: Stop, encode videos, and upload. +### Recording depth + +Intel RealSense cameras (`type: intelrealsense`) record a depth stream when you set `use_depth: true`. Depth is quantized to 12-bit codes and stored as its own video. + +```bash +lerobot-record \ + ... \ + --robot.cameras="{ head: {type: intelrealsense, serial_number_or_name: \"0123456789\", width: 640, height: 480, fps: 30, use_depth: true} }" \ + --dataset.repo_id=${HF_USER}/so101_depth_test \ + --dataset.single_task="put the red brick in a bowl" \ + --dataset.depth_encoder.depth_min=0.01 \ + --dataset.depth_encoder.depth_max=10.0 \ + --dataset.depth_encoder.shift=0.0 \ + --dataset.depth_encoder.use_log=true +``` + +### Video encoding parameters + +RGB and depth streams are encoded independently via the `--dataset.rgb_encoder.*` and `--dataset.depth_encoder.*` keys. + +```bash +lerobot-record \ + ... \ + --dataset.rgb_encoder.vcodec=h264 \ + --dataset.rgb_encoder.pix_fmt=yuv420p \ + --dataset.rgb_encoder.crf=23 \ + --dataset.depth_encoder.vcodec=hevc \ + --dataset.depth_encoder.extra_options='{"x265-params": "lossless=1"}' +``` + ### Training Depending on your hardware training the policy might take a few hours. That's how you train simple `ACT` policy: @@ -120,6 +150,14 @@ lerobot-train \ --steps=20000 ``` +No local GPU? Add `--job.target=` (e.g. `a10g-small`) to either command and `lerobot-train` runs it on [Hugging Face Jobs](https://huggingface.co/docs/hub/jobs) instead — it uploads a local-only dataset for you and pushes the trained model. List flavors with `hf jobs hardware`. + +To resume, point `--config_path` at a checkpoint and add `--resume=true`. It accepts a local path or a Hub repo id (the latest checkpoint is fetched), and works locally or on a job by adding `--job.target=`: + +```bash +lerobot-train --config_path=${HF_USER}/policy_test --resume=true --job.target=a10g-small +``` + ### Inference Inference means running the trained policy/model on a robot. For that we use `lerobot-rollout`. You will need to provide a path to your policy. It can be a local path or a path to Hugging Face for example "lerobot/folding_latest". Your cameras configuration needs to match what was used when collecting the dataset. Duration is in seconds if unspecified, it will run forever. diff --git a/docs/source/earthrover_mini_plus.mdx b/docs/source/earthrover_mini_plus.mdx index 508c0e3a9..f3b324093 100644 --- a/docs/source/earthrover_mini_plus.mdx +++ b/docs/source/earthrover_mini_plus.mdx @@ -194,7 +194,7 @@ lerobot-record \ --dataset.single_task="Navigate around obstacles" \ --dataset.streaming_encoding=true \ --dataset.encoder_threads=2 \ - # --dataset.camera_encoder.vcodec=auto \ + # --dataset.rgb_encoder.vcodec=auto \ --display_data=true ``` diff --git a/docs/source/groot.mdx b/docs/source/groot.mdx index a10b5e369..3ab202fb2 100644 --- a/docs/source/groot.mdx +++ b/docs/source/groot.mdx @@ -124,7 +124,7 @@ lerobot-rollout\ --dataset.single_task="Grab and handover the red cube to the other arm" \ --dataset.streaming_encoding=true \ --dataset.encoder_threads=2 \ - # --dataset.camera_encoder.vcodec=auto \ + # --dataset.rgb_encoder.vcodec=auto \ --policy.path=/groot-bimanual \ # your trained model --duration=600 ``` diff --git a/docs/source/hardware_guide.mdx b/docs/source/hardware_guide.mdx index 0998344ec..5f236d3e8 100644 --- a/docs/source/hardware_guide.mdx +++ b/docs/source/hardware_guide.mdx @@ -96,3 +96,4 @@ Notes: - The leading `nvidia-smi` is a quick sanity check that CUDA is visible inside the container — useful to fail fast if the flavor or driver mismatched. - The default Job timeout is 30 minutes; pass `--timeout 4h` (or longer) for real training. - `--flavor` maps onto the table above: `t4-small`/`t4-medium` (T4, ACT only), `l4x1`/`l4x4` (L4 24 GB), `a10g-small/large/largex2/largex4` (A10G 24 GB scaled out), `a100-large` (A100). For the current full catalogue + pricing see [https://huggingface.co/docs/hub/jobs](https://huggingface.co/docs/hub/jobs). +- Prefer not to write the `hf jobs run` wrapper yourself? `lerobot-train` can submit the job for you: just add `--job.target=` to a normal training command and it handles dataset upload, log streaming, and the final model push. See the [imitation-learning training guide](./il_robots). diff --git a/docs/source/hilserl.mdx b/docs/source/hilserl.mdx index 76e985cfe..09a370f3d 100644 --- a/docs/source/hilserl.mdx +++ b/docs/source/hilserl.mdx @@ -719,7 +719,7 @@ Example configuration for training the [reward classifier](https://huggingface.c "num_workers": 4, "steps": 5000, "log_freq": 10, - "eval_freq": 1000, + "env_eval_freq": 1000, "save_freq": 1000, "save_checkpoint": true, "seed": 2, diff --git a/docs/source/hope_jr.mdx b/docs/source/hope_jr.mdx index 1f3b08fd7..c29a9f216 100644 --- a/docs/source/hope_jr.mdx +++ b/docs/source/hope_jr.mdx @@ -232,7 +232,7 @@ lerobot-record \ --dataset.private=true \ --dataset.streaming_encoding=true \ --dataset.encoder_threads=2 \ - # --dataset.camera_encoder.vcodec=auto \ + # --dataset.rgb_encoder.vcodec=auto \ --display_data=true ``` @@ -278,6 +278,6 @@ lerobot-record \ --dataset.num_episodes=10 \ --dataset.streaming_encoding=true \ --dataset.encoder_threads=2 \ - # --dataset.camera_encoder.vcodec=auto \ + # --dataset.rgb_encoder.vcodec=auto \ --policy.path=outputs/train/hopejr_hand/checkpoints/last/pretrained_model ``` diff --git a/docs/source/il_robots.mdx b/docs/source/il_robots.mdx index 53ae5af82..178db13bb 100644 --- a/docs/source/il_robots.mdx +++ b/docs/source/il_robots.mdx @@ -207,7 +207,7 @@ lerobot-record \ --dataset.num_episodes=5 \ --dataset.single_task="Grab the black cube" \ --dataset.streaming_encoding=true \ - # --dataset.camera_encoder.vcodec=auto \ + # --dataset.rgb_encoder.vcodec=auto \ --dataset.encoder_threads=2 ``` @@ -390,9 +390,17 @@ Set the flow of data recording using command-line arguments: Control the data recording flow using keyboard shortcuts: -- Press **Right Arrow (`→`)**: Early stop the current episode or reset time and move to the next. -- Press **Left Arrow (`←`)**: Cancel the current episode and re-record it. -- Press **Escape (`ESC`)**: Immediately stop the session, encode videos, and upload the dataset. +- Press **Right Arrow (`→`)** or **`n`**: Early stop the current episode or reset time and move to the next. +- Press **Left Arrow (`←`)** or **`r`**: Cancel the current episode and re-record it. +- Press **Escape (`ESC`)** or **`q`**: Immediately stop the session, encode videos, and upload the dataset. + + + +These control-flow shortcuts work on **X11, Wayland, and headless/SSH** sessions. When a global keyboard backend isn't available (Wayland, a headless machine, or macOS without Accessibility permission), `lerobot-record` automatically reads the same keys from the terminal — launch it from an interactive terminal and keep it focused. You can also use the letter equivalents **`n`** (next, same as `→`), **`r`** (re-record, same as `←`) and **`q`** (quit, same as `ESC`). No `$DISPLAY` setup is required. + +This applies to the recording control flow only. Keyboard **teleoperation** (driving the robot with the keyboard) still needs a global key backend, so it works only on an X11 session, a Windows desktop, or macOS with Accessibility/Input Monitoring granted — not on Wayland or headless sessions. + + #### Tips for gathering data @@ -406,7 +414,7 @@ If you want to dive deeper into this important topic, you can check out the [blo #### Troubleshooting: -- On Linux, if the left and right arrow keys and escape key don't have any effect during data recording, make sure you've set the `$DISPLAY` environment variable. See [pynput limitations](https://pynput.readthedocs.io/en/latest/limitations.html#linux). +- On Linux, the recording control-flow keys (arrow keys, Escape) work on X11, Wayland, and headless/SSH sessions as long as `lerobot-record` runs in an interactive terminal — no `$DISPLAY` setup is needed. If the keys have no effect, make sure you are in an interactive (TTY) terminal, not a piped/non-TTY session, and that it is focused; the letter equivalents `n` / `r` / `q` also work. Keyboard _teleoperation_ (as opposed to the recording control flow) still requires a global key backend — an X11 session, a Windows desktop, or macOS with Accessibility/Input Monitoring granted — and is unavailable on Wayland or headless machines. See [pynput limitations](https://pynput.readthedocs.io/en/latest/limitations.html#linux). ## Visualize a dataset @@ -506,6 +514,12 @@ lerobot-train \ --resume=true ``` +`--config_path` also accepts a **Hub repo id**: if a run pushed its checkpoints to the Hub (with `--save_checkpoint_to_hub=true`), you can resume straight from the repo — its latest checkpoint is downloaded and training continues, restoring the optimizer, scheduler, step counter and data order: + +```bash +lerobot-train --config_path=${HF_USER}/my_policy --resume=true +``` + If you do not want to push your model to the hub after training use `--policy.push_to_hub=false`. Additionally you can provide extra `tags` or specify a `license` for your model or make the model repo `private` by adding this: `--policy.private=true --policy.tags=\[ppo,rl\] --policy.license=mit` @@ -518,7 +532,9 @@ If your local computer doesn't have a powerful GPU you could utilize Google Cola Hugging Face jobs let's you easily select hardware and run the training in the cloud. So if you don't have a powerful GPU or you need more VRAM or just want to train a model much faster use HF Jobs! It's pay as you go and you simply pay for each second of use, you can see the pricing and additional information [here](https://huggingface.co/docs/hub/jobs). -To run the training use this command: +> **Tip:** if you just want to launch a standard training run, you can skip building the command below and use the integrated **Train on HF Jobs via `--job.target`** flow described further down — `lerobot-train` then submits the job, uploads a local-only dataset for you, and streams the logs. + +To run the training manually use this command: @@ -591,6 +607,51 @@ Once the training is started you can go to [Jobs](https://huggingface.co/setting After training the model will be pushed to hub and you can use it as any other model with LeRobot. +#### Train on HF Jobs via `--job.target` (integrated CLI) + +`lerobot-train` runs locally by default. To run on a HuggingFace GPU without constructing the Docker command yourself, pass `--job.target` with a hardware flavor name: + +```bash +lerobot-train \ + --dataset.repo_id=${HF_USER}/so101_test \ + --policy.type=act \ + --policy.repo_id=${HF_USER}/my_policy \ + --job.target=a10g-small +``` + +List available flavors and pricing with `hf jobs hardware`. The run streams its logs to your terminal; press Ctrl-C to detach (the job keeps running in the cloud). Re-attach or cancel with: + +```bash +hf jobs logs +hf jobs cancel +``` + +If your dataset exists only locally (not yet on the Hub), it is automatically pushed to a **private** Hub repo so the job can download it by `repo_id` (nothing is made public). The trained model is pushed to the model repo at the end of the run. To also push every intermediate checkpoint to the Hub as it is saved (so you can monitor progress mid-run), add `--save_checkpoint_to_hub=true` — this requires a runtime image that includes this feature. + +Every job (and any dataset pushed by the run) is tagged `lerobot` so it's easy to find on the Hub. Add your own with `--job.tags '["my-tag"]'`. + +By default the job is capped at `2d` (48h) of wall-clock. Override it with an HF Jobs duration string, e.g. `--job.timeout=4h` to fail faster or `--job.timeout=7d` for a longer run. + +> **Note:** the model repo is created up front (it holds the staged training config the job runs from). If a run fails before the model is pushed, that repo is left on the Hub so you can inspect it — it is not deleted automatically, so repeated failures can leave empty repos behind. Remove one with `hf repo delete `. + +**Prerequisites:** run `hf auth login` before submitting. For Weights & Biases integration, run `wandb login` or set `WANDB_API_KEY` on your machine — the key is forwarded to the job automatically. + +**Resuming on a job.** Adding `--job.target` to a resume command runs the resume in the cloud — the same command works locally or remotely. The checkpoint repo is the source of truth, and new checkpoints continue the lineage in the same repo: + +```bash +# resume a Hub run on a job (its checkpoints are already on the Hub) +lerobot-train --config_path=${HF_USER}/my_policy --resume=true --job.target=a10g-small + +# resume a LOCAL run on a job — the checkpoint is uploaded to a private Hub repo first, +# then the job resumes from it (a local-only dataset is uploaded the same way) +lerobot-train \ + --config_path=outputs/train/act_so101_test/checkpoints/last/pretrained_model/train_config.json \ + --resume=true \ + --job.target=a10g-small +``` + +Job settings come from the current command, so override `--job.target`, `--job.timeout`, etc. as needed; for the resumed run to itself be resumable later, keep `--save_checkpoint_to_hub=true`. + #### Upload policy checkpoints Once training is done, upload the latest checkpoint with: @@ -612,6 +673,8 @@ hf upload ${HF_USER}/act_so101_test${CKPT} \ Use `lerobot-rollout` to deploy a trained policy on your robot. You can choose different strategies depending on your needs: +The examples below load the model from `--policy.path`. To pin a specific pushed version — useful once `--save_checkpoint_to_hub=true` has committed several checkpoints — add `--policy.pretrained_revision` with a commit hash, branch, or tag. Each pushed checkpoint is tagged with its step (e.g. `--policy.pretrained_revision=010000`), so you can recover a checkpoint by step without looking up its commit sha. + ```bash diff --git a/docs/source/lekiwi.mdx b/docs/source/lekiwi.mdx index 7e7c1a680..739073b65 100644 --- a/docs/source/lekiwi.mdx +++ b/docs/source/lekiwi.mdx @@ -319,7 +319,7 @@ If you want to dive deeper into this important topic, you can check out the [blo #### Troubleshooting: -- On Linux, if the left and right arrow keys and escape key don't have any effect during data recording, make sure you've set the `$DISPLAY` environment variable. See [pynput limitations](https://pynput.readthedocs.io/en/latest/limitations.html#linux). +- On Linux, the recording control-flow keys (arrow keys, Escape) work on X11, Wayland, and headless/SSH sessions as long as you run the recording from an interactive terminal (keep it focused) — no `$DISPLAY` setup is needed; the letter equivalents `n` / `r` / `q` also work. Note that **keyboard teleoperation of the LeKiwi base** is different: it relies on a global key backend and therefore works only on an X11 session, a Windows desktop, or macOS with Accessibility/Input Monitoring granted — not on Wayland or headless machines. See [pynput limitations](https://pynput.readthedocs.io/en/latest/limitations.html#linux). ## Replay an episode diff --git a/docs/source/lerobot-dataset-v3.mdx b/docs/source/lerobot-dataset-v3.mdx index 21cb232d3..0647af0b0 100644 --- a/docs/source/lerobot-dataset-v3.mdx +++ b/docs/source/lerobot-dataset-v3.mdx @@ -44,7 +44,7 @@ lerobot-record \ --dataset.num_episodes=5 \ --dataset.single_task="Grab the black cube" \ --dataset.streaming_encoding=true \ - # --dataset.camera_encoder.vcodec=auto \ + # --dataset.rgb_encoder.vcodec=auto \ --dataset.encoder_threads=2 ``` diff --git a/docs/source/libero.mdx b/docs/source/libero.mdx index 043348690..b95af1d27 100644 --- a/docs/source/libero.mdx +++ b/docs/source/libero.mdx @@ -143,7 +143,7 @@ lerobot-train \ --batch_size=4 \ --eval.batch_size=1 \ --eval.n_episodes=1 \ - --eval_freq=1000 + --env_eval_freq=1000 ``` ## Reproducing published results diff --git a/docs/source/libero_plus.mdx b/docs/source/libero_plus.mdx index 4249bf49e..b065649fa 100644 --- a/docs/source/libero_plus.mdx +++ b/docs/source/libero_plus.mdx @@ -173,7 +173,7 @@ lerobot-train \ --batch_size=4 \ --eval.batch_size=1 \ --eval.n_episodes=1 \ - --eval_freq=1000 + --env_eval_freq=1000 ``` ## Relationship to LIBERO diff --git a/docs/source/metaworld.mdx b/docs/source/metaworld.mdx index 8e629dea9..b7accdfa2 100644 --- a/docs/source/metaworld.mdx +++ b/docs/source/metaworld.mdx @@ -120,11 +120,11 @@ lerobot-train \ --batch_size=4 \ --eval.batch_size=1 \ --eval.n_episodes=1 \ - --eval_freq=1000 + --env_eval_freq=1000 ``` ## Practical tips - Use the one-hot task conditioning for multi-task training (MT10/MT50 conventions) so policies have explicit task context. - Inspect the dataset task descriptions and the `info["is_success"]` keys when writing post-processing or logging so your success metrics line up with the benchmark. -- Adjust `batch_size`, `steps`, and `eval_freq` to match your compute budget. +- Adjust `batch_size`, `steps`, and `env_eval_freq` to match your compute budget. diff --git a/docs/source/molmoact2.mdx b/docs/source/molmoact2.mdx index c6ae24e9e..9eb449ca9 100644 --- a/docs/source/molmoact2.mdx +++ b/docs/source/molmoact2.mdx @@ -17,7 +17,7 @@ the paper, see [allenai/molmoact2](https://github.com/allenai/molmoact2). Install LeRobot with the MolmoAct2 optional dependencies: ```bash -pip install -e ".[molmoact2]" +uv sync --locked --extra molmoact2 ``` To run the models in this repository, you need an NVIDIA GPU. The measurements @@ -46,8 +46,8 @@ The repo has been tested with Ubuntu 22.04. To use MolmoAct2 in a LeRobot training config, set: -```python -policy.type=molmoact2 +```bash +--policy.type=molmoact2 ``` ## Training @@ -103,7 +103,7 @@ accelerate launch \ --batch_size=32 \ --num_workers=4 \ --log_freq=20 \ - --eval_freq=-1 \ + --env_eval_freq=-1 \ --save_checkpoint=true \ --save_freq=2000 ``` @@ -142,7 +142,7 @@ accelerate launch \ --batch_size=32 \ --num_workers=4 \ --log_freq=20 \ - --eval_freq=-1 \ + --env_eval_freq=-1 \ --save_checkpoint=true \ --save_freq=2000 ``` @@ -386,6 +386,68 @@ These results demonstrate MolmoAct2's strong performance across diverse robotic manipulation tasks. To reproduce them, follow the instructions in the LIBERO evaluation section. +## Hardware Deployment (lerobot-rollout) + +LeRobot-format checkpoints are available on the Hub for direct use with +`lerobot-rollout`. Each checkpoint uses specific camera names that must +match your robot's camera configuration. + +### Camera naming convention + +Each checkpoint expects specific `observation.images.*` keys. +If your robot cameras have different names, use `--rename_map` to map them: + +| Checkpoint | Camera keys | Description | +| ----------------------------- | ---------------------- | ------------------------ | +| MolmoAct2-LIBERO-LeRobot | `image`, `wrist_image` | LIBERO sim cameras | +| MolmoAct2-BimanualYAM-LeRobot | `top`, `left`, `right` | YAM 3-camera setup | +| MolmoAct2-DROID-LeRobot | `cam0`, `cam1` | External + wrist | +| MolmoAct2-SO100_101-LeRobot | `cam0`, `cam1` | Primary + secondary view | + +Example with an SO-100 robot using top and side cameras: + +```bash +lerobot-rollout \ + --policy.path=lerobot/MolmoAct2-SO100_101-LeRobot \ + --rename_map='{"observation.images.top": "observation.images.cam0", "observation.images.side": "observation.images.cam1"}' \ + --robot.type=so100_follower \ + --robot.port=/dev/ttyACM0 \ + --robot.cameras='{ + top: {type: opencv, index_or_path: 0, width: 640, height: 480, fps: 30}, + side: {type: opencv, index_or_path: 2, width: 640, height: 480, fps: 30} + }' \ + --task="pick up the red cube" --duration=30 +``` + +To use a wrist camera instead, just change the rename mapping: + +```bash +--rename_map='{"observation.images.top": "observation.images.cam0", "observation.images.wrist": "observation.images.cam1"}' +``` + +### Joint frame transform (SO-100/101 zero-shot) + + +The MolmoAct2-SO100_101 checkpoint was trained on data that uses a different +joint calibration convention than LeRobot >= 0.5.0. Without a frame +correction, the arm may move in the wrong direction. + +This affects both **zero-shot deployment** and **fine-tuning** from the +original checkpoint. The pretrained weights expect the old convention, so +all joint data (observations and actions) must be transformed to match. + +The converted LeRobot checkpoint (`lerobot/MolmoAct2-SO100_101-LeRobot`) +already includes this correction in its processor pipeline. If you convert +or fine-tune the checkpoint yourself, set the following in the policy config (`configuration_molmoact2.py`): + +- `joint_signs`: `[1, -1, 1, 1, 1, 1]` (flips shoulder_lift direction) +- `joint_offsets`: `[0, 90, 90, 0, 0, 0]` (shifts shoulder_lift and elbow_flex by 90°) + +See the [backward compatibility guide](./backwardcomp) for details on the +calibration change. + + + ## Differences From the Original Implementation This LeRobot port is intended to match MolmoAct2 behavior while using LeRobot's diff --git a/docs/source/multi_gpu_training.mdx b/docs/source/multi_gpu_training.mdx index d7369e8f8..7c212364e 100644 --- a/docs/source/multi_gpu_training.mdx +++ b/docs/source/multi_gpu_training.mdx @@ -95,7 +95,7 @@ If you want to scale your hyperparameters when using multiple GPUs, you should d accelerate launch --num_processes=2 $(which lerobot-train) \ --optimizer.lr=2e-4 \ --dataset.repo_id=lerobot/pusht \ - --policy=act + --policy.type=act ``` **Training Steps Scaling:** @@ -110,9 +110,64 @@ accelerate launch --num_processes=2 $(which lerobot-train) \ --batch_size=8 \ --steps=50000 \ --dataset.repo_id=lerobot/pusht \ - --policy=act + --policy.type=act ``` +## Training Large Models with FSDP + +DDP replicates the full model on every GPU, so a model that doesn't fit on one GPU won't fit under +DDP either. For large models, use **FSDP** (Fully Sharded Data Parallel), which shards parameters, +gradients, and optimizer state across GPUs. See the [accelerate FSDP guide](https://huggingface.co/docs/accelerate/usage_guides/fsdp) for background. + +An example on how to launch LeRobot training with FSDP across 4 GPUs (1 machine): + +```bash +accelerate launch --config_file fsdp.yaml --num_processes=4 $(which lerobot-train) \ + --dataset.repo_id=${HF_USER}/my_dataset \ + --policy.type= \ + --output_dir=outputs/train/my_policy_fsdp +``` + +A minimal `fsdp.yaml` (FSDP1; shards params/grads/optimizer — ZeRO-3-equivalent): + +```yaml +compute_environment: LOCAL_MACHINE +distributed_type: FSDP +mixed_precision: bf16 +num_machines: 1 +num_processes: 4 +fsdp_config: + fsdp_version: 1 + fsdp_sharding_strategy: FULL_SHARD # params + grads + optimizer (ZeRO-3) + fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP + fsdp_transformer_layer_cls_to_wrap: # repeated block class to shard + fsdp_use_orig_params: true # required: optimizer is built pre-prepare + fsdp_state_dict_type: FULL_STATE_DICT +``` + +Set `fsdp_transformer_layer_cls_to_wrap` to your model's repeated transformer-block class so each +block is sharded as its own unit. `fsdp_use_orig_params: true` is required because LeRobot builds the +optimizer before `accelerator.prepare()`. + +### FSDP checkpoints + +LeRobot gathers the full state dict across all ranks and the main process writes it as a single +`model.safetensors`, loadable as usual with `Policy.from_pretrained(...)`. Two things to look out for: + +- **Checkpoints store fp32 weights.** Under mixed precision (`bf16`/`fp16`) FSDP keeps an fp32 master + copy, and the checkpoint saves it (~2× the bf16 size on disk) so training can resume consistently + with the fp32 optimizer state; `from_pretrained` casts back to the policy dtype on load. FSDP-specific + caveat: an fp32 checkpoint is materialized in full precision on the target device _before_ casting, + so loading it for inference on a tight GPU can OOM even when the bf16 model would fit — load on CPU + first, or cast `model.safetensors` to the deployment dtype offline. +- The sharded optimizer state is gathered into a full (world-size-independent) state dict and saved + alongside the model in the same `optimizer_state.safetensors` / `optimizer_param_groups.json` + format as single-GPU training, so **resume-from-checkpoint is supported** with `--resume=true`. + Resume reshards both the model and the optimizer state to the _current_ FSDP topology, so you can + resume an FSDP checkpoint on a different number of GPUs. Note that the data sampler is only + sample-exact when the world size and batch size match the original run (a warning is logged + otherwise); the optimizer/model state itself is unaffected. + ## Notes - The `--policy.use_amp` flag in `lerobot-train` is only used when **not** running with accelerate. When using accelerate, mixed precision is controlled by accelerate's configuration. diff --git a/docs/source/multi_task_dit.mdx b/docs/source/multi_task_dit.mdx index 450d8a9f2..ebe46489a 100644 --- a/docs/source/multi_task_dit.mdx +++ b/docs/source/multi_task_dit.mdx @@ -314,7 +314,7 @@ lerobot-train \ --steps=30000 \ --save_freq=1000 \ --log_freq=100 \ - --eval_freq=1000 \ + --env_eval_freq=1000 \ --policy.type=multi_task_dit \ --policy.device=cuda \ --policy.horizon=32 \ diff --git a/docs/source/pi0fast.mdx b/docs/source/pi0fast.mdx index f7272acc5..15dff8071 100644 --- a/docs/source/pi0fast.mdx +++ b/docs/source/pi0fast.mdx @@ -96,7 +96,7 @@ lerobot-train \ --policy.type=pi0_fast \ --output_dir=./outputs/pi0fast_training \ --job_name=pi0fast_training \ - --policy.pretrained_path=lerobot/pi0_fast_base \ + --policy.pretrained_path=lerobot/pi0fast-base \ --policy.dtype=bfloat16 \ --policy.gradient_checkpointing=true \ --policy.chunk_size=10 \ @@ -187,7 +187,7 @@ lerobot-train \ --dataset.repo_id=lerobot/libero \ --output_dir=outputs/libero_pi0fast \ --job_name=libero_pi0fast \ - --policy.path=lerobot/pi0fast_base \ + --policy.path=lerobot/pi0fast-base \ --policy.dtype=bfloat16 \ --steps=100000 \ --save_freq=20000 \ diff --git a/docs/source/reachy2.mdx b/docs/source/reachy2.mdx index 4b08569db..7f975af43 100644 --- a/docs/source/reachy2.mdx +++ b/docs/source/reachy2.mdx @@ -161,7 +161,7 @@ lerobot-record \ --dataset.private=true \ --dataset.streaming_encoding=true \ --dataset.encoder_threads=2 \ - # --dataset.camera_encoder.vcodec=auto \ + # --dataset.rgb_encoder.vcodec=auto \ --display_data=true ``` @@ -203,7 +203,7 @@ lerobot-record \ --dataset.private=true \ --dataset.streaming_encoding=true \ --dataset.encoder_threads=2 \ - # --dataset.camera_encoder.vcodec=auto \ + # --dataset.rgb_encoder.vcodec=auto \ --display_data=true ``` diff --git a/docs/source/robocasa.mdx b/docs/source/robocasa.mdx index f6a784e72..5a335a484 100644 --- a/docs/source/robocasa.mdx +++ b/docs/source/robocasa.mdx @@ -166,7 +166,7 @@ lerobot-train \ --output_dir=./outputs/smolvla_robocasa_CloseFridge \ --steps=100000 \ --batch_size=4 \ - --eval_freq=5000 \ + --env_eval_freq=5000 \ --eval.batch_size=1 \ --eval.n_episodes=5 \ --save_freq=10000 diff --git a/docs/source/so101.mdx b/docs/source/so101.mdx index 1274b8282..5b4ed0985 100644 --- a/docs/source/so101.mdx +++ b/docs/source/so101.mdx @@ -122,7 +122,7 @@ The video below shows the sequence of steps for setting the motor ids. #### Follower -Connect the usb cable from your computer and the power supply to the follower arm's controller board. Then, run the following command or run the API example with the port you got from the previous step. You'll also need to give your leader arm a name with the `id` parameter. +Connect the usb cable from your computer and the power supply to the follower arm's controller board. Then, run the following command or run the API example with the port you got from the previous step. You'll also need to give your follower arm a name with the `id` parameter. diff --git a/docs/source/streaming_video_encoding.mdx b/docs/source/streaming_video_encoding.mdx index 96e049eb3..0be32b717 100644 --- a/docs/source/streaming_video_encoding.mdx +++ b/docs/source/streaming_video_encoding.mdx @@ -17,7 +17,7 @@ This makes `save_episode()` near-instant (the video is already encoded by the ti | Parameter | CLI Flag | Type | Default | Description | | ----------------------- | --------------------------------- | ------------- | ------------- | ----------------------------------------------------------------- | | `streaming_encoding` | `--dataset.streaming_encoding` | `bool` | `True` | Enable real-time encoding during capture | -| `vcodec` | `--dataset.camera_encoder.vcodec` | `str` | `"libsvtav1"` | Video codec. `"auto"` detects best HW encoder | +| `vcodec` | `--dataset.rgb_encoder.vcodec` | `str` | `"libsvtav1"` | Video codec. `"auto"` detects best HW encoder | | `encoder_threads` | `--dataset.encoder_threads` | `int \| None` | `None` (auto) | Threads per encoder instance. `None` will leave the vcoded decide | | `encoder_queue_maxsize` | `--dataset.encoder_queue_maxsize` | `int` | `30` | Max buffered frames per camera (~1s at 30fps). Consumes RAM | @@ -82,15 +82,15 @@ Use HW encoding when: ### Available HW Encoders -| Encoder | Platform | Hardware | CLI Value | -| ------------------- | ------------- | ------------------------------------------------------------------------------------------------ | --------------------------------------------------- | -| `h264_videotoolbox` | macOS | Apple Silicon / Intel | `--dataset.camera_encoder.vcodec=h264_videotoolbox` | -| `hevc_videotoolbox` | macOS | Apple Silicon / Intel | `--dataset.camera_encoder.vcodec=hevc_videotoolbox` | -| `h264_nvenc` | Linux/Windows | NVIDIA GPU | `--dataset.camera_encoder.vcodec=h264_nvenc` | -| `hevc_nvenc` | Linux/Windows | NVIDIA GPU | `--dataset.camera_encoder.vcodec=hevc_nvenc` | -| `h264_vaapi` | Linux | Intel/AMD GPU | `--dataset.camera_encoder.vcodec=h264_vaapi` | -| `h264_qsv` | Linux/Windows | Intel Quick Sync | `--dataset.camera_encoder.vcodec=h264_qsv` | -| `auto` | Any | Probes the system for available HW encoders. Falls back to `libsvtav1` if no HW encoder is found | `--dataset.camera_encoder.vcodec=auto` | +| Encoder | Platform | Hardware | CLI Value | +| ------------------- | ------------- | ------------------------------------------------------------------------------------------------ | ------------------------------------------------ | +| `h264_videotoolbox` | macOS | Apple Silicon / Intel | `--dataset.rgb_encoder.vcodec=h264_videotoolbox` | +| `hevc_videotoolbox` | macOS | Apple Silicon / Intel | `--dataset.rgb_encoder.vcodec=hevc_videotoolbox` | +| `h264_nvenc` | Linux/Windows | NVIDIA GPU | `--dataset.rgb_encoder.vcodec=h264_nvenc` | +| `hevc_nvenc` | Linux/Windows | NVIDIA GPU | `--dataset.rgb_encoder.vcodec=hevc_nvenc` | +| `h264_vaapi` | Linux | Intel/AMD GPU | `--dataset.rgb_encoder.vcodec=h264_vaapi` | +| `h264_qsv` | Linux/Windows | Intel Quick Sync | `--dataset.rgb_encoder.vcodec=h264_qsv` | +| `auto` | Any | Probes the system for available HW encoders. Falls back to `libsvtav1` if no HW encoder is found | `--dataset.rgb_encoder.vcodec=auto` | > [!NOTE] > In order to use the HW accelerated encoders you might need to upgrade your GPU drivers. @@ -100,15 +100,15 @@ Use HW encoding when: ## 5. Troubleshooting -| Symptom | Likely Cause | Fix | -| ------------------------------------------------------------------ | -------------------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| System freezes or choppy robot movement or Rerun visualization lag | CPU starved (100% load usage) | Close other apps, reduce encoding throughput, lower `encoder_threads`, use `h264`, use `display_data=False`. If the CPU continues to be at 100% then it might be insufficient for your setup, consider `--dataset.streaming_encoding=false` or HW encoding (`--dataset.camera_encoder.vcodec=auto`) | -| "Encoder queue full" warnings or dropped frames in dataset | Encoder can't keep up (Queue overflow) | If CPU is not at 100%: Increase `encoder_threads`, increase `encoder_queue_maxsize` or use HW encoding (`--dataset.camera_encoder.vcodec=auto`). | -| High RAM usage | Queue filling faster than encoding | `encoder_threads` too low or CPU insufficient. Reduce `encoder_queue_maxsize` or use HW encoding | -| Large video files | Using HW encoder or H.264 | Expected trade-off. Switch to `libsvtav1` if CPU allows | -| `save_episode()` still slow | `streaming_encoding` is `False` | Set `--dataset.streaming_encoding=true` | -| Encoder thread crash | Codec not available or invalid settings | Check `vcodec` is installed, try `--dataset.camera_encoder.vcodec=auto` | -| Recorded dataset is missing frames | CPU/GPU starvation or occasional load spikes | If ~5% of frames are missing, your system is likely overloaded — follow the recommendations above. If fewer frames are missing (~2%), they are probably due to occasional transient load spikes (often at startup) and can be considered expected. | +| Symptom | Likely Cause | Fix | +| ------------------------------------------------------------------ | -------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | +| System freezes or choppy robot movement or Rerun visualization lag | CPU starved (100% load usage) | Close other apps, reduce encoding throughput, lower `encoder_threads`, use `h264`, use `display_data=False`. If the CPU continues to be at 100% then it might be insufficient for your setup, consider `--dataset.streaming_encoding=false` or HW encoding (`--dataset.rgb_encoder.vcodec=auto`) | +| "Encoder queue full" warnings or dropped frames in dataset | Encoder can't keep up (Queue overflow) | If CPU is not at 100%: Increase `encoder_threads`, increase `encoder_queue_maxsize` or use HW encoding (`--dataset.rgb_encoder.vcodec=auto`). | +| High RAM usage | Queue filling faster than encoding | `encoder_threads` too low or CPU insufficient. Reduce `encoder_queue_maxsize` or use HW encoding | +| Large video files | Using HW encoder or H.264 | Expected trade-off. Switch to `libsvtav1` if CPU allows | +| `save_episode()` still slow | `streaming_encoding` is `False` | Set `--dataset.streaming_encoding=true` | +| Encoder thread crash | Codec not available or invalid settings | Check `vcodec` is installed, try `--dataset.rgb_encoder.vcodec=auto` | +| Recorded dataset is missing frames | CPU/GPU starvation or occasional load spikes | If ~5% of frames are missing, your system is likely overloaded — follow the recommendations above. If fewer frames are missing (~2%), they are probably due to occasional transient load spikes (often at startup) and can be considered expected. | ## 6. Recommended Configurations @@ -146,7 +146,7 @@ On very constrained systems, streaming encoding may compete too heavily with the # 2camsx 640x480x3 @30fps: Requires some tuning. # Use H.264, disable streaming, consider batching encoding -lerobot-record --dataset.camera_encoder.vcodec=h264 --dataset.streaming_encoding=false ... +lerobot-record --dataset.rgb_encoder.vcodec=h264 --dataset.streaming_encoding=false ... ``` ## 7. Closing note diff --git a/docs/source/using_dataset_tools.mdx b/docs/source/using_dataset_tools.mdx index 49247a6c1..e9299d298 100644 --- a/docs/source/using_dataset_tools.mdx +++ b/docs/source/using_dataset_tools.mdx @@ -11,8 +11,9 @@ LeRobot provides several utilities for manipulating datasets: 3. **Merge Datasets** - Combine multiple datasets into one. The datasets must have identical features, and episodes are concatenated in the order specified in `repo_ids` 4. **Add Features** - Add new features to a dataset 5. **Remove Features** - Remove features from a dataset -6. **Convert to Video** - Convert image-based datasets to video format for efficient storage -7. **Show the Info of Datasets** - Show the summary of datasets information such as number of episode etc. +6. **Convert to Video** - Convert image-based datasets to video format for efficient storage (RGB and depth cameras are encoded with separate encoders) +7. **Re-encode Videos** - Re-encode an existing video dataset's RGB and/or depth streams with new encoder settings +8. **Show the Info of Datasets** - Show the summary of datasets information such as number of episode etc. The core implementation is in `lerobot.datasets.dataset_tools`. An example script detailing how to use the tools API is available in `examples/dataset/use_dataset_tools.py`. @@ -117,10 +118,19 @@ lerobot-edit-dataset \ --repo_id lerobot/pusht_image \ --operation.type convert_image_to_video \ --operation.output_dir outputs/pusht_video \ - --operation.camera_encoder.vcodec libsvtav1 \ - --operation.camera_encoder.pix_fmt yuv420p \ - --operation.camera_encoder.g 2 \ - --operation.camera_encoder.crf 30 + --operation.rgb_encoder.vcodec libsvtav1 \ + --operation.rgb_encoder.pix_fmt yuv420p \ + --operation.rgb_encoder.g 2 \ + --operation.rgb_encoder.crf 30 + +# Convert a dataset that includes depth maps, customizing the depth encoder +lerobot-edit-dataset \ + --repo_id lerobot/pusht_image \ + --operation.type convert_image_to_video \ + --operation.output_dir outputs/pusht_video \ + --operation.depth_encoder.depth_min 0.01 \ + --operation.depth_encoder.depth_max 10.0 \ + --operation.depth_encoder.use_log true # Convert only specific episodes lerobot-edit-dataset \ @@ -147,11 +157,42 @@ lerobot-edit-dataset \ **Parameters:** - `output_dir`: Custom output directory (optional - by default uses `new_repo_id` or `{repo_id}_video`) -- `camera_encoder`: Video encoder settings — all sub-fields accessible via `--operation.camera_encoder.. See [Video Encoding Parameters](./video_encoding_parameters) for more details. +- `rgb_encoder`: Video encoder settings applied to RGB cameras — all sub-fields accessible via `--operation.rgb_encoder.`. See [Video Encoding Parameters](./video_encoding_parameters) for more details. +- `depth_encoder`: Video encoder settings applied to depth-map cameras (e.g. from an Intel RealSense). In addition to the standard encoder fields it exposes the depth quantization knobs (`depth_min`, `depth_max`, `shift`, `use_log`), accessible via `--operation.depth_encoder.`. These quantization settings are persisted to the dataset metadata so depth can be dequantized back to physical units on load. See the [Depth streams](./video_encoding_parameters#depth-streams) section for details. - `episode_indices`: List of specific episodes to convert (default: all episodes) - `num_workers`: Number of parallel workers for processing (default: 4) -**Note:** The resulting dataset will be a proper LeRobotDataset with all cameras encoded as videos in the `videos/` directory, with parquet files containing only metadata (no raw image data). All episodes, stats, and tasks are preserved. +**Note:** The resulting dataset will be a proper LeRobotDataset with all cameras encoded as videos in the `videos/` directory, with parquet files containing only metadata (no raw image data). Depth-map cameras are detected automatically and routed to the `depth_encoder`, while RGB cameras use the `rgb_encoder`. All episodes, stats, and tasks are preserved. + +#### Re-encode Videos + +Re-encode the videos of an existing video dataset with different encoder settings, without going back to raw frames. RGB videos use the `rgb_encoder` and depth videos use the `depth_encoder`. Provide only the encoder(s) you want to re-encode; the other stream type is left untouched. + +```bash +# Re-encode all RGB videos with new settings (saves to lerobot/pusht_reencoded by default) +lerobot-edit-dataset \ + --repo_id lerobot/pusht \ + --operation.type reencode_videos \ + --operation.rgb_encoder.vcodec h264 \ + --operation.rgb_encoder.pix_fmt yuv420p \ + --operation.rgb_encoder.crf 23 + +# Re-encode both RGB and depth videos in a dataset with depth maps +lerobot-edit-dataset \ + --repo_id lerobot/pusht_depth \ + --operation.type reencode_videos \ + --operation.rgb_encoder.vcodec h264 \ + --operation.depth_encoder.crf 50 +``` + +**Parameters:** + +- `rgb_encoder`: Encoder settings applied to every RGB video. Omit to skip re-encoding RGB videos. +- `depth_encoder`: Encoder settings applied to every depth video. Omit to skip re-encoding depth videos. +- `num_workers`: Number of parallel workers for processing. + +> [!NOTE] +> When re-encoding depth videos, the existing depth quantization parameters (`depth_min`, `depth_max`, `shift`, `use_log`) and the `is_depth_map` flag are **preserved** — re-encoding only changes the codec/quality of the stored stream, not how depth is dequantized on load. ### Show the information of datasets diff --git a/docs/source/video_encoding_parameters.mdx b/docs/source/video_encoding_parameters.mdx index 0b5b99b2b..132d25056 100644 --- a/docs/source/video_encoding_parameters.mdx +++ b/docs/source/video_encoding_parameters.mdx @@ -2,15 +2,15 @@ When video storage is enabled, LeRobot stores each camera stream as an **MP4** file instead of saving one image file per timestep. Video encoding compresses across time, which usually cuts dataset size and I/O compared to a pile of PNG, while keeping MP4 — a format every player and loader understands. -Encoding frames into an MP4 is a full FFmpeg pipeline: choice of encoder, pixel format, GOP/keyframes, quality vs. speed, and optional extra encoder flags. Most of these knobs are user-tunable through `camera_encoder`, a nested `VideoEncoderConfig` (`lerobot.configs.video.VideoEncoderConfig`) passed through PyAV. +Encoding frames into an MP4 is a full FFmpeg pipeline: choice of encoder, pixel format, GOP/keyframes, quality vs. speed, and optional extra encoder flags. Most of these knobs are user-tunable through `rgb_encoder`, a nested `RGBEncoderConfig` (`lerobot.configs.video.RGBEncoderConfig`) passed through PyAV. -You can set these parameters from the CLI with `--dataset.camera_encoder.` (e.g. with `lerobot-record` or `lerobot-rollout`). The same block applies to every camera video stream in that run. +You can set these parameters from the CLI with `--dataset.rgb_encoder.` (e.g. with `lerobot-record` or `lerobot-rollout`). The same block applies to every camera video stream in that run. - Video storage must be on for `camera_encoder` to have any effect — + Video storage must be on for `rgb_encoder` to have any effect — `use_videos=True` in Python APIs, or `--dataset.video=true` on the CLI (the - recording default). With video off, inputs stay as images and `camera_encoder` - is ignored. + recording default). With video off, inputs stay as images and `rgb_encoder` is + ignored. For details on **when** frames are written vs. encoded (streaming vs. post-episode), queues, and other top-level `--dataset.*` switches, see [Streaming Video Encoding](./streaming_video_encoding). For an encoding-parameter comparison and experiments, see the [video-benchmark Space](https://huggingface.co/spaces/lerobot/video-benchmark). @@ -33,9 +33,9 @@ lerobot-record \ --dataset.single_task="Grab the cube" \ --dataset.streaming_encoding=true \ --dataset.encoder_threads=2 \ - --dataset.camera_encoder.vcodec=h264 \ - --dataset.camera_encoder.preset=fast \ - --dataset.camera_encoder.extra_options={"tune": "film", "profile:v": "high", "bf": 2} \ + --dataset.rgb_encoder.vcodec=h264 \ + --dataset.rgb_encoder.preset=fast \ + --dataset.rgb_encoder.extra_options={"tune": "film", "profile:v": "high", "bf": 2} \ --display_data=true ``` @@ -50,7 +50,7 @@ Only override these parameters if you have a specific reason to, and measure the -All flags below are prefixed with `--dataset.camera_encoder.` on the CLI. +All flags below are prefixed with `--dataset.rgb_encoder.` on the CLI. | Parameter | Type | Default | Description | | --------------- | ---------------- | ------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | @@ -65,6 +65,77 @@ All flags below are prefixed with `--dataset.camera_encoder.` on the CLI. --- +## Depth streams + +Depth maps (Intel RealSense, Reachy 2) are stored as their **own video streams** alongside the RGB streams. Raw depth (`uint16` millimetres or `float32` metres) can't survive an 8-bit codec, so LeRobot **quantizes** each map to a 12-bit code (`[0, 4095]`) — logarithmically by default, to match the `1/depth` error profile of depth sensors — then packs it into a high-bit-depth pixel format (`gray12le`) and encodes it with a 12-bit codec. + +```mermaid +flowchart LR + A["Raw depth (uint16 mm / float32 m)"] --> B["Clip to depth_min, depth_max"] + B --> C["Quantize to 12-bit code 0–4095 (log or linear)"] + C --> D["Pack into gray12le"] + D --> E["Encode video (hevc Main 12)"] + E --> F[("MP4 + metadata: depth_min/max, shift, use_log")] + F -. "load time (depth_output_unit)" .-> G["Dequantize to mm or m"] + + classDef input fill:#e3f2fd,stroke:#1565c0,color:#0d47a1; + classDef encode fill:#ede7f6,stroke:#5e35b1,color:#311b92; + classDef store fill:#fff8e1,stroke:#f9a825,color:#e65100; + classDef load fill:#e8f5e9,stroke:#2e7d32,color:#1b5e20; + + class A input; + class B,C,D,E encode; + class F store; + class G load; +``` + +Configure the depth pipeline through a parallel **`depth_encoder`** block (`DepthEncoderConfig`). It shares every `RGBEncoderConfig` field (`vcodec`, `pix_fmt`, `crf`, …) and adds four quantizer knobs, set via `--dataset.depth_encoder.`: + +```bash +lerobot-record \ + ... \ + --dataset.depth_encoder.vcodec=hevc \ + --dataset.depth_encoder.depth_min=0.05 \ + --dataset.depth_encoder.depth_max=5.0 \ + --dataset.depth_encoder.use_log=true +``` + +| Parameter | Type | Default | Description | +| --------------- | ------- | ------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------- | +| `vcodec` | `str` | `"hevc"` | HEVC Main 12 (a 12-bit-capable codec, MP4-compatible). | +| `extra_options` | `dict` | `{"x265-params": "lossless=1"}` | **Depth defaults to lossless** (exact round-trip); `crf` is ignored. Pass `extra_options={}` and set `crf` for a smaller lossy stream. | +| `pix_fmt` | `str` | `"gray12le"` | Single-channel 12-bit pixel format used to carry the quantized codes. | +| `depth_min` | `float` | `0.01` | Depth in metres mapped to quantum `0`. Values below are clipped on decode. | +| `depth_max` | `float` | `10.0` | Depth in metres mapped to quantum `4095`. Values above are clipped on decode. | +| `shift` | `float` | `3.5` | Pre-log offset (metres) used in logarithmic quantization for numerical stability near zero. Must satisfy `depth_min + shift > 0`. | +| `use_log` | `bool` | `True` | If `true`, quantize in log-space (recommended for typical depth sensors). Set to `false` for uniform/linear quantization. | + +> [!TIP] +> `depth_min`, `depth_max`, and `shift` are always interpreted in **metres**, regardless of the input depth's unit. Inputs are auto-detected: integer arrays (e.g. `uint16` millimetres straight from a RealSense) are treated as millimetres, floating arrays as metres. +> Pick `depth_min` / `depth_max` to bracket the actual working range of your sensor — quanta outside that range saturate, which can crush detail at the boundaries. + +Depth features are flagged with `"is_depth_map": true` in `meta/info.json`, and their quantizer settings (`video.depth_min`, `video.depth_max`, `video.shift`, `video.use_log`) are persisted — which is what lets depth be **dequantized back to physical units** on load. + +### Output unit at load time + +`depth_encoder` is a **record-time** concern. The unit that depth maps are dequantized to on _load_ (e.g. during training) is set separately by the read-time flag `--dataset.depth_output_unit`: + +```bash +lerobot-train \ + --dataset.repo_id=/ \ + --dataset.depth_output_unit=m \ + --policy.type=act +``` + +| Parameter | Type | Default | Description | +| ------------------- | ----- | ------- | -------------------------------------------------------------------------------------------- | +| `depth_output_unit` | `str` | `"mm"` | Physical unit depth maps are dequantized to on load: `"mm"` (millimetres) or `"m"` (metres). | + +> [!TIP] +> This is purely a decode-time presentation choice — it does **not** alter the stored video or its metadata, so the same dataset can be read as `mm` or `m` without re-encoding. It has no effect on datasets without depth cameras. + +--- + ## Persistence in dataset metadata After the first episode of a video stream is encoded, the encoder configuration is **persisted into the dataset metadata** (`meta/info.json`) under each video feature, alongside the values probed from the file itself. For a video feature `observation.images.`, the layout in `info.json` is: @@ -82,7 +153,7 @@ After the first episode of a video stream is encoded, the encoder configuration "video.pix_fmt": "yuv420p", "video.fps": 30, "video.channels": 3, - "video.is_depth_map": false, + "is_depth_map": false, "video.g": 2, "video.crf": 30, "video.preset": "fast", @@ -97,12 +168,12 @@ After the first episode of a video stream is encoded, the encoder configuration Two sources contribute to the `info` block: -- **Stream-derived** (read back from the encoded MP4 with PyAV): `video.height`, `video.width`, `video.codec`, `video.pix_fmt`, `video.fps`, `video.channels`, `video.is_depth_map`, plus `audio.*` if an audio stream is present. -- **Encoder-derived** (taken from `VideoEncoderConfig`): `video.g`, `video.crf`, `video.preset`, `video.fast_decode`, `video.video_backend`, `video.extra_options`. +- **Stream-derived** (read back from the encoded MP4 with PyAV): `video.height`, `video.width`, `video.codec`, `video.pix_fmt`, `video.fps`, `video.channels`, `is_depth_map`, plus `audio.*` if an audio stream is present. +- **Encoder-derived** (taken from `RGBEncoderConfig` or `DepthEncoderConfig`): `video.g`, `video.crf`, `video.preset`, `video.fast_decode`, `video.video_backend`, `video.extra_options`. This block is populated **once**, from the **first** episode. It assumes every - episode in the dataset was encoded with the same `camera_encoder`. Changing + episode in the dataset was encoded with the same `rgb_encoder`. Changing encoder settings partway through a recording is not supported — the `info.json` will only reflect the parameters used for the first episode. diff --git a/docs/source/vlabench.mdx b/docs/source/vlabench.mdx index da579d674..9d45da4ec 100644 --- a/docs/source/vlabench.mdx +++ b/docs/source/vlabench.mdx @@ -165,7 +165,7 @@ lerobot-train \ --output_dir=./outputs/smolvla_vlabench_primitive \ --steps=100000 \ --batch_size=4 \ - --eval_freq=5000 \ + --env_eval_freq=5000 \ --eval.batch_size=1 \ --eval.n_episodes=1 \ --save_freq=10000 diff --git a/examples/lekiwi/evaluate.py b/examples/lekiwi/evaluate.py index 3ddcd1f14..13bb6ac28 100644 --- a/examples/lekiwi/evaluate.py +++ b/examples/lekiwi/evaluate.py @@ -17,7 +17,7 @@ import logging import time -from lerobot.common.control_utils import init_keyboard_listener, predict_action +from lerobot.common.control_utils import predict_action from lerobot.datasets import LeRobotDataset from lerobot.policies import make_pre_post_processors from lerobot.policies.act import ACTPolicy @@ -26,6 +26,7 @@ from lerobot.processor import make_default_processors from lerobot.robots.lekiwi import LeKiwiClient, LeKiwiClientConfig from lerobot.utils.constants import ACTION, OBS_STR from lerobot.utils.feature_utils import build_dataset_frame, hw_to_dataset_features +from lerobot.utils.keyboard_input import init_keyboard_listener from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.utils import log_say from lerobot.utils.visualization_utils import init_rerun, log_rerun_data diff --git a/examples/lekiwi/record.py b/examples/lekiwi/record.py index 2c581f5ff..f62a9eb49 100644 --- a/examples/lekiwi/record.py +++ b/examples/lekiwi/record.py @@ -14,7 +14,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from lerobot.common.control_utils import init_keyboard_listener from lerobot.datasets import LeRobotDataset from lerobot.processor import make_default_processors from lerobot.robots.lekiwi import LeKiwiClient, LeKiwiClientConfig @@ -23,6 +22,7 @@ from lerobot.teleoperators.keyboard import KeyboardTeleop, KeyboardTeleopConfig from lerobot.teleoperators.so_leader import SO100Leader, SO100LeaderConfig from lerobot.utils.constants import ACTION, OBS_STR from lerobot.utils.feature_utils import hw_to_dataset_features +from lerobot.utils.keyboard_input import init_keyboard_listener from lerobot.utils.utils import log_say from lerobot.utils.visualization_utils import init_rerun diff --git a/examples/phone_to_so100/evaluate.py b/examples/phone_to_so100/evaluate.py index e859123d0..d1fb4de67 100644 --- a/examples/phone_to_so100/evaluate.py +++ b/examples/phone_to_so100/evaluate.py @@ -18,7 +18,7 @@ import logging import time from lerobot.cameras.opencv import OpenCVCameraConfig -from lerobot.common.control_utils import init_keyboard_listener, predict_action +from lerobot.common.control_utils import predict_action from lerobot.configs import FeatureType, PolicyFeature from lerobot.datasets import LeRobotDataset, aggregate_pipeline_dataset_features, create_initial_features from lerobot.model.kinematics import RobotKinematics @@ -41,6 +41,7 @@ from lerobot.robots.so_follower.robot_kinematic_processor import ( from lerobot.types import RobotAction, RobotObservation from lerobot.utils.constants import ACTION, OBS_STR from lerobot.utils.feature_utils import build_dataset_frame, combine_feature_dicts +from lerobot.utils.keyboard_input import init_keyboard_listener from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.utils import log_say from lerobot.utils.visualization_utils import init_rerun, log_rerun_data diff --git a/examples/phone_to_so100/record.py b/examples/phone_to_so100/record.py index 87b8e49fd..612e94ab9 100644 --- a/examples/phone_to_so100/record.py +++ b/examples/phone_to_so100/record.py @@ -15,7 +15,6 @@ # limitations under the License. from lerobot.cameras.opencv import OpenCVCameraConfig -from lerobot.common.control_utils import init_keyboard_listener from lerobot.datasets import LeRobotDataset, aggregate_pipeline_dataset_features, create_initial_features from lerobot.model.kinematics import RobotKinematics from lerobot.processor import ( @@ -39,6 +38,7 @@ from lerobot.teleoperators.phone.config_phone import PhoneOS from lerobot.teleoperators.phone.phone_processor import MapPhoneActionToRobotAction from lerobot.types import RobotAction, RobotObservation from lerobot.utils.feature_utils import combine_feature_dicts +from lerobot.utils.keyboard_input import init_keyboard_listener from lerobot.utils.utils import log_say from lerobot.utils.visualization_utils import init_rerun diff --git a/examples/so100_to_so100_EE/evaluate.py b/examples/so100_to_so100_EE/evaluate.py index 63def68d0..2a2022623 100644 --- a/examples/so100_to_so100_EE/evaluate.py +++ b/examples/so100_to_so100_EE/evaluate.py @@ -18,7 +18,7 @@ import logging import time from lerobot.cameras.opencv import OpenCVCameraConfig -from lerobot.common.control_utils import init_keyboard_listener, predict_action +from lerobot.common.control_utils import predict_action from lerobot.configs import FeatureType, PolicyFeature from lerobot.datasets import LeRobotDataset, aggregate_pipeline_dataset_features, create_initial_features from lerobot.model.kinematics import RobotKinematics @@ -41,6 +41,7 @@ from lerobot.robots.so_follower.robot_kinematic_processor import ( from lerobot.types import RobotAction, RobotObservation from lerobot.utils.constants import ACTION, OBS_STR from lerobot.utils.feature_utils import build_dataset_frame, combine_feature_dicts +from lerobot.utils.keyboard_input import init_keyboard_listener from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.utils import log_say from lerobot.utils.visualization_utils import init_rerun, log_rerun_data diff --git a/examples/so100_to_so100_EE/record.py b/examples/so100_to_so100_EE/record.py index a0b92da3b..3706ee4f5 100644 --- a/examples/so100_to_so100_EE/record.py +++ b/examples/so100_to_so100_EE/record.py @@ -16,7 +16,6 @@ from lerobot.cameras.opencv import OpenCVCameraConfig -from lerobot.common.control_utils import init_keyboard_listener from lerobot.datasets import LeRobotDataset, aggregate_pipeline_dataset_features, create_initial_features from lerobot.model.kinematics import RobotKinematics from lerobot.processor import ( @@ -36,6 +35,7 @@ from lerobot.scripts.lerobot_record import record_loop from lerobot.teleoperators.so_leader import SO100Leader, SO100LeaderConfig from lerobot.types import RobotAction, RobotObservation from lerobot.utils.feature_utils import combine_feature_dicts +from lerobot.utils.keyboard_input import init_keyboard_listener from lerobot.utils.utils import log_say from lerobot.utils.visualization_utils import init_rerun diff --git a/pyproject.toml b/pyproject.toml index 6fb2e4173..318f89042 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -124,7 +124,7 @@ hardware = [ "lerobot[deepdiff-dep]", ] viz = [ - "rerun-sdk>=0.24.0,<0.27.0", + "rerun-sdk>=0.24.0,<0.34.0", ] # ── User-facing composite extras (map to CLI scripts) ───── # lerobot-record, lerobot-replay, lerobot-calibrate, lerobot-teleoperate, etc. @@ -140,7 +140,14 @@ av-dep = ["av>=15.0.0,<16.0.0"] pygame-dep = ["pygame>=2.5.1,<2.7.0"] # NOTE: 0.9.16 links against liburdfdom_sensor.so.4, which is unavailable on Ubuntu 24.04 # (noble ships urdfdom 3.x). Cap below 0.9.16 until system urdfdom 4.x is broadly available. -placo-dep = ["placo>=0.9.6,<0.9.16"] +# +# NOTE: placo pulls in pin (Pinocchio), whose binary wheels dlopen specific cmeel sonames +# (liburdfdom_sensor.so.4.0, libtinyxml2.so.10) but declare only `>=` floors on their cmeel +# packages. The 2026-05-21 major bumps (cmeel-urdfdom 6.0.0 -> .so.6, cmeel-tinyxml2 11.0.0 +# -> .so.11) ship newer sonames, so left unpinned the resolver grabs them and `import placo` +# fails at load with "liburdfdom_sensor.so.4.0: cannot open shared object file" (see #3755). +# There is no cmeel-urdfdom 5.x; <5 selects the 4.x ABI the placo/pin wheels are built against. +placo-dep = ["placo>=0.9.6,<0.9.16", "cmeel-urdfdom>=4,<5", "cmeel-tinyxml2<11"] transformers-dep = ["transformers>=5.4.0,<5.6.0"] grpcio-dep = ["grpcio>=1.73.1,<2.0.0", "protobuf>=6.31.1,<8.0.0"] accelerate-dep = ["accelerate>=1.14.0,<2.0.0"] diff --git a/src/lerobot/annotations/steerable_pipeline/frames.py b/src/lerobot/annotations/steerable_pipeline/frames.py index a6c904673..5a6a5879c 100644 --- a/src/lerobot/annotations/steerable_pipeline/frames.py +++ b/src/lerobot/annotations/steerable_pipeline/frames.py @@ -36,7 +36,7 @@ from typing import Any, Protocol import PIL.Image import torch -from lerobot.configs.video import VideoEncoderConfig +from lerobot.configs import RGBEncoderConfig from lerobot.datasets.video_utils import decode_video_frames, reencode_video from .reader import EpisodeRecord, snap_to_frame @@ -164,7 +164,9 @@ class VideoFrameProvider: # only for video-stored cameras. Image-stored cameras (also in # ``camera_keys``) would KeyError, so restrict the list — and the # default — to video keys. - keys = list(self._meta.video_keys) + # Depth cameras are excluded from the annotation pipeline for now. + depth_keys = set(self._meta.depth_keys) + keys = [key for key in self._meta.video_keys if key not in depth_keys] # Last-resort fallback: if metadata didn't surface any video keys but # the caller explicitly named a camera (``--vlm.camera_key=...``), # trust them — the key is by definition known to exist on the dataset. @@ -276,12 +278,12 @@ class VideoFrameProvider: from_timestamp = float(ep[f"videos/{self.camera_key}/from_timestamp"]) to_timestamp = float(ep[f"videos/{self.camera_key}/to_timestamp"]) src = self.root / self._meta.get_video_file_path(record.episode_index, self.camera_key) - encoder = VideoEncoderConfig(vcodec="h264", pix_fmt="yuv420p", g=None, crf=23, preset="ultrafast") + encoder = RGBEncoderConfig(vcodec="h264", pix_fmt="yuv420p", g=None, crf=23, preset="ultrafast") try: reencode_video( src, out_path, - camera_encoder=encoder, + video_encoder=encoder, overwrite=True, start_time_s=from_timestamp, end_time_s=to_timestamp, diff --git a/src/lerobot/async_inference/helpers.py b/src/lerobot/async_inference/helpers.py index 4931c68c5..54f0ca69f 100644 --- a/src/lerobot/async_inference/helpers.py +++ b/src/lerobot/async_inference/helpers.py @@ -105,8 +105,9 @@ def raw_observation_to_observation( def prepare_image(image: torch.Tensor) -> torch.Tensor: - """Minimal preprocessing to turn int8 images to float32 in [0, 1], and create a memory-contiguous tensor""" - image = image.type(torch.float32) / 255 + """Minimal preprocessing to turn RGB uint8 images to float32 in [0, 1], and create a memory-contiguous tensor""" + if image.dtype == torch.uint8: + image = image.type(torch.float32) / 255 image = image.contiguous() return image diff --git a/src/lerobot/cameras/opencv/camera_opencv.py b/src/lerobot/cameras/opencv/camera_opencv.py index b3c20e8dd..e50d24c01 100644 --- a/src/lerobot/cameras/opencv/camera_opencv.py +++ b/src/lerobot/cameras/opencv/camera_opencv.py @@ -436,7 +436,7 @@ class OpenCVCamera(Camera): Internal loop run by the background thread for asynchronous reading. On each iteration: - 1. Reads a color frame + 1. Reads a color frame (blocking call) 2. Stores result in latest_frame and updates timestamp (thread-safe) 3. Sets new_frame_event to notify listeners @@ -485,6 +485,8 @@ class OpenCVCamera(Camera): if self.thread is not None and self.thread.is_alive(): self.thread.join(timeout=2.0) + if self.thread.is_alive(): + logger.warning(f"{self} read thread did not terminate within timeout.") self.thread = None self.stop_event = None diff --git a/src/lerobot/cameras/realsense/camera_realsense.py b/src/lerobot/cameras/realsense/camera_realsense.py index 80008e9f9..29cb1e5e0 100644 --- a/src/lerobot/cameras/realsense/camera_realsense.py +++ b/src/lerobot/cameras/realsense/camera_realsense.py @@ -128,6 +128,7 @@ class RealSenseCamera(Camera): self.fps = config.fps self.color_mode = config.color_mode + self.use_rgb = config.use_rgb self.use_depth = config.use_depth self.warmup_s = config.warmup_s @@ -195,12 +196,15 @@ class RealSenseCamera(Camera): # NOTE(Steven/Caroline): Enforcing at least one second of warmup as RS cameras need a bit of time before the first read. If we don't wait, the first read from the warmup will raise. self.warmup_s = max(self.warmup_s, 1) + warmup_read = self.async_read if self.use_rgb else self.async_read_depth start_time = time.time() while time.time() - start_time < self.warmup_s: - self.async_read(timeout_ms=self.warmup_s * 1000) + warmup_read(timeout_ms=self.warmup_s * 1000) time.sleep(0.1) with self.frame_lock: - if self.latest_color_frame is None or self.use_depth and self.latest_depth_frame is None: + if (self.use_rgb and self.latest_color_frame is None) or ( + self.use_depth and self.latest_depth_frame is None + ): raise ConnectionError(f"{self} failed to capture frames during warmup.") logger.info(f"{self} connected.") @@ -268,13 +272,13 @@ class RealSenseCamera(Camera): ) if len(found_devices) > 1: - serial_numbers = [dev["serial_number"] for dev in found_devices] + serial_numbers = [dev["id"] for dev in found_devices] raise ValueError( f"Multiple RealSense cameras found with name '{name}'. " f"Please use a unique serial number instead. Found SNs: {serial_numbers}" ) - serial_number = str(found_devices[0]["serial_number"]) + serial_number = str(found_devices[0]["id"]) return serial_number def _configure_rs_pipeline_config(self, rs_config: Any) -> None: @@ -282,15 +286,17 @@ class RealSenseCamera(Camera): rs.config.enable_device(rs_config, self.serial_number) if self.width and self.height and self.fps: - rs_config.enable_stream( - rs.stream.color, self.capture_width, self.capture_height, rs.format.rgb8, self.fps - ) + if self.use_rgb: + rs_config.enable_stream( + rs.stream.color, self.capture_width, self.capture_height, rs.format.rgb8, self.fps + ) if self.use_depth: rs_config.enable_stream( rs.stream.depth, self.capture_width, self.capture_height, rs.format.z16, self.fps ) else: - rs_config.enable_stream(rs.stream.color) + if self.use_rgb: + rs_config.enable_stream(rs.stream.color) if self.use_depth: rs_config.enable_stream(rs.stream.depth) @@ -298,8 +304,9 @@ class RealSenseCamera(Camera): def _configure_capture_settings(self) -> None: """Sets fps, width, and height from device stream if not already configured. - Uses the color stream profile to update unset attributes. Handles rotation by - swapping width/height when needed. Original capture dimensions are always stored. + Uses the color stream profile (or the depth stream profile when the color + stream is disabled) to update unset attributes. Handles rotation by swapping + width/height when needed. Original capture dimensions are always stored. Raises: DeviceNotConnectedError: If device is not connected. @@ -308,7 +315,8 @@ class RealSenseCamera(Camera): if self.rs_profile is None: raise RuntimeError(f"{self}: rs_profile must be initialized before use.") - stream = self.rs_profile.get_stream(rs.stream.color).as_video_stream_profile() + rs_stream = rs.stream.color if self.use_rgb else rs.stream.depth + stream = self.rs_profile.get_stream(rs_stream).as_video_stream_profile() if self.fps is None: self.fps = stream.fps() @@ -323,6 +331,14 @@ class RealSenseCamera(Camera): self.width, self.height = actual_width, actual_height self.capture_width, self.capture_height = actual_width, actual_height + def _read(self, read_depth: bool = False) -> NDArray[Any]: + """Shared helper for :meth:`read`/:meth:`read_depth`: wait for a fresh color or depth frame.""" + if self.thread is None or not self.thread.is_alive(): + raise RuntimeError(f"{self} read thread is not running.") + + self.new_frame_event.clear() + return self._async_read(timeout_ms=10000, read_depth=read_depth) + @check_if_not_connected def read_depth(self, timeout_ms: int = 200) -> NDArray[Any]: """ @@ -332,8 +348,8 @@ class RealSenseCamera(Camera): from the camera hardware via the RealSense pipeline. Returns: - np.ndarray: The depth map as a NumPy array (height, width) - of type `np.uint16` (raw depth values in millimeters) and rotation. + np.ndarray: The depth map as a NumPy array (height, width, 1) + of type `np.uint16` (raw depth values in millimeters). Raises: DeviceNotConnectedError: If the camera is not connected. @@ -349,20 +365,7 @@ class RealSenseCamera(Camera): f"Failed to capture depth frame '.read_depth()'. Depth stream is not enabled for {self}." ) - if self.thread is None or not self.thread.is_alive(): - raise RuntimeError(f"{self} read thread is not running.") - - self.new_frame_event.clear() - - _ = self.async_read(timeout_ms=10000) - - with self.frame_lock: - depth_map = self.latest_depth_frame - - if depth_map is None: - raise RuntimeError("No depth frame available. Ensure camera is streaming.") - - return depth_map + return self._read(read_depth=True) def _read_from_hardware(self): if self.rs_pipeline is None: @@ -405,12 +408,10 @@ class RealSenseCamera(Camera): f"{self} read() timeout_ms parameter is deprecated and will be removed in future versions." ) - if self.thread is None or not self.thread.is_alive(): - raise RuntimeError(f"{self} read thread is not running.") + if not self.use_rgb: + raise RuntimeError(f"{self}: cannot read color — camera was configured with use_rgb=False.") - self.new_frame_event.clear() - - frame = self.async_read(timeout_ms=10000) + frame = self._read() read_duration_ms = (time.perf_counter() - start_time) * 1e3 logger.debug(f"{self} read took: {read_duration_ms:.1f}ms") @@ -465,8 +466,8 @@ class RealSenseCamera(Camera): Internal loop run by the background thread for asynchronous reading. On each iteration: - 1. Reads a color frame with 500ms timeout - 2. Stores result in latest_frame and updates timestamp (thread-safe) + 1. Reads a color/depth frame (blocking call with 10s timeout) + 2. Stores result in latest_color_frame/latest_depth_frame and updates timestamp (thread-safe) 3. Sets new_frame_event to notify listeners Stops on DeviceNotConnectedError, logs other errors and continues. @@ -479,19 +480,24 @@ class RealSenseCamera(Camera): while not stop_event.is_set(): try: frame = self._read_from_hardware() - color_frame_raw = frame.get_color_frame() - color_frame = np.asanyarray(color_frame_raw.get_data()) - processed_color_frame = self._postprocess_image(color_frame) + + if self.use_rgb: + color_frame_raw = frame.get_color_frame() + color_frame = np.asanyarray(color_frame_raw.get_data()) + processed_color_frame = self._postprocess_image(color_frame) if self.use_depth: depth_frame_raw = frame.get_depth_frame() depth_frame = np.asanyarray(depth_frame_raw.get_data()) processed_depth_frame = self._postprocess_image(depth_frame, depth_frame=True) + if processed_depth_frame.ndim == 2: # (H, W) -> (H, W, 1) + processed_depth_frame = processed_depth_frame[..., np.newaxis] capture_time = time.perf_counter() with self.frame_lock: - self.latest_color_frame = processed_color_frame + if self.use_rgb: + self.latest_color_frame = processed_color_frame if self.use_depth: self.latest_depth_frame = processed_depth_frame self.latest_timestamp = capture_time @@ -523,6 +529,8 @@ class RealSenseCamera(Camera): if self.thread is not None and self.thread.is_alive(): self.thread.join(timeout=2.0) + if self.thread.is_alive(): # pragma: no cover + logger.warning(f"{self} read thread did not terminate within timeout.") self.thread = None self.stop_event = None @@ -533,7 +541,26 @@ class RealSenseCamera(Camera): self.latest_timestamp = None self.new_frame_event.clear() - # NOTE(Steven): Missing implementation for depth for now + def _async_read(self, timeout_ms: float, read_depth: bool = False) -> NDArray[Any]: + """Shared helper for :meth:`async_read`/:meth:`async_read_depth`: return the latest buffered frame.""" + if self.thread is None or not self.thread.is_alive(): + raise RuntimeError(f"{self} read thread is not running.") + + if not self.new_frame_event.wait(timeout=timeout_ms / 1000.0): + raise TimeoutError( + f"Timed out waiting for frame from camera {self} after {timeout_ms} ms. " + f"Read thread alive: {self.thread.is_alive()}." + ) + + with self.frame_lock: + frame = self.latest_depth_frame if read_depth else self.latest_color_frame + self.new_frame_event.clear() + + if frame is None: + raise RuntimeError(f"Internal error: Event set but no frame available for {self}.") + + return frame + @check_if_not_connected def async_read(self, timeout_ms: float = 200) -> NDArray[Any]: """ @@ -558,25 +585,31 @@ class RealSenseCamera(Camera): RuntimeError: If the background thread died unexpectedly or another error occurs. """ + if not self.use_rgb: + raise RuntimeError(f"{self}: cannot read color — camera was configured with use_rgb=False.") + + return self._async_read(timeout_ms=timeout_ms) + + def _read_latest(self, max_age_ms: int, read_depth: bool = False) -> NDArray[Any]: + """Shared helper for :meth:`read_latest`/:meth:`read_latest_depth`: peek the latest buffered frame.""" if self.thread is None or not self.thread.is_alive(): raise RuntimeError(f"{self} read thread is not running.") - if not self.new_frame_event.wait(timeout=timeout_ms / 1000.0): - raise TimeoutError( - f"Timed out waiting for frame from camera {self} after {timeout_ms} ms. " - f"Read thread alive: {self.thread.is_alive()}." - ) - with self.frame_lock: - frame = self.latest_color_frame - self.new_frame_event.clear() + frame = self.latest_depth_frame if read_depth else self.latest_color_frame + timestamp = self.latest_timestamp - if frame is None: - raise RuntimeError(f"Internal error: Event set but no frame available for {self}.") + if frame is None or timestamp is None: + raise RuntimeError(f"{self} has not captured any frames yet.") + + age_ms = (time.perf_counter() - timestamp) * 1e3 + if age_ms > max_age_ms: + raise TimeoutError( + f"{self} latest frame is too old: {age_ms:.1f} ms (max allowed: {max_age_ms} ms)." + ) return frame - # NOTE(Steven): Missing implementation for depth for now @check_if_not_connected def read_latest(self, max_age_ms: int = 500) -> NDArray[Any]: """Return the most recent (color) frame captured immediately (Peeking). @@ -593,24 +626,48 @@ class RealSenseCamera(Camera): DeviceNotConnectedError: If the camera is not connected. RuntimeError: If the camera is connected but has not captured any frames yet. """ + if not self.use_rgb: + raise RuntimeError(f"{self}: cannot read color — camera was configured with use_rgb=False.") - if self.thread is None or not self.thread.is_alive(): - raise RuntimeError(f"{self} read thread is not running.") + return self._read_latest(max_age_ms=max_age_ms) - with self.frame_lock: - frame = self.latest_color_frame - timestamp = self.latest_timestamp + @check_if_not_connected + def async_read_depth(self, timeout_ms: float = 200) -> NDArray[np.uint16]: + """Read the latest depth frame asynchronously, in millimeters. - if frame is None or timestamp is None: - raise RuntimeError(f"{self} has not captured any frames yet.") + Mirrors :meth:`async_read` but returns the depth stream rather than the + color stream. Output is ``np.uint16`` of shape ``(H, W, 1)``, where each + pixel is the distance from the sensor in millimeters. - age_ms = (time.perf_counter() - timestamp) * 1e3 - if age_ms > max_age_ms: - raise TimeoutError( - f"{self} latest frame is too old: {age_ms:.1f} ms (max allowed: {max_age_ms} ms)." - ) + Raises: + DeviceNotConnectedError: If the camera is not connected. + RuntimeError: If ``use_depth`` is ``False`` for this camera, or if + the background read thread is not running. + TimeoutError: If no frame becomes available within ``timeout_ms``. + """ + if not self.use_depth: + raise RuntimeError(f"{self}: cannot read depth — camera was configured with use_depth=False.") - return frame + return self._async_read(timeout_ms=timeout_ms, read_depth=True) + + @check_if_not_connected + def read_latest_depth(self, max_age_ms: int = 500) -> NDArray[Any]: + """Return the most recent depth frame in millimeters (peeking). + + Non-blocking counterpart of :meth:`read_latest` for the depth stream. + Output is ``np.uint16`` of shape ``(H, W, 1)``, where each pixel is the + distance from the sensor in millimeters. + + Raises: + DeviceNotConnectedError: If the camera is not connected. + RuntimeError: If ``use_depth`` is ``False`` for this camera, or if + no depth frame has been captured yet. + TimeoutError: If the latest depth frame is older than ``max_age_ms``. + """ + if not self.use_depth: + raise RuntimeError(f"{self}: cannot read depth — camera was configured with use_depth=False.") + + return self._read_latest(max_age_ms=max_age_ms, read_depth=True) def disconnect(self) -> None: """ diff --git a/src/lerobot/cameras/realsense/configuration_realsense.py b/src/lerobot/cameras/realsense/configuration_realsense.py index 71b083b00..018675195 100644 --- a/src/lerobot/cameras/realsense/configuration_realsense.py +++ b/src/lerobot/cameras/realsense/configuration_realsense.py @@ -42,12 +42,14 @@ class RealSenseCameraConfig(CameraConfig): height: Requested frame height in pixels for the color stream. serial_number_or_name: Unique serial number or human-readable name to identify the camera. color_mode: Color mode for image output (RGB or BGR). Defaults to RGB. + use_rgb: Whether to enable the color stream. Defaults to True. use_depth: Whether to enable depth stream. Defaults to False. rotation: Image rotation setting (0°, 90°, 180°, or 270°). Defaults to no rotation. warmup_s: Time reading frames before returning from connect (in seconds) Note: - Either name or serial_number must be specified. + - At least one of `use_rgb` or `use_depth` must be enabled. - Depth stream configuration (if enabled) will use the same FPS as the color stream. - The actual resolution and FPS may be adjusted by the camera to the nearest supported mode. - For `fps`, `width` and `height`, either all of them need to be set, or none of them. @@ -55,6 +57,7 @@ class RealSenseCameraConfig(CameraConfig): serial_number_or_name: str color_mode: ColorMode = ColorMode.RGB + use_rgb: bool = True use_depth: bool = False rotation: Cv2Rotation = Cv2Rotation.NO_ROTATION warmup_s: int = 1 @@ -63,6 +66,9 @@ class RealSenseCameraConfig(CameraConfig): self.color_mode = ColorMode(self.color_mode) self.rotation = Cv2Rotation(self.rotation) + if not self.use_rgb and not self.use_depth: + raise ValueError("At least one of `use_rgb` or `use_depth` must be enabled.") + values = (self.fps, self.width, self.height) if any(v is not None for v in values) and any(v is None for v in values): raise ValueError( diff --git a/src/lerobot/cameras/zmq/camera_zmq.py b/src/lerobot/cameras/zmq/camera_zmq.py index f3df17814..cd32a117b 100644 --- a/src/lerobot/cameras/zmq/camera_zmq.py +++ b/src/lerobot/cameras/zmq/camera_zmq.py @@ -293,6 +293,8 @@ class ZMQCamera(Camera): if self.thread is not None and self.thread.is_alive(): self.thread.join(timeout=2.0) + if self.thread.is_alive(): + logger.warning(f"{self} read thread did not terminate within timeout.") self.thread = None self.stop_event = None diff --git a/src/lerobot/common/control_utils.py b/src/lerobot/common/control_utils.py index ddaf77d26..e3130643d 100644 --- a/src/lerobot/common/control_utils.py +++ b/src/lerobot/common/control_utils.py @@ -17,12 +17,9 @@ from __future__ import annotations ######################################################################################## # Utilities ######################################################################################## -import logging import time -import traceback from contextlib import nullcontext from copy import copy -from functools import cache from typing import TYPE_CHECKING, Any import numpy as np @@ -43,34 +40,6 @@ from lerobot.robots import Robot from lerobot.types import PolicyAction -@cache -def is_headless(): - """ - Detects if the Python script is running in a headless environment (e.g., without a display). - - This function attempts to import `pynput`, a library that requires a graphical environment. - If the import fails, it assumes the environment is headless. The result is cached to avoid - re-running the check. - - Returns: - True if the environment is determined to be headless, False otherwise. - """ - try: - import pynput # noqa - - return False - except Exception: - print( - "Error trying to import pynput. Switching to headless mode. " - "As a result, the video stream from the cameras won't be shown, " - "and you won't be able to change the control flow with keyboards. " - "For more info, see traceback below.\n" - ) - traceback.print_exc() - print() - return True - - def predict_action( observation: dict[str, np.ndarray], policy: PreTrainedPolicy, @@ -122,59 +91,6 @@ def predict_action( return action -def init_keyboard_listener(): - """ - Initializes a non-blocking keyboard listener for real-time user interaction. - - This function sets up a listener for specific keys (right arrow, left arrow, escape) to control - the program flow during execution, such as stopping recording or exiting loops. It gracefully - handles headless environments where keyboard listening is not possible. - - Returns: - A tuple containing: - - The `pynput.keyboard.Listener` instance, or `None` if in a headless environment. - - A dictionary of event flags (e.g., `exit_early`) that are set by key presses. - """ - # Allow to exit early while recording an episode or resetting the environment, - # by tapping the right arrow key '->'. This might require a sudo permission - # to allow your terminal to monitor keyboard events. - events = {} - events["exit_early"] = False - events["rerecord_episode"] = False - events["stop_recording"] = False - - if is_headless(): - logging.warning( - "Headless environment detected. On-screen cameras display and keyboard inputs will not be available." - ) - listener = None - return listener, events - - # Only import pynput if not in a headless environment - from pynput import keyboard - - def on_press(key): - try: - if key == keyboard.Key.right: - print("Right arrow key pressed. Exiting loop...") - events["exit_early"] = True - elif key == keyboard.Key.left: - print("Left arrow key pressed. Exiting loop and rerecord the last episode...") - events["rerecord_episode"] = True - events["exit_early"] = True - elif key == keyboard.Key.esc: - print("Escape key pressed. Stopping data recording...") - events["stop_recording"] = True - events["exit_early"] = True - except Exception as e: - print(f"Error handling key press: {e}") - - listener = keyboard.Listener(on_press=on_press) - listener.start() - - return listener, events - - def sanity_check_dataset_name(repo_id, policy_cfg): """ Validates the dataset repository name against the presence of a policy configuration. diff --git a/src/lerobot/common/train_utils.py b/src/lerobot/common/train_utils.py index 2d23b4003..b26196f14 100644 --- a/src/lerobot/common/train_utils.py +++ b/src/lerobot/common/train_utils.py @@ -15,12 +15,14 @@ # limitations under the License. from pathlib import Path +from huggingface_hub import HfApi, snapshot_download from torch.optim import Optimizer from torch.optim.lr_scheduler import LRScheduler from lerobot.configs.train import TrainPipelineConfig from lerobot.optim import ( load_optimizer_state, + load_optimizer_state_dict, load_scheduler_state, save_optimizer_state, save_scheduler_state, @@ -34,6 +36,7 @@ from lerobot.utils.constants import ( TRAINING_STATE_DIR, TRAINING_STEP, ) +from lerobot.utils.hub import find_latest_hub_checkpoint from lerobot.utils.io_utils import load_json, write_json from lerobot.utils.random_utils import load_rng_state, save_rng_state @@ -98,6 +101,8 @@ def save_checkpoint( postprocessor: PolicyProcessorPipeline | None = None, num_processes: int | None = None, batch_size: int | None = None, + model_state_dict: dict | None = None, + optim_state_dict: dict | None = None, ) -> None: """This function creates the following directory structure: @@ -127,9 +132,18 @@ def save_checkpoint( resume. Defaults to None (not recorded). batch_size (int | None, optional): Per-process batch size to record for sample-exact resume. Defaults to None (not recorded). + model_state_dict: Pre-gathered full (unsharded) model state dict. Required under FSDP, + where `policy.state_dict()` would return sharded tensors; the caller gathers it via a + cross-rank collective and passes it here so rank 0 can write it directly. It holds + FSDP's fp32 master weights and is saved as-is (the loader casts to the policy dtype on + read). When None (DDP / single-GPU), the model is saved the normal way. Defaults to None. + optim_state_dict: Pre-gathered full (unsharded) optimizer state dict. Required under FSDP + (gathered alongside `model_state_dict` via `gather_fsdp_state_dicts`); saved in the same + safetensors format as the single-GPU path. When None, `optimizer.state_dict()` is used. + Defaults to None. """ pretrained_dir = checkpoint_dir / PRETRAINED_MODEL_DIR - policy.save_pretrained(pretrained_dir) + policy.save_pretrained(pretrained_dir, state_dict=model_state_dict) cfg.save_pretrained(pretrained_dir) if cfg.peft is not None: # When using PEFT, policy.save_pretrained will only write the adapter weights + config, not the @@ -140,7 +154,13 @@ def save_checkpoint( if postprocessor is not None: postprocessor.save_pretrained(pretrained_dir) save_training_state( - checkpoint_dir, step, optimizer, scheduler, num_processes=num_processes, batch_size=batch_size + checkpoint_dir, + step, + optimizer, + scheduler, + num_processes=num_processes, + batch_size=batch_size, + optim_state_dict=optim_state_dict, ) @@ -151,6 +171,7 @@ def save_training_state( scheduler: LRScheduler | None = None, num_processes: int | None = None, batch_size: int | None = None, + optim_state_dict: dict | None = None, ) -> None: """ Saves the training step, optimizer state, scheduler state, and rng state. @@ -164,19 +185,21 @@ def save_training_state( Defaults to None. num_processes (int | None, optional): Distributed world size to record. Defaults to None. batch_size (int | None, optional): Per-process batch size to record. Defaults to None. + optim_state_dict: Pre-gathered full optimizer state dict (for FSDP). Saved instead of + `optimizer.state_dict()` when provided. Defaults to None. """ save_dir = checkpoint_dir / TRAINING_STATE_DIR save_dir.mkdir(parents=True, exist_ok=True) save_training_step(train_step, save_dir, num_processes=num_processes, batch_size=batch_size) save_rng_state(save_dir) if optimizer is not None: - save_optimizer_state(optimizer, save_dir) + save_optimizer_state(optimizer, save_dir, optim_state_dict=optim_state_dict) if scheduler is not None: save_scheduler_state(scheduler, save_dir) def load_training_state( - checkpoint_dir: Path, optimizer: Optimizer, scheduler: LRScheduler | None + checkpoint_dir: Path, optimizer: Optimizer, scheduler: LRScheduler | None, load_optimizer: bool = True ) -> tuple[int, Optimizer, LRScheduler | None]: """ Loads the training step, optimizer state, scheduler state, and rng state. @@ -186,6 +209,10 @@ def load_training_state( checkpoint_dir (Path): The checkpoint directory. Should contain a 'training_state' dir. optimizer (Optimizer): The optimizer to load the state_dict to. scheduler (LRScheduler | None): The scheduler to load the state_dict to (can be None). + load_optimizer (bool, optional): Whether to load the optimizer state from disk. Defaults to + True. Set to False under FSDP, where the sharded optimizer state must be loaded after + `accelerator.prepare()` via `load_fsdp_optimizer_state` (the optimizer is returned + untouched here). Raises: NotADirectoryError: If 'checkpoint_dir' doesn't contain a 'training_state' dir @@ -200,8 +227,119 @@ def load_training_state( load_rng_state(training_state_dir) step = load_training_step(training_state_dir) - optimizer = load_optimizer_state(optimizer, training_state_dir) + if load_optimizer: + optimizer = load_optimizer_state(optimizer, training_state_dir) if scheduler is not None: scheduler = load_scheduler_state(scheduler, training_state_dir) return step, optimizer, scheduler + + +def gather_fsdp_state_dicts(model, optimizer) -> tuple[dict, dict]: + """Gather the full (unsharded) model and optimizer state dicts under FSDP. + + `model.state_dict()` and `FSDP.optim_state_dict(...)` are cross-rank collectives, so this must be + called on *every* rank with the prepared (FSDP-wrapped) `model` and `optimizer`. With + `rank0_only=True` and `offload_to_cpu=True`, every rank runs the all-gather but only rank 0 + materializes the full dicts (the others get empty dicts) and they are kept on CPU to bound GPU + memory. The returned optimizer state dict is keyed by parameter FQNs and is world-size + independent; `load_fsdp_optimizer_state` reshards it on resume. + + Returns: + (model_state_dict, optim_state_dict): full dicts on rank 0, empty dicts on other ranks. + """ + from torch.distributed.fsdp import ( + FullOptimStateDictConfig, + FullStateDictConfig, + FullyShardedDataParallel as FSDP, # noqa F401 + StateDictType, + ) + + state_cfg = FullStateDictConfig(offload_to_cpu=True, rank0_only=True) + optim_cfg = FullOptimStateDictConfig(offload_to_cpu=True, rank0_only=True) + with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, state_cfg, optim_cfg): + model_state_dict = model.state_dict() + optim_state_dict = FSDP.optim_state_dict(model, optimizer) + return model_state_dict, optim_state_dict + + +def load_fsdp_optimizer_state(model, optimizer, checkpoint_dir: Path) -> None: + """Load the FSDP optimizer state (saved as safetensors) and reshard it into the optimizer. + + This is a cross-rank collective and must be called on every rank *after* `accelerator.prepare()` + with the prepared (FSDP-wrapped) `model` and `optimizer`. The saved state is the full, + world-size-independent optimizer state (keyed by parameter FQNs); `FSDP.optim_state_dict_to_load` + reshards it to the current FSDP topology, so resume on a different number of GPUs works. + """ + from torch.distributed.fsdp import ( + FullOptimStateDictConfig, + FullStateDictConfig, + FullyShardedDataParallel as FSDP, # noqa F401 + StateDictType, + ) + + # Every rank reads the same full state from the (shared) checkpoint dir, so rank0_only=False. + full_osd = load_optimizer_state_dict(checkpoint_dir / TRAINING_STATE_DIR) + state_cfg = FullStateDictConfig(rank0_only=False) + optim_cfg = FullOptimStateDictConfig(rank0_only=False) + with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, state_cfg, optim_cfg): + sharded_osd = FSDP.optim_state_dict_to_load(model=model, optim=optimizer, optim_state_dict=full_osd) + optimizer.load_state_dict(sharded_osd) + + +def push_checkpoint_to_hub( + checkpoint_dir: Path, + repo_id: str, + *, + private: bool | None = None, +) -> None: + """Upload a saved checkpoint directory to the Hub under checkpoints//. + + Called once per save step when save_checkpoint_to_hub is enabled, so a + timed-out or crashed run still leaves recoverable checkpoints on the Hub. + The model repo is created idempotently, and the commit is tagged with the + checkpoint step so a checkpoint can be recovered with + --policy.pretrained_revision= instead of a commit sha. + """ + api = HfApi() + api.create_repo(repo_id=repo_id, repo_type="model", private=private, exist_ok=True) + commit = api.upload_folder( + folder_path=str(checkpoint_dir), + repo_id=repo_id, + repo_type="model", + path_in_repo=f"checkpoints/{checkpoint_dir.name}", + commit_message=f"checkpoint {checkpoint_dir.name}", + ) + api.create_tag( + repo_id=repo_id, + tag=checkpoint_dir.name, + revision=commit.oid, + repo_type="model", + exist_ok=True, + ) + + +def resolve_resume_checkpoint(repo_id: str, output_dir: Path) -> Path: + """Download the latest checkpoint of a Hub training repo into a local run dir. + + The symmetric counterpart to `push_checkpoint_to_hub`: given a model repo holding + `checkpoints//{pretrained_model,training_state}` subtrees, download the highest-numbered step + into `output_dir/checkpoints//`, recreate the local `last` symlink, and return that local + checkpoint dir. Used to resume training from the Hub on a machine (or HF Jobs pod) that does not + have the original local run dir. + """ + latest = find_latest_hub_checkpoint(repo_id) + if latest is None: + raise FileNotFoundError( + f"No checkpoint found in '{repo_id}' under '{CHECKPOINTS_DIR}/'. " + "Was the run trained with --save_checkpoint_to_hub?" + ) + snapshot_download( + repo_id=repo_id, + repo_type="model", + allow_patterns=f"{latest}/*", + local_dir=str(output_dir), + ) + checkpoint_dir = output_dir / latest + update_last_checkpoint(checkpoint_dir) + return checkpoint_dir diff --git a/src/lerobot/common/wandb_utils.py b/src/lerobot/common/wandb_utils.py index b782cd751..c229b5eaa 100644 --- a/src/lerobot/common/wandb_utils.py +++ b/src/lerobot/common/wandb_utils.py @@ -180,24 +180,26 @@ class WandBLogger: self._wandb_custom_step_key.add(new_custom_key) self._wandb.define_metric(new_custom_key, hidden=True) + batch_data = {} for k, v in d.items(): + # Skip the custom step key here, it's added to the batch below. + if custom_step_key is not None and k == custom_step_key: + continue + if not isinstance(v, (int | float | str)): logging.warning( f'WandB logging of key "{k}" was ignored as its type "{type(v)}" is not handled by this wrapper.' ) continue - # Do not log the custom step key itself. - if self._wandb_custom_step_key is not None and k in self._wandb_custom_step_key: - continue + batch_data[f"{mode}/{k}"] = v + if batch_data: if custom_step_key is not None: - value_custom_step = d[custom_step_key] - data = {f"{mode}/{k}": v, f"{mode}/{custom_step_key}": value_custom_step} - self._wandb.log(data) - continue - - self._wandb.log(data={f"{mode}/{k}": v}, step=step) + batch_data[f"{mode}/{custom_step_key}"] = d[custom_step_key] + self._wandb.log(batch_data) + else: + self._wandb.log(data=batch_data, step=step) def log_video(self, video_path: str, step: int, mode: str = "train"): if mode not in {"train", "eval"}: diff --git a/src/lerobot/configs/__init__.py b/src/lerobot/configs/__init__.py index be4491811..168b367db 100644 --- a/src/lerobot/configs/__init__.py +++ b/src/lerobot/configs/__init__.py @@ -22,7 +22,7 @@ Import them directly: ``from lerobot.configs.train import TrainPipelineConfig`` """ from .dataset import DatasetRecordConfig -from .default import DatasetConfig, EvalConfig, PeftConfig, WandBConfig +from .default import DatasetConfig, EvalConfig, JobConfig, PeftConfig, WandBConfig from .policies import PreTrainedConfig from .recipe import MessageTurn, TrainingRecipe, load_recipe from .types import ( @@ -33,10 +33,15 @@ from .types import ( RTCAttentionSchedule, ) from .video import ( + DEFAULT_DEPTH_UNIT, VALID_VIDEO_CODECS, VIDEO_ENCODER_INFO_KEYS, + DepthEncoderConfig, + RGBEncoderConfig, VideoEncoderConfig, - camera_encoder_defaults, + depth_encoder_defaults, + encoder_config_from_video_info, + rgb_encoder_defaults, ) __all__ = [ @@ -50,6 +55,7 @@ __all__ = [ "DatasetRecordConfig", "DatasetConfig", "EvalConfig", + "JobConfig", "MessageTurn", "PeftConfig", "PreTrainedConfig", @@ -57,9 +63,15 @@ __all__ = [ "WandBConfig", "load_recipe", "VideoEncoderConfig", + "RGBEncoderConfig", + "DepthEncoderConfig", # Defaults - "camera_encoder_defaults", + "rgb_encoder_defaults", + "depth_encoder_defaults", + # Factories + "encoder_config_from_video_info", # Constants + "DEFAULT_DEPTH_UNIT", "VALID_VIDEO_CODECS", "VIDEO_ENCODER_INFO_KEYS", ] diff --git a/src/lerobot/configs/dataset.py b/src/lerobot/configs/dataset.py index c40c0fae2..7d30ca038 100644 --- a/src/lerobot/configs/dataset.py +++ b/src/lerobot/configs/dataset.py @@ -18,7 +18,7 @@ from dataclasses import dataclass, field from datetime import datetime from pathlib import Path -from .video import VideoEncoderConfig, camera_encoder_defaults +from .video import DepthEncoderConfig, RGBEncoderConfig, depth_encoder_defaults, rgb_encoder_defaults @dataclass @@ -58,8 +58,10 @@ class DatasetRecordConfig: # Set to 1 for immediate encoding (default behavior), or higher for batched encoding video_encoding_batch_size: int = 1 # Video encoder settings for camera MP4s (codec, quality, GOP, etc.). Tuned via CLI nested keys, - # e.g. ``--dataset.camera_encoder.vcodec=h264`` (see ``VideoEncoderConfig``). - camera_encoder: VideoEncoderConfig = field(default_factory=camera_encoder_defaults) + # e.g. ``--dataset.rgb_encoder.vcodec=h264`` (see ``RGBEncoderConfig``). + rgb_encoder: RGBEncoderConfig = field(default_factory=rgb_encoder_defaults) + # Video encoder settings for depth-map MP4s (codec, quality, GOP, etc.). Tuned via CLI nested keys. + depth_encoder: DepthEncoderConfig = field(default_factory=depth_encoder_defaults) # Enable streaming video encoding: encode frames in real-time during capture instead # of writing PNG images first. Makes save_episode() near-instant. More info in the documentation: https://huggingface.co/docs/lerobot/streaming_video_encoding streaming_encoding: bool = False diff --git a/src/lerobot/configs/default.py b/src/lerobot/configs/default.py index b809e71d9..38991a665 100644 --- a/src/lerobot/configs/default.py +++ b/src/lerobot/configs/default.py @@ -19,6 +19,8 @@ from dataclasses import dataclass, field from lerobot.transforms import ImageTransformsConfig from lerobot.utils.import_utils import get_safe_default_video_backend +from .video import DEFAULT_DEPTH_UNIT, DEPTH_METER_UNIT, DEPTH_MILLIMETER_UNIT + @dataclass class DatasetConfig: @@ -35,12 +37,23 @@ class DatasetConfig: revision: str | None = None use_imagenet_stats: bool = True video_backend: str = field(default_factory=get_safe_default_video_backend) - # When True, video frames are returned as uint8 tensors (0-255) instead of float32 (0.0-1.0). + # When True, RGB video frames are returned as uint8 tensors (0-255) instead of float32 (0.0-1.0). # This reduces memory and speeds up DataLoader IPC. The training pipeline handles the conversion. return_uint8: bool = False + # Physical unit depth maps are dequantized to at load time: "mm" (millimeters) or "m" (metres). + # Has no effect on datasets without depth cameras. + depth_output_unit: str = DEFAULT_DEPTH_UNIT streaming: bool = False + # Fraction of episodes held out per task for offline evaluation (0.0 = disabled). + eval_split: float = 0.0 def __post_init__(self) -> None: + if self.depth_output_unit not in (DEPTH_METER_UNIT, DEPTH_MILLIMETER_UNIT): + raise ValueError( + f"depth_output_unit must be '{DEPTH_METER_UNIT}' or '{DEPTH_MILLIMETER_UNIT}', got {self.depth_output_unit!r}" + ) + if not (0.0 <= self.eval_split < 1.0): + raise ValueError(f"eval_split must be in [0.0, 1.0), got {self.eval_split}") if self.episodes is not None: if any(ep < 0 for ep in self.episodes): raise ValueError( @@ -73,8 +86,17 @@ class EvalConfig: # `use_async_envs` specifies whether to use asynchronous environments (multiprocessing). # Defaults to True; automatically downgraded to SyncVectorEnv when batch_size=1. use_async_envs: bool = True + # Whether to record eval rollouts as a LeRobot dataset on disk. + recording: bool = False + # If set, push recorded eval datasets to the Hub under this repo id (one repo per task, + # suffixed by task and env index). Requires recording=true. + recording_repo_id: str | None = None + # Whether the pushed recording repositories should be private. + recording_private: bool = False def __post_init__(self) -> None: + if self.recording_repo_id is not None and not self.recording: + raise ValueError("eval.recording_repo_id requires eval.recording=true.") if self.batch_size == 0: self.batch_size = self._auto_batch_size() if self.batch_size > self.n_episodes: @@ -123,3 +145,35 @@ class PeftConfig: # If None, the PEFT library defaults to alpha=8, which may dampen high-rank adapters. # Common values are r (alpha == rank) or 2*r. lora_alpha: int | None = None + + +@dataclass +class JobConfig: + # Where training runs. None (omitted) or "local" runs on this machine. + # Any other value is an HF Jobs flavor and submits the run to HF Jobs. + # List available flavors + pricing with `hf jobs hardware` command. + target: str | None = None + # Runtime image for the remote job (ignored for local runs). + image: str = "huggingface/lerobot-gpu:latest" + # Max wall-clock for the remote job as an HF Jobs duration string (e.g. "2h"). + # Defaults to "2d": We pass an explicit, generous cap instead. Set a smaller + # value to fail fast, or a larger one for long runs. + timeout: str | None = "2d" + # Submit and exit instead of streaming the job logs in the foreground. + detach: bool = False + # Extra tags attached to the HF job and to any dataset this run pushes to the + # Hub. A "lerobot" tag is always added; e.g. --job.tags '["lelab"]' adds more. + tags: list[str] = field(default_factory=list) + + # Two entry points to the same predicate: the staticmethod tests a raw target string + # straight from argv (before any JobConfig exists, to decide dispatch early), while the + # property is the ergonomic accessor for code that already holds a config instance. + @staticmethod + def is_remote_target(target: str | None) -> bool: + """True when `target` names an HF Jobs flavor rather than a local run.""" + return target not in (None, "local") + + @property + def is_remote(self) -> bool: + """True when training should run on HF Jobs rather than this machine.""" + return self.is_remote_target(self.target) diff --git a/src/lerobot/configs/policies.py b/src/lerobot/configs/policies.py index 91701af6d..b0f003519 100644 --- a/src/lerobot/configs/policies.py +++ b/src/lerobot/configs/policies.py @@ -79,6 +79,8 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC): # type: igno # Either the repo ID of a model hosted on the Hub or a path to a directory containing weights # saved using `Policy.save_pretrained`. If not provided, the policy is initialized from scratch. pretrained_path: Path | None = None + # Optional Hub revision (commit hash, branch, or tag) to pin the pretrained model version. + pretrained_revision: str | None = None def __post_init__(self) -> None: if not self.device or not is_torch_device_available(self.device): diff --git a/src/lerobot/configs/rewards.py b/src/lerobot/configs/rewards.py index 7e99e7f71..92490bc9f 100644 --- a/src/lerobot/configs/rewards.py +++ b/src/lerobot/configs/rewards.py @@ -56,6 +56,8 @@ class RewardModelConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC): device: str | None = None pretrained_path: str | None = None + # Optional Hub revision (commit hash, branch, or tag) to pin the pretrained reward model version. + pretrained_revision: str | None = None push_to_hub: bool = False repo_id: str | None = None diff --git a/src/lerobot/configs/train.py b/src/lerobot/configs/train.py index bac1a946b..e3d354691 100644 --- a/src/lerobot/configs/train.py +++ b/src/lerobot/configs/train.py @@ -26,11 +26,12 @@ from huggingface_hub.errors import HfHubHTTPError from lerobot import envs from lerobot.optim import LRSchedulerConfig, OptimizerConfig -from lerobot.utils.hub import HubMixin +from lerobot.utils.constants import PRETRAINED_MODEL_DIR +from lerobot.utils.hub import HubMixin, find_latest_hub_checkpoint from lerobot.utils.sample_weighting import SampleWeightingConfig from . import parser -from .default import DatasetConfig, EvalConfig, PeftConfig, WandBConfig +from .default import DatasetConfig, EvalConfig, JobConfig, PeftConfig, WandBConfig from .policies import PreTrainedConfig from .rewards import RewardModelConfig @@ -83,10 +84,11 @@ class TrainPipelineConfig(HubMixin): # with the same value for `dir` its contents will be overwritten unless you set `resume` to true. output_dir: Path | None = None job_name: str | None = None - # Set `resume` to true to resume a previous run. In order for this to work, you will need to make sure - # `dir` is the directory of an existing run with at least one checkpoint in it. - # Note that when resuming a run, the default behavior is to use the configuration from the checkpoint, - # regardless of what's provided with the training command at the time of resumption. + # Set `resume` to true to resume a previous run. Pass `--config_path` pointing at either a local + # checkpoint's train_config.json or a Hub repo id holding `checkpoints//` subtrees (the + # latest checkpoint is downloaded and resumed from). Note that when resuming, the default behavior + # is to use the configuration from the checkpoint, regardless of what's provided with the training + # command at the time of resumption (CLI `--*` flags still override). resume: bool = False # `seed` is used for training (eg: model initialization, dataset shuffling) # AND for the evaluation environments. @@ -100,8 +102,13 @@ class TrainPipelineConfig(HubMixin): prefetch_factor: int = 4 persistent_workers: bool = True steps: int = 100_000 - eval_freq: int = 20_000 + # Run policy in the simulation environment every N steps to measure reward/success (0 = disabled). + env_eval_freq: int = 20_000 log_freq: int = 200 + # Compute eval loss on held-out episodes every N steps (0 = disabled). Requires eval_split > 0. + eval_steps: int = 0 + # Cap on total eval samples, split uniformly across tasks (0 = use all held-out data). + max_eval_samples: int = 0 tolerance_s: float = 1e-4 save_checkpoint: bool = True # Checkpoint is saved every `save_freq` training iterations and after the last training step. @@ -113,6 +120,13 @@ class TrainPipelineConfig(HubMixin): wandb: WandBConfig = field(default_factory=WandBConfig) peft: PeftConfig | None = None + # Where to run training (local default, or an HF Jobs flavor). See JobConfig. + job: JobConfig = field(default_factory=JobConfig) + # Push each saved checkpoint to the Hub (policy.repo_id) as it is written, not + # just the final model (useful to monitor progress mid-run). Optional; the + # final model is pushed regardless. Works the same locally and remotely. + save_checkpoint_to_hub: bool = False + # Sample weighting configuration (e.g., for RA-BC training) sample_weighting: SampleWeightingConfig | None = None @@ -132,10 +146,17 @@ class TrainPipelineConfig(HubMixin): return self.reward_model # type: ignore[return-value] return self.policy # type: ignore[return-value] - def validate(self) -> None: - # HACK: We parse again the cli args here to get the pretrained paths if there was some. - policy_path = parser.get_path_arg("policy") + def _resolve_pretrained_from_cli(self) -> None: + """Resolve the pretrained source passed on the CLI into a loaded config. + + The pretrained paths (`--policy.path`, `--reward_model.path`) and + `--config_path` are only recoverable by re-reading the CLI args: draccus + has already consumed them by the time `validate()` runs, so they are not + reflected on `self`. Exactly one source applies, in priority order: + reward-model path, policy path, then resume. + """ reward_model_path = parser.get_path_arg("reward_model") + policy_path = parser.get_path_arg("policy") if reward_model_path: cli_overrides = parser.get_cli_overrides("reward_model") @@ -144,31 +165,54 @@ class TrainPipelineConfig(HubMixin): ) self.reward_model.pretrained_path = str(Path(reward_model_path)) elif policy_path: - yaml_overrides = parser.get_yaml_overrides("policy") - cli_overrides = parser.get_cli_overrides("policy") or [] - self.policy = PreTrainedConfig.from_pretrained( - policy_path, cli_overrides=yaml_overrides + cli_overrides - ) + overrides = parser.get_yaml_overrides("policy") + (parser.get_cli_overrides("policy") or []) + self.policy = PreTrainedConfig.from_pretrained(policy_path, cli_overrides=overrides) self.policy.pretrained_path = Path(policy_path) elif self.resume: - config_path = parser.parse_arg("config_path") - if not config_path: - raise ValueError( - f"A config_path is expected when resuming a run. Please specify path to {TRAIN_CONFIG_NAME}" - ) + self._resolve_resume_checkpoint() - if not Path(config_path).resolve().exists(): - raise NotADirectoryError( - f"{config_path=} is expected to be a local path. " - "Resuming from the hub is not supported for now." - ) + def _resolve_resume_checkpoint(self) -> None: + """Point the trainable config at the checkpoint named by `--config_path`. + `config_path` is either a local path (to a checkpoint's train_config.json or its + pretrained_model/ dir) or a Hub repo id. For a Hub repo, the latest checkpoint is downloaded + into a fresh local run dir and resumed from there. The download is skipped when dispatching to + an HF Job (`job.is_remote`): the pod performs it when it runs the resume locally, and + `submit_to_hf` resolves the source repo for the remote command. + """ + config_path = parser.parse_arg("config_path") + if not config_path: + raise ValueError( + f"A config_path is expected when resuming a run. Please specify path to {TRAIN_CONFIG_NAME}" + ) + + if Path(config_path).resolve().exists(): policy_dir = Path(config_path).parent - if self.policy is not None: - self.policy.pretrained_path = policy_dir - if self.reward_model is not None: - self.reward_model.pretrained_path = str(policy_dir) self.checkpoint_path = policy_dir.parent + elif self.job.is_remote: + return + else: + from lerobot.common.train_utils import resolve_resume_checkpoint + + # `self.output_dir` was loaded from the checkpoint's config and points at the original + # run's (now-absent) local dir. Resume into a fresh local dir instead, unless the user + # passed --output_dir explicitly. + cli_output_dir = parser.parse_arg("output_dir") + if cli_output_dir: + self.output_dir = Path(cli_output_dir) + else: + now = dt.datetime.now() + self.output_dir = Path("outputs/train") / f"{now:%Y-%m-%d}/{now:%H-%M-%S}_resume" + self.checkpoint_path = resolve_resume_checkpoint(config_path, self.output_dir) + policy_dir = self.checkpoint_path / PRETRAINED_MODEL_DIR + + if self.policy is not None: + self.policy.pretrained_path = policy_dir + if self.reward_model is not None: + self.reward_model.pretrained_path = str(policy_dir) + + def validate(self) -> None: + self._resolve_pretrained_from_cli() if self.policy is None and self.reward_model is None: raise ValueError( @@ -208,9 +252,22 @@ class TrainPipelineConfig(HubMixin): self.optimizer = active_cfg.get_optimizer_preset() self.scheduler = active_cfg.get_scheduler_preset() - if hasattr(active_cfg, "push_to_hub") and active_cfg.push_to_hub and not active_cfg.repo_id: + if self.eval_steps > 0 and self.dataset.eval_split == 0.0: + raise ValueError("eval_steps > 0 requires dataset.eval_split > 0.0 to hold out eval data.") + + # Remote runs auto-generate the repo_id in submit_to_hf (the policy may only be + # resolved here, from --policy.path), so don't demand it up front for them. + if ( + hasattr(active_cfg, "push_to_hub") + and active_cfg.push_to_hub + and not active_cfg.repo_id + and not self.job.is_remote + ): raise ValueError("'repo_id' argument missing. Please specify it to push the model to the hub.") + if self.save_checkpoint_to_hub and not (self.policy is not None and self.policy.repo_id): + raise ValueError("save_checkpoint_to_hub requires --policy.repo_id.") + @classmethod def __get_path_fields__(cls) -> list[str]: """Keys for draccus pretrained-path loading.""" @@ -247,22 +304,30 @@ class TrainPipelineConfig(HubMixin): elif Path(model_id).is_file(): config_file = model_id else: + dl_kwargs = { + "repo_id": model_id, + "revision": revision, + "cache_dir": cache_dir, + "force_download": force_download, + "proxies": proxies, + "resume_download": resume_download, + "token": token, + "local_files_only": local_files_only, + } try: - config_file = hf_hub_download( - repo_id=model_id, - filename=TRAIN_CONFIG_NAME, - revision=revision, - cache_dir=cache_dir, - force_download=force_download, - proxies=proxies, - resume_download=resume_download, - token=token, - local_files_only=local_files_only, - ) + config_file = hf_hub_download(filename=TRAIN_CONFIG_NAME, **dl_kwargs) except HfHubHTTPError as e: - raise FileNotFoundError( - f"{TRAIN_CONFIG_NAME} not found on the HuggingFace Hub in {model_id}" - ) from e + # No root train_config.json: this is a repo of periodic checkpoints from an + # interrupted run. Fall back to the latest checkpoint's config so the run can be + # resumed straight from the repo with `--config_path=`. + latest = find_latest_hub_checkpoint(model_id, token=token, revision=revision) + if latest is None: + raise FileNotFoundError( + f"{TRAIN_CONFIG_NAME} not found on the HuggingFace Hub in {model_id}" + ) from e + config_file = hf_hub_download( + filename=f"{latest}/{PRETRAINED_MODEL_DIR}/{TRAIN_CONFIG_NAME}", **dl_kwargs + ) cli_args = kwargs.pop("cli_args", []) # Legacy RA-BC migration only applies to framework-saved checkpoints (always JSON). diff --git a/src/lerobot/configs/video.py b/src/lerobot/configs/video.py index bf2471453..3ea834508 100644 --- a/src/lerobot/configs/video.py +++ b/src/lerobot/configs/video.py @@ -20,7 +20,7 @@ from __future__ import annotations import logging from dataclasses import dataclass, field -from typing import Any +from typing import Any, ClassVar, Self from lerobot.utils.import_utils import require_package @@ -40,7 +40,6 @@ VALID_VIDEO_CODECS: frozenset[str] = frozenset({"h264", "hevc", "libsvtav1", "au # Aliases for legacy video codec names. VIDEO_CODECS_ALIASES: dict[str, str] = {"av1": "libsvtav1"} - LIBSVTAV1_DEFAULT_PRESET: int = 12 # Keys persisted under ``features[*]["info"]`` as ``video.`` (from :class:`VideoEncoderConfig`). @@ -52,40 +51,45 @@ VIDEO_ENCODER_INFO_KEYS: frozenset[str] = frozenset( f"video.{name}" for name in VIDEO_ENCODER_INFO_FIELD_NAMES ) +# Default depth quantization and encoding parameters. +DEPTH_QUANT_BITS: int = 12 +DEPTH_QMAX: int = (1 << DEPTH_QUANT_BITS) - 1 # 4095 + +DEFAULT_DEPTH_MIN: float = 0.01 +DEFAULT_DEPTH_MAX: float = 10.0 +DEFAULT_DEPTH_SHIFT: float = 3.5 +DEFAULT_DEPTH_USE_LOG: bool = True +DEFAULT_DEPTH_PIX_FMT: str = "gray12le" + +DEPTH_METER_UNIT: str = "m" +DEPTH_MILLIMETER_UNIT: str = "mm" +DEFAULT_DEPTH_UNIT: str = DEPTH_MILLIMETER_UNIT + +# Depth-specific tuning fields persisted under ``features[*]["info"]`` as ``video.``. +DEPTH_ENCODER_INFO_FIELD_NAMES: frozenset[str] = frozenset({"depth_min", "depth_max", "shift", "use_log"}) + @dataclass class VideoEncoderConfig: - """Video encoder configuration. + """Video encoder configuration.""" - Attributes: - vcodec: Video encoder name. ``"auto"`` is resolved during - construction (HW encoder if available, else ``libsvtav1``). - pix_fmt: Pixel format (e.g. ``"yuv420p"``). - g: GOP size (keyframe interval). - crf: Quality level — mapped to the native quality parameter of the - codec (``crf`` for software, ``qp`` for NVENC/VAAPI, - ``q:v`` for VideoToolbox, ``global_quality`` for QSV). - preset: Speed/quality preset. Accepted type is per-codec. - fast_decode: Fast-decode tuning. For ``libsvtav1`` this is a level (0-2) - embedded in ``svtav1-params``. For ``h264`` and ``hevc`` non-zero values - set ``tune=fastdecode``. Ignored for other codecs. - video_backend: Python to be used for encoding. Only ``"pyav"`` - is currently supported. - extra_options: Free-form dictionary of additional video encoder options - (e.g. ``{"tune": "film", "profile:v": "high", "bf": 2}``). - """ - - vcodec: str = "libsvtav1" # TODO(CarolinePascal): rename to codec ? - pix_fmt: str = "yuv420p" - g: int | None = 2 - crf: int | float | None = 30 - preset: int | str | None = None - fast_decode: int = 0 + vcodec: str = "libsvtav1" # Video codec name. "auto" picks a hardware codec if available, else libsvtav1. + pix_fmt: str = "yuv420p" # Pixel format (e.g. yuv420p). + g: int | None = 2 # GOP size (keyframe interval). + crf: int | float | None = 30 # Quality level. Lower means better quality and larger files. + preset: int | str | None = None # Speed/quality preset. Accepted values are codec-specific. + fast_decode: int = 0 # Fast-decode tuning. Accepted values are codec-specific, 0 disables it. # TODO(CarolinePascal): add torchcodec support + find a way to unify the # two backends (encoding and decoding). - video_backend: str = "pyav" + video_backend: str = "pyav" # Encoding backend. Only "pyav" is currently supported. + # Extra codec options merged last, e.g. {"tune": "film"}. extra_options: dict[str, Any] = field(default_factory=dict) + # Source-data channel count this encoder is expected to handle. ``None`` + # disables the pix_fmt channel-count check; concrete subclasses set it + # (3 for RGB, 1 for depth, etc.). + _DEFAULT_CHANNELS: ClassVar[int | None] = None + def __post_init__(self) -> None: self.resolve_vcodec() # Empty-constructor ergonomics: ``VideoEncoderConfig()`` must "just work". @@ -94,9 +98,9 @@ class VideoEncoderConfig: self.validate() @classmethod - def from_video_info(cls, video_info: dict | None) -> VideoEncoderConfig: - """Reconstruct a :class:`VideoEncoderConfig` from a video feature's ``info`` block. - Missing or ``None`` values fall back to the class defaults. + def _kwargs_from_video_info(cls, video_info: dict | None) -> dict[str, Any]: + """Parse the ``video.*`` keys of a feature ``info`` block into + constructor kwargs. """ video_info = video_info or {} kwargs: dict[str, Any] = {} @@ -115,7 +119,15 @@ class VideoEncoderConfig: continue kwargs[field_name] = value - return cls(**kwargs) + return kwargs + + @classmethod + def from_video_info(cls, video_info: dict | None) -> Self: + """Reconstruct an encoder config from a video feature's ``info`` block. + + Missing or ``None`` values fall back to the class defaults. + """ + return cls(**cls._kwargs_from_video_info(video_info)) def detect_available_encoders(self, encoders: list[str] | str) -> list[str]: """Return the subset of available encoders based on the specified video backend. @@ -138,7 +150,9 @@ class VideoEncoderConfig: require_package("av", extra="dataset") from lerobot.datasets import check_video_encoder_parameters_pyav - check_video_encoder_parameters_pyav(self.vcodec, self.pix_fmt, self.get_codec_options()) + check_video_encoder_parameters_pyav( + self.vcodec, self.pix_fmt, self.get_codec_options(), channels=self._DEFAULT_CHANNELS + ) def resolve_vcodec(self) -> None: """Check ``vcodec`` and, when it is ``"auto"``, pick a concrete encoder. @@ -230,6 +244,79 @@ class VideoEncoderConfig: return opts -def camera_encoder_defaults() -> VideoEncoderConfig: - """Return a :class:`VideoEncoderConfig` with RGB-camera defaults.""" - return VideoEncoderConfig() +@dataclass +class RGBEncoderConfig(VideoEncoderConfig): + """Encoder configuration for RGB camera streams. + + Identical to :class:`VideoEncoderConfig` but declares the 3-channel + source-data layout so ``pix_fmt`` is validated against RGB inputs. + """ + + _DEFAULT_CHANNELS: ClassVar[int] = 3 + + +def rgb_encoder_defaults() -> RGBEncoderConfig: + """Return a :class:`RGBEncoderConfig` with RGB-camera defaults.""" + return RGBEncoderConfig() + + +@dataclass +class DepthEncoderConfig(VideoEncoderConfig): + """Encoder configuration for depth-map streams. + + Inherits the full :class:`VideoEncoderConfig` surface (codec, GOP, CRF, + preset, ``extra_options``…) and adds the parameters of the depth quantizer. + Defaults flip ``vcodec`` to ``"hevc"`` (Main 12 profile) and ``pix_fmt`` to + ``"gray12le"``. + """ + + vcodec: str = "hevc" # Video codec name. Defaults to HEVC Main 12 (a 12-bit-capable codec). + pix_fmt: str = "gray12le" # Pixel format. Defaults to 12-bit grayscale. + extra_options: dict[str, Any] = field(default_factory=lambda: {"x265-params": "lossless=1"}) + + depth_min: float = DEFAULT_DEPTH_MIN # Minimum depth in meters, mapped to the lowest quantum. + depth_max: float = DEFAULT_DEPTH_MAX # Maximum depth in meters, mapped to the highest quantum. + shift: float = DEFAULT_DEPTH_SHIFT # Pre-log offset in meters for numerical stability near zero. + use_log: bool = DEFAULT_DEPTH_USE_LOG # Use logarithmic quantization (True) or linear (False). + + _DEFAULT_CHANNELS: ClassVar[int] = 1 + + @classmethod + def _kwargs_from_video_info(cls, video_info: dict | None) -> dict[str, Any]: + """Layer the depth-specific tuning (``depth_min`` / ``depth_max`` / + ``shift`` / ``use_log``) on top of the base parser. Missing keys + fall back to the class defaults. + """ + kwargs = super()._kwargs_from_video_info(video_info) + video_info = video_info or {} + for name in DEPTH_ENCODER_INFO_FIELD_NAMES: + value = video_info.get(f"video.{name}") + if value is not None: + kwargs[name] = value + return kwargs + + +def depth_encoder_defaults() -> DepthEncoderConfig: + """Return a :class:`DepthEncoderConfig` with depth-camera defaults.""" + return DepthEncoderConfig() + + +def encoder_config_from_video_info(video_info: dict | None) -> VideoEncoderConfig: + """Build the appropriate encoder config from a feature's ``info`` block. + + Dispatches to :class:`DepthEncoderConfig` when the dict marks the feature + as a depth map and to :class:`RGBEncoderConfig` + otherwise. + + Args: + video_info: A feature's ``info`` dict as persisted in ``info.json``, + or ``None`` (treated as an empty dict). + + Returns: + A :class:`DepthEncoderConfig` for depth features, otherwise a + :class:`RGBEncoderConfig`. + """ + video_info = video_info or {} + is_depth = bool(video_info.get("is_depth_map") or video_info.get("video.is_depth_map")) + cls: type[VideoEncoderConfig] = DepthEncoderConfig if is_depth else RGBEncoderConfig + return cls.from_video_info(video_info) diff --git a/src/lerobot/datasets/__init__.py b/src/lerobot/datasets/__init__.py index bd12a7248..7715a115e 100644 --- a/src/lerobot/datasets/__init__.py +++ b/src/lerobot/datasets/__init__.py @@ -35,7 +35,7 @@ from .dataset_tools import ( remove_feature, split_dataset, ) -from .factory import make_dataset, resolve_delta_timestamps +from .factory import make_dataset, make_train_eval_datasets, resolve_delta_timestamps from .image_writer import safe_stop_image_writer from .io_utils import load_episodes, write_stats from .language import ( @@ -89,6 +89,7 @@ __all__ = [ "get_feature_stats", "load_episodes", "make_dataset", + "make_train_eval_datasets", "merge_datasets", "modify_features", "modify_tasks", diff --git a/src/lerobot/datasets/compute_stats.py b/src/lerobot/datasets/compute_stats.py index 09765c130..88f7ea226 100644 --- a/src/lerobot/datasets/compute_stats.py +++ b/src/lerobot/datasets/compute_stats.py @@ -242,12 +242,12 @@ def sample_images(image_paths: list[str]) -> np.ndarray: images = None for i, idx in enumerate(sampled_indices): path = image_paths[idx] - # we load as uint8 to reduce memory usage + # we load RGB images as uint8 to reduce memory usage; depth keeps its native dtype img = load_image_as_numpy(path, dtype=np.uint8, channel_first=True) img = auto_downsample_height_width(img) if images is None: - images = np.empty((len(sampled_indices), *img.shape), dtype=np.uint8) + images = np.empty((len(sampled_indices), *img.shape), dtype=img.dtype) images[i] = img @@ -506,8 +506,10 @@ def compute_episode_stats( Each statistics dictionary contains min, max, mean, std, count, and quantiles. Note: - Image statistics are normalized to [0,1] range and have shape (3,1,1) for - per-channel values when dtype is 'image' or 'video'. + For 'image'/'video' features, stats are computed per channel and kept with a + leading channel axis (e.g. shape (3, 1, 1) for RGB). RGB stats are divided by + 255 to land in [0, 1]; depth maps (features flagged with ``is_depth_map``) skip + this rescaling and remain in their stored units. """ if quantile_list is None: quantile_list = DEFAULT_QUANTILES @@ -531,8 +533,12 @@ def compute_episode_stats( ) if features[key]["dtype"] in ["image", "video"]: + normalization_factor = ( + 255.0 if not (features[key].get("info") or {}).get("is_depth_map", False) else 1.0 + ) ep_stats[key] = { - k: v if k == "count" else np.squeeze(v / 255.0, axis=0) for k, v in ep_stats[key].items() + k: v if k == "count" else np.squeeze(v / normalization_factor, axis=0) + for k, v in ep_stats[key].items() } return ep_stats @@ -552,8 +558,10 @@ def _validate_stat_value(value: np.ndarray, key: str, feature_key: str) -> None: if key == "count" and value.shape != (1,): raise ValueError(f"Shape of 'count' must be (1), but is {value.shape} instead.") - if "image" in feature_key and key != "count" and value.shape != (3, 1, 1): - raise ValueError(f"Shape of quantile '{key}' must be (3,1,1), but is {value.shape} instead.") + if "image" in feature_key and key != "count" and value.shape not in ((3, 1, 1), (1, 1, 1)): + raise ValueError( + f"Shape of quantile '{key}' must be (3,1,1) or (1,1,1) but is {value.shape} instead." + ) def _assert_type_and_shape(stats_list: list[dict[str, dict]]): diff --git a/src/lerobot/datasets/dataset_metadata.py b/src/lerobot/datasets/dataset_metadata.py index b496e4f65..ea329668c 100644 --- a/src/lerobot/datasets/dataset_metadata.py +++ b/src/lerobot/datasets/dataset_metadata.py @@ -14,7 +14,8 @@ # See the License for the specific language governing permissions and # limitations under the License. import contextlib -from collections.abc import Callable +import logging +from collections.abc import Callable, Iterable from copy import deepcopy from pathlib import Path @@ -338,6 +339,25 @@ class LeRobotDatasetMetadata: """Keys to access visual modalities stored as videos.""" return [key for key, ft in self.features.items() if ft["dtype"] == "video"] + @property + def depth_keys(self) -> list[str]: + """Keys to access depth-map modalities stored as videos or images. + + A depth key is a feature whose ``info`` dict carries ``"is_depth_map": True`` + (or the legacy ``"video.is_depth_map"`` inside ``info`` or ``video_info``). + """ + + def _is_depth(ft: dict) -> bool: + info = ft.get("info") or {} + video_info = ft.get("video_info") or {} + return ( + info.get("is_depth_map", False) + or info.get("video.is_depth_map", False) + or video_info.get("video.is_depth_map", False) + ) + + return [key for key, ft in self.features.items() if _is_depth(ft)] + @property def camera_keys(self) -> list[str]: """Keys to access visual modalities (regardless of their storage method).""" @@ -581,29 +601,48 @@ class LeRobotDatasetMetadata: def update_video_info( self, video_key: str | None = None, - camera_encoder: VideoEncoderConfig | None = None, + video_encoder: VideoEncoderConfig | None = None, + preserve_keys: Iterable[str] | None = None, ) -> None: - """Populate per-feature video info in ``info.json``. + """Populate or refresh per-feature video info in ``info.json``. Warning: this function writes info from first episode videos, implicitly assuming that all videos have been encoded the same way. Also, this means it assumes the first episode exists. + Always re-probes the videos and overwrites existing info for every recomputed + key. ``preserve_keys`` lists keys whose existing values must be kept (e.g. + data-intrinsic entries like ``is_depth_map`` and depth quantization params) + instead of being recomputed. + Args: video_key: If provided, only update this video key. Otherwise update all video keys in the dataset. - camera_encoder: Encoder configuration used to produce the + video_encoder: Encoder configuration used to produce the videos. When provided, its fields are recorded as ``video.`` entries alongside the stream-derived ``video.*`` entries (see :func:`get_video_info`). + preserve_keys: Keys whose existing values are kept instead of being + recomputed. ``None`` (default) recomputes every key. """ if video_key is not None and video_key not in self.video_keys: raise ValueError(f"Video key {video_key} not found in dataset") video_keys = [video_key] if video_key is not None else self.video_keys + preserve_set = set(preserve_keys or ()) for key in video_keys: - if not self.features[key].get("info", None): - video_path = self.root / self.video_path.format(video_key=key, chunk_index=0, file_index=0) - self.info.features[key]["info"] = get_video_info(video_path, camera_encoder=camera_encoder) + existing = self.features[key].get("info") or {} + video_path = self.root / self.video_path.format(video_key=key, chunk_index=0, file_index=0) + new_info = get_video_info(video_path, video_encoder=video_encoder) + # Drop preserved keys so the existing values win on merge. + new_info = {k: v for k, v in new_info.items() if k not in preserve_set} + merged = {**existing, **new_info} + # Migrate the legacy depth marker to the canonical key. + if "video.is_depth_map" in merged: + logging.warning( + f"Migrating legacy 'video.is_depth_map' to 'is_depth_map' for feature {key!r}." + ) + merged.setdefault("is_depth_map", merged.pop("video.is_depth_map")) + self.info.features[key]["info"] = merged def update_chunk_settings( self, diff --git a/src/lerobot/datasets/dataset_reader.py b/src/lerobot/datasets/dataset_reader.py index d7289ac48..e8e07301e 100644 --- a/src/lerobot/datasets/dataset_reader.py +++ b/src/lerobot/datasets/dataset_reader.py @@ -22,7 +22,10 @@ from pathlib import Path import datasets import torch +from lerobot.configs import DEFAULT_DEPTH_UNIT, DepthEncoderConfig + from .dataset_metadata import LeRobotDatasetMetadata +from .depth_utils import dequantize_depth from .feature_utils import ( check_delta_timestamps, get_delta_indices, @@ -51,6 +54,7 @@ class DatasetReader: delta_timestamps: dict[str, list[float]] | None, image_transforms: Callable | None, return_uint8: bool = False, + depth_output_unit: str = DEFAULT_DEPTH_UNIT, ): """Initialize the reader with metadata, filtering, and transform config. @@ -68,6 +72,10 @@ class DatasetReader: relative timestamp offsets for temporal context windows. image_transforms: Optional torchvision v2 transform applied to visual features. + return_uint8: If True, return RGB video frames as raw uint8 tensors + instead of normalized float32. + depth_output_unit: Physical unit depth maps are dequantized to + (``"m"`` or ``"mm"``). Defaults to ``"mm"``. """ self._meta = meta self.root = root @@ -78,6 +86,7 @@ class DatasetReader: raise TypeError("image_transforms must be callable or None.") self._image_transforms = image_transforms self._return_uint8 = return_uint8 + self._depth_output_unit = depth_output_unit self.hf_dataset: datasets.Dataset | None = None self._absolute_to_relative_idx: dict[int, int] | None = None @@ -88,6 +97,11 @@ class DatasetReader: check_delta_timestamps(delta_timestamps, meta.fps, tolerance_s) self.delta_indices = get_delta_indices(delta_timestamps, meta.fps) + self._depth_encoder_configs: dict[str, DepthEncoderConfig] = { + vid_key: DepthEncoderConfig.from_video_info(self._meta.features[vid_key].get("info")) + for vid_key in self._meta.depth_keys + } + def set_image_transforms(self, image_transforms: Callable | None) -> None: """Replace the transform applied to visual observations.""" if image_transforms is not None and not callable(image_transforms): @@ -259,7 +273,18 @@ class DatasetReader: self._tolerance_s, self._video_backend, return_uint8=self._return_uint8, + is_depth=vid_key in self._meta.depth_keys, ) + if vid_key in self._meta.depth_keys: + depth_encoder = self._depth_encoder_configs[vid_key] + frames = dequantize_depth( + frames, + depth_min=depth_encoder.depth_min, + depth_max=depth_encoder.depth_max, + shift=depth_encoder.shift, + use_log=depth_encoder.use_log, + output_unit=self._depth_output_unit, + ) return vid_key, frames.squeeze(0) items = list(query_timestamps.items()) @@ -299,8 +324,9 @@ class DatasetReader: item = {**video_frames, **item} if self._image_transforms is not None: - image_keys = self._meta.camera_keys - for cam in image_keys: + for cam in self._meta.camera_keys: + if cam in self._meta.depth_keys: + continue item[cam] = self._image_transforms(item[cam]) # Add task as a string diff --git a/src/lerobot/datasets/dataset_tools.py b/src/lerobot/datasets/dataset_tools.py index 9aca859b4..31e075d7c 100644 --- a/src/lerobot/datasets/dataset_tools.py +++ b/src/lerobot/datasets/dataset_tools.py @@ -37,7 +37,15 @@ import pyarrow.parquet as pq import torch from tqdm import tqdm -from lerobot.configs import VideoEncoderConfig, camera_encoder_defaults +from lerobot.configs import ( + DepthEncoderConfig, + RGBEncoderConfig, + VideoEncoderConfig, + depth_encoder_defaults, + encoder_config_from_video_info, + rgb_encoder_defaults, +) +from lerobot.configs.video import DEPTH_ENCODER_INFO_FIELD_NAMES from lerobot.utils.constants import ACTION, HF_LEROBOT_HOME, OBS_IMAGE, OBS_STATE from lerobot.utils.utils import flatten_dict @@ -48,6 +56,7 @@ from .compute_stats import ( compute_relative_action_stats, ) from .dataset_metadata import LeRobotDatasetMetadata +from .image_writer import write_image from .io_utils import ( get_parquet_file_size_in_mb, load_episodes, @@ -62,12 +71,13 @@ from .utils import ( DEFAULT_DATA_FILE_SIZE_IN_MB, DEFAULT_DATA_PATH, DEFAULT_EPISODES_PATH, + DEPTH_FILE_PATTERN, + IMAGE_FILE_PATTERN, VIDEO_DIR, update_chunk_file_indices, ) from .video_utils import ( encode_video_frames, - get_video_info, reencode_video, ) @@ -601,7 +611,7 @@ def _keep_episodes_from_video_with_av( output_path: Path, episodes_to_keep: list[tuple[int, int]], fps: float, - camera_encoder: VideoEncoderConfig, + video_encoder: VideoEncoderConfig, ) -> None: """Keep only specified episodes from a video file using PyAV. @@ -615,7 +625,7 @@ def _keep_episodes_from_video_with_av( Ranges are half-open intervals: [start_frame, end_frame), where start_frame is inclusive and end_frame is exclusive. fps: Frame rate of the video. - camera_encoder: Video encoder settings used to re-encode the kept frames. + video_encoder: Video encoder settings used to re-encode the kept frames. """ from fractions import Fraction @@ -640,13 +650,13 @@ def _keep_episodes_from_video_with_av( # Convert fps to Fraction for PyAV compatibility. fps_fraction = Fraction(fps).limit_denominator(1000) - codec_options = camera_encoder.get_codec_options(as_strings=True) - v_out = out.add_stream(camera_encoder.vcodec, rate=fps_fraction, options=codec_options) + codec_options = video_encoder.get_codec_options(as_strings=True) + v_out = out.add_stream(video_encoder.vcodec, rate=fps_fraction, options=codec_options) # PyAV type stubs don't distinguish video streams from audio/subtitle streams. v_out.width = v_in.codec_context.width v_out.height = v_in.codec_context.height - v_out.pix_fmt = camera_encoder.pix_fmt + v_out.pix_fmt = video_encoder.pix_fmt # Set time_base to match the frame rate for proper timestamp handling. v_out.time_base = Fraction(1, int(fps)) @@ -733,7 +743,7 @@ def _copy_and_reindex_videos( for video_key in src_dataset.meta.video_keys: logging.info(f"Processing videos for {video_key}") - camera_encoder = VideoEncoderConfig.from_video_info( + video_encoder = encoder_config_from_video_info( src_dataset.meta.info.features.get(video_key, {}).get("info") ) @@ -817,7 +827,7 @@ def _copy_and_reindex_videos( dst_video_path, episodes_to_keep_ranges, src_dataset.meta.fps, - camera_encoder, + video_encoder, ) cumulative_ts = 0.0 @@ -874,11 +884,11 @@ def _copy_and_reindex_episodes_metadata( episode_meta.update(video_metadata[new_idx]) # Extract episode statistics from parquet metadata. - # Note (maractingi): When pandas/pyarrow serializes numpy arrays with shape (3, 1, 1) to parquet, + # When pandas/pyarrow serializes numpy arrays with shape (C, 1, 1) to parquet, # they are being deserialized as nested object arrays like: # array([array([array([0.])]), array([array([0.])]), array([array([0.])])]) # This happens particularly with image/video statistics. We need to detect and flatten - # these nested structures back to proper (3, 1, 1) arrays so aggregate_stats can process them. + # these nested structures back to proper (C, 1, 1) arrays so aggregate_stats can process them. episode_stats = {} for key in src_episode_full: if key.startswith("stats/"): @@ -894,15 +904,16 @@ def _copy_and_reindex_episodes_metadata( if feature_name in src_dataset.meta.features: feature_dtype = src_dataset.meta.features[feature_name]["dtype"] if feature_dtype in ["image", "video"] and stat_name != "count": + # Stats are channel-first (C, 1, 1) if isinstance(value, np.ndarray) and value.dtype == object: flat_values = [] for item in value: while isinstance(item, np.ndarray): item = item.flatten()[0] flat_values.append(item) - value = np.array(flat_values, dtype=np.float64).reshape(3, 1, 1) - elif isinstance(value, np.ndarray) and value.shape == (3,): - value = value.reshape(3, 1, 1) + value = np.array(flat_values, dtype=np.float64).reshape(-1, 1, 1) + elif isinstance(value, np.ndarray) and value.ndim == 1: + value = value.reshape(-1, 1, 1) episode_stats[feature_name][stat_name] = value @@ -1153,15 +1164,15 @@ def _save_episode_images_for_video( # Get all items for this episode episode_dataset = imgs_dataset.select(range(from_idx, to_idx)) + is_depth = img_key in dataset.meta.depth_keys + frame_pattern = DEPTH_FILE_PATTERN if is_depth else IMAGE_FILE_PATTERN + # Define function to save a single image def save_single_image(i_item_tuple): i, item = i_item_tuple - img = item[img_key] - # Use frame-XXXXXX.png format to match encode_video_frames expectations - img.save(str(imgs_dir / f"frame-{i:06d}.png"), quality=100) + write_image(item[img_key], imgs_dir / frame_pattern.format(frame_index=i)) return i - # Save images with proper naming convention for encode_video_frames (frame-XXXXXX.png) items = list(enumerate(episode_dataset)) with ThreadPoolExecutor(max_workers=num_workers) as executor: @@ -1193,13 +1204,14 @@ def _save_batch_episodes_images( hf_dataset = dataset.hf_dataset.with_format(None) imgs_dataset = hf_dataset.select_columns(img_key) + is_depth = img_key in dataset.meta.depth_keys + frame_pattern = DEPTH_FILE_PATTERN if is_depth else IMAGE_FILE_PATTERN + # Define function to save a single image with global frame index # Defined once outside the loop to avoid repeated closure creation def save_single_image(i_item_tuple, base_frame_idx, img_key_param): i, item = i_item_tuple - img = item[img_key_param] - # Use global frame index for naming - img.save(str(imgs_dir / f"frame-{base_frame_idx + i:06d}.png"), quality=100) + write_image(item[img_key_param], imgs_dir / frame_pattern.format(frame_index=base_frame_idx + i)) return i episode_durations = [] @@ -1290,7 +1302,7 @@ def _estimate_frame_size_via_calibration( episode_indices: list[int], temp_dir: Path, fps: int, - camera_encoder: VideoEncoderConfig, + video_encoder: VideoEncoderConfig, num_calibration_frames: int = 30, ) -> float: """Estimate MB per frame by encoding a small calibration sample. @@ -1304,7 +1316,7 @@ def _estimate_frame_size_via_calibration( episode_indices: List of episode indices being processed. temp_dir: Temporary directory for calibration files. fps: Frames per second for video encoding. - camera_encoder: Video encoder settings used for calibration encoding. + video_encoder: Video encoder settings used for calibration encoding. num_calibration_frames: Number of frames to use for calibration (default: 30). Returns: @@ -1329,10 +1341,11 @@ def _estimate_frame_size_via_calibration( hf_dataset = dataset.hf_dataset.with_format(None) sample_indices = range(from_idx, from_idx + num_frames) - # Save calibration frames + # Save calibration frames using the suffix/format the encoder expects. + is_depth = img_key in dataset.meta.depth_keys + frame_pattern = DEPTH_FILE_PATTERN if is_depth else IMAGE_FILE_PATTERN for i, idx in enumerate(sample_indices): - img = hf_dataset[idx][img_key] - img.save(str(calibration_dir / f"frame-{i:06d}.png"), quality=100) + write_image(hf_dataset[idx][img_key], calibration_dir / frame_pattern.format(frame_index=i)) # Encode calibration video calibration_video_path = calibration_dir / "calibration.mp4" @@ -1340,7 +1353,7 @@ def _estimate_frame_size_via_calibration( imgs_dir=calibration_dir, video_path=calibration_video_path, fps=fps, - camera_encoder=camera_encoder, + video_encoder=video_encoder, overwrite=True, ) @@ -1613,6 +1626,7 @@ def recompute_stats( raise ValueError(f"No parquet files found in {data_dir}") all_episode_stats = [] + # TODO: enable image and video stats re-computation numeric_keys = [k for k, v in features_to_compute.items() if v["dtype"] not in ["image", "video"]] for parquet_path in tqdm(parquet_files, desc="Computing stats from data files"): @@ -1658,7 +1672,8 @@ def convert_image_to_video_dataset( dataset: LeRobotDataset, output_dir: Path | None = None, repo_id: str | None = None, - camera_encoder: VideoEncoderConfig | None = None, + rgb_encoder: RGBEncoderConfig | None = None, + depth_encoder: DepthEncoderConfig | None = None, episode_indices: list[int] | None = None, num_workers: int = 4, max_episodes_per_batch: int | None = None, @@ -1670,21 +1685,32 @@ def convert_image_to_video_dataset( LeRobot dataset structure with videos stored in chunked MP4 files. Args: - dataset: The source LeRobot dataset with images - 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: Edited dataset identifier. Equivalent to new_repo_id in EditDatasetConfig. - camera_encoder: Video encoder settings - (``None`` uses :func:`~lerobot.configs.camera_encoder_defaults`). - episode_indices: List of episode indices to convert (None = all episodes) - num_workers: Number of threads for parallel processing (default: 4) - max_episodes_per_batch: Maximum episodes per video batch to avoid memory issues (None = no limit) - max_frames_per_batch: Maximum frames per video batch to avoid memory issues (None = no limit) + dataset: The source LeRobot dataset with images. + output_dir: Root directory where the converted dataset will be stored. When + ``None``, defaults to ``$HF_LEROBOT_HOME/repo_id``. Equivalent to + ``new_root`` in ``EditDatasetConfig``. + repo_id: Converted dataset identifier. Equivalent to ``new_repo_id`` in + ``EditDatasetConfig``. + rgb_encoder: Video encoder settings applied to RGB cameras. When ``None``, + :func:`~lerobot.configs.video.rgb_encoder_defaults` is used. + depth_encoder: Video encoder settings applied to depth-map cameras, including + the quantization parameters persisted to the dataset metadata. When + ``None``, :func:`~lerobot.configs.video.depth_encoder_defaults` is used. + episode_indices: Episode indices to convert. When ``None``, all episodes are + converted. + num_workers: Number of threads for parallel processing. + max_episodes_per_batch: Maximum episodes per video batch, to bound memory use. + ``None`` means no limit. + max_frames_per_batch: Maximum frames per video batch, to bound memory use. + ``None`` means no limit. Returns: - New LeRobotDataset with images encoded as videos + A new :class:`LeRobotDataset` with images encoded as videos. """ - if camera_encoder is None: - camera_encoder = camera_encoder_defaults() + if rgb_encoder is None: + rgb_encoder = rgb_encoder_defaults() + if depth_encoder is None: + depth_encoder = depth_encoder_defaults() # Check that it's an image dataset if len(dataset.meta.video_keys) > 0: @@ -1709,10 +1735,7 @@ def convert_image_to_video_dataset( logging.info( f"Converting {len(episode_indices)} episodes with {len(img_keys)} cameras from {dataset.repo_id}" ) - logging.info( - f"Video codec: {camera_encoder.vcodec}, pixel format: {camera_encoder.pix_fmt}, " - f"GOP: {camera_encoder.g}, CRF: {camera_encoder.crf}" - ) + logging.info(f"RGB video encoder: {rgb_encoder}, depth video encoder: {depth_encoder}") # Create new features dict, converting image features to video features new_features = {} @@ -1774,6 +1797,8 @@ def convert_image_to_video_dataset( episode_lengths = {ep_idx: dataset.meta.episodes["length"][ep_idx] for ep_idx in episode_indices} for img_key in tqdm(img_keys, desc="Processing cameras"): + target_encoder = depth_encoder if img_key in dataset.meta.depth_keys else rgb_encoder + # Estimate size per frame by encoding a small calibration sample # This provides accurate compression ratio for the specific codec parameters size_per_frame_mb = _estimate_frame_size_via_calibration( @@ -1782,7 +1807,7 @@ def convert_image_to_video_dataset( episode_indices=episode_indices, temp_dir=temp_dir, fps=fps, - camera_encoder=camera_encoder, + video_encoder=target_encoder, ) logging.info(f"Processing camera: {img_key}") @@ -1824,7 +1849,7 @@ def convert_image_to_video_dataset( imgs_dir=imgs_dir, video_path=video_path, fps=fps, - camera_encoder=camera_encoder, + video_encoder=target_encoder, overwrite=True, ) @@ -1863,16 +1888,11 @@ def convert_image_to_video_dataset( new_meta.info.total_tasks = dataset.meta.total_tasks new_meta.info.splits = {"train": f"0:{len(episode_indices)}"} - # Update video info for all image keys (now videos) - # We need to manually set video info since update_video_info() checks video_keys first + # Update video info for all image keys (now videos). They are registered as + # video features above, so update_video_info populates their (still-empty) info. for img_key in img_keys: - if not new_meta.features[img_key].get("info", None): - video_path = new_meta.root / new_meta.video_path.format( - video_key=img_key, chunk_index=0, file_index=0 - ) - new_meta.info.features[img_key]["info"] = get_video_info( - video_path, camera_encoder=camera_encoder - ) + target_encoder = depth_encoder if img_key in dataset.meta.depth_keys else rgb_encoder + new_meta.update_video_info(video_key=img_key, video_encoder=target_encoder) write_info(new_meta.info, new_meta.root) @@ -1899,11 +1919,11 @@ def convert_image_to_video_dataset( def _reencode_video_worker(args: tuple) -> Path: """Picklable worker for :func:`reencode_dataset`'s process pool.""" - video_path, camera_encoder, encoder_threads = args + video_path, video_encoder, encoder_threads = args reencode_video( input_video_path=video_path, output_video_path=video_path, - camera_encoder=camera_encoder, + video_encoder=video_encoder, encoder_threads=encoder_threads, overwrite=True, ) @@ -1912,7 +1932,8 @@ def _reencode_video_worker(args: tuple) -> Path: def reencode_dataset( dataset: LeRobotDataset, - camera_encoder: VideoEncoderConfig, + rgb_encoder: RGBEncoderConfig | None = None, + depth_encoder: DepthEncoderConfig | None = None, encoder_threads: int | None = None, num_workers: int | None = None, ) -> LeRobotDataset: @@ -1923,8 +1944,11 @@ def reencode_dataset( Args: dataset: An existing :class:`LeRobotDataset` whose videos will be re-encoded. - camera_encoder: Target encoder configuration applied to every video - file. + rgb_encoder: Target encoder configuration applied to every RGB video + file. If ``None``, re-encoding is skipped for RGB videos. + depth_encoder: Target encoder configuration applied to every depth video + file. If ``None``, re-encoding is skipped for depth videos. + Quantization parameters will not override the ones in the current dataset. encoder_threads: Per-encoder thread count forwarded to :func:`reencode_video`. ``None`` lets the codec decide. num_workers: Number of parallel processes. ``None`` or ``0`` means @@ -1936,23 +1960,35 @@ def reencode_dataset( on disk. """ meta = dataset.meta - video_paths_list = [] + video_keys_encoders_dict = {} + video_keys_paths_dict = {} + + if rgb_encoder is None and depth_encoder is None: + raise ValueError("Either rgb_encoder or depth_encoder must be provided") # Only re-encode if the videos are not already encoded with the given video encoding parameters for video_key in meta.video_keys: current_info = meta.info.features[video_key].get("info", {}) - current_encoder = VideoEncoderConfig.from_video_info(current_info) - if current_encoder != camera_encoder: - video_paths_list.extend((meta.root / VIDEO_DIR / video_key).rglob("*.mp4")) + current_encoder = encoder_config_from_video_info(current_info) + target_encoder = depth_encoder if video_key in meta.depth_keys else rgb_encoder + if target_encoder is None: + logging.info(f"No encoder provided for {video_key} video. Skipping re-encoding.") + elif current_encoder != target_encoder: + video_keys_paths_dict[video_key] = list((meta.root / VIDEO_DIR / video_key).rglob("*.mp4")) + video_keys_encoders_dict[video_key] = target_encoder else: - logging.info(f"{video_key} videos are already encoded with {camera_encoder}. Nothing to do.") + logging.info(f"{video_key} videos are already encoded with {target_encoder}. Nothing to do.") - if len(video_paths_list) == 0: + if len(video_keys_paths_dict) == 0: logging.warning("Dataset has no videos to re-encode.") return dataset - logging.info(f"Re-encoding {len(video_paths_list)} video file(s) with {camera_encoder}") + logging.info(f"Re-encoding {sum(len(paths) for paths in video_keys_paths_dict.values())} video file(s).") - worker_args = [(vp, camera_encoder, encoder_threads) for vp in video_paths_list] + worker_args = [ + (path, encoder, encoder_threads) + for video_key, encoder in video_keys_encoders_dict.items() + for path in video_keys_paths_dict[video_key] + ] if num_workers and num_workers > 1: with ProcessPoolExecutor(max_workers=num_workers) as pool: futures = [pool.submit(_reencode_video_worker, args) for args in worker_args] @@ -1966,10 +2002,15 @@ def reencode_dataset( for args in tqdm(worker_args, desc="Re-encoding videos"): _reencode_video_worker(args) - # Refresh video info in metadata for every video key. - for vid_key in meta.video_keys: - video_path = meta.root / meta.get_video_file_path(0, vid_key) - meta.info.features[vid_key]["info"] = get_video_info(video_path, camera_encoder=camera_encoder) + # Refresh video info in metadata for every re-encoded key. Re-encoding only + # changes codec/container params, so for depth videos we preserve ``is_depth_map`` + # and the depth quantization params (``video.depth_min`` / ``video.depth_max`` / + # ...), which describe the data rather than the codec and must survive a transcode. + # RGB videos pass an empty set: still a refresh, but nothing to preserve. + depth_preserve_keys = {"is_depth_map", *(f"video.{n}" for n in DEPTH_ENCODER_INFO_FIELD_NAMES)} + for video_key, encoder in video_keys_encoders_dict.items(): + preserve_keys = depth_preserve_keys if video_key in meta.depth_keys else set() + meta.update_video_info(video_key=video_key, video_encoder=encoder, preserve_keys=preserve_keys) write_info(meta.info, meta.root) logging.info("Dataset metadata updated.") diff --git a/src/lerobot/datasets/dataset_writer.py b/src/lerobot/datasets/dataset_writer.py index 633c00c1a..1aee1497c 100644 --- a/src/lerobot/datasets/dataset_writer.py +++ b/src/lerobot/datasets/dataset_writer.py @@ -31,7 +31,13 @@ import PIL.Image import pyarrow.parquet as pq import torch -from lerobot.configs import VideoEncoderConfig, camera_encoder_defaults +from lerobot.configs import ( + DepthEncoderConfig, + RGBEncoderConfig, + VideoEncoderConfig, + depth_encoder_defaults, + rgb_encoder_defaults, +) from .compute_stats import compute_episode_stats from .dataset_metadata import LeRobotDatasetMetadata @@ -48,6 +54,7 @@ from .io_utils import ( write_info, ) from .utils import ( + DEFAULT_DEPTH_PATH, DEFAULT_EPISODES_PATH, DEFAULT_IMAGE_PATH, update_chunk_file_indices, @@ -67,17 +74,22 @@ def _encode_video_worker( episode_index: int, root: Path, fps: int, - camera_encoder: VideoEncoderConfig | None = None, + video_encoder: VideoEncoderConfig | None = None, encoder_threads: int | None = None, ) -> Path: temp_path = Path(tempfile.mkdtemp(dir=root)) / f"{video_key}_{episode_index:03d}.mp4" - fpath = DEFAULT_IMAGE_PATH.format(image_key=video_key, episode_index=episode_index, frame_index=0) + path_template = ( + DEFAULT_DEPTH_PATH + if video_encoder is not None and isinstance(video_encoder, DepthEncoderConfig) + else DEFAULT_IMAGE_PATH + ) + fpath = path_template.format(image_key=video_key, episode_index=episode_index, frame_index=0) img_dir = (root / fpath).parent encode_video_frames( img_dir, temp_path, fps, - camera_encoder=camera_encoder, + video_encoder=video_encoder, encoder_threads=encoder_threads, overwrite=True, ) @@ -96,7 +108,8 @@ class DatasetWriter: self, meta: LeRobotDatasetMetadata, root: Path, - camera_encoder: VideoEncoderConfig | None, + rgb_encoder: RGBEncoderConfig | None, + depth_encoder: DepthEncoderConfig | None, encoder_threads: int | None, batch_encoding_size: int, streaming_encoder: StreamingVideoEncoder | None = None, @@ -108,8 +121,11 @@ class DatasetWriter: meta: Dataset metadata instance (used for feature schema, chunk settings, and episode persistence). root: Local dataset root directory. - camera_encoder: Video encoder settings applied to all cameras. - ``None`` uses :func:`~lerobot.configs.camera_encoder_defaults`. + rgb_encoder: Video encoder settings applied to RGB cameras. When + ``None``, :func:`~lerobot.configs.video.rgb_encoder_defaults` is used. + depth_encoder: Video encoder settings applied to depth cameras, including + the quantization parameters. When ``None``, + :func:`~lerobot.configs.video.depth_encoder_defaults` is used. encoder_threads: Number of encoder threads (global). ``None`` lets the codec decide. batch_encoding_size: Number of episodes to accumulate before @@ -120,7 +136,8 @@ class DatasetWriter: """ self._meta = meta self._root = root - self._camera_encoder = camera_encoder or camera_encoder_defaults() + self._rgb_encoder = rgb_encoder or rgb_encoder_defaults() + self._depth_encoder = depth_encoder or depth_encoder_defaults() self._encoder_threads = encoder_threads self._batch_encoding_size = batch_encoding_size self._streaming_encoder = streaming_encoder @@ -145,7 +162,8 @@ class DatasetWriter: return ep_buffer def _get_image_file_path(self, episode_index: int, image_key: str, frame_index: int) -> Path: - fpath = DEFAULT_IMAGE_PATH.format( + path_template = DEFAULT_DEPTH_PATH if image_key in self._meta.depth_keys else DEFAULT_IMAGE_PATH + fpath = path_template.format( image_key=image_key, episode_index=episode_index, frame_index=frame_index ) return self._root / fpath @@ -195,6 +213,7 @@ class DatasetWriter: if frame_index == 0 and self._streaming_encoder is not None: self._streaming_encoder.start_episode( video_keys=list(self._meta.video_keys), + depth_video_keys=list(self._meta.depth_keys), temp_dir=self._root, ) @@ -282,10 +301,13 @@ class DatasetWriter: if use_streaming: streaming_results = self._streaming_encoder.finish_episode() for video_key in self._meta.video_keys: + normalization_factor = 255.0 if video_key not in self._meta.depth_keys else 1.0 temp_path, video_stats = streaming_results[video_key] if video_stats is not None: ep_stats[video_key] = { - k: v if k == "count" else np.squeeze(v.reshape(1, -1, 1, 1) / 255.0, axis=0) + k: v + if k == "count" + else np.squeeze(v.reshape(1, -1, 1, 1) / normalization_factor, axis=0) for k, v in video_stats.items() } ep_metadata.update(self._save_episode_video(video_key, episode_index, temp_path=temp_path)) @@ -300,7 +322,7 @@ class DatasetWriter: episode_index, self._root, self._meta.fps, - self._camera_encoder, + self._depth_encoder if video_key in self._meta.depth_keys else self._rgb_encoder, self._encoder_threads, ): video_key for video_key in self._meta.video_keys @@ -511,7 +533,12 @@ class DatasetWriter: # Update video info (only needed when first episode is encoded) if episode_index == 0: - self._meta.update_video_info(video_key, camera_encoder=self._camera_encoder) + self._meta.update_video_info( + video_key, + video_encoder=self._depth_encoder + if video_key in self._meta.depth_keys + else self._rgb_encoder, + ) write_info(self._meta.info, self._meta.root) metadata = { @@ -578,13 +605,14 @@ class DatasetWriter: self.image_writer.wait_until_done() def _encode_temporary_episode_video(self, video_key: str, episode_index: int) -> Path: - """Use ffmpeg to convert frames stored as png into mp4 videos.""" + """Use ffmpeg to convert frames stored as png/tiff into mp4 videos.""" + is_depth = video_key in self._meta.depth_keys return _encode_video_worker( video_key, episode_index, self._root, self._meta.fps, - self._camera_encoder, + self._depth_encoder if is_depth else self._rgb_encoder, self._encoder_threads, ) diff --git a/src/lerobot/datasets/depth_utils.py b/src/lerobot/datasets/depth_utils.py new file mode 100644 index 000000000..801c86a09 --- /dev/null +++ b/src/lerobot/datasets/depth_utils.py @@ -0,0 +1,268 @@ +#!/usr/bin/env python + +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Depth encoding/decoding helpers for :class:`DepthEncoderConfig`. +""" + +import math +from typing import Literal + +import av +import numpy as np +import torch +from numpy.typing import NDArray + +from lerobot.configs.video import ( + DEFAULT_DEPTH_MAX, + DEFAULT_DEPTH_MIN, + DEFAULT_DEPTH_PIX_FMT, + DEFAULT_DEPTH_SHIFT, + DEFAULT_DEPTH_USE_LOG, + DEPTH_METER_UNIT, + DEPTH_MILLIMETER_UNIT, + DEPTH_QMAX, +) + +from .image_writer import squeeze_single_channel +from .pyav_utils import write_u16_plane + +_MM_PER_METRE = 1000.0 +_UINT16_MAX = 65535 + + +def _validate_log_quant_params(depth_min: float, shift: float) -> None: + """Ensure ``log(depth_min + shift)`` is finite.""" + if depth_min + shift <= 0: + raise ValueError( + f"depth_min + shift must be positive for logarithmic quantization, " + f"got depth_min={depth_min} + shift={shift} = {depth_min + shift}" + ) + + +def _depth_input_to_float32_and_unit( + depth: NDArray[np.integer] | NDArray[np.floating], + input_unit: Literal["auto", DEPTH_METER_UNIT, DEPTH_MILLIMETER_UNIT], +) -> tuple[NDArray[np.float32], Literal[DEPTH_METER_UNIT, DEPTH_MILLIMETER_UNIT]]: + """Convert depth to float32 in the chosen unit, and return the resolved unit.""" + resolved_unit = ( + (DEPTH_METER_UNIT if np.issubdtype(depth.dtype, np.floating) else DEPTH_MILLIMETER_UNIT) + if input_unit == "auto" + else input_unit + ) + return depth.astype(np.float32, order="K"), resolved_unit + + +def quantize_depth( + depth: NDArray[np.uint16] | NDArray[np.float32] | torch.Tensor, + depth_min: float = DEFAULT_DEPTH_MIN, + depth_max: float = DEFAULT_DEPTH_MAX, + shift: float = DEFAULT_DEPTH_SHIFT, + use_log: bool = DEFAULT_DEPTH_USE_LOG, + pix_fmt: str = DEFAULT_DEPTH_PIX_FMT, + video_backend: str | None = "pyav", + input_unit: Literal["auto", DEPTH_METER_UNIT, DEPTH_MILLIMETER_UNIT] = "auto", +) -> NDArray[np.uint16] | av.VideoFrame: + """Quantize depth to 12-bit codes (``uint16``, values ``0…DEPTH_QMAX``). + + Depth maps are packed into 12-bit integer frames so they fit in standard + high-bit-depth pixel formats (e.g. ``yuv420p12le`` / ``gray12le``) + and can be encoded by widely supported video codecs (e.g. HEVC Main 12). + Logarithmic quantization is the default because it allocates more quanta + to near-range depth, which matches the (1/depth) error profile of typical + depth sensors. Math is ported from BEHAVIOR-1K's ``obs_utils.py``. + + **Input units**: + + - ``input_unit="auto"`` (default): infer from dtype (floating = m, non-floating = mm). + - ``input_unit="mm"``: interpret input values as millimetres. + - ``input_unit="m"``: interpret input values as metres. + + Quantization math runs in the **resolved input unit**. + + ``depth_min``, ``depth_max``, and ``shift`` are always in **metres**. + + Args: + depth: Depth map; ``torch.Tensor`` is moved to CPU for conversion. + depth_min: Depth (metres) at quantum ``0``. + depth_max: Depth (metres) at quantum :data:`DEPTH_QMAX`. + shift: Depth shift (metres); used in log mode. Must satisfy ``depth_min + shift > 0``. + use_log: If ``True`` (default), quantize in log space. + video_backend: Video backend to use for encoding. Defaults to "pyav". + input_unit: Input unit policy (``"auto"``, ``"mm"``, ``"m"``). + + Returns: + ``numpy.ndarray``, ``dtype=uint16``, same shape as ``depth``, values in + ``[0, DEPTH_QMAX]``. + + Raises: + ValueError: If ``input_unit`` is not ``"auto"``, ``"mm"``, or ``"m"``. + ValueError: If ``use_log=True`` and ``depth_min + shift <= 0``. + """ + if input_unit not in ("auto", DEPTH_METER_UNIT, DEPTH_MILLIMETER_UNIT): + raise ValueError( + f"input_unit must be 'auto', '{DEPTH_METER_UNIT}', or '{DEPTH_MILLIMETER_UNIT}', got {input_unit!r}" + ) + + if isinstance(depth, torch.Tensor): + depth = depth.detach().cpu().numpy() + + # Squeeze single-channel dim: (H, W, 1) or (1, H, W) → (H, W) + depth = squeeze_single_channel(depth) + + depth_f, resolved_unit = _depth_input_to_float32_and_unit(depth, input_unit=input_unit) + + # Convert depth_min, depth_max, and shift to the resolved input unit. + depth_min_u = ( + np.float32(depth_min) if resolved_unit == DEPTH_METER_UNIT else np.float32(depth_min * _MM_PER_METRE) + ) + depth_max_u = ( + np.float32(depth_max) if resolved_unit == DEPTH_METER_UNIT else np.float32(depth_max * _MM_PER_METRE) + ) + shift_u = np.float32(shift) if resolved_unit == DEPTH_METER_UNIT else np.float32(shift * _MM_PER_METRE) + + # Normalization and quantization is performed in the resolved input unit. + if use_log: + _validate_log_quant_params(depth_min, shift) + log_min = math.log(float(depth_min_u + shift_u)) + log_max = math.log(float(depth_max_u + shift_u)) + norm = (np.log(depth_f + shift_u) - log_min) / (log_max - log_min) + else: + norm = (depth_f - depth_min_u) / (depth_max_u - depth_min_u) + + quantized = np.rint(norm * DEPTH_QMAX).clip(0, DEPTH_QMAX).astype(np.uint16, copy=False) + + if video_backend == "pyav": + frame = av.VideoFrame.from_ndarray(quantized, format=pix_fmt) + write_u16_plane(frame.planes[0], quantized) + return frame + else: + return quantized + + +def dequantize_depth( + quantized: NDArray[np.uint16] | av.VideoFrame | torch.Tensor, + depth_min: float = DEFAULT_DEPTH_MIN, + depth_max: float = DEFAULT_DEPTH_MAX, + shift: float = DEFAULT_DEPTH_SHIFT, + use_log: bool = DEFAULT_DEPTH_USE_LOG, + pix_fmt: str = DEFAULT_DEPTH_PIX_FMT, + output_unit: Literal[DEPTH_METER_UNIT, DEPTH_MILLIMETER_UNIT] = DEPTH_MILLIMETER_UNIT, + output_tensor: bool = True, + output_channel_last: bool = False, +) -> NDArray[np.uint16] | NDArray[np.float32] | torch.Tensor: + """Inverse of :func:`quantize_depth`. + + Decoding inverts the same normalized code mapping as :func:`quantize_depth` + using ``depth_min`` / ``depth_max`` / ``shift`` (in metres), then returns + the requested output unit. Tuning arguments **must match** :func:`quantize_depth`. + + Accepted input layouts : + + - ``(H, W, 1)`` or ``(H, W)`` — single frame with channel-last. + - ``(..., 1, H, W)`` — batched frames with channel-first. + - ``(..., H, W, 1)`` — batched frames with channel-last. + Output layout is determined by ``output_channel_last``. + + Args: + quantized: 12-bit codes in ``[0, DEPTH_QMAX]``. ``np.ndarray``, + ``av.VideoFrame``, or ``torch.Tensor`` (any integer or float dtype). + depth_min, depth_max, shift, use_log: Same as :func:`quantize_depth` (metres). + pix_fmt: Pixel format used to extract the plane from an ``av.VideoFrame``. + output_unit: ``"mm"`` returns ``uint16`` millimetres (rint, clip + ``[0, 65535]``) when returning a numpy array, or ``float32`` mm when + ``output_tensor=True``. ``"m"`` returns ``float32`` metres in + ``[depth_min, depth_max]``. + output_tensor: If True, return a ``torch.Tensor`` instead of a numpy array. + + Returns: + Depth map in the requested unit and dtype. + + Raises: + ValueError: If ``output_unit`` is not ``"m"`` or ``"mm"``. + ValueError: If ``use_log=True`` and ``depth_min + shift <= 0``. + """ + if output_unit not in (DEPTH_METER_UNIT, DEPTH_MILLIMETER_UNIT): + raise ValueError( + f"output_unit must be '{DEPTH_METER_UNIT}' or '{DEPTH_MILLIMETER_UNIT}', got {output_unit!r}" + ) + if use_log: + _validate_log_quant_params(depth_min, shift) + + if isinstance(quantized, av.VideoFrame): + quantized = quantized.to_ndarray(format=pix_fmt) + + # Compute the scale and offset first. + depth_min_m = float(depth_min) + depth_max_m = float(depth_max) + shift_m = float(shift) + if use_log: + log_min = math.log(depth_min_m + shift_m) + log_max = math.log(depth_max_m + shift_m) + scale = (log_max - log_min) / DEPTH_QMAX + offset = log_min + else: + scale = (depth_max_m - depth_min_m) / DEPTH_QMAX + offset = depth_min_m + + # ── Torch path: stay on the input device, single fp32 allocation. ──────── + if isinstance(quantized, torch.Tensor): + if quantized.ndim >= 3: + # Drop the single-channel dimension so the math runs on (..., H, W). + quantized = quantized.squeeze(-3) if quantized.shape[-3] == 1 else quantized.squeeze(-1) + + # Single allocation we own; everything else is in-place. + buf = quantized.to(dtype=torch.float32, copy=True) + buf.mul_(scale).add_(offset) + if use_log: + buf.exp_().sub_(shift_m) + buf.clamp_(depth_min_m, depth_max_m) + buf.unsqueeze_(-1) if output_channel_last else buf.unsqueeze_(-3) + + if output_unit == DEPTH_METER_UNIT: + return buf if output_tensor else buf.cpu().numpy() + + # mm path: round + clamp in float32, skipping the uint16 round-trip + # when returning a tensor (torch.uint16 is poorly supported). + buf.mul_(_MM_PER_METRE).round_().clamp_(0.0, _UINT16_MAX) + if output_tensor: + return buf + return buf.cpu().numpy().astype(np.uint16, copy=False) + + # ── NumPy path: single fp32 allocation, ``out=`` for in-place math. ───── + arr = np.asarray(quantized) + if arr.ndim >= 3: + # Drop the single-channel dimension so the math runs on (..., H, W). + arr = np.squeeze(arr, axis=-3) if arr.shape[-3] == 1 else np.squeeze(arr, axis=-1) + + buf = np.empty(arr.shape, dtype=np.float32) + np.multiply(arr, scale, out=buf) + np.add(buf, offset, out=buf) + if use_log: + np.exp(buf, out=buf) + np.subtract(buf, shift_m, out=buf) + np.clip(buf, depth_min_m, depth_max_m, out=buf) + buf = np.expand_dims(buf, axis=-1) if output_channel_last else np.expand_dims(buf, axis=-3) + + if output_unit == DEPTH_METER_UNIT: + return torch.from_numpy(buf) if output_tensor else buf + + np.multiply(buf, _MM_PER_METRE, out=buf) + np.rint(buf, out=buf) + np.clip(buf, 0.0, _UINT16_MAX, out=buf) + if output_tensor: + # torch.uint16 support is very limited; return float32 millimetres. + return torch.from_numpy(buf) + return buf.astype(np.uint16, copy=False) diff --git a/src/lerobot/datasets/factory.py b/src/lerobot/datasets/factory.py index cbbe83dc8..da7b4365a 100644 --- a/src/lerobot/datasets/factory.py +++ b/src/lerobot/datasets/factory.py @@ -14,6 +14,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import logging +import math from pprint import pformat import torch @@ -96,6 +97,7 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas revision=cfg.dataset.revision, video_backend=cfg.dataset.video_backend, return_uint8=True, + depth_output_unit=cfg.dataset.depth_output_unit, tolerance_s=cfg.tolerance_s, ) else: @@ -126,7 +128,87 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas if cfg.dataset.use_imagenet_stats: for key in dataset.meta.camera_keys: + if key in dataset.meta.depth_keys: + continue # Exclude depth keys from ImageNet stats for stats_type, stats in IMAGENET_STATS.items(): dataset.meta.stats[key][stats_type] = torch.tensor(stats, dtype=torch.float32) return dataset + + +def make_train_eval_datasets( + cfg: TrainPipelineConfig, +) -> tuple[LeRobotDataset | MultiLeRobotDataset, LeRobotDataset | None]: + """Create train and optional eval datasets by splitting episodes based on eval_split. + + The last ceil(n_episodes * eval_split) episodes per task are held out for evaluation. + If eval_split == 0.0, returns (full_dataset, None). + """ + full_dataset = make_dataset(cfg) + + if cfg.dataset.eval_split == 0.0: + return full_dataset, None + + base_episodes = ( + full_dataset.episodes if full_dataset.episodes is not None else list(range(full_dataset.num_episodes)) + ) + + episode_tasks = full_dataset.meta.episodes["tasks"] + task_to_episodes: dict[str, list[int]] = {} + for ep_idx in base_episodes: + task_key = episode_tasks[ep_idx][0] if episode_tasks[ep_idx] else "" + task_to_episodes.setdefault(task_key, []).append(ep_idx) + + train_episodes, eval_episodes = [], [] + for eps in task_to_episodes.values(): + n_eval = math.ceil(len(eps) * cfg.dataset.eval_split) + train_episodes.extend(eps[: len(eps) - n_eval]) + eval_episodes.extend(eps[len(eps) - n_eval :]) + + if not train_episodes: + raise ValueError( + f"eval_split={cfg.dataset.eval_split} leaves 0 training episodes from {len(base_episodes)} total." + ) + + logging.info( + f"Train/eval split: {len(train_episodes)} train, {len(eval_episodes)} eval " + f"(eval_split={cfg.dataset.eval_split}, {len(task_to_episodes)} tasks)" + ) + + delta_timestamps = resolve_delta_timestamps(cfg.trainable_config, full_dataset.meta) + + train_image_transforms = ( + ImageTransforms(cfg.dataset.image_transforms) if cfg.dataset.image_transforms.enable else None + ) + + train_dataset = LeRobotDataset( + cfg.dataset.repo_id, + root=cfg.dataset.root, + episodes=train_episodes, + delta_timestamps=delta_timestamps, + image_transforms=train_image_transforms, + revision=cfg.dataset.revision, + video_backend=cfg.dataset.video_backend, + return_uint8=True, + tolerance_s=cfg.tolerance_s, + ) + + eval_dataset = LeRobotDataset( + cfg.dataset.repo_id, + root=cfg.dataset.root, + episodes=eval_episodes, + delta_timestamps=delta_timestamps, + image_transforms=None, + revision=cfg.dataset.revision, + video_backend=cfg.dataset.video_backend, + return_uint8=True, + tolerance_s=cfg.tolerance_s, + ) + + if cfg.dataset.use_imagenet_stats: + for ds in (train_dataset, eval_dataset): + for key in ds.meta.camera_keys: + for stats_type, stats in IMAGENET_STATS.items(): + ds.meta.stats[key][stats_type] = torch.tensor(stats, dtype=torch.float32) + + return train_dataset, eval_dataset diff --git a/src/lerobot/datasets/feature_utils.py b/src/lerobot/datasets/feature_utils.py index 56264408f..343b2fdcc 100644 --- a/src/lerobot/datasets/feature_utils.py +++ b/src/lerobot/datasets/feature_utils.py @@ -336,7 +336,7 @@ def validate_feature_image_or_video( Args: name (str): The name of the feature. - expected_shape (list[str]): The expected shape (C, H, W). + expected_shape (list[str]): The expected shape, e.g. (C, H, W) or (H, W, C). value: The image data to validate. Returns: diff --git a/src/lerobot/datasets/image_writer.py b/src/lerobot/datasets/image_writer.py index 8fb5804a5..41790b46a 100644 --- a/src/lerobot/datasets/image_writer.py +++ b/src/lerobot/datasets/image_writer.py @@ -41,11 +41,51 @@ def safe_stop_image_writer(func): return wrapper -def image_array_to_pil_image(image_array: np.ndarray, range_check: bool = True) -> PIL.Image.Image: - # TODO(aliberts): handle 1 channel and 4 for depth images - if image_array.ndim != 3: - raise ValueError(f"The array has {image_array.ndim} dimensions, but 3 is expected for an image.") +def squeeze_single_channel(array: np.ndarray) -> np.ndarray: + """Drop a leading or trailing singleton channel dim: ``(1, H, W)`` / ``(H, W, 1)`` -> ``(H, W)``. + Unlike ``array.squeeze()``, this only removes the channel axis, never an ``H`` or ``W`` of size 1. + """ + if array.ndim == 3: + if array.shape[0] == 1: + return array[0] + if array.shape[-1] == 1: + return array[..., 0] + return array + + +def image_array_to_pil_image(image_array: np.ndarray, range_check: bool = True) -> PIL.Image.Image: + """Convert a NumPy array to a PIL Image, preserving precision for grayscale. + + Behaviour by shape: + + - ``(H, W)`` or ``(1, H, W)`` / ``(H, W, 1)``: single-channel grayscale. + The native dtype is preserved using the matching PIL mode + (``I;16`` / ``F``). This is the path used for raw depth maps (no rescaling, clamping, or downcasting) + - ``(3, H, W)`` / ``(H, W, 3)``: RGB. Channels-first inputs are transposed + to channels-last. Float inputs in ``[0, 1]`` are scaled to ``uint8`` + (existing behaviour, gated by ``range_check``). + + Other shapes / channel counts raise ``NotImplementedError`` or + ``ValueError``. + """ + # TODO(CarolinePascal): 4 dimensions RGB-D images + if image_array.ndim not in (2, 3): + raise ValueError(f"The array has {image_array.ndim} dimensions, but 2 or 3 is expected for an image.") + + # Squeeze 3D single-channel inputs to 2D so depth maps work whether the + # caller emits (H, W), (1, H, W), or (H, W, 1). + image_array = squeeze_single_channel(image_array) + + if image_array.ndim == 2: + if image_array.dtype not in [np.uint16, np.float32]: + raise ValueError( + f"Unsupported single-channel image dtype: {image_array.dtype}. " + f"Supported dtypes: {sorted(str(d) for d in [np.uint16, np.float32])}." + ) + return PIL.Image.fromarray(np.ascontiguousarray(image_array)) + + # 3D path: must be RGB (3 channels), channels-first or channels-last. if image_array.shape[0] == 3: # Transpose from pytorch convention (C, H, W) to (H, W, C) image_array = image_array.transpose(1, 2, 0) @@ -71,13 +111,29 @@ def image_array_to_pil_image(image_array: np.ndarray, range_check: bool = True) return PIL.Image.fromarray(image_array) +def save_kwargs_for_path(fpath: Path, compress_level: int) -> dict: + """Pick the right format-specific kwargs for :meth:`PIL.Image.Image.save`. + + PNG uses ``compress_level`` (0-9, zlib). TIFF uses ``compression`` (raw) for lossless raw depth maps. + """ + suffix = Path(fpath).suffix.lower() + if suffix == ".png": + return {"compress_level": compress_level} + if suffix in (".tif", ".tiff"): + return {"compression": "raw"} + else: + raise ValueError(f"Unsupported image file extension: {suffix}") + + def write_image(image: np.ndarray | PIL.Image.Image, fpath: Path, compress_level: int = 1): """ Saves a NumPy array or PIL Image to a file. This function handles both NumPy arrays and PIL Image objects, converting the former to a PIL Image before saving. It includes error handling for - the save operation. + the save operation. The output format is inferred from the *fpath* + extension: ``.png`` → PNG with ``compress_level``, ``.tiff`` / ``.tif`` + → lossless raw depth maps (TIFF). Args: image (np.ndarray | PIL.Image.Image): The image data to save. @@ -101,7 +157,7 @@ def write_image(image: np.ndarray | PIL.Image.Image, fpath: Path, compress_level img = image else: raise TypeError(f"Unsupported image type: {type(image)}") - img.save(fpath, compress_level=compress_level) + img.save(fpath, **save_kwargs_for_path(fpath, compress_level)) except Exception as e: logger.error("Error writing image %s: %s", fpath, e) diff --git a/src/lerobot/datasets/io_utils.py b/src/lerobot/datasets/io_utils.py index b6344942c..868a114f5 100644 --- a/src/lerobot/datasets/io_utils.py +++ b/src/lerobot/datasets/io_utils.py @@ -154,7 +154,7 @@ def cast_stats_to_numpy(stats: dict) -> dict[str, dict[str, np.ndarray]]: Returns: dict: The statistics dictionary with values cast to numpy arrays. """ - stats = {key: np.array(value) for key, value in flatten_dict(stats).items()} + stats = {key: np.atleast_1d(np.array(value)) for key, value in flatten_dict(stats).items()} return unflatten_dict(stats) @@ -226,28 +226,50 @@ def load_image_as_numpy( Args: fpath (str | Path): Path to the image file. dtype (np.dtype): The desired data type of the output array. If floating, - pixels are scaled to [0, 1]. + pixels are scaled to [0, 1]. Only used for RGB images. channel_first (bool): If True, converts the image to (C, H, W) format. Otherwise, it remains in (H, W, C) format. Returns: np.ndarray: The image as a numpy array. """ - img = PILImage.open(fpath).convert("RGB") - img_array = np.array(img, dtype=dtype) + is_depth = fpath.endswith(".tiff") or fpath.endswith(".tif") + if is_depth: + # Preserve the native depth dtype (uint16 -> "I;16", float32 -> "F"). + img = PILImage.open(fpath) + img_array = np.array(img) + else: + img = PILImage.open(fpath).convert("RGB") + img_array = np.array(img, dtype=dtype) + if np.issubdtype(dtype, np.floating): + img_array /= 255.0 if channel_first: # (H, W, C) -> (C, H, W) - img_array = np.transpose(img_array, (2, 0, 1)) - if np.issubdtype(dtype, np.floating): - img_array /= 255.0 + img_array = img_array[np.newaxis, ...] if img_array.ndim == 2 else np.transpose(img_array, (2, 0, 1)) return img_array +# PIL modes for 16-bit unsigned depth maps. +UINT16_PIL_MODES = {"I;16", "I;16B", "I;16L"} + + +def pil_to_chw_tensor(img: PILImage.Image) -> torch.Tensor: + """Convert a PIL image to a channel-first tensor. + + ``uint16`` depth maps become ``float32 (1, H, W)`` in native units (``ToTensor`` + would overflow them to ``int16``); all other modes use the standard ``ToTensor`` path. + """ + if img.mode in UINT16_PIL_MODES: + return torch.from_numpy(np.array(img, dtype=np.float32))[None, ...] + return transforms.ToTensor()(img) + + def hf_transform_to_torch(items_dict: dict[str, list[Any]]) -> dict[str, list[torch.Tensor | str]]: """Convert a batch from a Hugging Face dataset to torch tensors. This transform function converts items from Hugging Face dataset format (pyarrow) - to torch tensors. Importantly, images are converted from PIL objects (H, W, C, uint8) - to a torch image representation (C, H, W, float32) in the range [0, 1]. Other + to torch tensors. RGB images are converted from PIL objects (H, W, C, uint8) + to a torch image representation (C, H, W, float32) in the range [0, 1]. Depth + maps are returned as float32 (1, H, W) in their native units. Other types are converted to torch.tensor. Args: @@ -262,8 +284,7 @@ def hf_transform_to_torch(items_dict: dict[str, list[Any]]) -> dict[str, list[to continue first_item = items_dict[key][0] if isinstance(first_item, PILImage.Image): - to_tensor = transforms.ToTensor() - items_dict[key] = [to_tensor(img) for img in items_dict[key]] + items_dict[key] = [pil_to_chw_tensor(img) for img in items_dict[key]] elif first_item is None or isinstance(first_item, dict): pass else: @@ -329,7 +350,11 @@ def item_to_torch(item: dict) -> dict: """ skip_keys = {"task", *LANGUAGE_COLUMNS} for key, val in item.items(): - if isinstance(val, (np.ndarray | list)) and key not in skip_keys: + if key in skip_keys: + continue + if isinstance(val, PILImage.Image): + item[key] = pil_to_chw_tensor(val) + elif isinstance(val, (np.ndarray | list)): # Convert numpy arrays and lists to torch tensors item[key] = torch.tensor(val) return item diff --git a/src/lerobot/datasets/lerobot_dataset.py b/src/lerobot/datasets/lerobot_dataset.py index d1e65fef1..f600f1804 100644 --- a/src/lerobot/datasets/lerobot_dataset.py +++ b/src/lerobot/datasets/lerobot_dataset.py @@ -24,7 +24,7 @@ import torch.utils from huggingface_hub import HfApi, snapshot_download from huggingface_hub.errors import RevisionNotFoundError -from lerobot.configs import VideoEncoderConfig +from lerobot.configs import DEFAULT_DEPTH_UNIT, DepthEncoderConfig, RGBEncoderConfig from lerobot.utils.constants import HF_LEROBOT_HUB_CACHE from .dataset_metadata import CODEBASE_VERSION, LeRobotDatasetMetadata @@ -58,8 +58,10 @@ class LeRobotDataset(torch.utils.data.Dataset): download_videos: bool = True, video_backend: str | None = None, return_uint8: bool = False, + depth_output_unit: str = DEFAULT_DEPTH_UNIT, batch_encoding_size: int = 1, - camera_encoder: VideoEncoderConfig | None = None, + rgb_encoder: RGBEncoderConfig | None = None, + depth_encoder: DepthEncoderConfig | None = None, encoder_threads: int | None = None, streaming_encoding: bool = False, encoder_queue_maxsize: int = 30, @@ -183,8 +185,11 @@ class LeRobotDataset(torch.utils.data.Dataset): You can also use the 'pyav' decoder used by Torchvision, which used to be the default option, or 'video_reader' which is another decoder of Torchvision. batch_encoding_size (int, optional): Number of episodes to accumulate before batch encoding videos. Set to 1 for immediate encoding (default), or higher for batched encoding. Defaults to 1. - camera_encoder (VideoEncoderConfig | None, optional): Video encoder settings for cameras - (codec, quality, etc.). When ``None``, :func:`~lerobot.configs.video.camera_encoder_defaults` + rgb_encoder (RGBEncoderConfig | None, optional): Video encoder settings for cameras + (codec, quality, etc.). When ``None``, :func:`~lerobot.configs.video.rgb_encoder_defaults` + is used by the writer. + depth_encoder (DepthEncoderConfig | None, optional): Video encoder settings for depth cameras + (codec, quality, etc.). When ``None``, :func:`~lerobot.configs.video.depth_encoder_defaults` is used by the writer. encoder_threads (int | None, optional): Number of encoder threads (global). ``None`` lets the codec decide. @@ -206,6 +211,7 @@ class LeRobotDataset(torch.utils.data.Dataset): self.revision = revision if revision else CODEBASE_VERSION self._video_backend = video_backend if video_backend else get_safe_default_video_backend() self._return_uint8 = return_uint8 + self._depth_output_unit = depth_output_unit self._batch_encoding_size = batch_encoding_size self._encoder_threads = encoder_threads @@ -246,6 +252,7 @@ class LeRobotDataset(torch.utils.data.Dataset): delta_timestamps=delta_timestamps, image_transforms=image_transforms, return_uint8=self._return_uint8, + depth_output_unit=self._depth_output_unit, ) self.image_transforms = image_transforms @@ -271,14 +278,16 @@ class LeRobotDataset(torch.utils.data.Dataset): if streaming_encoding and len(self.meta.video_keys) > 0: streaming_enc = self._build_streaming_encoder( self.meta.fps, - camera_encoder, + rgb_encoder, + depth_encoder, encoder_queue_maxsize, encoder_threads, ) self.writer = DatasetWriter( meta=self.meta, root=self.root, - camera_encoder=camera_encoder, + rgb_encoder=rgb_encoder, + depth_encoder=depth_encoder, encoder_threads=encoder_threads, batch_encoding_size=batch_encoding_size, streaming_encoder=streaming_enc, @@ -314,19 +323,22 @@ class LeRobotDataset(torch.utils.data.Dataset): delta_timestamps=self.delta_timestamps, image_transforms=self.image_transforms, return_uint8=self._return_uint8, + depth_output_unit=self._depth_output_unit, ) return self.reader @staticmethod def _build_streaming_encoder( fps: int, - camera_encoder: VideoEncoderConfig | None, + rgb_encoder: RGBEncoderConfig | None, + depth_encoder: DepthEncoderConfig | None, encoder_queue_maxsize: int, encoder_threads: int | None, ) -> StreamingVideoEncoder: return StreamingVideoEncoder( fps=fps, - camera_encoder=camera_encoder, + rgb_encoder=rgb_encoder, + depth_encoder=depth_encoder, queue_maxsize=encoder_queue_maxsize, encoder_threads=encoder_threads, ) @@ -369,6 +381,18 @@ class LeRobotDataset(torch.utils.data.Dataset): self.reader.load_and_activate() return self.reader.hf_dataset + @property + def absolute_to_relative_idx(self) -> dict[int, int] | None: + """Mapping from absolute frame indices to HF dataset row positions. + + Non-None only for episode-filtered datasets where absolute indices + (from metadata) differ from row positions in the loaded HF dataset. + """ + reader = self._ensure_reader() + if reader.hf_dataset is None: + reader.load_and_activate() + return reader._absolute_to_relative_idx + # ── Writer-delegated methods ────────────────────────────────────── def add_frame(self, frame: dict) -> None: @@ -643,7 +667,8 @@ class LeRobotDataset(torch.utils.data.Dataset): image_writer_threads: int = 0, video_backend: str | None = None, batch_encoding_size: int = 1, - camera_encoder: VideoEncoderConfig | None = None, + rgb_encoder: RGBEncoderConfig | None = None, + depth_encoder: DepthEncoderConfig | None = None, metadata_buffer_size: int = 10, streaming_encoding: bool = False, encoder_queue_maxsize: int = 30, @@ -674,8 +699,10 @@ class LeRobotDataset(torch.utils.data.Dataset): video_backend: Video decoding backend (used when reading back). batch_encoding_size: Number of episodes to accumulate before batch-encoding videos. ``1`` means encode immediately. - camera_encoder: Video encoder settings for cameras (codec, quality, etc.). - When ``None``, :func:`~lerobot.configs.video.camera_encoder_defaults` is used. + rgb_encoder: Video encoder settings for cameras (codec, quality, etc.). + When ``None``, :func:`~lerobot.configs.video.rgb_encoder_defaults` is used. + depth_encoder: Video encoder settings for depth cameras (codec, quality, etc.). + When ``None``, :func:`~lerobot.configs.video.depth_encoder_defaults` is used. encoder_threads: Number of encoder threads (global). ``None`` lets the codec decide. metadata_buffer_size: Number of episode metadata records to buffer @@ -710,6 +737,7 @@ class LeRobotDataset(torch.utils.data.Dataset): obj.episodes = None obj._video_backend = video_backend if video_backend is not None else get_safe_default_video_backend() obj._return_uint8 = False + obj._depth_output_unit = DEFAULT_DEPTH_UNIT obj._batch_encoding_size = batch_encoding_size obj._encoder_threads = encoder_threads @@ -719,12 +747,13 @@ class LeRobotDataset(torch.utils.data.Dataset): streaming_enc = None if streaming_encoding and len(obj.meta.video_keys) > 0: streaming_enc = cls._build_streaming_encoder( - fps, camera_encoder, encoder_queue_maxsize, encoder_threads + fps, rgb_encoder, depth_encoder, encoder_queue_maxsize, encoder_threads ) obj.writer = DatasetWriter( meta=obj.meta, root=obj.root, - camera_encoder=camera_encoder, + rgb_encoder=rgb_encoder, + depth_encoder=depth_encoder, encoder_threads=encoder_threads, batch_encoding_size=batch_encoding_size, streaming_encoder=streaming_enc, @@ -747,7 +776,8 @@ class LeRobotDataset(torch.utils.data.Dataset): force_cache_sync: bool = False, video_backend: str | None = None, batch_encoding_size: int = 1, - camera_encoder: VideoEncoderConfig | None = None, + rgb_encoder: RGBEncoderConfig | None = None, + depth_encoder: DepthEncoderConfig | None = None, encoder_threads: int | None = None, image_writer_processes: int = 0, image_writer_threads: int = 0, @@ -775,8 +805,10 @@ class LeRobotDataset(torch.utils.data.Dataset): video_backend: Video decoding backend for reading back data. batch_encoding_size: Number of episodes to accumulate before batch-encoding videos. - camera_encoder: Video encoder settings for cameras (codec, quality, etc.). - When ``None``, :func:`~lerobot.configs.video.camera_encoder_defaults` is used. + rgb_encoder: Video encoder settings for cameras (codec, quality, etc.). + When ``None``, :func:`~lerobot.configs.video.rgb_encoder_defaults` is used. + depth_encoder: Video encoder settings for depth cameras (codec, quality, etc.). + When ``None``, :func:`~lerobot.configs.video.depth_encoder_defaults` is used. encoder_threads: Number of encoder threads (global). ``None`` lets the codec decide. image_writer_processes: Subprocesses for async image writing. @@ -804,6 +836,7 @@ class LeRobotDataset(torch.utils.data.Dataset): obj.episodes = None obj._video_backend = video_backend if video_backend else get_safe_default_video_backend() obj._return_uint8 = False + obj._depth_output_unit = DEFAULT_DEPTH_UNIT obj._batch_encoding_size = batch_encoding_size if obj._requested_root is not None: @@ -823,12 +856,13 @@ class LeRobotDataset(torch.utils.data.Dataset): streaming_enc = None if streaming_encoding and len(obj.meta.video_keys) > 0: streaming_enc = cls._build_streaming_encoder( - obj.meta.fps, camera_encoder, encoder_queue_maxsize, encoder_threads + obj.meta.fps, rgb_encoder, depth_encoder, encoder_queue_maxsize, encoder_threads ) obj.writer = DatasetWriter( meta=obj.meta, root=obj.root, - camera_encoder=camera_encoder, + rgb_encoder=rgb_encoder, + depth_encoder=depth_encoder, encoder_threads=encoder_threads, batch_encoding_size=batch_encoding_size, streaming_encoder=streaming_enc, diff --git a/src/lerobot/datasets/pyav_utils.py b/src/lerobot/datasets/pyav_utils.py index d291f8b40..7b7d1e5de 100644 --- a/src/lerobot/datasets/pyav_utils.py +++ b/src/lerobot/datasets/pyav_utils.py @@ -24,6 +24,7 @@ import logging from typing import Any import av +import numpy as np logger = logging.getLogger(__name__) @@ -31,6 +32,34 @@ FFMPEG_NUMERIC_OPTION_TYPES = ("INT", "INT64", "UINT64", "FLOAT", "DOUBLE") FFMPEG_INTEGER_OPTION_TYPES = ("INT", "INT64", "UINT64") +def write_u16_plane(plane: av.video.plane.VideoPlane, src: np.ndarray, fill_value: int | None = None) -> None: + """Copy a 2D ``uint16`` image into the plane's memory buffer, row by row. + + For speed, each row is padded to a wider size than ``width``, so the true row width in + memory is ``plane.line_size`` (bytes), not ``width``. Copying as one straight stream + would skew the image, so we write only the first ``width`` columns of each row and + leave the padding untouched. + + Args: + plane: Destination 16-bit plane. + src: Source image, shape ``(height, width)``, dtype ``uint16``. + fill_value: If given, every pixel (padding included) is set to this first, so the + padding holds clean data instead of garbage. + """ + height, width = src.shape + stride_u16 = plane.line_size // np.dtype(np.uint16).itemsize + dst = np.frombuffer(plane, dtype=np.uint16).reshape(height, stride_u16) + if fill_value is not None: + dst.fill(fill_value) + dst[:, :width] = src + + +@functools.cache +def get_pix_fmt_channels(pix_fmt: str) -> int: + """Return the number of components (channels) for *pix_fmt*.""" + return len(av.VideoFormat(pix_fmt).components) + + @functools.cache def get_codec(vcodec: str) -> av.codec.Codec | None: """PyAV write-mode ``Codec`` for *vcodec*, or ``None`` if unavailable.""" @@ -92,7 +121,7 @@ def _check_option_value(vcodec: str, label: str, value: Any, opt: av.option.Opti f"{label}={value!r} is not numeric; codec {vcodec!r} expects a number for this option." ) from e elif isinstance(value, (float, int)): - num_val = value + num_val = float(value) else: raise ValueError( f"{label}={value!r} is not numeric; codec {vcodec!r} expects a number for this option." @@ -142,6 +171,16 @@ def _check_pixel_format(vcodec: str, pix_fmt: str) -> None: ) +def _check_pix_fmt_channels(pix_fmt: str, channels: int) -> None: + """Ensure *pix_fmt* can carry at least *channels* components.""" + pix_fmt_channels = get_pix_fmt_channels(pix_fmt) + if pix_fmt_channels < channels: + raise ValueError( + f"pix_fmt={pix_fmt!r} carries only {pix_fmt_channels} component(s) " + f"but the source data has {channels} channel(s)." + ) + + def _check_codec_options(vcodec: str, codec_options: dict[str, Any]) -> None: """Validate merged encoder options (typed) against the codec's published AVOptions.""" supported_options = _get_codec_options_by_name(vcodec) @@ -156,12 +195,18 @@ def _check_codec_options(vcodec: str, codec_options: dict[str, Any]) -> None: _check_option_value(vcodec, key, value, supported_options[key]) -def check_video_encoder_parameters_pyav(vcodec: str, pix_fmt: str, codec_options: dict[str, Any]) -> None: +def check_video_encoder_parameters_pyav( + vcodec: str, + pix_fmt: str, + codec_options: dict[str, Any], + channels: int | None = None, +) -> None: """Verify *config* is compatible with the bundled FFmpeg build. Checks pixel format, abstract tuning-field compatibility, and each merged encoder option from :meth:`~lerobot.configs.video.VideoEncoderConfig.get_codec_options` against PyAV (including numeric ``extra_options`` present in that dict). + When given, additionally verify that *pix_fmt* carries as many components as the source data channels. No-op when ``config.vcodec`` isn't in the local FFmpeg build. Raises: @@ -171,4 +216,6 @@ def check_video_encoder_parameters_pyav(vcodec: str, pix_fmt: str, codec_options if not options: raise ValueError(f"Codec {vcodec!r} is not available in the bundled FFmpeg build") _check_pixel_format(vcodec, pix_fmt) + if channels is not None: + _check_pix_fmt_channels(pix_fmt, channels) _check_codec_options(vcodec, codec_options) diff --git a/src/lerobot/datasets/sampler.py b/src/lerobot/datasets/sampler.py index af85dff9b..aee6ce46d 100644 --- a/src/lerobot/datasets/sampler.py +++ b/src/lerobot/datasets/sampler.py @@ -53,6 +53,7 @@ class EpisodeAwareSampler: drop_n_last_frames: int = 0, shuffle: bool = False, seed: int = 0, + absolute_to_relative_idx: dict[int, int] | None = None, ): """ Args: @@ -107,6 +108,7 @@ class EpisodeAwareSampler: self.seed = seed self._epoch = 0 self._start_index = 0 + self._absolute_to_relative = absolute_to_relative_idx @property def indices(self) -> list[int]: @@ -132,7 +134,10 @@ class EpisodeAwareSampler: def _frame_index(self, position: int) -> int: episode = int(np.searchsorted(self._cum_lengths, position, side="right")) position_in_episode = position - (int(self._cum_lengths[episode - 1]) if episode > 0 else 0) - return int(self._starts[episode]) + position_in_episode + absolute_idx = int(self._starts[episode]) + position_in_episode + if self._absolute_to_relative is not None: + return self._absolute_to_relative[absolute_idx] + return absolute_idx def __iter__(self) -> Iterator[int]: # Advance epoch state eagerly, not on first consumption of the generator. diff --git a/src/lerobot/datasets/streaming_dataset.py b/src/lerobot/datasets/streaming_dataset.py index 3c1e4a73c..4c4ae59bf 100644 --- a/src/lerobot/datasets/streaming_dataset.py +++ b/src/lerobot/datasets/streaming_dataset.py @@ -22,9 +22,11 @@ import numpy as np import torch from datasets import load_dataset +from lerobot.configs import DEFAULT_DEPTH_UNIT, DepthEncoderConfig from lerobot.utils.constants import HF_LEROBOT_HOME, LOOKAHEAD_BACKTRACKTABLE, LOOKBACK_BACKTRACKTABLE from .dataset_metadata import CODEBASE_VERSION, LeRobotDatasetMetadata +from .depth_utils import dequantize_depth from .feature_utils import get_delta_indices from .io_utils import item_to_torch from .utils import ( @@ -35,6 +37,7 @@ from .utils import ( ) from .video_utils import ( VideoDecoderCache, + decode_video_frames, decode_video_frames_torchcodec, ) @@ -252,6 +255,7 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): rng: np.random.Generator | None = None, shuffle: bool = True, return_uint8: bool = False, + depth_output_unit: str = DEFAULT_DEPTH_UNIT, ): """Initialize a StreamingLeRobotDataset. @@ -272,6 +276,8 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): seed (int, optional): Reproducibility random seed. rng (np.random.Generator | None, optional): Random number generator. shuffle (bool, optional): Whether to shuffle the dataset across exhaustions. Defaults to True. + depth_output_unit (str, optional): Physical unit depth maps are dequantized to ("m" or "mm"). + Defaults to "mm". """ super().__init__() self.repo_id = repo_id @@ -290,6 +296,7 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): self.streaming = streaming self.buffer_size = buffer_size self._return_uint8 = return_uint8 + self._depth_output_unit = depth_output_unit # We cache the video decoders to avoid re-initializing them at each frame (avoiding a ~10x slowdown) self.video_decoder_cache = None @@ -306,6 +313,11 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): # Check version check_version_compatibility(self.repo_id, self.meta._version, CODEBASE_VERSION) + self._depth_encoder_configs: dict[str, DepthEncoderConfig] = { + vid_key: DepthEncoderConfig.from_video_info(self.meta.features[vid_key].get("info")) + for vid_key in self.meta.depth_keys + } + self.delta_timestamps = None self.delta_indices = None @@ -554,13 +566,34 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): for video_key, query_ts in query_timestamps.items(): root = self.meta.url_root if self.streaming and not self.streaming_from_local else self.root video_path = f"{root}/{self.meta.get_video_file_path(ep_idx, video_key)}" - frames = decode_video_frames_torchcodec( - video_path, - query_ts, - self.tolerance_s, - decoder_cache=self.video_decoder_cache, - return_uint8=self._return_uint8, - ) + if video_key in self.meta.depth_keys: + # Depth maps are 12-bit quantized and only decodable via pyav; dequantize back + # to physical units to match the non-streaming reader. + frames = decode_video_frames( + video_path, + query_ts, + self.tolerance_s, + backend="pyav", + return_uint8=False, + is_depth=True, + ) + depth_encoder = self._depth_encoder_configs[video_key] + frames = dequantize_depth( + frames, + depth_min=depth_encoder.depth_min, + depth_max=depth_encoder.depth_max, + shift=depth_encoder.shift, + use_log=depth_encoder.use_log, + output_unit=self._depth_output_unit, + ) + else: + frames = decode_video_frames_torchcodec( + video_path, + query_ts, + self.tolerance_s, + decoder_cache=self.video_decoder_cache, + return_uint8=self._return_uint8, + ) item[video_key] = frames.squeeze(0) if len(query_ts) == 1 else frames diff --git a/src/lerobot/datasets/utils.py b/src/lerobot/datasets/utils.py index de91978ea..d30761515 100644 --- a/src/lerobot/datasets/utils.py +++ b/src/lerobot/datasets/utils.py @@ -87,11 +87,14 @@ DATA_DIR = "data" VIDEO_DIR = "videos" CHUNK_FILE_PATTERN = "chunk-{chunk_index:03d}/file-{file_index:03d}" +IMAGE_FILE_PATTERN = "frame-{frame_index:06d}.png" +DEPTH_FILE_PATTERN = "frame-{frame_index:06d}.tiff" DEFAULT_TASKS_PATH = "meta/tasks.parquet" DEFAULT_EPISODES_PATH = EPISODES_DIR + "/" + CHUNK_FILE_PATTERN + ".parquet" DEFAULT_DATA_PATH = DATA_DIR + "/" + CHUNK_FILE_PATTERN + ".parquet" DEFAULT_VIDEO_PATH = VIDEO_DIR + "/{video_key}/" + CHUNK_FILE_PATTERN + ".mp4" -DEFAULT_IMAGE_PATH = "images/{image_key}/episode-{episode_index:06d}/frame-{frame_index:06d}.png" +DEFAULT_IMAGE_PATH = "images/{image_key}/episode-{episode_index:06d}/" + IMAGE_FILE_PATTERN +DEFAULT_DEPTH_PATH = "images/{image_key}/episode-{episode_index:06d}/" + DEPTH_FILE_PATTERN LEGACY_EPISODES_PATH = "meta/episodes.jsonl" LEGACY_EPISODES_STATS_PATH = "meta/episodes_stats.jsonl" diff --git a/src/lerobot/datasets/video_utils.py b/src/lerobot/datasets/video_utils.py index ca90fba45..ef3005dd8 100644 --- a/src/lerobot/datasets/video_utils.py +++ b/src/lerobot/datasets/video_utils.py @@ -39,11 +39,17 @@ from datasets.features.features import register_feature from PIL import Image from lerobot.configs import ( + DepthEncoderConfig, + RGBEncoderConfig, VideoEncoderConfig, - camera_encoder_defaults, + depth_encoder_defaults, + rgb_encoder_defaults, ) from lerobot.utils.import_utils import get_safe_default_video_backend +from .depth_utils import quantize_depth +from .pyav_utils import get_pix_fmt_channels + logger = logging.getLogger(__name__) @@ -53,6 +59,7 @@ def decode_video_frames( tolerance_s: float, backend: str | None = None, return_uint8: bool = False, + is_depth: bool = False, ) -> torch.Tensor: """ Decodes video frames using the specified backend. @@ -64,23 +71,35 @@ def decode_video_frames( backend (str, optional): Backend to use for decoding. Defaults to "torchcodec" when available in the platform; otherwise, defaults to "pyav". The legacy value "video_reader" is accepted for one release as an alias for "pyav" and will be removed in a future version. - return_uint8 (bool): If True, return raw uint8 frames without float32 normalization. + return_uint8 (bool): For RGB videos, if True return raw uint8 frames without float32 normalization. This reduces memory for DataLoader IPC; normalization can be done on GPU afterward. + is_depth (bool): Set to True if the video is a depth map (1 channel, uint12). Returns: - torch.Tensor: Decoded frames (float32 in [0,1] by default, or uint8 if return_uint8=True). + torch.Tensor: Decoded frames (RGB: float32 in [0,1] by default, or uint8 if return_uint8=True, Depth: uint12). Currently supports torchcodec on cpu and pyav. """ + if backend != "pyav" and is_depth: + logger.debug("Decoding depth maps is only supported with the 'pyav' backend, falling back to pyav.") + # We do not actually return uint8 here, but we avoid the 255 normalization step. + return decode_video_frames_pyav( + video_path, timestamps, tolerance_s, return_uint8=False, is_depth=True + ) + if backend is None: backend = get_safe_default_video_backend() if backend == "torchcodec": return decode_video_frames_torchcodec(video_path, timestamps, tolerance_s, return_uint8=return_uint8) elif backend == "pyav": - return decode_video_frames_pyav(video_path, timestamps, tolerance_s, return_uint8=return_uint8) + return decode_video_frames_pyav( + video_path, timestamps, tolerance_s, return_uint8=return_uint8, is_depth=is_depth + ) elif backend == "video_reader": logger.warning("backend='video_reader' is deprecated and now aliases to 'pyav'.") - return decode_video_frames_pyav(video_path, timestamps, tolerance_s, return_uint8=return_uint8) + return decode_video_frames_pyav( + video_path, timestamps, tolerance_s, return_uint8=return_uint8, is_depth=is_depth + ) else: raise ValueError(f"Unsupported video backend: {backend}") @@ -91,6 +110,7 @@ def decode_video_frames_pyav( tolerance_s: float, log_loaded_timestamps: bool = False, return_uint8: bool = False, + is_depth: bool = False, ) -> torch.Tensor: """Loads frames associated to the requested timestamps of a video using PyAV. @@ -109,8 +129,9 @@ def decode_video_frames_pyav( tolerance_s: Allowed deviation in seconds between a queried timestamp and the closest decoded frame. log_loaded_timestamps: When True, log every decoded frame's timestamp at INFO level. - return_uint8: When True, return raw uint8 frames (C, H, W). Otherwise, return float32 in - [0, 1] range. + return_uint8: For RGB videos, if True return raw uint8 frames (C, H, W). + Otherwise, return float32 in [0, 1] range. + is_depth: Set to True if the video is a depth map (1 channel, uint12). Returns: torch.Tensor of shape (len(timestamps), C, H, W). @@ -132,7 +153,13 @@ def decode_video_frames_pyav( # https://pyav.basswood-io.com/docs/stable/api/container.html#av.container.InputContainer.seek with av.open(video_path) as container: stream = container.streams.video[0] - container.seek(int(first_ts * av.time_base), backward=True) + # Seek to the nearest keyframe at or before `first_ts` with a 1 frame margin + container.seek( + round(first_ts / stream.time_base) - 1, + backward=True, + any_frame=False, + stream=stream, + ) for frame in container.decode(stream): if frame.pts is None: @@ -140,9 +167,13 @@ def decode_video_frames_pyav( current_ts = float(frame.pts * stream.time_base) if log_loaded_timestamps: logger.info(f"frame loaded at timestamp={current_ts:.4f}") - # Convert to CHW uint8 to match torchcodec's output layout. - arr = frame.to_ndarray(format="rgb24") # H, W, 3 - loaded_frames.append(torch.from_numpy(arr).permute(2, 0, 1).contiguous()) + if is_depth: + arr = frame.to_ndarray(format="gray12le") # (H, W) uint12 + loaded_frames.append(torch.from_numpy(arr).unsqueeze(0).contiguous()) + else: + arr = frame.to_ndarray(format="rgb24") # (H, W, 3) + # Convert to CHW uint8 to match torchcodec's output layout. + loaded_frames.append(torch.from_numpy(arr).permute(2, 0, 1).contiguous()) loaded_ts.append(current_ts) if current_ts >= last_ts: break @@ -185,7 +216,7 @@ def decode_video_frames_pyav( f"number of queried timestamps ({len(timestamps)})" ) - if return_uint8: + if return_uint8 or is_depth: return closest_frames # convert to the pytorch format which is float32 in [0,1] range (and channel first) @@ -406,17 +437,38 @@ def encode_video_frames( imgs_dir: Path | str, video_path: Path | str, fps: int, - camera_encoder: VideoEncoderConfig | None = None, + video_encoder: VideoEncoderConfig | None = None, encoder_threads: int | None = None, *, log_level: int | None = av.logging.WARNING, overwrite: bool = False, ) -> None: - """More info on ffmpeg arguments tuning on `benchmark/video/README.md`""" - if camera_encoder is None: - camera_encoder = camera_encoder_defaults() - vcodec = camera_encoder.vcodec - pix_fmt = camera_encoder.pix_fmt + """Encode a directory of image frames into an MP4 video. + + When ``video_encoder`` is a :class:`~lerobot.configs.video.DepthEncoderConfig`, + frames are read from ``.tiff`` files and quantized to 12-bit depth codes using the + encoder's ``depth_min`` / ``depth_max`` / ``shift`` / ``use_log``; otherwise ``.png`` + RGB frames are encoded directly. + + Args: + imgs_dir: Directory containing the frames to encode, named ``frame-000000`` + onwards (``.png`` for RGB, ``.tiff`` for depth). + video_path: Output path for the encoded ``.mp4`` file. + fps: Frame rate of the output video. + video_encoder: Encoder settings (codec, pixel format, quality, ...). When + ``None``, :func:`rgb_encoder_defaults` is used. Pass a + :class:`~lerobot.configs.video.DepthEncoderConfig` to encode depth frames. + encoder_threads: Per-encoder thread count forwarded to the codec. ``None`` + lets the codec decide. + log_level: libav log level to set while encoding, or ``None`` to leave the + current logging configuration unchanged. + overwrite: When ``False`` and ``video_path`` already exists, skip encoding and + log a warning. When ``True``, re-encode and replace the existing file. + """ + if video_encoder is None: + video_encoder = rgb_encoder_defaults() + vcodec = video_encoder.vcodec + pix_fmt = video_encoder.pix_fmt video_path = Path(video_path) imgs_dir = Path(imgs_dir) @@ -428,17 +480,19 @@ def encode_video_frames( video_path.parent.mkdir(parents=True, exist_ok=True) # Get input frames - template = "frame-" + ("[0-9]" * 6) + ".png" + is_depth = isinstance(video_encoder, DepthEncoderConfig) + suffix = ".png" if not is_depth else ".tiff" + template = "frame-" + ("[0-9]" * 6) + suffix input_list = sorted( glob.glob(str(imgs_dir / template)), key=lambda x: int(x.split("-")[-1].split(".")[0]) ) if len(input_list) == 0: - raise FileNotFoundError(f"No images found in {imgs_dir}.") + raise FileNotFoundError(f"No images with suffix {suffix} found in {imgs_dir}.") with Image.open(input_list[0]) as dummy_image: width, height = dummy_image.size - video_options = camera_encoder.get_codec_options(encoder_threads, as_strings=True) + video_options = video_encoder.get_codec_options(encoder_threads, as_strings=True) # Set logging level if log_level is not None: @@ -455,8 +509,19 @@ def encode_video_frames( # Loop through input frames and encode them for input_data in input_list: with Image.open(input_data) as input_image: - input_image = input_image.convert("RGB") - input_frame = av.VideoFrame.from_image(input_image) + if is_depth: + input_frame = quantize_depth( + np.array(input_image), + depth_min=video_encoder.depth_min, + depth_max=video_encoder.depth_max, + shift=video_encoder.shift, + use_log=video_encoder.use_log, + pix_fmt=video_encoder.pix_fmt, + video_backend="pyav", + ) + else: + input_image = input_image.convert("RGB") + input_frame = av.VideoFrame.from_image(input_image) packet = output_stream.encode(input_frame) if packet: output.mux(packet) @@ -477,7 +542,7 @@ def encode_video_frames( def reencode_video( input_video_path: Path | str, output_video_path: Path | str, - camera_encoder: VideoEncoderConfig | None = None, + video_encoder: VideoEncoderConfig | None = None, encoder_threads: int | None = None, log_level: int | None = av.logging.WARNING, overwrite: bool = False, @@ -489,7 +554,7 @@ def reencode_video( Args: input_video_path: Existing video file to read. output_video_path: Path for the re-encoded file. - camera_encoder: Encoder configuration. Defaults to :func:`camera_encoder_defaults`. + video_encoder: Encoder configuration. Defaults to :func:`rgb_encoder_defaults`. encoder_threads: Optional thread count forwarded to :meth:`VideoEncoderConfig.get_codec_options`. log_level: libav log level while encoding, or ``None`` to leave logging unchanged. Defaults to WARNING. overwrite: When ``False`` and ``output_video_path`` already exists, skip and log a warning. @@ -497,7 +562,7 @@ def reencode_video( end_time_s: When set, trim the output to end at this timestamp (seconds, exclusive). """ - camera_encoder = camera_encoder or camera_encoder_defaults() + video_encoder = video_encoder or rgb_encoder_defaults() if (start_time_s is not None and start_time_s < 0) or (end_time_s is not None and end_time_s < 0): raise ValueError(f"Trim times must be non-negative, got start={start_time_s}, end={end_time_s}.") @@ -512,9 +577,9 @@ def reencode_video( output_video_path.parent.mkdir(parents=True, exist_ok=True) - video_options = camera_encoder.get_codec_options(encoder_threads, as_strings=True) - vcodec = camera_encoder.vcodec - pix_fmt = camera_encoder.pix_fmt + video_options = video_encoder.get_codec_options(encoder_threads, as_strings=True) + vcodec = video_encoder.vcodec + pix_fmt = video_encoder.pix_fmt with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp_named_file: tmp_output_video_path = tmp_named_file.name @@ -696,22 +761,21 @@ class _CameraEncoderThread(threading.Thread): self, video_path: Path, fps: int, - vcodec: str, - pix_fmt: str, - codec_options: dict[str, str], + video_encoder: VideoEncoderConfig, frame_queue: queue.Queue, result_queue: queue.Queue, stop_event: threading.Event, + encoder_threads: int | None = None, ): super().__init__(daemon=True) self.video_path = video_path self.fps = fps - self.vcodec = vcodec - self.pix_fmt = pix_fmt - self.codec_options = codec_options + self.video_encoder = video_encoder + self.is_depth = isinstance(video_encoder, DepthEncoderConfig) self.frame_queue = frame_queue self.result_queue = result_queue self.stop_event = stop_event + self.encoder_threads = encoder_threads def run(self) -> None: from .compute_stats import RunningQuantileStats, auto_downsample_height_width @@ -736,12 +800,12 @@ class _CameraEncoderThread(threading.Thread): # Sentinel: flush and close break - # Ensure HWC uint8 numpy array + # Ensure HWC (RGB or depth) uint8 (RGB only) numpy array if isinstance(frame_data, np.ndarray): - if frame_data.ndim == 3 and frame_data.shape[0] == 3: + if frame_data.ndim == 3 and frame_data.shape[0] in (1, 3): # CHW -> HWC frame_data = frame_data.transpose(1, 2, 0) - if frame_data.dtype != np.uint8: + if not self.is_depth and frame_data.dtype != np.uint8: frame_data = (frame_data * 255).astype(np.uint8) # Open container on first frame (to get width/height) @@ -749,15 +813,29 @@ class _CameraEncoderThread(threading.Thread): height, width = frame_data.shape[:2] Path(self.video_path).parent.mkdir(parents=True, exist_ok=True) container = av.open(str(self.video_path), "w") - output_stream = container.add_stream(self.vcodec, self.fps, options=self.codec_options) - output_stream.pix_fmt = self.pix_fmt + output_stream = container.add_stream( + self.video_encoder.vcodec, + self.fps, + options=self.video_encoder.get_codec_options(self.encoder_threads, as_strings=True), + ) + output_stream.pix_fmt = self.video_encoder.pix_fmt output_stream.width = width output_stream.height = height output_stream.time_base = Fraction(1, self.fps) # Encode frame with explicit timestamps - pil_img = Image.fromarray(frame_data) - video_frame = av.VideoFrame.from_image(pil_img) + if not self.is_depth: + pil_img = Image.fromarray(frame_data) + video_frame = av.VideoFrame.from_image(pil_img) + else: + video_frame = quantize_depth( + frame_data, + depth_min=self.video_encoder.depth_min, + depth_max=self.video_encoder.depth_max, + shift=self.video_encoder.shift, + use_log=self.video_encoder.use_log, + video_backend=self.video_encoder.video_backend, + ) video_frame.pts = frame_count video_frame.time_base = Fraction(1, self.fps) packet = output_stream.encode(video_frame) @@ -815,22 +893,27 @@ class StreamingVideoEncoder: def __init__( self, fps: int, - camera_encoder: VideoEncoderConfig | None = None, + rgb_encoder: RGBEncoderConfig | None = None, + depth_encoder: DepthEncoderConfig | None = None, queue_maxsize: int = 30, encoder_threads: int | None = None, ): """ Args: fps: Frames per second for the output videos. - camera_encoder: Video encoder settings applied to all cameras. - When ``None``, :func:`camera_encoder_defaults` is used. - encoder_threads: Number of encoder threads (global setting). - ``None`` lets the codec decide. + rgb_encoder: Video encoder settings applied to all RGB cameras. + When ``None``, :func:`rgb_encoder_defaults` is used. + depth_encoder: Video encoder settings applied to all depth cameras, + including the depth quantization parameters. When ``None``, + :func:`depth_encoder_defaults` is used. queue_maxsize: Max frames to buffer per camera before back-pressure drops frames. + encoder_threads: Number of encoder threads (global setting). + ``None`` lets the codec decide. """ self.fps = fps - self._camera_encoder = camera_encoder or camera_encoder_defaults() + self._rgb_encoder = rgb_encoder or rgb_encoder_defaults() + self._depth_encoder = depth_encoder or depth_encoder_defaults() self._encoder_threads = encoder_threads self.queue_maxsize = queue_maxsize @@ -843,18 +926,25 @@ class StreamingVideoEncoder: self._episode_active = False self._closed = False - def start_episode(self, video_keys: list[str], temp_dir: Path) -> None: + def start_episode( + self, video_keys: list[str], temp_dir: Path, depth_video_keys: list[str] | None = None + ) -> None: """Start encoder threads for a new episode. Args: video_keys: List of video feature keys (e.g. ["observation.images.laptop"]) temp_dir: Base directory for temporary MP4 files + depth_video_keys: List of video or image feature keys that carry depth maps (e.g. + ["observation.images.laptop_depth"]). Defaults to ``[]`` (no depth keys). """ if self._episode_active: self.cancel_episode() self._dropped_frames.clear() + if depth_video_keys is None: + depth_video_keys = [] + for video_key in video_keys: frame_queue: queue.Queue = queue.Queue(maxsize=self.queue_maxsize) result_queue: queue.Queue = queue.Queue(maxsize=1) @@ -863,17 +953,15 @@ class StreamingVideoEncoder: temp_video_dir = Path(tempfile.mkdtemp(dir=temp_dir)) video_path = temp_video_dir / f"{video_key.replace('/', '_')}_streaming.mp4" - vcodec = self._camera_encoder.vcodec - codec_options = self._camera_encoder.get_codec_options(self._encoder_threads, as_strings=True) + encoder = self._depth_encoder if video_key in depth_video_keys else self._rgb_encoder encoder_thread = _CameraEncoderThread( video_path=video_path, fps=self.fps, - vcodec=vcodec, - pix_fmt=self._camera_encoder.pix_fmt, - codec_options=codec_options, + video_encoder=encoder, frame_queue=frame_queue, result_queue=result_queue, stop_event=stop_event, + encoder_threads=self._encoder_threads, ) encoder_thread.start() @@ -1080,15 +1168,23 @@ def get_audio_info(video_path: Path | str) -> dict: def get_video_info( video_path: Path | str, - camera_encoder: VideoEncoderConfig | None = None, + video_encoder: VideoEncoderConfig | None = None, ) -> dict: """Build the ``video.*`` / ``audio.*`` info dict persisted in ``info.json``. Args: video_path: Path to the encoded video file to probe. - camera_encoder: If provided, record the exact encoder settings used to encode this + video_encoder: If provided, record the exact encoder settings used to encode this video. Stream-derived values take precedence — encoder fields are only written for keys - not already populated from the video file itself. + not already populated from the video file itself. When a + :class:`~lerobot.configs.video.DepthEncoderConfig` is passed, the depth + quantization parameters (``depth_min`` / ``depth_max`` / ``shift`` / + ``use_log``) are recorded so frames can be dequantized on read. + + Returns: + The ``video.*`` / ``audio.*`` info dict, including ``is_depth_map`` which is + ``True`` only when ``video_encoder`` is a + :class:`~lerobot.configs.video.DepthEncoderConfig`. """ logging.getLogger("libav").setLevel(av.logging.WARNING) @@ -1106,13 +1202,10 @@ def get_video_info( video_info["video.width"] = video_stream.width video_info["video.codec"] = video_stream.codec.canonical_name video_info["video.pix_fmt"] = video_stream.pix_fmt - video_info["video.is_depth_map"] = False # Calculate fps from r_frame_rate video_info["video.fps"] = int(video_stream.base_rate) - - pixel_channels = get_video_pixel_channels(video_stream.pix_fmt) - video_info["video.channels"] = pixel_channels + video_info["video.channels"] = get_pix_fmt_channels(video_stream.pix_fmt) # Reset logging level av.logging.restore_default_callback() @@ -1121,27 +1214,18 @@ def get_video_info( video_info.update(**get_audio_info(video_path)) # Add additional encoder configuration if provided - if camera_encoder is not None: - for field_name, field_value in asdict(camera_encoder).items(): + if video_encoder is not None: + for field_name, field_value in asdict(video_encoder).items(): # vcodec is already populated from the video stream if field_name == "vcodec": continue video_info.setdefault(f"video.{field_name}", field_value) + video_info["is_depth_map"] = isinstance(video_encoder, DepthEncoderConfig) + return video_info -def get_video_pixel_channels(pix_fmt: str) -> int: - if "gray" in pix_fmt or "depth" in pix_fmt or "monochrome" in pix_fmt: - return 1 - elif "rgba" in pix_fmt or "yuva" in pix_fmt: - return 4 - elif "rgb" in pix_fmt or "yuv" in pix_fmt: - return 3 - else: - raise ValueError("Unknown format") - - def get_video_duration_in_s(video_path: Path | str) -> float: """ Get the duration of a video file in seconds using PyAV. @@ -1202,10 +1286,13 @@ class VideoEncodingManager: img_dir = self.dataset.root / "images" if img_dir.exists(): png_files = list(img_dir.rglob("*.png")) - if len(png_files) == 0: + tiff_files = list(img_dir.rglob("*.tiff")) + if len(png_files) == 0 and len(tiff_files) == 0: shutil.rmtree(img_dir) logger.debug("Cleaned up empty images directory") else: - logger.debug(f"Images directory is not empty, containing {len(png_files)} PNG files") + logger.debug( + f"Images directory is not empty, containing {len(png_files)} PNG and {len(tiff_files)} TIFF files" + ) return False # Don't suppress the original exception diff --git a/src/lerobot/envs/utils.py b/src/lerobot/envs/utils.py index 6e6f352e9..8b9c4f94b 100644 --- a/src/lerobot/envs/utils.py +++ b/src/lerobot/envs/utils.py @@ -126,6 +126,26 @@ def preprocess_observation(observations: dict[str, np.ndarray]) -> dict[str, Ten if "camera_obs" in observations: return_observations[f"{OBS_STR}.camera_obs"] = observations["camera_obs"] + # Pass through any remaining ndarray/tensor keys not already handled above, + # so env plugins can expose extra observation keys via get_env_processors(). + _handled = {"pixels", "environment_state", "agent_pos", "robot_state", "policy", "camera_obs"} + for key, value in observations.items(): + if key in _handled: + continue + target = f"{OBS_STR}.{key}" + if target in return_observations: + continue + if isinstance(value, np.ndarray): + val = torch.from_numpy(value).float() + if val.dim() == 1: + val = val.unsqueeze(0) + return_observations[target] = val + elif isinstance(value, Tensor): + val = value.float() + if val.dim() == 1: + val = val.unsqueeze(0) + return_observations[target] = val + return return_observations diff --git a/src/lerobot/jobs/__init__.py b/src/lerobot/jobs/__init__.py new file mode 100644 index 000000000..674b98b85 --- /dev/null +++ b/src/lerobot/jobs/__init__.py @@ -0,0 +1,23 @@ +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from lerobot.utils.import_utils import require_package + +# LeRobotDataset (imported at module top in dataset.py) pulls in heavy dataset deps; +# guard the optional dependency here so importing this package fails loudly if it's missing. +require_package("datasets", extra="dataset") + +from .hf import submit_to_hf + +__all__ = ["submit_to_hf"] diff --git a/src/lerobot/jobs/dataset.py b/src/lerobot/jobs/dataset.py new file mode 100644 index 000000000..497f8445e --- /dev/null +++ b/src/lerobot/jobs/dataset.py @@ -0,0 +1,53 @@ +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Make a training dataset reachable from an HF Job pod. + +The pod can't see the host's ~/.cache/huggingface/lerobot, so the dataset has to +live on the Hub: the pod downloads it by repo_id at train time (the forwarded +HF_TOKEN covers private datasets). A dataset already on the Hub is used as-is; a +local-only dataset is pushed to a PRIVATE repo first (never public). +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from lerobot.datasets import LeRobotDataset +from lerobot.utils.constants import HF_LEROBOT_HOME + +if TYPE_CHECKING: + from huggingface_hub import HfApi + + +def ensure_dataset_available(repo_id: str, *, api: HfApi, tags: list[str] | None = None) -> None: + """Ensure repo_id resolves on the Hub, pushing a local-only dataset privately first. + + `tags` are attached to the dataset only when we push it (an already-on-Hub + dataset is left untouched). Raises RuntimeError if the dataset is neither on + the Hub nor in the local cache. + """ + if api.repo_exists(repo_id, repo_type="dataset"): + return + + local_present = (HF_LEROBOT_HOME / repo_id / "meta" / "info.json").is_file() + if not local_present: + raise RuntimeError( + f"Dataset '{repo_id}' is not in the local cache ({HF_LEROBOT_HOME}) and could not be " + f"reached on the Hub — it may not exist, or be private and inaccessible with your " + f"token. Record or download it first, or run `hf auth login`." + ) + + print(f"[dataset] '{repo_id}' is local-only; pushing to a PRIVATE Hub repo...") + LeRobotDataset(repo_id).push_to_hub(private=True, tags=tags) + print(f"[dataset] '{repo_id}' uploaded (private). The job will download it by repo_id.") diff --git a/src/lerobot/jobs/hf.py b/src/lerobot/jobs/hf.py new file mode 100644 index 000000000..645666412 --- /dev/null +++ b/src/lerobot/jobs/hf.py @@ -0,0 +1,425 @@ +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Run a lerobot training on HF Jobs (HuggingFace GPUs). + +Ported and simplified from lelab's runners/hf_cloud.py: no UI log queue, no +registry — just submit and stream to stdout. +""" + +from __future__ import annotations + +import copy +import datetime as dt +import json +import netrc +import os +import re +import signal +import sys +import tempfile +import threading +from pathlib import Path +from typing import TYPE_CHECKING + +import httpx +from huggingface_hub import ( + HfApi, + create_repo, + fetch_job_logs, + get_token, + inspect_job, + run_job, + upload_file, +) + +from lerobot.common.train_utils import push_checkpoint_to_hub +from lerobot.configs import parser + +from .dataset import ensure_dataset_available + +if TYPE_CHECKING: + from lerobot.configs.train import TrainPipelineConfig + +_SLUG_RE = re.compile(r"[^a-zA-Z0-9._-]+") + +_TERMINAL_STAGES = {"COMPLETED", "CANCELED", "ERROR", "DELETED"} + +# huggingface_hub 1.x runs on httpx: transient HTTP/transport failures surface as +# httpx.HTTPError and socket-level errors as OSError. Catching only these keeps real +# bugs (TypeError, AttributeError, ...) from being silently retried or counted as +# job failures. +_TRANSIENT_NET_ERRORS = (OSError, httpx.HTTPError) + +# Always attached to remote jobs and pushed datasets so LeRobot-originated work +# is identifiable on the Hub; callers (e.g. LeLab) add their own via --job.tags. +LEROBOT_TAG = "lerobot" + + +def resolve_job_tags(extra: list[str] | None) -> list[str]: + """Return the tag list for a run: the lerobot tag plus any extras, deduped, order-stable.""" + tags = [LEROBOT_TAG, *(extra or [])] + seen: set[str] = set() + return [t for t in tags if not (t in seen or seen.add(t))] + + +def resolve_wandb_api_key() -> str | None: + """Host's wandb key for forwarding to the job: $WANDB_API_KEY, else ~/.netrc.""" + key = os.environ.get("WANDB_API_KEY") + if key: + return key + try: + rc = netrc.netrc() + except (FileNotFoundError, netrc.NetrcParseError, OSError): + return None + auth = rc.authenticators("api.wandb.ai") + if auth is None: + return None + _login, _account, password = auth + return password or None + + +def build_repo_id(username: str, job_name: str, now: dt.datetime) -> str: + """Generate the model repo id for a remote run: /_.""" + slug = _SLUG_RE.sub("-", job_name).strip("-") or "train" + stamp = now.strftime("%Y-%m-%d_%H-%M-%S") + return f"{username}/{slug}_{stamp}" + + +def build_remote_config_file(cfg, repo_id: str, dest: Path, tags: list[str] | None = None) -> Path: + """Write a train_config.json for the pod, with remote overrides applied. + + The pod runs `lerobot-train --config_path=` and downloads the dataset + by repo_id into its own cache. Client-only fields are stripped so the config + is accepted by the trainer image: `job` (pure client orchestration) is always + removed, and `save_checkpoint_to_hub` is removed unless explicitly enabled — + older lerobot images reject unknown keys, so the default keeps the config + compatible with the released `lerobot-gpu` image. `tags` are merged into + policy.tags so the trained model the pod pushes carries them too. + """ + remote = copy.deepcopy(cfg) + remote.policy.push_to_hub = True + remote.policy.repo_id = repo_id + # Don't pin the client's resolved device (e.g. "mps"); let the pod auto-detect its GPU. + remote.policy.device = None + # Drop any host-local dataset root; the pod resolves the dataset by repo_id. + remote.dataset.root = None + if tags: + existing = list(remote.policy.tags or []) + remote.policy.tags = existing + [t for t in tags if t not in existing] + + # Encode to the canonical, pod-parseable dict, then drop the keys the released + # trainer image doesn't know about. + data = remote.to_dict() + data.pop("job", None) + if not remote.save_checkpoint_to_hub: + data.pop("save_checkpoint_to_hub", None) + + dest.parent.mkdir(parents=True, exist_ok=True) + dest.write_text(json.dumps(data, indent=4)) + return dest + + +def _stage_config_on_hub(cfg, repo_id: str, token: str, tags: list[str] | None = None) -> str: + """Upload train_config.json to the model repo and return the repo_id for --config_path.""" + create_repo(repo_id, repo_type="model", private=True, exist_ok=True, token=token) + with tempfile.TemporaryDirectory() as tmp: + config_path = build_remote_config_file(cfg, repo_id, Path(tmp) / "train_config.json", tags=tags) + upload_file( + path_or_fileobj=config_path, + path_in_repo="train_config.json", + repo_id=repo_id, + repo_type="model", + token=token, + ) + return repo_id + + +def _tail_logs( + job_id: str, + done: threading.Event, + success_marker: str | None = None, + success_event: threading.Event | None = None, +) -> None: + """Stream job logs to stdout, reconnecting on dropped streams until done is set. + + Each reconnect re-fetches the full buffered log, so we track how many lines + were already printed and skip them — otherwise a fast-failing job's traceback + gets reprinted on every reconnect. + + When `success_marker` appears in a line, set `success_event` and `done` so the + caller can finish as soon as the trained model lands on the Hub, rather than + waiting out the platform's post-run finalization (which can add ~30s). + """ + printed = 0 + while not done.is_set(): + try: + seen = 0 + for line in fetch_job_logs(job_id=job_id, follow=True): + seen += 1 + if seen <= printed: + continue # already shown on a previous connection + printed = seen + # fetch_job_logs yields SSE data without trailing newlines, so add one + # per entry — otherwise all log lines concatenate onto a single line. + print(line.rstrip("\n"), flush=True) + if success_marker and success_event is not None and success_marker in line: + success_event.set() + done.set() + return + if done.is_set(): + return + # Stream closed cleanly. Wait a moment so the status poller can mark + # the job terminal before we reconnect (avoids re-tailing the buffer). + if done.wait(3): + return + except _TRANSIENT_NET_ERRORS: + if done.wait(2): + return + + +def _poll_until_done( + job_id: str, + done: threading.Event, + poll_interval: float = 5.0, + status_holder: dict | None = None, + max_failures: int = 6, +) -> str | None: + """Poll inspect_job until a terminal stage or until `done` is set. + + Returns the terminal stage string, or None if `done` was set first (detach) + or after `max_failures` consecutive inspect_job errors. When a terminal stage + is reached and `status_holder` is given, records `status_holder["message"]` + (the platform's status message, e.g. "Job timeout"). + """ + failures = 0 + while not done.is_set(): + try: + info = inspect_job(job_id=job_id) + failures = 0 + # `stage` is an enum in some huggingface_hub versions and a plain str in others. + stage = getattr(info.status.stage, "value", info.status.stage) + if stage in _TERMINAL_STAGES: + if status_holder is not None: + status_holder["message"] = getattr(info.status, "message", None) + done.set() + return stage + except _TRANSIENT_NET_ERRORS: + failures += 1 + if failures >= max_failures: + done.set() + return None + done.wait(poll_interval) + return None + + +def _pod_forwarded_args( + argv: list[str], drop_names: tuple[str, ...] = (), drop_prefixes: tuple[str, ...] = () +) -> list[str]: + """User CLI overrides to replay on the pod, minus flags the submitter sets itself. + + Handles both `--name=value` and `--name value` forms. Forwarding the user's overrides (e.g. + `--steps`, `--save_checkpoint_to_hub`) makes a remote resume behave like the same local command. + """ + out: list[str] = [] + skip_next = False + for i, tok in enumerate(argv): + if skip_next: + skip_next = False + continue + name = tok.split("=", 1)[0] + if name in drop_names or any(name.startswith(p) for p in drop_prefixes): + if "=" not in tok and i + 1 < len(argv) and not argv[i + 1].startswith("--"): + skip_next = True # also drop the space-separated value + continue + out.append(tok) + return out + + +def _build_resume_job(cfg: TrainPipelineConfig, username: str) -> tuple[str, list[str]]: + """Resolve the model repo and pod command to resume a run on a job. + + A Hub `config_path` is resumed from directly: its checkpoint config already targets that repo, + so new checkpoints continue the lineage there. A local `config_path` has its checkpoint uploaded + to a new PRIVATE repo first, and the resumed run is forced to push back to it. The pod command + always carries `--job.target=local` so the checkpoint's saved `job.target` can't make the pod + re-dispatch itself. + """ + config_path = parser.parse_arg("config_path") + forwarded = _pod_forwarded_args( + sys.argv[1:], + drop_names=("--config_path", "--policy.repo_id", "--policy.push_to_hub", "--dataset.root"), + drop_prefixes=("--job.",), + ) + + if Path(config_path).exists(): + # Local checkpoint: stage it on the Hub so the pod can resume from it, and push back there. + # Resolve so a `last` symlink uploads under its real step name (digit), which the pod's + # latest-checkpoint lookup keys on. + checkpoint_dir = Path(cfg.checkpoint_path).resolve() + source_repo = build_repo_id(username, cfg.job_name or "train", dt.datetime.now(dt.UTC)) + push_checkpoint_to_hub(checkpoint_dir, source_repo, private=True) + extra = [f"--policy.repo_id={source_repo}", "--policy.push_to_hub=true"] + else: + source_repo = config_path + extra = [] + + command = [ + "lerobot-train", + *forwarded, + f"--config_path={source_repo}", + "--job.target=local", + *extra, + ] + return source_repo, command + + +def submit_to_hf(cfg: TrainPipelineConfig) -> None: + """Submit a training job to HF Jobs infrastructure. + + Validates cfg, resolves credentials, ensures the dataset is on the Hub, then either stages a + sanitized config (fresh run) or resumes from a checkpoint repo, submits the job, and tails logs + until completion or detaches immediately. Ctrl-C detaches without cancelling the remote job. + """ + token = get_token() + if not token: + raise RuntimeError("Not logged in to Hugging Face. Run `hf auth login` first.") + + api = HfApi(token=token) + user_info = api.whoami(token=token) + username = user_info["name"] + + now = dt.datetime.now(dt.UTC) + fresh_repo_id: str | None = None + if not cfg.resume: + # Resolve the model repo and mark it for push BEFORE validate(): validate() requires repo_id + # to be set whenever push_to_hub is True. (A resume reuses the checkpoint's repo instead.) + if cfg.policy is not None: + base_name = cfg.job_name or cfg.policy.type + fresh_repo_id = cfg.policy.repo_id or build_repo_id(username, base_name, now) + cfg.policy.repo_id = fresh_repo_id + cfg.policy.push_to_hub = True + else: + # Path-based policy is resolved inside validate(); fall back to a generic slug. + fresh_repo_id = build_repo_id(username, cfg.job_name or "train", now) + + cfg.validate() + + if cfg.is_reward_model_training: + raise ValueError( + "Remote training via --job.target only supports policy training, not reward models. " + "Run reward-model training locally." + ) + + secrets: dict[str, str] = {"HF_TOKEN": token} + if cfg.wandb.enable: + wandb_key = resolve_wandb_api_key() + if wandb_key is None: + raise ValueError( + "wandb is enabled but no WANDB_API_KEY found. " + "Set it via `export WANDB_API_KEY=...` or add it to ~/.netrc." + ) + secrets["WANDB_API_KEY"] = wandb_key + + tags = resolve_job_tags(cfg.job.tags) + # The dataset must be reachable from the pod for both fresh and resumed runs; a local-only + # dataset is pushed PRIVATE here. Hoisted before the resume/fresh branch since it applies to both. + ensure_dataset_available(cfg.dataset.repo_id, api=api, tags=tags) + + if cfg.resume: + repo_id, command = _build_resume_job(cfg, username) + else: + config_repo_id = _stage_config_on_hub(cfg, fresh_repo_id, token, tags=tags) + repo_id = fresh_repo_id + command = ["lerobot-train", f"--config_path={config_repo_id}"] + + print(f"Submitting job to HF Jobs (flavor={cfg.job.target}, image={cfg.job.image}) ...") + job_info = run_job( + image=cfg.job.image, + command=command, + flavor=cfg.job.target, + secrets=secrets, + timeout=cfg.job.timeout, + # HF Jobs labels are key/value; expose each tag as a queryable label. + labels=dict.fromkeys(tags, "true"), + ) + job_id = job_info.id + job_url = getattr(job_info, "url", None) + print(f"Job submitted: {job_id}") + if job_url: + print(f" Job page: {job_url}") + print(f" Model repo: https://huggingface.co/{repo_id}") + print(f" Monitor: hf jobs logs {job_id}") + print(f" Cancel: hf jobs cancel {job_id}") + + if cfg.job.detach: + return + + done = threading.Event() + detached = threading.Event() + pushed_ok = threading.Event() + stage_holder: dict[str, str | None] = {} + + def _poll() -> None: + stage_holder["stage"] = _poll_until_done(job_id, done, status_holder=stage_holder) + + poll_thread = threading.Thread(target=_poll, daemon=True) + poll_thread.start() + # Finish as soon as the model is pushed, rather than waiting out the platform's + # post-run finalization before the job stage flips to COMPLETED. This matches the + # exact log line emitted by PreTrainedPolicy.push_model_to_hub — the two must stay + # in sync. If it ever stops matching we just fall back to stage-based completion + # (~30s slower), so the contract is an optimization, not a correctness requirement. + success_marker = f"Model pushed to https://huggingface.co/{repo_id}" + log_thread = threading.Thread( + target=_tail_logs, args=(job_id, done, success_marker, pushed_ok), daemon=True + ) + log_thread.start() + + def _detach(sig, frame): + detached.set() + done.set() + print("\nDetached. Job is still running.") + print(f" Monitor: hf jobs logs {job_id}") + print(f" Cancel: hf jobs cancel {job_id}") + + # signal.signal only works on the main thread; when called from a worker thread + # (e.g. an orchestration framework) skip the Ctrl-C-detaches-instead-of-cancels + # handler rather than crashing with ValueError. + install_sigint = threading.current_thread() is threading.main_thread() + original_sigint = signal.getsignal(signal.SIGINT) if install_sigint else None + if install_sigint: + signal.signal(signal.SIGINT, _detach) + try: + # Timeout-based join so SIGINT is delivered to the main thread promptly. + while poll_thread.is_alive(): + poll_thread.join(timeout=0.5) + log_thread.join(timeout=5) + finally: + if install_sigint: + signal.signal(signal.SIGINT, original_sigint) + + if detached.is_set(): + return + + if pushed_ok.is_set(): + print(f"\nTraining complete — model pushed to https://huggingface.co/{repo_id}") + return + + stage = stage_holder.get("stage") + if stage != "COMPLETED": + message = stage_holder.get("message") + detail = f" ({message})" if message else "" + raise RuntimeError( + f"Job {job_id} ended with stage={stage}{detail}. Check logs: hf jobs logs {job_id}" + ) diff --git a/src/lerobot/optim/__init__.py b/src/lerobot/optim/__init__.py index 46676027b..2d564c25f 100644 --- a/src/lerobot/optim/__init__.py +++ b/src/lerobot/optim/__init__.py @@ -20,6 +20,7 @@ from .optimizers import ( SGDConfig as SGDConfig, XVLAAdamWConfig as XVLAAdamWConfig, load_optimizer_state, + load_optimizer_state_dict, save_optimizer_state, ) from .schedulers import ( @@ -50,6 +51,7 @@ __all__ = [ "VQBeTSchedulerConfig", # State management "load_optimizer_state", + "load_optimizer_state_dict", "load_scheduler_state", "save_optimizer_state", "save_scheduler_state", diff --git a/src/lerobot/optim/optimizers.py b/src/lerobot/optim/optimizers.py index 0bdd7a37e..0a462e1aa 100644 --- a/src/lerobot/optim/optimizers.py +++ b/src/lerobot/optim/optimizers.py @@ -27,7 +27,7 @@ from lerobot.utils.constants import ( OPTIMIZER_PARAM_GROUPS, OPTIMIZER_STATE, ) -from lerobot.utils.io_utils import deserialize_json_into_object, write_json +from lerobot.utils.io_utils import deserialize_json_into_object, load_json, write_json from lerobot.utils.utils import flatten_dict, unflatten_dict # Type alias for parameters accepted by optimizer build() methods. @@ -281,28 +281,37 @@ class MultiAdamConfig(OptimizerConfig): def save_optimizer_state( - optimizer: torch.optim.Optimizer | dict[str, torch.optim.Optimizer], save_dir: Path + optimizer: torch.optim.Optimizer | dict[str, torch.optim.Optimizer], + save_dir: Path, + optim_state_dict: dict | None = None, ) -> None: """Save optimizer state to disk. Args: optimizer: Either a single optimizer or a dictionary of optimizers. save_dir: Directory to save the optimizer state. + optim_state_dict: Pre-gathered optimizer state dict (for FSDP, where the sharded state must + be gathered across ranks first). If provided, it is saved directly instead of calling + ``optimizer.state_dict()``. Only supported for a single optimizer. Defaults to None. """ if isinstance(optimizer, dict): # Handle dictionary of optimizers + if optim_state_dict is not None: + raise ValueError("optim_state_dict is not supported for a dict of optimizers") for name, opt in optimizer.items(): optimizer_dir = save_dir / name optimizer_dir.mkdir(exist_ok=True, parents=True) _save_single_optimizer_state(opt, optimizer_dir) else: # Handle single optimizer - _save_single_optimizer_state(optimizer, save_dir) + _save_single_optimizer_state(optimizer, save_dir, optim_state_dict=optim_state_dict) -def _save_single_optimizer_state(optimizer: torch.optim.Optimizer, save_dir: Path) -> None: +def _save_single_optimizer_state( + optimizer: torch.optim.Optimizer, save_dir: Path, optim_state_dict: dict | None = None +) -> None: """Save a single optimizer's state to disk.""" - state = optimizer.state_dict() + state = dict(optim_state_dict) if optim_state_dict is not None else optimizer.state_dict() param_groups = state.pop("param_groups") flat_state = flatten_dict(state) save_file(flat_state, save_dir / OPTIMIZER_STATE) @@ -356,3 +365,19 @@ def _load_single_optimizer_state(optimizer: torch.optim.Optimizer, save_dir: Pat optimizer.load_state_dict(loaded_state_dict) return optimizer + + +def load_optimizer_state_dict(save_dir: Path) -> dict: + """Read a saved optimizer state dict (safetensors + json) back into a plain dict. + + Unlike `load_optimizer_state`, this does not load into an optimizer and preserves the original + ``state`` keys verbatim (e.g. FSDP parameter FQNs, which are not integer-castable). It is used by + the FSDP resume path, where the full state must be resharded via `FSDP.optim_state_dict_to_load` + before being loaded into the (sharded) optimizer. + """ + flat_state = load_file(save_dir / OPTIMIZER_STATE) + state = unflatten_dict(flat_state) + return { + "state": state.get("state", {}), + "param_groups": load_json(save_dir / OPTIMIZER_PARAM_GROUPS), + } diff --git a/src/lerobot/policies/act/modeling_act.py b/src/lerobot/policies/act/modeling_act.py index 5651fbfb1..1432b68a5 100644 --- a/src/lerobot/policies/act/modeling_act.py +++ b/src/lerobot/policies/act/modeling_act.py @@ -148,7 +148,7 @@ class ACTPolicy(PreTrainedPolicy): l1_loss = (abs_err * valid_mask).sum() / num_valid.clamp_min(1) loss_dict = {"l1_loss": l1_loss.item()} - if self.config.use_vae: + if self.config.use_vae and log_sigma_x2_hat is not None: # Calculate Dₖₗ(latent_pdf || standard_normal). Note: After computing the KL-divergence for # each dimension independently, we sum over the latent dimension to get the total # KL-divergence per batch element, then take the mean over the batch. diff --git a/src/lerobot/policies/diffusion/modeling_diffusion.py b/src/lerobot/policies/diffusion/modeling_diffusion.py index 9fbe1f703..8758a7e29 100644 --- a/src/lerobot/policies/diffusion/modeling_diffusion.py +++ b/src/lerobot/policies/diffusion/modeling_diffusion.py @@ -101,11 +101,23 @@ class DiffusionPolicy(PreTrainedPolicy): @torch.no_grad() def predict_action_chunk(self, batch: dict[str, Tensor], noise: Tensor | None = None) -> Tensor: - """Predict a chunk of actions given environment observations.""" - # stack n latest observations from the queue - batch = {k: torch.stack(list(self._queues[k]), dim=1) for k in batch if k in self._queues} - actions = self.diffusion.generate_actions(batch, noise=noise) + """Predict a chunk of actions given environment observations. + Supports two modes: + - Online (queues populated via select_action): stacks observations from internal queues. + - Offline (empty queues, e.g. dataloader batch): uses the batch directly. + """ + queues_populated = any(len(q) > 0 for q in self._queues.values()) + if queues_populated: + batch = {k: torch.stack(list(self._queues[k]), dim=1) for k in batch if k in self._queues} + else: + batch = dict(batch) + if self.config.image_features: + for key in self.config.image_features: + if batch[key].ndim == 4: + batch[key] = batch[key].unsqueeze(1) + batch[OBS_IMAGES] = torch.stack([batch[key] for key in self.config.image_features], dim=-4) + actions = self.diffusion.generate_actions(batch, noise=noise) return actions @torch.no_grad() diff --git a/src/lerobot/policies/factory.py b/src/lerobot/policies/factory.py index a42b38ba4..b82eaeb72 100644 --- a/src/lerobot/policies/factory.py +++ b/src/lerobot/policies/factory.py @@ -252,6 +252,7 @@ class ProcessorConfigKwargs(TypedDict, total=False): def make_pre_post_processors( policy_cfg: PreTrainedConfig, pretrained_path: str | None = None, + pretrained_revision: str | None = None, **kwargs: Unpack[ProcessorConfigKwargs], ) -> tuple[ PolicyProcessorPipeline[dict[str, Any], dict[str, Any]], @@ -309,6 +310,7 @@ def make_pre_post_processors( overrides=kwargs.get("preprocessor_overrides", {}), to_transition=batch_to_transition, to_output=transition_to_batch, + revision=pretrained_revision, ) postprocessor = PolicyProcessorPipeline.from_pretrained( pretrained_model_name_or_path=pretrained_path, @@ -318,6 +320,7 @@ def make_pre_post_processors( overrides=kwargs.get("postprocessor_overrides", {}), to_transition=policy_action_to_transition, to_output=transition_to_policy_action, + revision=pretrained_revision, ) _reconnect_relative_absolute_steps(preprocessor, postprocessor) return preprocessor, postprocessor @@ -557,6 +560,7 @@ def make_policy( # Load a pretrained policy and override the config if needed (for example, if there are inference-time # hyperparameters that we want to vary). kwargs["pretrained_name_or_path"] = cfg.pretrained_path + kwargs["revision"] = cfg.pretrained_revision policy = policy_cls.from_pretrained(**kwargs) elif cfg.pretrained_path and cfg.use_peft: # Load a pretrained PEFT model on top of the policy. The pretrained path points to the folder/repo diff --git a/src/lerobot/policies/molmoact2/README.md b/src/lerobot/policies/molmoact2/README.md index ef419516d..9756785d9 120000 --- a/src/lerobot/policies/molmoact2/README.md +++ b/src/lerobot/policies/molmoact2/README.md @@ -1 +1 @@ -../../../../docs/source/policy_molmoact2_README.md \ No newline at end of file +../../../../docs/source/molmoact2.mdx \ No newline at end of file diff --git a/src/lerobot/policies/molmoact2/__init__.py b/src/lerobot/policies/molmoact2/__init__.py index bfef53bb2..a4e7695c2 100644 --- a/src/lerobot/policies/molmoact2/__init__.py +++ b/src/lerobot/policies/molmoact2/__init__.py @@ -1,5 +1,3 @@ -#!/usr/bin/env python - # Copyright 2026 The Allen Institute for Artificial Intelligence and The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/src/lerobot/policies/molmoact2/configuration_molmoact2.py b/src/lerobot/policies/molmoact2/configuration_molmoact2.py index de2585281..bf9437ba9 100644 --- a/src/lerobot/policies/molmoact2/configuration_molmoact2.py +++ b/src/lerobot/policies/molmoact2/configuration_molmoact2.py @@ -1,5 +1,3 @@ -#!/usr/bin/env python - # Copyright 2026 The Allen Institute for Artificial Intelligence and The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -16,16 +14,9 @@ from __future__ import annotations -import json -import math -import os -from contextlib import suppress from dataclasses import dataclass, field -from pathlib import Path from typing import Any -from huggingface_hub import snapshot_download - from lerobot.configs import FeatureType, NormalizationMode, PolicyFeature, PreTrainedConfig from lerobot.optim import ( AdamWConfig, @@ -37,146 +28,6 @@ from lerobot.utils.constants import ACTION, OBS_STATE from ..rtc.configuration_rtc import RTCConfig -MOLMOACT2_DEFAULT_NUM_IMAGES = 2 -MOLMOACT2_IMAGE_TOKENS_PER_IMAGE = 196 -MOLMOACT2_FIXED_PROMPT_TOKEN_BUDGET = 80 -MOLMOACT2_TASK_TOKEN_BUDGET = 32 -MOLMOACT2_SEQUENCE_LENGTH_MARGIN = 32 -MOLMOACT2_SEQUENCE_LENGTH_MULTIPLE = 64 -MOLMOACT2_DISCRETE_ACTION_WRAPPER_TOKENS = 4 -MOLMOACT2_MIN_DISCRETE_ACTION_TOKENS_PER_STEP = 6 -MOLMOACT2_DISCRETE_ACTION_TOKENS_PER_DIM = 0.95 - - -def _hf_token() -> str | None: - return os.environ.get("HF_TOKEN") or os.environ.get("HF_ACCESS_TOKEN") - - -def _resolve_checkpoint_location( - checkpoint_path: str, - *, - revision: str | None = None, - force_download: bool = False, -) -> str: - checkpoint_path = str(checkpoint_path or "").strip() - if not checkpoint_path: - raise ValueError("MolmoAct2 policy requires `checkpoint_path`.") - local_path = Path(checkpoint_path).expanduser() - if local_path.exists(): - return str(local_path) - return snapshot_download( - repo_id=checkpoint_path, - repo_type="model", - revision=revision, - force_download=force_download, - ignore_patterns=["*.py", "*.pyc", "__pycache__/*"], - token=_hf_token(), - ) - - -def _load_hf_norm_metadata_for_tag( - checkpoint_path: str, - *, - revision: str | None, - force_download: bool, - norm_tag: str | None, -) -> dict[str, Any]: - norm_tag = str(norm_tag or "").strip() - if not norm_tag: - return {} - checkpoint_location = Path( - _resolve_checkpoint_location( - checkpoint_path, - revision=revision, - force_download=force_download, - ) - ) - norm_stats_filename = "norm_stats.json" - config_path = checkpoint_location / "config.json" - if config_path.exists(): - with suppress(OSError, json.JSONDecodeError): - norm_stats_filename = str( - json.loads(config_path.read_text()).get("norm_stats_filename") or norm_stats_filename - ) - stats_path = checkpoint_location / norm_stats_filename - if not stats_path.exists(): - raise FileNotFoundError( - f"MolmoAct2 HF checkpoint is missing {norm_stats_filename!r}; cannot resolve norm_tag={norm_tag!r}." - ) - payload = json.loads(stats_path.read_text()) - metadata_by_tag = payload.get("metadata_by_tag") - if not isinstance(metadata_by_tag, dict): - raise ValueError(f"MolmoAct2 norm stats file {stats_path} has no metadata_by_tag mapping.") - metadata = metadata_by_tag.get(norm_tag) - if not isinstance(metadata, dict): - available = sorted(str(tag) for tag in metadata_by_tag) - raise ValueError(f"Unknown MolmoAct2 norm_tag={norm_tag!r}. Available tags: {available}.") - return metadata - - -@LRSchedulerConfig.register_subclass("molmoact2_cosine_decay_with_warmup") -@dataclass -class MolmoAct2CosineDecayWithWarmupSchedulerConfig(CosineDecayWithWarmupSchedulerConfig): - """MolmoAct2-local cosine scheduler with optional decay-step auto-match. - - LeRobot's generic cosine scheduler keeps an explicit integer decay length. - For MolmoAct2, leaving num_decay_steps unset means "decay across this run's - training steps"; build() is the first point where num_training_steps is known. - """ - - num_decay_steps: int | None - - def build(self, optimizer, num_training_steps: int): - return CosineDecayWithWarmupSchedulerConfig( - peak_lr=self.peak_lr, - decay_lr=self.decay_lr, - num_warmup_steps=self.num_warmup_steps, - num_decay_steps=num_training_steps if self.num_decay_steps is None else self.num_decay_steps, - ).build(optimizer, num_training_steps=num_training_steps) - - -def _round_up(value: int, multiple: int) -> int: - return int(math.ceil(value / multiple) * multiple) - - -def infer_molmoact2_max_sequence_length( - *, - num_images: int, - state_dim: int, - action_dim: int, - action_horizon: int, - include_discrete_action: bool, -) -> int: - """Infer the padded text/image sequence cap from MolmoAct2's fixed token layout.""" - if num_images < 1: - num_images = MOLMOACT2_DEFAULT_NUM_IMAGES - if state_dim < 0: - state_dim = 0 - if action_dim < 1: - action_dim = 1 - if action_horizon < 1: - action_horizon = 1 - - image_tokens = num_images * MOLMOACT2_IMAGE_TOKENS_PER_IMAGE - prompt_tokens = ( - MOLMOACT2_FIXED_PROMPT_TOKEN_BUDGET - + MOLMOACT2_TASK_TOKEN_BUDGET - + state_dim - + MOLMOACT2_SEQUENCE_LENGTH_MARGIN - ) - action_tokens = 0 - if include_discrete_action: - action_tokens_per_step = max( - MOLMOACT2_MIN_DISCRETE_ACTION_TOKENS_PER_STEP, - math.ceil(action_dim * MOLMOACT2_DISCRETE_ACTION_TOKENS_PER_DIM), - ) - action_tokens = MOLMOACT2_DISCRETE_ACTION_WRAPPER_TOKENS + action_horizon * action_tokens_per_step - - return _round_up( - image_tokens + prompt_tokens + action_tokens, - MOLMOACT2_SEQUENCE_LENGTH_MULTIPLE, - ) - @PreTrainedConfig.register_subclass("molmoact2") @dataclass @@ -228,6 +79,15 @@ class MolmoAct2Config(PreTrainedConfig): eval_seed: int | None = None rtc_config: RTCConfig | None = None + # Joint frame transform for cross-calibration compatibility. + # Some MolmoAct2 checkpoints were trained on data using a different joint + # convention than the current LeRobot calibration. Set both to apply a + # sign/offset correction at runtime (state before model, action after). + # See: https://huggingface.co/docs/lerobot/backwardcomp + # Default is None (no transform). Both must be set together. + joint_signs: list[float] | None = None + joint_offsets: list[float] | None = None + # Default is full finetuning with gradients from the action expert flowing into the VLM. enable_lora_vlm: bool = False lora_rank: int = 64 @@ -255,7 +115,7 @@ class MolmoAct2Config(PreTrainedConfig): optimizer_grad_clip_norm: float = 1.0 scheduler_warmup_steps: int = 200 - scheduler_decay_steps: int | None = None + scheduler_decay_steps: int = 100_000 scheduler_decay_lr: float = 1e-6 normalization_mapping: dict[str, NormalizationMode] = field( @@ -272,6 +132,10 @@ class MolmoAct2Config(PreTrainedConfig): def __post_init__(self) -> None: super().__post_init__() + if (self.joint_signs is None) != (self.joint_offsets is None): + raise ValueError("joint_signs and joint_offsets must both be set or both be None.") + if self.joint_signs is not None and len(self.joint_signs) != len(self.joint_offsets): + raise ValueError("joint_signs and joint_offsets must have the same length.") if self.action_mode not in {"continuous", "discrete", "both"}: raise ValueError( f"Unsupported action_mode={self.action_mode!r}. " @@ -333,41 +197,6 @@ class MolmoAct2Config(PreTrainedConfig): if self.max_sequence_length is not None and self.max_sequence_length < 1: raise ValueError(f"max_sequence_length must be >= 1 or None, got {self.max_sequence_length}.") - def inferred_max_sequence_length( - self, - *, - num_images: int | None = None, - state_dim: int | None = None, - action_dim: int | None = None, - action_horizon: int | None = None, - include_discrete_action: bool | None = None, - ) -> int: - if self.max_sequence_length is not None: - return int(self.max_sequence_length) - - if num_images is None: - num_images = len(self.image_keys) or len(self.image_features) or MOLMOACT2_DEFAULT_NUM_IMAGES - if state_dim is None: - state_feature = self.robot_state_feature - state_dim = int(state_feature.shape[0]) if state_feature is not None else 0 - if action_dim is None: - action_feature = self.action_feature - action_dim = ( - int(action_feature.shape[0]) if action_feature is not None else self.expected_max_action_dim - ) - if action_horizon is None: - action_horizon = self.chunk_size - if include_discrete_action is None: - include_discrete_action = self.action_mode in {"discrete", "both"} - - return infer_molmoact2_max_sequence_length( - num_images=int(num_images), - state_dim=int(state_dim), - action_dim=int(action_dim), - action_horizon=int(action_horizon), - include_discrete_action=bool(include_discrete_action), - ) - @property def observation_delta_indices(self) -> None: return None @@ -390,7 +219,7 @@ class MolmoAct2Config(PreTrainedConfig): ) def get_scheduler_preset(self) -> LRSchedulerConfig | None: - return MolmoAct2CosineDecayWithWarmupSchedulerConfig( + return CosineDecayWithWarmupSchedulerConfig( peak_lr=self.optimizer_lr, decay_lr=self.scheduler_decay_lr, num_warmup_steps=self.scheduler_warmup_steps, @@ -426,94 +255,3 @@ class MolmoAct2Config(PreTrainedConfig): shape=(self.expected_max_action_dim,), ) self.output_features[ACTION] = action_feature - - def apply_norm_tag_metadata(self) -> None: - if not str(self.norm_tag or "").strip(): - return - metadata = _load_hf_norm_metadata_for_tag( - self.checkpoint_path, - revision=self.checkpoint_revision, - force_download=bool(self.checkpoint_force_download), - norm_tag=self.norm_tag, - ) - if metadata.get("action_horizon") is not None: - self.chunk_size = int(metadata["action_horizon"]) - if metadata.get("n_action_steps") is not None: - self.n_action_steps = int(metadata["n_action_steps"]) - if not self.setup_type and metadata.get("setup_type") is not None: - self.setup_type = str(metadata["setup_type"]) - if not self.control_mode and metadata.get("control_mode") is not None: - self.control_mode = str(metadata["control_mode"]) - - def saved_policy_action_mode(self) -> str | None: - pretrained_path = getattr(self, "pretrained_path", None) - if pretrained_path is None: - return None - config_path = Path(pretrained_path) / "config.json" - if not config_path.exists(): - return None - try: - mode = json.loads(config_path.read_text()).get("action_mode") - except (OSError, json.JSONDecodeError): - return None - if mode in {"continuous", "discrete", "both"}: - return str(mode) - return None - - def training_action_mode(self, saved_policy_action_mode: str | None = None) -> str: - return saved_policy_action_mode or self.action_mode - - def validate_inference_action_mode(self, saved_policy_action_mode: str | None = None) -> None: - requested_mode = self.inference_action_mode - if requested_mode is None: - return - training_mode = self.training_action_mode(saved_policy_action_mode) - if requested_mode == "continuous" and training_mode == "discrete": - raise ValueError( - "MolmoAct2 checkpoint was trained with action_mode='discrete' and cannot run " - "continuous inference." - ) - if requested_mode == "discrete" and training_mode == "continuous": - raise ValueError( - "MolmoAct2 checkpoint was trained with action_mode='continuous' and cannot run " - "discrete inference. Train with action_mode='both' or action_mode='discrete' first." - ) - - def validate_checkpoint_action_mode( - self, - checkpoint_action_mode: str, - *, - has_action_expert: bool, - ) -> None: - if self.action_mode == "both" and checkpoint_action_mode != "both": - raise ValueError( - f"action_mode='both' requires checkpoint action_mode='both', got {checkpoint_action_mode!r}." - ) - if self.action_mode == "discrete" and checkpoint_action_mode not in {"discrete", "both"}: - raise ValueError( - f"action_mode='discrete' requires checkpoint action_mode in {{'discrete', 'both'}}, " - f"got {checkpoint_action_mode!r}." - ) - if self.action_mode in {"continuous", "both"} and not has_action_expert: - raise ValueError("Continuous MolmoAct2 training requires an action expert checkpoint.") - - def resolve_inference_action_mode( - self, - requested_mode: str | None, - saved_policy_action_mode: str | None = None, - ) -> str: - training_mode = self.training_action_mode(saved_policy_action_mode) - if requested_mode is None: - requested_mode = self.inference_action_mode - if requested_mode is None: - raise ValueError( - "MolmoAct2 inference requires `inference_action_mode` to be set explicitly " - "to either 'continuous' or 'discrete'." - ) - if requested_mode not in {"continuous", "discrete"}: - raise ValueError("MolmoAct2 inference_action_mode must be either 'continuous' or 'discrete'.") - if requested_mode == "continuous" and training_mode == "discrete": - raise ValueError("MolmoAct2 action_mode='discrete' checkpoint cannot run continuous inference.") - if requested_mode == "discrete" and training_mode == "continuous": - raise ValueError("MolmoAct2 action_mode='continuous' checkpoint cannot run discrete inference.") - return requested_mode diff --git a/src/lerobot/policies/molmoact2/modeling_molmoact2.py b/src/lerobot/policies/molmoact2/modeling_molmoact2.py index f86be0904..2cc85ab02 100644 --- a/src/lerobot/policies/molmoact2/modeling_molmoact2.py +++ b/src/lerobot/policies/molmoact2/modeling_molmoact2.py @@ -1,5 +1,3 @@ -#!/usr/bin/env python - # Copyright 2026 The Allen Institute for Artificial Intelligence and The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -14,9 +12,22 @@ # See the License for the specific language governing permissions and # limitations under the License. +"""MolmoAct2 policy for LeRobot. + +MolmoAct2 is a VLM-based robotics policy from Allen AI that combines a +Molmo vision-language backbone with a per-layer flow-matching action expert +for continuous action generation, plus an optional discrete action token +head. This module wraps the vendored HF model implementation +(``molmoact2_hf_model/``) into the LeRobot ``PreTrainedPolicy`` interface. + +Paper: https://allenai.org/blog/molmoact2 +Code: https://github.com/allenai/molmoact2 +""" + from __future__ import annotations import json +import logging import os import types from collections import deque @@ -35,13 +46,58 @@ from lerobot.utils.constants import ACTION from lerobot.utils.import_utils import _scipy_available, _transformers_available, require_package from ..rtc.modeling_rtc import RTCProcessor -from .configuration_molmoact2 import MolmoAct2Config, _hf_token, _resolve_checkpoint_location +from .configuration_molmoact2 import MolmoAct2Config + +logger = logging.getLogger(__name__) + + +def _hf_token() -> str | None: + return os.environ.get("HF_TOKEN") or os.environ.get("HF_ACCESS_TOKEN") + + +def _resolve_checkpoint_location( + checkpoint_path: str, + *, + revision: str | None = None, + force_download: bool = False, +) -> str: + """Resolve a checkpoint path to a local directory, downloading from Hub if needed.""" + checkpoint_path = str(checkpoint_path or "").strip() + if not checkpoint_path: + raise ValueError("MolmoAct2 policy requires `checkpoint_path`.") + from pathlib import Path + + local_path = Path(checkpoint_path).expanduser() + if local_path.exists(): + return str(local_path) + from huggingface_hub import snapshot_download + + return snapshot_download( + repo_id=checkpoint_path, + repo_type="model", + revision=revision, + force_download=force_download, + ignore_patterns=["*.py", "*.pyc", "__pycache__/*"], + token=_hf_token(), + ) + + +def _torch_dtype(dtype: str) -> torch.dtype: + """Convert a dtype name string to a torch.dtype.""" + if dtype == "float32": + return torch.float32 + if dtype == "bfloat16": + return torch.bfloat16 + if dtype == "float16": + return torch.float16 + raise ValueError(f"Unsupported dtype: {dtype}") + if TYPE_CHECKING or _transformers_available: from transformers.utils import SAFE_WEIGHTS_INDEX_NAME, SAFE_WEIGHTS_NAME - from .hf_model.configuration_molmoact2 import MolmoAct2Config as HFMolmoAct2Config - from .hf_model.modeling_molmoact2 import MolmoAct2ForConditionalGeneration + from .molmoact2_hf_model.configuration_molmoact2 import MolmoAct2Config as HFMolmoAct2Config + from .molmoact2_hf_model.modeling_molmoact2 import MolmoAct2ForConditionalGeneration else: SAFE_WEIGHTS_INDEX_NAME = "model.safetensors.index.json" SAFE_WEIGHTS_NAME = "model.safetensors" @@ -49,7 +105,7 @@ else: MolmoAct2ForConditionalGeneration = None if TYPE_CHECKING or (_transformers_available and _scipy_available): - from .hf_model.action_tokenizer import UniversalActionProcessor + from .molmoact2_hf_model.action_tokenizer import UniversalActionProcessor else: UniversalActionProcessor = None @@ -70,6 +126,156 @@ _MODEL_INPUT_KEYS = { } +def _load_hf_norm_metadata_for_tag( + checkpoint_path: str, + *, + revision: str | None, + force_download: bool, + norm_tag: str | None, +) -> dict[str, Any]: + """Read per-tag metadata from the checkpoint's ``norm_stats.json``.""" + norm_tag = str(norm_tag or "").strip() + if not norm_tag: + return {} + from contextlib import suppress + from pathlib import Path + + checkpoint_location = Path( + _resolve_checkpoint_location( + checkpoint_path, + revision=revision, + force_download=force_download, + ) + ) + norm_stats_filename = "norm_stats.json" + config_path = checkpoint_location / "config.json" + if config_path.exists(): + with suppress(OSError, json.JSONDecodeError): + norm_stats_filename = str( + json.loads(config_path.read_text()).get("norm_stats_filename") or norm_stats_filename + ) + stats_path = checkpoint_location / norm_stats_filename + if not stats_path.exists(): + raise FileNotFoundError( + f"MolmoAct2 HF checkpoint is missing {norm_stats_filename!r}; cannot resolve norm_tag={norm_tag!r}." + ) + payload = json.loads(stats_path.read_text()) + metadata_by_tag = payload.get("metadata_by_tag") + if not isinstance(metadata_by_tag, dict): + raise ValueError(f"MolmoAct2 norm stats file {stats_path} has no metadata_by_tag mapping.") + metadata = metadata_by_tag.get(norm_tag) + if not isinstance(metadata, dict): + available = sorted(str(tag) for tag in metadata_by_tag) + raise ValueError(f"Unknown MolmoAct2 norm_tag={norm_tag!r}. Available tags: {available}.") + return metadata + + +def _apply_norm_tag_metadata(config: MolmoAct2Config) -> None: + """Populate config fields from the checkpoint's norm-tag metadata.""" + if not str(config.norm_tag or "").strip(): + return + metadata = _load_hf_norm_metadata_for_tag( + config.checkpoint_path, + revision=config.checkpoint_revision, + force_download=bool(config.checkpoint_force_download), + norm_tag=config.norm_tag, + ) + if metadata.get("action_horizon") is not None: + config.chunk_size = int(metadata["action_horizon"]) + if metadata.get("n_action_steps") is not None: + config.n_action_steps = int(metadata["n_action_steps"]) + if not config.setup_type and metadata.get("setup_type") is not None: + config.setup_type = str(metadata["setup_type"]) + if not config.control_mode and metadata.get("control_mode") is not None: + config.control_mode = str(metadata["control_mode"]) + + +def _saved_policy_action_mode(config: MolmoAct2Config) -> str | None: + """Read the action mode from a LeRobot-saved checkpoint's ``config.json``.""" + from pathlib import Path + + pretrained_path = getattr(config, "pretrained_path", None) + if pretrained_path is None: + return None + config_path = Path(pretrained_path) / "config.json" + if not config_path.exists(): + return None + try: + mode = json.loads(config_path.read_text()).get("action_mode") + except (OSError, json.JSONDecodeError): + return None + if mode in {"continuous", "discrete", "both"}: + return str(mode) + return None + + +def _training_action_mode(config: MolmoAct2Config, saved_policy_action_mode: str | None = None) -> str: + return saved_policy_action_mode or config.action_mode + + +def _validate_inference_action_mode( + config: MolmoAct2Config, saved_policy_action_mode: str | None = None +) -> None: + """Check that the requested inference mode is compatible with the training mode.""" + requested_mode = config.inference_action_mode + if requested_mode is None: + return + training_mode = _training_action_mode(config, saved_policy_action_mode) + if requested_mode == "continuous" and training_mode == "discrete": + raise ValueError( + "MolmoAct2 checkpoint was trained with action_mode='discrete' and cannot run " + "continuous inference." + ) + if requested_mode == "discrete" and training_mode == "continuous": + raise ValueError( + "MolmoAct2 checkpoint was trained with action_mode='continuous' and cannot run " + "discrete inference. Train with action_mode='both' or action_mode='discrete' first." + ) + + +def _validate_checkpoint_action_mode( + config: MolmoAct2Config, + checkpoint_action_mode: str, + *, + has_action_expert: bool, +) -> None: + """Check that the checkpoint's action mode is compatible with the config.""" + if config.action_mode == "both" and checkpoint_action_mode != "both": + raise ValueError( + f"action_mode='both' requires checkpoint action_mode='both', got {checkpoint_action_mode!r}." + ) + if config.action_mode == "discrete" and checkpoint_action_mode not in {"discrete", "both"}: + raise ValueError( + f"action_mode='discrete' requires checkpoint action_mode in {{'discrete', 'both'}}, " + f"got {checkpoint_action_mode!r}." + ) + if config.action_mode in {"continuous", "both"} and not has_action_expert: + raise ValueError("Continuous MolmoAct2 training requires an action expert checkpoint.") + + +def _resolve_inference_action_mode( + config: MolmoAct2Config, + requested_mode: str | None, + saved_policy_action_mode: str | None = None, +) -> str: + """Resolve the final inference action mode, validating compatibility.""" + training_mode = _training_action_mode(config, saved_policy_action_mode) + if requested_mode is None: + requested_mode = config.inference_action_mode + if requested_mode is None: + raise ValueError( + "MolmoAct2 inference requires `inference_action_mode` to be set explicitly " + "to either 'continuous' or 'discrete'." + ) + if requested_mode not in {"continuous", "discrete"}: + raise ValueError("MolmoAct2 inference_action_mode must be either 'continuous' or 'discrete'.") + if requested_mode == "continuous" and training_mode == "discrete": + raise ValueError("MolmoAct2 action_mode='discrete' checkpoint cannot run continuous inference.") + if requested_mode == "discrete" and training_mode == "continuous": + raise ValueError("MolmoAct2 action_mode='continuous' checkpoint cannot run discrete inference.") + return requested_mode + + def _strict_load_safetensors_weights(model: torch.nn.Module, checkpoint_location: str) -> None: index_path = os.path.join(checkpoint_location, SAFE_WEIGHTS_INDEX_NAME) single_file_path = os.path.join(checkpoint_location, SAFE_WEIGHTS_NAME) @@ -103,16 +309,6 @@ def _strict_load_safetensors_weights(model: torch.nn.Module, checkpoint_location ) -def _torch_dtype(dtype: str) -> torch.dtype: - if dtype == "float32": - return torch.float32 - if dtype == "bfloat16": - return torch.bfloat16 - if dtype == "float16": - return torch.float16 - raise ValueError(f"Unsupported dtype: {dtype}") - - def _sample_beta_timesteps( *, batch_size: int, @@ -136,7 +332,180 @@ def _sample_beta_timesteps( return time_offset + scale * samples +def _mask_discrete_action_spans( + *, + input_ids: Tensor, + mask: Tensor, + start_token_id: int | None, + end_token_id: int | None, +) -> Tensor: + if start_token_id is None or end_token_id is None: + return mask + mask = mask.clone() + for batch_idx in range(input_ids.shape[0]): + row = input_ids[batch_idx] + starts = (row == int(start_token_id)).nonzero(as_tuple=False).flatten().tolist() + ends = (row == int(end_token_id)).nonzero(as_tuple=False).flatten().tolist() + end_ptr = 0 + for start in starts: + while end_ptr < len(ends) and ends[end_ptr] < start: + end_ptr += 1 + if end_ptr >= len(ends): + mask[batch_idx, start:] = False + break + end = int(ends[end_ptr]) + mask[batch_idx, start : end + 1] = False + end_ptr += 1 + return mask + + +def _drop_trivial_attention_mask(model_inputs: dict[str, Tensor]) -> dict[str, Tensor]: + attention_mask = model_inputs.get("attention_mask") + if torch.is_tensor(attention_mask) and bool(attention_mask.to(dtype=torch.bool).all().item()): + model_inputs = dict(model_inputs) + model_inputs.pop("attention_mask", None) + return model_inputs + + +def _expand_mask(mask: Tensor | None, num_flow_timesteps: int) -> Tensor | None: + if mask is None: + return None + return ( + mask.unsqueeze(1) + .expand(-1, num_flow_timesteps, *([-1] * (mask.ndim - 1))) + .reshape(mask.shape[0] * num_flow_timesteps, *mask.shape[1:]) + ) + + +def _action_dim_valid_mask(target: Tensor, action_dim_is_pad: Tensor | None) -> Tensor | None: + if action_dim_is_pad is None: + return None + mask = ~action_dim_is_pad.to(device=target.device, dtype=torch.bool) + if mask.ndim == 1: + mask = mask.unsqueeze(0) + if mask.shape[-1] != target.shape[-1]: + raise ValueError( + f"action_dim_is_pad width {mask.shape[-1]} does not match target width {target.shape[-1]}." + ) + if mask.shape[0] == 1 and target.shape[0] != 1: + mask = mask.expand(target.shape[0], -1) + if mask.shape[0] != target.shape[0]: + raise ValueError( + f"action_dim_is_pad batch {mask.shape[0]} does not match target batch {target.shape[0]}." + ) + while mask.ndim < target.ndim: + mask = mask.unsqueeze(1) + return mask + + +def _mask_action_dim_tensor(tensor: Tensor, action_dim_is_pad: Tensor | None) -> Tensor: + if action_dim_is_pad is None: + return tensor + valid_mask = _action_dim_valid_mask(tensor, action_dim_is_pad) + if valid_mask is None: + return tensor + return tensor.masked_fill(~valid_mask, 0) + + +def _apply_action_dim_padding_mask(loss: Tensor, action_dim_is_pad: Tensor | None) -> Tensor: + valid_mask = _action_dim_valid_mask(loss, action_dim_is_pad) + if valid_mask is None: + return loss + valid = valid_mask.to(dtype=loss.dtype) + denom = valid.sum(dim=-1).clamp_min(1.0) + return (loss * valid).sum(dim=-1) / denom + + +def _apply_action_chunk_padding_mask(loss: Tensor, action_horizon_is_pad: Tensor | None) -> Tensor: + if action_horizon_is_pad is None: + return loss + valid_action = ( + (~action_horizon_is_pad.to(device=loss.device, dtype=torch.bool)).unsqueeze(1).unsqueeze(-1) + ) + return loss * valid_action + + +def _combine_rollout_seeds(first_seed: int, batch_size: int) -> int: + seed = 0 + for idx in range(batch_size): + seed = (seed + (idx + 1) * (first_seed + idx)) % (2**63 - 1) + return seed + + +def _rollout_task_signature(batch: dict[str, Any]) -> tuple[Any, ...] | None: + task = batch.get("task") + if task is None: + task = batch.get("observation.language") + if task is None: + return None + if isinstance(task, str): + return (task,) + if isinstance(task, (list, tuple)): + return tuple(str(item) for item in task) + return (str(task),) + + +def _extract_discrete_token_bins( + generated_ids: list[int], + start_token_id: int, + end_token_id: int, + token_id_to_bin: dict[int, int], +) -> list[int]: + start_idx = None + end_idx = None + for idx, token_id in enumerate(generated_ids): + if token_id == start_token_id: + start_idx = idx + break + if start_idx is not None: + for idx in range(start_idx + 1, len(generated_ids)): + if generated_ids[idx] == end_token_id: + end_idx = idx + break + span_start = 0 if start_idx is None else start_idx + 1 + span_end = len(generated_ids) if end_idx is None else end_idx + return [ + int(token_id_to_bin[token_id]) + for token_id in generated_ids[span_start:span_end] + if token_id in token_id_to_bin + ] + + +def _weighted_mean(values: Tensor, weights: Tensor | None) -> Tensor: + if weights is None: + return values.mean() + weights = weights.to(device=values.device, dtype=values.dtype) + return torch.dot(values, weights) / weights.sum().clamp_min(1.0) + + +def _weighted_per_example( + values: Tensor, + weights: Tensor | None, + example_indices: Tensor, + batch_size: int, +) -> Tensor: + values = values.float() + if weights is None: + weights = torch.ones_like(values) + else: + weights = weights.to(device=values.device, dtype=values.dtype) + loss_sum = torch.zeros(batch_size, device=values.device, dtype=torch.float32) + weight_sum = torch.zeros(batch_size, device=values.device, dtype=torch.float32) + loss_sum.scatter_add_(0, example_indices, values * weights) + weight_sum.scatter_add_(0, example_indices, weights) + global_weight_sum = weight_sum.sum().clamp_min(1.0) + return loss_sum * float(batch_size) / global_weight_sum + + class MolmoAct2Policy(PreTrainedPolicy): + """MolmoAct2 policy wrapping the vendored HF model for LeRobot. + + Supports three training modes via ``config.action_mode``: + ``"continuous"`` (flow-matching only), ``"discrete"`` (autoregressive + token prediction only), or ``"both"`` (joint loss). At inference, + ``config.inference_action_mode`` selects which head generates actions. + """ + config_class = MolmoAct2Config name = "molmoact2" @@ -149,10 +518,10 @@ class MolmoAct2Policy(PreTrainedPolicy): **kwargs, ): super().__init__(config, *inputs, **kwargs) - self.config.apply_norm_tag_metadata() + _apply_norm_tag_metadata(self.config) self.config.validate_features() del inputs, kwargs, dataset_stats, dataset_meta - self._checkpoint_action_mode = self.config.saved_policy_action_mode() + self._checkpoint_action_mode = _saved_policy_action_mode(self.config) self._action_queue: deque[Tensor] = deque(maxlen=self.config.n_action_steps) self._rollout_action_generator: torch.Generator | None = None self._rollout_task_key: tuple[Any, ...] | None = None @@ -160,7 +529,7 @@ class MolmoAct2Policy(PreTrainedPolicy): self.rtc_processor: RTCProcessor | None = None self.action_tokenizer: Any | None = None self._load_hf_model() - self.config.validate_inference_action_mode(self._checkpoint_action_mode) + _validate_inference_action_mode(self.config, self._checkpoint_action_mode) if self.config.enable_lora_vlm: self._apply_lora_adapters() self.init_rtc_processor() @@ -212,7 +581,8 @@ class MolmoAct2Policy(PreTrainedPolicy): "`policy.checkpoint_force_download=true` after the updated files are pushed." ) checkpoint_action_mode = str(self.model.config.action_mode) - self.config.validate_checkpoint_action_mode( + _validate_checkpoint_action_mode( + self.config, checkpoint_action_mode, has_action_expert=bool(getattr(self.model.config, "add_action_expert", False)), ) @@ -226,6 +596,7 @@ class MolmoAct2Policy(PreTrainedPolicy): self.train(self.training) def reset(self) -> None: + """Clear the action queue and rollout generator between episodes.""" self._action_queue = deque(maxlen=self.config.n_action_steps) self._rollout_action_generator = None @@ -334,6 +705,7 @@ class MolmoAct2Policy(PreTrainedPolicy): param.requires_grad = False def get_optim_params(self) -> list[dict[str, Any]]: + """Return optimizer param groups with per-component learning rates.""" vit_params: list[Tensor] = [] connector_params: list[Tensor] = [] action_expert_params: list[Tensor] = [] @@ -419,33 +791,6 @@ class MolmoAct2Policy(PreTrainedPolicy): return int(value) raise RuntimeError("MolmoAct2 could not resolve an action generation horizon.") - @staticmethod - def _mask_discrete_action_spans( - *, - input_ids: Tensor, - mask: Tensor, - start_token_id: int | None, - end_token_id: int | None, - ) -> Tensor: - if start_token_id is None or end_token_id is None: - return mask - mask = mask.clone() - for batch_idx in range(input_ids.shape[0]): - row = input_ids[batch_idx] - starts = (row == int(start_token_id)).nonzero(as_tuple=False).flatten().tolist() - ends = (row == int(end_token_id)).nonzero(as_tuple=False).flatten().tolist() - end_ptr = 0 - for start in starts: - while end_ptr < len(ends) and ends[end_ptr] < start: - end_ptr += 1 - if end_ptr >= len(ends): - mask[batch_idx, start:] = False - break - end = int(ends[end_ptr]) - mask[batch_idx, start : end + 1] = False - end_ptr += 1 - return mask - def _encoder_attention_mask_for_action_expert( self, *, @@ -470,21 +815,13 @@ class MolmoAct2Policy(PreTrainedPolicy): eos_token_id = getattr(self.model.config, "eos_token_id", None) if eos_token_id is not None: mask &= input_ids != int(eos_token_id) - return self._mask_discrete_action_spans( + return _mask_discrete_action_spans( input_ids=input_ids, mask=mask, start_token_id=getattr(self.model.config, "action_start_token_id", None), end_token_id=getattr(self.model.config, "action_end_token_id", None), ) - @staticmethod - def _drop_trivial_attention_mask(model_inputs: dict[str, Tensor]) -> dict[str, Tensor]: - attention_mask = model_inputs.get("attention_mask") - if torch.is_tensor(attention_mask) and bool(attention_mask.to(dtype=torch.bool).all().item()): - model_inputs = dict(model_inputs) - model_inputs.pop("attention_mask", None) - return model_inputs - def _load_discrete_action_tokenizer(self) -> Any: if self.action_tokenizer is None: require_package("transformers", extra="molmoact2") @@ -498,27 +835,7 @@ class MolmoAct2Policy(PreTrainedPolicy): return self.action_tokenizer def _resolve_inference_action_mode(self, requested_mode: str | None) -> str: - return self.config.resolve_inference_action_mode(requested_mode, self._checkpoint_action_mode) - - @staticmethod - def _combine_rollout_seeds(first_seed: int, batch_size: int) -> int: - seed = 0 - for idx in range(batch_size): - seed = (seed + (idx + 1) * (first_seed + idx)) % (2**63 - 1) - return seed - - @staticmethod - def _rollout_task_signature(batch: dict[str, Any]) -> tuple[Any, ...] | None: - task = batch.get("task") - if task is None: - task = batch.get("observation.language") - if task is None: - return None - if isinstance(task, str): - return (task,) - if isinstance(task, (list, tuple)): - return tuple(str(item) for item in task) - return (str(task),) + return _resolve_inference_action_mode(self.config, requested_mode, self._checkpoint_action_mode) def _rollout_generator_for_inputs( self, @@ -532,7 +849,7 @@ class MolmoAct2Policy(PreTrainedPolicy): if self._rollout_action_generator is not None: return self._rollout_action_generator - task_signature = self._rollout_task_signature(batch) + task_signature = _rollout_task_signature(batch) if task_signature != self._rollout_task_key: self._rollout_task_key = task_signature self._rollout_index_for_task = 0 @@ -545,72 +862,10 @@ class MolmoAct2Policy(PreTrainedPolicy): device if device.type == "cuda" and torch.cuda.is_available() else torch.device("cpu") ) generator = torch.Generator(device=generator_device) - generator.manual_seed(self._combine_rollout_seeds(first_seed, batch_size)) + generator.manual_seed(_combine_rollout_seeds(first_seed, batch_size)) self._rollout_action_generator = generator return generator - @staticmethod - def _expand_mask(mask: Tensor | None, num_flow_timesteps: int) -> Tensor | None: - if mask is None: - return None - return ( - mask.unsqueeze(1) - .expand(-1, num_flow_timesteps, *([-1] * (mask.ndim - 1))) - .reshape(mask.shape[0] * num_flow_timesteps, *mask.shape[1:]) - ) - - @staticmethod - def _action_dim_valid_mask(target: Tensor, action_dim_is_pad: Tensor | None) -> Tensor | None: - if action_dim_is_pad is None: - return None - mask = ~action_dim_is_pad.to(device=target.device, dtype=torch.bool) - if mask.ndim == 1: - mask = mask.unsqueeze(0) - if mask.shape[-1] != target.shape[-1]: - raise ValueError( - f"action_dim_is_pad width {mask.shape[-1]} does not match target width {target.shape[-1]}." - ) - if mask.shape[0] == 1 and target.shape[0] != 1: - mask = mask.expand(target.shape[0], -1) - if mask.shape[0] != target.shape[0]: - raise ValueError( - f"action_dim_is_pad batch {mask.shape[0]} does not match target batch {target.shape[0]}." - ) - while mask.ndim < target.ndim: - mask = mask.unsqueeze(1) - return mask - - @classmethod - def _mask_action_dim_tensor(cls, tensor: Tensor, action_dim_is_pad: Tensor | None) -> Tensor: - if not cls._mask_enabled_static(action_dim_is_pad): - return tensor - valid_mask = cls._action_dim_valid_mask(tensor, action_dim_is_pad) - if valid_mask is None: - return tensor - return tensor.masked_fill(~valid_mask, 0) - - @staticmethod - def _mask_enabled_static(action_dim_is_pad: Tensor | None) -> bool: - return action_dim_is_pad is not None - - @classmethod - def _apply_action_dim_padding_mask(cls, loss: Tensor, action_dim_is_pad: Tensor | None) -> Tensor: - valid_mask = cls._action_dim_valid_mask(loss, action_dim_is_pad) - if valid_mask is None: - return loss - valid = valid_mask.to(dtype=loss.dtype) - denom = valid.sum(dim=-1).clamp_min(1.0) - return (loss * valid).sum(dim=-1) / denom - - @staticmethod - def _apply_action_chunk_padding_mask(loss: Tensor, action_horizon_is_pad: Tensor | None) -> Tensor: - if action_horizon_is_pad is None: - return loss - valid_action = ( - (~action_horizon_is_pad.to(device=loss.device, dtype=torch.bool)).unsqueeze(1).unsqueeze(-1) - ) - return loss * valid_action - def _prepare_flow_matching_tensors( self, *, @@ -649,7 +904,7 @@ class MolmoAct2Policy(PreTrainedPolicy): ) if self.config.mask_action_dim_padding: - actions = self._mask_action_dim_tensor(actions, action_dim_is_pad) + actions = _mask_action_dim_tensor(actions, action_dim_is_pad) expected_noise_shape = (batch_size, num_flow_timesteps, actions.shape[1], actions.shape[2]) if noise is None: @@ -661,7 +916,7 @@ class MolmoAct2Policy(PreTrainedPolicy): f"flow noise must have shape {expected_noise_shape}, got {tuple(noise.shape)}." ) if self.config.mask_action_dim_padding: - noise = self._mask_action_dim_tensor(noise, action_dim_is_pad) + noise = _mask_action_dim_tensor(noise, action_dim_is_pad) t_broadcast = timesteps.view(batch_size, num_flow_timesteps, 1, 1) actions_expanded = actions.unsqueeze(1).expand(-1, num_flow_timesteps, -1, -1) @@ -789,7 +1044,7 @@ class MolmoAct2Policy(PreTrainedPolicy): valid_action = None if action_attention_mask is not None: valid_action = action_attention_mask.to(device=device, dtype=actions.dtype).unsqueeze(-1) - valid_action = self._expand_mask(valid_action, num_flow_timesteps) + valid_action = _expand_mask(valid_action, num_flow_timesteps) rope_cache = None if len(action_expert.blocks) > 0 and action_expert.blocks[0].self_attn.rope is not None: @@ -804,14 +1059,14 @@ class MolmoAct2Policy(PreTrainedPolicy): batch_size, actions.dtype, ) - cross_mask = self._expand_mask(cross_mask, num_flow_timesteps) + cross_mask = _expand_mask(cross_mask, num_flow_timesteps) self_mask = action_expert._build_self_attention_mask( action_attention_mask, actions.shape[1], device, actions.dtype, ) - self_mask = self._expand_mask(self_mask, num_flow_timesteps) + self_mask = _expand_mask(self_mask, num_flow_timesteps) conditioning = self._action_time_conditioning(action_expert, timesteps_flat) action_hidden = action_expert.action_embed(xt_flat) @@ -871,8 +1126,8 @@ class MolmoAct2Policy(PreTrainedPolicy): if k_norm is not None: k_ctx = k_norm(k_ctx.transpose(1, 2)).transpose(1, 2) if num_flow_timesteps != 1: - k_ctx = self._expand_mask(k_ctx, num_flow_timesteps) - v_ctx = self._expand_mask(v_ctx, num_flow_timesteps) + k_ctx = _expand_mask(k_ctx, num_flow_timesteps) + v_ctx = _expand_mask(v_ctx, num_flow_timesteps) next_action_hidden = action_block( layer_action_hidden, @@ -912,9 +1167,9 @@ class MolmoAct2Policy(PreTrainedPolicy): ) loss = F.mse_loss(pred_velocity, target_velocity, reduction="none") - loss = self._apply_action_chunk_padding_mask(loss, batch.get("action_horizon_is_pad")) + loss = _apply_action_chunk_padding_mask(loss, batch.get("action_horizon_is_pad")) if self.config.mask_action_dim_padding: - loss = self._apply_action_dim_padding_mask(loss, batch.get("action_dim_is_pad")) + loss = _apply_action_dim_padding_mask(loss, batch.get("action_dim_is_pad")) loss = loss.reshape(batch_size, -1).mean(dim=1) if reduction == "mean": loss = loss.mean() @@ -933,32 +1188,6 @@ class MolmoAct2Policy(PreTrainedPolicy): example_weights[nonempty] = 2.0 / torch.sqrt(token_counts[nonempty]) return example_weights[:, None].expand_as(valid_positions)[valid_positions].to(dtype=torch.float32) - @staticmethod - def _weighted_mean(values: Tensor, weights: Tensor | None) -> Tensor: - if weights is None: - return values.mean() - weights = weights.to(device=values.device, dtype=values.dtype) - return torch.dot(values, weights) / weights.sum().clamp_min(1.0) - - @staticmethod - def _weighted_per_example( - values: Tensor, - weights: Tensor | None, - example_indices: Tensor, - batch_size: int, - ) -> Tensor: - values = values.float() - if weights is None: - weights = torch.ones_like(values) - else: - weights = weights.to(device=values.device, dtype=values.dtype) - loss_sum = torch.zeros(batch_size, device=values.device, dtype=torch.float32) - weight_sum = torch.zeros(batch_size, device=values.device, dtype=torch.float32) - loss_sum.scatter_add_(0, example_indices, values * weights) - weight_sum.scatter_add_(0, example_indices, weights) - global_weight_sum = weight_sum.sum().clamp_min(1.0) - return loss_sum * float(batch_size) / global_weight_sum - def _discrete_loss_from_backbone_outputs( self, batch: dict[str, Tensor], @@ -992,56 +1221,28 @@ class MolmoAct2Policy(PreTrainedPolicy): token_weights = self._discrete_token_weights(valid_positions) if reduction == "none": example_indices = valid_positions.nonzero(as_tuple=False)[:, 0].to(device=hidden_states.device) - ce_loss = self._weighted_per_example( + ce_loss = _weighted_per_example( token_ce_loss, token_weights, example_indices, int(labels.shape[0]), ) else: - ce_loss = self._weighted_mean(token_ce_loss, token_weights) + ce_loss = _weighted_mean(token_ce_loss, token_weights) if not self.config.softmax_auxiliary_loss: return ce_loss, None if reduction == "none": - z_loss = self.config.softmax_auxiliary_loss_scale * self._weighted_per_example( + z_loss = self.config.softmax_auxiliary_loss_scale * _weighted_per_example( log_z.pow(2), token_weights, example_indices, int(labels.shape[0]), ) else: - z_loss = self.config.softmax_auxiliary_loss_scale * self._weighted_mean( - log_z.pow(2), token_weights - ) + z_loss = self.config.softmax_auxiliary_loss_scale * _weighted_mean(log_z.pow(2), token_weights) return ce_loss, z_loss - @staticmethod - def _extract_discrete_token_bins( - generated_ids: list[int], - start_token_id: int, - end_token_id: int, - token_id_to_bin: dict[int, int], - ) -> list[int]: - start_idx = None - end_idx = None - for idx, token_id in enumerate(generated_ids): - if token_id == start_token_id: - start_idx = idx - break - if start_idx is not None: - for idx in range(start_idx + 1, len(generated_ids)): - if generated_ids[idx] == end_token_id: - end_idx = idx - break - span_start = 0 if start_idx is None else start_idx + 1 - span_end = len(generated_ids) if end_idx is None else end_idx - return [ - int(token_id_to_bin[token_id]) - for token_id in generated_ids[span_start:span_end] - if token_id in token_id_to_bin - ] - def _action_token_id_to_bin(self) -> dict[int, int]: method = getattr(self.model, "_action_token_id_to_bin", None) if callable(method): @@ -1179,7 +1380,7 @@ class MolmoAct2Policy(PreTrainedPolicy): chunks: list[Tensor] = [] for token_row in generated_token_ids: generated_ids = [int(token_id) for token_id in token_row.detach().cpu().tolist()] - discrete_token_ids = self._extract_discrete_token_bins( + discrete_token_ids = _extract_discrete_token_bins( generated_ids, int(self.model.config.action_start_token_id), int(self.model.config.action_end_token_id), @@ -1218,7 +1419,7 @@ class MolmoAct2Policy(PreTrainedPolicy): model_inputs: dict[str, Tensor], action_dim: int, ) -> Tensor: - model_inputs = self._drop_trivial_attention_mask(model_inputs) + model_inputs = _drop_trivial_attention_mask(model_inputs) max_steps = self._discrete_generation_max_steps() static_cache, attention_bias = self._make_discrete_ar_graph_decode_inputs( model_inputs, @@ -1294,7 +1495,7 @@ class MolmoAct2Policy(PreTrainedPolicy): generator=generator, ) if self.config.mask_action_dim_padding: - trajectory = self._mask_action_dim_tensor(trajectory, action_dim_is_pad) + trajectory = _mask_action_dim_tensor(trajectory, action_dim_is_pad) action_context = action_expert.prepare_context( encoder_kv_states=encoder_kv_states, @@ -1327,7 +1528,7 @@ class MolmoAct2Policy(PreTrainedPolicy): modulation=step_modulation, ) if mask_enabled: - velocity = self._mask_action_dim_tensor(velocity, action_dim_is_pad) + velocity = _mask_action_dim_tensor(velocity, action_dim_is_pad) return velocity if self._rtc_enabled(): @@ -1352,7 +1553,7 @@ class MolmoAct2Policy(PreTrainedPolicy): trajectory = trajectory + dt * velocity if mask_enabled: - trajectory = self._mask_action_dim_tensor(trajectory, action_dim_is_pad) + trajectory = _mask_action_dim_tensor(trajectory, action_dim_is_pad) if self.rtc_processor is not None and self.rtc_processor.is_debug_enabled(): self.rtc_processor.track(time=float(flow_timestep[0].item()), x_t=trajectory, v_t=velocity) @@ -1363,6 +1564,7 @@ class MolmoAct2Policy(PreTrainedPolicy): batch: dict[str, Tensor], reduction: str = "mean", ) -> tuple[Tensor, dict[str, Any]]: + """Compute training loss (flow-matching and/or discrete token loss).""" if reduction not in {"mean", "none"}: raise ValueError(f"Unsupported reduction={reduction!r}. Expected 'mean' or 'none'.") model_inputs = self._model_inputs(batch) @@ -1422,6 +1624,7 @@ class MolmoAct2Policy(PreTrainedPolicy): @torch.no_grad() def predict_action_chunk(self, batch: dict[str, Tensor], **kwargs) -> Tensor: + """Generate an action chunk via continuous flow matching or discrete AR decoding.""" if "action_mode" in kwargs: raise TypeError( "MolmoAct2 predict_action_chunk got unexpected keyword argument 'action_mode'; " @@ -1476,6 +1679,7 @@ class MolmoAct2Policy(PreTrainedPolicy): @torch.no_grad() def select_action(self, batch: dict[str, Tensor], **kwargs) -> Tensor: + """Pop one action step from the queue, regenerating the chunk when empty.""" if self._rtc_enabled(): raise AssertionError("RTC is not supported for select_action, use it with predict_action_chunk") self.eval() diff --git a/src/lerobot/policies/molmoact2/hf_model/__init__.py b/src/lerobot/policies/molmoact2/molmoact2_hf_model/__init__.py similarity index 94% rename from src/lerobot/policies/molmoact2/hf_model/__init__.py rename to src/lerobot/policies/molmoact2/molmoact2_hf_model/__init__.py index 39b15cb3a..4436c9fda 100644 --- a/src/lerobot/policies/molmoact2/hf_model/__init__.py +++ b/src/lerobot/policies/molmoact2/molmoact2_hf_model/__init__.py @@ -1,5 +1,3 @@ -#!/usr/bin/env python - # Copyright 2026 The Allen Institute for Artificial Intelligence and The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -13,5 +11,3 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. - -# ruff: noqa diff --git a/src/lerobot/policies/molmoact2/hf_model/action_tokenizer.py b/src/lerobot/policies/molmoact2/molmoact2_hf_model/action_tokenizer.py similarity index 96% rename from src/lerobot/policies/molmoact2/hf_model/action_tokenizer.py rename to src/lerobot/policies/molmoact2/molmoact2_hf_model/action_tokenizer.py index f7dacbce6..11a228731 100644 --- a/src/lerobot/policies/molmoact2/hf_model/action_tokenizer.py +++ b/src/lerobot/policies/molmoact2/molmoact2_hf_model/action_tokenizer.py @@ -1,5 +1,3 @@ -#!/usr/bin/env python - # Copyright 2026 The Allen Institute for Artificial Intelligence and The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -14,23 +12,19 @@ # See the License for the specific language governing permissions and # limitations under the License. -# ruff: noqa - import logging -import os from pathlib import Path from typing import ClassVar import numpy as np from tokenizers import ByteLevelBPETokenizer from tokenizers.trainers import BpeTrainer -from huggingface_hub import snapshot_download from transformers import PreTrainedTokenizerFast from transformers.processing_utils import ProcessorMixin +from ..modeling_molmoact2 import _hf_token -def _hf_token() -> str | None: - return os.environ.get("HF_TOKEN") or os.environ.get("HF_ACCESS_TOKEN") +logger = logging.getLogger(__name__) def _resolve_tokenizer_location( @@ -42,6 +36,8 @@ def _resolve_tokenizer_location( local_path = Path(str(tokenizer_path)).expanduser() if local_path.exists(): return str(local_path) + from huggingface_hub import snapshot_download + return snapshot_download( repo_id=str(tokenizer_path), repo_type="model", @@ -134,9 +130,8 @@ class UniversalActionProcessor(ProcessorMixin): ), ( f"Decoded DCT coefficients have shape {decoded_dct_coeff.shape}, expected ({self.time_horizon}, {self.action_dim})" ) - except Exception as e: - print(f"Error decoding tokens: {e}") - print(f"Tokens: {token}") + except Exception: + logger.warning("Error decoding tokens: %s", token, exc_info=True) decoded_dct_coeff = np.zeros((self.time_horizon, self.action_dim)) decoded_actions.append(idct(decoded_dct_coeff / self.scale, axis=0, norm="ortho")) return np.stack(decoded_actions) diff --git a/src/lerobot/policies/molmoact2/hf_model/configuration_molmoact2.py b/src/lerobot/policies/molmoact2/molmoact2_hf_model/configuration_molmoact2.py similarity index 99% rename from src/lerobot/policies/molmoact2/hf_model/configuration_molmoact2.py rename to src/lerobot/policies/molmoact2/molmoact2_hf_model/configuration_molmoact2.py index 29da68c14..df5449bef 100644 --- a/src/lerobot/policies/molmoact2/hf_model/configuration_molmoact2.py +++ b/src/lerobot/policies/molmoact2/molmoact2_hf_model/configuration_molmoact2.py @@ -1,5 +1,3 @@ -#!/usr/bin/env python - # Copyright 2026 The Allen Institute for Artificial Intelligence and The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -14,13 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -# ruff: noqa """ MolmoAct2 configuration """ -from typing import Optional, Any +from typing import Any from transformers import PretrainedConfig from transformers.modeling_rope_utils import rope_config_validation diff --git a/src/lerobot/policies/molmoact2/hf_model/image_processing_molmoact2.py b/src/lerobot/policies/molmoact2/molmoact2_hf_model/image_processing_molmoact2.py similarity index 98% rename from src/lerobot/policies/molmoact2/hf_model/image_processing_molmoact2.py rename to src/lerobot/policies/molmoact2/molmoact2_hf_model/image_processing_molmoact2.py index a172c8477..acc709cb5 100644 --- a/src/lerobot/policies/molmoact2/hf_model/image_processing_molmoact2.py +++ b/src/lerobot/policies/molmoact2/molmoact2_hf_model/image_processing_molmoact2.py @@ -1,5 +1,3 @@ -#!/usr/bin/env python - # Copyright 2026 The Allen Institute for Artificial Intelligence and The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -14,33 +12,28 @@ # See the License for the specific language governing permissions and # limitations under the License. -# ruff: noqa """Image processor class for MolmoAct2""" -from typing import Optional, Union -import numpy as np import einops +import numpy as np import torch import torchvision.transforms - +from transformers.feature_extraction_utils import BatchFeature +from transformers.image_processing_utils import BaseImageProcessor, get_size_dict +from transformers.image_transforms import convert_to_rgb from transformers.image_utils import ( IMAGENET_STANDARD_MEAN, IMAGENET_STANDARD_STD, ImageInput, PILImageResampling, make_flat_list_of_images, - valid_images, to_numpy_array, + valid_images, ) -from transformers.image_transforms import convert_to_rgb from transformers.processing_utils import ImagesKwargs -from transformers.image_processing_utils import BaseImageProcessor, get_size_dict -from transformers.utils import logging -from transformers.feature_extraction_utils import BatchFeature from transformers.utils import TensorType, logging - logger = logging.get_logger(__name__) @@ -73,8 +66,8 @@ def resize_image( )(image) resized = torch.clip(resized, 0.0, 1.0).to(dtype) else: - assert image.dtype == torch.uint8, "SigLIP expects float images or uint8 images, but got {}".format( - image.dtype + assert image.dtype == torch.uint8, ( + f"SigLIP expects float images or uint8 images, but got {image.dtype}" ) in_min = 0.0 in_max = 255.0 @@ -96,7 +89,6 @@ def resize_image( def select_tiling(h, w, patch_size, max_num_crops): """Divide in image of size [w, h] in up to max_num_patches of size patch_size""" original_size = np.stack([h, w]) # [1, 2] - original_res = h * w tilings = [] for i in range(1, max_num_crops + 1): for j in range(1, max_num_crops + 1): @@ -406,13 +398,17 @@ class MolmoAct2ImageProcessor(BaseImageProcessor): image_std: float | list[float] | None = None, do_convert_rgb: bool = True, max_crops: int = 8, - overlap_margins: list[int] = [4, 4], + overlap_margins: list[int] | None = None, crop_mode: str = "overlap-and-resize-c2", patch_size: int = 14, - pooling_size: list[int] = [2, 2], + pooling_size: list[int] | None = None, **kwargs, ) -> None: super().__init__(**kwargs) + if overlap_margins is None: + overlap_margins = [4, 4] + if pooling_size is None: + pooling_size = [2, 2] size = size if size is not None else {"height": 378, "width": 378} size = get_size_dict(size, default_to_square=True) self.size = size diff --git a/src/lerobot/policies/molmoact2/hf_model/inference.py b/src/lerobot/policies/molmoact2/molmoact2_hf_model/inference.py similarity index 99% rename from src/lerobot/policies/molmoact2/hf_model/inference.py rename to src/lerobot/policies/molmoact2/molmoact2_hf_model/inference.py index 2c0243880..428800a8c 100644 --- a/src/lerobot/policies/molmoact2/hf_model/inference.py +++ b/src/lerobot/policies/molmoact2/molmoact2_hf_model/inference.py @@ -1,5 +1,3 @@ -#!/usr/bin/env python - # Copyright 2026 The Allen Institute for Artificial Intelligence and The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -14,16 +12,15 @@ # See the License for the specific language governing permissions and # limitations under the License. -# ruff: noqa """Inference utilities for MolmoAct2""" -from dataclasses import dataclass -from typing import Any, Optional, Tuple from collections.abc import Iterable, Sequence +from dataclasses import dataclass +from typing import Any import torch -from torch.nn import functional as F +from torch.nn import functional as F # noqa: N812 from transformers.cache_utils import Cache from transformers.configuration_utils import PretrainedConfig @@ -679,7 +676,7 @@ def _clone_static_inputs(inputs: _ActionFlowInputs) -> _ActionFlowInputs: def _copy_context_(dst: Any, src: Any) -> None: - for (dst_k, dst_v), (src_k, src_v) in zip(dst.kv_contexts, src.kv_contexts): + for (dst_k, dst_v), (src_k, src_v) in zip(dst.kv_contexts, src.kv_contexts, strict=False): dst_k.copy_(src_k) dst_v.copy_(src_v) if src.cross_mask is not None: @@ -689,7 +686,7 @@ def _copy_context_(dst: Any, src: Any) -> None: if src.valid_action is not None: dst.valid_action.copy_(src.valid_action) if src.rope_cache is not None: - for dst_tensor, src_tensor in zip(dst.rope_cache, src.rope_cache): + for dst_tensor, src_tensor in zip(dst.rope_cache, src.rope_cache, strict=False): dst_tensor.copy_(src_tensor) diff --git a/src/lerobot/policies/molmoact2/hf_model/modeling_molmoact2.py b/src/lerobot/policies/molmoact2/molmoact2_hf_model/modeling_molmoact2.py similarity index 99% rename from src/lerobot/policies/molmoact2/hf_model/modeling_molmoact2.py rename to src/lerobot/policies/molmoact2/molmoact2_hf_model/modeling_molmoact2.py index 4c36b04c8..e2edbe68d 100644 --- a/src/lerobot/policies/molmoact2/hf_model/modeling_molmoact2.py +++ b/src/lerobot/policies/molmoact2/molmoact2_hf_model/modeling_molmoact2.py @@ -1,5 +1,3 @@ -#!/usr/bin/env python - # Copyright 2026 The Allen Institute for Artificial Intelligence and The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -14,24 +12,25 @@ # See the License for the specific language governing permissions and # limitations under the License. -# ruff: noqa """Modeling code for MolmoAct2""" +# ruff: noqa: N806 + import json import math import os import re +from collections.abc import Callable, Mapping, Sequence from copy import deepcopy from dataclasses import dataclass -from typing import Any, Dict, List, Optional, Tuple, Union -from collections.abc import Callable, Mapping, Sequence +from typing import Any, Optional import numpy as np import torch import torch.utils.checkpoint from torch import nn -from torch.nn import functional as F +from torch.nn import functional as F # noqa: N812 from torch.nn.attention import SDPBackend, sdpa_kernel from transformers.activations import ACT2FN from transformers.cache_utils import Cache, DynamicCache @@ -647,7 +646,7 @@ class ActionExpert(nn.Module): f"got {len(encoder_kv_states)}." ) kv_contexts = [] - for block, (k_in, v_in) in zip(self.blocks, encoder_kv_states): + for block, (k_in, v_in) in zip(self.blocks, encoder_kv_states, strict=False): k_ctx = self._project_kv_tensor(k_in, self.context_k_proj) v_ctx = self._project_kv_tensor(v_in, self.context_v_proj) k_norm = block.cross_attn.k_norm @@ -732,7 +731,7 @@ class ActionExpert(nn.Module): timesteps: Sequence[torch.Tensor], ) -> Sequence[ActionExpertStepModulation]: cache = [] - for idx, step_t in enumerate(timesteps): + for _idx, step_t in enumerate(timesteps): conditioning = self._time_conditioning(step_t) block_modulations = [] for block in self.blocks: @@ -786,8 +785,8 @@ class ActionExpert(nn.Module): x = self.action_embed(actions) if context.valid_action is not None: x = x * context.valid_action - for idx, (block, kv_context, block_modulation) in enumerate( - zip(self.blocks, context.kv_contexts, block_modulations) + for _idx, (block, kv_context, block_modulation) in enumerate( + zip(self.blocks, context.kv_contexts, block_modulations, strict=False) ): x = block( x, @@ -2874,7 +2873,7 @@ class MolmoAct2Model(MolmoAct2PreTrainedModel): depth_mask=depth_mask, encoder_attention_mask=encoder_attention_mask, ) - for gate, source in zip(gate_head, sources) + for gate, source in zip(gate_head, sources, strict=False) ] return gates, depth_mask gate = self._depth_gate_from_source( @@ -4458,7 +4457,7 @@ class MolmoAct2ForConditionalGeneration(MolmoAct2PreTrainedModel, GenerationMixi ```python >>> from PIL import Image >>> import requests - >>> from lerobot.policies.molmoact2.hf_model.modeling_molmoact2 import MolmoAct2ForConditionalGeneration + >>> from lerobot.policies.molmoact2.molmoact2_hf_model.modeling_molmoact2 import MolmoAct2ForConditionalGeneration >>> from lerobot.policies.molmoact2.processor_molmoact2 import _load_local_molmoact2_processor >>> model = MolmoAct2ForConditionalGeneration.from_pretrained("...") diff --git a/src/lerobot/policies/molmoact2/hf_model/processing_molmoact2.py b/src/lerobot/policies/molmoact2/molmoact2_hf_model/processing_molmoact2.py similarity index 95% rename from src/lerobot/policies/molmoact2/hf_model/processing_molmoact2.py rename to src/lerobot/policies/molmoact2/molmoact2_hf_model/processing_molmoact2.py index 7b8775faa..6a73d2465 100644 --- a/src/lerobot/policies/molmoact2/hf_model/processing_molmoact2.py +++ b/src/lerobot/policies/molmoact2/molmoact2_hf_model/processing_molmoact2.py @@ -1,5 +1,3 @@ -#!/usr/bin/env python - # Copyright 2026 The Allen Institute for Artificial Intelligence and The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -14,45 +12,39 @@ # See the License for the specific language governing permissions and # limitations under the License. -# ruff: noqa """ Processor class for MolmoAct2. """ -from typing import Optional, Union -import dataclasses - import numpy as np - +from transformers import AutoTokenizer +from transformers.feature_extraction_utils import BatchFeature from transformers.image_utils import ImageInput -from transformers.video_utils import VideoInput from transformers.processing_utils import ( - Unpack, ProcessingKwargs, ProcessorMixin, + Unpack, ) -from transformers.feature_extraction_utils import BatchFeature -from transformers.tokenization_utils_base import TextInput, PreTokenizedInput +from transformers.tokenization_utils_base import PreTokenizedInput, TextInput from transformers.utils import logging +from transformers.video_utils import VideoInput -from transformers import AutoTokenizer -from .image_processing_molmoact2 import MolmoAct2ImagesKwargs, MolmoAct2ImageProcessor -from .video_processing_molmoact2 import MolmoAct2VideoProcessorKwargs, MolmoAct2VideoProcessor - +from .image_processing_molmoact2 import MolmoAct2ImageProcessor, MolmoAct2ImagesKwargs +from .video_processing_molmoact2 import MolmoAct2VideoProcessor, MolmoAct2VideoProcessorKwargs logger = logging.get_logger(__name__) # Special tokens, these should be present in any tokenizer we use since the preprocessor uses them -IMAGE_PATCH_TOKEN = f"" # Where to insert high-res tokens -IMAGE_LOW_RES_TOKEN = f"" # Where to insert low-res tokens -IM_START_TOKEN = f"" -LOW_RES_IMAGE_START_TOKEN = f"" -FRAME_START_TOKEN = f"" -IM_END_TOKEN = f"" -FRAME_END_TOKEN = f"" -IM_COL_TOKEN = f"" +IMAGE_PATCH_TOKEN = "" # nosec B105 # Where to insert high-res tokens +IMAGE_LOW_RES_TOKEN = "" # nosec B105 # Where to insert low-res tokens +IM_START_TOKEN = "" # nosec B105 +LOW_RES_IMAGE_START_TOKEN = "" # nosec B105 +FRAME_START_TOKEN = "" # nosec B105 +IM_END_TOKEN = "" # nosec B105 +FRAME_END_TOKEN = "" # nosec B105 +IM_COL_TOKEN = "" # nosec B105 IMAGE_PROMPT = "<|image|>" VIDEO_PROMPT = "<|video|>" @@ -224,7 +216,7 @@ class MolmoAct2Processor(ProcessorMixin): input_ids = input_ids[None, :] attention_mask = attention_mask[None, :] - B, S = input_ids.shape + B, S = input_ids.shape # noqa: N806 # Handle zero-length sequence if S == 0: @@ -364,7 +356,7 @@ class MolmoAct2Processor(ProcessorMixin): assert num_videos in {0, 1}, "At most one video is supported for now" video_grids_i = video_grids[index : index + num_videos] metadata_i = video_metadata[index : index + num_videos] - for video_grid, metadata in zip(video_grids_i, metadata_i): + for video_grid, metadata in zip(video_grids_i, metadata_i, strict=False): video_string = self.get_video_string( video_grid, metadata.timestamps, diff --git a/src/lerobot/policies/molmoact2/hf_model/video_processing_molmoact2.py b/src/lerobot/policies/molmoact2/molmoact2_hf_model/video_processing_molmoact2.py similarity index 98% rename from src/lerobot/policies/molmoact2/hf_model/video_processing_molmoact2.py rename to src/lerobot/policies/molmoact2/molmoact2_hf_model/video_processing_molmoact2.py index 644d5a691..bf4e44dde 100644 --- a/src/lerobot/policies/molmoact2/hf_model/video_processing_molmoact2.py +++ b/src/lerobot/policies/molmoact2/molmoact2_hf_model/video_processing_molmoact2.py @@ -1,5 +1,3 @@ -#!/usr/bin/env python - # Copyright 2026 The Allen Institute for Artificial Intelligence and The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -14,25 +12,23 @@ # See the License for the specific language governing permissions and # limitations under the License. -# ruff: noqa """Video processor class for MolmoAct2""" -from functools import partial import os import warnings +from collections.abc import Callable from contextlib import redirect_stdout +from functools import partial from io import BytesIO from urllib.parse import urlparse -from typing import Optional, Union -from collections.abc import Callable +import einops import numpy as np import requests -import einops import torch import torchvision.transforms - +from transformers.feature_extraction_utils import BatchFeature from transformers.image_utils import ( IMAGENET_STANDARD_MEAN, IMAGENET_STANDARD_STD, @@ -41,27 +37,24 @@ from transformers.image_utils import ( SizeDict, validate_kwargs, ) -from transformers.video_utils import ( - VideoInput, - is_valid_video, - make_batched_videos, - make_batched_metadata, - VideoMetadata, -) from transformers.processing_utils import Unpack, VideosKwargs -from transformers.video_processing_utils import BaseVideoProcessor -from transformers.utils import logging -from transformers.feature_extraction_utils import BatchFeature from transformers.utils import ( + TensorType, is_av_available, is_decord_available, is_torchcodec_available, is_yt_dlp_available, - TensorType, logging, to_numpy, ) - +from transformers.video_processing_utils import BaseVideoProcessor +from transformers.video_utils import ( + VideoInput, + VideoMetadata, + is_valid_video, + make_batched_metadata, + make_batched_videos, +) logger = logging.get_logger(__name__) @@ -102,8 +95,8 @@ def resize_image( )(image) resized = torch.clip(resized, 0.0, 1.0).to(dtype) else: - assert image.dtype == torch.uint8, "SigLIP expects float images or uint8 images, but got {}".format( - image.dtype + assert image.dtype == torch.uint8, ( + f"SigLIP expects float images or uint8 images, but got {image.dtype}" ) in_min = 0.0 in_max = 255.0 @@ -548,9 +541,8 @@ def get_target_fps( step_size = max(int(video_fps / target_fps), 1) num_frames_sampled_at_fps = int(total_frames / step_size) if num_frames_sampled == 0: - if "uniform" in frame_sample_mode: - if num_frames_sampled_at_fps > max_frames: - break + if "uniform" in frame_sample_mode and num_frames_sampled_at_fps > max_frames: + break selected_target_fps = target_fps num_frames_sampled = num_frames_sampled_at_fps @@ -779,13 +771,15 @@ class MolmoAct2VideoProcessor(BaseVideoProcessor): elif is_torchcodec_available(): warnings.warn( "`decord` is not installed and cannot be used to decode the video by default. " - "Falling back to `torchcodec`." + "Falling back to `torchcodec`.", + stacklevel=2, ) backend = "torchcodec" else: warnings.warn( "`decord` is not installed and cannot be used to decode the video by default. " - "Falling back to `PyAV`." + "Falling back to `PyAV`.", + stacklevel=2, ) backend = "pyav" @@ -795,7 +789,8 @@ class MolmoAct2VideoProcessor(BaseVideoProcessor): *[ self.fetch_videos(x, sample_timestamps_fn=sample_timestamps_fn) for x in video_url_or_urls - ] + ], + strict=False, ) ) else: @@ -821,7 +816,7 @@ class MolmoAct2VideoProcessor(BaseVideoProcessor): assert video_metadata[0].fps is not None, "FPS must be provided for video input" sampled_videos = [] sampled_metadata = [] - for video, metadata in zip(videos, video_metadata): + for video, metadata in zip(videos, video_metadata, strict=False): indices = sample_indices_fn(metadata=metadata) metadata.frames_indices = indices sampled_videos.append(video[indices]) @@ -985,11 +980,11 @@ class MolmoAct2VideoProcessor(BaseVideoProcessor): pixel_values_videos = np.concatenate(batch_crops, 0) video_token_pooling = np.concatenate(batch_pooled_patches_idx, 0) - data = dict( - pixel_values_videos=pixel_values_videos, - video_token_pooling=video_token_pooling, - video_grids=video_grids, - ) + data = { + "pixel_values_videos": pixel_values_videos, + "video_token_pooling": video_token_pooling, + "video_grids": video_grids, + } return BatchFeature(data, tensor_type=return_tensors) diff --git a/src/lerobot/policies/molmoact2/processor_molmoact2.py b/src/lerobot/policies/molmoact2/processor_molmoact2.py index 6c7a3ed5c..d2db817ef 100644 --- a/src/lerobot/policies/molmoact2/processor_molmoact2.py +++ b/src/lerobot/policies/molmoact2/processor_molmoact2.py @@ -1,5 +1,3 @@ -#!/usr/bin/env python - # Copyright 2026 The Allen Institute for Artificial Intelligence and The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -14,10 +12,18 @@ # See the License for the specific language governing permissions and # limitations under the License. +"""MolmoAct2 pre/post processing pipeline. + +Builds the multimodal prompt (images, discretised state, task text), +tokenises it via the vendored MolmoAct2 processor, and handles quantile +normalisation with optional per-dimension gripper masking. +""" + from __future__ import annotations import json -import os +import logging +import math import re from contextlib import suppress from copy import deepcopy @@ -27,7 +33,6 @@ from typing import TYPE_CHECKING, Any import numpy as np import torch -from huggingface_hub import snapshot_download from torch import Tensor from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature @@ -54,14 +59,71 @@ from lerobot.utils.constants import ( ) from lerobot.utils.import_utils import _scipy_available, _transformers_available, require_package -from .configuration_molmoact2 import MolmoAct2Config, infer_molmoact2_max_sequence_length +from .configuration_molmoact2 import MolmoAct2Config +from .modeling_molmoact2 import _hf_token, _resolve_checkpoint_location + +logger = logging.getLogger(__name__) + +MOLMOACT2_DEFAULT_NUM_IMAGES = 2 +MOLMOACT2_IMAGE_TOKENS_PER_IMAGE = 196 +MOLMOACT2_FIXED_PROMPT_TOKEN_BUDGET = 80 +MOLMOACT2_TASK_TOKEN_BUDGET = 32 +MOLMOACT2_SEQUENCE_LENGTH_MARGIN = 32 +MOLMOACT2_SEQUENCE_LENGTH_MULTIPLE = 64 +MOLMOACT2_DISCRETE_ACTION_WRAPPER_TOKENS = 4 +MOLMOACT2_MIN_DISCRETE_ACTION_TOKENS_PER_STEP = 6 +MOLMOACT2_DISCRETE_ACTION_TOKENS_PER_DIM = 0.95 + + +def _round_up(value: int, multiple: int) -> int: + return int(math.ceil(value / multiple) * multiple) + + +def infer_molmoact2_max_sequence_length( + *, + num_images: int, + state_dim: int, + action_dim: int, + action_horizon: int, + include_discrete_action: bool, +) -> int: + """Infer the padded text/image sequence cap from MolmoAct2's fixed token layout.""" + if num_images < 1: + num_images = MOLMOACT2_DEFAULT_NUM_IMAGES + if state_dim < 0: + state_dim = 0 + if action_dim < 1: + action_dim = 1 + if action_horizon < 1: + action_horizon = 1 + + image_tokens = num_images * MOLMOACT2_IMAGE_TOKENS_PER_IMAGE + prompt_tokens = ( + MOLMOACT2_FIXED_PROMPT_TOKEN_BUDGET + + MOLMOACT2_TASK_TOKEN_BUDGET + + state_dim + + MOLMOACT2_SEQUENCE_LENGTH_MARGIN + ) + action_tokens = 0 + if include_discrete_action: + action_tokens_per_step = max( + MOLMOACT2_MIN_DISCRETE_ACTION_TOKENS_PER_STEP, + math.ceil(action_dim * MOLMOACT2_DISCRETE_ACTION_TOKENS_PER_DIM), + ) + action_tokens = MOLMOACT2_DISCRETE_ACTION_WRAPPER_TOKENS + action_horizon * action_tokens_per_step + + return _round_up( + image_tokens + prompt_tokens + action_tokens, + MOLMOACT2_SEQUENCE_LENGTH_MULTIPLE, + ) + if TYPE_CHECKING or _transformers_available: from transformers import Qwen2Tokenizer - from .hf_model.image_processing_molmoact2 import MolmoAct2ImageProcessor - from .hf_model.processing_molmoact2 import MolmoAct2Processor - from .hf_model.video_processing_molmoact2 import MolmoAct2VideoProcessor + from .molmoact2_hf_model.image_processing_molmoact2 import MolmoAct2ImageProcessor + from .molmoact2_hf_model.processing_molmoact2 import MolmoAct2Processor + from .molmoact2_hf_model.video_processing_molmoact2 import MolmoAct2VideoProcessor else: Qwen2Tokenizer = None MolmoAct2ImageProcessor = None @@ -69,7 +131,7 @@ else: MolmoAct2VideoProcessor = None if TYPE_CHECKING or (_transformers_available and _scipy_available): - from .hf_model.action_tokenizer import UniversalActionProcessor + from .molmoact2_hf_model.action_tokenizer import UniversalActionProcessor else: UniversalActionProcessor = None @@ -97,32 +159,6 @@ _QUESTION_PREFIX_PATTERNS = tuple( ) -def _hf_token() -> str | None: - return os.environ.get("HF_TOKEN") or os.environ.get("HF_ACCESS_TOKEN") - - -def _resolve_checkpoint_location( - checkpoint_path: str, - *, - revision: str | None = None, - force_download: bool = False, -) -> str: - checkpoint_path = str(checkpoint_path or "").strip() - if not checkpoint_path: - raise ValueError("MolmoAct2 policy requires `checkpoint_path`.") - local_path = Path(checkpoint_path).expanduser() - if local_path.exists(): - return str(local_path) - return snapshot_download( - repo_id=checkpoint_path, - repo_type="model", - revision=revision, - force_download=force_download, - ignore_patterns=["*.py", "*.pyc", "__pycache__/*"], - token=_hf_token(), - ) - - def _load_hf_norm_stats_for_tag( checkpoint_path: str, *, @@ -969,6 +1005,93 @@ class MolmoAct2PackInputsProcessorStep(ProcessorStep): return features +@ProcessorStepRegistry.register(name="molmoact2_state_frame_transform") +@dataclass +class MolmoAct2StateFrameTransformStep(ProcessorStep): + """Convert robot state from arm frame to model frame before normalization. + + Required for zero-shot deployment of MolmoAct2-SO100_101 on SO-100/101 + arms calibrated with LeRobot >= 0.5.0 (v3.0 convention). The checkpoint + was trained on data using a different joint convention (sign flip on + shoulder_lift, 90 deg offset on shoulder_lift and elbow_flex). + + No-op when joint_signs and joint_offsets are None (default), so this + step has no effect on fine-tuned models or other embodiments. + + state_model = signs * arm_state + offsets + + See: https://huggingface.co/docs/lerobot/backwardcomp + """ + + joint_signs: list[float] | None = None + joint_offsets: list[float] | None = None + + def __call__(self, transition: EnvTransition) -> EnvTransition: + if self.joint_signs is None or self.joint_offsets is None: + return transition + observation = transition.get(TransitionKey.OBSERVATION) + if not isinstance(observation, dict) or OBS_STATE not in observation: + return transition + transition = transition.copy() + observation = observation.copy() + state = torch.as_tensor(observation[OBS_STATE], dtype=torch.float32).clone() + n = len(self.joint_signs) + signs = torch.tensor(self.joint_signs, dtype=torch.float32, device=state.device) + offsets = torch.tensor(self.joint_offsets, dtype=torch.float32, device=state.device) + state[..., :n] = signs * state[..., :n] + offsets + observation[OBS_STATE] = state + transition[TransitionKey.OBSERVATION] = observation + return transition + + def transform_features( + self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]] + ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]: + return features + + def get_config(self) -> dict[str, Any]: + return {"joint_signs": self.joint_signs, "joint_offsets": self.joint_offsets} + + +@ProcessorStepRegistry.register(name="molmoact2_action_frame_transform") +@dataclass +class MolmoAct2ActionFrameTransformStep(ProcessorStep): + """Convert model action from model frame back to arm frame after unnormalization. + + Inverse of MolmoAct2StateFrameTransformStep. Required for zero-shot + MolmoAct2-SO100_101 on SO-100/101 arms. No-op when both fields are None. + + action_arm = signs * (model_action - offsets) + + See: https://huggingface.co/docs/lerobot/backwardcomp + """ + + joint_signs: list[float] | None = None + joint_offsets: list[float] | None = None + + def __call__(self, transition: EnvTransition) -> EnvTransition: + if self.joint_signs is None or self.joint_offsets is None: + return transition + action = transition.get(TransitionKey.ACTION) + if action is None: + return transition + transition = transition.copy() + action = torch.as_tensor(action, dtype=torch.float32).clone() + n = len(self.joint_signs) + signs = torch.tensor(self.joint_signs, dtype=torch.float32, device=action.device) + offsets = torch.tensor(self.joint_offsets, dtype=torch.float32, device=action.device) + action[..., :n] = signs * (action[..., :n] - offsets) + transition[TransitionKey.ACTION] = action + return transition + + def transform_features( + self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]] + ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]: + return features + + def get_config(self) -> dict[str, Any]: + return {"joint_signs": self.joint_signs, "joint_offsets": self.joint_offsets} + + @ProcessorStepRegistry.register(name="molmoact2_clamp_action") @dataclass class MolmoAct2ClampActionProcessorStep(ProcessorStep): @@ -1031,6 +1154,10 @@ def make_molmoact2_pre_post_processors( input_steps: list[ProcessorStep] = [ RenameObservationsProcessorStep(rename_map={}), AddBatchDimensionProcessorStep(), + MolmoAct2StateFrameTransformStep( + joint_signs=config.joint_signs, + joint_offsets=config.joint_offsets, + ), MolmoAct2MaskedNormalizerProcessorStep( features={**config.input_features, **config.output_features}, norm_map=config.normalization_mapping, @@ -1066,6 +1193,10 @@ def make_molmoact2_pre_post_processors( norm_map=config.normalization_mapping, stats=masked_dataset_stats, ), + MolmoAct2ActionFrameTransformStep( + joint_signs=config.joint_signs, + joint_offsets=config.joint_offsets, + ), DeviceProcessorStep(device="cpu"), ] diff --git a/src/lerobot/policies/pretrained.py b/src/lerobot/policies/pretrained.py index a69487f3f..aea5f1b08 100644 --- a/src/lerobot/policies/pretrained.py +++ b/src/lerobot/policies/pretrained.py @@ -23,7 +23,7 @@ from typing import TypedDict, TypeVar, Unpack import packaging import safetensors -from huggingface_hub import HfApi, ModelCard, ModelCardData, hf_hub_download +from huggingface_hub import HfApi, ModelCard, ModelCardData, hf_hub_download, save_torch_state_dict from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE from huggingface_hub.errors import HfHubHTTPError from safetensors.torch import load_model as load_model_as_safetensor, save_model as save_model_as_safetensor @@ -129,10 +129,43 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC): if not getattr(cls, "name", None): raise TypeError(f"Class {cls.__name__} must define 'name'") - def _save_pretrained(self, save_directory: Path) -> None: + def save_pretrained( + self, + save_directory: str | Path, + *, + state_dict: dict[str, Tensor] | None = None, + repo_id: str | None = None, + push_to_hub: bool = False, + card_kwargs: dict | None = None, + **push_to_hub_kwargs, + ) -> str | None: + """Save the policy to a directory (and optionally push to the Hub). + + Overrides `HubMixin.save_pretrained` to add a `state_dict` argument (mirroring + `transformers.PreTrainedModel.save_pretrained`). Under FSDP, `self.state_dict()` would + return sharded tensors, so the caller gathers the full state dict via a cross-rank + collective and passes it here for `_save_pretrained` to write directly. + """ + save_directory = Path(save_directory) + save_directory.mkdir(parents=True, exist_ok=True) + self._save_pretrained(save_directory, state_dict=state_dict) + if push_to_hub: + if repo_id is None: + repo_id = save_directory.name + return self.push_to_hub(repo_id=repo_id, card_kwargs=card_kwargs, **push_to_hub_kwargs) + return None + + def _save_pretrained(self, save_directory: Path, state_dict: dict[str, Tensor] | None = None) -> None: self.config._save_pretrained(save_directory) model_to_save = self.module if hasattr(self, "module") else self - save_model_as_safetensor(model_to_save, str(save_directory / SAFETENSORS_SINGLE_FILE)) + if state_dict is None: + save_model_as_safetensor(model_to_save, str(save_directory / SAFETENSORS_SINGLE_FILE)) + return + # A pre-gathered (e.g. FSDP full) state dict was supplied: write it directly. + # `save_torch_state_dict` discards shared-tensor duplicates just like `save_model` does; + # pin `max_shard_size` above the total size so the output stays a single `model.safetensors` + total_bytes = sum(t.numel() * t.element_size() for t in state_dict.values()) + save_torch_state_dict(state_dict, str(save_directory), max_shard_size=max(total_bytes, 1)) @classmethod def from_pretrained( @@ -270,6 +303,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC): self, cfg: TrainPipelineConfig, peft_model=None, + state_dict: dict[str, Tensor] | None = None, ): api = HfApi() repo_id = api.create_repo( @@ -287,7 +321,8 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC): peft_model.save_pretrained(saved_path) self.config.save_pretrained(saved_path) else: - self.save_pretrained(saved_path) # Calls _save_pretrained and stores model tensors + # Calls _save_pretrained and stores model tensors + self.save_pretrained(saved_path, state_dict=state_dict) card = self.generate_model_card( cfg.dataset.repo_id, self.config.type, self.config.license, self.config.tags, cfg=cfg @@ -305,6 +340,9 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC): ignore_patterns=["*.tmp", "*.log"], ) + # Contract: lerobot.jobs.hf.submit_to_hf watches for this exact + # "Model pushed to " line to end a remote run early. Keep the wording + # and URL format in sync (it falls back to status polling if they drift). logging.info(f"Model pushed to {commit_info.repo_url.url}") def generate_model_card( diff --git a/src/lerobot/policies/utils.py b/src/lerobot/policies/utils.py index c37127813..f465fcff8 100644 --- a/src/lerobot/policies/utils.py +++ b/src/lerobot/policies/utils.py @@ -126,7 +126,8 @@ def prepare_observation_for_inference( for name in observation: observation[name] = torch.from_numpy(observation[name]) if "image" in name: - observation[name] = observation[name].type(torch.float32) / 255 + if observation[name].dtype == torch.uint8: + observation[name] = observation[name].type(torch.float32) / 255 observation[name] = observation[name].permute(2, 0, 1).contiguous() observation[name] = observation[name].unsqueeze(0) observation[name] = observation[name].to(device) diff --git a/src/lerobot/policies/vla_jepa/modeling_vla_jepa.py b/src/lerobot/policies/vla_jepa/modeling_vla_jepa.py index 45d83e652..9c689d3c7 100644 --- a/src/lerobot/policies/vla_jepa/modeling_vla_jepa.py +++ b/src/lerobot/policies/vla_jepa/modeling_vla_jepa.py @@ -17,12 +17,10 @@ from __future__ import annotations import logging from collections import deque from pathlib import Path -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any -import numpy as np import torch import torch.nn.functional as F # noqa: N812 -from PIL import Image from torch import Tensor, nn from lerobot.policies.pretrained import PreTrainedPolicy, T @@ -55,12 +53,13 @@ class VLAJEPAModel(nn.Module): - DiT-B: flow-matching action head for future action prediction - V-JEPA: world model for video frame prediction - Input: List[dict] native format (same as original starVLA) - - "image": List[PIL.Image] (multi-view images) - - "video": np.ndarray [V, T, H, W, 3] - - "lang": str (task instruction) - - "action": np.ndarray [T, action_dim] (optional, training only) - - "state": np.ndarray [1, state_dim] (optional) + Inputs are batched tensors kept on the model device + - images: List[List[Tensor [C, H, W]]] (float [0,1]) — per sample, per view (Qwen messages) + - instructions: List[str] + - videos: Tensor [B, V, T, C, H, W] (float [0,1], world model only) + - actions: Tensor [B, T, action_dim] (optional, training only) + - state: Tensor [B, 1, state_dim] (optional) + - action_is_pad: Tensor [B, T] (optional) """ def __init__(self, config: VLAJEPAConfig) -> None: @@ -75,6 +74,11 @@ class VLAJEPAModel(nn.Module): self.action_tokens, self.action_token_ids, self.embodied_action_token_id = ( self.qwen.expand_tokenizer() ) + self.register_buffer( + "_action_token_ids_t", + torch.tensor(self.action_token_ids, dtype=torch.long), + persistent=False, + ) # Action head (flow-matching DiT) self.action_model = VLAJEPAActionHead(config, cross_attention_dim=self.qwen.model.config.hidden_size) @@ -161,166 +165,123 @@ class VLAJEPAModel(nn.Module): # ---- Native VLA-JEPA forward (follows original VLA_JEPA.py) ---- - def forward(self, examples: list[dict]) -> dict[str, Tensor]: - """ - Native forward pass following original starVLA VLA_JEPA.forward. - - Args: - examples: List of per-sample dicts with keys: - "image" : List[PIL.Image] — multi-view images - "video" : np.ndarray [V, T, H, W, 3] - "lang" : str — task instruction - "action" : np.ndarray [T, action_dim] (optional) - "state" : np.ndarray [1, state_dim] (optional) - - Returns: - dict with "action_loss" and "wm_loss" keys (scalar Tensors). - """ - # Unpack native format (same pattern as original VLA_JEPA.py) - batch_images = [ex["image"] for ex in examples] # List[List[PIL.Image]] - batch_videos = [ex["video"] for ex in examples] # List[np.ndarray] - instructions = [ex["lang"] for ex in examples] # List[str] - has_action = "action" in examples[0] and examples[0]["action"] is not None - actions = [ex["action"] for ex in examples] if has_action else None - has_state = "state" in examples[0] and examples[0]["state"] is not None - state = [ex["state"] for ex in examples] if has_state else None - action_is_pad = ( - [ex["action_is_pad"] for ex in examples] - if has_action and "action_is_pad" in examples[0] and examples[0]["action_is_pad"] is not None - else None - ) - - # Stack videos: [B, V, T, H, W, 3] -> [B, V, T, 3, H, W] - batch_videos = np.stack(batch_videos) - batch_videos = batch_videos.transpose(0, 1, 2, 5, 3, 4) # [B, V, T, 3, H, W] - - # Adjust number of views for the world model: - # - fewer views than expected: duplicate the first view to fill up - # - more views than expected: keep only the first num_views_world_model views - num_views_world_model = self.config.jepa_tubelet_size - if batch_videos.shape[1] < num_views_world_model: - num_missing_views = num_views_world_model - batch_videos.shape[1] - first_view = np.repeat(batch_videos[:, :1], num_missing_views, axis=1) - batch_videos = np.concatenate([batch_videos, first_view], axis=1) - elif batch_videos.shape[1] > num_views_world_model: - batch_videos = batch_videos[:, :num_views_world_model] - - # ---- Step 1: QwenVL encode (same as original) ---- + def _encode_qwen( + self, images: list[list[Tensor]], instructions: list[str], *, need_action_tokens: bool + ) -> tuple[Tensor, Tensor, Tensor | None]: + """Run Qwen and gather the embodied-action (and optionally action) token hidden states.""" qwen_inputs = self.qwen.build_inputs( - images=batch_images, + images=images, instructions=instructions, action_prompt=self.replace_prompt, embodied_prompt=self.embodied_replace_prompt, ) - - # Locate embodied-action tokens (always needed for action head) - embodied_mask = qwen_inputs["input_ids"] == self.embodied_action_token_id - embodied_indices = embodied_mask.nonzero(as_tuple=True) - - # Locate action tokens (only needed for world model predictor) - if self.config.enable_world_model: - action_mask = torch.isin( - qwen_inputs["input_ids"], - torch.tensor(self.action_token_ids, device=qwen_inputs["input_ids"].device), - ) - action_indices = action_mask.nonzero(as_tuple=True) + input_ids = qwen_inputs["input_ids"] + embodied_idx = (input_ids == self.embodied_action_token_id).nonzero(as_tuple=True) + action_idx = None + if need_action_tokens: + action_mask = torch.isin(input_ids, self._action_token_ids_t) + action_idx = action_mask.nonzero(as_tuple=True) device_type = next(self.parameters()).device.type - with torch.autocast(device_type=device_type, dtype=torch.bfloat16): last_hidden = self._qwen_last_decoder_hidden(qwen_inputs) # [B, seq_len, H] b, _, h = last_hidden.shape + embodied_action_tokens = last_hidden[embodied_idx[0], embodied_idx[1], :].view(b, -1, h) + action_tokens = ( + last_hidden[action_idx[0], action_idx[1], :].view(b, -1, h) + if action_idx is not None + else None + ) + return embodied_action_tokens, action_tokens - if self.config.enable_world_model: - action_tokens = last_hidden[action_indices[0], action_indices[1], :].view(b, -1, h) + def _world_model_loss(self, videos: Tensor, action_tokens: Tensor) -> Tensor: + """JEPA encode + predictor L1 loss. `videos` is [B, V, T, C, H, W] float in [0, 1].""" + # Match the world model's expected view count: pad with the first view, or trim extras. + num_views = self.config.jepa_tubelet_size + if videos.shape[1] < num_views: + missing = num_views - videos.shape[1] + videos = torch.cat([videos, videos[:, :1].repeat(1, missing, 1, 1, 1, 1)], dim=1) + elif videos.shape[1] > num_views: + videos = videos[:, :num_views] - embodied_action_tokens = last_hidden[embodied_indices[0], embodied_indices[1], :].view(b, -1, h) + b, v, t_frames, c, h_img, w_img = videos.shape + flat = videos.reshape(b * v, t_frames, c, h_img, w_img) + # Fast (torchvision) video processor on-device, do_rescale=False (frames already in [0, 1]). + video_pixels = self.video_processor( + videos=list(flat), + return_tensors="pt", + device=self.video_encoder.device, + do_rescale=False, + )["pixel_values_videos"] # [B*V, T, C, H, W] - # ---- Step 2+3: JEPA Encoder + Predictor ---- - device_wm = last_hidden.device - if not self.config.enable_world_model: - wm_loss = torch.tensor(0.0, device=device_wm) + with torch.no_grad(): + video_embeddings = self.video_encoder.get_vision_features(pixel_values_videos=video_pixels) + # Merge views: [B*V, ...] -> [B, ..., V*embed_dim] + video_embeddings = torch.cat(torch.chunk(video_embeddings, chunks=v, dim=0), dim=2) + + tubelet_size = self.video_encoder.config.tubelet_size + # num_video_frames raw frames → t_enc_total temporal positions after tubelet compression + t_enc_total = self.config.num_video_frames // tubelet_size + if t_enc_total < 2: + return torch.zeros((), device=video_embeddings.device) + + # Shift-by-one JEPA split: input_states = positions 0..T-2, gt_states = positions 1..T-1 + t_enc_ctx = t_enc_total - 1 + tokens_per_frame = video_embeddings.shape[1] // t_enc_total + input_states = video_embeddings[:, : tokens_per_frame * t_enc_ctx, :] + gt_states = video_embeddings[:, tokens_per_frame:, :] + + expected_actions = t_enc_ctx * self.config.num_action_tokens_per_timestep + if action_tokens.shape[1] < expected_actions: + pad = action_tokens[:, -1:].repeat(1, expected_actions - action_tokens.shape[1], 1) + action_tokens = torch.cat([action_tokens, pad], dim=1) + + predicted_states = self.video_predictor( + input_states.float(), action_tokens[:, :expected_actions].float() + ) + return F.l1_loss(predicted_states, gt_states.float(), reduction="mean") + + def _action_loss( + self, + embodied_action_tokens: Tensor, + actions: Tensor, + state: Tensor | None, + action_is_pad: Tensor | None, + ) -> Tensor: + """Flow-matching action-head loss, repeated over `repeated_diffusion_steps`.""" + device_type = next(self.parameters()).device.type + with torch.autocast(device_type=device_type, dtype=torch.float32): + r = self.config.repeated_diffusion_steps + horizon = self.config.chunk_size + actions_target = actions[:, -horizon:, :].to(torch.float32).repeat(r, 1, 1) + embodied = embodied_action_tokens.repeat(r, 1, 1) + state_rep = state.to(embodied_action_tokens.dtype).repeat(r, 1, 1) if state is not None else None + pad_rep = action_is_pad[:, -horizon:].repeat(r, 1) if action_is_pad is not None else None + return self.action_model(embodied, actions_target, state_rep, pad_rep) + + def forward( + self, + images: list[list[Tensor]], + instructions: list[str], + videos: Tensor | None = None, + actions: Tensor | None = None, + state: Tensor | None = None, + action_is_pad: Tensor | None = None, + ) -> dict[str, Tensor]: + """Native forward: Qwen encode → optional world-model loss → optional action-head loss.""" + embodied_action_tokens, action_tokens = self._encode_qwen( + images, instructions, need_action_tokens=self.config.enable_world_model + ) + + if self.config.enable_world_model and videos is not None: + wm_loss = self._world_model_loss(videos, action_tokens) else: - b, v, t_frames, c, h_img, w_img = batch_videos.shape - batch_videos_flat = batch_videos.reshape(b * v, t_frames, c, h_img, w_img) + wm_loss = torch.zeros((), device=embodied_action_tokens.device) - video_pixels = self.video_processor(videos=list(batch_videos_flat), return_tensors="pt")[ - "pixel_values_videos" - ].to(self.video_encoder.device) # [B*V, T, C, H, W] - - with torch.no_grad(): - video_embeddings = self.video_encoder.get_vision_features(pixel_values_videos=video_pixels) - # Merge views: [B*V, ...] -> [B, ..., V*embed_dim] - video_embeddings = torch.cat(torch.chunk(video_embeddings, chunks=v, dim=0), dim=2) - - tubelet_size = self.video_encoder.config.tubelet_size - device_wm = video_embeddings.device - # num_video_frames raw frames → t_enc_total temporal positions after tubelet compression - t_enc_total = self.config.num_video_frames // tubelet_size - - if t_enc_total < 2: - wm_loss = torch.tensor(0.0, device=device_wm) - else: - # Shift-by-one JEPA split (matches original VLA_JEPA.py lines 231-232): - # input_states: positions 0..T-2, gt_states: positions 1..T-1 - t_enc_ctx = t_enc_total - 1 - tokens_per_frame = video_embeddings.shape[1] // t_enc_total - - input_states = video_embeddings[:, : tokens_per_frame * t_enc_ctx, :] - gt_states = video_embeddings[:, tokens_per_frame:, :] - - expected_actions = t_enc_ctx * self.config.num_action_tokens_per_timestep - if action_tokens.shape[1] < expected_actions: - pad = action_tokens[:, -1:].repeat(1, expected_actions - action_tokens.shape[1], 1) - action_tokens = torch.cat([action_tokens, pad], dim=1) - - predicted_states = self.video_predictor( - input_states.float(), - action_tokens[:, :expected_actions].float(), - ) - - wm_loss = F.l1_loss(predicted_states, gt_states.float(), reduction="mean") - - if not has_action: + if actions is None: return {"wm_loss": wm_loss} - # ---- Step 4: Action Head ---- - with torch.autocast(device_type=device_type, dtype=torch.float32): - actions_tensor = torch.tensor( - np.array(actions), device=last_hidden.device, dtype=torch.float32 - ) # [B, T_full, action_dim] - action_horizon = self.config.chunk_size - actions_target = actions_tensor[:, -action_horizon:, :] - - state_tensor = None - if state is not None: - state_tensor = torch.tensor( - np.array(state), device=last_hidden.device, dtype=last_hidden.dtype - ) # [B, 1, state_dim] - - repeated_diffusion_steps = self.config.repeated_diffusion_steps - actions_target = actions_target.repeat(repeated_diffusion_steps, 1, 1) - embodied_action_tokens = embodied_action_tokens.repeat(repeated_diffusion_steps, 1, 1) - if state_tensor is not None: - state_tensor = state_tensor.repeat(repeated_diffusion_steps, 1, 1) - - action_is_pad_rep = None - if action_is_pad is not None: - pad_tensor = torch.stack( - [ - p.to(actions_target.device) - if isinstance(p, Tensor) - else torch.tensor(p, device=actions_target.device) - for p in action_is_pad - ] - ) # [B, T_full] - pad_tensor = pad_tensor[:, -action_horizon:] # [B, action_horizon] - action_is_pad_rep = pad_tensor.repeat(repeated_diffusion_steps, 1) # [B*R, action_horizon] - - action_loss = self.action_model( - embodied_action_tokens, actions_target, state_tensor, action_is_pad_rep - ) - + action_loss = self._action_loss(embodied_action_tokens, actions, state, action_is_pad) return {"action_loss": action_loss, "wm_loss": wm_loss * self.config.world_model_loss_weight} # ---- Native predict_action (follows original VLA_JEPA.predict_action) ---- @@ -328,58 +289,23 @@ class VLAJEPAModel(nn.Module): @torch.no_grad() def predict_action( self, - batch_images: list[list[Image.Image]], + images: list[list[Tensor]], instructions: list[str], - state: np.ndarray | None = None, - ) -> np.ndarray: - """ - Native action prediction following original VLA_JEPA.predict_action. - - Args: - batch_images: List of samples; each is List[PIL.Image] (multi-view). - instructions: Task instructions, one per sample. - state: Optional [B, state_dim] numpy array. - - Returns: - np.ndarray [B, action_horizon, action_dim] — predicted actions. - """ + state: Tensor | None = None, + ) -> Tensor: + """Predict an action chunk. `images` is per-sample, per-view float [0,1] [C, H, W] tensors.""" if self.config.resize_images_to is not None: height, width = self.config.resize_images_to - resampling = getattr(Image, "Resampling", Image).BOX - batch_images = [ - [image.resize((width, height), resample=resampling) for image in sample_images] - for sample_images in batch_images + images = [ + [F.interpolate(img[None], size=(height, width), mode="area")[0] for img in views] + for views in images ] - qwen_inputs = self.qwen.build_inputs( - images=batch_images, - instructions=instructions, - action_prompt=self.replace_prompt, - embodied_prompt=self.embodied_replace_prompt, + embodied_action_tokens, _ = self._encode_qwen(images, instructions, need_action_tokens=False) + return self.action_model.predict_action( + embodied_action_tokens.float(), state.float() if state is not None else None ) - embodied_mask = qwen_inputs["input_ids"] == self.embodied_action_token_id - embodied_indices = embodied_mask.nonzero(as_tuple=True) - - device_type = next(self.parameters()).device.type - - with torch.autocast(device_type=device_type, dtype=torch.bfloat16): - last_hidden = self._qwen_last_decoder_hidden(qwen_inputs) # [B, seq_len, H] - b, _, h = last_hidden.shape - embodied_action_tokens = last_hidden[embodied_indices[0], embodied_indices[1], :].view(b, -1, h) - - state_tensor = None - if state is not None: - state_tensor = torch.from_numpy(np.array(state)).to( - device=last_hidden.device, dtype=last_hidden.dtype - ) - - pred_actions = self.action_model.predict_action( - embodied_action_tokens.float(), state_tensor.float() if state_tensor is not None else None - ) # [B, action_horizon, action_dim] - - return pred_actions.detach().cpu().numpy() - # ============================================================================ # LeRobot Adapter Layer - converts between LeRobot batch format and native VLA-JEPA format @@ -390,9 +316,9 @@ class VLAJEPAPolicy(PreTrainedPolicy): """ LeRobot adapter for VLA-JEPA. - Converts LeRobot's standard batch format (dict[str, Tensor]) to the native - VLA-JEPA format (List[dict]), calls the native model, and converts outputs - back to LeRobot format. + Converts LeRobot's standard batch format (dict[str, Tensor]) to the batched tensors + the native model expects (keeping everything on-device), calls the native model, and + converts outputs back to LeRobot format. """ config_class = VLAJEPAConfig @@ -419,9 +345,8 @@ class VLAJEPAPolicy(PreTrainedPolicy): # ---- Format Conversion: LeRobot → Native ---- - def _prepare_model_inputs(self, batch: dict[str, Tensor]) -> list[dict]: - """ - Convert LeRobot batch format to native VLA-JEPA examples format. + def _prepare_model_inputs(self, batch: dict[str, Tensor], training=True) -> dict[str, Any]: + """Convert a LeRobot batch to the model's batched, on-device inputs. LeRobot format: batch = { @@ -431,65 +356,25 @@ class VLAJEPAPolicy(PreTrainedPolicy): "task": str | List[str], (optional instruction) } - Native format (List[dict]): - { - "image": List[PIL.Image], # multi-view images per sample - "video": np.ndarray [V, T, H, W, 3], - "lang": str, # task instruction - "action": np.ndarray [T, action_dim], # optional - "state": np.ndarray [1, state_dim], # optional - } + Returns the kwargs for `VLAJEPAModel.forward` / `.predict_action` (everything stays + on the batch device; no per-sample shredding): `images` (per-sample, per-view list for + Qwen messages), `instructions`, and the batched `videos` / `actions` / `state` / + `action_is_pad` when present. """ - # Determine batch size from the first image feature image_keys = list(self.config.image_features.keys()) if not image_keys: raise ValueError("VLAJEPA requires at least one image feature.") - first_key = image_keys[0] - first_tensor = batch[first_key] - batch_size = first_tensor.shape[0] + batch_size = batch[image_keys[0]].shape[0] - # ---- Collect images per sample ---- - # images_per_sample[b][v] = PIL.Image for view v - images_per_sample: list[list[Image.Image]] = [[] for _ in range(batch_size)] + # Current-frame image per view ([B, C, H, W]); regroup per sample for Qwen messages. + frames = [] for key in image_keys: - tensor = batch[key] # [B, C, H, W] or [B, T, C, H, W] - if tensor.ndim == 5: - # observation_delta_indices = [0, 1, ..., num_video_frames-1] - # index 0 is the current observation (delta=0) - tensor = tensor[:, 0] - for b in range(batch_size): - images_per_sample[b].append(self.model.qwen.tensor_to_pil(tensor[b])) + t = batch[key] + if t.ndim == 5: # [B, T, C, H, W] -> current observation (delta=0) + t = t[:, 0] + frames.append(self.model.qwen.to_pixel_values(t)) + images = [[frame[b] for frame in frames] for b in range(batch_size)] - # ---- Collect videos per sample ---- - # Build video arrays: for each sample, stack views as [V, T, H, W, 3] - # Check whether any image feature has a time dimension - video_source = None - for k in image_keys: - if k in batch: - video_source = batch[k] # Use first available for shape inspection - break - - if video_source is None: - raise ValueError("No image data found in batch for video construction.") - - videos_per_sample = [] - for b in range(batch_size): - sample_views = [] - for k in image_keys: - t = batch[k][b] # [C, H, W] or [T, C, H, W] - if t.ndim == 3: - t = t.unsqueeze(0) # [1, C, H, W] - # Convert to [T, H, W, 3] numpy - t_np = t.permute(0, 2, 3, 1).detach().cpu().float().numpy() - # Clamp to [0, 255] - if t_np.max() <= 1.0: - t_np = t_np * 255.0 - t_np = np.rint(t_np.clip(0, 255)).astype(np.uint8) - sample_views.append(t_np) - # Stack views: [V, T, H, W, 3] - videos_per_sample.append(np.stack(sample_views, axis=0)) - - # ---- Collect instructions ---- tasks = batch.get("task") if tasks is None: instructions = ["Execute the robot action."] * batch_size @@ -498,52 +383,32 @@ class VLAJEPAPolicy(PreTrainedPolicy): else: instructions = list(tasks) - # ---- Collect actions (training only) ---- - actions_list = None - action_is_pad_list = None - actions_tensor = batch.get(ACTION) - if actions_tensor is not None: - if actions_tensor.ndim == 2: - actions_tensor = actions_tensor.unsqueeze(1) - actions_list = [actions_tensor[b].detach().cpu().float().numpy() for b in range(batch_size)] - action_is_pad_tensor = batch.get("action_is_pad") - if action_is_pad_tensor is not None: - action_is_pad_list = [action_is_pad_tensor[b].detach().cpu() for b in range(batch_size)] + inputs: dict[str, Any] = {"images": images, "instructions": instructions} - # ---- Collect state ---- - state_list = None - state_tensor = batch.get(OBS_STATE) - if state_tensor is not None: - if state_tensor.ndim > 2: - state_tensor = state_tensor[:, -1, :] - if state_tensor.ndim == 2: - state_tensor = state_tensor.unsqueeze(1) # [B, 1, state_dim] - state_list = [state_tensor[b].detach().cpu().float().numpy() for b in range(batch_size)] + # Videos [B, V, T, C, H, W] - only assembled during training when the world model consumes them. + if self.model.config.enable_world_model and training: + views = [batch[k].unsqueeze(1) if batch[k].ndim == 4 else batch[k] for k in image_keys] + inputs["videos"] = self.model.qwen.to_pixel_values(torch.stack(views, dim=1)) - # ---- Assemble native examples ---- - examples = [] - for b in range(batch_size): - example = { - "image": images_per_sample[b], - "video": videos_per_sample[b], - "lang": instructions[b], - } - if actions_list is not None: - example["action"] = actions_list[b] - if action_is_pad_list is not None: - example["action_is_pad"] = action_is_pad_list[b] - if state_list is not None: - example["state"] = state_list[b] - examples.append(example) + actions = batch.get(ACTION) + if actions is not None: + inputs["actions"] = (actions.unsqueeze(1) if actions.ndim == 2 else actions).float() + if (pad := batch.get("action_is_pad")) is not None: + inputs["action_is_pad"] = pad - return examples + state = batch.get(OBS_STATE) + if state is not None: + if state.ndim > 2: + state = state[:, -1, :] + inputs["state"] = (state.unsqueeze(1) if state.ndim == 2 else state).float() # [B, 1, dim] + + return inputs # ---- LeRobot Policy Interface ---- def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict]: """LeRobot train forward: convert → native forward → aggregate losses.""" - examples = self._prepare_model_inputs(batch) - native_output = self.model.forward(examples) + native_output = self.model.forward(**self._prepare_model_inputs(batch, training=True)) ref = next(iter(native_output.values())) zero = torch.zeros((), device=ref.device, dtype=ref.dtype) @@ -561,16 +426,9 @@ class VLAJEPAPolicy(PreTrainedPolicy): self.eval() self._queues = populate_queues(self._queues, batch, exclude_keys=[ACTION]) - examples = self._prepare_model_inputs(batch) - batch_images = [ex["image"] for ex in examples] - instructions = [ex["lang"] for ex in examples] - - state_np = None - if "state" in examples[0] and examples[0]["state"] is not None: - state_np = np.stack([ex["state"] for ex in examples]) - - actions_np = self.model.predict_action(batch_images, instructions, state_np) - return torch.from_numpy(actions_np).to(device=self.config.device, dtype=torch.float32) + inputs = self._prepare_model_inputs(batch, training=False) + actions = self.model.predict_action(inputs["images"], inputs["instructions"], inputs.get("state")) + return actions.to(device=self.config.device, dtype=torch.float32) @torch.no_grad() def select_action(self, batch: dict[str, Tensor], noise: Tensor | None = None) -> Tensor: diff --git a/src/lerobot/policies/vla_jepa/qwen_interface.py b/src/lerobot/policies/vla_jepa/qwen_interface.py index 24f530efc..bcad1f558 100644 --- a/src/lerobot/policies/vla_jepa/qwen_interface.py +++ b/src/lerobot/policies/vla_jepa/qwen_interface.py @@ -17,9 +17,7 @@ from __future__ import annotations from collections.abc import Sequence from typing import TYPE_CHECKING -import numpy as np import torch -from PIL import Image from lerobot.utils.import_utils import _transformers_available @@ -78,7 +76,7 @@ class Qwen3VLInterface(torch.nn.Module): def build_inputs( self, - images: Sequence[Sequence[Image.Image]], + images: Sequence[Sequence[torch.Tensor]], instructions: Sequence[str], action_prompt: str, embodied_prompt: str, @@ -94,24 +92,42 @@ class Qwen3VLInterface(torch.nn.Module): content.append({"type": "text", "text": prompt}) messages.append([{"role": "user", "content": content}]) + # The Qwen image processor is a torchvision-backed fast processor: passing the + # images as GPU tensors (with `device`) keeps the whole vision pipeline on-device + # and avoids a GPU->CPU->GPU roundtrip. The image tensors are forwarded through + # apply_chat_template untouched into Qwen3VLProcessor.__call__. + # do_rescale=False: images already arrive as float in [0, 1] (the dataset decoder + # yields float32/255 and VISUAL normalization is IDENTITY), so we skip the + # processor's /255 rescale instead of round-tripping through uint8. batch_inputs = self.processor.apply_chat_template( messages, tokenize=True, add_generation_prompt=True, return_dict=True, - processor_kwargs={"padding": True, "return_tensors": "pt"}, + processor_kwargs={ + "padding": True, + "return_tensors": "pt", + "device": self.model.device, + "do_rescale": False, + }, ) return batch_inputs.to(self.model.device) @staticmethod - def tensor_to_pil(image_tensor: torch.Tensor) -> Image.Image: - image = image_tensor.detach().cpu() - if image.ndim == 3 and image.shape[0] in (1, 3): - image = image.permute(1, 2, 0) - image = image.float() - if image.max() <= 1.0: - image = image * 255.0 - image = image.clamp(0, 255).round().to(torch.uint8).numpy() - if image.shape[-1] == 1: - image = np.repeat(image, 3, axis=-1) - return Image.fromarray(image) + def to_pixel_values(image_tensor: torch.Tensor) -> torch.Tensor: + """Prepare an image/video tensor for the fast processors (used with do_rescale=False). + + The dataset decoder yields float32 in [0, 1] (channels-first) and VISUAL + normalization is IDENTITY, so the tensor already arrives in [0, 1]; we pass it + through as float and let the processors normalize (no rescale, no uint8 + quantization). A single channel is expanded to 3 to match the RGB processors. + + Works for any channels-first layout (channel dim is -3): [C, H, W], [B, C, H, W], + [T, C, H, W], [B, V, T, C, H, W], ... + """ + image = image_tensor.detach().float() + if image.shape[-3] == 1: + repeats = [1] * image.ndim + repeats[-3] = 3 + image = image.repeat(*repeats) + return image diff --git a/src/lerobot/rewards/factory.py b/src/lerobot/rewards/factory.py index 474c466de..c791555f8 100644 --- a/src/lerobot/rewards/factory.py +++ b/src/lerobot/rewards/factory.py @@ -133,6 +133,7 @@ def make_reward_model(cfg: RewardModelConfig, **kwargs) -> PreTrainedRewardModel if cfg.pretrained_path: kwargs["pretrained_name_or_path"] = cfg.pretrained_path + kwargs["revision"] = cfg.pretrained_revision reward_model = reward_cls.from_pretrained(**kwargs) else: reward_model = reward_cls(**kwargs) diff --git a/src/lerobot/robots/hope_jr/hope_jr_arm.py b/src/lerobot/robots/hope_jr/hope_jr_arm.py index 4918bcae3..b606a4fe7 100644 --- a/src/lerobot/robots/hope_jr/hope_jr_arm.py +++ b/src/lerobot/robots/hope_jr/hope_jr_arm.py @@ -66,9 +66,14 @@ class HopeJrArm(Robot): @property def _cameras_ft(self) -> dict[str, tuple]: - return { - cam: (self.config.cameras[cam].height, self.config.cameras[cam].width, 3) for cam in self.cameras - } + features: dict[str, tuple] = {} + for cam in self.cameras: + cfg = self.config.cameras[cam] + if getattr(cfg, "use_rgb", True): + features[cam] = (cfg.height, cfg.width, 3) + if getattr(cfg, "use_depth", False): + features[f"{cam}_depth"] = (cfg.height, cfg.width, 1) + return features @cached_property def observation_features(self) -> dict[str, type | tuple]: @@ -139,10 +144,17 @@ class HopeJrArm(Robot): # Capture images from cameras for cam_key, cam in self.cameras.items(): - start = time.perf_counter() - obs_dict[cam_key] = cam.read_latest() - dt_ms = (time.perf_counter() - start) * 1e3 - logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms") + if getattr(cam, "use_rgb", True): + start = time.perf_counter() + obs_dict[cam_key] = cam.read_latest() + dt_ms = (time.perf_counter() - start) * 1e3 + logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms") + + if getattr(cam, "use_depth", False): + start = time.perf_counter() + obs_dict[f"{cam_key}_depth"] = cam.read_latest_depth() + dt_ms = (time.perf_counter() - start) * 1e3 + logger.debug(f"{self} read {cam_key} depth: {dt_ms:.1f}ms") return obs_dict diff --git a/src/lerobot/robots/hope_jr/hope_jr_hand.py b/src/lerobot/robots/hope_jr/hope_jr_hand.py index 566628724..ce70e7e13 100644 --- a/src/lerobot/robots/hope_jr/hope_jr_hand.py +++ b/src/lerobot/robots/hope_jr/hope_jr_hand.py @@ -102,9 +102,14 @@ class HopeJrHand(Robot): @property def _cameras_ft(self) -> dict[str, tuple]: - return { - cam: (self.config.cameras[cam].height, self.config.cameras[cam].width, 3) for cam in self.cameras - } + features: dict[str, tuple] = {} + for cam in self.cameras: + cfg = self.config.cameras[cam] + if getattr(cfg, "use_rgb", True): + features[cam] = (cfg.height, cfg.width, 3) + if getattr(cfg, "use_depth", False): + features[f"{cam}_depth"] = (cfg.height, cfg.width, 1) + return features @cached_property def observation_features(self) -> dict[str, type | tuple]: @@ -170,10 +175,17 @@ class HopeJrHand(Robot): # Capture images from cameras for cam_key, cam in self.cameras.items(): - start = time.perf_counter() - obs_dict[cam_key] = cam.read_latest() - dt_ms = (time.perf_counter() - start) * 1e3 - logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms") + if getattr(cam, "use_rgb", True): + start = time.perf_counter() + obs_dict[cam_key] = cam.read_latest() + dt_ms = (time.perf_counter() - start) * 1e3 + logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms") + + if getattr(cam, "use_depth", False): + start = time.perf_counter() + obs_dict[f"{cam_key}_depth"] = cam.read_latest_depth() + dt_ms = (time.perf_counter() - start) * 1e3 + logger.debug(f"{self} read {cam_key} depth: {dt_ms:.1f}ms") return obs_dict diff --git a/src/lerobot/robots/koch_follower/koch_follower.py b/src/lerobot/robots/koch_follower/koch_follower.py index 3f40ac738..de6f9c4a3 100644 --- a/src/lerobot/robots/koch_follower/koch_follower.py +++ b/src/lerobot/robots/koch_follower/koch_follower.py @@ -68,9 +68,14 @@ class KochFollower(Robot): @property def _cameras_ft(self) -> dict[str, tuple]: - return { - cam: (self.config.cameras[cam].height, self.config.cameras[cam].width, 3) for cam in self.cameras - } + features: dict[str, tuple] = {} + for cam in self.cameras: + cfg = self.config.cameras[cam] + if getattr(cfg, "use_rgb", True): + features[cam] = (cfg.height, cfg.width, 3) + if getattr(cfg, "use_depth", False): + features[f"{cam}_depth"] = (cfg.height, cfg.width, 1) + return features @cached_property def observation_features(self) -> dict[str, type | tuple]: @@ -192,10 +197,17 @@ class KochFollower(Robot): # Capture images from cameras for cam_key, cam in self.cameras.items(): - start = time.perf_counter() - obs_dict[cam_key] = cam.read_latest() - dt_ms = (time.perf_counter() - start) * 1e3 - logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms") + if getattr(cam, "use_rgb", True): + start = time.perf_counter() + obs_dict[cam_key] = cam.read_latest() + dt_ms = (time.perf_counter() - start) * 1e3 + logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms") + + if getattr(cam, "use_depth", False): + start = time.perf_counter() + obs_dict[f"{cam_key}_depth"] = cam.read_latest_depth() + dt_ms = (time.perf_counter() - start) * 1e3 + logger.debug(f"{self} read {cam_key} depth: {dt_ms:.1f}ms") return obs_dict diff --git a/src/lerobot/robots/lekiwi/lekiwi.py b/src/lerobot/robots/lekiwi/lekiwi.py index b73ebeab9..3712a64d3 100644 --- a/src/lerobot/robots/lekiwi/lekiwi.py +++ b/src/lerobot/robots/lekiwi/lekiwi.py @@ -72,6 +72,12 @@ class LeKiwi(Robot): ) self.arm_motors = [motor for motor in self.bus.motors if motor.startswith("arm")] self.base_motors = [motor for motor in self.bus.motors if motor.startswith("base")] + depth_cameras = [name for name, cfg in config.cameras.items() if getattr(cfg, "use_depth", False)] + if depth_cameras: + raise NotImplementedError( + f"Depth cameras are not supported on LeKiwi (got depth-enabled cameras: {depth_cameras}). " + "The host/client transport only carries color frames." + ) self.cameras = make_cameras_from_configs(config.cameras) @property diff --git a/src/lerobot/robots/lekiwi/lekiwi_client.py b/src/lerobot/robots/lekiwi/lekiwi_client.py index fd43e84fe..1bc3dadc4 100644 --- a/src/lerobot/robots/lekiwi/lekiwi_client.py +++ b/src/lerobot/robots/lekiwi/lekiwi_client.py @@ -44,6 +44,13 @@ class LeKiwiClient(Robot): self.id = config.id self.robot_type = config.type + depth_cameras = [name for name, cfg in config.cameras.items() if getattr(cfg, "use_depth", False)] + if depth_cameras: + raise NotImplementedError( + f"Depth cameras are not supported on LeKiwi (got depth-enabled cameras: {depth_cameras}). " + "The host/client transport only carries color frames." + ) + self.remote_ip = config.remote_ip self.port_zmq_cmd = config.port_zmq_cmd self.port_zmq_observations = config.port_zmq_observations diff --git a/src/lerobot/robots/omx_follower/omx_follower.py b/src/lerobot/robots/omx_follower/omx_follower.py index c30eec97a..b2cfb52e9 100644 --- a/src/lerobot/robots/omx_follower/omx_follower.py +++ b/src/lerobot/robots/omx_follower/omx_follower.py @@ -68,9 +68,14 @@ class OmxFollower(Robot): @property def _cameras_ft(self) -> dict[str, tuple]: - return { - cam: (self.config.cameras[cam].height, self.config.cameras[cam].width, 3) for cam in self.cameras - } + features: dict[str, tuple] = {} + for cam in self.cameras: + cfg = self.config.cameras[cam] + if getattr(cfg, "use_rgb", True): + features[cam] = (cfg.height, cfg.width, 3) + if getattr(cfg, "use_depth", False): + features[f"{cam}_depth"] = (cfg.height, cfg.width, 1) + return features @cached_property def observation_features(self) -> dict[str, type | tuple]: @@ -175,10 +180,17 @@ class OmxFollower(Robot): # Capture images from cameras for cam_key, cam in self.cameras.items(): - start = time.perf_counter() - obs_dict[cam_key] = cam.read_latest() - dt_ms = (time.perf_counter() - start) * 1e3 - logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms") + if getattr(cam, "use_rgb", True): + start = time.perf_counter() + obs_dict[cam_key] = cam.read_latest() + dt_ms = (time.perf_counter() - start) * 1e3 + logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms") + + if getattr(cam, "use_depth", False): + start = time.perf_counter() + obs_dict[f"{cam_key}_depth"] = cam.read_latest_depth() + dt_ms = (time.perf_counter() - start) * 1e3 + logger.debug(f"{self} read {cam_key} depth: {dt_ms:.1f}ms") return obs_dict diff --git a/src/lerobot/robots/openarm_follower/openarm_follower.py b/src/lerobot/robots/openarm_follower/openarm_follower.py index 020f24052..e2c7c8cf5 100644 --- a/src/lerobot/robots/openarm_follower/openarm_follower.py +++ b/src/lerobot/robots/openarm_follower/openarm_follower.py @@ -101,9 +101,14 @@ class OpenArmFollower(Robot): @property def _cameras_ft(self) -> dict[str, tuple]: """Camera features for observation space.""" - return { - cam: (self.config.cameras[cam].height, self.config.cameras[cam].width, 3) for cam in self.cameras - } + features: dict[str, tuple] = {} + for cam in self.cameras: + cfg = self.config.cameras[cam] + if getattr(cfg, "use_rgb", True): + features[cam] = (cfg.height, cfg.width, 3) + if getattr(cfg, "use_depth", False): + features[f"{cam}_depth"] = (cfg.height, cfg.width, 1) + return features @cached_property def observation_features(self) -> dict[str, type | tuple]: @@ -242,10 +247,17 @@ class OpenArmFollower(Robot): # Capture images from cameras for cam_key, cam in self.cameras.items(): - start = time.perf_counter() - obs_dict[cam_key] = cam.read_latest() - dt_ms = (time.perf_counter() - start) * 1e3 - logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms") + if getattr(cam, "use_rgb", True): + start = time.perf_counter() + obs_dict[cam_key] = cam.read_latest() + dt_ms = (time.perf_counter() - start) * 1e3 + logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms") + + if getattr(cam, "use_depth", False): + start = time.perf_counter() + obs_dict[f"{cam_key}_depth"] = cam.read_latest_depth() + dt_ms = (time.perf_counter() - start) * 1e3 + logger.debug(f"{self} read {cam_key} depth: {dt_ms:.1f}ms") dt_ms = (time.perf_counter() - start) * 1e3 logger.debug(f"{self} get_observation took: {dt_ms:.1f}ms") diff --git a/src/lerobot/robots/rebot_b601_follower/rebot_b601_follower.py b/src/lerobot/robots/rebot_b601_follower/rebot_b601_follower.py index ec00f4aa9..bf989702b 100644 --- a/src/lerobot/robots/rebot_b601_follower/rebot_b601_follower.py +++ b/src/lerobot/robots/rebot_b601_follower/rebot_b601_follower.py @@ -80,9 +80,14 @@ class RebotB601Follower(Robot): @property def _cameras_ft(self) -> dict[str, tuple]: - return { - cam: (self.config.cameras[cam].height, self.config.cameras[cam].width, 3) for cam in self.cameras - } + features: dict[str, tuple] = {} + for cam in self.cameras: + cfg = self.config.cameras[cam] + if getattr(cfg, "use_rgb", True): + features[cam] = (cfg.height, cfg.width, 3) + if getattr(cfg, "use_depth", False): + features[f"{cam}_depth"] = (cfg.height, cfg.width, 1) + return features @cached_property def observation_features(self) -> dict[str, type | tuple]: @@ -213,10 +218,17 @@ class RebotB601Follower(Robot): logger.debug(f"{self} read state: {dt_ms:.1f}ms") for cam_key, cam in self.cameras.items(): - start = time.perf_counter() - obs_dict[cam_key] = cam.read_latest() - dt_ms = (time.perf_counter() - start) * 1e3 - logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms") + if getattr(cam, "use_rgb", True): + start = time.perf_counter() + obs_dict[cam_key] = cam.read_latest() + dt_ms = (time.perf_counter() - start) * 1e3 + logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms") + + if getattr(cam, "use_depth", False): + start = time.perf_counter() + obs_dict[f"{cam_key}_depth"] = cam.read_latest_depth() + dt_ms = (time.perf_counter() - start) * 1e3 + logger.debug(f"{self} read {cam_key} depth: {dt_ms:.1f}ms") return obs_dict diff --git a/src/lerobot/robots/so_follower/so_follower.py b/src/lerobot/robots/so_follower/so_follower.py index 0651f566c..c6e67fafe 100644 --- a/src/lerobot/robots/so_follower/so_follower.py +++ b/src/lerobot/robots/so_follower/so_follower.py @@ -68,9 +68,13 @@ class SOFollower(Robot): @property def _cameras_ft(self) -> dict[str, tuple]: - return { - cam: (self.config.cameras[cam].height, self.config.cameras[cam].width, 3) for cam in self.cameras - } + features: dict[str, tuple] = {} + for cam in self.cameras: + if getattr(self.cameras[cam], "use_rgb", True): + features[cam] = (self.cameras[cam].height, self.cameras[cam].width, 3) + if getattr(self.cameras[cam], "use_depth", False): + features[f"{cam}_depth"] = (self.cameras[cam].height, self.cameras[cam].width, 1) + return features @cached_property def observation_features(self) -> dict[str, type | tuple]: @@ -185,10 +189,17 @@ class SOFollower(Robot): # Capture images from cameras for cam_key, cam in self.cameras.items(): - start = time.perf_counter() - obs_dict[cam_key] = cam.read_latest() - dt_ms = (time.perf_counter() - start) * 1e3 - logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms") + if getattr(cam, "use_rgb", True): + start = time.perf_counter() + obs_dict[cam_key] = cam.read_latest() + dt_ms = (time.perf_counter() - start) * 1e3 + logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms") + + if getattr(cam, "use_depth", False): + start = time.perf_counter() + obs_dict[f"{cam_key}_depth"] = cam.read_latest_depth() + dt_ms = (time.perf_counter() - start) * 1e3 + logger.debug(f"{self} read {cam_key} depth: {dt_ms:.1f}ms") return obs_dict diff --git a/src/lerobot/robots/unitree_g1/unitree_g1.py b/src/lerobot/robots/unitree_g1/unitree_g1.py index 25ec32716..5b8be0941 100644 --- a/src/lerobot/robots/unitree_g1/unitree_g1.py +++ b/src/lerobot/robots/unitree_g1/unitree_g1.py @@ -222,9 +222,14 @@ class UnitreeG1(Robot): @property def _cameras_ft(self) -> dict[str, tuple]: - return { - cam: (self.config.cameras[cam].height, self.config.cameras[cam].width, 3) for cam in self.cameras - } + features: dict[str, tuple] = {} + for cam in self.cameras: + cfg = self.config.cameras[cam] + if getattr(cfg, "use_rgb", True): + features[cam] = (cfg.height, cfg.width, 3) + if getattr(cfg, "use_depth", False): + features[f"{cam}_depth"] = (cfg.height, cfg.width, 1) + return features @cached_property def observation_features(self) -> dict[str, type | tuple]: @@ -458,7 +463,10 @@ class UnitreeG1(Robot): # Cameras - read images from ZMQ cameras for cam_name, cam in self._cameras.items(): - obs[cam_name] = cam.read_latest() + if getattr(cam, "use_rgb", True): + obs[cam_name] = cam.read_latest() + if getattr(cam, "use_depth", False): + obs[f"{cam_name}_depth"] = cam.read_latest_depth() return obs diff --git a/src/lerobot/rollout/context.py b/src/lerobot/rollout/context.py index 5e3b60674..863dc1058 100644 --- a/src/lerobot/rollout/context.py +++ b/src/lerobot/rollout/context.py @@ -320,7 +320,9 @@ def build_rollout_context( raise ValueError( f"Visual feature mismatch between policy and robot hardware.\n" f"Policy expects: {expected_visuals}\n" - f"Robot provides: {provided_visuals}" + f"Robot provides: {provided_visuals}\n" + f"Use --rename_map to map camera names, e.g. " + f"""--rename_map='{{"observation.images.top": "observation.images.cam0"}}'""" ) # --- 5. Dataset ------------- @@ -332,7 +334,8 @@ def build_rollout_context( cfg.dataset.repo_id, root=cfg.dataset.root, batch_encoding_size=cfg.dataset.video_encoding_batch_size, - camera_encoder=cfg.dataset.camera_encoder, + rgb_encoder=cfg.dataset.rgb_encoder, + depth_encoder=cfg.dataset.depth_encoder, streaming_encoding=cfg.dataset.streaming_encoding, encoder_queue_maxsize=cfg.dataset.encoder_queue_maxsize, encoder_threads=cfg.dataset.encoder_threads, @@ -372,7 +375,8 @@ def build_rollout_context( image_writer_threads=cfg.dataset.num_image_writer_threads_per_camera * len(robot.cameras if hasattr(robot, "cameras") else []), batch_encoding_size=cfg.dataset.video_encoding_batch_size, - camera_encoder=cfg.dataset.camera_encoder, + rgb_encoder=cfg.dataset.rgb_encoder, + depth_encoder=cfg.dataset.depth_encoder, streaming_encoding=cfg.dataset.streaming_encoding, encoder_queue_maxsize=cfg.dataset.encoder_queue_maxsize, encoder_threads=cfg.dataset.encoder_threads, diff --git a/src/lerobot/rollout/strategies/dagger.py b/src/lerobot/rollout/strategies/dagger.py index b46828f0c..2c306de68 100644 --- a/src/lerobot/rollout/strategies/dagger.py +++ b/src/lerobot/rollout/strategies/dagger.py @@ -47,8 +47,6 @@ from __future__ import annotations import contextlib import enum import logging -import os -import sys import time from concurrent.futures import Future, ThreadPoolExecutor from threading import Event, Lock @@ -58,7 +56,6 @@ import numpy as np from lerobot.common.control_utils import ( follower_smooth_move_to, - is_headless, teleop_smooth_move_to, teleop_supports_feedback, ) @@ -66,7 +63,7 @@ from lerobot.datasets import VideoEncodingManager from lerobot.datasets.utils import DEFAULT_VIDEO_FILE_SIZE_IN_MB from lerobot.utils.constants import ACTION, OBS_STR from lerobot.utils.feature_utils import build_dataset_frame -from lerobot.utils.import_utils import _pynput_available +from lerobot.utils.keyboard_input import create_key_listener from lerobot.utils.pedal import start_pedal_listener from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.utils import log_say @@ -75,19 +72,6 @@ from ..configs import DAggerKeyboardConfig, DAggerPedalConfig, DAggerStrategyCon from ..context import RolloutContext from .core import RolloutStrategy, estimate_max_episode_seconds, safe_push_to_hub, send_next_action -PYNPUT_AVAILABLE = _pynput_available -keyboard = None -if PYNPUT_AVAILABLE: - try: - if ("DISPLAY" not in os.environ) and ("linux" in sys.platform): - logging.info("No DISPLAY set. Skipping pynput import.") - PYNPUT_AVAILABLE = False - else: - from pynput import keyboard - except Exception as e: - PYNPUT_AVAILABLE = False - logging.info(f"Could not import pynput: {e}") - logger = logging.getLogger(__name__) @@ -201,73 +185,42 @@ class DAggerEvents: def _init_dagger_keyboard(events: DAggerEvents, cfg: DAggerKeyboardConfig): - """Initialise keyboard listener with DAgger 3-key controls. + """Initialise a keyboard listener for DAgger's 3 controls. - Returns the pynput Listener (or ``None`` in headless mode or when - pynput is unavailable). + Backend selection (pynput on X11 / trusted-macOS / Windows, a terminal reader on + Wayland / headless TTY) is delegated to :func:`create_key_listener`. Returns the + listener (exposing ``stop()``) or ``None`` when no keyboard backend is usable. """ - if not PYNPUT_AVAILABLE or is_headless(): - logger.warning("Headless environment or pynput unavailable — keyboard controls disabled") - return None - - # Map config key names to pynput Key objects for special keys - special_keys = { - "space": keyboard.Key.space, - "tab": keyboard.Key.tab, - "enter": keyboard.Key.enter, - } - - def _resolve_key(key) -> str | None: - """Resolve a pynput key event to a config-comparable string.""" - if key == keyboard.Key.esc: - return "esc" - for name, pynput_key in special_keys.items(): - if key == pynput_key: - return name - if hasattr(key, "char") and key.char: - return key.char - return None - - # Build mapping: resolved key string -> DAgger event name + # Map config key names to DAgger event names. key_to_event = { cfg.pause_resume: "pause_resume", cfg.correction: "correction", } - def on_press(key): - try: - resolved = _resolve_key(key) - if resolved is None: - return - if resolved == "esc": - logger.info("Stop recording...") - events.stop_recording.set() - return - if resolved in key_to_event: - events.request_transition(key_to_event[resolved]) - if resolved == cfg.upload: - events.upload_requested.set() - if resolved == cfg.success: - events.mark_success() - logger.info("Episode marked as SUCCESS") - if resolved == cfg.failure: - events.mark_failure() - logger.info("Episode marked as FAILURE") - except Exception as e: - logger.debug("Key error: %s", e) + def dispatch(name: str) -> None: + """Apply a resolved key name to the DAgger events.""" + if name == "esc": + logger.info("Stop recording...") + events.stop_recording.set() + return + if name in key_to_event: + events.request_transition(key_to_event[name]) + if name == cfg.upload: + events.upload_requested.set() + if name == cfg.success: + events.mark_success() + logger.info("Episode marked as SUCCESS") + if name == cfg.failure: + events.mark_failure() + logger.info("Episode marked as FAILURE") - listener = keyboard.Listener(on_press=on_press) - listener.start() - logger.info( - "DAgger keyboard listener started (pause_resume='%s', correction='%s', " - "upload='%s', success='%s', failure='%s', ESC=stop)", - cfg.pause_resume, - cfg.correction, - cfg.upload, - cfg.success, - cfg.failure, + return create_key_listener( + dispatch, + controls_help=( + f"pause_resume='{cfg.pause_resume}', correction='{cfg.correction}', " + f"upload='{cfg.upload}', success='{cfg.success}', failure='{cfg.failure}', ESC=stop" + ), ) - return listener def _init_dagger_pedal(events: DAggerEvents, cfg: DAggerPedalConfig): @@ -364,7 +317,7 @@ class DAggerStrategy(RolloutStrategy): logger.info("Stopping DAgger recording") log_say("Stopping DAgger recording", play_sounds) - if self._listener is not None and not is_headless(): + if self._listener is not None: logger.info("Stopping keyboard listener") self._listener.stop() diff --git a/src/lerobot/rollout/strategies/episodic.py b/src/lerobot/rollout/strategies/episodic.py index e925fb2ea..e70e66787 100644 --- a/src/lerobot/rollout/strategies/episodic.py +++ b/src/lerobot/rollout/strategies/episodic.py @@ -35,14 +35,13 @@ import time from lerobot.common.control_utils import ( follower_smooth_move_to, - init_keyboard_listener, - is_headless, teleop_smooth_move_to, teleop_supports_feedback, ) from lerobot.datasets import VideoEncodingManager from lerobot.utils.constants import ACTION, OBS_STR from lerobot.utils.feature_utils import build_dataset_frame +from lerobot.utils.keyboard_input import init_keyboard_listener from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.utils import log_say from lerobot.utils.visualization_utils import log_rerun_data @@ -307,7 +306,7 @@ class EpisodicStrategy(RolloutStrategy): log_say("Stop recording", play_sounds, blocking=True) - if not is_headless() and self._listener is not None: + if self._listener is not None: self._listener.stop() if ctx.data.dataset is not None: diff --git a/src/lerobot/rollout/strategies/highlight.py b/src/lerobot/rollout/strategies/highlight.py index baff70da7..385a9e2b6 100644 --- a/src/lerobot/rollout/strategies/highlight.py +++ b/src/lerobot/rollout/strategies/highlight.py @@ -18,17 +18,14 @@ from __future__ import annotations import contextlib import logging -import os -import sys import time from concurrent.futures import Future, ThreadPoolExecutor from threading import Event as ThreadingEvent, Lock -from lerobot.common.control_utils import is_headless from lerobot.datasets import VideoEncodingManager from lerobot.utils.constants import ACTION, OBS_STR from lerobot.utils.feature_utils import build_dataset_frame -from lerobot.utils.import_utils import _pynput_available, require_package +from lerobot.utils.keyboard_input import create_key_listener from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.utils import log_say @@ -37,19 +34,6 @@ from ..context import RolloutContext from ..ring_buffer import RolloutRingBuffer from .core import RolloutStrategy, safe_push_to_hub, send_next_action -PYNPUT_AVAILABLE = _pynput_available -keyboard = None -if PYNPUT_AVAILABLE: - try: - if ("DISPLAY" not in os.environ) and ("linux" in sys.platform): - logging.info("No DISPLAY set. Skipping pynput import.") - PYNPUT_AVAILABLE = False - else: - from pynput import keyboard - except Exception as e: - PYNPUT_AVAILABLE = False - logging.info(f"Could not import pynput: {e}") - logger = logging.getLogger(__name__) @@ -72,7 +56,6 @@ class HighlightStrategy(RolloutStrategy): def __init__(self, config: HighlightStrategyConfig): super().__init__(config) - require_package("pynput", extra="pynput-dep") self._ring: RolloutRingBuffer | None = None self._listener = None self._save_requested = ThreadingEvent() @@ -234,30 +217,27 @@ class HighlightStrategy(RolloutStrategy): logger.info("Highlight strategy teardown complete") def _setup_keyboard(self, shutdown_event: ThreadingEvent) -> None: - """Set up keyboard listener for save and push keys.""" - if is_headless(): - logger.warning("Headless environment — highlight keys unavailable") - return + """Set up a keyboard listener for the save and push keys. - try: - save_key = self.config.save_key - push_key = self.config.push_key + Backend selection (pynput on X11 / trusted-macOS / Windows, a terminal reader on + Wayland / headless TTY) is delegated to :func:`create_key_listener`. + """ + save_key = self.config.save_key + push_key = self.config.push_key - def on_press(key): - with contextlib.suppress(Exception): - if hasattr(key, "char") and key.char == save_key: - self._save_requested.set() - elif hasattr(key, "char") and key.char == push_key: - self._push_requested.set() - elif key == keyboard.Key.esc: - self._save_requested.clear() - shutdown_event.set() + def dispatch(name: str) -> None: + """Apply a resolved key name to the highlight events.""" + if name == save_key: + self._save_requested.set() + elif name == push_key: + self._push_requested.set() + elif name == "esc": + self._save_requested.clear() + shutdown_event.set() - self._listener = keyboard.Listener(on_press=on_press) - self._listener.start() - logger.info("Keyboard listener started (save='%s', push='%s', ESC=stop)", save_key, push_key) - except ImportError: - logger.warning("pynput not available — keyboard listener disabled") + self._listener = create_key_listener( + dispatch, controls_help=f"save='{save_key}', push='{push_key}', ESC=stop" + ) def _background_push(self, dataset, cfg) -> None: """Queue a Hub push on the single-worker executor.""" diff --git a/src/lerobot/scripts/lerobot_dataset_viz.py b/src/lerobot/scripts/lerobot_dataset_viz.py index d07a2767d..22a7208d4 100644 --- a/src/lerobot/scripts/lerobot_dataset_viz.py +++ b/src/lerobot/scripts/lerobot_dataset_viz.py @@ -77,15 +77,68 @@ from lerobot.utils.constants import ACTION, DONE, OBS_STATE, REWARD from lerobot.utils.utils import init_logging +def get_feature_names(dataset: LeRobotDataset, key: str) -> list[str]: + """Return per-dimension names for a feature from the dataset metadata. + + Only flat-list ``names`` metadata is used. Dict-style ``names`` and missing names fall back to ``{key}_{i}`` indices. + """ + feature = dataset.features[key] + dim = feature["shape"][-1] + + names = feature.get("names") + if isinstance(names, list) and len(names) == dim: + return [str(name) for name in names] + + return [f"{key}_{d}" for d in range(dim)] + + +def check_chw_float32(frame: torch.Tensor) -> None: + """ + Check if a frame is a channel-first, float32 tensor. + """ + assert frame.dtype == torch.float32 + assert frame.ndim == 3 + c, h, w = frame.shape + assert c < h and c < w, f"expect channel first images, but instead {frame.shape}" + + def to_hwc_uint8_numpy(chw_float32_torch: torch.Tensor) -> np.ndarray: - assert chw_float32_torch.dtype == torch.float32 - assert chw_float32_torch.ndim == 3 - c, h, w = chw_float32_torch.shape - assert c < h and c < w, f"expect channel first images, but instead {chw_float32_torch.shape}" + check_chw_float32(chw_float32_torch) hwc_uint8_numpy = (chw_float32_torch * 255).type(torch.uint8).permute(1, 2, 0).numpy() return hwc_uint8_numpy +def build_blueprint_from_dataset(dataset: LeRobotDataset): + """Build a Rerun blueprint laying out camera images and time series for the given dataset. + + Camera images and scalar signals (action, state, reward, done, success) are arranged in a grid. + The per-dimension series names for ``action`` and ``state`` are applied directly + via blueprint overrides. + """ + import rerun as rr + import rerun.blueprint as rrb + + views = [rrb.Spatial2DView(origin=key, name=key) for key in dataset.meta.camera_keys] + + # Style multi-dimensional signals (action, state) with per-dimension names. + for origin, key in ((ACTION, ACTION), ("state", OBS_STATE)): + if key in dataset.features: + names = get_feature_names(dataset, key) + styling = rr.SeriesLines(names=names) + views.append(rrb.TimeSeriesView(origin=origin, name=origin, overrides={origin: styling})) + for key in (DONE, REWARD, "next.success"): + if key in dataset.features: + views.append(rrb.TimeSeriesView(origin=key, name=key)) + + return rrb.Blueprint(rrb.Grid(*views)) + + +def to_hwc_uint16_numpy(chw_float32_torch: torch.Tensor) -> np.ndarray: + check_chw_float32(chw_float32_torch) + hwc_uint16_numpy = chw_float32_torch.round().type(torch.uint16).permute(1, 2, 0).numpy() + return hwc_uint16_numpy + + def visualize_dataset( dataset: LeRobotDataset, episode_index: int, @@ -124,7 +177,8 @@ def visualize_dataset( import rerun as rr spawn_local_viewer = mode == "local" and not save - rr.init(f"{repo_id}/episode_{episode_index}", spawn=spawn_local_viewer) + blueprint = build_blueprint_from_dataset(dataset) + rr.init(f"{repo_id}/episode_{episode_index}", spawn=spawn_local_viewer, default_blueprint=blueprint) # Manually call python garbage collector after `rr.init` to avoid hanging in a blocking flush # when iterating on a dataloader with `num_workers` > 0 @@ -138,30 +192,46 @@ def visualize_dataset( logging.info("Logging to Rerun") + # Use the dataset's q01/q99 depth statistics for robust depth range bounds + depth_ranges = {} + for key in dataset.meta.depth_keys: + stats = dataset.meta.stats[key] + lo = stats["q01"] if "q01" in stats else stats["min"] + hi = stats["q99"] if "q99" in stats else stats["max"] + depth_ranges[key] = (float(np.asarray(lo).item()), float(np.asarray(hi).item())) + first_index = None for batch in tqdm.tqdm(dataloader, total=len(dataloader)): if first_index is None: first_index = batch["index"][0].item() + # iterate over the batch for i in range(len(batch["index"])): rr.set_time("frame_index", sequence=batch["index"][i].item() - first_index) rr.set_time("timestamp", timestamp=batch["timestamp"][i].item()) - # display each camera image + # display each camera image (or depth map) for key in dataset.meta.camera_keys: - img = to_hwc_uint8_numpy(batch[key][i]) - img_entity = rr.Image(img).compress() if display_compressed_images else rr.Image(img) - rr.log(key, entity=img_entity) + if key in dataset.meta.depth_keys: + depth = to_hwc_uint16_numpy(batch[key][i]) + depth_entity = rr.DepthImage( + depth, + colormap=rr.components.Colormap.Viridis, + depth_range=depth_ranges[key], + ) + rr.log(key, entity=depth_entity) + else: + img = to_hwc_uint8_numpy(batch[key][i]) + img_entity = rr.Image(img).compress() if display_compressed_images else rr.Image(img) + rr.log(key, entity=img_entity) - # display each dimension of action space (e.g. actuators command) + # display the action space (e.g. actuators command) if ACTION in batch: - for dim_idx, val in enumerate(batch[ACTION][i]): - rr.log(f"{ACTION}/{dim_idx}", rr.Scalars(val.item())) + rr.log(ACTION, rr.Scalars(batch[ACTION][i].numpy())) - # display each dimension of observed state space (e.g. agent position in joint space) + # display the observed state space (e.g. agent position in joint space) if OBS_STATE in batch: - for dim_idx, val in enumerate(batch[OBS_STATE][i]): - rr.log(f"state/{dim_idx}", rr.Scalars(val.item())) + rr.log("state", rr.Scalars(batch[OBS_STATE][i].numpy())) if DONE in batch: rr.log(DONE, rr.Scalars(batch[DONE][i].item())) @@ -172,9 +242,8 @@ def visualize_dataset( if "next.success" in batch: rr.log("next.success", rr.Scalars(batch["next.success"][i].item())) + # save .rrd locally if mode == "local" and save: - # save .rrd locally - output_dir = Path(output_dir) output_dir.mkdir(parents=True, exist_ok=True) repo_id_str = repo_id.replace("/", "_") rrd_path = output_dir / f"{repo_id_str}_episode_{episode_index}.rrd" @@ -182,7 +251,7 @@ def visualize_dataset( return rrd_path elif mode == "distant": - # stop the process from exiting since it is serving the websocket connection + # Keep the process alive while it serves the gRPC/web connection. try: while True: time.sleep(1) @@ -297,12 +366,14 @@ def main(): ) logging.warning("Setting grpc_port to ws_port value.") kwargs["grpc_port"] = kwargs.pop("ws_port") + else: + kwargs.pop("ws_port") # Always remove ws_port from kwargs init_logging() logging.info("Loading dataset") dataset = LeRobotDataset(repo_id, episodes=[args.episode_index], root=root, tolerance_s=tolerance_s) - visualize_dataset(dataset, **vars(args)) + visualize_dataset(dataset, **kwargs) if __name__ == "__main__": diff --git a/src/lerobot/scripts/lerobot_edit_dataset.py b/src/lerobot/scripts/lerobot_edit_dataset.py index eaadf47de..42dce438f 100644 --- a/src/lerobot/scripts/lerobot_edit_dataset.py +++ b/src/lerobot/scripts/lerobot_edit_dataset.py @@ -133,6 +133,15 @@ Convert image dataset to video format and save locally: --new_root /path/to/output/pusht_video \ --operation.type convert_image_to_video +Convert image dataset (with depth maps) to video format, customizing the depth encoder: + lerobot-edit-dataset \ + --repo_id lerobot/pusht_image \ + --new_root /path/to/output/pusht_video \ + --operation.type convert_image_to_video \ + --operation.depth_encoder.depth_min 0.01 \ + --operation.depth_encoder.depth_max 10.0 \ + --operation.depth_encoder.use_log true + Convert image dataset to video format and save with new repo_id: lerobot-edit-dataset \ --repo_id lerobot/pusht_image \ @@ -190,17 +199,17 @@ Re-encode all videos in a dataset (saves to lerobot/pusht_reencoded by default): lerobot-edit-dataset \ --repo_id lerobot/pusht \ --operation.type reencode_videos \ - --operation.camera_encoder.vcodec h264 \ - --operation.camera_encoder.pix_fmt yuv420p \ - --operation.camera_encoder.crf 23 + --operation.rgb_encoder.vcodec h264 \ + --operation.rgb_encoder.pix_fmt yuv420p \ + --operation.rgb_encoder.crf 23 Re-encode videos into a new dataset using 4 parallel processes: lerobot-edit-dataset \ --repo_id lerobot/pusht \ --new_repo_id lerobot/pusht_h264 \ --operation.type reencode_videos \ - --operation.camera_encoder.vcodec h264 \ - --operation.camera_encoder.crf 23 \ + --operation.rgb_encoder.vcodec h264 \ + --operation.rgb_encoder.crf 23 \ --operation.num_workers 4 Re-encode videos in-place (overwrites original dataset): @@ -208,9 +217,16 @@ Re-encode videos in-place (overwrites original dataset): --repo_id lerobot/pusht \ --new_repo_id lerobot/pusht \ --operation.type reencode_videos \ - --operation.camera_encoder.vcodec h264 \ + --operation.rgb_encoder.vcodec h264 \ --operation.overwrite true +Re-encode both RGB and depth videos in a dataset (depth quantization params are preserved): + lerobot-edit-dataset \ + --repo_id lerobot/pusht_depth \ + --operation.type reencode_videos \ + --operation.rgb_encoder.vcodec h264 \ + --operation.depth_encoder.extra_options '{"x265-params": "lossless=1"}' + Using JSON config file: lerobot-edit-dataset \ --config_path path/to/edit_config.json @@ -225,7 +241,13 @@ from pathlib import Path import draccus -from lerobot.configs import VideoEncoderConfig, camera_encoder_defaults, parser +from lerobot.configs import ( + DepthEncoderConfig, + RGBEncoderConfig, + depth_encoder_defaults, + parser, + rgb_encoder_defaults, +) from lerobot.datasets import ( LeRobotDataset, convert_image_to_video_dataset, @@ -287,7 +309,8 @@ class ModifyTasksConfig(OperationConfig): @dataclass class ConvertImageToVideoConfig(OperationConfig): output_dir: str | None = None - camera_encoder: VideoEncoderConfig = field(default_factory=camera_encoder_defaults) + rgb_encoder: RGBEncoderConfig = field(default_factory=rgb_encoder_defaults) + depth_encoder: DepthEncoderConfig = field(default_factory=depth_encoder_defaults) episode_indices: list[int] | None = None num_workers: int = 4 max_episodes_per_batch: int | None = None @@ -308,7 +331,8 @@ class RecomputeStatsConfig(OperationConfig): @OperationConfig.register_subclass("reencode_videos") @dataclass class ReencodeVideosConfig(OperationConfig): - camera_encoder: VideoEncoderConfig = field(default_factory=camera_encoder_defaults) + rgb_encoder: RGBEncoderConfig = field(default_factory=rgb_encoder_defaults) + depth_encoder: DepthEncoderConfig = field(default_factory=depth_encoder_defaults) num_workers: int = 0 encoder_threads: int | None = None overwrite: bool = False @@ -601,7 +625,8 @@ def handle_convert_image_to_video(cfg: EditDatasetConfig) -> None: dataset=dataset, output_dir=output_dir, repo_id=output_repo_id, - camera_encoder=getattr(cfg.operation, "camera_encoder", None) or camera_encoder_defaults(), + rgb_encoder=getattr(cfg.operation, "rgb_encoder", None) or rgb_encoder_defaults(), + depth_encoder=getattr(cfg.operation, "depth_encoder", None) or depth_encoder_defaults(), episode_indices=getattr(cfg.operation, "episode_indices", None), num_workers=getattr(cfg.operation, "num_workers", 4), max_episodes_per_batch=getattr(cfg.operation, "max_episodes_per_batch", None), @@ -719,10 +744,14 @@ def handle_reencode_videos(cfg: EditDatasetConfig) -> None: shutil.copytree(input_root, output_root) dataset = LeRobotDataset(output_repo_id, root=output_root) - logging.info(f"Re-encoding videos in {output_repo_id} with {cfg.operation.camera_encoder}") + logging.info( + f"Re-encoding videos in {output_repo_id} with RGB encoder {cfg.operation.rgb_encoder} " + f"and depth encoder {cfg.operation.depth_encoder}" + ) reencode_dataset( dataset, - camera_encoder=cfg.operation.camera_encoder, + rgb_encoder=cfg.operation.rgb_encoder, + depth_encoder=cfg.operation.depth_encoder, encoder_threads=cfg.operation.encoder_threads, num_workers=cfg.operation.num_workers, ) diff --git a/src/lerobot/scripts/lerobot_eval.py b/src/lerobot/scripts/lerobot_eval.py index d45483d21..1ec4ea75f 100644 --- a/src/lerobot/scripts/lerobot_eval.py +++ b/src/lerobot/scripts/lerobot_eval.py @@ -72,8 +72,9 @@ from termcolor import colored from torch import Tensor, nn from tqdm import trange -from lerobot.configs import parser +from lerobot.configs import FeatureType, parser from lerobot.configs.eval import EvalPipelineConfig +from lerobot.datasets.lerobot_dataset import LeRobotDataset from lerobot.envs import ( check_env_attributes_and_types, close_envs, @@ -84,7 +85,7 @@ from lerobot.envs import ( from lerobot.policies import PreTrainedPolicy, make_policy, make_pre_post_processors from lerobot.processor import PolicyProcessorPipeline from lerobot.types import PolicyAction -from lerobot.utils.constants import ACTION, DONE, OBS_STR, REWARD +from lerobot.utils.constants import ACTION, DONE, OBS_IMAGE, OBS_IMAGES, OBS_STR, REWARD from lerobot.utils.device_utils import get_safe_torch_device from lerobot.utils.import_utils import register_third_party_plugins from lerobot.utils.io_utils import write_video @@ -95,6 +96,65 @@ from lerobot.utils.utils import ( ) +def _env_features_to_dataset_features(env_features: dict) -> dict: + """Convert EnvConfig.features to the dict format expected by LeRobotDataset.create().""" + features = {} + for key, ft in env_features.items(): + shape = tuple(ft.shape) + if ft.type is FeatureType.VISUAL: + features[key] = {"dtype": "video", "shape": shape, "names": ["height", "width", "channel"]} + else: + features[key] = {"dtype": "float32", "shape": shape, "names": None} + features["next.reward"] = {"dtype": "float32", "shape": (1,), "names": None} + features["next.success"] = {"dtype": "bool", "shape": (1,), "names": None} + features["next.done"] = {"dtype": "bool", "shape": (1,), "names": None} + return features + + +def _build_raw_frame( + raw_obs: dict, + env_idx: int, + action: np.ndarray, + reward: float, + success: bool, + done: bool, + task: str, + env_features: dict, +) -> dict: + """Build a dataset frame from raw env observations for one env index. + + Keys in the frame match the keys in env_features so they align with the + dataset schema created by _env_features_to_dataset_features(). + """ + frame: dict[str, Any] = {} + for key in env_features: + if key == ACTION: + continue + if key.startswith("next."): + continue + if "pixels" in raw_obs and isinstance(raw_obs["pixels"], dict): + for cam_name, img in raw_obs["pixels"].items(): + candidate = f"{OBS_IMAGES}.{cam_name}" + if candidate == key: + frame[key] = img[env_idx] + if key in frame: + continue + if "pixels" in raw_obs and not isinstance(raw_obs["pixels"], dict) and key in ("pixels", OBS_IMAGE): + frame[key] = raw_obs["pixels"][env_idx] + continue + if key in raw_obs and isinstance(raw_obs[key], np.ndarray): + val = raw_obs[key][env_idx] + if val.dtype == np.float64: + val = val.astype(np.float32) + frame[key] = val + frame[ACTION] = action + frame["next.reward"] = np.atleast_1d(np.float32(reward)) + frame["next.success"] = np.atleast_1d(np.bool_(success)) + frame["next.done"] = np.atleast_1d(np.bool_(done)) + frame["task"] = task + return frame + + def rollout( env: gym.vector.VectorEnv, policy: PreTrainedPolicy, @@ -105,6 +165,10 @@ def rollout( seeds: list[int] | None = None, return_observations: bool = False, render_callback: Callable[[gym.vector.VectorEnv], None] | None = None, + recording_dir: Path | None = None, + env_features: dict | None = None, + recording_repo_id: str | None = None, + recording_private: bool = False, ) -> dict: """Run a batched policy rollout once through a batch of environments. @@ -145,6 +209,33 @@ def rollout( if render_callback is not None: render_callback(env) + recording_datasets: list[LeRobotDataset] | None = None + raw_observation = None + task_desc = "" + if recording_dir is not None and env_features is not None: + features = _env_features_to_dataset_features(env_features) + fps = env.unwrapped.metadata.get("render_fps", 30) + recording_datasets = [] + multi_env = env.num_envs > 1 + base_repo_id = recording_repo_id or "eval_recording" + for i in range(env.num_envs): + root = str(recording_dir / f"env_{i}") if multi_env else str(recording_dir) + repo_id = f"{base_repo_id}_env_{i}" if multi_env else base_repo_id + recording_datasets.append( + LeRobotDataset.create( + repo_id=repo_id, + fps=fps, + features=features, + root=root, + use_videos=True, + ) + ) + raw_observation = deepcopy(observation) + try: + task_desc = list(env.call("task_description"))[0] + except (AttributeError, NotImplementedError): + task_desc = "" + all_observations = [] all_actions = [] all_rewards = [] @@ -162,80 +253,112 @@ def rollout( leave=False, ) check_env_attributes_and_types(env) - while not np.all(done) and step < max_steps: - # Numpy array to tensor and changing dictionary keys to LeRobot policy format. - observation = preprocess_observation(observation) - if return_observations: - all_observations.append(deepcopy(observation)) + try: + while not np.all(done) and step < max_steps: + # Numpy array to tensor and changing dictionary keys to LeRobot policy format. + observation = preprocess_observation(observation) + if return_observations: + all_observations.append(deepcopy(observation)) - # Infer "task" from sub-environments (prefer natural language description). - # env.call() works with both SyncVectorEnv and AsyncVectorEnv. - try: - observation["task"] = list(env.call("task_description")) - except (AttributeError, NotImplementedError): + # Infer "task" from sub-environments (prefer natural language description). + # env.call() works with both SyncVectorEnv and AsyncVectorEnv. try: - observation["task"] = list(env.call("task")) + observation["task"] = list(env.call("task_description")) except (AttributeError, NotImplementedError): - observation["task"] = [""] * env.num_envs + try: + observation["task"] = list(env.call("task")) + except (AttributeError, NotImplementedError): + observation["task"] = [""] * env.num_envs - # Apply environment-specific preprocessing (e.g., LiberoProcessorStep for LIBERO) - observation = env_preprocessor(observation) + # Apply environment-specific preprocessing (e.g., LiberoProcessorStep for LIBERO) + observation = env_preprocessor(observation) - observation = preprocessor(observation) - with torch.inference_mode(): - action = policy.select_action(observation) - action = postprocessor(action) + observation = preprocessor(observation) + with torch.inference_mode(): + action = policy.select_action(observation) + action = postprocessor(action) - action_transition = {ACTION: action} - action_transition = env_postprocessor(action_transition) - action = action_transition[ACTION] + action_transition = {ACTION: action} + action_transition = env_postprocessor(action_transition) + action = action_transition[ACTION] - # Convert to CPU / numpy. - action_numpy: np.ndarray = action.to("cpu").numpy() - assert action_numpy.ndim == 2, "Action dimensions should be (batch, action_dim)" + # Convert to CPU / numpy. + action_numpy: np.ndarray = action.to("cpu").numpy() + assert action_numpy.ndim == 2, "Action dimensions should be (batch, action_dim)" - # Apply the next action. - observation, reward, terminated, truncated, info = env.step(action_numpy) - if render_callback is not None: - render_callback(env) + # Apply the next action. + observation, reward, terminated, truncated, info = env.step(action_numpy) + if render_callback is not None: + render_callback(env) - # VectorEnv stores is_success in `info["final_info"][env_index]["is_success"]`. "final_info" isn't - # available if none of the envs finished. - if "final_info" in info: - final_info = info["final_info"] - if not isinstance(final_info, dict): - raise RuntimeError( - "Unsupported `final_info` format: expected dict (Gymnasium >= 1.0). " - "You're likely using an older version of gymnasium (< 1.0). Please upgrade." + # VectorEnv stores is_success in `info["final_info"][env_index]["is_success"]`. "final_info" isn't + # available if none of the envs finished. + if "final_info" in info: + final_info = info["final_info"] + if not isinstance(final_info, dict): + raise RuntimeError( + "Unsupported `final_info` format: expected dict (Gymnasium >= 1.0). " + "You're likely using an older version of gymnasium (< 1.0). Please upgrade." + ) + successes = final_info["is_success"].tolist() + elif "is_success" in info: + is_success = info["is_success"] + successes = ( + is_success.tolist() + if hasattr(is_success, "tolist") + else [bool(is_success)] * env.num_envs ) - successes = final_info["is_success"].tolist() - elif "is_success" in info: - is_success = info["is_success"] - successes = ( - is_success.tolist() if hasattr(is_success, "tolist") else [bool(is_success)] * env.num_envs + else: + successes = [False] * env.num_envs + + if recording_datasets is not None and raw_observation is not None: + prev_done = done.copy() + for env_idx in range(env.num_envs): + if prev_done[env_idx]: + continue + frame = _build_raw_frame( + raw_observation, + env_idx, + action_numpy[env_idx], + reward[env_idx], + successes[env_idx], + bool(terminated[env_idx] | truncated[env_idx]), + task_desc, + recording_datasets[env_idx].features, + ) + recording_datasets[env_idx].add_frame(frame) + if terminated[env_idx] or truncated[env_idx]: + recording_datasets[env_idx].save_episode() + raw_observation = deepcopy(observation) + + # Keep track of which environments are done so far. + # Mark the episode as done if we reach the maximum step limit. + # This ensures that the rollout always terminates cleanly at `max_steps`, + # and allows logging/saving (e.g., videos) to be triggered consistently. + done = terminated | truncated | done + if step + 1 == max_steps: + done = np.ones_like(done, dtype=bool) + + all_actions.append(torch.from_numpy(action_numpy)) + all_rewards.append(torch.from_numpy(reward)) + all_dones.append(torch.from_numpy(done)) + all_successes.append(torch.tensor(successes)) + + step += 1 + running_success_rate = ( + einops.reduce(torch.stack(all_successes, dim=1), "b n -> b", "any").numpy().mean() ) - else: - successes = [False] * env.num_envs - - # Keep track of which environments are done so far. - # Mark the episode as done if we reach the maximum step limit. - # This ensures that the rollout always terminates cleanly at `max_steps`, - # and allows logging/saving (e.g., videos) to be triggered consistently. - done = terminated | truncated | done - if step + 1 == max_steps: - done = np.ones_like(done, dtype=bool) - - all_actions.append(torch.from_numpy(action_numpy)) - all_rewards.append(torch.from_numpy(reward)) - all_dones.append(torch.from_numpy(done)) - all_successes.append(torch.tensor(successes)) - - step += 1 - running_success_rate = ( - einops.reduce(torch.stack(all_successes, dim=1), "b n -> b", "any").numpy().mean() - ) - progbar.set_postfix({"running_success_rate": f"{running_success_rate.item() * 100:.1f}%"}) - progbar.update() + progbar.set_postfix({"running_success_rate": f"{running_success_rate.item() * 100:.1f}%"}) + progbar.update() + finally: + if recording_datasets is not None: + for ds in recording_datasets: + ds.finalize() + if recording_repo_id is not None: + if ds.num_episodes > 0: + ds.push_to_hub(private=recording_private) + else: + logging.warning("No episodes recorded for %s — skipping push to hub.", ds.repo_id) # Track the final observation. if return_observations: @@ -273,6 +396,10 @@ def eval_policy( videos_dir: Path | None = None, return_episode_data: bool = False, start_seed: int | None = None, + recording_dir: Path | None = None, + env_features: dict | None = None, + recording_repo_id: str | None = None, + recording_private: bool = False, ) -> dict: """ Args: @@ -361,6 +488,10 @@ def eval_policy( seeds=list(seeds) if seeds else None, return_observations=return_episode_data, render_callback=render_frame if max_episodes_rendered > 0 else None, + recording_dir=recording_dir, + env_features=env_features, + recording_repo_id=recording_repo_id, + recording_private=recording_private, ) # Figure out where in each rollout sequence the first done condition was encountered (results after @@ -563,6 +694,10 @@ def eval_main(cfg: EvalPipelineConfig): # Create environment-specific preprocessor and postprocessor (e.g., for LIBERO environments) env_preprocessor, env_postprocessor = make_env_pre_post_processors(env_cfg=cfg.env, policy_cfg=cfg.policy) + recording_dir = Path(cfg.output_dir) / "recordings" if cfg.eval.recording else None + max_episodes_rendered = 0 if cfg.eval.recording else 10 + videos_dir = None if cfg.eval.recording else Path(cfg.output_dir) / "videos" + with torch.no_grad(), torch.autocast(device_type=device.type) if cfg.policy.use_amp else nullcontext(): info = eval_policy_all( envs=envs, @@ -572,10 +707,15 @@ def eval_main(cfg: EvalPipelineConfig): preprocessor=preprocessor, postprocessor=postprocessor, n_episodes=cfg.eval.n_episodes, - max_episodes_rendered=10, - videos_dir=Path(cfg.output_dir) / "videos", + max_episodes_rendered=max_episodes_rendered, + videos_dir=videos_dir, + return_episode_data=False, start_seed=cfg.seed, max_parallel_tasks=cfg.env.max_parallel_tasks, + recording_dir=recording_dir, + env_features=cfg.env.features if cfg.eval.recording else None, + recording_repo_id=cfg.eval.recording_repo_id, + recording_private=cfg.eval.recording_private, ) print("Overall Aggregated Metrics:") print(info["overall"]) @@ -618,6 +758,10 @@ def eval_one( videos_dir: Path | None, return_episode_data: bool, start_seed: int | None, + recording_dir: Path | None = None, + env_features: dict | None = None, + recording_repo_id: str | None = None, + recording_private: bool = False, ) -> TaskMetrics: """Evaluates one task_id of one suite using the provided vec env.""" @@ -635,6 +779,10 @@ def eval_one( videos_dir=task_videos_dir, return_episode_data=return_episode_data, start_seed=start_seed, + recording_dir=recording_dir, + env_features=env_features, + recording_repo_id=recording_repo_id, + recording_private=recording_private, ) per_episode = task_result["per_episode"] @@ -661,6 +809,10 @@ def run_one( videos_dir: Path | None, return_episode_data: bool, start_seed: int | None, + recording_dir: Path | None = None, + env_features: dict | None = None, + recording_repo_id: str | None = None, + recording_private: bool = False, ): """ Run eval_one for a single (task_group, task_id, env). @@ -672,7 +824,13 @@ def run_one( task_videos_dir = videos_dir / f"{task_group}_{task_id}" task_videos_dir.mkdir(parents=True, exist_ok=True) - # Call the existing eval_one (assumed to return TaskMetrics-like dict) + task_recording_dir = None + task_repo_id = None + if recording_dir is not None and env_features is not None: + task_recording_dir = recording_dir / f"{task_group}_{task_id}" + if recording_repo_id is not None: + task_repo_id = f"{recording_repo_id}_{task_group}_{task_id}" + metrics = eval_one( env, policy=policy, @@ -685,8 +843,12 @@ def run_one( videos_dir=task_videos_dir, return_episode_data=return_episode_data, start_seed=start_seed, + recording_dir=task_recording_dir, + env_features=env_features, + recording_repo_id=task_repo_id, + recording_private=recording_private, ) - # ensure we always provide video_paths key to simplify accumulation + if max_episodes_rendered > 0: metrics.setdefault("video_paths", []) return task_group, task_id, metrics @@ -702,6 +864,10 @@ def eval_policy_all( n_episodes: int, *, max_episodes_rendered: int = 0, + recording_dir: Path | None = None, + env_features: dict | None = None, + recording_repo_id: str | None = None, + recording_private: bool = False, videos_dir: Path | None = None, return_episode_data: bool = False, start_seed: int | None = None, @@ -761,6 +927,10 @@ def eval_policy_all( videos_dir=videos_dir, return_episode_data=return_episode_data, start_seed=start_seed, + recording_dir=recording_dir, + env_features=env_features, + recording_repo_id=recording_repo_id, + recording_private=recording_private, ) if max_parallel_tasks <= 1: diff --git a/src/lerobot/scripts/lerobot_record.py b/src/lerobot/scripts/lerobot_record.py index 0deb54b90..b759d86e0 100644 --- a/src/lerobot/scripts/lerobot_record.py +++ b/src/lerobot/scripts/lerobot_record.py @@ -79,9 +79,9 @@ lerobot-record \\ --dataset.single_task="Grab the cube" \\ --dataset.streaming_encoding=true \\ --dataset.encoder_threads=2 \\ - --dataset.camera_encoder.vcodec=h264 \\ - --dataset.camera_encoder.preset=fast \\ - --dataset.camera_encoder.extra_options={"tune": "film", "profile:v": "high", "bf": 2} \\ + --dataset.rgb_encoder.vcodec=h264 \\ + --dataset.rgb_encoder.preset=fast \\ + --dataset.rgb_encoder.extra_options={"tune": "film", "profile:v": "high", "bf": 2} \\ --display_data=true ``` """ @@ -96,11 +96,7 @@ from lerobot.cameras.opencv import OpenCVCameraConfig # noqa: F401 from lerobot.cameras.reachy2_camera import Reachy2CameraConfig # noqa: F401 from lerobot.cameras.realsense import RealSenseCameraConfig # noqa: F401 from lerobot.cameras.zmq import ZMQCameraConfig # noqa: F401 -from lerobot.common.control_utils import ( - init_keyboard_listener, - is_headless, - sanity_check_dataset_robot_compatibility, -) +from lerobot.common.control_utils import sanity_check_dataset_robot_compatibility from lerobot.configs import parser from lerobot.configs.dataset import DatasetRecordConfig from lerobot.datasets import ( @@ -155,6 +151,7 @@ from lerobot.teleoperators.keyboard import KeyboardTeleop from lerobot.utils.constants import ACTION, OBS_STR from lerobot.utils.feature_utils import build_dataset_frame, combine_feature_dicts from lerobot.utils.import_utils import register_third_party_plugins +from lerobot.utils.keyboard_input import init_keyboard_listener from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.utils import ( init_logging, @@ -403,7 +400,8 @@ def record( cfg.dataset.repo_id, root=cfg.dataset.root, batch_encoding_size=cfg.dataset.video_encoding_batch_size, - camera_encoder=cfg.dataset.camera_encoder, + rgb_encoder=cfg.dataset.rgb_encoder, + depth_encoder=cfg.dataset.depth_encoder, encoder_threads=cfg.dataset.encoder_threads, streaming_encoding=cfg.dataset.streaming_encoding, encoder_queue_maxsize=cfg.dataset.encoder_queue_maxsize, @@ -432,7 +430,8 @@ def record( image_writer_processes=cfg.dataset.num_image_writer_processes, image_writer_threads=cfg.dataset.num_image_writer_threads_per_camera * len(robot.cameras), batch_encoding_size=cfg.dataset.video_encoding_batch_size, - camera_encoder=cfg.dataset.camera_encoder, + rgb_encoder=cfg.dataset.rgb_encoder, + depth_encoder=cfg.dataset.depth_encoder, encoder_threads=cfg.dataset.encoder_threads, streaming_encoding=cfg.dataset.streaming_encoding, encoder_queue_maxsize=cfg.dataset.encoder_queue_maxsize, @@ -446,7 +445,7 @@ def record( if not cfg.dataset.streaming_encoding: logging.info( - "Streaming encoding is disabled. If you have capable hardware, consider enabling it for way faster episode saving. --dataset.streaming_encoding=true --dataset.encoder_threads=2 # --dataset.camera_encoder.vcodec=auto. More info in the documentation: https://huggingface.co/docs/lerobot/streaming_video_encoding" + "Streaming encoding is disabled. If you have capable hardware, consider enabling it for way faster episode saving. --dataset.streaming_encoding=true --dataset.encoder_threads=2 # --dataset.rgb_encoder.vcodec=auto. More info in the documentation: https://huggingface.co/docs/lerobot/streaming_video_encoding" ) with VideoEncodingManager(dataset): @@ -508,7 +507,7 @@ def record( if teleop and teleop.is_connected: teleop.disconnect() - if not is_headless() and listener: + if listener is not None: listener.stop() if cfg.dataset.push_to_hub: diff --git a/src/lerobot/scripts/lerobot_rollout.py b/src/lerobot/scripts/lerobot_rollout.py index 8515c4cc9..daee87bbe 100644 --- a/src/lerobot/scripts/lerobot_rollout.py +++ b/src/lerobot/scripts/lerobot_rollout.py @@ -142,9 +142,9 @@ Usage examples --robot.port=/dev/ttyACM0 \\ --task="pick up cube" --duration=60 \\ --display_data=true \\ - --dataset.camera_encoder.vcodec=h264 \\ - --dataset.camera_encoder.preset=fast \\ - --dataset.camera_encoder.extra_options={"tune": "film", "profile:v": "high", "bf": 2} + --dataset.rgb_encoder.vcodec=h264 \\ + --dataset.rgb_encoder.preset=fast \\ + --dataset.rgb_encoder.extra_options={"tune": "film", "profile:v": "high", "bf": 2} """ import logging diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index 70a5e9e9d..f2a152df9 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/src/lerobot/scripts/lerobot_train.py @@ -20,6 +20,7 @@ Requires: pip install 'lerobot[training]' (includes dataset + accelerate + wand import dataclasses import logging +import sys import time from contextlib import nullcontext from pprint import pformat @@ -34,19 +35,24 @@ from torch.optim import Optimizer from tqdm import tqdm from lerobot.common.train_utils import ( + gather_fsdp_state_dicts, get_step_checkpoint_dir, get_step_identifier, + load_fsdp_optimizer_state, load_training_batch_size, load_training_num_processes, load_training_state, + push_checkpoint_to_hub, save_checkpoint, update_last_checkpoint, ) from lerobot.common.wandb_utils import WandBLogger -from lerobot.configs import parser +from lerobot.configs import JobConfig, parser from lerobot.configs.train import TrainPipelineConfig -from lerobot.datasets import EpisodeAwareSampler, compute_sampler_state, make_dataset +from lerobot.datasets import EpisodeAwareSampler, compute_sampler_state +from lerobot.datasets.factory import make_train_eval_datasets from lerobot.envs import close_envs, make_env, make_env_pre_post_processors +from lerobot.jobs import submit_to_hf from lerobot.optim.factory import make_optimizer_and_scheduler from lerobot.policies import PreTrainedPolicy, make_policy, make_pre_post_processors from lerobot.rewards import make_reward_pre_post_processors @@ -185,10 +191,14 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): cfg: A `TrainPipelineConfig` object containing all training configurations. accelerator: Optional Accelerator instance. If None, one will be created automatically. """ + if cfg.job.is_remote: + return submit_to_hf(cfg) + from lerobot.utils.import_utils import require_package require_package("accelerate", extra="training") from accelerate import Accelerator + from accelerate.utils import DistributedDataParallelKwargs, DistributedType cfg.validate() @@ -197,8 +207,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): # We set step_scheduler_with_optimizer=False to prevent accelerate from adjusting the lr_scheduler steps based on the num_processes # We set find_unused_parameters=True to handle models with conditional computation if accelerator is None: - from accelerate.utils import DistributedDataParallelKwargs - ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True) # Accelerate auto-detects the device based on the available hardware and ignores the policy.device setting. # Force the device to be CPU when the active config's device is set to CPU (works for both policy and reward model training). @@ -244,19 +252,19 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): # LeRobotDataset skips its snapshot_download when try_load() succeeds, so no rank re-downloads. if is_main_process: logging.info("Creating dataset") - dataset = make_dataset(cfg) + dataset, eval_dataset = make_train_eval_datasets(cfg) accelerator.wait_for_everyone() # Other ranks read from the shared copy populated by the main process. if not is_main_process: - dataset = make_dataset(cfg) + dataset, eval_dataset = make_train_eval_datasets(cfg) # Create environment used for evaluating checkpoints during training on simulation data. # On real-world data, no need to create an environment as evaluations are done outside train.py, # using the eval.py instead, with gym_dora environment and dora-rs. eval_env = None - if cfg.eval_freq > 0 and cfg.env is not None and is_main_process: + if cfg.env_eval_freq > 0 and cfg.env is not None and is_main_process: logging.info("Creating env") eval_env = make_env(cfg.env, n_envs=cfg.eval.batch_size, use_async_envs=cfg.eval.use_async_envs) @@ -345,6 +353,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): preprocessor, postprocessor = make_pre_post_processors( policy_cfg=cfg.policy, pretrained_path=processor_pretrained_path, + pretrained_revision=getattr(cfg.policy, "pretrained_revision", None), **processor_kwargs, ) @@ -370,7 +379,12 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): step = 0 # number of policy updates (forward + backward + optim) if cfg.resume: - step, optimizer, lr_scheduler = load_training_state(cfg.checkpoint_path, optimizer, lr_scheduler) + # Under FSDP the optimizer state is sharded and must be loaded after `accelerator.prepare()` + # (see load_fsdp_optimizer_state below), so skip the optimizer here and load it then. + is_fsdp = accelerator.distributed_type == DistributedType.FSDP + step, optimizer, lr_scheduler = load_training_state( + cfg.checkpoint_path, optimizer, lr_scheduler, load_optimizer=not is_fsdp + ) num_learnable_params = sum(p.numel() for p in policy.parameters() if p.requires_grad) num_total_params = sum(p.numel() for p in policy.parameters()) @@ -406,6 +420,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): drop_n_last_frames=getattr(active_cfg, "drop_n_last_frames", 0), shuffle=True, seed=cfg.seed if cfg.seed is not None else 0, + absolute_to_relative_idx=dataset.absolute_to_relative_idx, ) if cfg.resume and step > 0: # The resume offset depends on the (num_processes, batch_size) that produced `step`, so @@ -455,11 +470,49 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): persistent_workers=cfg.persistent_workers and cfg.num_workers > 0, ) + # Build eval dataloader if a held-out split exists + eval_dataloader = None + if eval_dataset is not None: + eval_ds = eval_dataset + if cfg.max_eval_samples > 0 and hasattr(eval_dataset, "hf_dataset"): + task_arr = eval_dataset.hf_dataset.data.column("task_index").to_numpy() + unique_tasks = sorted(set(task_arr.tolist())) + per_task = max(1, cfg.max_eval_samples // len(unique_tasks)) + selected: list[int] = [] + for t in unique_tasks: + frames = (task_arr == t).nonzero()[0][:per_task] + selected.extend(frames.tolist()) + eval_ds = torch.utils.data.Subset(eval_dataset, selected) + + eval_collate_fn = lerobot_collate_fn if dataset.meta.has_language_columns else None + eval_dataloader = torch.utils.data.DataLoader( + eval_ds, + batch_size=cfg.batch_size, + shuffle=False, + num_workers=cfg.num_workers, + pin_memory=device.type == "cuda", + drop_last=False, + collate_fn=eval_collate_fn, + prefetch_factor=cfg.prefetch_factor if cfg.num_workers > 0 else None, + persistent_workers=cfg.persistent_workers and cfg.num_workers > 0, + ) + # Prepare everything with accelerator accelerator.wait_for_everyone() - policy, optimizer, dataloader, lr_scheduler = accelerator.prepare( - policy, optimizer, dataloader, lr_scheduler - ) + if eval_dataloader is not None: + policy, optimizer, dataloader, lr_scheduler, eval_dataloader = accelerator.prepare( + policy, optimizer, dataloader, lr_scheduler, eval_dataloader + ) + else: + policy, optimizer, dataloader, lr_scheduler = accelerator.prepare( + policy, optimizer, dataloader, lr_scheduler + ) + + # FSDP optimizer state is sharded across ranks, so it can only be loaded once the optimizer and + # model are FSDP-wrapped (i.e. after `prepare`). Collective: every rank must participate. + if cfg.resume and accelerator.distributed_type == DistributedType.FSDP: + load_fsdp_optimizer_state(policy, optimizer, cfg.checkpoint_path) + dl_iter = cycle(dataloader) policy.train() @@ -534,7 +587,8 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): train_tracker.step() is_log_step = cfg.log_freq > 0 and step % cfg.log_freq == 0 is_saving_step = step % cfg.save_freq == 0 or step == cfg.steps - is_eval_step = cfg.eval_freq > 0 and step % cfg.eval_freq == 0 + is_env_eval_step = cfg.env_eval_freq > 0 and step % cfg.env_eval_freq == 0 + is_eval_step = cfg.eval_steps > 0 and eval_dataloader is not None and step % cfg.eval_steps == 0 if is_log_step: # Collective reduce must run on every rank, before the main-process gate below. @@ -557,7 +611,38 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): wandb_logger.log_dict(wandb_log_dict, step) train_tracker.reset_averages() + if is_eval_step: + policy.eval() + eval_loss_sum = 0.0 + n_eval_batches = 0 + with torch.no_grad(), accelerator.autocast(): + for eval_batch in eval_dataloader: + for cam_key in dataset.meta.camera_keys: + if cam_key in eval_batch and eval_batch[cam_key].dtype == torch.uint8: + eval_batch[cam_key] = eval_batch[cam_key].to(dtype=torch.float32) / 255.0 + eval_batch = preprocessor(eval_batch) + loss, _ = policy.forward(eval_batch) + eval_loss_sum += loss.item() + n_eval_batches += 1 + eval_loss = eval_loss_sum / max(n_eval_batches, 1) + eval_loss = torch.tensor(eval_loss, device=device) + eval_loss = accelerator.reduce(eval_loss, reduction="mean").item() + policy.train() + + if is_main_process: + logging.info(f"step {step}: eval_loss={eval_loss:.4f}") + if wandb_logger: + wandb_logger.log_dict({"eval_loss": eval_loss}, step=step, mode="eval") + if cfg.save_checkpoint and is_saving_step: + # Under FSDP, gathering the full model + optimizer state dicts is a cross-rank collective, + # so all ranks must participate; rank 0 then writes the materialized dicts. For DDP / + # single-GPU the state dicts are saved the normal way inside save_checkpoint. + is_fsdp = accelerator.distributed_type == DistributedType.FSDP + if is_fsdp: + model_state_dict, optim_state_dict = gather_fsdp_state_dicts(policy, optimizer) + else: + model_state_dict, optim_state_dict = None, None if is_main_process: logging.info(f"Checkpoint policy after step {step}") checkpoint_dir = get_step_checkpoint_dir(cfg.output_dir, cfg.steps, step) @@ -572,14 +657,22 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): postprocessor=postprocessor, num_processes=accelerator.num_processes, batch_size=cfg.batch_size, + model_state_dict=model_state_dict, + optim_state_dict=optim_state_dict, ) update_last_checkpoint(checkpoint_dir) + if cfg.save_checkpoint_to_hub: + push_checkpoint_to_hub( + checkpoint_dir, + cfg.policy.repo_id, + private=cfg.policy.private, + ) if wandb_logger: wandb_logger.log_policy(checkpoint_dir) accelerator.wait_for_everyone() - if cfg.env and is_eval_step: + if cfg.env and is_env_eval_step: if is_main_process: step_id = get_step_identifier(step, cfg.steps) logging.info(f"Eval policy at step {step}") @@ -634,6 +727,8 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): if eval_env: close_envs(eval_env) + is_fsdp = accelerator.distributed_type == DistributedType.FSDP + model_state_dict = accelerator.get_state_dict(policy) if is_fsdp else None if is_main_process: logging.info("End of training") @@ -643,7 +738,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): if not cfg.is_reward_model_training and cfg.policy.use_peft: unwrapped_model.push_model_to_hub(cfg, peft_model=unwrapped_model) else: - unwrapped_model.push_model_to_hub(cfg) + unwrapped_model.push_model_to_hub(cfg, state_dict=model_state_dict) preprocessor.push_to_hub(active_cfg.repo_id) postprocessor.push_to_hub(active_cfg.repo_id) @@ -652,8 +747,25 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): accelerator.end_training() +def _remote_target_in_argv() -> bool: + """True when the CLI requests a remote HF Jobs run (--job.target=).""" + target = None + args = sys.argv[1:] + for i, tok in enumerate(args): + if tok == "--job.target" and i + 1 < len(args): + target = args[i + 1] + elif tok.startswith("--job.target="): + target = tok.split("=", 1)[1] + return JobConfig.is_remote_target(target) + + def main(): register_third_party_plugins() + if _remote_target_in_argv(): + # The policy device is resolved on the remote pod, not here, so silence the + # client-side "Device '...' is not available" warning PreTrainedConfig emits + # while parsing the config (it fires before train() can dispatch remotely). + logging.getLogger("lerobot.configs.policies").setLevel(logging.ERROR) train() diff --git a/src/lerobot/teleoperators/gamepad/gamepad_utils.py b/src/lerobot/teleoperators/gamepad/gamepad_utils.py index c1531ca84..22dbb7cca 100644 --- a/src/lerobot/teleoperators/gamepad/gamepad_utils.py +++ b/src/lerobot/teleoperators/gamepad/gamepad_utils.py @@ -18,6 +18,7 @@ import logging from typing import TYPE_CHECKING from lerobot.utils.import_utils import _hidapi_available, _pygame_available, require_package +from lerobot.utils.keyboard_input import pynput_can_capture from ..utils import TeleopEvents @@ -123,6 +124,15 @@ class KeyboardController(InputController): def start(self): """Start the keyboard listener.""" + if not pynput_can_capture(): + logging.warning( + "Keyboard control is unavailable in this environment. pynput cannot capture keys " + "on Wayland or headless machines, or on macOS without Accessibility / Input " + "Monitoring permission. Keyboard motion will be inactive." + ) + self.running = False + return + from pynput import keyboard def on_press(key): diff --git a/src/lerobot/teleoperators/keyboard/teleop_keyboard.py b/src/lerobot/teleoperators/keyboard/teleop_keyboard.py index 801789bcb..872cc7a26 100644 --- a/src/lerobot/teleoperators/keyboard/teleop_keyboard.py +++ b/src/lerobot/teleoperators/keyboard/teleop_keyboard.py @@ -15,8 +15,6 @@ # limitations under the License. import logging -import os -import sys import time from queue import Queue from typing import Any @@ -24,6 +22,7 @@ from typing import Any from lerobot.types import RobotAction from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected from lerobot.utils.import_utils import _pynput_available, require_package +from lerobot.utils.keyboard_input import pynput_can_capture from ..teleoperator import Teleoperator from ..utils import TeleopEvents @@ -37,14 +36,10 @@ PYNPUT_AVAILABLE = _pynput_available keyboard = None if PYNPUT_AVAILABLE: try: - if ("DISPLAY" not in os.environ) and ("linux" in sys.platform): - logging.info("No DISPLAY set. Skipping pynput import.") - PYNPUT_AVAILABLE = False - else: - from pynput import keyboard + from pynput import keyboard except Exception as e: PYNPUT_AVAILABLE = False - logging.info(f"Could not import pynput: {e}") + logging.info("Could not import pynput keyboard backend: %s", e) class KeyboardTeleop(Teleoperator): @@ -88,7 +83,7 @@ class KeyboardTeleop(Teleoperator): @check_if_already_connected def connect(self) -> None: - if PYNPUT_AVAILABLE: + if PYNPUT_AVAILABLE and pynput_can_capture(): logging.info("pynput is available - enabling local keyboard listener.") self.listener = keyboard.Listener( on_press=self._on_press, @@ -96,7 +91,13 @@ class KeyboardTeleop(Teleoperator): ) self.listener.start() else: - logging.info("pynput not available - skipping local keyboard listener.") + logging.warning( + "Keyboard teleoperation is unavailable in this environment. pynput can only " + "capture key events on an X11 session (Linux), a Windows desktop, or macOS with " + "Accessibility / Input Monitoring granted - not on Wayland or headless machines. " + "This keyboard teleoperator will produce no actions; use an X11 session, a " + "gamepad, or a leader-arm teleoperator instead." + ) self.listener = None def calibrate(self) -> None: diff --git a/src/lerobot/utils/feature_utils.py b/src/lerobot/utils/feature_utils.py index 2a4886234..38516d6ab 100644 --- a/src/lerobot/utils/feature_utils.py +++ b/src/lerobot/utils/feature_utils.py @@ -51,7 +51,9 @@ def hw_to_dataset_features( This function takes a dictionary describing hardware outputs (like joint states or camera image shapes) and formats it into the standard LeRobot feature - specification. + specification. Single-channel cameras (shape ``(H, W, 1)``) are flagged as depth + maps via ``info["is_depth_map"] = True``; three-channel cameras ``(H, W, 3)`` are + treated as RGB. Args: hw_features (dict): Dictionary mapping feature names to their type (float for @@ -61,7 +63,7 @@ def hw_to_dataset_features( use_video (bool): If True, image features are marked as "video", otherwise "image". Returns: - dict: A LeRobot features dictionary. + dict: A LeRobot features dictionary. Depth cameras carry ``info["is_depth_map"] = True``. """ features = {} joint_fts = { @@ -69,6 +71,7 @@ def hw_to_dataset_features( for key, ftype in hw_features.items() if ftype is float or (isinstance(ftype, PolicyFeature) and ftype.type != FeatureType.VISUAL) } + # TODO(CarolinePascal): we should not rely on the shape to determine if a feature is a camera ! cam_fts = {key: shape for key, shape in hw_features.items() if isinstance(shape, tuple)} if joint_fts and prefix == ACTION: @@ -86,11 +89,19 @@ def hw_to_dataset_features( } for key, shape in cam_fts.items(): - features[f"{prefix}.images.{key}"] = { - "dtype": "video" if use_video else "image", - "shape": shape, - "names": ["height", "width", "channels"], - } + dtype = "video" if use_video else "image" + if len(shape) == 3 and shape[2] in (1, 3): + features[f"{prefix}.images.{key}"] = { + "dtype": dtype, + "shape": shape, + "names": ["height", "width", "channels"], + "info": {"is_depth_map": shape[2] == 1}, + } + else: + raise ValueError( + f"Camera feature '{key}' has shape {shape}. " + f"Expected a 3-tuple (H, W, C), e.g. (480, 640, 3) for RGB or (480, 640, 1) for depth." + ) _validate_feature_names(features) return features @@ -149,11 +160,11 @@ def dataset_to_policy_features(features: dict[str, dict]) -> dict[str, PolicyFea type = FeatureType.VISUAL if len(shape) != 3: raise ValueError(f"Number of dimensions of {key} != 3 (shape={shape})") - - names = ft["names"] - # Backward compatibility for "channel" which is an error introduced in LeRobotDataset v2.0 for ported datasets. - if names[2] in ["channel", "channels"]: # (h, w, c) -> (c, h, w) - shape = (shape[2], shape[0], shape[1]) + else: + names = ft["names"] + # Backward compatibility for "channel" which is an error introduced in LeRobotDataset v2.0 for ported datasets. + if names[2] in ["channel", "channels"]: # (h, w, c) -> (c, h, w) + shape = (shape[2], shape[0], shape[1]) elif key == OBS_ENV_STATE: type = FeatureType.ENV elif key.startswith(OBS_STR): diff --git a/src/lerobot/utils/hub.py b/src/lerobot/utils/hub.py index 566701b31..57eb819f8 100644 --- a/src/lerobot/utils/hub.py +++ b/src/lerobot/utils/hub.py @@ -20,9 +20,33 @@ from typing import Any, TypeVar from huggingface_hub import HfApi from huggingface_hub.utils import validate_hf_hub_args +from .constants import CHECKPOINTS_DIR + T = TypeVar("T", bound="HubMixin") +def find_latest_hub_checkpoint( + repo_id: str, + *, + token: str | bool | None = None, + revision: str | None = None, +) -> str | None: + """Repo-relative path of the most recent checkpoint in a training repo. + + Training runs push checkpoints to ``checkpoints//`` (see + ``push_checkpoint_to_hub``). This lists those step dirs and returns + ``checkpoints/``, or ``None`` if the repo has no checkpoints. + """ + files = HfApi().list_repo_files(repo_id=repo_id, repo_type="model", revision=revision, token=token) + prefix = f"{CHECKPOINTS_DIR}/" + steps = { + name for f in files if f.startswith(prefix) and (name := f[len(prefix) :].split("/", 1)[0]).isdigit() + } + if not steps: + return None + return f"{CHECKPOINTS_DIR}/{max(steps, key=int)}" + + class HubMixin: """ A Mixin containing the functionality to push an object to the hub. diff --git a/src/lerobot/utils/import_utils.py b/src/lerobot/utils/import_utils.py index 5dbce2c5b..b0d894c04 100644 --- a/src/lerobot/utils/import_utils.py +++ b/src/lerobot/utils/import_utils.py @@ -216,9 +216,15 @@ def register_third_party_plugins() -> None: This function uses `importlib.metadata` to find packages installed in the environment (including editable installs) starting with 'lerobot_robot_', 'lerobot_camera_', - 'lerobot_teleoperator_', or 'lerobot_policy_' and imports them. + 'lerobot_teleoperator_', 'lerobot_policy_', or 'lerobot_env_' and imports them. """ - prefixes = ("lerobot_robot_", "lerobot_camera_", "lerobot_teleoperator_", "lerobot_policy_") + prefixes = ( + "lerobot_robot_", + "lerobot_camera_", + "lerobot_teleoperator_", + "lerobot_policy_", + "lerobot_env_", + ) imported: list[str] = [] failed: list[str] = [] diff --git a/src/lerobot/utils/keyboard_input.py b/src/lerobot/utils/keyboard_input.py new file mode 100644 index 000000000..00c0f53ec --- /dev/null +++ b/src/lerobot/utils/keyboard_input.py @@ -0,0 +1,440 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Display-independent keyboard input for interactive controls. + +This module centralizes everything related to *discrete* keyboard controls +(end-episode-early, re-record, stop, and the rollout strategies' custom keys): + +* environment detection — :func:`is_headless`, :func:`is_wayland`, + :func:`pynput_can_capture` (the single predicate every call-site should use to + decide whether ``pynput`` can actually capture keys here); +* a shared key mapping — :func:`apply_recording_control`; and +* two interchangeable backends behind one ``(listener, events)`` contract: + the ``pynput`` global listener (X11 / trusted-macOS / Windows) and a + standard-library :class:`TerminalKeyListener` that reads the controlling TTY + (Wayland / headless-SSH-with-TTY / macOS without Accessibility permission). + +NOTE: *continuous* key-state teleoperation ("hold a key to keep moving") is +deliberately NOT served here. A terminal in cbreak mode delivers only key-down +bytes — there is no key-release event — so the held-key model cannot be +reproduced. Those teleoperators stay on ``pynput`` and use +:func:`pynput_can_capture` to warn instead of silently doing nothing. +""" + +from __future__ import annotations + +import atexit +import contextlib +import logging +import os +import platform +import select +import sys +import threading +import time +from collections.abc import Callable +from functools import cache +from typing import TYPE_CHECKING + +from .import_utils import _pynput_available + +logger = logging.getLogger(__name__) + +# POSIX-only terminal modules (absent on Windows, where the pynput backend is used). +if TYPE_CHECKING: + import termios + import tty + + _TERMIOS_AVAILABLE = True +else: + try: + import termios + import tty + + _TERMIOS_AVAILABLE = True + except ImportError: # POSIX-only modules; unavailable on Windows + termios = tty = None + _TERMIOS_AVAILABLE = False + +keyboard = None +if _pynput_available: + try: + from pynput import keyboard + except Exception as e: # e.g. no reachable X display on a headless Linux box + logger.info("Could not import pynput keyboard backend: %s", e) + + +@cache +def is_headless() -> bool: + """Return ``True`` when no display server is available. + + * Linux: headless when neither ``DISPLAY`` (X11) nor ``WAYLAND_DISPLAY`` is set. + * macOS / Windows: a display is always assumed to be present. A genuinely GUI-less + Mac/Windows CI host would be misclassified but it doesn't matter, because the + sys.stdin.isatty() gate returns None there regardless. + """ + if platform.system() == "Linux": + return not (os.environ.get("DISPLAY") or os.environ.get("WAYLAND_DISPLAY")) + return False + + +@cache +def is_wayland() -> bool: + """Return ``True`` when running under a Wayland session. + + ``pynput`` relies on an X11 backend. Under Wayland it still imports (XWayland + is usually present and ``$DISPLAY`` is set) but cannot capture *global* + hotkeys, so the documented arrow/Esc shortcuts silently do nothing. This case + is invisible to :func:`is_headless`, hence the dedicated check. + """ + return os.environ.get("XDG_SESSION_TYPE", "").lower() == "wayland" or bool( + os.environ.get("WAYLAND_DISPLAY") + ) + + +@cache +def pynput_can_capture() -> bool: + """Return ``True`` when a ``pynput`` global listener can actually capture keys. + + This is the single predicate every keyboard call-site should use to choose + between the ``pynput`` backend and a fallback. It is intentionally + conservative: + + * Linux: only a real X11 session (a display is present *and* it is not Wayland). + * macOS: ``True`` here — Accessibility / Input-Monitoring permission + (``IS_TRUSTED``) can only be confirmed at runtime *after* starting a + listener, so :func:`init_keyboard_listener` refines this with + :func:`pynput_listener_is_trusted`. + * Windows: ``True`` (the low-level global hook needs no special permission). + + Always ``False`` when ``pynput`` is not installed. + """ + if not _pynput_available: + return False + if platform.system() == "Linux": + return not is_headless() and not is_wayland() + return True + + +def pynput_listener_is_trusted(listener, timeout_s: float = 1.0) -> bool: + """Best-effort check that a freshly started ``pynput`` listener can capture. + + On macOS, ``pynput`` sets ``listener.IS_TRUSTED`` on its *listener thread* + once the Quartz event tap is created; the class default is ``False``. We + therefore wait for the thread to either flip it ``True`` (trusted) or for a + short timeout to elapse (untrusted — it stays ``False`` forever). This biases + toward the common trusted case (returns as soon as the flag flips) and only + pays the full ``timeout_s`` on an already-broken untrusted machine. + + On non-macOS backends the attribute is absent and capture is assumed to work. + """ + if platform.system() != "Darwin": + return True + deadline = time.perf_counter() + timeout_s + while time.perf_counter() < deadline: + if getattr(listener, "IS_TRUSTED", False): + return True + time.sleep(0.02) + return bool(getattr(listener, "IS_TRUSTED", False)) + + +def apply_recording_control(control: str, events: dict) -> None: + """Apply a recording control-flow key press to the shared ``events`` dict. + + Centralizes the mapping so the ``pynput`` and terminal backends behave + identically. ``control`` is one of ``"right"`` (end the loop early), ``"left"`` + (re-record the last episode), or ``"esc"`` (stop recording). + """ + if control == "right": + print("Right arrow key pressed. Exiting loop...") + events["exit_early"] = True + elif control == "left": + print("Left arrow key pressed. Exiting loop and rerecord the last episode...") + events["rerecord_episode"] = True + events["exit_early"] = True + elif control == "esc": + print("Escape key pressed. Stopping data recording...") + events["stop_recording"] = True + events["exit_early"] = True + + +# Terminal arrow keys arrive as a 3-byte escape sequence whose *final* byte identifies +# the direction. Two encodings exist depending on the terminal's cursor-key mode — CSI +# ("ESC [ X") and SS3 ("ESC O X", common over SSH/tmux) — but both share the same final +# byte, so this single table decodes either. Looked up by TerminalKeyListener._parse; +# an unknown final byte yields None (sequence ignored). +_ARROW_FINAL_BYTES = {"A": "up", "B": "down", "C": "right", "D": "left"} + + +class TerminalKeyListener: + """Display-independent keyboard listener that reads keys from the controlling TTY. + + Used as the Wayland / headless / macOS-untrusted equivalent of the ``pynput`` + listener for *discrete* controls. It puts the terminal into cbreak mode with + echo disabled and reads bytes on a daemon thread, decoding them into logical + key names that are passed to ``on_key``: + + * arrow keys (``ESC [ C`` / ``ESC O C`` …) -> ``"right"`` / ``"left"`` / ``"up"`` / ``"down"`` + * a bare ``ESC`` -> ``"esc"`` + * Enter / Tab / Space / Backspace -> ``"enter"`` / ``"tab"`` / ``"space"`` / ``"backspace"`` + * any other printable byte -> that character (e.g. ``"n"``, ``"s"``) + + Only key-down events are produced (terminals have no key-release), so this is + suitable for discrete commands but NOT for continuous "hold-to-move" teleop. + + The terminal is restored on :meth:`stop` and also via an ``atexit`` hook, so a + crash or Ctrl-C never leaves the shell in a no-echo cbreak state. POSIX-only + (``termios`` / ``tty`` / ``select``); those modules are imported lazily so this + file stays importable on Windows (where ``pynput`` is used instead). + """ + + def __init__(self, on_key: Callable[[str], None]): + self._on_key = on_key + self._running = False + self._thread: threading.Thread | None = None + self._fd: int | None = None + self._old_attrs = None + + def _read_char(self, timeout: float) -> str | None: + """Return one character from stdin within ``timeout`` seconds, or ``None``.""" + if self._fd is None: + return None + ready, _, _ = select.select([self._fd], [], [], timeout) + if not ready: + return None + try: + data = os.read(self._fd, 1) + except OSError: + return None + if not data: + return None + return data.decode(errors="ignore") + + def _parse(self, ch: str) -> str | None: + """Decode one (possibly multi-byte) key starting at ``ch`` into a key name.""" + if ch == "\x1b": + # Possible CSI / SS3 escape sequence (arrow keys) or a bare ESC. Use + # short follow-up reads so a lone ESC is not mistaken for a sequence. + ch2 = self._read_char(timeout=0.02) + if ch2 is None: + return "esc" + if ch2 in ("[", "O"): + ch3 = self._read_char(timeout=0.02) + return _ARROW_FINAL_BYTES.get(ch3 or "") + # Some other escape sequence (e.g. Alt+key); ignore it. + return None + if ch in ("\r", "\n"): + return "enter" + if ch == "\t": + return "tab" + if ch == " ": + return "space" + if ch in ("\x7f", "\x08"): + return "backspace" + if ch.isprintable(): + return ch + return None + + def _run(self) -> None: + while self._running: + ch = self._read_char(timeout=0.05) + if ch is None: + continue + name = self._parse(ch) + if name is None: + continue + try: + self._on_key(name) + except Exception as e: # never let a handler error kill the reader thread + logger.debug("Terminal key handler error: %s", e) + + def start(self) -> None: + """Switch the terminal to cbreak mode (echo off) and read keys on a daemon thread. + + No-op when stdin is not a TTY (piped/redirected input) or on platforms + without ``termios`` (e.g. Windows), so non-interactive runs are unaffected. + """ + if not sys.stdin.isatty(): + return + if not _TERMIOS_AVAILABLE: # POSIX-only modules (e.g. unavailable on Windows) + logger.warning("Terminal keyboard input is not supported on this platform.") + return + + self._fd = sys.stdin.fileno() + self._old_attrs = termios.tcgetattr(self._fd) + tty.setcbreak(self._fd) + # Explicitly disable ECHO so arrow-key escape sequences (e.g. ^[[C) are not + # echoed as garbage into the recording terminal. (Independent of the + # version-specific behavior of tty.setcbreak.) + new_attrs = termios.tcgetattr(self._fd) + new_attrs[3] &= ~termios.ECHO # index 3 == lflags + termios.tcsetattr(self._fd, termios.TCSADRAIN, new_attrs) + # Safety net: restore the terminal even if stop() is never reached (crash). + atexit.register(self.stop) + + self._running = True + self._thread = threading.Thread(target=self._run, daemon=True) + self._thread.start() + + def stop(self) -> None: + """Stop the reader thread and restore the original terminal attributes. + + Idempotent: safe to call multiple times (e.g. explicitly and via atexit). + """ + self._running = False + thread = self._thread + if thread is not None: + thread.join(timeout=0.5) + self._thread = None + if self._fd is not None and self._old_attrs is not None and _TERMIOS_AVAILABLE: + try: + termios.tcsetattr(self._fd, termios.TCSADRAIN, self._old_attrs) + finally: + self._old_attrs = None + with contextlib.suppress(Exception): + atexit.unregister(self.stop) + + +# Map pynput key objects to the same canonical names TerminalKeyListener emits, so a +# single dispatch works across both backends. Empty when pynput is unavailable. +if keyboard is not None: + _PYNPUT_KEY_NAMES = { + keyboard.Key.right: "right", + keyboard.Key.left: "left", + keyboard.Key.up: "up", + keyboard.Key.down: "down", + keyboard.Key.esc: "esc", + keyboard.Key.enter: "enter", + keyboard.Key.tab: "tab", + keyboard.Key.space: "space", + keyboard.Key.backspace: "backspace", + } +else: + _PYNPUT_KEY_NAMES = {} + + +def _resolve_pynput_key(key) -> str | None: + """Resolve a pynput key event to the canonical name TerminalKeyListener also emits. + + Special keys map through :data:`_PYNPUT_KEY_NAMES`; character keys fall back to their + ``.char`` (e.g. ``"n"``). Returns ``None`` for keys with no mapping and no character. + """ + name = _PYNPUT_KEY_NAMES.get(key) + if name is not None: + return name + # ``or None`` keeps the historical truthy-char semantics: an empty/None char is "no key". + return getattr(key, "char", None) or None + + +def create_key_listener(dispatch: Callable[[str], None], *, controls_help: str = ""): + """Start a keyboard listener that routes resolved key names to ``dispatch``. + + Shared backend selection used by recording and the rollout strategies: + + * the ``pynput`` global listener on X11 / trusted-macOS / Windows (on macOS the + listener's ``IS_TRUSTED`` flag is checked after start, and an untrusted listener is + stopped so the terminal backend is used instead); + * the stdlib :class:`TerminalKeyListener` on Wayland / headless sessions with a TTY; + * ``None`` when no backend is usable (non-interactive / piped runs). + + Both backends pass ``dispatch`` the same canonical key names ("right" / "left" / "up" / + "down" / "esc" / "enter" / "tab" / "space" / "backspace", or a character), so one + ``dispatch`` works regardless of backend. ``controls_help`` is an optional hint + appended to the log messages. + + Returns the listener (exposing ``.stop()``) or ``None``. + """ + suffix = f" ({controls_help})" if controls_help else "" + + if pynput_can_capture() and keyboard is not None: + + def on_press(key): + with contextlib.suppress(Exception): + name = _resolve_pynput_key(key) + if name is not None: + dispatch(name) + + listener = keyboard.Listener(on_press=on_press) + listener.start() + if pynput_listener_is_trusted(listener): + logger.info("Keyboard listener started%s.", suffix) + return listener + # macOS without Accessibility / Input-Monitoring permission: the listener never + # fires. Stop it and fall through to the terminal backend. + logger.warning( + "pynput keyboard listener is not trusted (missing macOS Accessibility / " + "Input Monitoring permission); falling back to terminal keyboard input." + ) + listener.stop() + + if sys.stdin.isatty(): + listener = TerminalKeyListener(dispatch) + listener.start() + logger.info("Using terminal keyboard input — keep this terminal focused%s.", suffix) + return listener + + logger.warning( + "Keyboard controls unavailable: no usable display (Wayland/headless) and stdin is " + "not an interactive terminal%s.", + suffix, + ) + return None + + +def init_keyboard_listener(): + """Initialize a non-blocking keyboard listener for interactive recording controls. + + Backend selection: + + * ``pynput`` global listener when :func:`pynput_can_capture` is true (real + X11, macOS, Windows). On macOS the listener's ``IS_TRUSTED`` flag is checked + after start; if the process lacks Accessibility / Input-Monitoring + permission, the listener is stopped and the terminal backend is used. + * a :class:`TerminalKeyListener` reading the controlling TTY when ``pynput`` + cannot capture (Wayland / headless-SSH / macOS-untrusted) *and* stdin is a TTY. + * otherwise no listener (non-interactive / piped runs) — recording relies on + the episode/reset timers (or Ctrl+C). + + Both backends accept the same controls: Right/Left/Esc, plus the single-byte letter + equivalents ``n`` (next), ``r`` (re-record) and ``q`` (quit). The letters are the most + reliable choice over high-latency SSH/VNC links, where arrow-key escape sequences can + be split, delayed, or intercepted by the terminal. + + Returns: + A tuple ``(listener, events)`` where ``listener`` exposes ``.stop()`` or is + ``None``, and ``events`` is the dict of flags (``exit_early``, + ``rerecord_episode``, ``stop_recording``) set by key presses. + """ + events = { + "exit_early": False, + "rerecord_episode": False, + "stop_recording": False, + } + + # Accept the single-byte letter equivalents n/r/q alongside the arrow/Esc keys: the + # letters are immune to the escape-sequence split/delay/interception that affects arrows + # over laggy SSH/VNC links. Case-insensitive so Shift+letter still works. + def on_key(name: str) -> None: + key = name.lower() + if key in ("right", "n"): + apply_recording_control("right", events) + elif key in ("left", "r"): + apply_recording_control("left", events) + elif key in ("esc", "q"): + apply_recording_control("esc", events) + # other keys (incl. up/down) are intentionally ignored + + listener = create_key_listener(on_key, controls_help="Right/Left/Esc, or n=next, r=re-record, q=quit") + return listener, events diff --git a/src/lerobot/utils/visualization_utils.py b/src/lerobot/utils/visualization_utils.py index d9d5bf6b5..a0f07f0c7 100644 --- a/src/lerobot/utils/visualization_utils.py +++ b/src/lerobot/utils/visualization_utils.py @@ -38,6 +38,8 @@ def init_rerun( require_package("rerun-sdk", extra="viz", import_name="rerun") import rerun as rr + log_rerun_data.blueprint = None # Reset blueprint cache for new session + batch_size = os.getenv("RERUN_FLUSH_NUM_BYTES", "8000") os.environ["RERUN_FLUSH_NUM_BYTES"] = batch_size rr.init(session_name) @@ -63,6 +65,41 @@ def _is_scalar(x): ) +def _build_blueprint(observation_paths: set[str], action_paths: set[str], image_paths: set[str]): + """Build a Rerun blueprint laying out camera images, observation and action scalars in separate views. + + Camera images, observation and action scalars are arranged in a grid. + """ + + # Safe + zero-overhead: `log_rerun_data` already ran the `require_package` guard and imported rerun. + import rerun.blueprint as rrb + + views = [rrb.Spatial2DView(origin=path, name=path) for path in sorted(image_paths)] + + if observation_paths: + views.append(rrb.TimeSeriesView(name="observation", contents=sorted(observation_paths))) + if action_paths: + views.append(rrb.TimeSeriesView(name="action", contents=sorted(action_paths))) + + return rrb.Blueprint(rrb.Grid(*views)) + + +def _ensure_blueprint(observation_paths: set[str], action_paths: set[str], image_paths: set[str]) -> None: + """Build and send the blueprint once, from the first observation and action data.""" + if getattr(log_rerun_data, "blueprint", None) is not None: + return + + if not (observation_paths or action_paths or image_paths): + return + + # Safe + zero-overhead: `log_rerun_data` already ran the `require_package` guard and imported rerun. + import rerun as rr + + blueprint = _build_blueprint(observation_paths, action_paths, image_paths) + log_rerun_data.blueprint = blueprint + rr.send_blueprint(blueprint) + + def log_rerun_data( observation: RobotObservation | None = None, action: RobotAction | None = None, @@ -76,11 +113,15 @@ def log_rerun_data( - Scalars values (floats, ints) are logged as `rr.Scalars`. - 3D NumPy arrays that resemble images (e.g., with 1, 3, or 4 channels first) are transposed from CHW to HWC format, (optionally) compressed to JPEG and logged as `rr.Image` or `rr.EncodedImage`. - - 1D NumPy arrays are logged as a series of individual scalars, with each element indexed. - - Other multi-dimensional arrays are flattened and logged as individual scalars. + - 1D NumPy arrays are logged as a single `rr.Scalars` batch under one entity path, so that every + dimension shares the same view instead of being split across one view per element. + - Multi-dimensional **action** arrays are flattened and logged as a single `rr.Scalars` batch. Keys are automatically namespaced with "observation." or "action." if not already present. + On the first call, a blueprint is built and sent so observation and action scalars get separate + time-series views and each image gets its own spatial view. + Args: observation: An optional dictionary containing observation data to log. action: An optional dictionary containing action data to log. @@ -90,6 +131,10 @@ def log_rerun_data( require_package("rerun-sdk", extra="viz", import_name="rerun") import rerun as rr + observation_paths: set[str] = set() + action_paths: set[str] = set() + image_paths: set[str] = set() + if observation: for k, v in observation.items(): if v is None: @@ -98,17 +143,22 @@ def log_rerun_data( if _is_scalar(v): rr.log(key, rr.Scalars(float(v))) + observation_paths.add(key) elif isinstance(v, np.ndarray): arr = v # Convert CHW -> HWC when needed if arr.ndim == 3 and arr.shape[0] in (1, 3, 4) and arr.shape[-1] not in (1, 3, 4): arr = np.transpose(arr, (1, 2, 0)) if arr.ndim == 1: - for i, vi in enumerate(arr): - rr.log(f"{key}_{i}", rr.Scalars(float(vi))) + rr.log(key, rr.Scalars(arr.astype(float))) + observation_paths.add(key) else: - img_entity = rr.Image(arr).compress() if compress_images else rr.Image(arr) + if arr.shape[-1] == 1: + img_entity = rr.DepthImage(arr, colormap=rr.components.Colormap.Viridis) + else: + img_entity = rr.Image(arr).compress() if compress_images else rr.Image(arr) rr.log(key, entity=img_entity, static=True) + image_paths.add(key) if action: for k, v in action.items(): @@ -118,12 +168,10 @@ def log_rerun_data( if _is_scalar(v): rr.log(key, rr.Scalars(float(v))) + action_paths.add(key) elif isinstance(v, np.ndarray): - if v.ndim == 1: - for i, vi in enumerate(v): - rr.log(f"{key}_{i}", rr.Scalars(float(vi))) - else: - # Fall back to flattening higher-dimensional arrays - flat = v.flatten() - for i, vi in enumerate(flat): - rr.log(f"{key}_{i}", rr.Scalars(float(vi))) + # Flatten any (incl. higher-dimensional) array into a single batched Scalars + rr.log(key, rr.Scalars(v.reshape(-1).astype(float))) + action_paths.add(key) + + _ensure_blueprint(observation_paths, action_paths, image_paths) diff --git a/tests/annotations/test_frames.py b/tests/annotations/test_frames.py index 5c9c58f7b..1a626533f 100644 --- a/tests/annotations/test_frames.py +++ b/tests/annotations/test_frames.py @@ -47,6 +47,7 @@ class _FakeMeta: def __init__(self, video_keys: list[str], image_keys: list[str], video_path: Path | None = None) -> None: self.video_keys = video_keys self.camera_keys = [*video_keys, *image_keys] + self.depth_keys = [] self._video_path = video_path self.episodes = {0: {f"videos/{key}/from_timestamp": 0.0 for key in video_keys}} @@ -208,14 +209,14 @@ def test_episode_clip_path_trims_via_reencode_video(tmp_path: Path, monkeypatch) def fake_reencode( input_video_path, output_video_path, - camera_encoder=None, + video_encoder=None, overwrite=False, start_time_s=None, end_time_s=None, ): captured.update( src=Path(input_video_path), - encoder=camera_encoder, + encoder=video_encoder, start_time_s=start_time_s, end_time_s=end_time_s, ) diff --git a/tests/configs/test_resume_from_hub.py b/tests/configs/test_resume_from_hub.py new file mode 100644 index 000000000..2cc9fd7ae --- /dev/null +++ b/tests/configs/test_resume_from_hub.py @@ -0,0 +1,68 @@ +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest + +import lerobot.configs.train as tc +from lerobot.configs.train import TrainPipelineConfig + + +class _FakeHTTPError(tc.HfHubHTTPError): + """HfHubHTTPError that can be raised without a real HTTP response object.""" + + def __init__(self): + pass + + +def test_from_pretrained_falls_back_to_latest_checkpoint_config(tmp_path, monkeypatch): + """A Hub repo with no root train_config.json (an interrupted run that only pushed + checkpoints/) resolves via the latest checkpoint's config.""" + # A real train_config.json written by save_pretrained, to be returned by the fallback. + parsed = tc.draccus.parse(TrainPipelineConfig, args=["--dataset.repo_id", "u/d"]) + cfg_file = tmp_path / "train_config.json" + parsed._save_pretrained(tmp_path) + assert cfg_file.is_file() + + calls = [] + + def fake_hf_hub_download(filename=None, **kwargs): + calls.append(filename) + if filename == "train_config.json": + raise _FakeHTTPError() # no root config + if filename == "checkpoints/000010/pretrained_model/train_config.json": + return str(cfg_file) + raise AssertionError(f"unexpected filename {filename}") + + monkeypatch.setattr(tc, "hf_hub_download", fake_hf_hub_download) + monkeypatch.setattr( + tc, "find_latest_hub_checkpoint", lambda repo_id, token=None, revision=None: "checkpoints/000010" + ) + + loaded = TrainPipelineConfig.from_pretrained("user/interrupted-run") + assert loaded.dataset.repo_id == "u/d" + # Tried the root config first, then fell back to the latest checkpoint's config. + assert calls == ["train_config.json", "checkpoints/000010/pretrained_model/train_config.json"] + + +def test_from_pretrained_raises_when_no_root_config_and_no_checkpoints(monkeypatch): + """No root config AND no checkpoints → a clear FileNotFoundError, not the raw HTTP error.""" + + def fake_hf_hub_download(filename=None, **kwargs): + raise _FakeHTTPError() + + monkeypatch.setattr(tc, "hf_hub_download", fake_hf_hub_download) + monkeypatch.setattr(tc, "find_latest_hub_checkpoint", lambda repo_id, token=None, revision=None: None) + + with pytest.raises(FileNotFoundError, match="train_config.json not found"): + TrainPipelineConfig.from_pretrained("user/empty-repo") diff --git a/tests/datasets/test_aggregate.py b/tests/datasets/test_aggregate.py index e9930575f..2fafd2777 100644 --- a/tests/datasets/test_aggregate.py +++ b/tests/datasets/test_aggregate.py @@ -29,7 +29,10 @@ from lerobot.configs import VIDEO_ENCODER_INFO_KEYS from lerobot.datasets.aggregate import aggregate_datasets from lerobot.datasets.feature_utils import features_equal_for_merge from lerobot.datasets.lerobot_dataset import LeRobotDataset -from tests.fixtures.constants import DUMMY_REPO_ID +from tests.fixtures.constants import ( + DUMMY_CAMERA_FEATURES_WITH_DEPTH, + DUMMY_REPO_ID, +) def assert_data_shards_one_row_group_per_episode(root): @@ -211,6 +214,26 @@ def assert_dataset_iteration_works(aggr_ds): pass +def assert_depth_keys_preserved(aggr_ds, ds_0, ds_1): + """Test that depth keys are correctly preserved after aggregation. + + Ensures that the ``is_depth_map`` marker on visual features survives + aggregation, so that downstream consumers (e.g. the dataset reader's + depth decoding path) keep working on the merged dataset. + """ + expected_depth_keys = set(ds_0.meta.depth_keys) + assert expected_depth_keys == set(ds_1.meta.depth_keys), ( + "Source datasets disagree on depth_keys; test setup is inconsistent" + ) + actual_depth_keys = set(aggr_ds.meta.depth_keys) + assert actual_depth_keys == expected_depth_keys, ( + f"Expected depth_keys {expected_depth_keys}, got {actual_depth_keys}" + ) + for key in expected_depth_keys: + info = aggr_ds.meta.info.features[key].get("info") or {} + assert info.get("is_depth_map") is True, f"Depth marker lost on feature {key!r} after aggregation" + + def assert_video_timestamps_within_bounds(aggr_ds): """Test that all video timestamps are within valid bounds for their respective video files. @@ -260,7 +283,11 @@ def assert_video_timestamps_within_bounds(aggr_ds): def test_aggregate_datasets(tmp_path, lerobot_dataset_factory): - """Test basic aggregation functionality with standard parameters.""" + """Test basic aggregation functionality with standard parameters. + + Source datasets include both RGB and depth video features so the same + aggregation flow is exercised on the ``is_depth_map`` branch. + """ ds_0_num_frames = 400 ds_1_num_frames = 800 ds_0_num_episodes = 10 @@ -272,14 +299,21 @@ def test_aggregate_datasets(tmp_path, lerobot_dataset_factory): repo_id=f"{DUMMY_REPO_ID}_0", total_episodes=ds_0_num_episodes, total_frames=ds_0_num_frames, + camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH, ) ds_1 = lerobot_dataset_factory( root=tmp_path / "test_1", repo_id=f"{DUMMY_REPO_ID}_1", total_episodes=ds_1_num_episodes, total_frames=ds_1_num_frames, + camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH, ) + # Confirm depth was actually wired into the source datasets so the + # rest of the assertions exercise the depth aggregation path. + assert len(ds_0.meta.depth_keys) > 0, "ds_0 should expose at least one depth key" + assert len(ds_1.meta.depth_keys) > 0, "ds_1 should expose at least one depth key" + aggregate_datasets( repo_ids=[ds_0.repo_id, ds_1.repo_id], roots=[ds_0.root, ds_1.root], @@ -306,6 +340,7 @@ def test_aggregate_datasets(tmp_path, lerobot_dataset_factory): assert_episode_indices_updated_correctly(aggr_ds, ds_0, ds_1) assert_video_frames_integrity(aggr_ds, ds_0, ds_1) assert_video_timestamps_within_bounds(aggr_ds) + assert_depth_keys_preserved(aggr_ds, ds_0, ds_1) assert_dataset_iteration_works(aggr_ds) @@ -423,7 +458,11 @@ def test_aggregate_incomplete_video_encoder_info_warns_and_nuls_encoders( def test_aggregate_with_low_threshold(tmp_path, lerobot_dataset_factory): - """Test aggregation with small file size limits to force file rotation/sharding.""" + """Test aggregation with small file size limits to force file rotation/sharding. + + Depth video features are included to verify that file rotation/concat + correctly handles depth-marked features alongside regular RGB ones. + """ ds_0_num_episodes = ds_1_num_episodes = 10 ds_0_num_frames = ds_1_num_frames = 400 @@ -432,14 +471,19 @@ def test_aggregate_with_low_threshold(tmp_path, lerobot_dataset_factory): repo_id=f"{DUMMY_REPO_ID}_small_0", total_episodes=ds_0_num_episodes, total_frames=ds_0_num_frames, + camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH, ) ds_1 = lerobot_dataset_factory( root=tmp_path / "small_1", repo_id=f"{DUMMY_REPO_ID}_small_1", total_episodes=ds_1_num_episodes, total_frames=ds_1_num_frames, + camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH, ) + assert len(ds_0.meta.depth_keys) > 0, "ds_0 should expose at least one depth key" + assert len(ds_1.meta.depth_keys) > 0, "ds_1 should expose at least one depth key" + # Use the new configurable parameters to force file rotation aggregate_datasets( repo_ids=[ds_0.repo_id, ds_1.repo_id], @@ -470,6 +514,7 @@ def test_aggregate_with_low_threshold(tmp_path, lerobot_dataset_factory): assert_episode_indices_updated_correctly(aggr_ds, ds_0, ds_1) assert_video_frames_integrity(aggr_ds, ds_0, ds_1) assert_video_timestamps_within_bounds(aggr_ds) + assert_depth_keys_preserved(aggr_ds, ds_0, ds_1) assert_dataset_iteration_works(aggr_ds) # Check that multiple files were actually created due to small size limits @@ -489,7 +534,8 @@ def test_video_timestamps_regression(tmp_path, lerobot_dataset_factory): """Regression test for video timestamp bug when merging datasets. This test specifically checks that video timestamps are correctly calculated - and accumulated when merging multiple datasets. + and accumulated when merging multiple datasets. Depth video features are + included so depth timestamps are also covered by the regression. """ datasets = [] for i in range(3): @@ -498,9 +544,13 @@ def test_video_timestamps_regression(tmp_path, lerobot_dataset_factory): repo_id=f"{DUMMY_REPO_ID}_regression_{i}", total_episodes=2, total_frames=100, + camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH, ) datasets.append(ds) + for i, ds in enumerate(datasets): + assert len(ds.meta.depth_keys) > 0, f"Dataset {i} should expose at least one depth key" + aggregate_datasets( repo_ids=[ds.repo_id for ds in datasets], roots=[ds.root for ds in datasets], @@ -517,12 +567,21 @@ def test_video_timestamps_regression(tmp_path, lerobot_dataset_factory): aggr_ds = LeRobotDataset(f"{DUMMY_REPO_ID}_regression_aggr", root=tmp_path / "regression_aggr") assert_video_timestamps_within_bounds(aggr_ds) + # Depth keys must survive the merge for the regression to cover the + # ``is_depth_map`` decoding branch. + assert set(aggr_ds.meta.depth_keys) == set(datasets[0].meta.depth_keys) + depth_keys = set(aggr_ds.meta.depth_keys) for i in range(len(aggr_ds)): item = aggr_ds[i] for key in aggr_ds.meta.video_keys: assert key in item, f"Video key {key} missing from item {i}" - assert item[key].shape[0] == 3, f"Expected 3 channels for video key {key}" + # Depth frames are single-channel (1, H, W) after dequantization; + # standard RGB frames keep the 3-channel layout. + expected_channels = 1 if key in depth_keys else 3 + assert item[key].shape[0] == expected_channels, ( + f"Expected {expected_channels} channels for video key {key}, got {item[key].shape}" + ) def assert_image_schema_preserved(aggr_ds): @@ -639,25 +698,31 @@ def test_aggregate_image_datasets(tmp_path, lerobot_dataset_factory): ds_0_num_episodes = 2 ds_1_num_episodes = 3 - # Create two image-based datasets (use_videos=False) + # Create two image-based datasets (use_videos=False) with a mix of RGB + # and depth-marked cameras so the depth path is exercised in image mode. ds_0 = lerobot_dataset_factory( root=tmp_path / "image_0", repo_id=f"{DUMMY_REPO_ID}_image_0", total_episodes=ds_0_num_episodes, total_frames=ds_0_num_frames, - use_videos=False, # Image-based dataset + use_videos=False, + camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH, ) ds_1 = lerobot_dataset_factory( root=tmp_path / "image_1", repo_id=f"{DUMMY_REPO_ID}_image_1", total_episodes=ds_1_num_episodes, total_frames=ds_1_num_frames, - use_videos=False, # Image-based dataset + use_videos=False, + camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH, ) # Verify source datasets have image keys assert len(ds_0.meta.image_keys) > 0, "ds_0 should have image keys" assert len(ds_1.meta.image_keys) > 0, "ds_1 should have image keys" + # And that the depth marker actually made it onto an image feature. + assert len(ds_0.meta.depth_keys) > 0, "ds_0 should expose at least one depth key" + assert len(ds_1.meta.depth_keys) > 0, "ds_1 should expose at least one depth key" # Aggregate the datasets aggregate_datasets( @@ -692,6 +757,7 @@ def test_aggregate_image_datasets(tmp_path, lerobot_dataset_factory): # Image-specific assertions assert_image_schema_preserved(aggr_ds) assert_image_frames_integrity(aggr_ds, ds_0, ds_1) + assert_depth_keys_preserved(aggr_ds, ds_0, ds_1) # Verify images can be accessed and have correct shape sample_item = aggr_ds[0] diff --git a/tests/datasets/test_compute_stats.py b/tests/datasets/test_compute_stats.py index 0f5abfb95..9f399b85c 100644 --- a/tests/datasets/test_compute_stats.py +++ b/tests/datasets/test_compute_stats.py @@ -35,7 +35,11 @@ from lerobot.utils.constants import OBS_IMAGE, OBS_STATE def mock_load_image_as_numpy(path, dtype, channel_first): - return np.ones((3, 32, 32), dtype=dtype) if channel_first else np.ones((32, 32, 3), dtype=dtype) + is_depth = "depth" in str(path) + channels = 1 if is_depth else 3 + out_dtype = np.uint16 if is_depth else dtype + arr = np.arange(channels * 32 * 32, dtype=out_dtype).reshape(channels, 32, 32) + return arr if channel_first else arr.transpose(1, 2, 0) @pytest.fixture @@ -168,22 +172,33 @@ def test_get_feature_stats_single_value(): def test_compute_episode_stats(): + depth_key = "observation.images.depth" episode_data = { OBS_IMAGE: [f"image_{i}.jpg" for i in range(100)], + depth_key: [f"depth_{i}.tiff" for i in range(100)], OBS_STATE: np.random.rand(100, 10), } features = { OBS_IMAGE: {"dtype": "image"}, + depth_key: {"dtype": "image", "info": {"is_depth_map": True}}, OBS_STATE: {"dtype": "numeric"}, } with patch("lerobot.datasets.compute_stats.load_image_as_numpy", side_effect=mock_load_image_as_numpy): stats = compute_episode_stats(episode_data, features) - assert OBS_IMAGE in stats and OBS_STATE in stats + assert OBS_IMAGE in stats and depth_key in stats and OBS_STATE in stats assert stats[OBS_IMAGE]["count"].item() == 100 + assert stats[depth_key]["count"].item() == 100 assert stats[OBS_STATE]["count"].item() == 100 assert stats[OBS_IMAGE]["mean"].shape == (3, 1, 1) + assert stats[depth_key]["mean"].shape == (1, 1, 1) + # Depth keeps raw values: max far exceeds 255, proving no /255 and no uint8 downcast. + assert stats[depth_key]["min"].item() == 0.0 + assert stats[depth_key]["max"].item() == 1023.0 + # RGB is normalized to [0, 1]. + np.testing.assert_allclose(stats[OBS_IMAGE]["min"], 0.0) + np.testing.assert_allclose(stats[OBS_IMAGE]["max"], 1.0) def test_assert_type_and_shape_valid(): @@ -618,25 +633,31 @@ def test_compute_episode_stats_with_custom_quantiles(): def test_compute_episode_stats_with_image_data(): """Test quantile computation with image features.""" image_paths = [f"image_{i}.jpg" for i in range(50)] + depth_paths = [f"depth_{i}.tiff" for i in range(50)] episode_data = { "observation.image": image_paths, + "observation.images.depth": depth_paths, "action": np.random.normal(0, 1, (50, 5)), } features = { "observation.image": {"dtype": "image"}, + "observation.images.depth": {"dtype": "image", "info": {"is_depth_map": True}}, "action": {"dtype": "float32", "shape": (5,)}, } with patch("lerobot.datasets.compute_stats.load_image_as_numpy", side_effect=mock_load_image_as_numpy): stats = compute_episode_stats(episode_data, features) - # Image quantiles should be normalized and have correct shape - assert "q01" in stats["observation.image"] - assert "q50" in stats["observation.image"] - assert "q99" in stats["observation.image"] - assert stats["observation.image"]["q01"].shape == (3, 1, 1) - assert stats["observation.image"]["q50"].shape == (3, 1, 1) - assert stats["observation.image"]["q99"].shape == (3, 1, 1) + # RGB image quantiles should be normalized and per-channel. + for q in ("q01", "q50", "q99"): + assert stats["observation.image"][q].shape == (3, 1, 1) + + # Depth quantiles are single-channel and kept in raw (un-normalized) units. + for q in ("q01", "q50", "q99"): + assert stats["observation.images.depth"][q].shape == (1, 1, 1) + # Depth max stays in raw units (not /255, not uint8-capped); RGB is normalized. + assert stats["observation.images.depth"]["max"].item() == 1023.0 + np.testing.assert_allclose(stats["observation.image"]["max"], 1.0) # Action quantiles should have correct shape assert stats["action"]["q01"].shape == (5,) diff --git a/tests/datasets/test_dataset_metadata.py b/tests/datasets/test_dataset_metadata.py index 171d8af8b..a1630f17d 100644 --- a/tests/datasets/test_dataset_metadata.py +++ b/tests/datasets/test_dataset_metadata.py @@ -59,11 +59,13 @@ def _make_dummy_stats(features: dict) -> dict: stats = {} for key, ft in features.items(): if ft["dtype"] in ("image", "video"): + channels = ft["shape"][-1] + stat_shape = (channels, 1, 1) stats[key] = { - "max": np.ones((3, 1, 1), dtype=np.float32), - "mean": np.full((3, 1, 1), 0.5, dtype=np.float32), - "min": np.zeros((3, 1, 1), dtype=np.float32), - "std": np.full((3, 1, 1), 0.25, dtype=np.float32), + "max": np.ones(stat_shape, dtype=np.float32), + "mean": np.full(stat_shape, 0.5, dtype=np.float32), + "min": np.zeros(stat_shape, dtype=np.float32), + "std": np.full(stat_shape, 0.25, dtype=np.float32), "count": np.array([5]), } elif ft["dtype"] in ("float32", "float64", "int64"): @@ -142,6 +144,45 @@ def test_create_without_videos_has_no_video_path(tmp_path): assert meta.video_keys == [] +@pytest.mark.parametrize( + ("marker_field", "marker_key"), + [ + ("info", "is_depth_map"), + ("info", "video.is_depth_map"), + ("video_info", "video.is_depth_map"), + ], + ids=["info.is_depth_map", "info.video.is_depth_map_legacy", "video_info.video.is_depth_map_legacy"], +) +def test_depth_keys_property_filters_by_marker(tmp_path, marker_field, marker_key): + """``depth_keys`` recognises the canonical and the two legacy marker variants.""" + depth_feature = { + "dtype": "video", + "shape": (64, 96, 1), + "names": ["height", "width", "channels"], + marker_field: {marker_key: True}, + } + features = { + **VIDEO_FEATURES, + "observation.images.laptop_depth": depth_feature, + } + meta = LeRobotDatasetMetadata.create( + repo_id="test/depth_keys", + fps=DEFAULT_FPS, + features=features, + root=tmp_path / f"depth_keys_{marker_field}_{marker_key.replace('.', '_')}", + ) + + assert set(meta.video_keys) == {"observation.images.laptop", "observation.images.laptop_depth"} + assert meta.depth_keys == ["observation.images.laptop_depth"] + + +def test_depth_keys_empty_when_no_marker(tmp_path): + meta = LeRobotDatasetMetadata.create( + repo_id="test/no_depth", fps=DEFAULT_FPS, features=VIDEO_FEATURES, root=tmp_path / "no_depth" + ) + assert meta.depth_keys == [] + + def test_create_raises_on_existing_directory(tmp_path): """create() raises if root directory already exists.""" root = tmp_path / "existing" diff --git a/tests/datasets/test_dataset_tools.py b/tests/datasets/test_dataset_tools.py index d36312920..c19e7f41f 100644 --- a/tests/datasets/test_dataset_tools.py +++ b/tests/datasets/test_dataset_tools.py @@ -24,7 +24,7 @@ import torch pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])") -from lerobot.configs import VideoEncoderConfig +from lerobot.configs import DepthEncoderConfig, RGBEncoderConfig from lerobot.datasets.dataset_tools import ( add_features, convert_image_to_video_dataset, @@ -37,7 +37,9 @@ from lerobot.datasets.dataset_tools import ( split_dataset, ) from lerobot.datasets.io_utils import load_info -from tests.datasets.test_video_encoding import _add_frames, require_h264, require_libsvtav1 +from tests.datasets.test_video_encoding import require_h264, require_hevc, require_libsvtav1 +from tests.fixtures.constants import DUMMY_DEPTH_FEATURES, DUMMY_DEPTH_KEY +from tests.fixtures.dataset_factories import add_frames @pytest.fixture @@ -1251,7 +1253,7 @@ def test_convert_image_to_video_dataset(tmp_path): dataset=source_dataset, output_dir=output_dir, repo_id="lerobot/pusht_video", - camera_encoder=VideoEncoderConfig( + rgb_encoder=RGBEncoderConfig( vcodec="libsvtav1", pix_fmt="yuv420p", g=2, @@ -1332,9 +1334,131 @@ def test_convert_image_to_video_dataset_subset_episodes(tmp_path): shutil.rmtree(output_dir) +@require_libsvtav1 +@require_hevc +def test_convert_image_to_video_dataset_depth(tmp_path, empty_lerobot_dataset_factory): + """Depth image features convert to depth videos using the depth encoder. + + Mirrors :func:`test_convert_image_to_video_dataset` but with a small local + image dataset that mixes an RGB camera with a depth camera, so the + ``depth_keys`` → ``depth_encoder`` routing and ``is_depth_map`` preservation + are exercised end-to-end. + """ + features = { + "action": {"dtype": "float32", "shape": (2,), "names": ["a", "b"]}, + "observation.images.cam": { + "dtype": "image", + "shape": (64, 96, 3), + "names": ["height", "width", "channels"], + }, + "observation.images.depth": { + "dtype": "image", + "shape": (64, 96, 1), + "names": ["height", "width", "channels"], + "info": {"is_depth_map": True}, + }, + } + source_dataset = empty_lerobot_dataset_factory( + root=tmp_path / "img_ds", + features=features, + use_videos=False, + ) + + add_frames(source_dataset, num_frames=4) + source_dataset.save_episode() + source_dataset.finalize() + + # Source is an image dataset with the depth marker on the depth camera. + assert len(source_dataset.meta.video_keys) == 0 + assert "observation.images.depth" in source_dataset.meta.depth_keys + + output_dir = tmp_path / "video_ds" + with ( + patch("lerobot.datasets.dataset_metadata.get_safe_version") as mock_get_safe_version, + patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download, + ): + mock_get_safe_version.return_value = "v3.0" + mock_snapshot_download.return_value = str(output_dir) + + # Use non-default quantization params so the persisted metadata must + # come from the depth encoder (not RGB encoder defaults). + depth_encoder = DepthEncoderConfig( + vcodec="hevc", + pix_fmt="gray12le", + g=2, + crf=30, + depth_min=0.05, + depth_max=8.0, + shift=2.0, + use_log=False, + ) + video_dataset = convert_image_to_video_dataset( + dataset=source_dataset, + output_dir=output_dir, + repo_id="dummy/depth_video", + rgb_encoder=RGBEncoderConfig(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30), + depth_encoder=depth_encoder, + num_workers=1, + ) + + # Both cameras are now videos, and the depth marker survived the conversion. + assert "observation.images.cam" in video_dataset.meta.video_keys + assert "observation.images.depth" in video_dataset.meta.video_keys + assert "observation.images.depth" in video_dataset.meta.depth_keys + assert "observation.images.cam" not in video_dataset.meta.depth_keys + + depth_path = video_dataset.root / video_dataset.meta.get_video_file_path(0, "observation.images.depth") + assert depth_path.exists(), f"Depth video file should exist: {depth_path}" + + # The persisted depth-video metadata must carry the depth quantization params + # from the depth encoder (so frames dequantize correctly on read), and the RGB + # camera must not be marked as a depth map. + persisted_info = load_info(video_dataset.root) + depth_info = persisted_info.features["observation.images.depth"]["info"] + assert depth_info["is_depth_map"] is True + assert DepthEncoderConfig.from_video_info(depth_info) == depth_encoder + + cam_info = persisted_info.features["observation.images.cam"]["info"] + assert cam_info.get("is_depth_map") is False + assert "video.codec" in cam_info + + # ─── reencode_dataset ───────────────────────────────────────────────── +@require_hevc +def test_reencode_dataset_depth_uses_depth_encoder(tmp_path, empty_lerobot_dataset_factory): + """Depth videos are re-encoded with the depth encoder and keep their depth metadata. + + Depth-focused companion to :func:`test_reencode_dataset_multi_key_multiprocessing`. + """ + initial_cfg = DepthEncoderConfig(vcodec="hevc", pix_fmt="gray12le", g=2, crf=30) + dataset = empty_lerobot_dataset_factory( + root=tmp_path / "ds", + features=DUMMY_DEPTH_FEATURES, + use_videos=True, + depth_encoder=initial_cfg, + ) + + add_frames(dataset, num_frames=4) + dataset.save_episode() + dataset.finalize() + + assert DUMMY_DEPTH_KEY in dataset.meta.depth_keys + + target_cfg = DepthEncoderConfig(vcodec="hevc", pix_fmt="gray12le", g=6, crf=23) + result = reencode_dataset(dataset, depth_encoder=target_cfg, num_workers=0) + + assert result is dataset + + persisted_info = load_info(dataset.root) + depth_info = persisted_info.features[DUMMY_DEPTH_KEY].get("info", {}) + # Re-encode applied the new codec parameters to the depth video ... + assert DepthEncoderConfig.from_video_info(depth_info) == target_cfg + # ... while preserving the depth marker. + assert depth_info["is_depth_map"] is True + + @require_libsvtav1 @require_h264 def test_reencode_dataset_multi_key_multiprocessing( @@ -1342,29 +1466,29 @@ def test_reencode_dataset_multi_key_multiprocessing( ): """Re-encode a two-camera dataset with num_workers=2 and verify metadata refresh.""" features = features_factory(use_videos=True) - initial_cfg = VideoEncoderConfig(vcodec="libsvtav1", g=2, crf=30, preset=12) + initial_cfg = RGBEncoderConfig(vcodec="libsvtav1", g=2, crf=30, preset=12) dataset = empty_lerobot_dataset_factory( root=tmp_path / "ds", features=features, use_videos=True, - camera_encoder=initial_cfg, + rgb_encoder=initial_cfg, ) - _add_frames(dataset, num_frames=4) + add_frames(dataset, num_frames=4) dataset.save_episode() - _add_frames(dataset, num_frames=4) + add_frames(dataset, num_frames=4) dataset.save_episode() dataset.finalize() assert len(dataset.meta.video_keys) == 2 - target_cfg = VideoEncoderConfig(vcodec="h264", g=6, crf=23, pix_fmt="yuv420p") + target_cfg = RGBEncoderConfig(vcodec="h264", g=6, crf=23, pix_fmt="yuv420p") - result = reencode_dataset(dataset, camera_encoder=target_cfg, num_workers=2) + result = reencode_dataset(dataset, rgb_encoder=target_cfg, num_workers=2) assert result is dataset persisted_info = load_info(dataset.root) for vk in dataset.meta.video_keys: - persisted_encoder = VideoEncoderConfig.from_video_info(persisted_info.features[vk].get("info", {})) + persisted_encoder = RGBEncoderConfig.from_video_info(persisted_info.features[vk].get("info", {})) assert persisted_encoder == target_cfg diff --git a/tests/datasets/test_dataset_writer.py b/tests/datasets/test_dataset_writer.py index 8670aeebc..17785ad74 100644 --- a/tests/datasets/test_dataset_writer.py +++ b/tests/datasets/test_dataset_writer.py @@ -53,8 +53,8 @@ def _make_frame(features: dict, task: str = "Dummy task") -> dict: # ── Existing encode_video_worker tests ─────────────────────────────── -def test_encode_video_worker_forwards_camera_encoder(tmp_path): - """_encode_video_worker forwards camera_encoder to encode_video_frames.""" +def test_encode_video_worker_forwards_video_encoder(tmp_path): + """_encode_video_worker forwards video_encoder to encode_video_frames.""" video_key = "observation.images.laptop" fpath = DEFAULT_IMAGE_PATH.format(image_key=video_key, episode_index=0, frame_index=0) img_dir = tmp_path / Path(fpath).parent @@ -74,16 +74,16 @@ def test_encode_video_worker_forwards_camera_encoder(tmp_path): 0, tmp_path, fps=30, - camera_encoder=VideoEncoderConfig(vcodec="h264", preset=None), + video_encoder=VideoEncoderConfig(vcodec="h264", preset=None), encoder_threads=4, ) - assert captured_kwargs["camera_encoder"].vcodec == "h264" + assert captured_kwargs["video_encoder"].vcodec == "h264" assert captured_kwargs["encoder_threads"] == 4 -def test_encode_video_worker_default_camera_encoder(tmp_path): - """_encode_video_worker passes None camera_encoder which encode_video_frames defaults.""" +def test_encode_video_worker_default_video_encoder(tmp_path): + """_encode_video_worker passes None video_encoder which encode_video_frames defaults.""" video_key = "observation.images.laptop" fpath = DEFAULT_IMAGE_PATH.format(image_key=video_key, episode_index=0, frame_index=0) img_dir = tmp_path / Path(fpath).parent @@ -100,7 +100,7 @@ def test_encode_video_worker_default_camera_encoder(tmp_path): with patch("lerobot.datasets.dataset_writer.encode_video_frames", side_effect=mock_encode): _encode_video_worker(video_key, 0, tmp_path, fps=30) - assert captured_kwargs["camera_encoder"] is None + assert captured_kwargs["video_encoder"] is None assert captured_kwargs["encoder_threads"] is None diff --git a/tests/datasets/test_datasets.py b/tests/datasets/test_datasets.py index 1d2fb1d55..225479814 100644 --- a/tests/datasets/test_datasets.py +++ b/tests/datasets/test_datasets.py @@ -1534,6 +1534,10 @@ def test_valid_video_codecs_constant(): assert "auto" in VALID_VIDEO_CODECS assert "h264_videotoolbox" in VALID_VIDEO_CODECS assert "h264_nvenc" in VALID_VIDEO_CODECS + assert "h264_vaapi" in VALID_VIDEO_CODECS + assert "h264_qsv" in VALID_VIDEO_CODECS + assert "hevc_videotoolbox" in VALID_VIDEO_CODECS + assert "hevc_nvenc" in VALID_VIDEO_CODECS assert len(VALID_VIDEO_CODECS) == 10 diff --git a/tests/datasets/test_depth.py b/tests/datasets/test_depth.py new file mode 100644 index 000000000..a075fa6b5 --- /dev/null +++ b/tests/datasets/test_depth.py @@ -0,0 +1,247 @@ +"""Tests for the depth-integration feature. + +Covers: +- ``depth_utils`` quantize/dequantize round-trips and backend agreement. +- Image-writer support for single-channel depth. +- Hardware-feature → depth flag routing. +- Feature-to-file-format routing through the dataset writer. + +Depth metadata detection on ``LeRobotDatasetMetadata.depth_keys`` lives in +``test_dataset_metadata.py``. Depth video encoding/decoding lives in +``test_video_encoding.py``. +""" + +from pathlib import Path + +import pytest + +pytest.importorskip("av", reason="av is required (install lerobot[dataset])") + +import av +import numpy as np +import PIL.Image +import torch + +from lerobot.configs import DepthEncoderConfig +from lerobot.configs.video import ( + DEFAULT_DEPTH_MAX, + DEFAULT_DEPTH_MIN, + DEPTH_METER_UNIT, + DEPTH_MILLIMETER_UNIT, + DEPTH_QMAX, +) +from lerobot.datasets.depth_utils import dequantize_depth, quantize_depth +from lerobot.datasets.image_writer import image_array_to_pil_image, write_image +from tests.fixtures.constants import ( + DEFAULT_FPS, + DUMMY_CAMERA_FEATURES, + DUMMY_CAMERA_FEATURES_WITH_DEPTH, + DUMMY_CHW, + DUMMY_DEPTH_CAMERA_FEATURES, + DUMMY_REPO_ID, +) +from tests.fixtures.dataset_factories import add_frames + +_, H, W = DUMMY_CHW + + +def _depth_metres_ramp() -> np.ndarray: + """Linearly-spaced float32 depth in metres covering the default range.""" + return np.linspace(DEFAULT_DEPTH_MIN, DEFAULT_DEPTH_MAX, H * W, dtype=np.float32).reshape(H, W) + + +# ── 1. Quantize / dequantize round-trips ────────────────────────────── + + +class TestQuantizeDequantize: + """Numerical contract of ``quantize_depth`` / ``dequantize_depth``.""" + + @pytest.mark.parametrize("use_log", [False, True]) + @pytest.mark.parametrize("output_unit", [DEPTH_METER_UNIT, DEPTH_MILLIMETER_UNIT]) + @pytest.mark.parametrize("output_channel_last", [False, True]) + def test_roundtrip(self, use_log, output_unit, output_channel_last): + """quantize → dequantize recovers depth; layout and unit are honored.""" + depth = _depth_metres_ramp() + quantized = quantize_depth(depth, use_log=use_log, video_backend=None) + recovered = dequantize_depth( + quantized, + use_log=use_log, + output_unit=output_unit, + output_tensor=False, + output_channel_last=output_channel_last, + ) + + expected_shape = (H, W, 1) if output_channel_last else (1, H, W) + assert recovered.shape == expected_shape + + recovered_m = recovered.astype(np.float32) + if output_unit == DEPTH_MILLIMETER_UNIT: + recovered_m = recovered_m / 1000.0 + recovered_2d = recovered_m[..., 0] if output_channel_last else recovered_m[0] + + if use_log: + # Log mode: tighter near-range error than far-range (the whole point). + near = depth < 1.0 + far = depth > 8.0 + err_near = np.abs(recovered_2d[near] - depth[near]) + err_far = np.abs(recovered_2d[far] - depth[far]) + assert err_near.mean() < err_far.mean() + else: + # Linear mode: bounded by quant step + 1 mm of unit-conversion rounding. + tol = (DEFAULT_DEPTH_MAX - DEFAULT_DEPTH_MIN) / DEPTH_QMAX + 1e-3 + np.testing.assert_allclose(recovered_2d, depth, atol=tol) + + @pytest.mark.parametrize("use_log", [False, True]) + @pytest.mark.parametrize("output_unit", [DEPTH_METER_UNIT, DEPTH_MILLIMETER_UNIT]) + def test_numpy_torch_agree(self, use_log, output_unit): + """Batched torch path produces the same values as the numpy path.""" + batch_size = 3 + per_frame = np.linspace(0, DEPTH_QMAX, H * W, dtype=np.uint16).reshape(H, W) + batch_np = np.broadcast_to(per_frame[None, None, ...], (batch_size, 1, H, W)).copy() + batch_t = torch.from_numpy(batch_np.astype(np.int32)) # torch.uint16 support is patchy. + + ref = dequantize_depth(batch_np, use_log=use_log, output_unit=output_unit, output_tensor=False) + out = dequantize_depth(batch_t, use_log=use_log, output_unit=output_unit, output_tensor=True) + + assert isinstance(out, torch.Tensor) + assert out.shape == (batch_size, 1, H, W) + # ``m``: float32 noise (~10 µm in log mode, after ``exp``) — still 200× below the ~2 mm quant step. + # ``mm`` + tensor stays in float32 (no uint16 round-trip), so allow 1 mm slop. + atol = 1e-5 if output_unit == DEPTH_METER_UNIT else 1.0 + np.testing.assert_allclose(out.cpu().numpy().astype(np.float64), ref.astype(np.float64), atol=atol) + + @pytest.mark.parametrize( + "input_shape,output_shape", + [ + ((H, W), (1, H, W)), + ((1, H, W), (1, H, W)), + ((H, W, 1), (1, H, W)), + ((3, 1, H, W), (3, 1, H, W)), + ((3, H, W, 1), (3, 1, H, W)), + ], + ) + def test_input_layouts_accepted(self, input_shape, output_shape): + """All documented input layouts decode to the channel-first default.""" + quantized = np.full(input_shape, DEPTH_QMAX // 2, dtype=np.uint16) + out = dequantize_depth(quantized, output_unit=DEPTH_METER_UNIT, output_tensor=False) + assert out.shape == output_shape + + def test_pyav_frame_roundtrip(self): + """quantize → av.VideoFrame → dequantize works.""" + depth = _depth_metres_ramp() + frame = quantize_depth(depth, use_log=False, video_backend="pyav") + assert isinstance(frame, av.VideoFrame) + + recovered = dequantize_depth(frame, use_log=False, output_unit=DEPTH_METER_UNIT, output_tensor=False) + assert recovered.shape == (1, H, W) + tol = (DEFAULT_DEPTH_MAX - DEFAULT_DEPTH_MIN) / DEPTH_QMAX + 1e-3 + np.testing.assert_allclose(recovered[0], depth, atol=tol) + + def test_invalid_log_params_raises(self): + with pytest.raises(ValueError, match=r"depth_min \+ shift must be positive"): + quantize_depth(_depth_metres_ramp(), depth_min=1.0, shift=-2.0, use_log=True, video_backend=None) + + +# ── 2. Image writer depth support ───────────────────────────────────── + + +class TestImageWriterDepth: + """``image_array_to_pil_image`` and ``write_image`` for depth maps.""" + + @pytest.mark.parametrize("dtype,expected_mode", [(np.uint16, "I;16"), (np.float32, "F")]) + @pytest.mark.parametrize("shape", [(H, W), (H, W, 1), (1, H, W)]) + def test_pil_depth_modes_and_squeeze(self, dtype, expected_mode, shape): + """Single-channel depth converts to PIL with the right mode and (W, H) size.""" + arr = np.zeros(shape, dtype=dtype) + img = image_array_to_pil_image(arr) + assert img.mode == expected_mode + assert img.size == (W, H) + + def test_write_image_tiff_roundtrip(self, tmp_path): + """uint16 depth round-trips through .tiff.""" + arr = np.arange(H * W, dtype=np.uint16).reshape(H, W) + fpath = tmp_path / "depth.tiff" + write_image(arr, fpath) + with PIL.Image.open(fpath) as loaded: + recovered = np.array(loaded) + np.testing.assert_array_equal(recovered, arr) + + +# ── 3. Hardware-feature → depth flag ────────────────────────────────── + + +class TestHwToDatasetFeaturesDepth: + """``hw_to_dataset_features`` flags single-channel cameras as depth.""" + + @pytest.mark.parametrize("channels,is_depth", [(1, True), (3, False)]) + def test_depth_marker_by_channels(self, channels, is_depth): + from lerobot.utils.feature_utils import hw_to_dataset_features + + features = hw_to_dataset_features({"cam": (480, 640, channels)}, prefix="observation") + assert features["observation.images.cam"]["info"]["is_depth_map"] is is_depth + + def test_invalid_channel_count_raises(self): + from lerobot.utils.feature_utils import hw_to_dataset_features + + with pytest.raises(ValueError, match="Expected a 3-tuple"): + hw_to_dataset_features({"cam": (480, 640, 2)}, prefix="observation") + + +# ── 4. Feature-to-file-format routing ──────────────────────────────── + + +# Keys derived from DUMMY_CAMERA_FEATURES_WITH_DEPTH; pick one RGB and the depth camera. +RGB_KEY = next(iter(DUMMY_CAMERA_FEATURES)) +DEPTH_KEY = next(iter(DUMMY_DEPTH_CAMERA_FEATURES)) + + +class TestFeatureFileRouting: + """Depth vs RGB features route to the correct file format.""" + + NUM_FRAMES = 5 + + def test_image_mode_depth_tiff_rgb_png(self, tmp_path, features_factory): + """Without video encoding: depth → .tiff, RGB → .png.""" + from lerobot.datasets.lerobot_dataset import LeRobotDataset + + features = features_factory(camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH, use_videos=False) + dataset = LeRobotDataset.create( + repo_id=DUMMY_REPO_ID, + fps=DEFAULT_FPS, + features=features, + root=tmp_path / "ds", + use_videos=False, + ) + + add_frames(dataset, num_frames=self.NUM_FRAMES) + + buf = dataset.writer.episode_buffer + assert all(Path(p).suffix == ".tiff" for p in buf[DEPTH_KEY]) + assert all(Path(p).suffix == ".png" for p in buf[RGB_KEY]) + + dataset.save_episode() + dataset.finalize() + + def test_video_mode_depth_uses_depth_encoder(self, tmp_path, features_factory): + """With streaming video encoding: depth → DepthEncoderConfig, RGB does not.""" + from lerobot.datasets.lerobot_dataset import LeRobotDataset + + features = features_factory(camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH, use_videos=True) + dataset = LeRobotDataset.create( + repo_id=DUMMY_REPO_ID, + fps=DEFAULT_FPS, + features=features, + root=tmp_path / "ds", + use_videos=True, + streaming_encoding=True, + ) + + add_frames(dataset, num_frames=self.NUM_FRAMES) + + encoder = dataset.writer._streaming_encoder + assert encoder is not None + assert isinstance(encoder._threads[DEPTH_KEY].video_encoder, DepthEncoderConfig) + assert not isinstance(encoder._threads[RGB_KEY].video_encoder, DepthEncoderConfig) + + dataset.save_episode() + dataset.finalize() diff --git a/tests/datasets/test_image_writer.py b/tests/datasets/test_image_writer.py index 916b8f017..1cf2cf75c 100644 --- a/tests/datasets/test_image_writer.py +++ b/tests/datasets/test_image_writer.py @@ -94,7 +94,7 @@ def test_image_array_to_pil_image_pytorch_format(img_array_factory): def test_image_array_to_pil_image_single_channel(img_array_factory): img_array = img_array_factory(channels=1) - with pytest.raises(NotImplementedError): + with pytest.raises(ValueError, match="Unsupported single-channel image dtype"): image_array_to_pil_image(img_array) @@ -344,7 +344,7 @@ def test_with_different_image_formats(tmp_path, img_array_factory): writer = AsyncImageWriter() try: image_array = img_array_factory() - formats = ["png", "jpeg", "bmp"] + formats = ["png", "tiff", "tif"] for fmt in formats: fpath = tmp_path / f"test_image.{fmt}" write_image(image_array, fpath) diff --git a/tests/datasets/test_streaming_video_encoder.py b/tests/datasets/test_streaming_video_encoder.py index b69f24254..1ffad6854 100644 --- a/tests/datasets/test_streaming_video_encoder.py +++ b/tests/datasets/test_streaming_video_encoder.py @@ -26,7 +26,7 @@ pytest.importorskip("av", reason="av is required (install lerobot[dataset])") import av # noqa: E402 -from lerobot.configs import VideoEncoderConfig +from lerobot.configs import RGBEncoderConfig from lerobot.datasets.pyav_utils import get_codec from lerobot.datasets.video_utils import ( StreamingVideoEncoder, @@ -57,13 +57,11 @@ class TestCameraEncoderThread: result_queue: queue.Queue = queue.Queue(maxsize=1) stop_event = threading.Event() - enc_cfg = VideoEncoderConfig(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13) + enc_cfg = RGBEncoderConfig(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13) encoder_thread = _CameraEncoderThread( video_path=video_path, fps=fps, - vcodec=enc_cfg.vcodec, - pix_fmt=enc_cfg.pix_fmt, - codec_options=enc_cfg.get_codec_options(as_strings=True), + video_encoder=enc_cfg, frame_queue=frame_queue, result_queue=result_queue, stop_event=stop_event, @@ -108,13 +106,11 @@ class TestCameraEncoderThread: result_queue: queue.Queue = queue.Queue(maxsize=1) stop_event = threading.Event() - enc_cfg = VideoEncoderConfig(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13) + enc_cfg = RGBEncoderConfig(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13) encoder_thread = _CameraEncoderThread( video_path=video_path, fps=fps, - vcodec=enc_cfg.vcodec, - pix_fmt=enc_cfg.pix_fmt, - codec_options=enc_cfg.get_codec_options(as_strings=True), + video_encoder=enc_cfg, frame_queue=frame_queue, result_queue=result_queue, stop_event=stop_event, @@ -142,13 +138,11 @@ class TestCameraEncoderThread: result_queue: queue.Queue = queue.Queue(maxsize=1) stop_event = threading.Event() - enc_cfg = VideoEncoderConfig(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13) + enc_cfg = RGBEncoderConfig(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13) encoder_thread = _CameraEncoderThread( video_path=video_path, fps=fps, - vcodec=enc_cfg.vcodec, - pix_fmt=enc_cfg.pix_fmt, - codec_options=enc_cfg.get_codec_options(as_strings=True), + video_encoder=enc_cfg, frame_queue=frame_queue, result_queue=result_queue, stop_event=stop_event, @@ -171,15 +165,15 @@ class TestCameraEncoderThread: class TestStreamingVideoEncoder: def _make_encoder_config(self, **kwargs): - """Helper to build a VideoEncoderConfig.""" - return VideoEncoderConfig(**kwargs) + """Helper to build an RGBEncoderConfig.""" + return RGBEncoderConfig(**kwargs) def test_single_camera_episode(self, tmp_path): """Test encoding a single camera episode.""" video_keys = [f"{OBS_IMAGES}.laptop"] encoder = StreamingVideoEncoder( fps=30, - camera_encoder=self._make_encoder_config( + rgb_encoder=self._make_encoder_config( vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13 ), ) @@ -211,7 +205,7 @@ class TestStreamingVideoEncoder: video_keys = [f"{OBS_IMAGES}.laptop", f"{OBS_IMAGES}.phone"] encoder = StreamingVideoEncoder( fps=30, - camera_encoder=self._make_encoder_config(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30), + rgb_encoder=self._make_encoder_config(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30), ) encoder.start_episode(video_keys, tmp_path) @@ -237,7 +231,7 @@ class TestStreamingVideoEncoder: video_keys = [f"{OBS_IMAGES}.cam"] encoder = StreamingVideoEncoder( fps=30, - camera_encoder=self._make_encoder_config(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30), + rgb_encoder=self._make_encoder_config(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30), ) for ep in range(3): @@ -263,7 +257,7 @@ class TestStreamingVideoEncoder: video_keys = [f"{OBS_IMAGES}.cam"] encoder = StreamingVideoEncoder( fps=30, - camera_encoder=self._make_encoder_config(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30), + rgb_encoder=self._make_encoder_config(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30), ) encoder.start_episode(video_keys, tmp_path) @@ -309,7 +303,7 @@ class TestStreamingVideoEncoder: video_keys = [f"{OBS_IMAGES}.cam"] encoder = StreamingVideoEncoder( fps=30, - camera_encoder=self._make_encoder_config( + rgb_encoder=self._make_encoder_config( vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13 ), ) @@ -346,7 +340,7 @@ class TestStreamingVideoEncoder: video_keys = [f"{OBS_IMAGES}.cam1", f"{OBS_IMAGES}.cam2"] encoder = StreamingVideoEncoder( fps=30, - camera_encoder=self._make_encoder_config(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30), + rgb_encoder=self._make_encoder_config(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30), ) encoder.start_episode(video_keys, tmp_path) @@ -375,7 +369,7 @@ class TestStreamingVideoEncoder: def test_encoder_threads_passed_to_thread(self, tmp_path): """Test that encoder_threads is stored and passed through to encoder threads.""" video_keys = [f"{OBS_IMAGES}.cam"] - cfg = VideoEncoderConfig( + cfg = RGBEncoderConfig( vcodec="libsvtav1", pix_fmt="yuv420p", g=2, @@ -383,7 +377,7 @@ class TestStreamingVideoEncoder: ) encoder = StreamingVideoEncoder( fps=30, - camera_encoder=cfg, + rgb_encoder=cfg, encoder_threads=2, ) assert encoder._encoder_threads == 2 @@ -391,7 +385,8 @@ class TestStreamingVideoEncoder: # Verify codec options include thread tuning for libsvtav1 (lp=…) thread = encoder._threads[f"{OBS_IMAGES}.cam"] - assert "svtav1-params" in thread.codec_options or "threads" in thread.codec_options + codec_opts = thread.video_encoder.get_codec_options(encoder_threads=thread.encoder_threads) + assert "svtav1-params" in codec_opts or "threads" in codec_opts # Feed some frames and finish to ensure it works end-to-end num_frames = 10 @@ -422,7 +417,7 @@ class TestStreamingVideoEncoder: video_keys = [f"{OBS_IMAGES}.cam"] encoder = StreamingVideoEncoder( fps=30, - camera_encoder=self._make_encoder_config( + rgb_encoder=self._make_encoder_config( vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13 ), queue_maxsize=1, diff --git a/tests/datasets/test_video_encoding.py b/tests/datasets/test_video_encoding.py index 2a35f3210..80819d665 100644 --- a/tests/datasets/test_video_encoding.py +++ b/tests/datasets/test_video_encoding.py @@ -26,7 +26,7 @@ pytest.importorskip("av", reason="av is required (install lerobot[dataset])") import av # noqa: E402 -from lerobot.configs import VALID_VIDEO_CODECS, VideoEncoderConfig +from lerobot.configs import VALID_VIDEO_CODECS, DepthEncoderConfig, RGBEncoderConfig, VideoEncoderConfig from lerobot.datasets.image_writer import write_image from lerobot.datasets.lerobot_dataset import LeRobotDataset from lerobot.datasets.pyav_utils import get_codec @@ -37,7 +37,15 @@ from lerobot.datasets.video_utils import ( get_video_info, reencode_video, ) -from tests.fixtures.constants import DUMMY_VIDEO_INFO +from tests.fixtures.constants import ( + DUMMY_DEPTH_FEATURES, + DUMMY_DEPTH_KEY, + DUMMY_DEPTH_VIDEO_INFO_FULL, + DUMMY_VIDEO_FEATURES, + DUMMY_VIDEO_INFO, + DUMMY_VIDEO_KEY, +) +from tests.fixtures.dataset_factories import add_frames # Per-codec skip markers — validation tests only fire when the codec is available @@ -48,19 +56,74 @@ def _require_encoder(vcodec: str) -> pytest.MarkDecorator: require_libsvtav1 = _require_encoder("libsvtav1") require_h264 = _require_encoder("h264") +require_hevc = _require_encoder("hevc") require_videotoolbox = _require_encoder("h264_videotoolbox") require_nvenc = _require_encoder("h264_nvenc") require_vaapi = _require_encoder("h264_vaapi") require_qsv = _require_encoder("h264_qsv") -# ─── VideoEncoderConfig / codec options ────────────────────────────── +TEST_ARTIFACTS_DIR = Path(__file__).parent.parent / "artifacts" / "encoded_videos" + + +def _write_color_frames(imgs_dir: Path, num_frames: int = 4, height: int = 64, width: int = 96) -> None: + imgs_dir.mkdir(parents=True, exist_ok=True) + for i in range(num_frames): + arr = np.random.randint(0, 256, (height, width, 3), dtype=np.uint8) + write_image(arr, imgs_dir / f"frame-{i:06d}.png") + + +def _write_depth_frames(imgs_dir: Path, num_frames: int = 4, height: int = 64, width: int = 96) -> None: + """Write synthetic uint16 depth TIFFs (millimetres) for depth encoder tests. + + Uses a smooth linear ramp + per-frame offset (not white noise) so HEVC Main 12 + on ``gray12le`` compresses well. Values span ~100 mm to 10 m, covering most + of the default ``[DEPTH_MIN, DEPTH_MAX]`` metres range after + ``quantize_depth(input_unit="auto"="mm")``. + """ + imgs_dir.mkdir(parents=True, exist_ok=True) + base = np.linspace(100.0, 10_000.0, height * width, dtype=np.float32).reshape(height, width) + for i in range(num_frames): + arr = (base + 50.0 * i).clip(0, 65535).astype(np.uint16) + write_image(arr, imgs_dir / f"frame-{i:06d}.tiff") + + +def _encode_video( + path: Path, + num_frames: int = 4, + fps: int = 30, + cfg: VideoEncoderConfig | None = None, + depth: bool = False, +) -> Path: + """Write synthetic frames to a temp dir and encode them to ``path``. + + ``depth=False`` writes uint8 RGB PNG noise and encodes with ``cfg`` + (defaulting to the library default). ``depth=True`` writes synthetic uint16 + depth TIFFs and encodes with ``cfg`` or a default :class:`DepthEncoderConfig` + (HEVC Main 12 / ``gray12le``). + """ + imgs_dir = path.parent / f"imgs_{path.stem}" + if depth: + _write_depth_frames(imgs_dir, num_frames=num_frames) + cfg = cfg or DepthEncoderConfig() + else: + _write_color_frames(imgs_dir, num_frames=num_frames) + encode_video_frames(imgs_dir, path, fps=fps, video_encoder=cfg, overwrite=True) + return path + + +def _read_feature_info(dataset: LeRobotDataset, key: str = DUMMY_VIDEO_KEY) -> dict: + info = json.loads((dataset.root / INFO_PATH).read_text()) + return info["features"][key]["info"] + + +# ─── RGBEncoderConfig / codec options ────────────────────────────── class TestCodecOptions: @require_libsvtav1 def test_libsvtav1_defaults(self): - cfg = VideoEncoderConfig() + cfg = RGBEncoderConfig() opts = cfg.get_codec_options() assert opts["g"] == 2 assert opts["crf"] == 30 @@ -68,12 +131,12 @@ class TestCodecOptions: @require_libsvtav1 def test_libsvtav1_custom_preset(self): - cfg = VideoEncoderConfig(preset=8) + cfg = RGBEncoderConfig(preset=8) assert cfg.get_codec_options()["preset"] == 8 @require_h264 def test_h264_options(self): - cfg = VideoEncoderConfig(vcodec="h264", g=10, crf=23, preset=None) + cfg = RGBEncoderConfig(vcodec="h264", g=10, crf=23, preset=None) opts = cfg.get_codec_options() assert opts["g"] == 10 assert opts["crf"] == 23 @@ -81,120 +144,120 @@ class TestCodecOptions: @require_videotoolbox def test_videotoolbox_options(self): - cfg = VideoEncoderConfig(vcodec="h264_videotoolbox", g=2, crf=30, preset=None) + cfg = RGBEncoderConfig(vcodec="h264_videotoolbox", g=2, crf=30, preset=None) opts = cfg.get_codec_options() assert opts["g"] == 2 assert opts["q:v"] == 40 assert "crf" not in opts - @_require_encoder("h264_nvenc") + @require_nvenc def test_nvenc_options(self): - cfg = VideoEncoderConfig(vcodec="h264_nvenc", g=2, crf=25, preset=None) + cfg = RGBEncoderConfig(vcodec="h264_nvenc", g=2, crf=25, preset=None) opts = cfg.get_codec_options() assert opts["rc"] == 0 assert opts["qp"] == 25 assert "crf" not in opts assert opts["g"] == 2 - @_require_encoder("h264_vaapi") + @require_vaapi def test_vaapi_options(self): - cfg = VideoEncoderConfig(vcodec="h264_vaapi", crf=28, preset=None) + cfg = RGBEncoderConfig(vcodec="h264_vaapi", crf=28, preset=None) assert cfg.get_codec_options()["qp"] == 28 - @_require_encoder("h264_qsv") + @require_qsv def test_qsv_options(self): - cfg = VideoEncoderConfig(vcodec="h264_qsv", crf=25, preset=None) + cfg = RGBEncoderConfig(vcodec="h264_qsv", crf=25, preset=None) assert cfg.get_codec_options()["global_quality"] == 25 @require_h264 def test_no_g_no_crf(self): - cfg = VideoEncoderConfig(vcodec="h264", g=None, crf=None, preset=None) + cfg = RGBEncoderConfig(vcodec="h264", g=None, crf=None, preset=None) opts = cfg.get_codec_options() assert "g" not in opts assert "crf" not in opts @require_libsvtav1 def test_encoder_threads_libsvtav1(self): - cfg = VideoEncoderConfig(fast_decode=0) + cfg = RGBEncoderConfig(fast_decode=0) opts = cfg.get_codec_options(encoder_threads=4) assert "lp=4" in opts.get("svtav1-params", "") @require_h264 def test_encoder_threads_h264(self): - cfg = VideoEncoderConfig(vcodec="h264", preset=None) + cfg = RGBEncoderConfig(vcodec="h264", preset=None) assert cfg.get_codec_options(encoder_threads=2)["threads"] == 2 @require_libsvtav1 def test_fast_decode_libsvtav1(self): - cfg = VideoEncoderConfig(fast_decode=1) + cfg = RGBEncoderConfig(fast_decode=1) opts = cfg.get_codec_options() assert "fast-decode=1" in opts.get("svtav1-params", "") @require_libsvtav1 def test_libsvtav1_fast_decode_clamped_to_svt_range(self): """Out-of-range fast_decode is clamped to [0, 2] in svtav1-params (SVT-AV1 FastDecode).""" - cfg = VideoEncoderConfig(fast_decode=100) + cfg = RGBEncoderConfig(fast_decode=100) assert "fast-decode=2" in cfg.get_codec_options().get("svtav1-params", "") - cfg_neg = VideoEncoderConfig(fast_decode=-5) + cfg_neg = RGBEncoderConfig(fast_decode=-5) assert "fast-decode=0" in cfg_neg.get_codec_options().get("svtav1-params", "") @require_h264 def test_fast_decode_h264(self): - cfg = VideoEncoderConfig(vcodec="h264", fast_decode=1, preset=None) + cfg = RGBEncoderConfig(vcodec="h264", fast_decode=1, preset=None) assert cfg.get_codec_options()["tune"] == "fastdecode" @require_libsvtav1 def test_pix_fmt_unsupported_raises(self): """Passing an unsupported pix_fmt is a hard error.""" with pytest.raises(ValueError, match="pix_fmt"): - VideoEncoderConfig(pix_fmt="yuv444p") # libsvtav1 only supports yuv420p variants + RGBEncoderConfig(pix_fmt="yuv444p") # libsvtav1 only supports yuv420p variants @require_libsvtav1 @require_h264 def test_preset_default_behaviour(self): """Empty constructor picks preset=12 (libsvtav1 path); other codecs stay None.""" - assert VideoEncoderConfig().preset == 12 - assert VideoEncoderConfig(vcodec="libsvtav1").preset == 12 - assert VideoEncoderConfig(vcodec="h264").preset is None - assert VideoEncoderConfig(vcodec="h264", preset=None).preset is None + assert RGBEncoderConfig().preset == 12 + assert RGBEncoderConfig(vcodec="libsvtav1").preset == 12 + assert RGBEncoderConfig(vcodec="h264").preset is None + assert RGBEncoderConfig(vcodec="h264", preset=None).preset is None @require_h264 def test_preset_string_on_h264(self): """h264 accepts string presets and forwards them to FFmpeg.""" - cfg = VideoEncoderConfig(vcodec="h264", preset="slow") + cfg = RGBEncoderConfig(vcodec="h264", preset="slow") assert cfg.get_codec_options()["preset"] == "slow" @require_videotoolbox def test_preset_on_videotoolbox_not_set(self): """videotoolbox has no preset option at all.""" - cfg = VideoEncoderConfig(vcodec="h264_videotoolbox", preset="slow") + cfg = RGBEncoderConfig(vcodec="h264_videotoolbox", preset="slow") assert "preset" not in cfg.get_codec_options() @require_libsvtav1 def test_libsvtav1_preset_out_of_range_raises(self): """libsvtav1 preset must sit in [-2, 13] as exposed by PyAV.""" with pytest.raises(ValueError, match="out of range"): - VideoEncoderConfig(vcodec="libsvtav1", preset=100) + RGBEncoderConfig(vcodec="libsvtav1", preset=100) with pytest.raises(ValueError, match="out of range"): - VideoEncoderConfig(vcodec="libsvtav1", preset=-3) + RGBEncoderConfig(vcodec="libsvtav1", preset=-3) @require_libsvtav1 def test_libsvtav1_crf_out_of_range_raises(self): """libsvtav1 crf must sit in [0, 63].""" with pytest.raises(ValueError, match="crf.*out of range"): - VideoEncoderConfig(vcodec="libsvtav1", crf=64) + RGBEncoderConfig(vcodec="libsvtav1", crf=64) @require_libsvtav1 def test_libsvtav1_crf_rejects_python_float(self): """libsvtav1 exposes ``crf`` as an INT AVOption; Python float must not pass validation.""" with pytest.raises(ValueError, match="float values are not allowed"): - VideoEncoderConfig(vcodec="libsvtav1", crf=2.5) + RGBEncoderConfig(vcodec="libsvtav1", crf=2.5) @require_libsvtav1 def test_libsvtav1_extra_crf_rejects_fractional_string(self): """INT options reject fractional values even when supplied only via ``extra_options``.""" with pytest.raises(ValueError, match="float values are not allowed"): - VideoEncoderConfig( + RGBEncoderConfig( vcodec="libsvtav1", crf=None, extra_options={"crf": "2.5"}, @@ -203,7 +266,7 @@ class TestCodecOptions: @require_libsvtav1 def test_libsvtav1_extra_crf_rejects_float(self): with pytest.raises(ValueError, match="float values are not allowed"): - VideoEncoderConfig( + RGBEncoderConfig( vcodec="libsvtav1", crf=None, extra_options={"crf": 2.5}, @@ -212,13 +275,13 @@ class TestCodecOptions: @require_h264 def test_h264_crf_accepts_float_and_int(self): """x264 exposes crf as a FLOAT option, so both int and float are accepted.""" - assert VideoEncoderConfig(vcodec="h264", crf=23).get_codec_options()["crf"] == 23 - assert VideoEncoderConfig(vcodec="h264", crf=23.5).get_codec_options()["crf"] == 23.5 + assert RGBEncoderConfig(vcodec="h264", crf=23).get_codec_options()["crf"] == 23 + assert RGBEncoderConfig(vcodec="h264", crf=23.5).get_codec_options()["crf"] == 23.5 @require_libsvtav1 def test_validate_is_rerunnable(self): """After mutating a field, validate() re-checks and surfaces new issues.""" - cfg = VideoEncoderConfig(vcodec="libsvtav1") + cfg = RGBEncoderConfig(vcodec="libsvtav1") cfg.preset = 100 # now out of range with pytest.raises(ValueError, match="out of range"): cfg.validate() @@ -227,58 +290,58 @@ class TestCodecOptions: class TestExtraOptions: @require_libsvtav1 def test_default_is_empty_dict(self): - cfg = VideoEncoderConfig() + cfg = RGBEncoderConfig() assert cfg.extra_options == {} @require_libsvtav1 def test_unknown_key_passes_through(self): """Keys not published as AVOptions are forwarded to FFmpeg.""" - cfg = VideoEncoderConfig(extra_options={"totally_made_up_option": "value"}) + cfg = RGBEncoderConfig(extra_options={"totally_made_up_option": "value"}) assert cfg.extra_options == {"totally_made_up_option": "value"} @require_libsvtav1 def test_numeric_value_in_range_ok(self): """libsvtav1 exposes ``qp`` as INT in [0, 63].""" - cfg = VideoEncoderConfig(extra_options={"qp": 30}) + cfg = RGBEncoderConfig(extra_options={"qp": 30}) assert cfg.extra_options == {"qp": 30} @require_libsvtav1 def test_numeric_out_of_range_raises(self): with pytest.raises(ValueError, match=r"qp=.*out of range"): - VideoEncoderConfig(extra_options={"qp": 999}) + RGBEncoderConfig(extra_options={"qp": 999}) @require_libsvtav1 def test_numeric_string_accepted_in_range(self): """Numeric strings are accepted for numeric options (mirrors FFmpeg).""" - cfg = VideoEncoderConfig(extra_options={"qp": "18"}) + cfg = RGBEncoderConfig(extra_options={"qp": "18"}) assert cfg.extra_options == {"qp": "18"} @require_libsvtav1 def test_numeric_string_out_of_range_raises(self): with pytest.raises(ValueError, match=r"qp=.*out of range"): - VideoEncoderConfig(extra_options={"qp": "999"}) + RGBEncoderConfig(extra_options={"qp": "999"}) @require_libsvtav1 def test_non_numeric_string_on_numeric_option_raises(self): with pytest.raises(ValueError, match=r"qp=.*not numeric"): - VideoEncoderConfig(extra_options={"qp": "medium"}) + RGBEncoderConfig(extra_options={"qp": "medium"}) @require_libsvtav1 def test_bool_on_numeric_option_raises(self): """``bool`` is explicitly rejected for numeric options.""" with pytest.raises(ValueError, match=r"qp=.*not numeric"): - VideoEncoderConfig(extra_options={"qp": True}) + RGBEncoderConfig(extra_options={"qp": True}) @require_h264 def test_string_option_passes_through_unchecked(self): """String-typed AVOptions are NOT enum-checked (too many accept freeform).""" - cfg = VideoEncoderConfig(vcodec="h264", preset=None, extra_options={"tune": "some-future-tune"}) + cfg = RGBEncoderConfig(vcodec="h264", preset=None, extra_options={"tune": "some-future-tune"}) assert cfg.extra_options == {"tune": "some-future-tune"} @require_libsvtav1 def test_merged_into_codec_options_and_stringified(self): """Typed merge by default; ``as_strings=True`` matches FFmpeg option dict.""" - cfg = VideoEncoderConfig(extra_options={"qp": 20}) + cfg = RGBEncoderConfig(extra_options={"qp": 20}) opts = cfg.get_codec_options() assert opts["qp"] == 20 assert isinstance(opts["qp"], int) @@ -287,25 +350,25 @@ class TestExtraOptions: @require_libsvtav1 def test_structured_fields_win_on_collision(self): """A colliding extra_options key is discarded; the structured field wins.""" - cfg = VideoEncoderConfig(crf=30, extra_options={"crf": 18}) + cfg = RGBEncoderConfig(crf=30, extra_options={"crf": 18}) assert cfg.get_codec_options()["crf"] == 30 class TestEncoderDetection: @require_h264 def test_explicit_codec_kept_when_available(self): - cfg = VideoEncoderConfig(vcodec="h264") + cfg = RGBEncoderConfig(vcodec="h264") assert cfg.vcodec == "h264" @require_videotoolbox def test_auto_picks_videotoolbox_when_available(self): """``h264_videotoolbox`` sits at the top of ``HW_VIDEO_CODECS`` so it wins when present.""" - cfg = VideoEncoderConfig(vcodec="auto") + cfg = RGBEncoderConfig(vcodec="auto") assert cfg.vcodec == "h264_videotoolbox" def test_invalid_codec_raises(self): with pytest.raises(ValueError, match="Invalid vcodec"): - VideoEncoderConfig(vcodec="not_a_real_codec") + RGBEncoderConfig(vcodec="not_a_real_codec") def test_hw_encoder_names_listed_as_valid(self): assert "auto" in VALID_VIDEO_CODECS @@ -313,59 +376,6 @@ class TestEncoderDetection: assert "h264_nvenc" in VALID_VIDEO_CODECS -TEST_ARTIFACTS_DIR = Path(__file__).parent.parent / "artifacts" / "encoded_videos" - -# Default video feature set used by persistence tests. -VIDEO_FEATURES = { - "observation.images.cam": { - "dtype": "video", - "shape": (64, 96, 3), - "names": ["height", "width", "channels"], - }, - "action": {"dtype": "float32", "shape": (2,), "names": ["a", "b"]}, -} -VIDEO_KEY = "observation.images.cam" - - -def _write_frames(imgs_dir: Path, num_frames: int = 4, height: int = 64, width: int = 96) -> None: - imgs_dir.mkdir(parents=True, exist_ok=True) - for i in range(num_frames): - arr = np.random.randint(0, 256, (height, width, 3), dtype=np.uint8) - write_image(arr, imgs_dir / f"frame-{i:06d}.png") - - -def _encode_video( - path: Path, num_frames: int = 4, fps: int = 30, cfg: VideoEncoderConfig | None = None -) -> Path: - imgs_dir = path.parent / f"imgs_{path.stem}" - _write_frames(imgs_dir, num_frames=num_frames) - encode_video_frames(imgs_dir, path, fps=fps, camera_encoder=cfg, overwrite=True) - return path - - -def _read_feature_info(dataset: LeRobotDataset) -> dict: - info = json.loads((dataset.root / INFO_PATH).read_text()) - return info["features"][VIDEO_KEY]["info"] - - -def _add_frames(dataset: LeRobotDataset, num_frames: int, video_keys: list[str] | None = None) -> None: - from lerobot.utils.constants import DEFAULT_FEATURES - - if video_keys is None: - video_keys = dataset.meta.video_keys - for _ in range(num_frames): - frame: dict = {"task": "test"} - for key, ft in dataset.meta.features.items(): - if key in DEFAULT_FEATURES: - continue - shape = ft["shape"] - if key in video_keys: - frame[key] = np.random.randint(0, 256, shape, dtype=np.uint8) - else: - frame[key] = np.zeros(shape, dtype=np.float32) - dataset.add_frame(frame) - - class TestGetVideoInfo: def test_returns_all_stream_fields(self): info = get_video_info(TEST_ARTIFACTS_DIR / "clip_4frames.mp4") @@ -375,7 +385,7 @@ class TestGetVideoInfo: assert info["video.pix_fmt"] == "yuv420p" assert info["video.fps"] == 30 assert info["video.channels"] == 3 - assert info["video.is_depth_map"] is False + assert info["is_depth_map"] is False assert info["has_audio"] is False assert "video.g" not in info assert "video.crf" not in info @@ -383,9 +393,9 @@ class TestGetVideoInfo: @require_libsvtav1 def test_merges_encoder_config_as_video_prefixed_entries(self): - cfg = VideoEncoderConfig(vcodec="libsvtav1", g=2, crf=30, preset=12) + cfg = RGBEncoderConfig(vcodec="libsvtav1", g=2, crf=30, preset=12) - info = get_video_info(TEST_ARTIFACTS_DIR / "clip_4frames.mp4", camera_encoder=cfg) + info = get_video_info(TEST_ARTIFACTS_DIR / "clip_4frames.mp4", video_encoder=cfg) assert info["video.g"] == 2 assert info["video.crf"] == 30 @@ -396,13 +406,18 @@ class TestGetVideoInfo: @require_libsvtav1 def test_stream_derived_keys_take_precedence_over_config(self): - cfg = VideoEncoderConfig(vcodec="libsvtav1", pix_fmt="yuv420p") + cfg = RGBEncoderConfig(vcodec="libsvtav1", pix_fmt="yuv420p") - info = get_video_info(TEST_ARTIFACTS_DIR / "clip_4frames.mp4", camera_encoder=cfg) + info = get_video_info(TEST_ARTIFACTS_DIR / "clip_4frames.mp4", video_encoder=cfg) assert info["video.codec"] # populated from stream, not from config's vcodec assert info["video.pix_fmt"] == "yuv420p" + def test_depth_encoder_config_sets_is_depth_map_true(self): + """A ``DepthEncoderConfig`` causes ``get_video_info`` to mark the stream as depth.""" + info = get_video_info(TEST_ARTIFACTS_DIR / "clip_4frames.mp4", video_encoder=DepthEncoderConfig()) + assert info["is_depth_map"] is True + class TestEncodeVideoFrames: @require_libsvtav1 @@ -434,7 +449,7 @@ class TestEncodeVideoFrames: def test_overwrite_false_skips_existing_file(self, tmp_path): imgs_dir = tmp_path / "imgs" - _write_frames(imgs_dir) + _write_color_frames(imgs_dir) video_path = tmp_path / "out.mp4" sentinel = b"pre-existing content" video_path.write_bytes(sentinel) @@ -446,7 +461,7 @@ class TestEncodeVideoFrames: @require_libsvtav1 def test_overwrite_true_replaces_existing_file(self, tmp_path): imgs_dir = tmp_path / "imgs" - _write_frames(imgs_dir) + _write_color_frames(imgs_dir) video_path = tmp_path / "out.mp4" video_path.write_bytes(b"stale content") @@ -458,10 +473,10 @@ class TestEncodeVideoFrames: @require_libsvtav1 def test_custom_encoder_config_fields_stored_in_info(self, tmp_path): """All stream-derived and encoder config fields are present after encoding.""" - cfg = VideoEncoderConfig(vcodec="libsvtav1", g=4, crf=25, preset=10) + cfg = RGBEncoderConfig(vcodec="libsvtav1", g=4, crf=25, preset=10) video_path = _encode_video(tmp_path / "out.mp4", num_frames=4, fps=30, cfg=cfg) - info = get_video_info(video_path, camera_encoder=cfg) + info = get_video_info(video_path, video_encoder=cfg) # Stream-derived assert info["video.height"] == 64 @@ -470,7 +485,7 @@ class TestEncodeVideoFrames: assert info["video.codec"] == "av1" assert info["video.pix_fmt"] == "yuv420p" assert info["video.fps"] == 30 - assert info["video.is_depth_map"] is False + assert info["is_depth_map"] is False assert info["has_audio"] is False # Encoder config assert info["video.g"] == 4 @@ -487,15 +502,15 @@ class TestReencodeVideo: def test_reencode_video(self, tmp_path): src = TEST_ARTIFACTS_DIR / "clip_4frames.mp4" out = tmp_path / "reencoded.mp4" - cfg = VideoEncoderConfig(vcodec="h264", g=6, crf=23, pix_fmt="yuv444p") - reencode_video(src, out, camera_encoder=cfg, overwrite=True) + cfg = RGBEncoderConfig(vcodec="h264", g=6, crf=23, pix_fmt="yuv444p") + reencode_video(src, out, video_encoder=cfg, overwrite=True) assert out.exists() with av.open(str(out)) as container: n_frames = sum(1 for _ in container.decode(video=0)) assert n_frames == 4 - info = get_video_info(out, camera_encoder=cfg) + info = get_video_info(out, video_encoder=cfg) assert info["video.codec"] == "h264" assert info["video.pix_fmt"] == "yuv444p" assert info["video.height"] == 64 @@ -508,8 +523,8 @@ class TestReencodeVideo: def test_reencode_video_trim_window(self, tmp_path): src = TEST_ARTIFACTS_DIR / "clip_6frames.mp4" out = tmp_path / "trim_window.mp4" - cfg = VideoEncoderConfig(vcodec="h264") - reencode_video(src, out, camera_encoder=cfg, start_time_s=0.05, end_time_s=0.12, overwrite=True) + cfg = RGBEncoderConfig(vcodec="h264") + reencode_video(src, out, video_encoder=cfg, start_time_s=0.05, end_time_s=0.12, overwrite=True) with av.open(str(out)) as container: frames = list(container.decode(video=0)) @@ -578,12 +593,12 @@ class TestEncoderConfigPersistence: @require_libsvtav1 def test_first_episode_save_persists_encoder_config(self, tmp_path, empty_lerobot_dataset_factory): - cfg = VideoEncoderConfig(vcodec="libsvtav1", g=2, crf=30, preset=12) + cfg = RGBEncoderConfig(vcodec="libsvtav1", g=2, crf=30, preset=12) dataset = empty_lerobot_dataset_factory( - root=tmp_path / "ds", features=VIDEO_FEATURES, use_videos=True, camera_encoder=cfg + root=tmp_path / "ds", features=DUMMY_VIDEO_FEATURES, use_videos=True, rgb_encoder=cfg ) - _add_frames(dataset, num_frames=4) + add_frames(dataset, num_frames=4) dataset.save_episode() dataset.finalize() @@ -601,16 +616,16 @@ class TestEncoderConfigPersistence: @require_libsvtav1 def test_second_episode_does_not_overwrite_encoder_fields(self, tmp_path, empty_lerobot_dataset_factory): - cfg = VideoEncoderConfig(vcodec="libsvtav1", g=2, crf=30, preset=12) + cfg = RGBEncoderConfig(vcodec="libsvtav1", g=2, crf=30, preset=12) dataset = empty_lerobot_dataset_factory( - root=tmp_path / "ds", features=VIDEO_FEATURES, use_videos=True, camera_encoder=cfg + root=tmp_path / "ds", features=DUMMY_VIDEO_FEATURES, use_videos=True, rgb_encoder=cfg ) - _add_frames(dataset, num_frames=4) + add_frames(dataset, num_frames=4) dataset.save_episode() first_info = dict(_read_feature_info(dataset)) - _add_frames(dataset, num_frames=4) + add_frames(dataset, num_frames=4) dataset.save_episode() dataset.finalize() @@ -618,13 +633,13 @@ class TestEncoderConfigPersistence: class TestFromVideoInfo: - """``VideoEncoderConfig.from_video_info`` reconstructs an encoder config + """``RGBEncoderConfig.from_video_info`` reconstructs an encoder config from the ``video.*`` keys persisted in a dataset's ``info.json``. """ @require_libsvtav1 def test_reconstructs_from_dummy_video_info(self): - cfg = VideoEncoderConfig.from_video_info(DUMMY_VIDEO_INFO) + cfg = RGBEncoderConfig.from_video_info(DUMMY_VIDEO_INFO) # Canonical stream codec ``"av1"`` is aliased to the encoder name. assert cfg.vcodec == "libsvtav1" @@ -636,4 +651,220 @@ class TestFromVideoInfo: assert cfg.video_backend == DUMMY_VIDEO_INFO["video.video_backend"] # ``{}`` placeholder (typical after a merge with disagreeing sources) # must not leak into the reconstructed config. - assert cfg.extra_options == VideoEncoderConfig().extra_options + assert cfg.extra_options == RGBEncoderConfig().extra_options + + +# ─── Depth-specific encoding tests ──────────────────────────────────── + + +class TestEncodeDepthVideoFrames: + """Depth mirror of :class:`TestEncodeVideoFrames`. + + Exercises ``encode_video_frames`` end-to-end through + :class:`DepthEncoderConfig` (HEVC Main 12 / ``gray12le``) on synthetic + uint16 depth TIFFs. + """ + + @require_hevc + def test_produces_readable_file(self, tmp_path): + video_path = _encode_video(tmp_path / "out.mp4", depth=True) + + assert video_path.exists() + info = get_video_info(video_path, video_encoder=DepthEncoderConfig()) + assert info["video.height"] == 64 + assert info["video.width"] == 96 + assert info["video.codec"] == "hevc" + assert info["video.pix_fmt"] == "gray12le" + assert info["video.channels"] == 1 + assert info["is_depth_map"] is True + + @require_hevc + def test_frame_count_and_duration_match_input(self, tmp_path): + num_frames = 10 + fps = 30 + video_path = _encode_video(tmp_path / "out.mp4", num_frames=num_frames, fps=fps, depth=True) + + with av.open(str(video_path)) as container: + stream = container.streams.video[0] + actual_frames = sum(1 for _ in container.decode(stream)) + duration = ( + float(stream.duration * stream.time_base) + if stream.duration is not None + else float(container.duration / av.time_base) + ) + + assert actual_frames == num_frames + assert abs(duration - num_frames / fps) < 0.1 + + def test_overwrite_false_skips_existing_file(self, tmp_path): + """Codec-agnostic: file-system semantics must hold even without an HEVC encoder.""" + imgs_dir = tmp_path / "imgs" + _write_depth_frames(imgs_dir) + video_path = tmp_path / "out.mp4" + sentinel = b"pre-existing depth content" + video_path.write_bytes(sentinel) + + encode_video_frames(imgs_dir, video_path, fps=30, video_encoder=DepthEncoderConfig(), overwrite=False) + + assert video_path.read_bytes() == sentinel + + @require_hevc + def test_overwrite_true_replaces_existing_file(self, tmp_path): + imgs_dir = tmp_path / "imgs" + _write_depth_frames(imgs_dir) + video_path = tmp_path / "out.mp4" + video_path.write_bytes(b"stale content") + + encode_video_frames(imgs_dir, video_path, fps=30, video_encoder=DepthEncoderConfig(), overwrite=True) + + info = get_video_info(video_path, video_encoder=DepthEncoderConfig()) + assert info["video.height"] == 64 + assert info["video.pix_fmt"] == "gray12le" + assert info["is_depth_map"] is True + + @require_hevc + def test_custom_encoder_config_fields_stored_in_info(self, tmp_path): + """All stream-derived and depth-encoder config fields are present after encoding.""" + cfg = DepthEncoderConfig( + vcodec="hevc", + pix_fmt="gray12le", + g=4, + crf=25, + extra_options={}, + depth_min=0.05, + depth_max=8.0, + shift=2.5, + use_log=False, + ) + video_path = _encode_video(tmp_path / "out.mp4", num_frames=4, fps=30, cfg=cfg, depth=True) + + info = get_video_info(video_path, video_encoder=cfg) + + # Stream-derived + assert info["video.height"] == 64 + assert info["video.width"] == 96 + assert info["video.channels"] == 1 + assert info["video.codec"] == "hevc" + assert info["video.pix_fmt"] == "gray12le" + assert info["video.fps"] == 30 + assert info["is_depth_map"] is True + assert info["has_audio"] is False + # Base encoder config + assert info["video.g"] == 4 + assert info["video.crf"] == 25 + assert info["video.fast_decode"] == 0 + assert info["video.video_backend"] == "pyav" + assert info["video.extra_options"] == {} + # Depth-specific tuning + assert info["video.depth_min"] == 0.05 + assert info["video.depth_max"] == 8.0 + assert info["video.shift"] == 2.5 + assert info["video.use_log"] is False + + +class TestDepthEncoderConfigPersistence: + """Depth mirror of :class:`TestEncoderConfigPersistence`. + + ``DepthEncoderConfig`` must be stored as ``video.`` entries + (including the depth-specific ``depth_min`` / ``depth_max`` / ``shift`` / + ``use_log``) under ``info["features"][]["info"]`` when the + first episode is saved. + """ + + @require_hevc + def test_first_episode_save_persists_depth_encoder_config(self, tmp_path, empty_lerobot_dataset_factory): + cfg = DepthEncoderConfig( + vcodec="hevc", + pix_fmt="gray12le", + g=2, + crf=30, + extra_options={}, + depth_min=0.05, + depth_max=8.0, + shift=2.5, + use_log=False, + ) + dataset = empty_lerobot_dataset_factory( + root=tmp_path / "ds", features=DUMMY_DEPTH_FEATURES, use_videos=True, depth_encoder=cfg + ) + + add_frames(dataset, num_frames=4) + dataset.save_episode() + dataset.finalize() + + info = _read_feature_info(dataset, key=DUMMY_DEPTH_KEY) + + # Stream-derived + assert info["video.height"] == 64 + assert info["video.width"] == 96 + assert info["video.fps"] == 30 + assert info["video.codec"] == "hevc" + assert info["video.pix_fmt"] == "gray12le" + assert info["is_depth_map"] is True + # Base encoder config + assert info["video.g"] == 2 + assert info["video.crf"] == 30 + assert info["video.fast_decode"] == 0 + assert info["video.video_backend"] == "pyav" + assert info["video.extra_options"] == {} + # Depth-specific tuning + assert info["video.depth_min"] == 0.05 + assert info["video.depth_max"] == 8.0 + assert info["video.shift"] == 2.5 + assert info["video.use_log"] is False + + @require_hevc + def test_second_episode_does_not_overwrite_depth_encoder_fields( + self, tmp_path, empty_lerobot_dataset_factory + ): + cfg = DepthEncoderConfig( + vcodec="hevc", + pix_fmt="gray12le", + g=2, + crf=30, + depth_min=0.05, + depth_max=8.0, + shift=2.5, + use_log=False, + ) + dataset = empty_lerobot_dataset_factory( + root=tmp_path / "ds", features=DUMMY_DEPTH_FEATURES, use_videos=True, depth_encoder=cfg + ) + + add_frames(dataset, num_frames=4) + dataset.save_episode() + first_info = dict(_read_feature_info(dataset, key=DUMMY_DEPTH_KEY)) + + add_frames(dataset, num_frames=4) + dataset.save_episode() + dataset.finalize() + + assert _read_feature_info(dataset, key=DUMMY_DEPTH_KEY) == first_info + + +class TestDepthFromVideoInfo: + """``DepthEncoderConfig.from_video_info`` reconstructs a depth encoder + config from the ``video.*`` keys persisted in a dataset's ``info.json``. + + Depth mirror of :class:`TestFromVideoInfo`. + """ + + @require_hevc + def test_reconstructs_from_dummy_depth_video_info(self): + cfg = DepthEncoderConfig.from_video_info(DUMMY_DEPTH_VIDEO_INFO_FULL) + + # No alias for ``"hevc"``; the canonical stream codec is reused as-is. + assert cfg.vcodec == "hevc" + assert cfg.pix_fmt == DUMMY_DEPTH_VIDEO_INFO_FULL["video.pix_fmt"] + assert cfg.g == DUMMY_DEPTH_VIDEO_INFO_FULL["video.g"] + assert cfg.crf == DUMMY_DEPTH_VIDEO_INFO_FULL["video.crf"] + assert cfg.fast_decode == DUMMY_DEPTH_VIDEO_INFO_FULL["video.fast_decode"] + assert cfg.video_backend == DUMMY_DEPTH_VIDEO_INFO_FULL["video.video_backend"] + # ``{}`` placeholder (typical after a merge with disagreeing sources) + # must not leak into the reconstructed config. + assert cfg.extra_options == DepthEncoderConfig().extra_options + # Depth-specific tuning round-trips through ``info.json``. + assert cfg.depth_min == DUMMY_DEPTH_VIDEO_INFO_FULL["video.depth_min"] + assert cfg.depth_max == DUMMY_DEPTH_VIDEO_INFO_FULL["video.depth_max"] + assert cfg.shift == DUMMY_DEPTH_VIDEO_INFO_FULL["video.shift"] + assert cfg.use_log == DUMMY_DEPTH_VIDEO_INFO_FULL["video.use_log"] diff --git a/tests/fixtures/constants.py b/tests/fixtures/constants.py index 4d578b503..d6f4f8ae5 100644 --- a/tests/fixtures/constants.py +++ b/tests/fixtures/constants.py @@ -39,12 +39,56 @@ DUMMY_VIDEO_INFO = { "video.crf": 30, "video.preset": 12, "video.fast_decode": 0, - "video.is_depth_map": False, + "is_depth_map": False, "has_audio": False, } DUMMY_CAMERA_FEATURES = { "laptop": {"shape": (64, 96, 3), "names": ["height", "width", "channels"], "info": DUMMY_VIDEO_INFO}, "phone": {"shape": (64, 96, 3), "names": ["height", "width", "channels"], "info": DUMMY_VIDEO_INFO}, } +DUMMY_DEPTH_VIDEO_INFO = { + **DUMMY_VIDEO_INFO, + "is_depth_map": True, +} +DUMMY_DEPTH_VIDEO_INFO_FULL = { + **{k: v for k, v in DUMMY_VIDEO_INFO.items() if k != "video.preset"}, + "video.codec": "hevc", + "video.pix_fmt": "gray12le", + "is_depth_map": True, + "video.depth_min": 0.05, + "video.depth_max": 8.0, + "video.shift": 2.5, + "video.use_log": True, +} +DUMMY_DEPTH_CAMERA_FEATURES = { + "laptop_depth": { + "shape": (64, 96, 1), + "names": ["height", "width", "channels"], + "info": DUMMY_DEPTH_VIDEO_INFO, + }, +} +DUMMY_CAMERA_FEATURES_WITH_DEPTH = {**DUMMY_CAMERA_FEATURES, **DUMMY_DEPTH_CAMERA_FEATURES} DUMMY_CHW = (3, 96, 128) DUMMY_HWC = (96, 128, 3) + +# Default video feature set used by video-encoding persistence tests. +DUMMY_VIDEO_FEATURES = { + "observation.images.cam": { + "dtype": "video", + "shape": (64, 96, 3), + "names": ["height", "width", "channels"], + }, + "action": {"dtype": "float32", "shape": (2,), "names": ["a", "b"]}, +} +DUMMY_VIDEO_KEY = "observation.images.cam" + +DUMMY_DEPTH_FEATURES = { + "observation.images.depth": { + "dtype": "video", + "shape": (64, 96, 1), + "names": ["height", "width", "channels"], + "info": {"is_depth_map": True}, + }, + "action": {"dtype": "float32", "shape": (2,), "names": ["a", "b"]}, +} +DUMMY_DEPTH_KEY = "observation.images.depth" diff --git a/tests/fixtures/dataset_factories.py b/tests/fixtures/dataset_factories.py index 2f4d41ff8..100922f9c 100644 --- a/tests/fixtures/dataset_factories.py +++ b/tests/fixtures/dataset_factories.py @@ -49,6 +49,39 @@ from tests.fixtures.constants import ( ) +def add_frames(dataset: LeRobotDataset, num_frames: int) -> None: + """Append ``num_frames`` synthetic frames to ``dataset``. + + Generates per-feature payloads from ``dataset.meta``: uint16 depth ramps for + keys in ``dataset.meta.depth_keys``, uint8 random noise for video/image keys, + and float32 zeros for everything else. ``DEFAULT_FEATURES`` (timestamp, + frame_index, ...) are auto-populated by ``add_frame`` and skipped here. + """ + video_keys = dataset.meta.video_keys + depth_keys = dataset.meta.depth_keys + # Smooth gradient base reused per (H, W) to keep depth frames cheap to + # encode (HEVC Main 12 hates white noise). + _depth_base_cache: dict[tuple[int, int], np.ndarray] = {} + for i in range(num_frames): + frame: dict = {"task": "test"} + for key, ft in dataset.meta.features.items(): + if key in DEFAULT_FEATURES: + continue + shape = ft["shape"] + if key in depth_keys: + h, w, _ = shape + base = _depth_base_cache.setdefault( + (h, w), + np.linspace(100.0, 10_000.0, h * w, dtype=np.float32).reshape(h, w, 1), + ) + frame[key] = (base + 50.0 * i).clip(0, 65535).astype(np.uint16) + elif key in video_keys: + frame[key] = np.random.randint(0, 256, shape, dtype=np.uint8) + else: + frame[key] = np.zeros(shape, dtype=np.float32) + dataset.add_frame(frame) + + class LeRobotDatasetFactory(Protocol): def __call__(self, *args, **kwargs) -> LeRobotDataset: ... @@ -485,10 +518,14 @@ def lerobot_dataset_factory( hf_dataset: datasets.Dataset | None = None, data_files_size_in_mb: float = DEFAULT_DATA_FILE_SIZE_IN_MB, chunks_size: int = DEFAULT_CHUNK_SIZE, + camera_features: dict | None = None, **kwargs, ) -> LeRobotDataset: # Instantiate objects if info is None: + info_kwargs = {} + if camera_features is not None: + info_kwargs["camera_features"] = camera_features info = info_factory( total_episodes=total_episodes, total_frames=total_frames, @@ -496,6 +533,7 @@ def lerobot_dataset_factory( use_videos=use_videos, data_files_size_in_mb=data_files_size_in_mb, chunks_size=chunks_size, + **info_kwargs, ) if stats is None: stats = stats_factory(features=info.features) diff --git a/tests/jobs/__init__.py b/tests/jobs/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/jobs/conftest.py b/tests/jobs/conftest.py new file mode 100644 index 000000000..419d2f83f --- /dev/null +++ b/tests/jobs/conftest.py @@ -0,0 +1,17 @@ +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Importing concrete policy configs registers their draccus `--policy.type` +# choices (e.g. "act") so tests can parse them. +from lerobot.policies.act.configuration_act import ACTConfig # noqa: F401 diff --git a/tests/jobs/test_dataset.py b/tests/jobs/test_dataset.py new file mode 100644 index 000000000..1f8b9b836 --- /dev/null +++ b/tests/jobs/test_dataset.py @@ -0,0 +1,66 @@ +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest.mock import MagicMock + +import pytest + +pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])") + +from lerobot.jobs.dataset import ensure_dataset_available + + +def _api_with_dataset(exists: bool): + api = MagicMock() + api.repo_exists.return_value = exists + return api + + +def _make_local_cache(tmp_path, repo_id: str) -> None: + """Create the minimal local-cache layout that ensure_dataset_available checks.""" + info = tmp_path / repo_id / "meta" / "info.json" + info.parent.mkdir(parents=True) + info.write_text("{}") + + +# Branch 1: dataset already on Hub → no push, no error (pod downloads by repo_id). +def test_dataset_already_on_hub_is_noop(): + api = _api_with_dataset(True) + assert ensure_dataset_available("user/ds", api=api) is None + api.repo_exists.assert_called_once_with("user/ds", repo_type="dataset") + + +# Branch 2: not on Hub but present locally → always push privately. +def test_dataset_local_only_uploads_privately(tmp_path, monkeypatch): + monkeypatch.setattr("lerobot.jobs.dataset.HF_LEROBOT_HOME", tmp_path) + _make_local_cache(tmp_path, "user/ds") + + api = _api_with_dataset(False) + mock_ds_cls = MagicMock() + monkeypatch.setattr("lerobot.jobs.dataset.LeRobotDataset", mock_ds_cls) + + assert ensure_dataset_available("user/ds", api=api, tags=["lerobot", "lelab"]) is None + + mock_ds_cls.assert_called_once_with("user/ds") + mock_ds_cls.return_value.push_to_hub.assert_called_once_with(private=True, tags=["lerobot", "lelab"]) + + +# Branch 3: not on Hub, NOT in local cache → RuntimeError. +def test_dataset_neither_on_hub_nor_local_raises(tmp_path, monkeypatch): + monkeypatch.setattr("lerobot.jobs.dataset.HF_LEROBOT_HOME", tmp_path) + # tmp_path is empty — no local cache. + + api = _api_with_dataset(False) + with pytest.raises(RuntimeError, match="not in the local cache"): + ensure_dataset_available("user/ds", api=api) diff --git a/tests/jobs/test_hf.py b/tests/jobs/test_hf.py new file mode 100644 index 000000000..3b275cb95 --- /dev/null +++ b/tests/jobs/test_hf.py @@ -0,0 +1,493 @@ +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import datetime as dt +import json +import threading +from types import SimpleNamespace + +import draccus +import httpx +import pytest + +pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])") + +from lerobot.configs.train import TrainPipelineConfig +from lerobot.jobs.hf import ( + _pod_forwarded_args, + _poll_until_done, + build_remote_config_file, + build_repo_id, + resolve_job_tags, + resolve_wandb_api_key, + submit_to_hf, +) + + +def test_resolve_job_tags_always_includes_lerobot_and_dedups(): + assert resolve_job_tags(None) == ["lerobot"] + assert resolve_job_tags([]) == ["lerobot"] + assert resolve_job_tags(["lelab"]) == ["lerobot", "lelab"] + # lerobot isn't duplicated if passed explicitly; order is stable. + assert resolve_job_tags(["lelab", "lerobot", "lelab"]) == ["lerobot", "lelab"] + + +def _fake_inspect(stage_value, *, as_enum=True): + # huggingface_hub returns `stage` as an enum (with `.value`) in some versions and a plain str in others. + stage = SimpleNamespace(value=stage_value) if as_enum else stage_value + return lambda job_id: SimpleNamespace(status=SimpleNamespace(stage=stage)) + + +@pytest.mark.parametrize("as_enum", [True, False], ids=["enum_stage", "str_stage"]) +def test_poll_until_done_returns_terminal_stage(monkeypatch, as_enum): + monkeypatch.setattr("lerobot.jobs.hf.inspect_job", _fake_inspect("COMPLETED", as_enum=as_enum)) + done = threading.Event() + assert _poll_until_done("j", done, poll_interval=0.01) == "COMPLETED" + assert done.is_set() + + +def test_poll_until_done_exits_when_done_already_set(monkeypatch): + # Non-terminal forever; with done pre-set the loop must not block and returns None. + monkeypatch.setattr("lerobot.jobs.hf.inspect_job", _fake_inspect("RUNNING")) + done = threading.Event() + done.set() + assert _poll_until_done("j", done, poll_interval=0.01) is None + + +def test_poll_until_done_gives_up_after_repeated_network_failures(monkeypatch): + monkeypatch.setattr( + "lerobot.jobs.hf.inspect_job", lambda job_id: (_ for _ in ()).throw(httpx.ConnectError("boom")) + ) + done = threading.Event() + result = _poll_until_done("j", done, poll_interval=0.001, max_failures=3) + assert result is None + assert done.is_set() + + +def test_poll_until_done_propagates_programming_errors(monkeypatch): + """A bug (e.g. TypeError) must surface, not be silently retried as a transient failure.""" + monkeypatch.setattr("lerobot.jobs.hf.inspect_job", lambda job_id: (_ for _ in ()).throw(TypeError("bug"))) + done = threading.Event() + with pytest.raises(TypeError): + _poll_until_done("j", done, poll_interval=0.001, max_failures=3) + + +def test_resolve_wandb_key_from_env(monkeypatch): + monkeypatch.setenv("WANDB_API_KEY", "abc123") + assert resolve_wandb_api_key() == "abc123" + + +def test_resolve_wandb_key_missing(monkeypatch, tmp_path): + monkeypatch.delenv("WANDB_API_KEY", raising=False) + monkeypatch.setenv("HOME", str(tmp_path)) # no ~/.netrc here + monkeypatch.setattr("netrc.netrc", lambda *a, **k: (_ for _ in ()).throw(FileNotFoundError())) + assert resolve_wandb_api_key() is None + + +def test_resolve_wandb_key_from_netrc(monkeypatch): + # No env var → fall back to the wandb credentials in ~/.netrc. + monkeypatch.delenv("WANDB_API_KEY", raising=False) + + class _FakeNetrc: + def authenticators(self, host): + assert host == "api.wandb.ai" + return ("login", "account", "netrc-secret") + + monkeypatch.setattr("netrc.netrc", lambda *a, **k: _FakeNetrc()) + assert resolve_wandb_api_key() == "netrc-secret" + + +def test_resolve_wandb_key_netrc_without_wandb_entry(monkeypatch): + # ~/.netrc exists but has no api.wandb.ai entry → None. + monkeypatch.delenv("WANDB_API_KEY", raising=False) + + class _FakeNetrc: + def authenticators(self, host): + return None + + monkeypatch.setattr("netrc.netrc", lambda *a, **k: _FakeNetrc()) + assert resolve_wandb_api_key() is None + + +def test_build_repo_id_sanitizes_and_timestamps(): + now = dt.datetime(2026, 6, 19, 10, 22, 3) + assert build_repo_id("alice", "act", now) == "alice/act_2026-06-19_10-22-03" + # Runs of illegal characters collapse to a single dash; edges are trimmed. + assert build_repo_id("alice", "my cool/run!!", now) == "alice/my-cool-run_2026-06-19_10-22-03" + # A name with nothing usable falls back to "train". + assert build_repo_id("alice", "///", now) == "alice/train_2026-06-19_10-22-03" + + +def test_pod_forwarded_args_drops_host_only_flags(): + """User overrides are replayed on the pod, minus flags that only make sense on the submitter. + + `--dataset.root` is a host-local path the pod can't read, so it must be dropped in both the + `--name=value` and `--name value` forms; unrelated overrides are forwarded untouched. + """ + argv = [ + "--config_path=u/d", + "--dataset.root=/local/data", + "--dataset.root", + "/other/local/data", + "--policy.repo_id=u/keep", + "--steps=10", + "--job.target=a10g-small", + ] + forwarded = _pod_forwarded_args( + argv, + drop_names=("--config_path", "--policy.repo_id", "--policy.push_to_hub", "--dataset.root"), + drop_prefixes=("--job.",), + ) + assert forwarded == ["--steps=10"] + + +def _minimal_cfg(): + return draccus.parse( + TrainPipelineConfig, + args=["--dataset.repo_id", "u/d", "--policy.type", "act", "--job.target", "a10g-small"], + ) + + +def test_validate_skips_repo_id_check_for_remote(): + """Remote runs auto-assign repo_id in submit_to_hf, so validate() must not demand it up front.""" + cfg = _minimal_cfg() # remote target, push_to_hub default True, no explicit repo_id + assert cfg.policy.repo_id is None + cfg.validate() # must not raise + + +def test_validate_requires_repo_id_for_local_push(): + """Local runs that push to the Hub still need an explicit repo_id.""" + cfg = draccus.parse( + TrainPipelineConfig, + args=["--dataset.repo_id", "u/d", "--policy.type", "act"], + ) + with pytest.raises(ValueError, match="repo_id"): + cfg.validate() + + +def test_build_remote_config_applies_overrides(tmp_path): + cfg = _minimal_cfg() + dest = tmp_path / "train_config.json" + out = build_remote_config_file(cfg, "u/run", dest) + assert out == dest + data = json.loads(dest.read_text()) + # `job` is client-only orchestration and must be stripped for the pod. + assert "job" not in data + # save_checkpoint_to_hub defaults off → omitted so older images accept the config. + assert "save_checkpoint_to_hub" not in data + assert data["policy"]["push_to_hub"] is True + assert data["policy"]["repo_id"] == "u/run" + assert data["policy"]["device"] is None # pod auto-detects its GPU + assert data["dataset"]["root"] is None # pod resolves the dataset by repo_id + # the caller's cfg must be left untouched (function works on a deep copy) + assert cfg.job.target == "a10g-small" + assert cfg.save_checkpoint_to_hub is False + + +def test_build_remote_config_includes_checkpoint_flag_when_enabled(tmp_path): + cfg = draccus.parse( + TrainPipelineConfig, + args=[ + "--dataset.repo_id", + "u/d", + "--policy.type", + "act", + "--job.target", + "a10g-small", + "--save_checkpoint_to_hub", + "true", + ], + ) + dest = tmp_path / "train_config.json" + build_remote_config_file(cfg, "u/run", dest) + data = json.loads(dest.read_text()) + # explicitly enabled → kept in the config (requires a matching trainer image). + assert data["save_checkpoint_to_hub"] is True + assert "job" not in data + + +def test_build_remote_config_merges_tags_into_policy(tmp_path): + cfg = _minimal_cfg() + dest = tmp_path / "train_config.json" + build_remote_config_file(cfg, "u/run", dest, tags=["lerobot", "lelab"]) + data = json.loads(dest.read_text()) + # tags propagate to the model the pod pushes. + assert data["policy"]["tags"] == ["lerobot", "lelab"] + + +def test_build_remote_config_merges_tags_without_duplicating(tmp_path): + cfg = _minimal_cfg() + cfg.policy.tags = ["existing", "lerobot"] + dest = tmp_path / "train_config.json" + build_remote_config_file(cfg, "u/run", dest, tags=["lerobot", "lelab"]) + data = json.loads(dest.read_text()) + # pre-existing policy tags are kept; only genuinely-new tags are appended (no dup "lerobot"). + assert data["policy"]["tags"] == ["existing", "lerobot", "lelab"] + + +def test_submit_requires_login(monkeypatch): + monkeypatch.setattr("lerobot.jobs.hf.get_token", lambda: None) + cfg = draccus.parse( + TrainPipelineConfig, + args=["--dataset.repo_id", "u/d", "--policy.type", "act", "--job.target", "a10g-small"], + ) + with pytest.raises(RuntimeError, match="hf auth login"): + submit_to_hf(cfg) + + +def test_submit_passes_validation_and_submits(monkeypatch): + """A type-based policy with no explicit repo_id is auto-assigned one and submitted.""" + from unittest.mock import MagicMock + + # Patch get_token + monkeypatch.setattr("lerobot.jobs.hf.get_token", lambda: "tok") + + # Patch HfApi so whoami returns alice + class FakeHfApi: + def __init__(self, token=None): + pass + + def whoami(self, token=None): + return {"name": "alice"} + + monkeypatch.setattr("lerobot.jobs.hf.HfApi", FakeHfApi) + + # ensure_dataset_available returns None; patch it out so no Hub access happens + # (hf.py imports it at module level, so patch it on lerobot.jobs.hf). + monkeypatch.setattr("lerobot.jobs.hf.ensure_dataset_available", lambda *a, **kw: None) + + # Patch _stage_config_on_hub to skip network + monkeypatch.setattr( + "lerobot.jobs.hf._stage_config_on_hub", + lambda cfg, repo_id, token, tags=None: repo_id, + ) + + # Patch run_job to return a fake job + fake_job = MagicMock() + fake_job.id = "job-123" + run_job_calls = [] + + def fake_run_job(**kwargs): + run_job_calls.append(kwargs) + return fake_job + + monkeypatch.setattr("lerobot.jobs.hf.run_job", fake_run_job) + + cfg = draccus.parse( + TrainPipelineConfig, + args=[ + "--dataset.repo_id", + "u/d", + "--policy.type", + "act", + "--job.target", + "a10g-small", + "--job.detach", + "true", + ], + ) + + # Must NOT raise (pre-fix this raised ValueError about missing repo_id) + submit_to_hf(cfg) + + assert len(run_job_calls) == 1, "run_job should have been called exactly once" + assert cfg.policy.repo_id is not None + assert cfg.policy.repo_id.startswith("alice/") + call = run_job_calls[0] + # The pod runs `lerobot-train --config_path=` on the requested flavor/image. + assert call["command"][0] == "lerobot-train" + assert call["command"][1].startswith("--config_path=") + assert call["flavor"] == "a10g-small" + assert call["image"] == "huggingface/lerobot-gpu:latest" + # The Hub token is forwarded so the pod can pull the (possibly private) dataset. + assert call["secrets"]["HF_TOKEN"] == "tok" + # Every job carries the lerobot tag as a queryable label. + assert call["labels"].get("lerobot") == "true" + + +def test_submit_rejects_reward_model_training(monkeypatch): + """Remote training only supports policies; reward-model runs fail fast with a clear error.""" + monkeypatch.setattr("lerobot.jobs.hf.get_token", lambda: "tok") + + class FakeHfApi: + def __init__(self, token=None): + pass + + def whoami(self, token=None): + return {"name": "alice"} + + monkeypatch.setattr("lerobot.jobs.hf.HfApi", FakeHfApi) + + cfg = _minimal_cfg() + cfg.reward_model = SimpleNamespace(type="reward") # marks this as reward-model training + monkeypatch.setattr(cfg, "validate", lambda: None) # skip pretrained-path resolution + + with pytest.raises(ValueError, match="reward model"): + submit_to_hf(cfg) + + +@pytest.mark.timeout(15) +def test_submit_returns_when_job_completes(monkeypatch): + """Non-detach path must RETURN (not hang) once the job reaches a terminal stage.""" + from types import SimpleNamespace + + monkeypatch.setattr("lerobot.jobs.hf.get_token", lambda: "tok") + + class FakeHfApi: + def __init__(self, token=None): + pass + + def whoami(self, token=None): + return {"name": "alice"} + + monkeypatch.setattr("lerobot.jobs.hf.HfApi", FakeHfApi) + monkeypatch.setattr("lerobot.jobs.hf.ensure_dataset_available", lambda *a, **kw: None) + monkeypatch.setattr( + "lerobot.jobs.hf._stage_config_on_hub", lambda cfg, repo_id, token, tags=None: repo_id + ) + monkeypatch.setattr("lerobot.jobs.hf.run_job", lambda **kw: SimpleNamespace(id="job-1", url="http://x")) + # Job is already COMPLETED on the first poll. + monkeypatch.setattr( + "lerobot.jobs.hf.inspect_job", + lambda job_id: SimpleNamespace( + status=SimpleNamespace(stage=SimpleNamespace(value="COMPLETED"), message=None) + ), + ) + # Log stream ends immediately. + monkeypatch.setattr("lerobot.jobs.hf.fetch_job_logs", lambda job_id, follow=True: iter(())) + + cfg = draccus.parse( + TrainPipelineConfig, + args=["--dataset.repo_id", "u/d", "--policy.type", "act", "--job.target", "a10g-small"], + ) + # Runs in the pytest main thread (signal handler install requires it); the + # @timeout marker fails the test instead of hanging if it regresses. + submit_to_hf(cfg) + + +@pytest.mark.timeout(15) +def test_submit_returns_on_model_pushed_marker(monkeypatch): + """Finish when the model-pushed log appears, even if the job stage never flips.""" + from types import SimpleNamespace + + monkeypatch.setattr("lerobot.jobs.hf.get_token", lambda: "tok") + + class FakeHfApi: + def __init__(self, token=None): + pass + + def whoami(self, token=None): + return {"name": "alice"} + + monkeypatch.setattr("lerobot.jobs.hf.HfApi", FakeHfApi) + monkeypatch.setattr("lerobot.jobs.hf.ensure_dataset_available", lambda *a, **kw: None) + monkeypatch.setattr( + "lerobot.jobs.hf._stage_config_on_hub", lambda cfg, repo_id, token, tags=None: repo_id + ) + monkeypatch.setattr("lerobot.jobs.hf.run_job", lambda **kw: SimpleNamespace(id="job-1", url="http://x")) + # Job stays RUNNING forever — only the log marker can end the command. + monkeypatch.setattr( + "lerobot.jobs.hf.inspect_job", + lambda job_id: SimpleNamespace( + status=SimpleNamespace(stage=SimpleNamespace(value="RUNNING"), message=None) + ), + ) + pushed_line = "INFO Model pushed to https://huggingface.co/alice/myrun" + monkeypatch.setattr("lerobot.jobs.hf.fetch_job_logs", lambda job_id, follow=True: iter([pushed_line])) + + cfg = draccus.parse( + TrainPipelineConfig, + args=[ + "--dataset.repo_id", + "u/d", + "--policy.type", + "act", + "--policy.repo_id", + "alice/myrun", + "--job.target", + "a10g-small", + ], + ) + # Must return via the model-pushed marker despite the perpetual RUNNING stage. + submit_to_hf(cfg) + + +def test_submit_raises_when_wandb_enabled_without_key(monkeypatch): + """wandb.enable with no key reachable anywhere fails fast, before submitting.""" + + monkeypatch.setattr("lerobot.jobs.hf.get_token", lambda: "tok") + + class FakeHfApi: + def __init__(self, token=None): + pass + + def whoami(self, token=None): + return {"name": "alice"} + + monkeypatch.setattr("lerobot.jobs.hf.HfApi", FakeHfApi) + monkeypatch.setattr("lerobot.jobs.hf.resolve_wandb_api_key", lambda: None) + + cfg = draccus.parse( + TrainPipelineConfig, + args=[ + "--dataset.repo_id", + "u/d", + "--policy.type", + "act", + "--job.target", + "a10g-small", + "--wandb.enable", + "true", + ], + ) + with pytest.raises(ValueError, match="WANDB_API_KEY"): + submit_to_hf(cfg) + + +@pytest.mark.timeout(15) +def test_submit_raises_when_job_ends_in_error(monkeypatch): + """A terminal non-COMPLETED stage with no model-pushed marker must raise with the status.""" + from types import SimpleNamespace + + monkeypatch.setattr("lerobot.jobs.hf.get_token", lambda: "tok") + + class FakeHfApi: + def __init__(self, token=None): + pass + + def whoami(self, token=None): + return {"name": "alice"} + + monkeypatch.setattr("lerobot.jobs.hf.HfApi", FakeHfApi) + monkeypatch.setattr("lerobot.jobs.hf.ensure_dataset_available", lambda *a, **kw: None) + monkeypatch.setattr( + "lerobot.jobs.hf._stage_config_on_hub", lambda cfg, repo_id, token, tags=None: repo_id + ) + monkeypatch.setattr("lerobot.jobs.hf.run_job", lambda **kw: SimpleNamespace(id="job-1", url="http://x")) + # Job fails: a terminal ERROR stage carrying the platform's status message. + monkeypatch.setattr( + "lerobot.jobs.hf.inspect_job", + lambda job_id: SimpleNamespace( + status=SimpleNamespace(stage=SimpleNamespace(value="ERROR"), message="Job timeout") + ), + ) + # Logs end without the model-pushed marker. + monkeypatch.setattr("lerobot.jobs.hf.fetch_job_logs", lambda job_id, follow=True: iter(())) + + cfg = draccus.parse( + TrainPipelineConfig, + args=["--dataset.repo_id", "u/d", "--policy.type", "act", "--job.target", "a10g-small"], + ) + with pytest.raises(RuntimeError, match=r"stage=ERROR \(Job timeout\)"): + submit_to_hf(cfg) diff --git a/tests/jobs/test_job_config.py b/tests/jobs/test_job_config.py new file mode 100644 index 000000000..20760fb18 --- /dev/null +++ b/tests/jobs/test_job_config.py @@ -0,0 +1,64 @@ +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import draccus +import pytest + +from lerobot.configs import JobConfig +from lerobot.configs.train import TrainPipelineConfig + + +def test_jobconfig_defaults_are_local(): + cfg = JobConfig() + assert cfg.target is None + assert cfg.is_remote is False + assert cfg.image == "huggingface/lerobot-gpu:latest" + assert cfg.timeout == "2d" + assert cfg.detach is False + + +def test_jobconfig_local_string_is_not_remote(): + assert JobConfig(target="local").is_remote is False + + +def test_jobconfig_flavor_is_remote(): + assert JobConfig(target="a10g-small").is_remote is True + + +def test_train_config_parses_job_target(): + parsed = draccus.parse( + TrainPipelineConfig, + args=["--dataset.repo_id", "u/d", "--policy.type", "act", "--job.target", "a10g-small"], + ) + assert parsed.job.target == "a10g-small" + assert parsed.job.is_remote is True + assert parsed.save_checkpoint_to_hub is False + + +def test_save_checkpoint_to_hub_requires_repo_id(): + cfg = draccus.parse( + TrainPipelineConfig, + args=[ + "--dataset.repo_id", + "u/d", + "--policy.type", + "act", + "--policy.push_to_hub", + "false", + "--save_checkpoint_to_hub", + "true", + ], + ) + with pytest.raises(ValueError, match="requires --policy.repo_id"): + cfg.validate() diff --git a/tests/optim/test_optimizers.py b/tests/optim/test_optimizers.py index d18565562..5b480f70d 100644 --- a/tests/optim/test_optimizers.py +++ b/tests/optim/test_optimizers.py @@ -20,6 +20,7 @@ from lerobot.optim.optimizers import ( MultiAdamConfig, SGDConfig, load_optimizer_state, + load_optimizer_state_dict, save_optimizer_state, ) from lerobot.utils.constants import ( @@ -65,6 +66,44 @@ def test_save_and_load_optimizer_state(model_params, optimizer, tmp_path): torch.testing.assert_close(optimizer.state_dict(), loaded_optimizer.state_dict()) +def test_save_and_load_fsdp_optimizer_state_dict_roundtrip(tmp_path): + """The FSDP full optimizer state dict is keyed by parameter FQNs (dotted strings), not the + integer indices of the single-GPU path. Verify it survives the safetensors save -> read + round-trip used by the FSDP save/resume path (save_optimizer_state(optim_state_dict=...) then + load_optimizer_state_dict), which the flatten/unflatten "/" separator must not corrupt.""" + full_osd = { + "state": { + "model.layers.0.weight": { + "step": torch.tensor(3.0), + "exp_avg": torch.randn(4, 4), + "exp_avg_sq": torch.randn(4, 4), + }, + "model.layers.0.bias": { + "step": torch.tensor(3.0), + "exp_avg": torch.randn(4), + "exp_avg_sq": torch.randn(4), + }, + }, + "param_groups": [ + {"lr": 1e-4, "betas": [0.9, 0.999], "eps": 1e-8, "weight_decay": 0.0, "params": [0, 1]} + ], + } + + save_optimizer_state( + torch.optim.Adam([torch.nn.Parameter(torch.randn(1))]), tmp_path, optim_state_dict=full_osd + ) + assert (tmp_path / OPTIMIZER_STATE).is_file() + assert (tmp_path / OPTIMIZER_PARAM_GROUPS).is_file() + + loaded = load_optimizer_state_dict(tmp_path) + # FQN keys must be preserved verbatim (not int-cast, not split on their dots). + assert set(loaded["state"].keys()) == set(full_osd["state"].keys()) + for fqn, sub in full_osd["state"].items(): + for k, v in sub.items(): + torch.testing.assert_close(loaded["state"][fqn][k], v) + assert loaded["param_groups"] == full_osd["param_groups"] + + @pytest.fixture def base_params_dict(): return { diff --git a/tests/policies/molmoact2/test_molmoact2.py b/tests/policies/molmoact2/test_molmoact2.py index 3631bcc9b..095b73180 100644 --- a/tests/policies/molmoact2/test_molmoact2.py +++ b/tests/policies/molmoact2/test_molmoact2.py @@ -1,5 +1,3 @@ -#!/usr/bin/env python - # Copyright 2026 The Allen Institute for Artificial Intelligence and The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -35,24 +33,27 @@ pytest.importorskip("scipy") from lerobot.configs import FeatureType, NormalizationMode, PolicyFeature from lerobot.policies import get_policy_class, make_policy_config from lerobot.policies.molmoact2 import ( - configuration_molmoact2 as molmoact2_config, modeling_molmoact2 as molmoact2_modeling, processor_molmoact2 as molmoact2_processor, ) -from lerobot.policies.molmoact2.configuration_molmoact2 import ( - MolmoAct2Config, - MolmoAct2CosineDecayWithWarmupSchedulerConfig, - infer_molmoact2_max_sequence_length, +from lerobot.policies.molmoact2.configuration_molmoact2 import MolmoAct2Config +from lerobot.policies.molmoact2.modeling_molmoact2 import ( + MolmoAct2Policy, + _apply_action_chunk_padding_mask, + _apply_action_dim_padding_mask, + _combine_rollout_seeds, ) -from lerobot.policies.molmoact2.modeling_molmoact2 import MolmoAct2Policy from lerobot.policies.molmoact2.processor_molmoact2 import ( + MolmoAct2ActionFrameTransformStep, MolmoAct2ClampNormalizedProcessorStep, MolmoAct2MaskedNormalizerProcessorStep, MolmoAct2MaskedUnnormalizerProcessorStep, MolmoAct2PackInputsProcessorStep, + MolmoAct2StateFrameTransformStep, _add_gripper_masks_to_stats, _build_discrete_state_string, _normalize_question_text, + infer_molmoact2_max_sequence_length, make_molmoact2_pre_post_processors, ) from lerobot.policies.rtc.configuration_rtc import RTCConfig @@ -71,34 +72,38 @@ def test_molmoact2_policy_registration(): assert cfg.per_episode_seed is False assert cfg.eval_seed is None assert cfg.normalize_language is True - assert cfg.get_scheduler_preset().num_decay_steps is None + assert cfg.get_scheduler_preset().num_decay_steps == 100_000 assert cfg.action_delta_indices == list(range(cfg.chunk_size)) assert get_policy_class("molmoact2") is MolmoAct2Policy def test_molmoact2_checkpoint_download_ignores_remote_python(monkeypatch): + import huggingface_hub + download_kwargs = {} def fake_snapshot_download(**kwargs): download_kwargs.update(kwargs) return "/tmp/downloaded-molmoact2" - monkeypatch.setattr(molmoact2_config, "snapshot_download", fake_snapshot_download) + monkeypatch.setattr(huggingface_hub, "snapshot_download", fake_snapshot_download) - checkpoint_location = molmoact2_config._resolve_checkpoint_location("allenai/MolmoAct2") + checkpoint_location = molmoact2_modeling._resolve_checkpoint_location("allenai/MolmoAct2") assert checkpoint_location == "/tmp/downloaded-molmoact2" assert download_kwargs["ignore_patterns"] == ["*.py", "*.pyc", "__pycache__/*"] -def test_molmoact2_scheduler_decay_steps_auto_match_training_steps(): +def test_molmoact2_scheduler_auto_scales_to_training_steps(): + from lerobot.optim import CosineDecayWithWarmupSchedulerConfig + param = torch.nn.Parameter(torch.ones(())) optimizer = torch.optim.AdamW([param], lr=0.001) - config = MolmoAct2CosineDecayWithWarmupSchedulerConfig( + config = CosineDecayWithWarmupSchedulerConfig( peak_lr=0.01, decay_lr=0.001, num_warmup_steps=10, - num_decay_steps=None, + num_decay_steps=100_000, ) scheduler = config.build(optimizer, num_training_steps=100) @@ -123,9 +128,7 @@ def test_molmoact2_rollout_generator_uses_eval_seed_per_task(): batch_size=3, device=torch.device("cpu"), ) - expected_first = torch.Generator().manual_seed( - MolmoAct2Policy._combine_rollout_seeds(first_seed=1000, batch_size=3) - ) + expected_first = torch.Generator().manual_seed(_combine_rollout_seeds(first_seed=1000, batch_size=3)) assert torch.allclose(torch.rand(4, generator=first), torch.rand(4, generator=expected_first)) policy.reset() @@ -134,9 +137,7 @@ def test_molmoact2_rollout_generator_uses_eval_seed_per_task(): batch_size=3, device=torch.device("cpu"), ) - expected_second = torch.Generator().manual_seed( - MolmoAct2Policy._combine_rollout_seeds(first_seed=1003, batch_size=3) - ) + expected_second = torch.Generator().manual_seed(_combine_rollout_seeds(first_seed=1003, batch_size=3)) assert torch.allclose(torch.rand(4, generator=second), torch.rand(4, generator=expected_second)) policy.reset() @@ -145,9 +146,7 @@ def test_molmoact2_rollout_generator_uses_eval_seed_per_task(): batch_size=3, device=torch.device("cpu"), ) - expected_new_task = torch.Generator().manual_seed( - MolmoAct2Policy._combine_rollout_seeds(first_seed=1000, batch_size=3) - ) + expected_new_task = torch.Generator().manual_seed(_combine_rollout_seeds(first_seed=1000, batch_size=3)) assert torch.allclose(torch.rand(4, generator=new_task), torch.rand(4, generator=expected_new_task)) @@ -537,36 +536,26 @@ def test_train_action_expert_only_requires_continuous_action_mode(): def test_molmoact2_sequence_length_is_inferred_from_fixed_token_budget(): - cfg = MolmoAct2Config( - action_mode="both", - chunk_size=10, - n_action_steps=10, - image_keys=["observation.images.image", "observation.images.wrist_image"], - input_features={OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(8,))}, - output_features={ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(7,))}, - ) - - assert cfg.max_sequence_length is None - assert cfg.inferred_max_sequence_length() == 640 - assert cfg.inferred_max_sequence_length(include_discrete_action=False) == 576 assert ( infer_molmoact2_max_sequence_length( - num_images=2, - state_dim=8, - action_dim=7, - action_horizon=30, - include_discrete_action=True, + num_images=2, state_dim=8, action_dim=7, action_horizon=10, include_discrete_action=True + ) + == 640 + ) + assert ( + infer_molmoact2_max_sequence_length( + num_images=2, state_dim=8, action_dim=7, action_horizon=10, include_discrete_action=False + ) + == 576 + ) + assert ( + infer_molmoact2_max_sequence_length( + num_images=2, state_dim=8, action_dim=7, action_horizon=30, include_discrete_action=True ) == 768 ) -def test_molmoact2_sequence_length_override_is_preserved(): - cfg = MolmoAct2Config(max_sequence_length=1024) - - assert cfg.inferred_max_sequence_length(num_images=2, state_dim=8, action_dim=7) == 1024 - - def test_train_action_expert_only_freezes_non_action_expert_params(): class DummyBackbone(torch.nn.Module): def __init__(self): @@ -939,6 +928,39 @@ def test_question_normalization_matches_release_prompt_style(): ) +def test_joint_frame_transform_round_trip(): + signs = [1.0, -1.0, 1.0, 1.0, 1.0, 1.0] + offsets = [0.0, 90.0, 90.0, 0.0, 0.0, 0.0] + original_state = torch.tensor([[10.0, -90.0, -120.0, 30.0, 0.0, -45.0]]) + + state_step = MolmoAct2StateFrameTransformStep(joint_signs=signs, joint_offsets=offsets) + action_step = MolmoAct2ActionFrameTransformStep(joint_signs=signs, joint_offsets=offsets) + + transition = { + TransitionKey.OBSERVATION: {OBS_STATE: original_state.clone()}, + } + transformed = state_step(transition) + model_state = transformed[TransitionKey.OBSERVATION][OBS_STATE] + + action_transition = {TransitionKey.ACTION: model_state.clone()} + recovered = action_step(action_transition) + recovered_state = recovered[TransitionKey.ACTION] + + assert torch.allclose(recovered_state, original_state) + + +def test_joint_frame_transform_noop_when_none(): + state_step = MolmoAct2StateFrameTransformStep(joint_signs=None, joint_offsets=None) + action_step = MolmoAct2ActionFrameTransformStep(joint_signs=None, joint_offsets=None) + state = torch.tensor([[10.0, -90.0, -120.0]]) + + state_transition = {TransitionKey.OBSERVATION: {OBS_STATE: state}} + assert state_step(state_transition) is state_transition + + action_transition = {TransitionKey.ACTION: state} + assert action_step(action_transition) is action_transition + + def test_action_padding_marks_only_real_dimensions(): step = object.__new__(MolmoAct2PackInputsProcessorStep) step.max_action_dim = 32 @@ -963,7 +985,7 @@ def test_action_dim_padding_loss_reduces_like_old_trainer(): ] ) - reduced = MolmoAct2Policy._apply_action_dim_padding_mask(loss, action_dim_is_pad) + reduced = _apply_action_dim_padding_mask(loss, action_dim_is_pad) expected = torch.stack( [ @@ -979,7 +1001,7 @@ def test_action_chunk_padding_keeps_old_mean_denominator(): loss = torch.ones(1, 2, 4, 3) action_horizon_is_pad = torch.tensor([[False, False, True, True]]) - masked = MolmoAct2Policy._apply_action_chunk_padding_mask(loss, action_horizon_is_pad) + masked = _apply_action_chunk_padding_mask(loss, action_horizon_is_pad) assert masked.mean().item() == 0.5 diff --git a/tests/policies/test_policies.py b/tests/policies/test_policies.py index e9388b3ed..285b87d4c 100644 --- a/tests/policies/test_policies.py +++ b/tests/policies/test_policies.py @@ -23,6 +23,7 @@ import torch pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])") +from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE from packaging import version from safetensors.torch import load_file @@ -300,6 +301,29 @@ def test_save_and_load_pretrained(dummy_dataset_metadata, tmp_path, policy_name: torch.testing.assert_close(list(policy.parameters()), list(loaded_policy.parameters()), rtol=0, atol=0) +def test_save_pretrained_with_state_dict(dummy_dataset_metadata, tmp_path): + """Exercise the FSDP checkpoint path: save_pretrained with a pre-gathered state_dict.""" + policy_cls = get_policy_class("act") + policy_cfg = make_policy_config("act") + features = dataset_to_policy_features(dummy_dataset_metadata.features) + policy_cfg.output_features = {key: ft for key, ft in features.items() if ft.type is FeatureType.ACTION} + policy_cfg.input_features = { + key: ft for key, ft in features.items() if key not in policy_cfg.output_features + } + policy = policy_cls(policy_cfg) + policy.to(policy_cfg.device) + + save_dir = tmp_path / "fsdp_state_dict" + policy.save_pretrained(save_dir, state_dict=policy.state_dict()) + + # A single, unsharded safetensors file (no sharded set + index). + assert (save_dir / SAFETENSORS_SINGLE_FILE).is_file() + assert not (save_dir / f"{SAFETENSORS_SINGLE_FILE}.index.json").exists() + + loaded_policy = policy_cls.from_pretrained(save_dir, config=policy_cfg) + torch.testing.assert_close(list(policy.parameters()), list(loaded_policy.parameters()), rtol=0, atol=0) + + @pytest.mark.parametrize("multikey", [True, False]) def test_multikey_construction(multikey: bool): """ diff --git a/tests/policies/vla_jepa/conftest.py b/tests/policies/vla_jepa/conftest.py index 5301b5bc7..dd40ca9ea 100644 --- a/tests/policies/vla_jepa/conftest.py +++ b/tests/policies/vla_jepa/conftest.py @@ -8,7 +8,6 @@ from types import SimpleNamespace import numpy as np import pytest import torch -from PIL import Image from torch import Tensor, nn from lerobot.configs.types import FeatureType, PolicyFeature @@ -191,7 +190,7 @@ class _FakeQwenInterface(nn.Module): def build_inputs( self, - images: list[list[Image.Image]], + images: list[list[Tensor]], instructions: list[str], action_prompt: str, embodied_prompt: str, @@ -214,12 +213,13 @@ class _FakeQwenInterface(nn.Module): } @staticmethod - def tensor_to_pil(image_tensor: Tensor) -> Image.Image: - image = image_tensor.detach().cpu() - if image.ndim == 3 and image.shape[0] in (1, 3): - image = image.permute(1, 2, 0) - image = (image.float().clamp(0, 1) * 255).to(torch.uint8).numpy() - return Image.fromarray(image) + def to_pixel_values(image_tensor: Tensor) -> Tensor: + image = image_tensor.detach().float() + if image.shape[-3] == 1: + repeats = [1] * image.ndim + repeats[-3] = 3 + image = image.repeat(*repeats) + return image class _FakeVideoEncoder(nn.Module): @@ -242,12 +242,14 @@ class _FakeVideoEncoder(nn.Module): class _FakeVideoProcessor: - def __call__(self, videos, return_tensors: str) -> dict[str, Tensor]: + def __call__(self, videos, return_tensors: str, device=None, **kwargs) -> dict[str, Tensor]: assert return_tensors == "pt" if isinstance(videos, list): pixel_values = torch.stack([torch.as_tensor(v) for v in videos]) else: pixel_values = torch.as_tensor(videos).unsqueeze(0) + if device is not None: + pixel_values = pixel_values.to(device) return {"pixel_values_videos": pixel_values} diff --git a/tests/policies/vla_jepa/test_vla_jepa.py b/tests/policies/vla_jepa/test_vla_jepa.py index 70194dd59..a3e24a660 100644 --- a/tests/policies/vla_jepa/test_vla_jepa.py +++ b/tests/policies/vla_jepa/test_vla_jepa.py @@ -211,40 +211,42 @@ def test_reset_clears_action_queue(patch_vla_jepa_external_models: None) -> None def test_prepare_model_inputs_training_format(patch_vla_jepa_external_models: None) -> None: - from PIL import Image - policy = VLAJEPAPolicy(make_config()) - examples = policy._prepare_model_inputs(make_train_batch()) + inputs = policy._prepare_model_inputs(make_train_batch()) - assert len(examples) == BATCH_SIZE - for ex in examples: - assert set(ex) >= {"image", "video", "lang", "action", "state"} - assert len(ex["image"]) == 1 and isinstance(ex["image"][0], Image.Image) - assert ex["video"].ndim == 5 and ex["video"].dtype == np.uint8 # [V,T,H,W,C] - assert ex["action"].shape == (ACTION_HORIZON, ACTION_DIM) - assert ex["state"].shape == (1, STATE_DIM) + assert set(inputs) >= {"images", "instructions", "videos", "actions", "state"} + # images: per-sample, per-view [C, H, W] float tensors (kept as a list for Qwen messages) + assert len(inputs["images"]) == BATCH_SIZE and len(inputs["images"][0]) == 1 + img = inputs["images"][0][0] + assert isinstance(img, torch.Tensor) and img.dtype == torch.float32 and img.ndim == 3 + assert len(inputs["instructions"]) == BATCH_SIZE + # videos: batched [B, V, T, C, H, W] float + assert inputs["videos"].ndim == 6 and inputs["videos"].shape[0] == BATCH_SIZE + assert inputs["videos"].dtype == torch.float32 + assert inputs["actions"].shape == (BATCH_SIZE, ACTION_HORIZON, ACTION_DIM) + assert inputs["state"].shape == (BATCH_SIZE, 1, STATE_DIM) def test_prepare_model_inputs_inference_omits_action(patch_vla_jepa_external_models: None) -> None: policy = VLAJEPAPolicy(make_config()) - for ex in policy._prepare_model_inputs(make_inference_batch()): - assert "action" not in ex - assert "image" in ex and "video" in ex and "lang" in ex + inputs = policy._prepare_model_inputs(make_inference_batch()) + assert "actions" not in inputs and "action_is_pad" not in inputs + assert {"images", "instructions", "state"} <= set(inputs) def test_prepare_model_inputs_missing_task_uses_default(patch_vla_jepa_external_models: None) -> None: policy = VLAJEPAPolicy(make_config()) batch = make_inference_batch() del batch["task"] - examples = policy._prepare_model_inputs(batch) - assert all(isinstance(ex["lang"], str) and len(ex["lang"]) > 0 for ex in examples) + instructions = policy._prepare_model_inputs(batch)["instructions"] + assert all(isinstance(s, str) and len(s) > 0 for s in instructions) def test_prepare_model_inputs_string_task_broadcast(patch_vla_jepa_external_models: None) -> None: policy = VLAJEPAPolicy(make_config()) batch = make_inference_batch() batch["task"] = "open the drawer" - assert all(ex["lang"] == "open the drawer" for ex in policy._prepare_model_inputs(batch)) + assert policy._prepare_model_inputs(batch)["instructions"] == ["open the drawer"] * BATCH_SIZE def test_prepare_model_inputs_no_state_omitted(patch_vla_jepa_external_models: None) -> None: @@ -253,7 +255,7 @@ def test_prepare_model_inputs_no_state_omitted(patch_vla_jepa_external_models: N policy = VLAJEPAPolicy(make_config()) batch = make_inference_batch() del batch[OBS_STATE] - assert all("state" not in ex for ex in policy._prepare_model_inputs(batch)) + assert "state" not in policy._prepare_model_inputs(batch) # --------------------------------------------------------------------------- @@ -446,14 +448,14 @@ def test_postprocessor_applied_after_predict_action_chunk( """ from lerobot.policies.vla_jepa.processor_vla_jepa import make_vla_jepa_pre_post_processors - raw_actions = np.zeros((BATCH_SIZE, ACTION_HORIZON, ACTION_DIM), dtype=np.float32) + raw_actions = torch.zeros((BATCH_SIZE, ACTION_HORIZON, ACTION_DIM), dtype=torch.float32) cfg = make_config() cfg.clip_normalized_actions = False cfg.binarize_gripper_action = False policy = VLAJEPAPolicy(cfg) policy.eval() - monkeypatch.setattr(policy.model, "predict_action", lambda *a, **kw: raw_actions.copy()) + monkeypatch.setattr(policy.model, "predict_action", lambda *a, **kw: raw_actions.clone()) dataset_stats = _make_dataset_stats() _, postprocessor = make_vla_jepa_pre_post_processors(cfg, dataset_stats) @@ -564,9 +566,9 @@ def test_single_view_is_duplicated_for_world_model(patch_vla_jepa_external_model original_processor = policy.model.video_processor class _CapturingProcessor: - def __call__(self, videos: list, return_tensors: str) -> dict: + def __call__(self, videos: list, return_tensors: str, **kwargs) -> dict: captured_videos.extend(videos) - return original_processor(videos=videos, return_tensors=return_tensors) + return original_processor(videos=videos, return_tensors=return_tensors, **kwargs) policy.model.video_processor = _CapturingProcessor() policy.forward(_make_multiview_train_batch(num_views=1)) @@ -587,9 +589,9 @@ def test_excess_views_trimmed_for_world_model(patch_vla_jepa_external_models: No original_processor = policy.model.video_processor class _CapturingProcessor: - def __call__(self, videos: list, return_tensors: str) -> dict: + def __call__(self, videos: list, return_tensors: str, **kwargs) -> dict: captured_videos.extend(videos) - return original_processor(videos=videos, return_tensors=return_tensors) + return original_processor(videos=videos, return_tensors=return_tensors, **kwargs) policy.model.video_processor = _CapturingProcessor() policy.forward(_make_multiview_train_batch(num_views=3)) diff --git a/tests/scripts/test_edit_dataset_parsing.py b/tests/scripts/test_edit_dataset_parsing.py index c90cffb38..22a3c1be2 100644 --- a/tests/scripts/test_edit_dataset_parsing.py +++ b/tests/scripts/test_edit_dataset_parsing.py @@ -27,6 +27,7 @@ from lerobot.scripts.lerobot_edit_dataset import ( MergeConfig, ModifyTasksConfig, OperationConfig, + ReencodeVideosConfig, RemoveFeatureConfig, SplitConfig, _validate_config, @@ -103,3 +104,47 @@ class TestOperationTypeParsing: ) resolved_name = OperationConfig.get_choice_name(type(cfg.operation)) assert resolved_name == type_name + + +class TestDepthEncoderParsing: + """Test that the depth encoder is exposed and parsed for video operations.""" + + def test_reencode_has_default_depth_encoder(self): + cfg = parse_cfg(["--repo_id", "test/repo", "--operation.type", "reencode_videos"]) + assert isinstance(cfg.operation, ReencodeVideosConfig) + # A depth encoder is configured by default so depth videos are re-encoded too. + assert cfg.operation.depth_encoder is not None + assert hasattr(cfg.operation.depth_encoder, "depth_min") + + def test_reencode_parses_depth_encoder_overrides(self): + cfg = parse_cfg( + [ + "--repo_id", + "test/repo", + "--operation.type", + "reencode_videos", + "--operation.depth_encoder.extra_options", + '{"x265-params": "lossless=1"}', + "--operation.depth_encoder.depth_max", + "12.0", + "--operation.depth_encoder.use_log", + "false", + ] + ) + assert cfg.operation.depth_encoder.extra_options == {"x265-params": "lossless=1"} + assert cfg.operation.depth_encoder.depth_max == 12.0 + assert cfg.operation.depth_encoder.use_log is False + + def test_convert_image_to_video_parses_depth_encoder_overrides(self): + cfg = parse_cfg( + [ + "--repo_id", + "test/repo", + "--operation.type", + "convert_image_to_video", + "--operation.depth_encoder.depth_min", + "0.05", + ] + ) + assert isinstance(cfg.operation, ConvertImageToVideoConfig) + assert cfg.operation.depth_encoder.depth_min == 0.05 diff --git a/tests/scripts/test_train_remote_dispatch.py b/tests/scripts/test_train_remote_dispatch.py new file mode 100644 index 000000000..50431da9e --- /dev/null +++ b/tests/scripts/test_train_remote_dispatch.py @@ -0,0 +1,67 @@ +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import sys + +import draccus +import pytest + +# Importing lerobot_train eagerly pulls in lerobot.datasets, which needs the +# `dataset` extra. The base CI tier runs without it, so skip the whole module there. +pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])") + +from lerobot.configs.train import TrainPipelineConfig # noqa: E402 +from lerobot.policies.act.configuration_act import ( + ACTConfig, # noqa: E402, F401 (registers --policy.type act) +) +from lerobot.scripts.lerobot_train import _remote_target_in_argv, train # noqa: E402 + + +def _set_argv(monkeypatch, *args): + monkeypatch.setattr(sys, "argv", ["lerobot-train", *args]) + + +def test_remote_target_detected_space_separated(monkeypatch): + _set_argv(monkeypatch, "--policy.type", "act", "--job.target", "a10g-small") + assert _remote_target_in_argv() is True + + +def test_remote_target_detected_equals(monkeypatch): + _set_argv(monkeypatch, "--job.target=t4-small") + assert _remote_target_in_argv() is True + + +def test_local_string_is_not_remote(monkeypatch): + _set_argv(monkeypatch, "--job.target", "local") + assert _remote_target_in_argv() is False + + +def test_no_target_is_not_remote(monkeypatch): + _set_argv(monkeypatch, "--policy.type", "act") + assert _remote_target_in_argv() is False + + +def test_train_dispatches_to_submit_when_remote(monkeypatch): + """A remote --job.target short-circuits train() to the HF Jobs submitter.""" + import lerobot.scripts.lerobot_train as train_module + + captured = [] + monkeypatch.setattr(train_module, "submit_to_hf", lambda cfg: captured.append(cfg) or "submitted") + cfg = draccus.parse( + TrainPipelineConfig, + args=["--dataset.repo_id", "u/d", "--policy.type", "act", "--job.target", "a10g-small"], + ) + # Returns the submitter's result and never enters the local training path. + assert train(cfg) == "submitted" + assert captured == [cfg] diff --git a/tests/training/test_multi_gpu.py b/tests/training/test_multi_gpu.py index 638dc3131..1a75d21cb 100644 --- a/tests/training/test_multi_gpu.py +++ b/tests/training/test_multi_gpu.py @@ -58,7 +58,46 @@ def download_dataset(repo_id, episodes): print(f"Dataset {repo_id} downloaded successfully") -def run_accelerate_training(config_args, num_processes=4, temp_dir=None): +def _write_multi_gpu_config(f, num_processes): + f.write("compute_environment: LOCAL_MACHINE\n") + f.write("distributed_type: MULTI_GPU\n") + f.write("mixed_precision: 'no'\n") + f.write(f"num_processes: {num_processes}\n") + f.write("use_cpu: false\n") + f.write("gpu_ids: all\n") + f.write("downcast_bf16: 'no'\n") + f.write("machine_rank: 0\n") + f.write("main_training_function: main\n") + f.write("num_machines: 1\n") + f.write("rdzv_backend: static\n") + f.write("same_network: true\n") + + +def _write_fsdp_config(f, num_processes): + # FSDP1 with FULL_SHARD (ZeRO-3-equivalent) and FULL_STATE_DICT, matching + # docs/source/multi_gpu_training.mdx. ACT's repeated transformer blocks are the wrap units; + # fsdp_use_orig_params is required because LeRobot builds the optimizer before prepare(). + f.write("compute_environment: LOCAL_MACHINE\n") + f.write("distributed_type: FSDP\n") + f.write("mixed_precision: 'no'\n") + f.write(f"num_processes: {num_processes}\n") + f.write("use_cpu: false\n") + f.write("gpu_ids: all\n") + f.write("machine_rank: 0\n") + f.write("main_training_function: main\n") + f.write("num_machines: 1\n") + f.write("rdzv_backend: static\n") + f.write("same_network: true\n") + f.write("fsdp_config:\n") + f.write(" fsdp_version: 1\n") + f.write(" fsdp_sharding_strategy: FULL_SHARD\n") + f.write(" fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP\n") + f.write(" fsdp_transformer_layer_cls_to_wrap: ACTEncoderLayer,ACTDecoderLayer\n") + f.write(" fsdp_use_orig_params: true\n") + f.write(" fsdp_state_dict_type: FULL_STATE_DICT\n") + + +def run_accelerate_training(config_args, num_processes=4, temp_dir=None, distributed_type="MULTI_GPU"): """ Helper function to run training with accelerate launch. @@ -66,6 +105,7 @@ def run_accelerate_training(config_args, num_processes=4, temp_dir=None): config_args: List of config arguments to pass to lerobot_train.py num_processes: Number of processes (GPUs) to use temp_dir: Temporary directory for outputs + distributed_type: "MULTI_GPU" (DDP) or "FSDP" — selects the generated accelerate config. Returns: subprocess.CompletedProcess result @@ -75,18 +115,10 @@ def run_accelerate_training(config_args, num_processes=4, temp_dir=None): # Write YAML config with open(config_path, "w") as f: - f.write("compute_environment: LOCAL_MACHINE\n") - f.write("distributed_type: MULTI_GPU\n") - f.write("mixed_precision: 'no'\n") - f.write(f"num_processes: {num_processes}\n") - f.write("use_cpu: false\n") - f.write("gpu_ids: all\n") - f.write("downcast_bf16: 'no'\n") - f.write("machine_rank: 0\n") - f.write("main_training_function: main\n") - f.write("num_machines: 1\n") - f.write("rdzv_backend: static\n") - f.write("same_network: true\n") + if distributed_type == "FSDP": + _write_fsdp_config(f, num_processes) + else: + _write_multi_gpu_config(f, num_processes) cmd = [ "accelerate", @@ -134,7 +166,7 @@ class TestMultiGPUTraining: f"--output_dir={output_dir}", "--batch_size=4", "--steps=10", - "--eval_freq=-1", + "--env_eval_freq=-1", "--log_freq=5", "--save_freq=10", "--seed=42", @@ -177,7 +209,7 @@ class TestMultiGPUTraining: f"--output_dir={output_dir}", "--batch_size=4", "--steps=20", - "--eval_freq=-1", + "--env_eval_freq=-1", "--log_freq=5", "--save_freq=10", "--seed=42", @@ -211,3 +243,66 @@ class TestMultiGPUTraining: # Verify optimizer state exists optimizer_state = training_state_dir / "optimizer_state.safetensors" assert optimizer_state.exists(), f"No optimizer state in checkpoint {checkpoint_dir}" + + def test_fsdp_optimizer_save_and_resume(self): + """ + Test that FSDP saves the (gathered) optimizer state and can resume from it. + + Trains a few steps under FSDP, verifies the gathered optimizer state is written next to the + rest of the training state, then resumes from the checkpoint for more steps and checks it + completes without shape/key errors in the FSDP optimizer load path. + """ + # Pre-download dataset to avoid race conditions + download_dataset("lerobot/pusht", episodes=[0]) + + with tempfile.TemporaryDirectory() as temp_dir: + output_dir = Path(temp_dir) / "outputs" + + config_args = [ + "--dataset.repo_id=lerobot/pusht", + "--dataset.episodes=[0]", + "--policy.type=act", + "--policy.device=cuda", + "--policy.push_to_hub=false", + f"--output_dir={output_dir}", + "--batch_size=4", + "--steps=10", + "--env_eval_freq=-1", + "--log_freq=5", + "--save_freq=10", + "--seed=42", + "--num_workers=0", + ] + + result = run_accelerate_training( + config_args, num_processes=2, temp_dir=temp_dir, distributed_type="FSDP" + ) + assert result.returncode == 0, ( + f"FSDP training failed:\nSTDOUT:\n{result.stdout}\n\nSTDERR:\n{result.stderr}" + ) + + # The gathered optimizer state must be written under FSDP (proves the save collective ran), + # in the same safetensors format as single-GPU training. + training_state_dir = output_dir / "checkpoints" / "last" / "training_state" + optimizer_state = training_state_dir / "optimizer_state.safetensors" + optimizer_param_groups = training_state_dir / "optimizer_param_groups.json" + assert optimizer_state.exists(), f"FSDP optimizer state not saved in {training_state_dir}" + assert optimizer_param_groups.exists(), ( + f"FSDP optimizer param groups not saved in {training_state_dir}" + ) + + # Resume from the checkpoint for more steps. A successful run proves load_fsdp_optimizer + # accepts the saved state and reshards it without shape/key errors. + resume_config = output_dir / "checkpoints" / "last" / "pretrained_model" / "train_config.json" + resume_args = [ + f"--config_path={resume_config}", + "--resume=true", + "--steps=20", + ] + resume_result = run_accelerate_training( + resume_args, num_processes=2, temp_dir=temp_dir, distributed_type="FSDP" + ) + assert resume_result.returncode == 0, ( + f"FSDP resume failed:\nSTDOUT:\n{resume_result.stdout}\n\nSTDERR:\n{resume_result.stderr}" + ) + assert "End of training" in resume_result.stdout or "End of training" in resume_result.stderr diff --git a/tests/utils/test_hub.py b/tests/utils/test_hub.py new file mode 100644 index 000000000..a55631aeb --- /dev/null +++ b/tests/utils/test_hub.py @@ -0,0 +1,54 @@ +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest.mock import MagicMock + +from lerobot.utils.hub import find_latest_hub_checkpoint + + +def _patch_list_files(monkeypatch, files): + api = MagicMock() + api.list_repo_files.return_value = files + # HfApi is imported into lerobot.utils.hub at module load, so patch it there. + monkeypatch.setattr("lerobot.utils.hub.HfApi", lambda *a, **k: api) + return api + + +def test_find_latest_hub_checkpoint_picks_highest_step(monkeypatch): + _patch_list_files( + monkeypatch, + [ + "README.md", + "checkpoints/000500/pretrained_model/model.safetensors", + "checkpoints/000500/training_state/training_step.json", + "checkpoints/020000/pretrained_model/model.safetensors", + "checkpoints/001000/training_state/training_step.json", + ], + ) + # Numeric max, not lexicographic — "020000" beats "001000"/"000500". + assert find_latest_hub_checkpoint("u/run") == "checkpoints/020000" + + +def test_find_latest_hub_checkpoint_ignores_non_step_entries(monkeypatch): + _patch_list_files( + monkeypatch, + ["checkpoints/last/pretrained_model/model.safetensors", "config.json"], + ) + # "last" (a symlink target name) is not a numeric step → no resolvable checkpoint. + assert find_latest_hub_checkpoint("u/run") is None + + +def test_find_latest_hub_checkpoint_none_when_no_checkpoints(monkeypatch): + _patch_list_files(monkeypatch, ["config.json", "model.safetensors"]) + assert find_latest_hub_checkpoint("u/run") is None diff --git a/tests/utils/test_keyboard_input.py b/tests/utils/test_keyboard_input.py new file mode 100644 index 000000000..2f0dee889 --- /dev/null +++ b/tests/utils/test_keyboard_input.py @@ -0,0 +1,228 @@ +#!/usr/bin/env python + +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for the display-independent keyboard input helpers. + +These cover the parts most likely to regress: the environment-detection decision +table (the heart of the Wayland/headless fix), the macOS trust probe, the control +mapping, the terminal escape-sequence parsing, and backend selection. They require +neither ``pynput`` nor a real terminal. +""" + +import io +import platform +import sys + +import pytest + +import lerobot.utils.keyboard_input as ki +from lerobot.utils.keyboard_input import ( + TerminalKeyListener, + apply_recording_control, + create_key_listener, + init_keyboard_listener, + is_headless, + is_wayland, + pynput_can_capture, + pynput_listener_is_trusted, +) + + +@pytest.fixture(autouse=True) +def _clear_detection_caches(): + """The detection helpers are ``@cache``-decorated; clear around each test.""" + for fn in (is_headless, is_wayland, pynput_can_capture): + fn.cache_clear() + yield + for fn in (is_headless, is_wayland, pynput_can_capture): + fn.cache_clear() + + +def _set_platform(monkeypatch, name): + monkeypatch.setattr(platform, "system", lambda: name) + + +def _set_tty(monkeypatch, is_tty): + stdin = io.StringIO("") + stdin.isatty = lambda: is_tty + monkeypatch.setattr(sys, "stdin", stdin) + + +# --- Environment detection (the core of the fix) --------------------------- +@pytest.mark.parametrize( + ("system", "env", "expected"), + [ + ("Linux", {}, True), # no display server + ("Linux", {"DISPLAY": ":0"}, False), # X11 + ("Linux", {"WAYLAND_DISPLAY": "wayland-0"}, False), # Wayland + ("Darwin", {}, False), # display always assumed present + ], +) +def test_is_headless(monkeypatch, system, env, expected): + _set_platform(monkeypatch, system) + monkeypatch.delenv("DISPLAY", raising=False) + monkeypatch.delenv("WAYLAND_DISPLAY", raising=False) + for key, value in env.items(): + monkeypatch.setenv(key, value) + assert is_headless() is expected + + +@pytest.mark.parametrize( + ("env", "expected"), + [ + ({"XDG_SESSION_TYPE": "wayland"}, True), + ({"WAYLAND_DISPLAY": "wayland-0"}, True), + ({"XDG_SESSION_TYPE": "x11"}, False), + ({}, False), + ], +) +def test_is_wayland(monkeypatch, env, expected): + monkeypatch.delenv("XDG_SESSION_TYPE", raising=False) + monkeypatch.delenv("WAYLAND_DISPLAY", raising=False) + for key, value in env.items(): + monkeypatch.setenv(key, value) + assert is_wayland() is expected + + +@pytest.mark.parametrize( + ("system", "env", "pynput_available", "expected"), + [ + ("Linux", {"DISPLAY": ":0"}, True, True), # X11 + ("Linux", {"DISPLAY": ":0", "WAYLAND_DISPLAY": "wayland-0"}, True, False), # Wayland + ("Linux", {}, True, False), # headless + ("Darwin", {}, True, True), + ("Linux", {"DISPLAY": ":0"}, False, False), # pynput not installed + ], +) +def test_pynput_can_capture(monkeypatch, system, env, pynput_available, expected): + _set_platform(monkeypatch, system) + monkeypatch.setattr(ki, "_pynput_available", pynput_available) + for var in ("DISPLAY", "WAYLAND_DISPLAY", "XDG_SESSION_TYPE"): + monkeypatch.delenv(var, raising=False) + for key, value in env.items(): + monkeypatch.setenv(key, value) + assert pynput_can_capture() is expected + + +# --- macOS trust probe ------------------------------------------------------ +class _FakeListener: + def __init__(self, is_trusted): + self.IS_TRUSTED = is_trusted + + +def test_pynput_listener_is_trusted(monkeypatch): + _set_platform(monkeypatch, "Linux") + assert pynput_listener_is_trusted(_FakeListener(False)) is True # non-macOS: always assumed ok + _set_platform(monkeypatch, "Darwin") + assert pynput_listener_is_trusted(_FakeListener(False), timeout_s=0.05) is False + + +# --- Control mapping -------------------------------------------------------- +def test_apply_recording_control(): + events = {"exit_early": False, "rerecord_episode": False, "stop_recording": False} + apply_recording_control("left", events) + assert events == {"exit_early": True, "rerecord_episode": True, "stop_recording": False} + apply_recording_control("esc", events) + assert events["stop_recording"] is True + apply_recording_control("up", events) # unknown control -> no-op (no error) + + +# --- Terminal escape-sequence parsing (the tricky bit) ---------------------- +def _drive(listener, byte_seq): + """Run the listener's read loop over a scripted list of bytes (no real terminal).""" + script = list(byte_seq) + + def fake_read(timeout): + if script: + return script.pop(0) + listener._running = False + return None + + listener._read_char = fake_read + listener._running = True + listener._run() + + +@pytest.mark.parametrize( + ("byte_seq", "expected"), + [ + (["\x1b", "[", "C"], ["right"]), # CSI arrow + (["\x1b", "O", "D"], ["left"]), # SS3 arrow (e.g. over SSH/tmux) + (["\x1b"], ["esc"]), # bare ESC + (["\x1b", "[", "A"], ["up"]), # decoded even though the record handler ignores it + (["n"], ["n"]), # letter passthrough + ], +) +def test_terminal_parsing(byte_seq, expected): + collected = [] + _drive(TerminalKeyListener(collected.append), byte_seq) + assert collected == expected + + +# --- Backend selection ------------------------------------------------------ +def test_init_selects_terminal_when_pynput_cannot_capture(monkeypatch): + monkeypatch.setattr(ki, "pynput_can_capture", lambda: False) + _set_tty(monkeypatch, is_tty=True) + monkeypatch.setattr(TerminalKeyListener, "start", lambda self: None) # avoid touching termios + listener, _ = init_keyboard_listener() + assert isinstance(listener, TerminalKeyListener) + + +def test_init_returns_none_without_tty(monkeypatch): + monkeypatch.setattr(ki, "pynput_can_capture", lambda: False) + _set_tty(monkeypatch, is_tty=False) + listener, _ = init_keyboard_listener() + assert listener is None + + +@pytest.mark.parametrize( + ("key", "flag"), + [("right", "exit_early"), ("r", "rerecord_episode"), ("q", "stop_recording")], +) +def test_init_terminal_key_routing(monkeypatch, key, flag): + """Arrows and their letter equivalents drive the same events (terminal backend).""" + monkeypatch.setattr(ki, "pynput_can_capture", lambda: False) + _set_tty(monkeypatch, is_tty=True) + monkeypatch.setattr(TerminalKeyListener, "start", lambda self: None) + listener, events = init_keyboard_listener() + listener._on_key(key) + assert events[flag] is True + + +# --- Shared factory + pynput key resolver ----------------------------------- +def test_resolve_pynput_key_char_fallback(): + """Unmapped keys fall back to ``.char`` (and yield None when there is none).""" + assert ki._resolve_pynput_key(type("K", (), {"char": "s"})()) == "s" + assert ki._resolve_pynput_key(type("K", (), {"char": None})()) is None + assert ki._resolve_pynput_key(type("K", (), {"char": ""})()) is None # empty char -> no key + + +def test_create_key_listener_routes_to_dispatch(monkeypatch): + """The terminal backend forwards canonical key names straight to ``dispatch``.""" + monkeypatch.setattr(ki, "pynput_can_capture", lambda: False) + _set_tty(monkeypatch, is_tty=True) + monkeypatch.setattr(TerminalKeyListener, "start", lambda self: None) + seen = [] + listener = create_key_listener(seen.append, controls_help="save='s'") + assert isinstance(listener, TerminalKeyListener) + listener._on_key("space") + assert seen == ["space"] + + +def test_create_key_listener_none_without_tty(monkeypatch): + monkeypatch.setattr(ki, "pynput_can_capture", lambda: False) + _set_tty(monkeypatch, is_tty=False) + assert create_key_listener(lambda name: None) is None diff --git a/tests/utils/test_train_utils.py b/tests/utils/test_train_utils.py index c171763c2..ccd769bd0 100644 --- a/tests/utils/test_train_utils.py +++ b/tests/utils/test_train_utils.py @@ -15,7 +15,9 @@ # limitations under the License. from pathlib import Path -from unittest.mock import Mock, patch +from unittest.mock import MagicMock, Mock, patch + +import pytest from lerobot.common.train_utils import ( get_step_checkpoint_dir, @@ -24,6 +26,7 @@ from lerobot.common.train_utils import ( load_training_num_processes, load_training_state, load_training_step, + push_checkpoint_to_hub, save_checkpoint, save_training_state, save_training_step, @@ -136,3 +139,87 @@ def test_save_load_training_state(tmp_path, optimizer, scheduler): assert loaded_step == 10 assert loaded_optimizer is optimizer assert loaded_scheduler is scheduler + + +def test_load_training_state_skip_optimizer(tmp_path, optimizer, scheduler): + # FSDP loads optimizer separately (after accelerator.prepare) + # load_training_state(load_optimizer=False) must restore step + scheduler but leave the + # optimizer untouched and never touch the on-disk optimizer state. + save_training_state(tmp_path, 10, optimizer, scheduler) + with patch("lerobot.common.train_utils.load_optimizer_state") as mock_load_optimizer_state: + loaded_step, loaded_optimizer, loaded_scheduler = load_training_state( + tmp_path, optimizer, scheduler, load_optimizer=False + ) + mock_load_optimizer_state.assert_not_called() + assert loaded_step == 10 + assert loaded_optimizer is optimizer + assert loaded_scheduler is scheduler + + +def test_push_checkpoint_to_hub_creates_repo_and_uploads(tmp_path, monkeypatch): + ckpt = tmp_path / "010000" + (ckpt / "pretrained_model").mkdir(parents=True) + api = MagicMock() + monkeypatch.setattr("lerobot.common.train_utils.HfApi", lambda *a, **k: api) + push_checkpoint_to_hub(ckpt, "user/run", private=True) + api.create_repo.assert_called_once() + assert api.create_repo.call_args.kwargs["private"] is True + assert api.create_repo.call_args.kwargs["repo_type"] == "model" + api.upload_folder.assert_called_once() + kwargs = api.upload_folder.call_args.kwargs + assert kwargs["repo_id"] == "user/run" + assert kwargs["repo_type"] == "model" + assert kwargs["path_in_repo"] == "checkpoints/010000" + assert kwargs["folder_path"] == str(ckpt) + assert kwargs["commit_message"] == "checkpoint 010000" + # A tag named after the checkpoint step is created so the checkpoint can be + # recovered with --policy.pretrained_revision instead of a commit sha. + api.create_tag.assert_called_once() + tag_kwargs = api.create_tag.call_args.kwargs + assert tag_kwargs["tag"] == "010000" + assert tag_kwargs["revision"] == api.upload_folder.return_value.oid + assert tag_kwargs["repo_type"] == "model" + assert tag_kwargs["exist_ok"] is True + + +def test_push_checkpoint_to_hub_defaults_to_hub_default_visibility(tmp_path, monkeypatch): + ckpt = tmp_path / "010000" + (ckpt / "pretrained_model").mkdir(parents=True) + api = MagicMock() + monkeypatch.setattr("lerobot.common.train_utils.HfApi", lambda *a, **k: api) + push_checkpoint_to_hub(ckpt, "user/run") + api.create_repo.assert_called_once() + assert api.create_repo.call_args.kwargs["private"] is None + + +def test_resolve_resume_checkpoint_downloads_latest_and_links(tmp_path, monkeypatch): + from lerobot.common import train_utils + + out = tmp_path / "run" + + def fake_snapshot_download(repo_id, repo_type, allow_patterns, local_dir): + # Mimic the Hub layout the real download materializes locally. + assert allow_patterns == "checkpoints/020000/*" + (Path(local_dir) / "checkpoints" / "020000" / "pretrained_model").mkdir(parents=True) + return local_dir + + monkeypatch.setattr("lerobot.common.train_utils.snapshot_download", fake_snapshot_download) + monkeypatch.setattr( + "lerobot.common.train_utils.find_latest_hub_checkpoint", lambda repo_id: "checkpoints/020000" + ) + + checkpoint_dir = train_utils.resolve_resume_checkpoint("u/run", out) + + assert checkpoint_dir == out / CHECKPOINTS_DIR / "020000" + last = out / CHECKPOINTS_DIR / LAST_CHECKPOINT_LINK + assert last.is_symlink() + # `last` points at the downloaded step dir. + assert (last.parent / last.readlink()).resolve() == checkpoint_dir.resolve() + + +def test_resolve_resume_checkpoint_raises_without_checkpoints(tmp_path, monkeypatch): + from lerobot.common import train_utils + + monkeypatch.setattr("lerobot.common.train_utils.find_latest_hub_checkpoint", lambda repo_id: None) + with pytest.raises(FileNotFoundError, match="No checkpoint"): + train_utils.resolve_resume_checkpoint("u/run", tmp_path / "run") diff --git a/tests/utils/test_visualization_utils.py b/tests/utils/test_visualization_utils.py index 63ff76c77..f62a697cd 100644 --- a/tests/utils/test_visualization_utils.py +++ b/tests/utils/test_visualization_utils.py @@ -30,46 +30,77 @@ from lerobot.utils.constants import OBS_STATE @pytest.fixture def mock_rerun(monkeypatch): """ - Provide a mock `rerun` module so tests don't depend on the real library. - Also reload the module-under-test so it binds to this mock `rr`. + Provide a mock `rerun` module (and `rerun.blueprint` submodule) so tests don't + depend on the real library. Also reload the module-under-test so it binds to + this mock `rr`. """ calls = [] + blueprints = [] class DummyScalar: def __init__(self, value): - self.value = float(value) + # Scalars may be built from a single float or from a 1D array batch. + self.value = value class DummyImage: def __init__(self, arr): self.arr = arr + def compress(self, *a, **k): + return self + + class DummyDepthImage: + def __init__(self, arr, colormap=None): + self.arr = arr + self.colormap = colormap + def dummy_log(key, obj=None, **kwargs): # Accept either positional `obj` or keyword `entity` and record remaining kwargs. if obj is None and "entity" in kwargs: obj = kwargs.pop("entity") calls.append((key, obj, kwargs)) + def dummy_send_blueprint(blueprint, *a, **k): + blueprints.append(blueprint) + + # Mock the `rerun.blueprint` submodule used to build the layout. + dummy_rrb = SimpleNamespace( + Spatial2DView=lambda origin=None, name=None: SimpleNamespace( + kind="Spatial2DView", origin=origin, name=name + ), + TimeSeriesView=lambda name=None, contents=None: SimpleNamespace( + kind="TimeSeriesView", name=name, contents=contents + ), + Grid=lambda *views: SimpleNamespace(kind="Grid", views=list(views)), + Blueprint=lambda root: SimpleNamespace(kind="Blueprint", root=root), + ) + dummy_rr = SimpleNamespace( __name__="rerun", __package__="rerun", __spec__=SimpleNamespace(name="rerun", submodule_search_locations=None), Scalars=DummyScalar, Image=DummyImage, + DepthImage=DummyDepthImage, + components=SimpleNamespace(Colormap=SimpleNamespace(Viridis="viridis")), log=dummy_log, + send_blueprint=dummy_send_blueprint, init=lambda *a, **k: None, spawn=lambda *a, **k: None, + blueprint=dummy_rrb, ) - # Inject fake module into sys.modules + # Inject fake modules into sys.modules (both `rerun` and `rerun.blueprint`). monkeypatch.setitem(sys.modules, "rerun", dummy_rr) + monkeypatch.setitem(sys.modules, "rerun.blueprint", dummy_rrb) # Now import and reload the module under test, to bind to our rerun mock import lerobot.utils.visualization_utils as vu importlib.reload(vu) - # Expose both the reloaded module and the call recorder - yield vu, calls + # Expose the reloaded module, the call recorder and the captured blueprints + yield vu, calls, blueprints def _keys(calls): @@ -92,8 +123,13 @@ def _kwargs_for(calls, key): raise KeyError(f"Key {key} not found in calls: {calls}") +def _views_by_kind(blueprint, kind): + """Return the views of a given kind from the (single) blueprint's grid.""" + return [v for v in blueprint.root.views if v.kind == kind] + + def test_log_rerun_data_envtransition_scalars_and_image(mock_rerun): - vu, calls = mock_rerun + vu, calls, blueprints = mock_rerun # Build EnvTransition dict obs = { @@ -103,7 +139,7 @@ def test_log_rerun_data_envtransition_scalars_and_image(mock_rerun): } act = { "action.throttle": 0.7, - # 1D array should log individual Scalars with suffix _i + # 1D array should be logged as a single Scalars batch under one entity path "action.vector": np.array([1.0, 2.0], dtype=np.float32), } transition = { @@ -120,31 +156,28 @@ def test_log_rerun_data_envtransition_scalars_and_image(mock_rerun): # - observation.state.temperature -> Scalars # - observation.camera -> Image (HWC) with static=True # - action.throttle -> Scalars - # - action.vector_0, action.vector_1 -> Scalars + # - action.vector -> single Scalars batch (no per-element suffix) expected_keys = { f"{OBS_STATE}.temperature", "observation.camera", "action.throttle", - "action.vector_0", - "action.vector_1", + "action.vector", } assert set(_keys(calls)) == expected_keys # Check scalar types and values temp_obj = _obj_for(calls, f"{OBS_STATE}.temperature") assert type(temp_obj).__name__ == "DummyScalar" - assert temp_obj.value == pytest.approx(25.0) + assert float(temp_obj.value) == pytest.approx(25.0) throttle_obj = _obj_for(calls, "action.throttle") assert type(throttle_obj).__name__ == "DummyScalar" - assert throttle_obj.value == pytest.approx(0.7) + assert float(throttle_obj.value) == pytest.approx(0.7) - v0 = _obj_for(calls, "action.vector_0") - v1 = _obj_for(calls, "action.vector_1") - assert type(v0).__name__ == "DummyScalar" - assert type(v1).__name__ == "DummyScalar" - assert v0.value == pytest.approx(1.0) - assert v1.value == pytest.approx(2.0) + # 1D vector logged as a single batched Scalars under one entity path + vec = _obj_for(calls, "action.vector") + assert type(vec).__name__ == "DummyScalar" + np.testing.assert_allclose(np.asarray(vec.value), [1.0, 2.0]) # Check image handling: CHW -> HWC img_obj = _obj_for(calls, "observation.camera") @@ -152,9 +185,24 @@ def test_log_rerun_data_envtransition_scalars_and_image(mock_rerun): assert img_obj.arr.shape == (10, 20, 3) # transposed assert _kwargs_for(calls, "observation.camera").get("static", False) is True # static=True for images + # A blueprint should have been built and sent exactly once, and cached on the function. + assert len(blueprints) == 1 + assert vu.log_rerun_data.blueprint is blueprints[0] + + bp = blueprints[0] + # One spatial view per image path + spatial_views = _views_by_kind(bp, "Spatial2DView") + assert {v.origin for v in spatial_views} == {"observation.camera"} + + # One time-series view each for observation and action scalars + ts_views = {v.name: v for v in _views_by_kind(bp, "TimeSeriesView")} + assert set(ts_views) == {"observation", "action"} + assert ts_views["observation"].contents == [f"{OBS_STATE}.temperature"] + assert ts_views["action"].contents == ["action.throttle", "action.vector"] + def test_log_rerun_data_plain_list_ordering_and_prefixes(mock_rerun): - vu, calls = mock_rerun + vu, calls, blueprints = mock_rerun # First dict without prefixes treated as observation # Second dict without prefixes treated as action @@ -173,14 +221,12 @@ def test_log_rerun_data_plain_list_ordering_and_prefixes(mock_rerun): # First dict was treated as observation, second as action vu.log_rerun_data(observation=obs_plain, action=act_plain) - # Expected keys with auto-prefixes + # Expected keys with auto-prefixes. The 1D vector is a single batched Scalars. expected = { "observation.temp", "observation.img", "action.throttle", - "action.vec_0", - "action.vec_1", - "action.vec_2", + "action.vec", } logged = set(_keys(calls)) assert logged == expected @@ -188,11 +234,11 @@ def test_log_rerun_data_plain_list_ordering_and_prefixes(mock_rerun): # Scalars t = _obj_for(calls, "observation.temp") assert type(t).__name__ == "DummyScalar" - assert t.value == pytest.approx(1.5) + assert float(t.value) == pytest.approx(1.5) throttle = _obj_for(calls, "action.throttle") assert type(throttle).__name__ == "DummyScalar" - assert throttle.value == pytest.approx(0.3) + assert float(throttle.value) == pytest.approx(0.3) # Image stays HWC img = _obj_for(calls, "observation.img") @@ -200,15 +246,23 @@ def test_log_rerun_data_plain_list_ordering_and_prefixes(mock_rerun): assert img.arr.shape == (5, 6, 3) assert _kwargs_for(calls, "observation.img").get("static", False) is True - # Vectors - for i, val in enumerate([9, 8, 7]): - o = _obj_for(calls, f"action.vec_{i}") - assert type(o).__name__ == "DummyScalar" - assert o.value == pytest.approx(val) + # Vector logged as a single batched Scalars under one entity path + vec = _obj_for(calls, "action.vec") + assert type(vec).__name__ == "DummyScalar" + np.testing.assert_allclose(np.asarray(vec.value), [9, 8, 7]) + + # Blueprint sent once with the expected view layout + assert len(blueprints) == 1 + bp = blueprints[0] + spatial_views = _views_by_kind(bp, "Spatial2DView") + assert {v.origin for v in spatial_views} == {"observation.img"} + ts_views = {v.name: v for v in _views_by_kind(bp, "TimeSeriesView")} + assert ts_views["observation"].contents == ["observation.temp"] + assert ts_views["action"].contents == ["action.throttle", "action.vec"] def test_log_rerun_data_kwargs_only(mock_rerun): - vu, calls = mock_rerun + vu, calls, blueprints = mock_rerun vu.log_rerun_data( observation={"observation.temp": 10.0, "observation.gray": np.zeros((8, 8, 1), dtype=np.uint8)}, @@ -222,13 +276,35 @@ def test_log_rerun_data_kwargs_only(mock_rerun): temp = _obj_for(calls, "observation.temp") assert type(temp).__name__ == "DummyScalar" - assert temp.value == pytest.approx(10.0) + assert float(temp.value) == pytest.approx(10.0) img = _obj_for(calls, "observation.gray") - assert type(img).__name__ == "DummyImage" + assert type(img).__name__ == "DummyDepthImage" # single-channel -> DepthImage assert img.arr.shape == (8, 8, 1) # remains HWC assert _kwargs_for(calls, "observation.gray").get("static", False) is True a = _obj_for(calls, "action.a") assert type(a).__name__ == "DummyScalar" - assert a.value == pytest.approx(1.0) + assert float(a.value) == pytest.approx(1.0) + + # Blueprint sent once, with a spatial view for the image and time-series views for scalars + assert len(blueprints) == 1 + bp = blueprints[0] + assert {v.origin for v in _views_by_kind(bp, "Spatial2DView")} == {"observation.gray"} + ts_views = {v.name: v for v in _views_by_kind(bp, "TimeSeriesView")} + assert ts_views["observation"].contents == ["observation.temp"] + assert ts_views["action"].contents == ["action.a"] + + +def test_log_rerun_data_blueprint_sent_only_once(mock_rerun): + """The blueprint is built from the first call and not resent on subsequent calls.""" + vu, calls, blueprints = mock_rerun + + vu.log_rerun_data(observation={"temp": 1.0}, action={"a": 2.0}) + assert len(blueprints) == 1 + first_blueprint = vu.log_rerun_data.blueprint + + vu.log_rerun_data(observation={"temp": 3.0}, action={"a": 4.0}) + # Still only one blueprint, and the cached one is unchanged. + assert len(blueprints) == 1 + assert vu.log_rerun_data.blueprint is first_blueprint diff --git a/uv.lock b/uv.lock index 743a71e45..5a76fcbf8 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 3 +revision = 2 requires-python = ">=3.12" resolution-markers = [ "(python_full_version >= '3.15' and platform_machine == 'AMD64' and sys_platform == 'linux') or (python_full_version >= '3.15' and platform_machine == 'x86_64' and sys_platform == 'linux')", @@ -51,7 +51,7 @@ resolution-markers = [ [[package]] name = "absl-py" version = "2.4.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/64/c7/8de93764ad66968d19329a7e0c147a2bb3c7054c554d4a119111b8f9440f/absl_py-2.4.0.tar.gz", hash = "sha256:8c6af82722b35cf71e0f4d1d47dcaebfff286e27110a99fc359349b247dfb5d4", size = 116543, upload-time = "2026-01-28T10:17:05.322Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/18/a6/907a406bb7d359e6a63f99c313846d9eec4f7e6f7437809e03aa00fa3074/absl_py-2.4.0-py3-none-any.whl", hash = "sha256:88476fd881ca8aab94ffa78b7b6c632a782ab3ba1cd19c9bd423abc4fb4cd28d", size = 135750, upload-time = "2026-01-28T10:17:04.19Z" }, @@ -60,7 +60,7 @@ wheels = [ [[package]] name = "accelerate" version = "1.14.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "huggingface-hub" }, { name = "numpy" }, @@ -68,7 +68,7 @@ dependencies = [ { name = "psutil" }, { name = "pyyaml" }, { name = "safetensors" }, - { name = "torch", version = "2.11.0", source = { registry = "https://pypi.registries.huggingface.tech/" }, marker = "sys_platform != 'linux'" }, + { name = "torch", version = "2.11.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'linux'" }, { name = "torch", version = "2.11.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" }, marker = "sys_platform == 'linux'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/8d/75/94cd5d389649578aca399e5aa822637eec18319a1dadc400ffe2f9a7493f/accelerate-1.14.0.tar.gz", hash = "sha256:41b9c4377a54e0b460a959b0defa1b736e4ca0a2373252d9a539964c2afe3c8d", size = 412167, upload-time = "2026-06-11T13:45:52.326Z" } @@ -79,7 +79,7 @@ wheels = [ [[package]] name = "aiohappyeyeballs" version = "2.6.2" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/33/c6/61a2d7b7572279226bb2e7f61d7a19ca7c90da0329c93fa0d560cbf288d8/aiohappyeyeballs-2.6.2.tar.gz", hash = "sha256:e202810ee718bd01fc6ef49e8ea53d023d5cb6b581076d7925aa499fa55dbe64", size = 22591, upload-time = "2026-05-20T15:12:24.631Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/5f/fc/a7bf5b6e4e617b45f90f2d9d2a68519c249c81dd4fc2658c7a2a61c4f4b7/aiohappyeyeballs-2.6.2-py3-none-any.whl", hash = "sha256:4708045e2d7a6c6bdf8aafa8ed39649eaf926a4543b54560659129e3365953c4", size = 15062, upload-time = "2026-05-20T15:12:23.328Z" }, @@ -88,7 +88,7 @@ wheels = [ [[package]] name = "aiohttp" version = "3.14.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "aiohappyeyeballs" }, { name = "aiosignal" }, @@ -188,7 +188,7 @@ wheels = [ [[package]] name = "aiosignal" version = "1.4.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "frozenlist" }, { name = "typing-extensions", marker = "python_full_version < '3.13'" }, @@ -201,7 +201,7 @@ wheels = [ [[package]] name = "annotated-doc" version = "0.0.4" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/57/ba/046ceea27344560984e26a590f90bc7f4a75b06701f653222458922b558c/annotated_doc-0.0.4.tar.gz", hash = "sha256:fbcda96e87e9c92ad167c2e53839e57503ecfda18804ea28102353485033faa4", size = 7288, upload-time = "2025-11-10T22:07:42.062Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/1e/d3/26bf1008eb3d2daa8ef4cacc7f3bfdc11818d111f7e2d0201bc6e3b49d45/annotated_doc-0.0.4-py3-none-any.whl", hash = "sha256:571ac1dc6991c450b25a9c2d84a3705e2ae7a53467b5d111c24fa8baabbed320", size = 5303, upload-time = "2025-11-10T22:07:40.673Z" }, @@ -210,7 +210,7 @@ wheels = [ [[package]] name = "annotated-types" version = "0.7.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/ee/67/531ea369ba64dcff5ec9c3402f9f51bf748cec26dde048a2f973a4eea7f5/annotated_types-0.7.0.tar.gz", hash = "sha256:aff07c09a53a08bc8cfccb9c85b05f1aa9a2a6f23728d790723543408344ce89", size = 16081, upload-time = "2024-05-20T21:33:25.928Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/78/b6/6307fbef88d9b5ee7421e68d78a9f162e0da4900bc5f5793f6d3d0e34fb8/annotated_types-0.7.0-py3-none-any.whl", hash = "sha256:1f02e8b43a8fbbc3f3e0d4f0f4bfc8131bcb4eebe8849b8e5c773f3a1c582a53", size = 13643, upload-time = "2024-05-20T21:33:24.1Z" }, @@ -219,13 +219,13 @@ wheels = [ [[package]] name = "antlr4-python3-runtime" version = "4.9.3" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/3e/38/7859ff46355f76f8d19459005ca000b6e7012f2f1ca597746cbcd1fbfe5e/antlr4-python3-runtime-4.9.3.tar.gz", hash = "sha256:f224469b4168294902bb1efa80a8bf7855f24c99aef99cbefc1bcd3cce77881b", size = 117034, upload-time = "2021-11-06T17:52:23.524Z" } [[package]] name = "anyio" version = "4.13.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "idna" }, { name = "typing-extensions", marker = "python_full_version < '3.13'" }, @@ -238,7 +238,7 @@ wheels = [ [[package]] name = "appnope" version = "0.1.4" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/35/5d/752690df9ef5b76e169e68d6a129fa6d08a7100ca7f754c89495db3c6019/appnope-0.1.4.tar.gz", hash = "sha256:1de3860566df9caf38f01f86f65e0e13e379af54f9e4bee1e66b48f2efffd1ee", size = 4170, upload-time = "2024-02-06T09:43:11.258Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/81/29/5ecc3a15d5a33e31b26c11426c45c501e439cb865d0bff96315d86443b78/appnope-0.1.4-py2.py3-none-any.whl", hash = "sha256:502575ee11cd7a28c0205f379b525beefebab9d161b7c964670864014ed7213c", size = 4321, upload-time = "2024-02-06T09:43:09.663Z" }, @@ -247,7 +247,7 @@ wheels = [ [[package]] name = "argon2-cffi" version = "25.1.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "argon2-cffi-bindings" }, ] @@ -259,7 +259,7 @@ wheels = [ [[package]] name = "argon2-cffi-bindings" version = "25.1.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "cffi" }, ] @@ -290,7 +290,7 @@ wheels = [ [[package]] name = "arrow" version = "1.4.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "python-dateutil" }, { name = "tzdata" }, @@ -303,7 +303,7 @@ wheels = [ [[package]] name = "ast-serialize" version = "0.5.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/81/9d/09e27731bd5864a9ce04e3244074e674bb8936bf62b45e0357248717adac/ast_serialize-0.5.0.tar.gz", hash = "sha256:5880091bfe6f4f986f22866375c2e884843e7a0b6343ae41aeea659613d879b6", size = 61157, upload-time = "2026-05-17T17:48:29.429Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/c0/9a/13dde51ba9e15f8b97957ab7cb0120d0e381524d651c6bd630b9c359227f/ast_serialize-0.5.0-cp314-cp314t-macosx_10_12_x86_64.whl", hash = "sha256:8f5c14f169eb0972c0c21bada5358b23d6047c76583b005234f865b11f1fa00a", size = 1183520, upload-time = "2026-05-17T17:47:30.831Z" }, @@ -343,7 +343,7 @@ wheels = [ [[package]] name = "asttokens" version = "3.0.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/be/a5/8e3f9b6771b0b408517c82d97aed8f2036509bc247d46114925e32fe33f0/asttokens-3.0.1.tar.gz", hash = "sha256:71a4ee5de0bde6a31d64f6b13f2293ac190344478f081c3d1bccfcf5eacb0cb7", size = 62308, upload-time = "2025-11-15T16:43:48.578Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/d2/39/e7eaf1799466a4aef85b6a4fe7bd175ad2b1c6345066aa33f1f58d4b18d0/asttokens-3.0.1-py3-none-any.whl", hash = "sha256:15a3ebc0f43c2d0a50eeafea25e19046c68398e487b9f1f5b517f7c0f40f976a", size = 27047, upload-time = "2025-11-15T16:43:16.109Z" }, @@ -352,7 +352,7 @@ wheels = [ [[package]] name = "async-lru" version = "2.3.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/e8/1f/989ecfef8e64109a489fff357450cb73fa73a865a92bd8c272170a6922c2/async_lru-2.3.0.tar.gz", hash = "sha256:89bdb258a0140d7313cf8f4031d816a042202faa61d0ab310a0a538baa1c24b6", size = 16332, upload-time = "2026-03-19T01:04:32.413Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/e5/e2/c2e3abf398f80732e58b03be77bde9022550d221dd8781bf586bd4d97cc1/async_lru-2.3.0-py3-none-any.whl", hash = "sha256:eea27b01841909316f2cc739807acea1c623df2be8c5cfad7583286397bb8315", size = 8403, upload-time = "2026-03-19T01:04:30.883Z" }, @@ -361,7 +361,7 @@ wheels = [ [[package]] name = "attrs" version = "26.1.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/9a/8e/82a0fe20a541c03148528be8cac2408564a6c9a0cc7e9171802bc1d26985/attrs-26.1.0.tar.gz", hash = "sha256:d03ceb89cb322a8fd706d4fb91940737b6642aa36998fe130a9bc96c985eff32", size = 952055, upload-time = "2026-03-19T14:22:25.026Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/64/b4/17d4b0b2a2dc85a6df63d1157e028ed19f90d4cd97c36717afef2bc2f395/attrs-26.1.0-py3-none-any.whl", hash = "sha256:c647aa4a12dfbad9333ca4e71fe62ddc36f4e63b2d260a37a8b83d2f043ac309", size = 67548, upload-time = "2026-03-19T14:22:23.645Z" }, @@ -370,7 +370,7 @@ wheels = [ [[package]] name = "av" version = "15.1.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/e9/c3/83e6e73d1592bc54436eae0bc61704ae0cff0c3cfbde7b58af9ed67ebb49/av-15.1.0.tar.gz", hash = "sha256:39cda2dc810e11c1938f8cb5759c41d6b630550236b3365790e67a313660ec85", size = 3774192, upload-time = "2025-08-30T04:41:56.076Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/d1/58/de78b276d20db6ffcd4371283df771721a833ba525a3d57e753d00a9fe79/av-15.1.0-cp312-cp312-macosx_13_0_arm64.whl", hash = "sha256:40c5df37f4c354ab8190c6fd68dab7881d112f527906f64ca73da4c252a58cee", size = 21760991, upload-time = "2025-08-30T04:40:00.801Z" }, @@ -406,7 +406,7 @@ wheels = [ [[package]] name = "babel" version = "2.18.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/7d/b2/51899539b6ceeeb420d40ed3cd4b7a40519404f9baf3d4ac99dc413a834b/babel-2.18.0.tar.gz", hash = "sha256:b80b99a14bd085fcacfa15c9165f651fbb3406e66cc603abf11c5750937c992d", size = 9959554, upload-time = "2026-02-01T12:30:56.078Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/77/f5/21d2de20e8b8b0408f0681956ca2c69f1320a3848ac50e6e7f39c6159675/babel-2.18.0-py3-none-any.whl", hash = "sha256:e2b422b277c2b9a9630c1d7903c2a00d0830c409c59ac8cae9081c92f1aeba35", size = 10196845, upload-time = "2026-02-01T12:30:53.445Z" }, @@ -415,7 +415,7 @@ wheels = [ [[package]] name = "bddl" version = "1.0.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "jupytext", marker = "sys_platform == 'linux'" }, { name = "networkx", marker = "sys_platform == 'linux'" }, @@ -427,7 +427,7 @@ sdist = { url = "https://files.pythonhosted.org/packages/5c/37/0211f82891a9f14ef [[package]] name = "beautifulsoup4" version = "4.15.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "soupsieve" }, { name = "typing-extensions" }, @@ -440,7 +440,7 @@ wheels = [ [[package]] name = "bleach" version = "6.4.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "webencodings" }, ] @@ -457,7 +457,7 @@ css = [ [[package]] name = "certifi" version = "2026.5.20" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/f3/ce/ee2ecad540810a79593028e88299baeae54d346cc7a0d94b6199988b89b1/certifi-2026.5.20.tar.gz", hash = "sha256:69dea482ab64caa7b9f6aba1c6bf48bb6a5448d1c0f1b17ab42ad8c763a5344d", size = 135422, upload-time = "2026-05-20T11:46:50.073Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/59/8c/57e832b7af6d7c5abe66eb3fbe3a3a32f4d11ea23a1aa7131371035be991/certifi-2026.5.20-py3-none-any.whl", hash = "sha256:3c52e209ba0a4ad7aebe60436a4ab349c39e1e602e8c134221e546902ad25897", size = 134134, upload-time = "2026-05-20T11:46:48.578Z" }, @@ -466,7 +466,7 @@ wheels = [ [[package]] name = "cffi" version = "2.0.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "pycparser", marker = "implementation_name != 'PyPy'" }, ] @@ -523,7 +523,7 @@ wheels = [ [[package]] name = "cfgv" version = "3.5.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/4e/b5/721b8799b04bf9afe054a3899c6cf4e880fcf8563cc71c15610242490a0c/cfgv-3.5.0.tar.gz", hash = "sha256:d5b1034354820651caa73ede66a6294d6e95c1b00acc5e9b098e917404669132", size = 7334, upload-time = "2025-11-19T20:55:51.612Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/db/3c/33bac158f8ab7f89b2e59426d5fe2e4f63f7ed25df84c036890172b412b5/cfgv-3.5.0-py2.py3-none-any.whl", hash = "sha256:a8dc6b26ad22ff227d2634a65cb388215ce6cc96bbcc5cfde7641ae87e8dacc0", size = 7445, upload-time = "2025-11-19T20:55:50.744Z" }, @@ -532,7 +532,7 @@ wheels = [ [[package]] name = "charset-normalizer" version = "3.4.7" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/e7/a1/67fe25fac3c7642725500a3f6cfe5821ad557c3abb11c9d20d12c7008d3e/charset_normalizer-3.4.7.tar.gz", hash = "sha256:ae89db9e5f98a11a4bf50407d4363e7b09b31e55bc117b4f7d80aab97ba009e5", size = 144271, upload-time = "2026-04-02T09:28:39.342Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/0c/eb/4fc8d0a7110eb5fc9cc161723a34a8a6c200ce3b4fbf681bc86feee22308/charset_normalizer-3.4.7-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:eca9705049ad3c7345d574e3510665cb2cf844c2f2dcfe675332677f081cbd46", size = 311328, upload-time = "2026-04-02T09:26:24.331Z" }, @@ -605,7 +605,7 @@ wheels = [ [[package]] name = "click" version = "8.4.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "colorama", marker = "sys_platform == 'win32'" }, ] @@ -617,7 +617,7 @@ wheels = [ [[package]] name = "cloudpickle" version = "3.1.2" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/27/fb/576f067976d320f5f0114a8d9fa1215425441bb35627b1993e5afd8111e5/cloudpickle-3.1.2.tar.gz", hash = "sha256:7fda9eb655c9c230dab534f1983763de5835249750e85fbcef43aaa30a9a2414", size = 22330, upload-time = "2025-11-03T09:25:26.604Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/88/39/799be3f2f0f38cc727ee3b4f1445fe6d5e4133064ec2e4115069418a5bb6/cloudpickle-3.1.2-py3-none-any.whl", hash = "sha256:9acb47f6afd73f60dc1df93bb801b472f05ff42fa6c84167d25cb206be1fbf4a", size = 22228, upload-time = "2025-11-03T09:25:25.534Z" }, @@ -626,7 +626,7 @@ wheels = [ [[package]] name = "cmake" version = "4.1.3" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/f2/17/f8f42ae205604319cc36f46d9929bd9bfbd83d3d02d6314c44fa97c42006/cmake-4.1.3.tar.gz", hash = "sha256:89f48ddc2570eb62447e33311cffc6dfeb09631bd0a19423d8a59cec8af030f1", size = 34998, upload-time = "2025-11-19T22:41:27.312Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/24/79/1bf4009d7ef16d62e0b92ddb78efeda830ca5903149abf9dc01d270c3d4e/cmake-4.1.3-py3-none-macosx_10_10_universal2.whl", hash = "sha256:3b6b25ce8fecc768881b36a1dfbca0013adac10a299c73e24cf4cbb99e4c37d6", size = 49246088, upload-time = "2025-11-19T22:40:28.02Z" }, @@ -652,7 +652,7 @@ wheels = [ [[package]] name = "cmeel" version = "0.60.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/54/46/ddc7df697e49cae32b1a97b3e1a3b47b815238f9059312f987bc62a2e756/cmeel-0.60.1.tar.gz", hash = "sha256:3e0b92eb933a3693ad3f1da8aae31defcbee5f25969daaf20e59c57d6a9474cf", size = 14972, upload-time = "2026-05-11T17:13:57.216Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/da/39/f2db2ff475d42222c70fe25c737028aeaafdd9c0aeba04e6b2dc66e7f93f/cmeel-0.60.1-py3-none-any.whl", hash = "sha256:3f92b68353a58b4d6b5a664ea96bf58a3fef0891a8ed570d3c153361bbbb94b7", size = 20612, upload-time = "2026-05-11T17:13:55.812Z" }, @@ -661,7 +661,7 @@ wheels = [ [[package]] name = "cmeel-assimp" version = "5.4.3.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "cmeel" }, { name = "cmeel-zlib" }, @@ -681,7 +681,7 @@ wheels = [ [[package]] name = "cmeel-boost" version = "1.87.0.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "cmeel" }, { name = "numpy" }, @@ -703,7 +703,7 @@ wheels = [ [[package]] name = "cmeel-console-bridge" version = "1.0.2.3" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "cmeel" }, ] @@ -722,7 +722,7 @@ wheels = [ [[package]] name = "cmeel-octomap" version = "1.10.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "cmeel" }, ] @@ -745,7 +745,7 @@ wheels = [ [[package]] name = "cmeel-qhull" version = "8.0.2.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "cmeel" }, ] @@ -768,40 +768,52 @@ wheels = [ [[package]] name = "cmeel-tinyxml2" -version = "11.0.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +version = "10.0.0" +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "cmeel" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/9b/96/4311533fee0a364bb605b585762f04c249f47857b33548a8ea837a7eb860/cmeel_tinyxml2-11.0.0.tar.gz", hash = "sha256:85d9c7680b3369af4c6b40a0dce70bbd84aa67832755622e57eb260cd95abe40", size = 645900, upload-time = "2026-05-21T11:49:32.652Z" } +sdist = { url = "https://files.pythonhosted.org/packages/28/9f/030eca702c485f7a641f975f167fa93164911b3329f005fb0730ff5e793f/cmeel_tinyxml2-10.0.0.tar.gz", hash = "sha256:00252aefc1c94a55b89f25ad08ee79fda2da8d1d94703e051598ddb52a9088fe", size = 645297, upload-time = "2025-02-06T10:29:00.106Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/79/f0/90c1640c53b623359d75ab1c70bdf19dc0afe82722bc5df57d09f8eaf83a/cmeel_tinyxml2-11.0.0-0-py3-none-macosx_10_9_x86_64.whl", hash = "sha256:b0bd974e549b8c444626671a8e645897603ebf5225734cbe04a9dd3461477754", size = 111719, upload-time = "2026-05-21T11:49:25.999Z" }, - { url = "https://files.pythonhosted.org/packages/56/40/166447150a31bc3b794ffb493d5a634f67ffbc75dd8b4c46373701b7ef15/cmeel_tinyxml2-11.0.0-0-py3-none-macosx_11_0_arm64.whl", hash = "sha256:9a1406f408262c37ae7c4566b1d67801c4b10c4980903fb1ef0ba45fa4407072", size = 109146, upload-time = "2026-05-21T11:49:27.829Z" }, - { url = "https://files.pythonhosted.org/packages/d8/ca/3cc665afe2d76999f15454bb3b2f7c05f0088ad7de35718648291a536fd9/cmeel_tinyxml2-11.0.0-0-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:6f830007917c3e36f26b27d170ce84a619a62f46104d3cce435dff0125dd665f", size = 157109, upload-time = "2026-05-21T11:49:29.358Z" }, - { url = "https://files.pythonhosted.org/packages/87/4e/dcc0d9756d93be734d824e2a570cc9ac68909a1d7d3b6fc87c2fb32726c0/cmeel_tinyxml2-11.0.0-0-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:18674156bd41f3993dc1d5199da04fa496674358daa6588090fb9f86c71917b0", size = 148825, upload-time = "2026-05-21T11:49:31.035Z" }, + { url = "https://files.pythonhosted.org/packages/fe/5d/bc3a932eb7996a0a789979426a9bb8a3948bf57f3f17bab87dddbef62433/cmeel_tinyxml2-10.0.0-0-py3-none-macosx_10_9_x86_64.whl", hash = "sha256:924499bb1b60b9a17bd001d12a9af88ddbee4ca888638ae684ba7f0f3ce49e87", size = 111913, upload-time = "2025-02-06T10:28:45.723Z" }, + { url = "https://files.pythonhosted.org/packages/92/bf/67d11e123313c034712896e94038291fe506bb099bdb75a136392002ffd0/cmeel_tinyxml2-10.0.0-0-py3-none-macosx_11_0_arm64.whl", hash = "sha256:26a1eb30c2a00bfc172e89ed015a18b8efb2b383546252ca8859574aed684686", size = 109487, upload-time = "2025-02-06T10:28:47.546Z" }, + { url = "https://files.pythonhosted.org/packages/ca/48/d8c81ce19b4b278ed0e8f81f93ae8670209bf3a9ac20141b9c386bb40cc7/cmeel_tinyxml2-10.0.0-0-py3-none-manylinux_2_17_i686.whl", hash = "sha256:53d86e02864c712f51f9a9adfcd8b6046b2ed51d44a0c34a8438d93b72b48325", size = 160118, upload-time = "2025-02-06T10:28:49.627Z" }, + { url = "https://files.pythonhosted.org/packages/87/4e/62193e27c9581f8ba7aeaeca7805632a64f2f4a824b1db37ad02ee953e8a/cmeel_tinyxml2-10.0.0-0-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:74112e2e9473afbf6ee2d25c9942553e9f6a40465e714533db72db48bc7658e1", size = 158477, upload-time = "2025-02-06T10:28:51.667Z" }, + { url = "https://files.pythonhosted.org/packages/14/f9/d0420c39e9ade99beeec61cd3abc68880fe6e14d85e9df292af8fabe65c8/cmeel_tinyxml2-10.0.0-0-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:ecd6e99caa2a06ac0d4b333b740c20fca526d0ca426f99eb5c0a0039117afdb6", size = 147025, upload-time = "2025-02-06T10:28:53.944Z" }, + { url = "https://files.pythonhosted.org/packages/66/9e/df63147fc162ab487217fa5596778ab7a81a82d9b3ce4236fd3a1e48cecb/cmeel_tinyxml2-10.0.0-0-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:30993fffb7032a45d5d3b1e5670cb879dad667a13144cd68c8f4e0371a8a3d2e", size = 150958, upload-time = "2025-02-06T10:28:55.301Z" }, + { url = "https://files.pythonhosted.org/packages/0e/a8/b03567275fd83f5af33ddb61de942689dec72c5b21bec01e6a5b11101aa5/cmeel_tinyxml2-10.0.0-0-py3-none-musllinux_1_2_i686.whl", hash = "sha256:8c09ede51784af54211a6225884dc7ddbb02ea1681656d173060c7ad2a5b9a3c", size = 160300, upload-time = "2025-02-06T10:28:57.189Z" }, + { url = "https://files.pythonhosted.org/packages/3a/ec/2781635b66c1059ca1243ae0f5a0410e171a5d8b8a71be3e34cb172f9f2d/cmeel_tinyxml2-10.0.0-0-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:3bd511d6d0758224efdebc23d3ead6e94f0755b04141ebf7d5493377829e8332", size = 149184, upload-time = "2025-02-06T10:28:58.734Z" }, ] [[package]] name = "cmeel-urdfdom" -version = "6.0.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +version = "4.0.1" +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "cmeel" }, { name = "cmeel-console-bridge" }, { name = "cmeel-tinyxml2" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/21/75/4e8aff079e98582aeeb8e752805081da0c2dea405e79bafeefb555defe9f/cmeel_urdfdom-6.0.0.tar.gz", hash = "sha256:65c0fdc6021300fc55b2d0c03ab64dedc328034a74e40498e671bc894bb1dcf7", size = 303688, upload-time = "2026-05-21T12:08:56.663Z" } +sdist = { url = "https://files.pythonhosted.org/packages/31/09/be81a5e7db56f34b6ccdbe7afe855c95a18c8439e173519e0146e9276a8c/cmeel_urdfdom-4.0.1.tar.gz", hash = "sha256:2e3f41e8483889e195b574acb326a4464cf11a3c0a8724031ac28bcda2223efc", size = 291511, upload-time = "2025-02-12T12:07:09.699Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/f2/40/51ba667135f01631179eee1614557193f8453740f248302d1b8b7f9f693e/cmeel_urdfdom-6.0.0-0-py3-none-macosx_10_9_x86_64.whl", hash = "sha256:53d55cebb137a6e4dac6c16fa53f2dc2b7b9b5cda644bd1637a5bb849cd96e52", size = 381501, upload-time = "2026-05-21T12:08:48.758Z" }, - { url = "https://files.pythonhosted.org/packages/3c/d1/2b49a8c940fa75abc13df9842c14e577e6a82d5854b6d52597ce3bb04894/cmeel_urdfdom-6.0.0-0-py3-none-macosx_11_0_arm64.whl", hash = "sha256:0ef424735bd30f4afa4d1b4ddca9b297498c43005ddd775c080e55f62e9e0466", size = 377159, upload-time = "2026-05-21T12:08:50.485Z" }, - { url = "https://files.pythonhosted.org/packages/db/ac/0efde3a48220b55707bafb6d2e2dcca562f99dcd5c2c15311f7696eeacce/cmeel_urdfdom-6.0.0-0-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:0436f5230f1484c8e583284ef48c7291b230ada3dc5fb2937941f582e72409ec", size = 506000, upload-time = "2026-05-21T12:08:52.273Z" }, - { url = "https://files.pythonhosted.org/packages/32/d4/dfd617e598100e4e53ae3d228a968facff80bae53038fb18e2dccb1ab03a/cmeel_urdfdom-6.0.0-0-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:7ab1be680a8ec866d5422c617b641d1f0e38774061df28b8b426fb26edce6337", size = 530049, upload-time = "2026-05-21T12:08:54.224Z" }, + { url = "https://files.pythonhosted.org/packages/be/d0/20147dd6bb723afc44a58d89ea624df2bad1bed7b898a2df112aaca4a479/cmeel_urdfdom-4.0.1-0-py3-none-macosx_10_9_x86_64.whl", hash = "sha256:2fe56939c6b47f6ec57021aac154123da47ecdcd79a217f3a5e3c4b705a07dee", size = 300860, upload-time = "2025-02-12T12:06:58.536Z" }, + { url = "https://files.pythonhosted.org/packages/8e/98/f832bca347e2d987c6b0ebb6930caf7b2c402535324aeed466b6aa2c4513/cmeel_urdfdom-4.0.1-0-py3-none-macosx_11_0_arm64.whl", hash = "sha256:00a0aba78b68c428b27abeed1db58d73e65319ed966911a0e97b37367442e756", size = 300616, upload-time = "2025-02-12T12:07:00.556Z" }, + { url = "https://files.pythonhosted.org/packages/cf/10/bf5765b6f388037cff166a754a0958ac2fee34ca3c0975ef64d0324e4647/cmeel_urdfdom-4.0.1-0-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:a701a8f9671331f11b18ecf37a6537db546a21e6a0e5d0ff53341fea0693ed7f", size = 385951, upload-time = "2025-02-12T12:07:02.556Z" }, + { url = "https://files.pythonhosted.org/packages/c3/82/cb3f8f587d293a17bdbea15b50cdaa4a1e28e04583eb4cb4821685b89466/cmeel_urdfdom-4.0.1-0-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:12e39fc388c077d79fc9b3841d3d972a1da90b90de754d3363194c1540e18abf", size = 399619, upload-time = "2025-02-12T12:07:04.388Z" }, + { url = "https://files.pythonhosted.org/packages/24/77/322d7ac92c692d8dfaeda9de2d937087d15e2b564dc457d656e5fde3991d/cmeel_urdfdom-4.0.1-0-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:c4a83925df1d5923c4485c3eb2b80b3a61b14f119ab724fb5bd04cec494690ee", size = 373969, upload-time = "2025-02-12T12:07:06.222Z" }, + { url = "https://files.pythonhosted.org/packages/9f/63/bdc6b55cc8bd99bb9dce6be801b30feffaa1c3841ecb7f4fe4d137424518/cmeel_urdfdom-4.0.1-0-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:4c4f44270971b3d05c45a4e21b1fb2df7e05a750363ae918f59532bff0bfe0e1", size = 388237, upload-time = "2025-02-12T12:07:08.326Z" }, + { url = "https://files.pythonhosted.org/packages/1d/2d/8463fc23230612daf4da1e31d3229f47708381f3ae4d1500f0f007ac0f92/cmeel_urdfdom-4.0.1-1-py3-none-macosx_10_9_x86_64.whl", hash = "sha256:f7535158f45992eb2ba79e90d9db1bf9adc3846d9c7ed3e7a8c1c4d5343afa37", size = 301006, upload-time = "2025-02-13T11:42:08.8Z" }, + { url = "https://files.pythonhosted.org/packages/0f/d5/c8cdf500e49300d85624cbc3ef804107ddcdc9c541b1d3f726bfb58a9fc1/cmeel_urdfdom-4.0.1-1-py3-none-macosx_11_0_arm64.whl", hash = "sha256:fef2a01a00d61d41b3d35dd4958bba973e9025c26eea1d3c9880932f4dba89a5", size = 300758, upload-time = "2025-02-13T11:42:10.449Z" }, + { url = "https://files.pythonhosted.org/packages/cf/b3/2f7bac1544113a7f8e0f6d8b1fab5e75c6a3d27ffbb584b03267251b2165/cmeel_urdfdom-4.0.1-1-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:7a52eb36950ce982014d99a55717ca29985da056e3705f20746f15d3244c1f7a", size = 386043, upload-time = "2025-02-13T11:42:11.923Z" }, + { url = "https://files.pythonhosted.org/packages/86/03/8bdeb36ba6a3e8125d523ecfc010403049e463fe589f9896858d4bdcaf1e/cmeel_urdfdom-4.0.1-1-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:9f3b9c80b10d7246821ff61c2573f799e3da23d483e6f7367ddcad8a48baf58f", size = 399719, upload-time = "2025-02-13T11:42:14.325Z" }, + { url = "https://files.pythonhosted.org/packages/3f/ed/43f99e7512460294cd8acc5753ba25f8a20bdf28d62e143eaf3ec7a28bb6/cmeel_urdfdom-4.0.1-1-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:2de69f47e8312cc09157624802d5bdaad6406443f863fb4b9ec62a19b4de3c72", size = 374073, upload-time = "2025-02-13T11:42:17.907Z" }, + { url = "https://files.pythonhosted.org/packages/17/c6/2e9bde6d7c02c1cf203ea896f8ce1afd441412f09b44830f1ee4a96d77de/cmeel_urdfdom-4.0.1-1-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:7708c1402de450fbeab21f7ca264a9a4676ed4c1cdf8d84d840bc5d057aac920", size = 388337, upload-time = "2025-02-13T11:42:19.657Z" }, ] [[package]] name = "cmeel-zlib" version = "1.3.2" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "cmeel" }, ] @@ -818,7 +830,7 @@ wheels = [ [[package]] name = "coal-library" version = "3.0.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "cmeel" }, { name = "cmeel-assimp" }, @@ -842,7 +854,7 @@ wheels = [ [[package]] name = "colorama" version = "0.4.6" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/d8/53/6f443c9a4a8358a93a6792e2acffb9d9d5cb0a5cfd8802644b7b1c9a02e4/colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44", size = 27697, upload-time = "2022-10-25T02:36:22.414Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6", size = 25335, upload-time = "2022-10-25T02:36:20.889Z" }, @@ -851,7 +863,7 @@ wheels = [ [[package]] name = "comm" version = "0.2.3" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/4c/13/7d740c5849255756bc17888787313b61fd38a0a8304fc4f073dfc46122aa/comm-0.2.3.tar.gz", hash = "sha256:2dc8048c10962d55d7ad693be1e7045d891b7ce8d999c97963a5e3e99c055971", size = 6319, upload-time = "2025-07-25T14:02:04.452Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/60/97/891a0971e1e4a8c5d2b20bbe0e524dc04548d2307fee33cdeba148fd4fc7/comm-0.2.3-py3-none-any.whl", hash = "sha256:c615d91d75f7f04f095b30d1c1711babd43bdc6419c1be9886a85f2f4e489417", size = 7294, upload-time = "2025-07-25T14:02:02.896Z" }, @@ -860,7 +872,7 @@ wheels = [ [[package]] name = "contourpy" version = "1.3.3" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "numpy" }, ] @@ -926,7 +938,7 @@ wheels = [ [[package]] name = "coverage" version = "7.14.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/54/fd/0ab2772530e946e1be1abd0bc09e647ec9b02e88f0867857601fefca8953/coverage-7.14.1.tar.gz", hash = "sha256:30c08f7d90415aa98b3c990385dea2939b0da55f38515e5b369b83655f8523be", size = 920132, upload-time = "2026-05-26T20:41:36.783Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/3d/b7/bdbb725ba02c5b42825b200c940f38b7a54fcad24627b7192f78f8110d76/coverage-7.14.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:a06c76364a9360e33d6d23769aefdf7f66f38e2ffb60ceb1baaa4989d83b695c", size = 220022, upload-time = "2026-05-26T20:39:03.702Z" }, @@ -1010,7 +1022,7 @@ wheels = [ [[package]] name = "cuda-bindings" version = "12.9.7" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "cuda-pathfinder", marker = "sys_platform == 'linux'" }, ] @@ -1030,7 +1042,7 @@ wheels = [ [[package]] name = "cuda-pathfinder" version = "1.5.5" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } wheels = [ { url = "https://files.pythonhosted.org/packages/11/c8/26f2e4aae92f11522a96043892ba39a90eac610d5242523aa863212bc1c7/cuda_pathfinder-1.5.5-py3-none-any.whl", hash = "sha256:0228c023f95d1480f143ef5c8922d27a2ab052087a942e81dc289c9eb8f91689", size = 51671, upload-time = "2026-05-27T01:21:25.413Z" }, ] @@ -1038,7 +1050,7 @@ wheels = [ [[package]] name = "cuda-toolkit" version = "12.8.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } wheels = [ { url = "https://files.pythonhosted.org/packages/d4/c8/7dce3a0b15b42a3b58e7d96eb22a687d3bf2c44e01d149a6874629cd9938/cuda_toolkit-12.8.1-py2.py3-none-any.whl", hash = "sha256:adc7906af4ecbf9a352f9dca5734eceb21daec281ccfcf5675e1d2f724fc2cba", size = 2283, upload-time = "2025-08-13T02:03:07.842Z" }, ] @@ -1081,7 +1093,7 @@ nvtx = [ [[package]] name = "cycler" version = "0.12.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/a9/95/a3dbbb5028f35eafb79008e7522a75244477d2838f38cbb722248dabc2a8/cycler-0.12.1.tar.gz", hash = "sha256:88bb128f02ba341da8ef447245a9e138fae777f6a23943da4540077d3601eb1c", size = 7615, upload-time = "2023-10-07T05:32:18.335Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/e7/05/c19819d5e3d95294a6f5947fb9b9629efb316b96de511b418c53d245aae6/cycler-0.12.1-py3-none-any.whl", hash = "sha256:85cef7cff222d8644161529808465972e51340599459b8ac3ccbac5a854e0d30", size = 8321, upload-time = "2023-10-07T05:32:16.783Z" }, @@ -1090,7 +1102,7 @@ wheels = [ [[package]] name = "datasets" version = "4.8.5" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "dill" }, { name = "filelock" }, @@ -1115,7 +1127,7 @@ wheels = [ [[package]] name = "debugpy" version = "1.8.21" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/f2/aa/12037145b7a56eaa5b29b41872f7a21b538e807e13f32c4d3c46e59be084/debugpy-1.8.21.tar.gz", hash = "sha256:a3c53278e84c94e11bd87c53970ec391d1a67396c8b22609fcac576520e611a6", size = 1697577, upload-time = "2026-06-01T19:30:35.156Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/a2/df/bf625547431a9cadc9f4cbfeda38866e2b17f6aed147b625377e87834449/debugpy-1.8.21-cp312-cp312-macosx_15_0_universal2.whl", hash = "sha256:9f96713896f39c3dff0ee841f47320c3f2983d33c341e009361bb0ebc79adc4e", size = 2483609, upload-time = "2026-06-01T19:30:50.794Z" }, @@ -1136,7 +1148,7 @@ wheels = [ [[package]] name = "decorator" version = "5.3.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/60/8b/32f9823da46cde7df2087faa08cd98d01b908f8dcab982cdba9c84e85355/decorator-5.3.1.tar.gz", hash = "sha256:4cbcdd55a6efadb9dbea26b858f4fb3264567b52d69ca0d25b721b553f60ea82", size = 58084, upload-time = "2026-05-18T06:03:28.057Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/05/7f/798705f5296a58ca505d600456748d1be48078eac8a7050d8a98bc9edb89/decorator-5.3.1-py3-none-any.whl", hash = "sha256:f47fe6fdbd2edd623ecfe36875d37aba411624e2670dd395dddae1358689bb3c", size = 10365, upload-time = "2026-05-18T06:03:26.517Z" }, @@ -1145,7 +1157,7 @@ wheels = [ [[package]] name = "decord" version = "0.6.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "numpy", marker = "(platform_machine != 'arm64' and platform_machine != 's390x' and sys_platform == 'darwin') or (platform_machine == 'AMD64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine != 's390x' and sys_platform != 'darwin' and sys_platform != 'linux')" }, ] @@ -1157,7 +1169,7 @@ wheels = [ [[package]] name = "deepdiff" version = "8.6.2" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "orderly-set" }, ] @@ -1169,7 +1181,7 @@ wheels = [ [[package]] name = "defusedxml" version = "0.7.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/0f/d5/c66da9b79e5bdb124974bfe172b4daf3c984ebd9c2a06e2b8a4dc7331c72/defusedxml-0.7.1.tar.gz", hash = "sha256:1bb3032db185915b62d7c6209c5a8792be6a32ab2fedacc84e01b52c51aa3e69", size = 75520, upload-time = "2021-03-08T10:59:26.269Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/07/6c/aa3f2f849e01cb6a001cd8554a88d4c77c5c1a31c95bdf1cf9301e6d9ef4/defusedxml-0.7.1-py2.py3-none-any.whl", hash = "sha256:a352e7e428770286cc899e2542b6cdaedb2b4953ff269a210103ec58f6198a61", size = 25604, upload-time = "2021-03-08T10:59:24.45Z" }, @@ -1178,7 +1190,7 @@ wheels = [ [[package]] name = "diffusers" version = "0.35.2" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "filelock" }, { name = "huggingface-hub" }, @@ -1197,7 +1209,7 @@ wheels = [ [[package]] name = "dill" version = "0.4.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/81/e1/56027a71e31b02ddc53c7d65b01e68edf64dea2932122fe7746a516f75d5/dill-0.4.1.tar.gz", hash = "sha256:423092df4182177d4d8ba8290c8a5b640c66ab35ec7da59ccfa00f6fa3eea5fa", size = 187315, upload-time = "2026-01-19T02:36:56.85Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/1e/77/dc8c558f7593132cf8fefec57c4f60c83b16941c574ac5f619abb3ae7933/dill-0.4.1-py3-none-any.whl", hash = "sha256:1e1ce33e978ae97fcfcff5638477032b801c46c7c65cf717f95fbc2248f79a9d", size = 120019, upload-time = "2026-01-19T02:36:55.663Z" }, @@ -1206,7 +1218,7 @@ wheels = [ [[package]] name = "distlib" version = "0.4.3" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/c9/02/bd72be9134d25ed783ecbbc38a539ffaefbf90c78418c7fb7229600dbac7/distlib-0.4.3.tar.gz", hash = "sha256:f152097224a0ae24be5a0f6bae1b9359af82133bce63f98a95f86cae1aede9ed", size = 615141, upload-time = "2026-06-12T08:04:52.847Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/02/08/9c41fb51ab5b43eb21674aff13df270e8ba6c4b29c8624e328dc7a9482af/distlib-0.4.3-py2.py3-none-any.whl", hash = "sha256:4b0ce306c966eb73bc3a7b6abad017c556dadd92c44701562cd528ac7fde4d5b", size = 470628, upload-time = "2026-06-12T08:04:50.506Z" }, @@ -1215,7 +1227,7 @@ wheels = [ [[package]] name = "distro" version = "1.9.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/fc/f8/98eea607f65de6527f8a2e8885fc8015d3e6f5775df186e443e0964a11c3/distro-1.9.0.tar.gz", hash = "sha256:2fa77c6fd8940f116ee1d6b94a2f90b13b5ea8d019b98bc8bafdcabcdd9bdbed", size = 60722, upload-time = "2023-12-24T09:54:32.31Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/12/b3/231ffd4ab1fc9d679809f356cebee130ac7daa00d6d6f3206dd4fd137e9e/distro-1.9.0-py3-none-any.whl", hash = "sha256:7bffd925d65168f85027d8da9af6bddab658135b840670a223589bc0c8ef02b2", size = 20277, upload-time = "2023-12-24T09:54:30.421Z" }, @@ -1224,7 +1236,7 @@ wheels = [ [[package]] name = "dm-control" version = "1.0.41" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "absl-py" }, { name = "dm-env" }, @@ -1250,7 +1262,7 @@ wheels = [ [[package]] name = "dm-env" version = "1.6" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "absl-py" }, { name = "dm-tree" }, @@ -1264,7 +1276,7 @@ wheels = [ [[package]] name = "dm-tree" version = "0.1.9" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "absl-py" }, { name = "attrs" }, @@ -1289,13 +1301,13 @@ wheels = [ [[package]] name = "docopt" version = "0.6.2" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/a2/55/8f8cab2afd404cf578136ef2cc5dfb50baa1761b68c9da1fb1e4eed343c9/docopt-0.6.2.tar.gz", hash = "sha256:49b3a825280bd66b3aa83585ef59c4a8c82f2c8a522dbe754a8bc8d08c85c491", size = 25901, upload-time = "2014-06-16T11:18:57.406Z" } [[package]] name = "draccus" version = "0.10.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "mergedeep" }, { name = "pyyaml" }, @@ -1311,7 +1323,7 @@ wheels = [ [[package]] name = "dynamixel-sdk" version = "3.8.4" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "pyserial" }, ] @@ -1323,7 +1335,7 @@ wheels = [ [[package]] name = "easydict" version = "1.13" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/24/9f/d18d6b5e19244788a6d09c14a8406376b4f4bfcc008e6d17a4f4c15362e8/easydict-1.13.tar.gz", hash = "sha256:b1135dedbc41c8010e2bc1f77ec9744c7faa42bce1a1c87416791449d6c87780", size = 6809, upload-time = "2024-03-04T12:04:41.251Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/05/ec/fa6963f1198172c2b75c9ab6ecefb3045991f92f75f5eb41b6621b198123/easydict-1.13-py3-none-any.whl", hash = "sha256:6b787daf4dcaf6377b4ad9403a5cee5a86adbc0ca9a5bcf5410e9902002aeac2", size = 6804, upload-time = "2024-03-04T12:04:39.508Z" }, @@ -1332,13 +1344,13 @@ wheels = [ [[package]] name = "egl-probe" version = "1.0.2" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/b3/c3/c5fea892e18631721a7d7e492a790bde8e3dc930861866bedb38762d0d7c/egl_probe-1.0.2.tar.gz", hash = "sha256:29bdca7b08da1e060cfb42cd46af8300a7ac4f3b1b2eeb16e545ea16d9a5ac93", size = 217495, upload-time = "2021-07-16T08:50:18.481Z" } [[package]] name = "eigenpy" version = "3.10.3" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "cmeel" }, { name = "cmeel-boost" }, @@ -1358,7 +1370,7 @@ wheels = [ [[package]] name = "einops" version = "0.8.2" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/2c/77/850bef8d72ffb9219f0b1aac23fbc1bf7d038ee6ea666f331fa273031aa2/einops-0.8.2.tar.gz", hash = "sha256:609da665570e5e265e27283aab09e7f279ade90c4f01bcfca111f3d3e13f2827", size = 56261, upload-time = "2026-01-26T04:13:17.638Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/2a/09/f8d8f8f31e4483c10a906437b4ce31bdf3d6d417b73fe33f1a8b59e34228/einops-0.8.2-py3-none-any.whl", hash = "sha256:54058201ac7087911181bfec4af6091bb59380360f069276601256a76af08193", size = 65638, upload-time = "2026-01-26T04:13:18.546Z" }, @@ -1367,7 +1379,7 @@ wheels = [ [[package]] name = "eiquadprog" version = "1.2.9" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "cmeel" }, { name = "cmeel-boost" }, @@ -1385,7 +1397,7 @@ wheels = [ [[package]] name = "etils" version = "1.14.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/26/ce/6e067242fde898841922ac6fc82b0bb2fe35c38e995880bdffdfbe30182a/etils-1.14.0.tar.gz", hash = "sha256:8136e7f4c4173cd0af0ca5481c4475152f0b8686192951eefa60ee8711e1ede4", size = 108127, upload-time = "2026-03-04T17:41:36.291Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/5a/3d/589663aeeacd59bb2f3e8596bfd3e81cf0fb18d70bb433199041f469771b/etils-1.14.0-py3-none-any.whl", hash = "sha256:b5df7341f54dbe1405a4450b2741207b4a8c279780402b45f87202b94dfc52b4", size = 172934, upload-time = "2026-03-04T17:41:35.01Z" }, @@ -1401,13 +1413,13 @@ epath = [ [[package]] name = "evdev" version = "1.9.3" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/a5/f5/397b61091120a9ca5001041dd7bf76c385b3bfd67a0e5bcb74b852bd22a4/evdev-1.9.3.tar.gz", hash = "sha256:2c140e01ac8437758fa23fe5c871397412461f42d421aa20241dc8fe8cfccbc9", size = 32723, upload-time = "2026-02-05T21:54:24.987Z" } [[package]] name = "executing" version = "2.2.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/cc/28/c14e053b6762b1044f34a13aab6859bbf40456d37d23aa286ac24cfd9a5d/executing-2.2.1.tar.gz", hash = "sha256:3632cc370565f6648cc328b32435bd120a1e4ebb20c77e3fdde9a13cd1e533c4", size = 1129488, upload-time = "2025-09-01T09:48:10.866Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/c1/ea/53f2148663b321f21b5a606bd5f191517cf40b7072c0497d3c92c4a13b1e/executing-2.2.1-py2.py3-none-any.whl", hash = "sha256:760643d3452b4d777d295bb167ccc74c64a81df23fb5e08eff250c425a4b2017", size = 28317, upload-time = "2025-09-01T09:48:08.5Z" }, @@ -1416,7 +1428,7 @@ wheels = [ [[package]] name = "faker" version = "34.0.2" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "python-dateutil" }, { name = "typing-extensions" }, @@ -1429,7 +1441,7 @@ wheels = [ [[package]] name = "farama-notifications" version = "0.0.6" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/ec/91/14397890dde30adc4bee6462158933806207bc5dd10d7b4d09d5c33845cf/farama_notifications-0.0.6.tar.gz", hash = "sha256:b19acac4bb41d76e59e03394b5dd165f4761c86fa327f56307a35cbee3b60158", size = 2517, upload-time = "2026-04-24T08:43:57.603Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/7c/f0/21f81892e4ed10f4ec3ef2e7cf8635fb76e7c0907c55d0da66be50094760/farama_notifications-0.0.6-py3-none-any.whl", hash = "sha256:f84839188efa1ce5bb361c2a84881b2dc2c0d0d7fb661ff00421820170930935", size = 2897, upload-time = "2026-04-24T08:43:56.785Z" }, @@ -1438,7 +1450,7 @@ wheels = [ [[package]] name = "fastapi" version = "0.137.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "annotated-doc" }, { name = "pydantic" }, @@ -1454,7 +1466,7 @@ wheels = [ [[package]] name = "fastjsonschema" version = "2.21.2" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/20/b5/23b216d9d985a956623b6bd12d4086b60f0059b27799f23016af04a74ea1/fastjsonschema-2.21.2.tar.gz", hash = "sha256:b1eb43748041c880796cd077f1a07c3d94e93ae84bba5ed36800a33554ae05de", size = 374130, upload-time = "2025-08-14T18:49:36.666Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/cb/a8/20d0723294217e47de6d9e2e40fd4a9d2f7c4b6ef974babd482a59743694/fastjsonschema-2.21.2-py3-none-any.whl", hash = "sha256:1c797122d0a86c5cace2e54bf4e819c36223b552017172f32c5c024a6b77e463", size = 24024, upload-time = "2025-08-14T18:49:34.776Z" }, @@ -1463,7 +1475,7 @@ wheels = [ [[package]] name = "feetech-servo-sdk" version = "1.0.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "pyserial" }, ] @@ -1472,7 +1484,7 @@ sdist = { url = "https://files.pythonhosted.org/packages/5f/8e/c53d6f9a8bf3a86a6 [[package]] name = "filelock" version = "3.29.4" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/e6/dc/be6cbe99670cd6e4ad387123647cb08e0c32975e223f82551e914c5568a6/filelock-3.29.4.tar.gz", hash = "sha256:10cdb3656fc44541cdf30652a93fb10ec6b05325620eb316bd26893e4201538a", size = 63028, upload-time = "2026-06-13T16:12:00.744Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/13/37/a065dc3bd6e49423a6532c642ca7378d3f467b1ef44c2800c937af7f9739/filelock-3.29.4-py3-none-any.whl", hash = "sha256:dac1648087d5115554850d113e7dd8c83ab2d38e3435dde2d4f163847e57b767", size = 42757, upload-time = "2026-06-13T16:11:59.582Z" }, @@ -1481,10 +1493,10 @@ wheels = [ [[package]] name = "flash-attn" version = "2.8.3.post1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "einops", marker = "platform_machine != 'arm64' or sys_platform != 'darwin'" }, - { name = "torch", version = "2.11.0", source = { registry = "https://pypi.registries.huggingface.tech/" }, marker = "(platform_machine != 'arm64' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'linux')" }, + { name = "torch", version = "2.11.0", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine != 'arm64' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'linux')" }, { name = "torch", version = "2.11.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" }, marker = "sys_platform == 'linux'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/01/7a/92a46e7cd6bbb4d7b2855a457c3b855df54a97af5656d98fc92e58e61065/flash_attn-2.8.3.post1.tar.gz", hash = "sha256:55d5103ed846da8b56e0797acf4bde07dee4b1c7e8907fcfc6699c203030c348", size = 8451657, upload-time = "2026-06-11T04:35:21.318Z" } @@ -1492,7 +1504,7 @@ sdist = { url = "https://files.pythonhosted.org/packages/01/7a/92a46e7cd6bbb4d7b [[package]] name = "flatbuffers" version = "25.12.19" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } wheels = [ { url = "https://files.pythonhosted.org/packages/e8/2d/d2a548598be01649e2d46231d151a6c56d10b964d94043a335ae56ea2d92/flatbuffers-25.12.19-py2.py3-none-any.whl", hash = "sha256:7634f50c427838bb021c2d66a3d1168e9d199b0607e6329399f04846d42e20b4", size = 26661, upload-time = "2025-12-19T23:16:13.622Z" }, ] @@ -1500,7 +1512,7 @@ wheels = [ [[package]] name = "fonttools" version = "4.63.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/84/69/c97f2c18e0db87d2c7b15da1974dace76ae938f1cfa22e2727a648b7ed43/fonttools-4.63.0.tar.gz", hash = "sha256:caeb583deeb5168e694b65cda8b4ee62abedfa66cf88488734466f2366b9c4e0", size = 3597189, upload-time = "2026-05-14T12:04:30.958Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/08/ef/b3c6b9b5be2f82416d73fe2ed2e96e2793cd80e7510bd6a17ca79cdd88ec/fonttools-4.63.0-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:37dd23e621e3b0aef1baa70a303b80aaf38449632cfc8fd2a55fb285bbccfc02", size = 2881131, upload-time = "2026-05-14T12:03:13.386Z" }, @@ -1541,7 +1553,7 @@ wheels = [ [[package]] name = "fqdn" version = "1.5.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/30/3e/a80a8c077fd798951169626cde3e239adeba7dab75deb3555716415bd9b0/fqdn-1.5.1.tar.gz", hash = "sha256:105ed3677e767fb5ca086a0c1f4bb66ebc3c100be518f0e0d755d9eae164d89f", size = 6015, upload-time = "2021-03-11T07:16:29.08Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/cf/58/8acf1b3e91c58313ce5cb67df61001fc9dcd21be4fadb76c1a2d540e09ed/fqdn-1.5.1-py3-none-any.whl", hash = "sha256:3a179af3761e4df6eb2e026ff9e1a3033d3587bf980a0b1b2e1e5d08d7358014", size = 9121, upload-time = "2021-03-11T07:16:28.351Z" }, @@ -1550,7 +1562,7 @@ wheels = [ [[package]] name = "frozenlist" version = "1.8.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/2d/f5/c831fac6cc817d26fd54c7eaccd04ef7e0288806943f7cc5bbf69f3ac1f0/frozenlist-1.8.0.tar.gz", hash = "sha256:3ede829ed8d842f6cd48fc7081d7a41001a56f1f38603f9d49bf3020d59a31ad", size = 45875, upload-time = "2025-10-06T05:38:17.865Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/69/29/948b9aa87e75820a38650af445d2ef2b6b8a6fab1a23b6bb9e4ef0be2d59/frozenlist-1.8.0-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:78f7b9e5d6f2fdb88cdde9440dc147259b62b9d3b019924def9f6478be254ac1", size = 87782, upload-time = "2025-10-06T05:36:06.649Z" }, @@ -1639,7 +1651,7 @@ wheels = [ [[package]] name = "fsspec" version = "2026.2.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/51/7c/f60c259dcbf4f0c47cc4ddb8f7720d2dcdc8888c8e5ad84c73ea4531cc5b/fsspec-2026.2.0.tar.gz", hash = "sha256:6544e34b16869f5aacd5b90bdf1a71acb37792ea3ddf6125ee69a22a53fb8bff", size = 313441, upload-time = "2026-02-05T21:50:53.743Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/e6/ab/fb21f4c939bb440104cc2b396d3be1d9b7a9fd3c6c2a53d98c45b3d7c954/fsspec-2026.2.0-py3-none-any.whl", hash = "sha256:98de475b5cb3bd66bedd5c4679e87b4fdfe1a3bf4d707b151b3c07e58c9a2437", size = 202505, upload-time = "2026-02-05T21:50:51.819Z" }, @@ -1653,7 +1665,7 @@ http = [ [[package]] name = "future" version = "1.0.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/a7/b2/4140c69c6a66432916b26158687e821ba631a4c9273c474343badf84d3ba/future-1.0.0.tar.gz", hash = "sha256:bd2968309307861edae1458a4f8a4f3598c03be43b97521076aebf5d94c07b05", size = 1228490, upload-time = "2024-02-21T11:52:38.461Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/da/71/ae30dadffc90b9006d77af76b393cb9dfbfc9629f339fc1574a1c52e6806/future-1.0.0-py3-none-any.whl", hash = "sha256:929292d34f5872e70396626ef385ec22355a1fae8ad29e1a734c3e43f9fbc216", size = 491326, upload-time = "2024-02-21T11:52:35.956Z" }, @@ -1662,7 +1674,7 @@ wheels = [ [[package]] name = "gitdb" version = "4.0.12" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "smmap" }, ] @@ -1674,7 +1686,7 @@ wheels = [ [[package]] name = "gitpython" version = "3.1.50" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "gitdb" }, ] @@ -1686,7 +1698,7 @@ wheels = [ [[package]] name = "glfw" version = "2.10.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/96/72/642d4f12f61816ac96777f7360d413e3977a7dd08237d196f02da681b186/glfw-2.10.0.tar.gz", hash = "sha256:801e55d8581b34df9aa2cfea43feb06ff617576e2a8cc5dac23ee75b26d10abe", size = 31475, upload-time = "2025-09-12T08:54:38.871Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/3d/1f/a9ce08b1173b0ab625ee92f0c47a5278b3e76fd367699880d8ee7d56c338/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.p39.p310.p311.p312.p313-none-macosx_10_6_intel.whl", hash = "sha256:5f365a8c94bcea71ec91327e7c16e7cf739128479a18b8c1241b004b40acc412", size = 105329, upload-time = "2025-09-12T08:54:27.938Z" }, @@ -1718,7 +1730,7 @@ wheels = [ [[package]] name = "grpcio" version = "1.73.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/79/e8/b43b851537da2e2f03fa8be1aef207e5cbfb1a2e014fbb6b40d24c177cd3/grpcio-1.73.1.tar.gz", hash = "sha256:7fce2cd1c0c1116cf3850564ebfc3264fba75d3c74a7414373f1238ea365ef87", size = 12730355, upload-time = "2025-06-26T01:53:24.622Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/b8/41/456caf570c55d5ac26f4c1f2db1f2ac1467d5bf3bcd660cba3e0a25b195f/grpcio-1.73.1-cp312-cp312-linux_armv7l.whl", hash = "sha256:921b25618b084e75d424a9f8e6403bfeb7abef074bb6c3174701e0f2542debcf", size = 5334621, upload-time = "2025-06-26T01:52:23.602Z" }, @@ -1746,7 +1758,7 @@ wheels = [ [[package]] name = "grpcio-tools" version = "1.73.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "grpcio" }, { name = "protobuf" }, @@ -1779,7 +1791,7 @@ wheels = [ [[package]] name = "gym-aloha" version = "0.1.4" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "dm-control" }, { name = "gymnasium" }, @@ -1794,7 +1806,7 @@ wheels = [ [[package]] name = "gym-hil" version = "0.1.14" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "gymnasium" }, { name = "hidapi" }, @@ -1811,7 +1823,7 @@ wheels = [ [[package]] name = "gym-pusht" version = "0.1.6" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "gymnasium" }, { name = "opencv-python" }, @@ -1828,7 +1840,7 @@ wheels = [ [[package]] name = "gymnasium" version = "1.3.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "cloudpickle" }, { name = "farama-notifications" }, @@ -1843,7 +1855,7 @@ wheels = [ [[package]] name = "h11" version = "0.16.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/01/ee/02a2c011bdab74c6fb3c75474d40b3052059d95df7e73351460c8588d963/h11-0.16.0.tar.gz", hash = "sha256:4e35b956cf45792e4caa5885e69fba00bdbc6ffafbfa020300e549b208ee5ff1", size = 101250, upload-time = "2025-04-24T03:35:25.427Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/04/4b/29cac41a4d98d144bf5f6d33995617b185d14b22401f75ca86f384e87ff1/h11-0.16.0-py3-none-any.whl", hash = "sha256:63cf8bbe7522de3bf65932fda1d9c2772064ffb3dae62d55932da54b31cb6c86", size = 37515, upload-time = "2025-04-24T03:35:24.344Z" }, @@ -1852,7 +1864,7 @@ wheels = [ [[package]] name = "h5py" version = "3.16.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "numpy", marker = "sys_platform == 'linux'" }, ] @@ -1879,7 +1891,7 @@ wheels = [ [[package]] name = "hebi-py" version = "2.11.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "numpy" }, { name = "pyyaml" }, @@ -1890,13 +1902,13 @@ sdist = { url = "https://files.pythonhosted.org/packages/ce/54/eeee3e2163d6c82de [[package]] name = "hf-egl-probe" version = "1.0.2" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/e6/3e/ffad88145b342d5a972711d0b47e202f2e51c9a436d9b6bdd55f10ed893e/hf_egl_probe-1.0.2.tar.gz", hash = "sha256:0cbc6352236f5c14f681328e2e0d2cb4d101828c81aa58793dbc9a91feee5e93", size = 217668, upload-time = "2025-11-03T15:58:15.592Z" } [[package]] name = "hf-libero" version = "0.1.4" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "bddl", marker = "sys_platform == 'linux'" }, { name = "cloudpickle", marker = "sys_platform == 'linux'" }, @@ -1924,7 +1936,7 @@ wheels = [ [[package]] name = "hf-xet" version = "1.5.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/4b/2d/57fd21d84d93efb4bd0b962383790e19dd1bc053501b4264c97903b4e83e/hf_xet-1.5.1.tar.gz", hash = "sha256:51ef4500dab3764b41135ee1381a4b62ce56fc54d4c92b719b59e597d6df5bf6", size = 876636, upload-time = "2026-06-08T23:02:53.897Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/64/ee/dd9ba7beae1005e54131b7d45263cc74c8a066d47d354e6d58ae9445a388/hf_xet-1.5.1-cp313-cp313t-macosx_10_12_x86_64.whl", hash = "sha256:dbf48c0d02cf0b2e568944330c60d9120c272dabe013bd892d48e25bc6797577", size = 4069485, upload-time = "2026-06-08T23:02:13.193Z" }, @@ -1956,7 +1968,7 @@ wheels = [ [[package]] name = "hidapi" version = "0.14.0.post4" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "setuptools" }, ] @@ -1987,7 +1999,7 @@ wheels = [ [[package]] name = "httpcore" version = "1.0.9" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "certifi" }, { name = "h11" }, @@ -2000,7 +2012,7 @@ wheels = [ [[package]] name = "httptools" version = "0.8.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/43/e5/d471fcb0e14523fe1c3f4ba58ca52480e7bd70ad7109a3846bc75892f7fb/httptools-0.8.0.tar.gz", hash = "sha256:6b2a32f18d97e16e90827d7a819ffa8dbd8cc245fc4e1fa9d1095b54ef4bd999", size = 271342, upload-time = "2026-05-25T22:17:48.841Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/14/88/1d21a36da8f5cb0fa49eafd4b169eba5608d57e75bbcf61845cbc6243216/httptools-0.8.0-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:880490234c10f70a9830743097e8958d6e4b9f5a0ffc24515023afeef984054d", size = 208247, upload-time = "2026-05-25T22:17:07.843Z" }, @@ -2036,7 +2048,7 @@ wheels = [ [[package]] name = "httpx" version = "0.28.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "anyio" }, { name = "certifi" }, @@ -2051,7 +2063,7 @@ wheels = [ [[package]] name = "huggingface-hub" version = "1.19.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "click" }, { name = "filelock" }, @@ -2072,7 +2084,7 @@ wheels = [ [[package]] name = "hydra-core" version = "1.3.3" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "antlr4-python3-runtime", marker = "sys_platform == 'linux'" }, { name = "omegaconf", marker = "sys_platform == 'linux'" }, @@ -2086,7 +2098,7 @@ wheels = [ [[package]] name = "identify" version = "2.6.19" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/52/63/51723b5f116cc04b061cb6f5a561790abf249d25931d515cd375e063e0f4/identify-2.6.19.tar.gz", hash = "sha256:6be5020c38fcb07da56c53733538a3081ea5aa70d36a156f83044bfbf9173842", size = 99567, upload-time = "2026-04-17T18:39:50.265Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/94/84/d9273cd09688070a6523c4aee4663a8538721b2b755c4962aafae0011e72/identify-2.6.19-py2.py3-none-any.whl", hash = "sha256:20e6a87f786f768c092a721ad107fc9df0eb89347be9396cadf3f4abbd1fb78a", size = 99397, upload-time = "2026-04-17T18:39:49.221Z" }, @@ -2095,7 +2107,7 @@ wheels = [ [[package]] name = "idna" version = "3.18" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/cd/63/9496c57188a2ee585e0f1db071d75089a11e98aa86eb99d9d7618fc1edce/idna-3.18.tar.gz", hash = "sha256:ffb385a7e039654cef1ab9ef32c6fafe283c0c0467bba1d9029738ce4a14a848", size = 196711, upload-time = "2026-06-02T14:34:07.794Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/1e/5e/d4e9f1a599fb8e573b7b87160658329fbf28d19eac2718f51fc3def3aa5a/idna-3.18-py3-none-any.whl", hash = "sha256:7f952cbe720b688055e3f87de14f5c3e5fdaa8bc3928985c4077ca689de849a2", size = 65455, upload-time = "2026-06-02T14:34:06.319Z" }, @@ -2104,7 +2116,7 @@ wheels = [ [[package]] name = "imageio" version = "2.37.3" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "numpy" }, { name = "pillow" }, @@ -2123,7 +2135,7 @@ ffmpeg = [ [[package]] name = "imageio-ffmpeg" version = "0.6.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/44/bd/c3343c721f2a1b0c9fc71c1aebf1966a3b7f08c2eea8ed5437a2865611d6/imageio_ffmpeg-0.6.0.tar.gz", hash = "sha256:e2556bed8e005564a9f925bb7afa4002d82770d6b08825078b7697ab88ba1755", size = 25210, upload-time = "2025-01-16T21:34:32.747Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/da/58/87ef68ac83f4c7690961bce288fd8e382bc5f1513860fc7f90a9c1c1c6bf/imageio_ffmpeg-0.6.0-py3-none-macosx_10_9_intel.macosx_10_9_x86_64.whl", hash = "sha256:9d2baaf867088508d4a3458e61eeb30e945c4ad8016025545f66c4b5aaef0a61", size = 24932969, upload-time = "2025-01-16T21:34:20.464Z" }, @@ -2137,7 +2149,7 @@ wheels = [ [[package]] name = "importlib-metadata" version = "9.0.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "zipp" }, ] @@ -2149,7 +2161,7 @@ wheels = [ [[package]] name = "iniconfig" version = "2.3.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/72/34/14ca021ce8e5dfedc35312d08ba8bf51fdd999c576889fc2c24cb97f4f10/iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730", size = 20503, upload-time = "2025-10-18T21:55:43.219Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" }, @@ -2158,7 +2170,7 @@ wheels = [ [[package]] name = "ipykernel" version = "6.31.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "appnope", marker = "sys_platform == 'darwin'" }, { name = "comm" }, @@ -2182,7 +2194,7 @@ wheels = [ [[package]] name = "ipython" version = "9.14.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "colorama", marker = "sys_platform == 'win32'" }, { name = "decorator" }, @@ -2204,7 +2216,7 @@ wheels = [ [[package]] name = "ipython-pygments-lexers" version = "1.1.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "pygments" }, ] @@ -2216,7 +2228,7 @@ wheels = [ [[package]] name = "ipywidgets" version = "8.1.8" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "comm" }, { name = "ipython" }, @@ -2232,7 +2244,7 @@ wheels = [ [[package]] name = "ischedule" version = "1.2.7" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/78/46/0d471c8318050c182badc6b6c3d4d05c1705508de57175e6e316de76e43c/ischedule-1.2.7.tar.gz", hash = "sha256:517ac2c717c8de7e39ded69a86ab974348e5273e466a7ad2af6a096668ff9ba5", size = 7421, upload-time = "2024-03-15T19:13:12.606Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/94/0c/937caa3558fa5787927585cce67af4446fec8996d4c28c88d7efdbad6b14/ischedule-1.2.7-py3-none-any.whl", hash = "sha256:2a327e57e3ea179d3d9675cef1e7a0719f2a936c61a9d1ae288c3a5f5156939d", size = 5273, upload-time = "2024-03-15T19:13:09.105Z" }, @@ -2241,7 +2253,7 @@ wheels = [ [[package]] name = "isoduration" version = "20.11.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "arrow" }, ] @@ -2253,7 +2265,7 @@ wheels = [ [[package]] name = "jedi" version = "0.20.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "parso" }, ] @@ -2265,7 +2277,7 @@ wheels = [ [[package]] name = "jinja2" version = "3.1.6" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "markupsafe" }, ] @@ -2277,7 +2289,7 @@ wheels = [ [[package]] name = "jiter" version = "0.15.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/66/b5/55f06bb281d92fb3cc86d14e1def2bd908bb77693183e7cb1f5a3c388b0c/jiter-0.15.0.tar.gz", hash = "sha256:4251acc80e2b7c9b7b8823456ea0fceeb0734dac2df7636d3c711b38476b5a76", size = 166640, upload-time = "2026-05-19T10:09:48.361Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/44/53/4f6bddbcde3c71e56d0aa1337ec95950f3d27dd4153e25aadf0feac71751/jiter-0.15.0-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:0e90a1c315a0226ec822d973817967f9223b7701546c8c2a7913e7ab0926294d", size = 308793, upload-time = "2026-05-19T10:07:35.25Z" }, @@ -2349,7 +2361,7 @@ wheels = [ [[package]] name = "json5" version = "0.14.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/9c/4b/6f8906aaf67d501e259b0adab4d312945bb7211e8b8d4dcc77c92320edaa/json5-0.14.0.tar.gz", hash = "sha256:b3f492fad9f6cdbced8b7d40b28b9b1c9701c5f561bef0d33b81c2ff433fefcb", size = 52656, upload-time = "2026-03-27T22:50:48.108Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/b8/42/cf027b4ac873b076189d935b135397675dac80cb29acb13e1ab86ad6c631/json5-0.14.0-py3-none-any.whl", hash = "sha256:56cf861bab076b1178eb8c92e1311d273a9b9acea2ccc82c276abf839ebaef3a", size = 36271, upload-time = "2026-03-27T22:50:47.073Z" }, @@ -2358,7 +2370,7 @@ wheels = [ [[package]] name = "jsonlines" version = "4.0.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "attrs" }, ] @@ -2370,7 +2382,7 @@ wheels = [ [[package]] name = "jsonpointer" version = "3.1.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/18/c7/af399a2e7a67fd18d63c40c5e62d3af4e67b836a2107468b6a5ea24c4304/jsonpointer-3.1.1.tar.gz", hash = "sha256:0b801c7db33a904024f6004d526dcc53bbb8a4a0f4e32bfd10beadf60adf1900", size = 9068, upload-time = "2026-03-23T22:32:32.458Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/9e/6a/a83720e953b1682d2d109d3c2dbb0bc9bf28cc1cbc205be4ef4be5da709d/jsonpointer-3.1.1-py3-none-any.whl", hash = "sha256:8ff8b95779d071ba472cf5bc913028df06031797532f08a7d5b602d8b2a488ca", size = 7659, upload-time = "2026-03-23T22:32:31.568Z" }, @@ -2379,7 +2391,7 @@ wheels = [ [[package]] name = "jsonschema" version = "4.26.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "attrs" }, { name = "jsonschema-specifications" }, @@ -2407,7 +2419,7 @@ format-nongpl = [ [[package]] name = "jsonschema-specifications" version = "2025.9.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "referencing" }, ] @@ -2419,7 +2431,7 @@ wheels = [ [[package]] name = "jupyter" version = "1.1.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "ipykernel" }, { name = "ipywidgets" }, @@ -2436,7 +2448,7 @@ wheels = [ [[package]] name = "jupyter-client" version = "8.9.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "jupyter-core" }, { name = "python-dateutil" }, @@ -2453,7 +2465,7 @@ wheels = [ [[package]] name = "jupyter-console" version = "6.6.3" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "ipykernel" }, { name = "ipython" }, @@ -2472,7 +2484,7 @@ wheels = [ [[package]] name = "jupyter-core" version = "5.9.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "platformdirs" }, { name = "traitlets" }, @@ -2485,7 +2497,7 @@ wheels = [ [[package]] name = "jupyter-events" version = "0.12.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "jsonschema", extra = ["format-nongpl"] }, { name = "packaging" }, @@ -2504,7 +2516,7 @@ wheels = [ [[package]] name = "jupyter-lsp" version = "2.3.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "jupyter-server" }, ] @@ -2516,7 +2528,7 @@ wheels = [ [[package]] name = "jupyter-server" version = "2.19.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "anyio" }, { name = "argon2-cffi" }, @@ -2545,7 +2557,7 @@ wheels = [ [[package]] name = "jupyter-server-terminals" version = "0.5.4" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "pywinpty", marker = "(os_name == 'nt' and platform_machine != 'arm64' and sys_platform == 'darwin') or (os_name == 'nt' and sys_platform != 'darwin' and sys_platform != 'linux')" }, { name = "terminado" }, @@ -2558,7 +2570,7 @@ wheels = [ [[package]] name = "jupyterlab" version = "4.5.8" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "async-lru" }, { name = "httpx" }, @@ -2582,7 +2594,7 @@ wheels = [ [[package]] name = "jupyterlab-pygments" version = "0.3.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/90/51/9187be60d989df97f5f0aba133fa54e7300f17616e065d1ada7d7646b6d6/jupyterlab_pygments-0.3.0.tar.gz", hash = "sha256:721aca4d9029252b11cfa9d185e5b5af4d54772bb8072f9b7036f4170054d35d", size = 512900, upload-time = "2023-11-23T09:26:37.44Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/b1/dd/ead9d8ea85bf202d90cc513b533f9c363121c7792674f78e0d8a854b63b4/jupyterlab_pygments-0.3.0-py3-none-any.whl", hash = "sha256:841a89020971da1d8693f1a99997aefc5dc424bb1b251fd6322462a1b8842780", size = 15884, upload-time = "2023-11-23T09:26:34.325Z" }, @@ -2591,7 +2603,7 @@ wheels = [ [[package]] name = "jupyterlab-server" version = "2.28.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "babel" }, { name = "jinja2" }, @@ -2609,7 +2621,7 @@ wheels = [ [[package]] name = "jupyterlab-widgets" version = "3.0.16" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/26/2d/ef58fed122b268c69c0aa099da20bc67657cdfb2e222688d5731bd5b971d/jupyterlab_widgets-3.0.16.tar.gz", hash = "sha256:423da05071d55cf27a9e602216d35a3a65a3e41cdf9c5d3b643b814ce38c19e0", size = 897423, upload-time = "2025-11-01T21:11:29.724Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/ab/b5/36c712098e6191d1b4e349304ef73a8d06aed77e56ceaac8c0a306c7bda1/jupyterlab_widgets-3.0.16-py3-none-any.whl", hash = "sha256:45fa36d9c6422cf2559198e4db481aa243c7a32d9926b500781c830c80f7ecf8", size = 914926, upload-time = "2025-11-01T21:11:28.008Z" }, @@ -2618,7 +2630,7 @@ wheels = [ [[package]] name = "jupytext" version = "1.19.3" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "markdown-it-py", marker = "sys_platform == 'linux'" }, { name = "mdit-py-plugins", marker = "sys_platform == 'linux'" }, @@ -2634,7 +2646,7 @@ wheels = [ [[package]] name = "kiwisolver" version = "1.5.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/d0/67/9c61eccb13f0bdca9307614e782fec49ffdde0f7a2314935d489fa93cd9c/kiwisolver-1.5.0.tar.gz", hash = "sha256:d4193f3d9dc3f6f79aaed0e5637f45d98850ebf01f7ca20e69457f3e8946b66a", size = 103482, upload-time = "2026-03-09T13:15:53.382Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/4d/b2/818b74ebea34dabe6d0c51cb1c572e046730e64844da6ed646d5298c40ce/kiwisolver-1.5.0-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:4e9750bc21b886308024f8a54ccb9a2cc38ac9fa813bf4348434e3d54f337ff9", size = 123158, upload-time = "2026-03-09T13:13:23.127Z" }, @@ -2720,7 +2732,7 @@ wheels = [ [[package]] name = "labmaze" version = "1.0.6" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "absl-py" }, { name = "numpy" }, @@ -2738,7 +2750,7 @@ wheels = [ [[package]] name = "lark" version = "1.3.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/da/34/28fff3ab31ccff1fd4f6c7c7b0ceb2b6968d8ea4950663eadcb5720591a0/lark-1.3.1.tar.gz", hash = "sha256:b426a7a6d6d53189d318f2b6236ab5d6429eaf09259f1ca33eb716eed10d2905", size = 382732, upload-time = "2025-10-27T18:25:56.653Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/82/3d/14ce75ef66813643812f3093ab17e46d3a206942ce7376d31ec2d36229e7/lark-1.3.1-py3-none-any.whl", hash = "sha256:c629b661023a014c37da873b4ff58a817398d12635d3bbb2c5a03be7fe5d1e12", size = 113151, upload-time = "2025-10-27T18:25:54.882Z" }, @@ -2747,7 +2759,7 @@ wheels = [ [[package]] name = "lazy-loader" version = "0.5" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "packaging" }, ] @@ -2774,9 +2786,9 @@ dependencies = [ { name = "safetensors" }, { name = "setuptools" }, { name = "termcolor" }, - { name = "torch", version = "2.11.0", source = { registry = "https://pypi.registries.huggingface.tech/" }, marker = "sys_platform != 'linux'" }, + { name = "torch", version = "2.11.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'linux'" }, { name = "torch", version = "2.11.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" }, marker = "sys_platform == 'linux'" }, - { name = "torchvision", version = "0.26.0", source = { registry = "https://pypi.registries.huggingface.tech/" }, marker = "sys_platform != 'linux'" }, + { name = "torchvision", version = "0.26.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'linux'" }, { name = "torchvision", version = "0.26.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" }, marker = "sys_platform == 'linux'" }, { name = "tqdm" }, ] @@ -2788,6 +2800,8 @@ accelerate-dep = [ all = [ { name = "accelerate" }, { name = "av" }, + { name = "cmeel-tinyxml2" }, + { name = "cmeel-urdfdom" }, { name = "contourpy" }, { name = "datasets" }, { name = "debugpy" }, @@ -2971,6 +2985,8 @@ hardware = [ ] hilserl = [ { name = "av" }, + { name = "cmeel-tinyxml2" }, + { name = "cmeel-urdfdom" }, { name = "datasets" }, { name = "grpcio" }, { name = "gym-hil" }, @@ -2993,6 +3009,8 @@ intelrealsense = [ { name = "pyrealsense2-macosx", marker = "sys_platform == 'darwin'" }, ] kinematics = [ + { name = "cmeel-tinyxml2" }, + { name = "cmeel-urdfdom" }, { name = "placo" }, ] lekiwi = [ @@ -3066,6 +3084,8 @@ pi = [ { name = "transformers" }, ] placo-dep = [ + { name = "cmeel-tinyxml2" }, + { name = "cmeel-urdfdom" }, { name = "placo" }, ] pusht = [ @@ -3102,9 +3122,6 @@ rebot = [ { name = "motorbridge" }, { name = "motorbridge-smart-servo" }, ] -recap = [ - { name = "transformers" }, -] robometer = [ { name = "peft" }, { name = "qwen-vl-utils" }, @@ -3189,6 +3206,8 @@ requires-dist = [ { name = "accelerate", marker = "extra == 'accelerate-dep'", specifier = ">=1.14.0,<2.0.0" }, { name = "av", marker = "extra == 'av-dep'", specifier = ">=15.0.0,<16.0.0" }, { name = "cmake", specifier = ">=3.29.0.1,<4.2.0" }, + { name = "cmeel-tinyxml2", marker = "extra == 'placo-dep'", specifier = "<11" }, + { name = "cmeel-urdfdom", marker = "extra == 'placo-dep'", specifier = ">=4,<5" }, { name = "contourpy", marker = "extra == 'matplotlib-dep'", specifier = ">=1.3.0,<2.0.0" }, { name = "datasets", marker = "extra == 'dataset'", specifier = ">=4.7.0,<5.0.0" }, { name = "debugpy", marker = "extra == 'dev'", specifier = ">=1.8.1,<1.9.0" }, @@ -3300,7 +3319,6 @@ requires-dist = [ { name = "lerobot", extras = ["qwen-vl-utils-dep"], marker = "extra == 'wallx'" }, { name = "lerobot", extras = ["reachy2"], marker = "extra == 'all'" }, { name = "lerobot", extras = ["rebot"], marker = "extra == 'all'" }, - { name = "lerobot", extras = ["recap"], marker = "extra == 'all'" }, { name = "lerobot", extras = ["robometer"], marker = "extra == 'all'" }, { name = "lerobot", extras = ["robstride"], marker = "extra == 'all'" }, { name = "lerobot", extras = ["sarm"], marker = "extra == 'all'" }, @@ -3324,7 +3342,6 @@ requires-dist = [ { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'multi-task-dit'" }, { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'peft'" }, { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'pi'" }, - { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'recap'" }, { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'robometer'" }, { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'sarm'" }, { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'smolvla'" }, @@ -3378,7 +3395,7 @@ requires-dist = [ { name = "qwen-vl-utils", marker = "extra == 'qwen-vl-utils-dep'", specifier = ">=0.0.11,<0.1.0" }, { name = "reachy2-sdk", marker = "extra == 'reachy2'", specifier = ">=1.0.15,<1.1.0" }, { name = "requests", specifier = ">=2.32.0,<3.0.0" }, - { name = "rerun-sdk", marker = "extra == 'viz'", specifier = ">=0.24.0,<0.27.0" }, + { name = "rerun-sdk", marker = "extra == 'viz'", specifier = ">=0.24.0,<0.34.0" }, { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.14.1" }, { name = "safetensors", specifier = ">=0.4.3,<1.0.0" }, { name = "scikit-image", marker = "extra == 'video-benchmark'", specifier = ">=0.23.2,<0.26.0" }, @@ -3400,12 +3417,12 @@ requires-dist = [ { name = "transformers", marker = "extra == 'transformers-dep'", specifier = ">=5.4.0,<5.6.0" }, { name = "wandb", marker = "extra == 'training'", specifier = ">=0.24.0,<0.28.0" }, ] -provides-extras = ["dataset", "training", "hardware", "viz", "core-scripts", "evaluation", "dataset-viz", "av-dep", "pygame-dep", "placo-dep", "transformers-dep", "grpcio-dep", "accelerate-dep", "can-dep", "peft-dep", "scipy-dep", "diffusers-dep", "qwen-vl-utils-dep", "matplotlib-dep", "pyserial-dep", "deepdiff-dep", "pynput-dep", "pyzmq-dep", "motorbridge-dep", "motorbridge-smart-servo-dep", "feetech", "dynamixel", "damiao", "robstride", "openarms", "gamepad", "hopejr", "lekiwi", "unitree-g1", "reachy2", "rebot", "kinematics", "intelrealsense", "phone", "diffusion", "wallx", "pi", "molmoact2", "smolvla", "multi-task-dit", "groot", "sarm", "robometer", "topreward", "recap", "xvla", "eo1", "hilserl", "vla-jepa", "async", "peft", "annotations", "dev", "notebook", "test", "video-benchmark", "aloha", "pusht", "libero", "metaworld", "all"] +provides-extras = ["dataset", "training", "hardware", "viz", "core-scripts", "evaluation", "dataset-viz", "av-dep", "pygame-dep", "placo-dep", "transformers-dep", "grpcio-dep", "accelerate-dep", "can-dep", "peft-dep", "scipy-dep", "diffusers-dep", "qwen-vl-utils-dep", "matplotlib-dep", "pyserial-dep", "deepdiff-dep", "pynput-dep", "pyzmq-dep", "motorbridge-dep", "motorbridge-smart-servo-dep", "feetech", "dynamixel", "damiao", "robstride", "openarms", "gamepad", "hopejr", "lekiwi", "unitree-g1", "reachy2", "rebot", "kinematics", "intelrealsense", "phone", "diffusion", "wallx", "pi", "molmoact2", "smolvla", "multi-task-dit", "groot", "sarm", "robometer", "topreward", "xvla", "eo1", "hilserl", "vla-jepa", "async", "peft", "annotations", "dev", "notebook", "test", "video-benchmark", "aloha", "pusht", "libero", "metaworld", "all"] [[package]] name = "librt" version = "0.11.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/40/08/9e7f6b5d2b5bed6ad055cdd5925f192bb403a51280f86b56554d9d0699a2/librt-0.11.0.tar.gz", hash = "sha256:075dc3ef4458a278e0195cbf6ac9d38808d9b906c5a6c7f7f79c3888276a3fb1", size = 200139, upload-time = "2026-05-10T18:17:25.138Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/8b/d0/07c77e067f0838949b43bd89232c29d72efebb9d2801a9750184eb706b71/librt-0.11.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:b87504f1690a23b9a2cca841191a04f83895d4fc2dd04df91d82b1a04ca2ad46", size = 144147, upload-time = "2026-05-10T18:15:53.227Z" }, @@ -3465,7 +3482,7 @@ wheels = [ [[package]] name = "llvmlite" version = "0.47.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/01/88/a8952b6d5c21e74cbf158515b779666f692846502623e9e3c39d8e8ba25f/llvmlite-0.47.0.tar.gz", hash = "sha256:62031ce968ec74e95092184d4b0e857e444f8fdff0b8f9213707699570c33ccc", size = 193614, upload-time = "2026-03-31T18:29:53.497Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/e6/4b/e3f2cd17822cf772a4a51a0a8080b0032e6d37b2dbe8cfb724eac4e31c52/llvmlite-0.47.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5853bf26160857c0c2573415ff4efe01c4c651e59e2c55c2a088740acfee51cd", size = 56275178, upload-time = "2026-03-31T18:28:48.342Z" }, @@ -3481,7 +3498,7 @@ wheels = [ [[package]] name = "lxml" version = "6.1.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/05/3b/aab6728cae887456f409b4d75e8a01856e4f04bd510de38052a47768b680/lxml-6.1.1.tar.gz", hash = "sha256:ba96ae44888e0185281e937633a743ea90d5a196c6000f82565ebb0580012d40", size = 4197430, upload-time = "2026-05-18T19:19:06.424Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/6a/6e/c4add832b6fc1e887125b96f880d7b9b70aae5248718e046b1704bcac4b9/lxml-6.1.1-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:104c09bda8d2a562824c0e319d0768ce26a779b7601e0931d33b09b53c392ef7", size = 8570821, upload-time = "2026-05-18T19:17:42.068Z" }, @@ -3561,7 +3578,7 @@ wheels = [ [[package]] name = "markdown" version = "3.10.2" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/2b/f4/69fa6ed85ae003c2378ffa8f6d2e3234662abd02c10d216c0ba96081a238/markdown-3.10.2.tar.gz", hash = "sha256:994d51325d25ad8aa7ce4ebaec003febcce822c3f8c911e3b17c52f7f589f950", size = 368805, upload-time = "2026-02-09T14:57:26.942Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/de/1f/77fa3081e4f66ca3576c896ae5d31c3002ac6607f9747d2e3aa49227e464/markdown-3.10.2-py3-none-any.whl", hash = "sha256:e91464b71ae3ee7afd3017d9f358ef0baf158fd9a298db92f1d4761133824c36", size = 108180, upload-time = "2026-02-09T14:57:25.787Z" }, @@ -3570,7 +3587,7 @@ wheels = [ [[package]] name = "markdown-it-py" version = "4.2.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "mdurl" }, ] @@ -3582,7 +3599,7 @@ wheels = [ [[package]] name = "markupsafe" version = "3.0.3" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/7e/99/7690b6d4034fffd95959cbe0c02de8deb3098cc577c67bb6a24fe5d7caa7/markupsafe-3.0.3.tar.gz", hash = "sha256:722695808f4b6457b320fdc131280796bdceb04ab50fe1795cd540799ebe1698", size = 80313, upload-time = "2025-09-27T18:37:40.426Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/5a/72/147da192e38635ada20e0a2e1a51cf8823d2119ce8883f7053879c2199b5/markupsafe-3.0.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:d53197da72cc091b024dd97249dfc7794d6a56530370992a5e1a08983ad9230e", size = 11615, upload-time = "2025-09-27T18:36:30.854Z" }, @@ -3645,7 +3662,7 @@ wheels = [ [[package]] name = "matplotlib" version = "3.11.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "contourpy" }, { name = "cycler" }, @@ -3699,7 +3716,7 @@ wheels = [ [[package]] name = "matplotlib-inline" version = "0.2.2" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "traitlets" }, ] @@ -3711,7 +3728,7 @@ wheels = [ [[package]] name = "mdit-py-plugins" version = "0.6.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "markdown-it-py", marker = "sys_platform == 'linux'" }, ] @@ -3723,7 +3740,7 @@ wheels = [ [[package]] name = "mdurl" version = "0.1.2" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/d6/54/cfe61301667036ec958cb99bd3efefba235e65cdeb9c84d24a8293ba1d90/mdurl-0.1.2.tar.gz", hash = "sha256:bb413d29f5eea38f31dd4754dd7377d4465116fb207585f97bf925588687c1ba", size = 8729, upload-time = "2022-08-14T12:40:10.846Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8", size = 9979, upload-time = "2022-08-14T12:40:09.779Z" }, @@ -3732,7 +3749,7 @@ wheels = [ [[package]] name = "mergedeep" version = "1.3.4" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/3a/41/580bb4006e3ed0361b8151a01d324fb03f420815446c7def45d02f74c270/mergedeep-1.3.4.tar.gz", hash = "sha256:0096d52e9dad9939c3d975a774666af186eda617e6ca84df4c94dec30004f2a8", size = 4661, upload-time = "2021-02-05T18:55:30.623Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/2c/19/04f9b178c2d8a15b076c8b5140708fa6ffc5601fb6f1e975537072df5b2a/mergedeep-1.3.4-py3-none-any.whl", hash = "sha256:70775750742b25c0d8f36c55aed03d24c3384d17c951b3175d898bd778ef0307", size = 6354, upload-time = "2021-02-05T18:55:29.583Z" }, @@ -3741,7 +3758,7 @@ wheels = [ [[package]] name = "meshcat" version = "0.3.2" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "ipython" }, { name = "numpy" }, @@ -3759,7 +3776,7 @@ wheels = [ [[package]] name = "metaworld" version = "3.0.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "gymnasium" }, { name = "imageio" }, @@ -3775,7 +3792,7 @@ wheels = [ [[package]] name = "mistune" version = "3.2.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/ca/84/620cc3f7e3adf6f5067e10f4dbae71295d8f9e16d5d3f9ef97c40f2f592c/mistune-3.2.1.tar.gz", hash = "sha256:7c8e5501d38bac1582e067e46c8343f17d57ea1aaa735823f3aba1fd59c88a28", size = 98003, upload-time = "2026-05-03T14:33:22.312Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/2a/7f/a946aa4f8752b37102b41e64dca18a1976ac705c3a0d1dfe74d820a02552/mistune-3.2.1-py3-none-any.whl", hash = "sha256:78cdb0ba5e938053ccf63651b352508d2efa9411dc8810bfb05f2dc5140c0048", size = 53749, upload-time = "2026-05-03T14:33:20.551Z" }, @@ -3784,7 +3801,7 @@ wheels = [ [[package]] name = "ml-dtypes" version = "0.5.4" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "numpy" }, ] @@ -3820,7 +3837,7 @@ wheels = [ [[package]] name = "mock-serial" version = "0.0.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/c6/d9/b9ec64510ffb4bccd7491c4583b8f24732ed4af6311c621822c3a9e47b46/mock_serial-0.0.1.tar.gz", hash = "sha256:9c92de7495ac375717bbbeb4993534079d2e634f3298d4c400420c4046add06e", size = 5355, upload-time = "2021-11-23T09:34:52.625Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/98/c2/8c1e6bf77cf62a10203a107179e34e0965fc5369386e0b7034a247ed054d/mock_serial-0.0.1-py3-none-any.whl", hash = "sha256:b6b8cc10c302354bf3ca270a3d4d6bf199c4bbe41478c65046db8f30ea967675", size = 6080, upload-time = "2021-11-23T09:34:51.108Z" }, @@ -3829,7 +3846,7 @@ wheels = [ [[package]] name = "motorbridge" version = "0.3.9" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/89/b2/36e84fc4274fb9fb288e50ab8fe9451ea6f31aed5b735f7234df515eb557/motorbridge-0.3.9.tar.gz", hash = "sha256:12948aa4d3662354744bfac5777243a64ac35c9bb990d42182efc5c8f179d40f", size = 34002, upload-time = "2026-05-26T09:28:52.458Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/0d/16/3dba18382393dda35c37a590c991e7d998c642ff04b4171f5cacea09f79b/motorbridge-0.3.9-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:6abc0111779da3ed6ecd8d14c4241e0a37f5adf9c5e997d360e9940cb17fcde9", size = 1204016, upload-time = "2026-05-26T09:28:41.065Z" }, @@ -3846,7 +3863,7 @@ wheels = [ [[package]] name = "motorbridge-smart-servo" version = "0.0.4" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/e6/56/45af87189dc49abbe46157b792b7c71f502a5f819f04e7485de0cfa52d9b/motorbridge_smart_servo-0.0.4.tar.gz", hash = "sha256:fb65f3f6e765e6b1915071c255caaf112fad3796fa1761aeee0132d15b8a0989", size = 20415, upload-time = "2026-05-08T09:24:57.563Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/e9/ee/bec4b3acf55cd18e7db83a6d951caccf699533dbd038c1f0b5f2d16d5208/motorbridge_smart_servo-0.0.4-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:8bc1f034fa9f96e23229a834db6e7cfe1368dba7b9a2a6f6dbd316448c4390dc", size = 304384, upload-time = "2026-05-08T09:24:52.619Z" }, @@ -3858,7 +3875,7 @@ wheels = [ [[package]] name = "mpmath" version = "1.3.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/e0/47/dd32fa426cc72114383ac549964eecb20ecfd886d1e5ccf5340b55b02f57/mpmath-1.3.0.tar.gz", hash = "sha256:7a28eb2a9774d00c7bc92411c19a89209d5da7c4c9a9e227be8330a23a25b91f", size = 508106, upload-time = "2023-03-07T16:47:11.061Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/43/e3/7d92a15f894aa0c9c4b49b8ee9ac9850d6e63b03c9c32c0367a13ae62209/mpmath-1.3.0-py3-none-any.whl", hash = "sha256:a0b2b9fe80bbcd81a6647ff13108738cfb482d481d826cc0e02f5b35e5c88d2c", size = 536198, upload-time = "2023-03-07T16:47:09.197Z" }, @@ -3867,7 +3884,7 @@ wheels = [ [[package]] name = "mujoco" version = "3.8.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "absl-py" }, { name = "etils", extra = ["epath"] }, @@ -3897,7 +3914,7 @@ wheels = [ [[package]] name = "multidict" version = "6.7.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/1a/c2/c2d94cbe6ac1753f3fc980da97b3d930efe1da3af3c9f5125354436c073d/multidict-6.7.1.tar.gz", hash = "sha256:ec6652a1bee61c53a3e5776b6049172c53b6aaba34f18c9ad04f82712bac623d", size = 102010, upload-time = "2026-01-26T02:46:45.979Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/8d/9c/f20e0e2cf80e4b2e4b1c365bf5fe104ee633c751a724246262db8f1a0b13/multidict-6.7.1-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:a90f75c956e32891a4eda3639ce6dd86e87105271f43d43442a3aedf3cddf172", size = 76893, upload-time = "2026-01-26T02:43:52.754Z" }, @@ -3996,7 +4013,7 @@ wheels = [ [[package]] name = "multiprocess" version = "0.70.19" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "dill" }, ] @@ -4013,7 +4030,7 @@ wheels = [ [[package]] name = "mypy" version = "2.1.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "ast-serialize" }, { name = "librt", marker = "platform_python_implementation != 'PyPy'" }, @@ -4057,7 +4074,7 @@ wheels = [ [[package]] name = "mypy-extensions" version = "1.1.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/a2/6e/371856a3fb9d31ca8dac321cda606860fa4548858c0cc45d9d1d4ca2628b/mypy_extensions-1.1.0.tar.gz", hash = "sha256:52e68efc3284861e772bbcd66823fde5ae21fd2fdb51c62a211403730b916558", size = 6343, upload-time = "2025-04-22T14:54:24.164Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/79/7b/2c79738432f5c924bef5071f933bcc9efd0473bac3b4aa584a6f7c1c8df8/mypy_extensions-1.1.0-py3-none-any.whl", hash = "sha256:1be4cccdb0f2482337c4743e60421de3a356cd97508abadd57d47403e94f5505", size = 4963, upload-time = "2025-04-22T14:54:22.983Z" }, @@ -4066,7 +4083,7 @@ wheels = [ [[package]] name = "nbclient" version = "0.11.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "jupyter-client" }, { name = "jupyter-core" }, @@ -4081,7 +4098,7 @@ wheels = [ [[package]] name = "nbconvert" version = "7.17.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "beautifulsoup4" }, { name = "bleach", extra = ["css"] }, @@ -4106,7 +4123,7 @@ wheels = [ [[package]] name = "nbformat" version = "5.10.4" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "fastjsonschema" }, { name = "jsonschema" }, @@ -4121,7 +4138,7 @@ wheels = [ [[package]] name = "nest-asyncio" version = "1.6.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/83/f8/51569ac65d696c8ecbee95938f89d4abf00f47d58d48f6fbabfe8f0baefe/nest_asyncio-1.6.0.tar.gz", hash = "sha256:6f172d5449aca15afd6c646851f4e31e02c598d553a667e38cafa997cfec55fe", size = 7418, upload-time = "2024-01-21T14:25:19.227Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/a0/c4/c2971a3ba4c6103a3d10c4b0f24f461ddc027f0f09763220cf35ca1401b3/nest_asyncio-1.6.0-py3-none-any.whl", hash = "sha256:87af6efd6b5e897c81050477ef65c62e2b2f35d51703cae01aff2905b1852e1c", size = 5195, upload-time = "2024-01-21T14:25:17.223Z" }, @@ -4130,7 +4147,7 @@ wheels = [ [[package]] name = "networkx" version = "3.6.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/6a/51/63fe664f3908c97be9d2e4f1158eb633317598cfa6e1fc14af5383f17512/networkx-3.6.1.tar.gz", hash = "sha256:26b7c357accc0c8cde558ad486283728b65b6a95d85ee1cd66bafab4c8168509", size = 2517025, upload-time = "2025-12-08T17:02:39.908Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl", hash = "sha256:d47fbf302e7d9cbbb9e2555a0d267983d2aa476bac30e90dfbe5669bd57f3762", size = 2068504, upload-time = "2025-12-08T17:02:38.159Z" }, @@ -4139,7 +4156,7 @@ wheels = [ [[package]] name = "ninja" version = "1.13.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/43/73/79a0b22fc731989c708068427579e840a6cf4e937fe7ae5c5d0b7356ac22/ninja-1.13.0.tar.gz", hash = "sha256:4a40ce995ded54d9dc24f8ea37ff3bf62ad192b547f6c7126e7e25045e76f978", size = 242558, upload-time = "2025-08-11T15:10:19.421Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/3c/74/d02409ed2aa865e051b7edda22ad416a39d81a84980f544f8de717cab133/ninja-1.13.0-py3-none-macosx_10_9_universal2.whl", hash = "sha256:fa2a8bfc62e31b08f83127d1613d10821775a0eb334197154c4d6067b7068ff1", size = 310125, upload-time = "2025-08-11T15:09:50.971Z" }, @@ -4165,7 +4182,7 @@ wheels = [ [[package]] name = "nodeenv" version = "1.10.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/24/bf/d1bda4f6168e0b2e9e5958945e01910052158313224ada5ce1fb2e1113b8/nodeenv-1.10.0.tar.gz", hash = "sha256:996c191ad80897d076bdfba80a41994c2b47c68e224c542b48feba42ba00f8bb", size = 55611, upload-time = "2025-12-20T14:08:54.006Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/88/b2/d0896bdcdc8d28a7fc5717c305f1a861c26e18c05047949fb371034d98bd/nodeenv-1.10.0-py2.py3-none-any.whl", hash = "sha256:5bb13e3eed2923615535339b3c620e76779af4cb4c6a90deccc9e36b274d3827", size = 23438, upload-time = "2025-12-20T14:08:52.782Z" }, @@ -4174,7 +4191,7 @@ wheels = [ [[package]] name = "notebook" version = "7.5.7" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "jupyter-server" }, { name = "jupyterlab" }, @@ -4190,7 +4207,7 @@ wheels = [ [[package]] name = "notebook-shim" version = "0.2.4" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "jupyter-server" }, ] @@ -4202,7 +4219,7 @@ wheels = [ [[package]] name = "num2words" version = "0.5.14" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "docopt" }, ] @@ -4214,7 +4231,7 @@ wheels = [ [[package]] name = "numba" version = "0.65.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "llvmlite", marker = "sys_platform == 'linux'" }, { name = "numpy", marker = "sys_platform == 'linux'" }, @@ -4234,7 +4251,7 @@ wheels = [ [[package]] name = "numpy" version = "2.2.6" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/76/21/7d2a95e4bba9dc13d043ee156a356c0a8f0c6309dff6b21b4d71a073b8a8/numpy-2.2.6.tar.gz", hash = "sha256:e29554e2bef54a90aa5cc07da6ce955accb83f21ab5de01a62c8478897b264fd", size = 20276440, upload-time = "2025-05-17T22:38:04.611Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/82/5d/c00588b6cf18e1da539b45d3598d3557084990dcc4331960c15ee776ee41/numpy-2.2.6-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:41c5a21f4a04fa86436124d388f6ed60a9343a6f767fced1a8a71c3fbca038ff", size = 20875348, upload-time = "2025-05-17T21:34:39.648Z" }, @@ -4272,7 +4289,7 @@ wheels = [ [[package]] name = "nvidia-cublas-cu12" version = "12.8.4.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } wheels = [ { url = "https://files.pythonhosted.org/packages/29/99/db44d685f0e257ff0e213ade1964fc459b4a690a73293220e98feb3307cf/nvidia_cublas_cu12-12.8.4.1-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:b86f6dd8935884615a0683b663891d43781b819ac4f2ba2b0c9604676af346d0", size = 590537124, upload-time = "2025-03-07T01:43:53.556Z" }, { url = "https://files.pythonhosted.org/packages/dc/61/e24b560ab2e2eaeb3c839129175fb330dfcfc29e5203196e5541a4c44682/nvidia_cublas_cu12-12.8.4.1-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:8ac4e771d5a348c551b2a426eda6193c19aa630236b418086020df5ba9667142", size = 594346921, upload-time = "2025-03-07T01:44:31.254Z" }, @@ -4281,7 +4298,7 @@ wheels = [ [[package]] name = "nvidia-cuda-cupti-cu12" version = "12.8.90" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } wheels = [ { url = "https://files.pythonhosted.org/packages/d5/1f/b3bd73445e5cb342727fd24fe1f7b748f690b460acadc27ea22f904502c8/nvidia_cuda_cupti_cu12-12.8.90-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:4412396548808ddfed3f17a467b104ba7751e6b58678a4b840675c56d21cf7ed", size = 9533318, upload-time = "2025-03-07T01:40:10.421Z" }, { url = "https://files.pythonhosted.org/packages/f8/02/2adcaa145158bf1a8295d83591d22e4103dbfd821bcaf6f3f53151ca4ffa/nvidia_cuda_cupti_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ea0cb07ebda26bb9b29ba82cda34849e73c166c18162d3913575b0c9db9a6182", size = 10248621, upload-time = "2025-03-07T01:40:21.213Z" }, @@ -4290,7 +4307,7 @@ wheels = [ [[package]] name = "nvidia-cuda-nvrtc-cu12" version = "12.8.93" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } wheels = [ { url = "https://files.pythonhosted.org/packages/05/6b/32f747947df2da6994e999492ab306a903659555dddc0fbdeb9d71f75e52/nvidia_cuda_nvrtc_cu12-12.8.93-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:a7756528852ef889772a84c6cd89d41dfa74667e24cca16bb31f8f061e3e9994", size = 88040029, upload-time = "2025-03-07T01:42:13.562Z" }, { url = "https://files.pythonhosted.org/packages/eb/d1/e50d0acaab360482034b84b6e27ee83c6738f7d32182b987f9c7a4e32962/nvidia_cuda_nvrtc_cu12-12.8.93-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:fc1fec1e1637854b4c0a65fb9a8346b51dd9ee69e61ebaccc82058441f15bce8", size = 43106076, upload-time = "2025-03-07T01:41:59.817Z" }, @@ -4299,7 +4316,7 @@ wheels = [ [[package]] name = "nvidia-cuda-runtime-cu12" version = "12.8.90" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } wheels = [ { url = "https://files.pythonhosted.org/packages/7c/75/f865a3b236e4647605ea34cc450900854ba123834a5f1598e160b9530c3a/nvidia_cuda_runtime_cu12-12.8.90-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:52bf7bbee900262ffefe5e9d5a2a69a30d97e2bc5bb6cc866688caa976966e3d", size = 965265, upload-time = "2025-03-07T01:39:43.533Z" }, { url = "https://files.pythonhosted.org/packages/0d/9b/a997b638fcd068ad6e4d53b8551a7d30fe8b404d6f1804abf1df69838932/nvidia_cuda_runtime_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:adade8dcbd0edf427b7204d480d6066d33902cab2a4707dcfc48a2d0fd44ab90", size = 954765, upload-time = "2025-03-07T01:40:01.615Z" }, @@ -4308,7 +4325,7 @@ wheels = [ [[package]] name = "nvidia-cudnn-cu12" version = "9.19.0.56" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "nvidia-cublas-cu12", marker = "sys_platform == 'linux'" }, ] @@ -4320,7 +4337,7 @@ wheels = [ [[package]] name = "nvidia-cufft-cu12" version = "11.3.3.83" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "nvidia-nvjitlink-cu12", marker = "sys_platform == 'linux'" }, ] @@ -4332,7 +4349,7 @@ wheels = [ [[package]] name = "nvidia-cufile-cu12" version = "1.13.1.3" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } wheels = [ { url = "https://files.pythonhosted.org/packages/bb/fe/1bcba1dfbfb8d01be8d93f07bfc502c93fa23afa6fd5ab3fc7c1df71038a/nvidia_cufile_cu12-1.13.1.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1d069003be650e131b21c932ec3d8969c1715379251f8d23a1860554b1cb24fc", size = 1197834, upload-time = "2025-03-07T01:45:50.723Z" }, { url = "https://files.pythonhosted.org/packages/1e/f5/5607710447a6fe9fd9b3283956fceeee8a06cda1d2f56ce31371f595db2a/nvidia_cufile_cu12-1.13.1.3-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:4beb6d4cce47c1a0f1013d72e02b0994730359e17801d395bdcbf20cfb3bb00a", size = 1120705, upload-time = "2025-03-07T01:45:41.434Z" }, @@ -4341,7 +4358,7 @@ wheels = [ [[package]] name = "nvidia-curand-cu12" version = "10.3.9.90" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } wheels = [ { url = "https://files.pythonhosted.org/packages/45/5e/92aa15eca622a388b80fbf8375d4760738df6285b1e92c43d37390a33a9a/nvidia_curand_cu12-10.3.9.90-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:dfab99248034673b779bc6decafdc3404a8a6f502462201f2f31f11354204acd", size = 63625754, upload-time = "2025-03-07T01:46:10.735Z" }, { url = "https://files.pythonhosted.org/packages/fb/aa/6584b56dc84ebe9cf93226a5cde4d99080c8e90ab40f0c27bda7a0f29aa1/nvidia_curand_cu12-10.3.9.90-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:b32331d4f4df5d6eefa0554c565b626c7216f87a06a4f56fab27c3b68a830ec9", size = 63619976, upload-time = "2025-03-07T01:46:23.323Z" }, @@ -4350,7 +4367,7 @@ wheels = [ [[package]] name = "nvidia-cusolver-cu12" version = "11.7.3.90" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "nvidia-cublas-cu12", marker = "sys_platform == 'linux'" }, { name = "nvidia-cusparse-cu12", marker = "sys_platform == 'linux'" }, @@ -4364,7 +4381,7 @@ wheels = [ [[package]] name = "nvidia-cusparse-cu12" version = "12.5.8.93" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "nvidia-nvjitlink-cu12", marker = "sys_platform == 'linux'" }, ] @@ -4376,7 +4393,7 @@ wheels = [ [[package]] name = "nvidia-cusparselt-cu12" version = "0.7.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } wheels = [ { url = "https://files.pythonhosted.org/packages/73/b9/598f6ff36faaece4b3c50d26f50e38661499ff34346f00e057760b35cc9d/nvidia_cusparselt_cu12-0.7.1-py3-none-manylinux2014_aarch64.whl", hash = "sha256:8878dce784d0fac90131b6817b607e803c36e629ba34dc5b433471382196b6a5", size = 283835557, upload-time = "2025-02-26T00:16:54.265Z" }, { url = "https://files.pythonhosted.org/packages/56/79/12978b96bd44274fe38b5dde5cfb660b1d114f70a65ef962bcbbed99b549/nvidia_cusparselt_cu12-0.7.1-py3-none-manylinux2014_x86_64.whl", hash = "sha256:f1bb701d6b930d5a7cea44c19ceb973311500847f81b634d802b7b539dc55623", size = 287193691, upload-time = "2025-02-26T00:15:44.104Z" }, @@ -4385,7 +4402,7 @@ wheels = [ [[package]] name = "nvidia-nccl-cu12" version = "2.28.9" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } wheels = [ { url = "https://files.pythonhosted.org/packages/08/c4/120d2dfd92dff2c776d68f361ff8705fdea2ca64e20b612fab0fd3f581ac/nvidia_nccl_cu12-2.28.9-py3-none-manylinux_2_18_aarch64.whl", hash = "sha256:50a36e01c4a090b9f9c47d92cec54964de6b9fcb3362d0e19b8ffc6323c21b60", size = 296766525, upload-time = "2025-11-18T05:49:16.094Z" }, { url = "https://files.pythonhosted.org/packages/4a/4e/44dbb46b3d1b0ec61afda8e84837870f2f9ace33c564317d59b70bc19d3e/nvidia_nccl_cu12-2.28.9-py3-none-manylinux_2_18_x86_64.whl", hash = "sha256:485776daa8447da5da39681af455aa3b2c2586ddcf4af8772495e7c532c7e5ab", size = 296782137, upload-time = "2025-11-18T05:49:34.248Z" }, @@ -4394,7 +4411,7 @@ wheels = [ [[package]] name = "nvidia-nvjitlink-cu12" version = "12.8.93" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } wheels = [ { url = "https://files.pythonhosted.org/packages/f6/74/86a07f1d0f42998ca31312f998bd3b9a7eff7f52378f4f270c8679c77fb9/nvidia_nvjitlink_cu12-12.8.93-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:81ff63371a7ebd6e6451970684f916be2eab07321b73c9d244dc2b4da7f73b88", size = 39254836, upload-time = "2025-03-07T01:49:55.661Z" }, { url = "https://files.pythonhosted.org/packages/2a/a2/8cee5da30d13430e87bf99bb33455d2724d0a4a9cb5d7926d80ccb96d008/nvidia_nvjitlink_cu12-12.8.93-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:adccd7161ace7261e01bb91e44e88da350895c270d23f744f0820c818b7229e7", size = 38386204, upload-time = "2025-03-07T01:49:43.612Z" }, @@ -4403,7 +4420,7 @@ wheels = [ [[package]] name = "nvidia-nvshmem-cu12" version = "3.4.5" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } wheels = [ { url = "https://files.pythonhosted.org/packages/1d/6a/03aa43cc9bd3ad91553a88b5f6fb25ed6a3752ae86ce2180221962bc2aa5/nvidia_nvshmem_cu12-3.4.5-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:0b48363fc6964dede448029434c6abed6c5e37f823cb43c3bcde7ecfc0457e15", size = 138936938, upload-time = "2025-09-06T00:32:05.589Z" }, { url = "https://files.pythonhosted.org/packages/b5/09/6ea3ea725f82e1e76684f0708bbedd871fc96da89945adeba65c3835a64c/nvidia_nvshmem_cu12-3.4.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:042f2500f24c021db8a06c5eec2539027d57460e1c1a762055a6554f72c369bd", size = 139103095, upload-time = "2025-09-06T00:32:31.266Z" }, @@ -4412,7 +4429,7 @@ wheels = [ [[package]] name = "nvidia-nvtx-cu12" version = "12.8.90" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } wheels = [ { url = "https://files.pythonhosted.org/packages/10/c0/1b303feea90d296f6176f32a2a70b5ef230f9bdeb3a72bddb0dc922dc137/nvidia_nvtx_cu12-12.8.90-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d7ad891da111ebafbf7e015d34879f7112832fc239ff0d7d776b6cb685274615", size = 91161, upload-time = "2025-03-07T01:42:23.922Z" }, { url = "https://files.pythonhosted.org/packages/a2/eb/86626c1bbc2edb86323022371c39aa48df6fd8b0a1647bc274577f72e90b/nvidia_nvtx_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5b17e2001cc0d751a5bc2c6ec6d26ad95913324a4adb86788c944f8ce9ba441f", size = 89954, upload-time = "2025-03-07T01:42:44.131Z" }, @@ -4421,7 +4438,7 @@ wheels = [ [[package]] name = "omegaconf" version = "2.3.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "antlr4-python3-runtime", marker = "sys_platform == 'linux'" }, { name = "pyyaml", marker = "sys_platform == 'linux'" }, @@ -4434,7 +4451,7 @@ wheels = [ [[package]] name = "onnx" version = "1.21.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "ml-dtypes" }, { name = "numpy" }, @@ -4464,7 +4481,7 @@ wheels = [ [[package]] name = "onnxruntime" version = "1.26.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "flatbuffers" }, { name = "numpy" }, @@ -4496,7 +4513,7 @@ wheels = [ [[package]] name = "openai" version = "1.109.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "anyio" }, { name = "distro" }, @@ -4515,7 +4532,7 @@ wheels = [ [[package]] name = "opencv-python" version = "4.13.0.92" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "numpy" }, ] @@ -4533,7 +4550,7 @@ wheels = [ [[package]] name = "opencv-python-headless" version = "4.13.0.92" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "numpy" }, ] @@ -4551,7 +4568,7 @@ wheels = [ [[package]] name = "orderly-set" version = "5.5.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/4a/88/39c83c35d5e97cc203e9e77a4f93bf87ec89cf6a22ac4818fdcc65d66584/orderly_set-5.5.0.tar.gz", hash = "sha256:e87185c8e4d8afa64e7f8160ee2c542a475b738bc891dc3f58102e654125e6ce", size = 27414, upload-time = "2025-07-10T20:10:55.885Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/12/27/fb8d7338b4d551900fa3e580acbe7a0cf655d940e164cb5c00ec31961094/orderly_set-5.5.0-py3-none-any.whl", hash = "sha256:46f0b801948e98f427b412fcabb831677194c05c3b699b80de260374baa0b1e7", size = 13068, upload-time = "2025-07-10T20:10:54.377Z" }, @@ -4560,7 +4577,7 @@ wheels = [ [[package]] name = "packaging" version = "25.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/a1/d4/1fc4078c65507b51b96ca8f8c3ba19e6a61c8253c72794544580a7b6c24d/packaging-25.0.tar.gz", hash = "sha256:d443872c98d677bf60f6a1f2f8c1cb748e8fe762d2bf9d3148b5599295b0fc4f", size = 165727, upload-time = "2025-04-19T11:48:59.673Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/20/12/38679034af332785aac8774540895e234f4d07f7545804097de4b666afd8/packaging-25.0-py3-none-any.whl", hash = "sha256:29572ef2b1f17581046b3a2227d5c611fb25ec70ca1ba8554b24b0e69331a484", size = 66469, upload-time = "2025-04-19T11:48:57.875Z" }, @@ -4569,7 +4586,7 @@ wheels = [ [[package]] name = "pandas" version = "2.3.3" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "numpy" }, { name = "python-dateutil" }, @@ -4616,7 +4633,7 @@ wheels = [ [[package]] name = "pandocfilters" version = "1.5.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/70/6f/3dd4940bbe001c06a65f88e36bad298bc7a0de5036115639926b0c5c0458/pandocfilters-1.5.1.tar.gz", hash = "sha256:002b4a555ee4ebc03f8b66307e287fa492e4a77b4ea14d3f934328297bb4939e", size = 8454, upload-time = "2024-01-18T20:08:13.726Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/ef/af/4fbc8cab944db5d21b7e2a5b8e9211a03a79852b1157e2c102fcc61ac440/pandocfilters-1.5.1-py2.py3-none-any.whl", hash = "sha256:93be382804a9cdb0a7267585f157e5d1731bbe5545a85b268d6f5fe6232de2bc", size = 8663, upload-time = "2024-01-18T20:08:11.28Z" }, @@ -4625,7 +4642,7 @@ wheels = [ [[package]] name = "parso" version = "0.8.7" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/30/4b/90c937815137d43ce71ba043cd3566221e9df6b9c805f24b5d138c9d40a7/parso-0.8.7.tar.gz", hash = "sha256:eaaac4c9fdd5e9e8852dc778d2d7405897ec510f2a298071453e5e3a07914bb1", size = 401824, upload-time = "2026-05-01T23:13:02.138Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/99/5d/8268b644392ee874ee82a635cd0df1773de230bde356c38de28e298392cc/parso-0.8.7-py2.py3-none-any.whl", hash = "sha256:a8926eb2a1b915486941fdbd31e86a4baf88fe8c210f25f2f35ecec5b574ca1c", size = 107025, upload-time = "2026-05-01T23:12:58.867Z" }, @@ -4634,7 +4651,7 @@ wheels = [ [[package]] name = "pathspec" version = "1.1.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/5a/82/42f767fc1c1143d6fd36efb827202a2d997a375e160a71eb2888a925aac1/pathspec-1.1.1.tar.gz", hash = "sha256:17db5ecd524104a120e173814c90367a96a98d07c45b2e10c2f3919fff91bf5a", size = 135180, upload-time = "2026-04-27T01:46:08.907Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/f1/d9/7fb5aa316bc299258e68c73ba3bddbc499654a07f151cba08f6153988714/pathspec-1.1.1-py3-none-any.whl", hash = "sha256:a00ce642f577bf7f473932318056212bc4f8bfdf53128c78bbd5af0b9b20b189", size = 57328, upload-time = "2026-04-27T01:46:07.06Z" }, @@ -4643,7 +4660,7 @@ wheels = [ [[package]] name = "peft" version = "0.19.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "accelerate" }, { name = "huggingface-hub" }, @@ -4652,7 +4669,7 @@ dependencies = [ { name = "psutil" }, { name = "pyyaml" }, { name = "safetensors" }, - { name = "torch", version = "2.11.0", source = { registry = "https://pypi.registries.huggingface.tech/" }, marker = "sys_platform != 'linux'" }, + { name = "torch", version = "2.11.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'linux'" }, { name = "torch", version = "2.11.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" }, marker = "sys_platform == 'linux'" }, { name = "tqdm" }, { name = "transformers" }, @@ -4665,7 +4682,7 @@ wheels = [ [[package]] name = "pexpect" version = "4.9.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "ptyprocess", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, ] @@ -4677,7 +4694,7 @@ wheels = [ [[package]] name = "pillow" version = "12.2.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/8c/21/c2bcdd5906101a30244eaffc1b6e6ce71a31bd0742a01eb89e660ebfac2d/pillow-12.2.0.tar.gz", hash = "sha256:a830b1a40919539d07806aa58e1b114df53ddd43213d9c8b75847eee6c0182b5", size = 46987819, upload-time = "2026-04-01T14:46:17.687Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/58/be/7482c8a5ebebbc6470b3eb791812fff7d5e0216c2be3827b30b8bb6603ed/pillow-12.2.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:2d192a155bbcec180f8564f693e6fd9bccff5a7af9b32e2e4bf8c9c69dbad6b5", size = 5308279, upload-time = "2026-04-01T14:43:13.246Z" }, @@ -4746,7 +4763,7 @@ wheels = [ [[package]] name = "pin" version = "3.4.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "cmeel" }, { name = "cmeel-boost" }, @@ -4768,7 +4785,7 @@ wheels = [ [[package]] name = "placo" version = "0.9.15" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "cmeel" }, { name = "eiquadprog" }, @@ -4792,7 +4809,7 @@ wheels = [ [[package]] name = "platformdirs" version = "4.10.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/d7/47/e4501f49c178ae1d9f4a75073fda4204f52647993f075a9db4d14930e0c5/platformdirs-4.10.0.tar.gz", hash = "sha256:31e761a6a0ca04faf7353ea759bdba55652be214725111e5aac52dfa29d4bef7", size = 31224, upload-time = "2026-05-28T03:32:53.587Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/81/e6/cd9575ac904136b3cbf7aa7ee819ef86eedb7274e46f230e94ea4342e729/platformdirs-4.10.0-py3-none-any.whl", hash = "sha256:fb516cdb12eb0d857d0cd85a7c57cea4d060bee4578d6cf5a14dfdf8cbf8784a", size = 22743, upload-time = "2026-05-28T03:32:52.175Z" }, @@ -4801,7 +4818,7 @@ wheels = [ [[package]] name = "pluggy" version = "1.6.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/f9/e2/3e91f31a7d2b083fe6ef3fa267035b518369d9511ffab804f839851d2779/pluggy-1.6.0.tar.gz", hash = "sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3", size = 69412, upload-time = "2025-05-15T12:30:07.975Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" }, @@ -4810,7 +4827,7 @@ wheels = [ [[package]] name = "pre-commit" version = "4.6.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "cfgv" }, { name = "identify" }, @@ -4826,7 +4843,7 @@ wheels = [ [[package]] name = "prometheus-client" version = "0.25.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/1b/fb/d9aa83ffe43ce1f19e557c0971d04b90561b0cfd50762aafb01968285553/prometheus_client-0.25.0.tar.gz", hash = "sha256:5e373b75c31afb3c86f1a52fa1ad470c9aace18082d39ec0d2f918d11cc9ba28", size = 86035, upload-time = "2026-04-09T19:53:42.359Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/8d/9b/d4b1e644385499c8346fa9b622a3f030dce14cd6ef8a1871c221a17a67e7/prometheus_client-0.25.0-py3-none-any.whl", hash = "sha256:d5aec89e349a6ec230805d0df882f3807f74fd6c1a2fa86864e3c2279059fed1", size = 64154, upload-time = "2026-04-09T19:53:41.324Z" }, @@ -4835,7 +4852,7 @@ wheels = [ [[package]] name = "prompt-toolkit" version = "3.0.52" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "wcwidth" }, ] @@ -4847,7 +4864,7 @@ wheels = [ [[package]] name = "propcache" version = "0.5.2" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/ec/44/c87281c333769159c50594f22610f77398a47ccbfbbf23074e744e86f87c/propcache-0.5.2.tar.gz", hash = "sha256:01c4fc7480cd0598bb4b57022df55b9ca296da7fc5a8760bd8451a7e63a7d427", size = 50208, upload-time = "2026-05-08T21:02:12.199Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/4a/cb/e27bc2b2737a0bb49962b275efa051e8f1c35a936df7d5139b6b658b7dc9/propcache-0.5.2-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:806719138ecd720339a12410fb9614ac9b2b2d3a5fdf8235d56981c36f4039ba", size = 95887, upload-time = "2026-05-08T21:00:11.277Z" }, @@ -4941,7 +4958,7 @@ wheels = [ [[package]] name = "protobuf" version = "6.32.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/c0/df/fb4a8eeea482eca989b51cffd274aac2ee24e825f0bf3cbce5281fa1567b/protobuf-6.32.0.tar.gz", hash = "sha256:a81439049127067fc49ec1d36e25c6ee1d1a2b7be930675f919258d03c04e7d2", size = 440614, upload-time = "2025-08-14T21:21:25.015Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/33/18/df8c87da2e47f4f1dcc5153a81cd6bca4e429803f4069a299e236e4dd510/protobuf-6.32.0-cp310-abi3-win32.whl", hash = "sha256:84f9e3c1ff6fb0308dbacb0950d8aa90694b0d0ee68e75719cb044b7078fe741", size = 424409, upload-time = "2025-08-14T21:21:12.366Z" }, @@ -4955,7 +4972,7 @@ wheels = [ [[package]] name = "psutil" version = "7.2.2" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/aa/c6/d1ddf4abb55e93cebc4f2ed8b5d6dbad109ecb8d63748dd2b20ab5e57ebe/psutil-7.2.2.tar.gz", hash = "sha256:0746f5f8d406af344fd547f1c8daa5f5c33dbc293bb8d6a16d80b4bb88f59372", size = 493740, upload-time = "2026-01-28T18:14:54.428Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/51/08/510cbdb69c25a96f4ae523f733cdc963ae654904e8db864c07585ef99875/psutil-7.2.2-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:2edccc433cbfa046b980b0df0171cd25bcaeb3a68fe9022db0979e7aa74a826b", size = 130595, upload-time = "2026-01-28T18:14:57.293Z" }, @@ -4983,7 +5000,7 @@ wheels = [ [[package]] name = "ptyprocess" version = "0.7.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/20/e5/16ff212c1e452235a90aeb09066144d0c5a6a8c0834397e03f5224495c4e/ptyprocess-0.7.0.tar.gz", hash = "sha256:5c5d0a3b48ceee0b48485e0c26037c0acd7d29765ca3fbb5cb3831d347423220", size = 70762, upload-time = "2020-12-28T15:15:30.155Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/22/a6/858897256d0deac81a172289110f31629fc4cee19b6f01283303e18c8db3/ptyprocess-0.7.0-py2.py3-none-any.whl", hash = "sha256:4b41f3967fce3af57cc7e94b888626c18bf37a083e3651ca8feeb66d492fef35", size = 13993, upload-time = "2020-12-28T15:15:28.35Z" }, @@ -4992,7 +5009,7 @@ wheels = [ [[package]] name = "pure-eval" version = "0.2.3" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/cd/05/0a34433a064256a578f1783a10da6df098ceaa4a57bbeaa96a6c0352786b/pure_eval-0.2.3.tar.gz", hash = "sha256:5f4e983f40564c576c7c8635ae88db5956bb2229d7e9237d03b3c0b0190eaf42", size = 19752, upload-time = "2024-07-21T12:58:21.801Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/8e/37/efad0257dc6e593a18957422533ff0f87ede7c9c6ea010a2177d738fb82f/pure_eval-0.2.3-py3-none-any.whl", hash = "sha256:1db8e35b67b3d218d818ae653e27f06c3aa420901fa7b081ca98cbedc874e0d0", size = 11842, upload-time = "2024-07-21T12:58:20.04Z" }, @@ -5001,7 +5018,7 @@ wheels = [ [[package]] name = "pyarrow" version = "24.0.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/91/13/13e1069b351bdc3881266e11147ffccf687505dbb0ea74036237f5d454a5/pyarrow-24.0.0.tar.gz", hash = "sha256:85fe721a14dd823aca09127acbb06c3ca723efbd436c004f16bca601b04dcc83", size = 1180261, upload-time = "2026-04-21T10:51:25.837Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/b4/a9/9686d9f07837f91f775e8932659192e02c74f9d8920524b480b85212cc68/pyarrow-24.0.0-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:6233c9ed9ab9d1db47de57d9753256d9dcffbf42db341576099f0fd9f6bf4810", size = 34981559, upload-time = "2026-04-21T10:47:22.17Z" }, @@ -5044,7 +5061,7 @@ wheels = [ [[package]] name = "pycparser" version = "3.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/1b/7d/92392ff7815c21062bea51aa7b87d45576f649f16458d78b7cf94b9ab2e6/pycparser-3.0.tar.gz", hash = "sha256:600f49d217304a5902ac3c37e1281c9fe94e4d0489de643a9504c5cdfdfc6b29", size = 103492, upload-time = "2026-01-21T14:26:51.89Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/0c/c3/44f3fbbfa403ea2a7c779186dc20772604442dde72947e7d01069cbe98e3/pycparser-3.0-py3-none-any.whl", hash = "sha256:b727414169a36b7d524c1c3e31839a521725078d7b2ff038656844266160a992", size = 48172, upload-time = "2026-01-21T14:26:50.693Z" }, @@ -5053,7 +5070,7 @@ wheels = [ [[package]] name = "pydantic" version = "2.13.4" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "annotated-types" }, { name = "pydantic-core" }, @@ -5068,7 +5085,7 @@ wheels = [ [[package]] name = "pydantic-core" version = "2.46.4" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "typing-extensions" }, ] @@ -5143,7 +5160,7 @@ wheels = [ [[package]] name = "pygame" version = "2.6.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/49/cc/08bba60f00541f62aaa252ce0cfbd60aebd04616c0b9574f755b583e45ae/pygame-2.6.1.tar.gz", hash = "sha256:56fb02ead529cee00d415c3e007f75e0780c655909aaa8e8bf616ee09c9feb1f", size = 14808125, upload-time = "2024-09-29T13:41:34.698Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/92/16/2c602c332f45ff9526d61f6bd764db5096ff9035433e2172e2d2cadae8db/pygame-2.6.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:4ee7f2771f588c966fa2fa8b829be26698c9b4836f82ede5e4edc1a68594942e", size = 13118279, upload-time = "2024-09-29T14:26:30.427Z" }, @@ -5165,7 +5182,7 @@ wheels = [ [[package]] name = "pygments" version = "2.20.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/c3/b2/bc9c9196916376152d655522fdcebac55e66de6603a76a02bca1b6414f6c/pygments-2.20.0.tar.gz", hash = "sha256:6757cd03768053ff99f3039c1a36d6c0aa0b263438fcab17520b30a303a82b5f", size = 4955991, upload-time = "2026-03-29T13:29:33.898Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/f4/7e/a72dd26f3b0f4f2bf1dd8923c85f7ceb43172af56d63c7383eb62b332364/pygments-2.20.0-py3-none-any.whl", hash = "sha256:81a9e26dd42fd28a23a2d169d86d7ac03b46e2f8b59ed4698fb4785f946d0176", size = 1231151, upload-time = "2026-03-29T13:29:30.038Z" }, @@ -5174,7 +5191,7 @@ wheels = [ [[package]] name = "pymunk" version = "6.11.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "cffi" }, ] @@ -5201,7 +5218,7 @@ wheels = [ [[package]] name = "pyngrok" version = "8.1.2" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "pyyaml" }, ] @@ -5213,7 +5230,7 @@ wheels = [ [[package]] name = "pynput" version = "1.8.2" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "evdev", marker = "'linux' in sys_platform" }, { name = "pyobjc-framework-applicationservices", marker = "sys_platform == 'darwin'" }, @@ -5229,7 +5246,7 @@ wheels = [ [[package]] name = "pyobjc-core" version = "12.2" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/2a/e8/a6cc12669211e7c9b29e8f26bf2159e67c7a73555dc229018abf46d8167a/pyobjc_core-12.2.tar.gz", hash = "sha256:51d7de4cfa32f508c6a7aac31f131b12d5e196a8dcf588e6e8d7e6337224f66d", size = 1062064, upload-time = "2026-05-30T12:29:55.417Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/24/be/4771f4fd786f0e1a2bd6d8931a72a5f3929b7bb1b28a1fe6ca8a08371c55/pyobjc_core-12.2-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:7677ed758a367bbbb5589d6f5276fb360a45c89168276c26162f61840b0fa03d", size = 6421145, upload-time = "2026-05-30T10:17:41.992Z" }, @@ -5244,7 +5261,7 @@ wheels = [ [[package]] name = "pyobjc-framework-applicationservices" version = "12.2" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "pyobjc-core", marker = "sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32'" }, { name = "pyobjc-framework-cocoa", marker = "sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32'" }, @@ -5265,7 +5282,7 @@ wheels = [ [[package]] name = "pyobjc-framework-cocoa" version = "12.2" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "pyobjc-core", marker = "sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32'" }, ] @@ -5283,7 +5300,7 @@ wheels = [ [[package]] name = "pyobjc-framework-coretext" version = "12.2" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "pyobjc-core", marker = "sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32'" }, { name = "pyobjc-framework-cocoa", marker = "sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32'" }, @@ -5303,7 +5320,7 @@ wheels = [ [[package]] name = "pyobjc-framework-quartz" version = "12.2" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "pyobjc-core", marker = "sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32'" }, { name = "pyobjc-framework-cocoa", marker = "sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32'" }, @@ -5322,7 +5339,7 @@ wheels = [ [[package]] name = "pyopengl" version = "3.1.10" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/6f/16/912b7225d56284859cd9a672827f18be43f8012f8b7b932bc4bd959a298e/pyopengl-3.1.10.tar.gz", hash = "sha256:c4a02d6866b54eb119c8e9b3fb04fa835a95ab802dd96607ab4cdb0012df8335", size = 1915580, upload-time = "2025-08-18T02:33:01.76Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/de/e4/1ba6f44e491c4eece978685230dde56b14d51a0365bc1b774ddaa94d14cd/pyopengl-3.1.10-py3-none-any.whl", hash = "sha256:794a943daced39300879e4e47bd94525280685f42dbb5a998d336cfff151d74f", size = 3194996, upload-time = "2025-08-18T02:32:59.902Z" }, @@ -5331,7 +5348,7 @@ wheels = [ [[package]] name = "pyparsing" version = "3.3.2" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/f3/91/9c6ee907786a473bf81c5f53cf703ba0957b23ab84c264080fb5a450416f/pyparsing-3.3.2.tar.gz", hash = "sha256:c777f4d763f140633dcb6d8a3eda953bf7a214dc4eff598413c070bcdc117cbc", size = 6851574, upload-time = "2026-01-21T03:57:59.36Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/10/bd/c038d7cc38edc1aa5bf91ab8068b63d4308c66c4c8bb3cbba7dfbc049f9c/pyparsing-3.3.2-py3-none-any.whl", hash = "sha256:850ba148bd908d7e2411587e247a1e4f0327839c40e2e5e6d05a007ecc69911d", size = 122781, upload-time = "2026-01-21T03:57:55.912Z" }, @@ -5340,7 +5357,7 @@ wheels = [ [[package]] name = "pyquaternion" version = "0.9.9" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "numpy" }, ] @@ -5352,7 +5369,7 @@ wheels = [ [[package]] name = "pyrealsense2" version = "2.56.5.9235" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } wheels = [ { url = "https://files.pythonhosted.org/packages/fc/20/a46e60a496a17f1cf0cbdffb45e66dc015756e4dbce83580fd569e53e178/pyrealsense2-2.56.5.9235-cp312-cp312-manylinux1_x86_64.whl", hash = "sha256:ba2c22981111adbefb169c39e023af4352a2409dfbff59f02c2404c68b82064b", size = 11062766, upload-time = "2025-07-28T14:59:26.551Z" }, { url = "https://files.pythonhosted.org/packages/22/b5/dd8349abac780aed774f65825fc2ed3ca832b0ad2bf3293262bfa9a517b2/pyrealsense2-2.56.5.9235-cp312-cp312-win_amd64.whl", hash = "sha256:e9c64b94cf6170a3ad60416ff1bf969df8aafe383d4bff14e0fa10b2459d885b", size = 7801788, upload-time = "2025-07-28T14:59:28.303Z" }, @@ -5363,7 +5380,7 @@ wheels = [ [[package]] name = "pyrealsense2-macosx" version = "2.56.5" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } wheels = [ { url = "https://files.pythonhosted.org/packages/45/42/df369b9c8943395e8166dfc2006163279a850d7990e9c7fb7a0f0bdc4bbb/pyrealsense2_macosx-2.56.5-cp312-cp312-macosx_15_0_arm64.whl", hash = "sha256:7a0c2e124d74da97c3a928f0dc084c51f1473745a863db08cc2191bd2bdf8623", size = 4509608, upload-time = "2026-02-16T09:56:06.215Z" }, { url = "https://files.pythonhosted.org/packages/6b/f4/b4a0741f8a5e77495e64a88985327515decd86c64553103bbf4b1e11a8b8/pyrealsense2_macosx-2.56.5-cp313-cp313-macosx_15_0_arm64.whl", hash = "sha256:606aa36debd6b1f6179d7ca866161c16d5b60c6a2dfd95eadd10f0d6eee56c02", size = 4510296, upload-time = "2026-02-16T09:56:07.639Z" }, @@ -5375,7 +5392,7 @@ wheels = [ [[package]] name = "pyserial" version = "3.5" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/1e/7d/ae3f0a63f41e4d2f6cb66a5b57197850f919f59e558159a4dd3a818f5082/pyserial-3.5.tar.gz", hash = "sha256:3c77e014170dfffbd816e6ffc205e9842efb10be9f58ec16d3e8675b4925cddb", size = 159125, upload-time = "2020-11-23T03:59:15.045Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/07/bc/587a445451b253b285629263eb51c2d8e9bcea4fc97826266d186f96f558/pyserial-3.5-py2.py3-none-any.whl", hash = "sha256:c4451db6ba391ca6ca299fb3ec7bae67a5c55dde170964c7a14ceefec02f2cf0", size = 90585, upload-time = "2020-11-23T03:59:13.41Z" }, @@ -5384,7 +5401,7 @@ wheels = [ [[package]] name = "pytest" version = "8.4.2" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "colorama", marker = "sys_platform == 'win32'" }, { name = "iniconfig" }, @@ -5400,7 +5417,7 @@ wheels = [ [[package]] name = "pytest-cov" version = "7.1.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "coverage" }, { name = "pluggy" }, @@ -5414,7 +5431,7 @@ wheels = [ [[package]] name = "pytest-timeout" version = "2.4.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "pytest" }, ] @@ -5426,7 +5443,7 @@ wheels = [ [[package]] name = "python-can" version = "4.6.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "packaging" }, { name = "typing-extensions" }, @@ -5440,7 +5457,7 @@ wheels = [ [[package]] name = "python-dateutil" version = "2.9.0.post0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "six" }, ] @@ -5452,7 +5469,7 @@ wheels = [ [[package]] name = "python-discovery" version = "1.4.2" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "filelock" }, { name = "platformdirs" }, @@ -5465,7 +5482,7 @@ wheels = [ [[package]] name = "python-dotenv" version = "1.2.2" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/82/ed/0301aeeac3e5353ef3d94b6ec08bbcabd04a72018415dcb29e588514bba8/python_dotenv-1.2.2.tar.gz", hash = "sha256:2c371a91fbd7ba082c2c1dc1f8bf89ca22564a087c2c287cd9b662adde799cf3", size = 50135, upload-time = "2026-03-01T16:00:26.196Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/0b/d7/1959b9648791274998a9c3526f6d0ec8fd2233e4d4acce81bbae76b44b2a/python_dotenv-1.2.2-py3-none-any.whl", hash = "sha256:1d8214789a24de455a8b8bd8ae6fe3c6b69a5e3d64aa8a8e5d68e694bbcb285a", size = 22101, upload-time = "2026-03-01T16:00:25.09Z" }, @@ -5474,7 +5491,7 @@ wheels = [ [[package]] name = "python-json-logger" version = "4.1.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/f7/ff/3cc9165fd44106973cd7ac9facb674a65ed853494592541d339bdc9a30eb/python_json_logger-4.1.0.tar.gz", hash = "sha256:b396b9e3ed782b09ff9d6e4f1683d46c83ad0d35d2e407c09a9ebbf038f88195", size = 17573, upload-time = "2026-03-29T04:39:56.805Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/27/be/0631a861af4d1c875f096c07d34e9a63639560a717130e7a87cbc82b7e3f/python_json_logger-4.1.0-py3-none-any.whl", hash = "sha256:132994765cf75bf44554be9aa49b06ef2345d23661a96720262716438141b6b2", size = 15021, upload-time = "2026-03-29T04:39:55.266Z" }, @@ -5483,7 +5500,7 @@ wheels = [ [[package]] name = "python-xlib" version = "0.33" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "six" }, ] @@ -5495,7 +5512,7 @@ wheels = [ [[package]] name = "pytz" version = "2026.2" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/ff/46/dd499ec9038423421951e4fad73051febaa13d2df82b4064f87af8b8c0c3/pytz-2026.2.tar.gz", hash = "sha256:0e60b47b29f21574376f218fe21abc009894a2321ea16c6754f3cad6eb7cdd6a", size = 320861, upload-time = "2026-05-04T01:35:29.667Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/ec/dd/96da98f892250475bdf2328112d7468abdd4acc7b902b6af23f4ed958ea0/pytz-2026.2-py2.py3-none-any.whl", hash = "sha256:04156e608bee23d3792fd45c94ae47fae1036688e75032eea2e3bf0323d1f126", size = 510141, upload-time = "2026-05-04T01:35:27.408Z" }, @@ -5504,7 +5521,7 @@ wheels = [ [[package]] name = "pywinpty" version = "3.0.5" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/a9/ef/2d27f30c59a67be7025b2d7858c8c2d282b74d66544b2384730b82de74fd/pywinpty-3.0.5.tar.gz", hash = "sha256:61db0db063de9865adbea66db294628f8577f608d9764a4c7d3384eeacc4e81b", size = 16223484, upload-time = "2026-06-11T00:11:58.93Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/45/34/942cc95ca4e26489875aa8a95192766247a687379ec29543eebe73ec945f/pywinpty-3.0.5-cp312-cp312-win_amd64.whl", hash = "sha256:d62946adf14b15b54c0b8d785f93fe18b04da23f4ad59e2e8c4612646e9abd23", size = 2090915, upload-time = "2026-06-10T23:43:14.98Z" }, @@ -5522,7 +5539,7 @@ wheels = [ [[package]] name = "pyyaml" version = "6.0.3" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/05/8e/961c0007c59b8dd7729d542c61a4d537767a59645b82a0b521206e1e25c2/pyyaml-6.0.3.tar.gz", hash = "sha256:d76623373421df22fb4cf8817020cbb7ef15c725b9d5e45f17e189bfc384190f", size = 130960, upload-time = "2025-09-25T21:33:16.546Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/d1/33/422b98d2195232ca1826284a76852ad5a86fe23e31b009c9886b2d0fb8b2/pyyaml-6.0.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:7f047e29dcae44602496db43be01ad42fc6f1cc0d8cd6c83d342306c32270196", size = 182063, upload-time = "2025-09-25T21:32:11.445Z" }, @@ -5568,7 +5585,7 @@ wheels = [ [[package]] name = "pyyaml-include" version = "1.4.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "pyyaml" }, ] @@ -5580,7 +5597,7 @@ wheels = [ [[package]] name = "pyzmq" version = "27.1.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "cffi", marker = "implementation_name == 'pypy'" }, ] @@ -5623,7 +5640,7 @@ wheels = [ [[package]] name = "qwen-vl-utils" version = "0.0.14" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "av" }, { name = "packaging" }, @@ -5638,7 +5655,7 @@ wheels = [ [[package]] name = "reachy2-sdk" version = "1.0.15" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "grpcio" }, { name = "numpy" }, @@ -5656,7 +5673,7 @@ wheels = [ [[package]] name = "reachy2-sdk-api" version = "1.0.21" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "grpcio" }, { name = "grpcio-tools" }, @@ -5670,7 +5687,7 @@ wheels = [ [[package]] name = "referencing" version = "0.37.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "attrs" }, { name = "rpds-py" }, @@ -5684,7 +5701,7 @@ wheels = [ [[package]] name = "regex" version = "2026.5.9" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/dc/0e/49aee608ad09480e7fd276898c99ec6192985fa331abe4eb3a986094490b/regex-2026.5.9.tar.gz", hash = "sha256:a8234aa23ec39894bfe4a3f1b85616a7032481964a13ac6fc9f10de4f6fca270", size = 416074, upload-time = "2026-05-09T23:15:19.37Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/50/9b/6550044bc44e17c84d312c031c2ec42fbdb6a4ec4e29093be3a172d08772/regex-2026.5.9-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:57eeeb05db7979413dec5438f2db21d7ecbba787cde7a711df1a6f6df672aa06", size = 490451, upload-time = "2026-05-09T23:12:34.72Z" }, @@ -5772,7 +5789,7 @@ wheels = [ [[package]] name = "requests" version = "2.34.2" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "certifi" }, { name = "charset-normalizer" }, @@ -5786,27 +5803,27 @@ wheels = [ [[package]] name = "rerun-sdk" -version = "0.26.2" -source = { registry = "https://pypi.registries.huggingface.tech/" } +version = "0.33.1" +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "attrs" }, { name = "numpy" }, { name = "pillow" }, + { name = "psutil" }, { name = "pyarrow" }, { name = "typing-extensions" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/4b/4a/767c20e1529d74d9be5b5e55c6c26b63a6918ef3c1709fc422d08a460114/rerun_sdk-0.26.2-cp39-abi3-macosx_10_12_x86_64.whl", hash = "sha256:3d4151c9a3484e112b53d1df90c8fa07397dc7b8bfbb420f09e011eff20f1ef2", size = 93349439, upload-time = "2025-10-27T11:34:10.745Z" }, - { url = "https://files.pythonhosted.org/packages/2b/3d/d8dd0af9c287a85d51ec99d69406cc4b94a9feb1d6f192d3bbcaac9f0b81/rerun_sdk-0.26.2-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:03977d2aba4966d9a70b682eca196123fda11408fecd733441ede9916c6341e2", size = 86323042, upload-time = "2025-10-27T11:34:17.995Z" }, - { url = "https://files.pythonhosted.org/packages/13/29/53d8d98799ab32418fd4ba6834d6a5749c31f56160d3c87f52a7219887e9/rerun_sdk-0.26.2-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:b6128c3c4f014cae5be18e4d37657c5932d1bcdb2ce5e9d4b488a6eed47f7437", size = 92677274, upload-time = "2025-10-27T11:34:22.601Z" }, - { url = "https://files.pythonhosted.org/packages/f5/86/0b9c8f56398b4fc85f8e99279907c258413a297e5603f8f2537fe5806e51/rerun_sdk-0.26.2-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:a6f97b60aaa7d4e8c6124a3f6b97ce9dbd09520050955f0e0bdacb72b0eb106a", size = 98768129, upload-time = "2025-10-27T11:34:27.36Z" }, - { url = "https://files.pythonhosted.org/packages/be/e7/99fc91c0f99f69d7d43e1db0a6f6cb8273ffc02111539bfc1fee43749bad/rerun_sdk-0.26.2-cp39-abi3-win_amd64.whl", hash = "sha256:a493ad6c8357022cba2ca6f8954a81d0faf984b0b22154eb1d976bfc7649df63", size = 84267089, upload-time = "2025-10-27T11:34:32.023Z" }, + { url = "https://files.pythonhosted.org/packages/16/07/380198590b194f1c17052d672865aa4a56e606eae47665f66edfb391999d/rerun_sdk-0.33.1-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:c0115b710289022587bd2e9ecb715f98d1e87dd7d8ac48e053324131d4addc89", size = 125706253, upload-time = "2026-06-22T09:04:18.974Z" }, + { url = "https://files.pythonhosted.org/packages/f8/eb/6741bbf6868175ab126aff58d372066241c6cd2fc1c4f82ed64069728e73/rerun_sdk-0.33.1-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:14451d31bc7bd0f7c6bcc9c1213ed679ab81b65ecd8b36eae99f738219897dc5", size = 135278374, upload-time = "2026-06-22T09:04:26.768Z" }, + { url = "https://files.pythonhosted.org/packages/04/28/fd3b832652900fe3415739e7411c8af8af4c44e9e1a9d55d79e37f7f9094/rerun_sdk-0.33.1-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:0f87da7a270614074aca37846350f9e257f65081345474748f578c7da64fdeba", size = 139565470, upload-time = "2026-06-22T09:04:34.941Z" }, + { url = "https://files.pythonhosted.org/packages/19/41/339920f5a6734054c07bcae543365a7ef3368ceee3eb67906e2e38bd1d67/rerun_sdk-0.33.1-cp310-abi3-win_amd64.whl", hash = "sha256:b2c2af67f3c2a85b282669d97b52596593fcfdd19bf57c423f18827c837a6e49", size = 120411717, upload-time = "2026-06-22T09:04:42.643Z" }, ] [[package]] name = "rfc3339-validator" version = "0.1.4" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "six" }, ] @@ -5818,7 +5835,7 @@ wheels = [ [[package]] name = "rfc3986-validator" version = "0.1.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/da/88/f270de456dd7d11dcc808abfa291ecdd3f45ff44e3b549ffa01b126464d0/rfc3986_validator-0.1.1.tar.gz", hash = "sha256:3d44bde7921b3b9ec3ae4e3adca370438eccebc676456449b145d533b240d055", size = 6760, upload-time = "2019-10-28T16:00:19.144Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/9e/51/17023c0f8f1869d8806b979a2bffa3f861f26a3f1a66b094288323fba52f/rfc3986_validator-0.1.1-py2.py3-none-any.whl", hash = "sha256:2f235c432ef459970b4306369336b9d5dbdda31b510ca1e327636e01f528bfa9", size = 4242, upload-time = "2019-10-28T16:00:13.976Z" }, @@ -5827,7 +5844,7 @@ wheels = [ [[package]] name = "rfc3987-syntax" version = "1.1.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "lark" }, ] @@ -5839,7 +5856,7 @@ wheels = [ [[package]] name = "rhoban-cmeel-jsoncpp" version = "1.9.4.9" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "cmeel" }, ] @@ -5866,7 +5883,7 @@ wheels = [ [[package]] name = "rich" version = "15.0.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "markdown-it-py" }, { name = "pygments" }, @@ -5879,7 +5896,7 @@ wheels = [ [[package]] name = "robomimic" version = "0.2.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "egl-probe", marker = "sys_platform == 'linux'" }, { name = "h5py", marker = "sys_platform == 'linux'" }, @@ -5899,7 +5916,7 @@ sdist = { url = "https://files.pythonhosted.org/packages/3d/c3/44b1d1ea4bcb4bbed [[package]] name = "robosuite" version = "1.4.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "mujoco", marker = "sys_platform == 'linux'" }, { name = "numba", marker = "sys_platform == 'linux'" }, @@ -5916,7 +5933,7 @@ wheels = [ [[package]] name = "rpds-py" version = "2026.5.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/2e/43/25a8dcd3feedd735039a8f0b5b7e3b118232b5eae288c4fd9ab200d41094/rpds_py-2026.5.1.tar.gz", hash = "sha256:07b24fea40541e28570e5b795a4a38fbdcd12550c06bd0748005ecc8116ca256", size = 64459, upload-time = "2026-05-28T12:02:13.232Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/d4/e7/a78582dc57caa592dcc7d4fb69b61390561e908eb3d2f5df5928a8e354c0/rpds_py-2026.5.1-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:3abe24a66e57adcfa645d718063a5fa5103ecc71ddbf26d78af8f9368018ff1d", size = 353040, upload-time = "2026-05-28T11:59:12.531Z" }, @@ -6026,7 +6043,7 @@ wheels = [ [[package]] name = "ruff" version = "0.15.17" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/8c/a9/3abdf488f1bf3d24c699415e454ed554a6350d5d89ce183be1ee0a3361ac/ruff-0.15.17.tar.gz", hash = "sha256:2ec446937fd16c8c4de2674a209cc5af64d9c6f17d21fbf1151054fa0bcf5219", size = 4743346, upload-time = "2026-06-11T17:54:47.663Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/db/4d/e11259f5da07cb6afb2d074c31bf09da9671993f7329d4f15d2fdc458301/ruff-0.15.17-py3-none-linux_armv6l.whl", hash = "sha256:d9feddb927fc68bd295f5eebc587a7e42cfaf9b65f60ca4a2386febff575da8f", size = 10856677, upload-time = "2026-06-11T17:54:49.533Z" }, @@ -6051,7 +6068,7 @@ wheels = [ [[package]] name = "safetensors" version = "0.8.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/45/06/f955dbbb1859e3bd23c8ac6141af5106e7ad5fedec4a3a6e3d60f94b7001/safetensors-0.8.0.tar.gz", hash = "sha256:fabaf3e0f18a6618d9b36560682562157f77c2b71fcffc7b432be2baed9d753d", size = 325846, upload-time = "2026-06-09T07:52:25.563Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/39/a0/f718cda65b05407d228f97602cf60dca269c979867aa5beb25410de26cd3/safetensors-0.8.0-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:c554f85858e05226d3c2828e32395e677434685d6d94594a41643361c5e837f0", size = 473568, upload-time = "2026-06-09T07:52:18.829Z" }, @@ -6075,7 +6092,7 @@ wheels = [ [[package]] name = "scikit-image" version = "0.25.2" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "imageio" }, { name = "lazy-loader" }, @@ -6104,7 +6121,7 @@ wheels = [ [[package]] name = "scipy" version = "1.17.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "numpy" }, ] @@ -6165,7 +6182,7 @@ wheels = [ [[package]] name = "send2trash" version = "2.1.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/c5/f0/184b4b5f8d00f2a92cf96eec8967a3d550b52cf94362dad1100df9e48d57/send2trash-2.1.0.tar.gz", hash = "sha256:1c72b39f09457db3c05ce1d19158c2cbef4c32b8bedd02c155e49282b7ea7459", size = 17255, upload-time = "2026-01-14T06:27:36.056Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/1c/78/504fdd027da3b84ff1aecd9f6957e65f35134534ccc6da8628eb71e76d3f/send2trash-2.1.0-py3-none-any.whl", hash = "sha256:0da2f112e6d6bb22de6aa6daa7e144831a4febf2a87261451c4ad849fe9a873c", size = 17610, upload-time = "2026-01-14T06:27:35.218Z" }, @@ -6174,7 +6191,7 @@ wheels = [ [[package]] name = "sentry-sdk" version = "2.62.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "certifi" }, { name = "urllib3" }, @@ -6187,7 +6204,7 @@ wheels = [ [[package]] name = "setuptools" version = "80.10.2" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/76/95/faf61eb8363f26aa7e1d762267a8d602a1b26d4f3a1e758e92cb3cb8b054/setuptools-80.10.2.tar.gz", hash = "sha256:8b0e9d10c784bf7d262c4e5ec5d4ec94127ce206e8738f29a437945fbc219b70", size = 1200343, upload-time = "2026-01-25T22:38:17.252Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/94/b8/f1f62a5e3c0ad2ff1d189590bfa4c46b4f3b6e49cef6f26c6ee4e575394d/setuptools-80.10.2-py3-none-any.whl", hash = "sha256:95b30ddfb717250edb492926c92b5221f7ef3fbcc2b07579bcd4a27da21d0173", size = 1064234, upload-time = "2026-01-25T22:38:15.216Z" }, @@ -6196,7 +6213,7 @@ wheels = [ [[package]] name = "shapely" version = "2.1.2" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "numpy" }, ] @@ -6247,7 +6264,7 @@ wheels = [ [[package]] name = "shellingham" version = "1.5.4" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/58/15/8b3609fd3830ef7b27b655beb4b4e9c62313a4e8da8c676e142cc210d58e/shellingham-1.5.4.tar.gz", hash = "sha256:8dbca0739d487e5bd35ab3ca4b36e11c4078f3a234bfce294b0a0291363404de", size = 10310, upload-time = "2023-10-24T04:13:40.426Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/e0/f9/0595336914c5619e5f28a1fb793285925a8cd4b432c9da0a987836c7f822/shellingham-1.5.4-py2.py3-none-any.whl", hash = "sha256:7ecfff8f2fd72616f7481040475a65b2bf8af90a56c89140852d1120324e8686", size = 9755, upload-time = "2023-10-24T04:13:38.866Z" }, @@ -6256,7 +6273,7 @@ wheels = [ [[package]] name = "six" version = "1.17.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/94/e7/b2c673351809dca68a0e064b6af791aa332cf192da575fd474ed7d6f16a2/six-1.17.0.tar.gz", hash = "sha256:ff70335d468e7eb6ec65b95b99d3a2836546063f63acc5171de367e834932a81", size = 34031, upload-time = "2024-12-04T17:35:28.174Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/b7/ce/149a00dd41f10bc29e5921b496af8b574d8413afcd5e30dfa0ed46c2cc5e/six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274", size = 11050, upload-time = "2024-12-04T17:35:26.475Z" }, @@ -6265,7 +6282,7 @@ wheels = [ [[package]] name = "smmap" version = "5.0.3" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/1f/ea/49c993d6dfdd7338c9b1000a0f36817ed7ec84577ae2e52f890d1a4ff909/smmap-5.0.3.tar.gz", hash = "sha256:4d9debb8b99007ae47165abc08670bd74cb74b5227dda7f643eccc4e9eb5642c", size = 22506, upload-time = "2026-03-09T03:43:26.1Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/c1/d4/59e74daffcb57a07668852eeeb6035af9f32cbfd7a1d2511f17d2fe6a738/smmap-5.0.3-py3-none-any.whl", hash = "sha256:c106e05d5a61449cf6ba9a1e650227ecfb141590d2a98412103ff35d89fc7b2f", size = 24390, upload-time = "2026-03-09T03:43:24.361Z" }, @@ -6274,7 +6291,7 @@ wheels = [ [[package]] name = "sniffio" version = "1.3.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/a2/87/a6771e1546d97e7e041b6ae58d80074f81b7d5121207425c964ddf5cfdbd/sniffio-1.3.1.tar.gz", hash = "sha256:f4324edc670a0f49750a81b895f35c3adb843cca46f0530f79fc1babb23789dc", size = 20372, upload-time = "2024-02-25T23:20:04.057Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/e9/44/75a9c9421471a6c4805dbf2356f7c181a29c1879239abab1ea2cc8f38b40/sniffio-1.3.1-py3-none-any.whl", hash = "sha256:2f6da418d1f1e0fddd844478f41680e794e6051915791a034ff65e5f100525a2", size = 10235, upload-time = "2024-02-25T23:20:01.196Z" }, @@ -6283,7 +6300,7 @@ wheels = [ [[package]] name = "soupsieve" version = "2.8.4" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/47/2c/0a5f6f8ee0d5589e48c7640213ed5175d52cf540a06725b628cc1a45d6ce/soupsieve-2.8.4.tar.gz", hash = "sha256:e121fd02e975c695e4e9e8774a5ee35d74714b59307868dcc5319ad2d9e3328e", size = 121110, upload-time = "2026-05-24T13:55:57.154Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/5e/f5/0c41cb68dcae6b7de4fac4188a3a9589e21fb31df21ea3a2e888db95e6c9/soupsieve-2.8.4-py3-none-any.whl", hash = "sha256:e7e6b0769c8f51ed59acab6e994b00621096cfb1c640a7509295987388fbaf65", size = 37304, upload-time = "2026-05-24T13:55:55.406Z" }, @@ -6292,7 +6309,7 @@ wheels = [ [[package]] name = "stack-data" version = "0.6.3" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "asttokens" }, { name = "executing" }, @@ -6306,7 +6323,7 @@ wheels = [ [[package]] name = "starlette" version = "1.3.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "anyio" }, { name = "typing-extensions", marker = "python_full_version < '3.13'" }, @@ -6319,7 +6336,7 @@ wheels = [ [[package]] name = "sympy" version = "1.14.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "mpmath" }, ] @@ -6331,7 +6348,7 @@ wheels = [ [[package]] name = "teleop" version = "0.1.5" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "fastapi" }, { name = "numpy" }, @@ -6349,7 +6366,7 @@ wheels = [ [[package]] name = "tensorboard" version = "2.20.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "absl-py", marker = "sys_platform == 'linux'" }, { name = "grpcio", marker = "sys_platform == 'linux'" }, @@ -6369,7 +6386,7 @@ wheels = [ [[package]] name = "tensorboard-data-server" version = "0.7.2" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } wheels = [ { url = "https://files.pythonhosted.org/packages/7a/13/e503968fefabd4c6b2650af21e110aa8466fe21432cd7c43a84577a89438/tensorboard_data_server-0.7.2-py3-none-any.whl", hash = "sha256:7e0610d205889588983836ec05dc098e80f97b7e7bbff7e994ebb78f578d0ddb", size = 2356, upload-time = "2023-10-23T21:23:32.16Z" }, { url = "https://files.pythonhosted.org/packages/73/c6/825dab04195756cf8ff2e12698f22513b3db2f64925bdd41671bfb33aaa5/tensorboard_data_server-0.7.2-py3-none-manylinux_2_31_x86_64.whl", hash = "sha256:ef687163c24185ae9754ed5650eb5bc4d84ff257aabdc33f0cc6f74d8ba54530", size = 6590363, upload-time = "2023-10-23T21:23:35.583Z" }, @@ -6378,7 +6395,7 @@ wheels = [ [[package]] name = "tensorboardx" version = "2.6.5" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "numpy", marker = "sys_platform == 'linux'" }, { name = "packaging", marker = "sys_platform == 'linux'" }, @@ -6392,7 +6409,7 @@ wheels = [ [[package]] name = "termcolor" version = "3.3.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/46/79/cf31d7a93a8fdc6aa0fbb665be84426a8c5a557d9240b6239e9e11e35fc5/termcolor-3.3.0.tar.gz", hash = "sha256:348871ca648ec6a9a983a13ab626c0acce02f515b9e1983332b17af7979521c5", size = 14434, upload-time = "2025-12-29T12:55:21.882Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/33/d1/8bb87d21e9aeb323cc03034f5eaf2c8f69841e40e4853c2627edf8111ed3/termcolor-3.3.0-py3-none-any.whl", hash = "sha256:cf642efadaf0a8ebbbf4bc7a31cec2f9b5f21a9f726f4ccbb08192c9c26f43a5", size = 7734, upload-time = "2025-12-29T12:55:20.718Z" }, @@ -6401,7 +6418,7 @@ wheels = [ [[package]] name = "terminado" version = "0.18.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "ptyprocess", marker = "os_name != 'nt'" }, { name = "pywinpty", marker = "(os_name == 'nt' and platform_machine != 'arm64' and sys_platform == 'darwin') or (os_name == 'nt' and sys_platform != 'darwin' and sys_platform != 'linux')" }, @@ -6415,7 +6432,7 @@ wheels = [ [[package]] name = "thop" version = "0.1.1.post2209072238" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "torch", version = "2.11.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" }, marker = "sys_platform == 'linux'" }, ] @@ -6426,7 +6443,7 @@ wheels = [ [[package]] name = "tifffile" version = "2026.6.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "numpy" }, ] @@ -6438,14 +6455,14 @@ wheels = [ [[package]] name = "timm" version = "1.0.27" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "huggingface-hub" }, { name = "pyyaml" }, { name = "safetensors" }, - { name = "torch", version = "2.11.0", source = { registry = "https://pypi.registries.huggingface.tech/" }, marker = "sys_platform != 'linux'" }, + { name = "torch", version = "2.11.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'linux'" }, { name = "torch", version = "2.11.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" }, marker = "sys_platform == 'linux'" }, - { name = "torchvision", version = "0.26.0", source = { registry = "https://pypi.registries.huggingface.tech/" }, marker = "sys_platform != 'linux'" }, + { name = "torchvision", version = "0.26.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'linux'" }, { name = "torchvision", version = "0.26.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" }, marker = "sys_platform == 'linux'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/08/54/ece85b0eef3700c90db8271a43669b05a0ebbe2edb1962329c34374a297e/timm-1.0.27.tar.gz", hash = "sha256:315dfe63186ca9fb7ff941268941231fd5be259f2b4bb4afa28560ae1015cb9a", size = 2439861, upload-time = "2026-05-08T19:38:36.844Z" } @@ -6456,7 +6473,7 @@ wheels = [ [[package]] name = "tinycss2" version = "1.5.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "webencodings" }, ] @@ -6468,7 +6485,7 @@ wheels = [ [[package]] name = "tokenizers" version = "0.22.2" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "huggingface-hub" }, ] @@ -6494,7 +6511,7 @@ wheels = [ [[package]] name = "toml" version = "0.10.2" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/be/ba/1f744cdc819428fc6b5084ec34d9b30660f6f9daaf70eead706e3203ec3c/toml-0.10.2.tar.gz", hash = "sha256:b3bda1d108d5dd99f4a20d24d9c348e91c4db7ab1b749200bded2f839ccbe68f", size = 22253, upload-time = "2020-11-01T01:40:22.204Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/44/6f/7120676b6d73228c96e17f1f794d8ab046fc910d781c8d151120c3f1569e/toml-0.10.2-py2.py3-none-any.whl", hash = "sha256:806143ae5bfb6a3c6e736a764057db0e6a0e05e338b5630894a5f779cabb4f9b", size = 16588, upload-time = "2020-11-01T01:40:20.672Z" }, @@ -6503,7 +6520,7 @@ wheels = [ [[package]] name = "torch" version = "2.11.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } resolution-markers = [ "python_full_version >= '3.15' and platform_machine == 'arm64' and sys_platform == 'darwin'", "python_full_version == '3.14.*' and platform_machine == 'arm64' and sys_platform == 'darwin'", @@ -6610,7 +6627,7 @@ wheels = [ [[package]] name = "torchcodec" version = "0.11.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } wheels = [ { url = "https://files.pythonhosted.org/packages/64/85/38f4843ff2a6bf7dfb71a153acd99024dadb96749965a67524c2f1cc1894/torchcodec-0.11.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:57056e91d1d883d0fb77ca7759e304be9c0bdb4ea0e37bde5c2e361347063b8c", size = 4368988, upload-time = "2026-04-14T18:24:51.46Z" }, { url = "https://files.pythonhosted.org/packages/4b/85/3b41034b0f1289423745f918ace2a1e1e86b9c578c2e2461b6afcbb5354a/torchcodec-0.11.1-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:f1aee486a84247fcaa67870ac5005aa8d382a9839e91e476fa71b5b3d9fda9b7", size = 2397532, upload-time = "2026-04-14T18:24:53.368Z" }, @@ -6629,10 +6646,10 @@ wheels = [ [[package]] name = "torchdiffeq" version = "0.2.5" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "scipy" }, - { name = "torch", version = "2.11.0", source = { registry = "https://pypi.registries.huggingface.tech/" }, marker = "sys_platform != 'linux'" }, + { name = "torch", version = "2.11.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'linux'" }, { name = "torch", version = "2.11.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" }, marker = "sys_platform == 'linux'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/87/ec/a40aa124660f0ee65e6760cb53df6a82ad91a1a3ef1da5e747f1336644dd/torchdiffeq-0.2.5.tar.gz", hash = "sha256:b50d3760d13fd138dcceac651f4b80396f44fefcebd037a033fecfeaa9cc12e7", size = 31197, upload-time = "2024-11-21T20:20:11.552Z" } @@ -6643,7 +6660,7 @@ wheels = [ [[package]] name = "torchvision" version = "0.26.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } resolution-markers = [ "python_full_version >= '3.15' and platform_machine == 'arm64' and sys_platform == 'darwin'", "python_full_version == '3.14.*' and platform_machine == 'arm64' and sys_platform == 'darwin'", @@ -6677,7 +6694,7 @@ resolution-markers = [ dependencies = [ { name = "numpy", marker = "sys_platform != 'linux'" }, { name = "pillow", marker = "sys_platform != 'linux'" }, - { name = "torch", version = "2.11.0", source = { registry = "https://pypi.registries.huggingface.tech/" }, marker = "sys_platform != 'linux'" }, + { name = "torch", version = "2.11.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'linux'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/ae/e7/56b47cc3b132aea90ccce22bcb8975dec688b002150012acc842846039d0/torchvision-0.26.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:c409e1c3fdebec7a3834465086dbda8bf7680eff79abf7fd2f10c6b59520a7a4", size = 1863502, upload-time = "2026-03-23T18:12:57.326Z" }, @@ -6735,7 +6752,7 @@ wheels = [ [[package]] name = "tornado" version = "6.5.7" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/64/24/95ec527ad67b76d59299e5465b3935d05e4294b7e0290a3924b7487df30b/tornado-6.5.7.tar.gz", hash = "sha256:66c513a76cda70d53907bc27cf1447557699c2e95aa48ba27a442ff61c3ddfc2", size = 519252, upload-time = "2026-06-08T17:34:51.232Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/02/dc/c7043cab6fed8ae159fc1923ce829ada35c4dbd797d408a43858ffaf9639/tornado-6.5.7-cp39-abi3-macosx_10_9_universal2.whl", hash = "sha256:148b2eb15c2c765a50796172c1e499649b35f30d2e3c3d3e15913cfa56bfb163", size = 448543, upload-time = "2026-06-08T17:34:38.052Z" }, @@ -6752,7 +6769,7 @@ wheels = [ [[package]] name = "tqdm" version = "4.68.2" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "colorama", marker = "sys_platform == 'win32'" }, ] @@ -6764,7 +6781,7 @@ wheels = [ [[package]] name = "traitlets" version = "5.15.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/57/a9/a2584b8313b89f94869ddb3c4074617a691de1812a614d2d50e32ca5a7a6/traitlets-5.15.1.tar.gz", hash = "sha256:7b1c07854fe25acb39e009bae49f11b79ff6cbb2f27999104e9110e7a6b53722", size = 163344, upload-time = "2026-06-03T12:26:06.181Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/96/8d/1080ee4c231f361b6ce4470d556c8c435b67c7e0753aaa641497ee92f88b/traitlets-5.15.1-py3-none-any.whl", hash = "sha256:770a53705f84b81ac107e83a1b3328ff2dae16094d8fc3cfc004e4b22dfd8e92", size = 85858, upload-time = "2026-06-03T12:26:04.395Z" }, @@ -6773,7 +6790,7 @@ wheels = [ [[package]] name = "transformers" version = "5.5.4" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "huggingface-hub" }, { name = "numpy" }, @@ -6793,7 +6810,7 @@ wheels = [ [[package]] name = "transforms3d" version = "0.4.2" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "numpy" }, ] @@ -6805,7 +6822,7 @@ wheels = [ [[package]] name = "triton" version = "3.6.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } wheels = [ { url = "https://files.pythonhosted.org/packages/17/5d/08201db32823bdf77a0e2b9039540080b2e5c23a20706ddba942924ebcd6/triton-3.6.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:374f52c11a711fd062b4bfbb201fd9ac0a5febd28a96fb41b4a0f51dde3157f4", size = 176128243, upload-time = "2026-01-20T16:16:07.857Z" }, { url = "https://files.pythonhosted.org/packages/ab/a8/cdf8b3e4c98132f965f88c2313a4b493266832ad47fb52f23d14d4f86bb5/triton-3.6.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:74caf5e34b66d9f3a429af689c1c7128daba1d8208df60e81106b115c00d6fca", size = 188266850, upload-time = "2026-01-20T16:00:43.041Z" }, @@ -6822,7 +6839,7 @@ wheels = [ [[package]] name = "typer" version = "0.25.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "annotated-doc" }, { name = "click" }, @@ -6837,7 +6854,7 @@ wheels = [ [[package]] name = "typing-extensions" version = "4.15.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/72/94/1a15dd82efb362ac84269196e94cf00f187f7ed21c242792a923cdb1c61f/typing_extensions-4.15.0.tar.gz", hash = "sha256:0cea48d173cc12fa28ecabc3b837ea3cf6f38c6d1136f85cbaaf598984861466", size = 109391, upload-time = "2025-08-25T13:49:26.313Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/18/67/36e9267722cc04a6b9f15c7f3441c2363321a3ea07da7ae0c0707beb2a9c/typing_extensions-4.15.0-py3-none-any.whl", hash = "sha256:f0fa19c6845758ab08074a0cfa8b7aecb71c999ca73d62883bc25cc018c4e548", size = 44614, upload-time = "2025-08-25T13:49:24.86Z" }, @@ -6846,7 +6863,7 @@ wheels = [ [[package]] name = "typing-inspect" version = "0.9.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "mypy-extensions" }, { name = "typing-extensions" }, @@ -6859,7 +6876,7 @@ wheels = [ [[package]] name = "typing-inspection" version = "0.4.2" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "typing-extensions" }, ] @@ -6871,7 +6888,7 @@ wheels = [ [[package]] name = "tzdata" version = "2026.2" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/ba/19/1b9b0e29f30c6d35cb345486df41110984ea67ae69dddbc0e8a100999493/tzdata-2026.2.tar.gz", hash = "sha256:9173fde7d80d9018e02a662e168e5a2d04f87c41ea174b139fbef642eda62d10", size = 198254, upload-time = "2026-04-24T15:22:08.651Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/ce/e4/dccd7f47c4b64213ac01ef921a1337ee6e30e8c6466046018326977efd95/tzdata-2026.2-py2.py3-none-any.whl", hash = "sha256:bbe9af844f658da81a5f95019480da3a89415801f6cc966806612cc7169bffe7", size = 349321, upload-time = "2026-04-24T15:22:05.876Z" }, @@ -6880,7 +6897,7 @@ wheels = [ [[package]] name = "u-msgpack-python" version = "2.8.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/36/9d/a40411a475e7d4838994b7f6bcc6bfca9acc5b119ce3a7503608c4428b49/u-msgpack-python-2.8.0.tar.gz", hash = "sha256:b801a83d6ed75e6df41e44518b4f2a9c221dc2da4bcd5380e3a0feda520bc61a", size = 18167, upload-time = "2023-05-18T09:28:12.187Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/b1/5e/512aeb40fd819f4660d00f96f5c7371ee36fc8c6b605128c5ee59e0b28c6/u_msgpack_python-2.8.0-py2.py3-none-any.whl", hash = "sha256:1d853d33e78b72c4228a2025b4db28cda81214076e5b0422ed0ae1b1b2bb586a", size = 10590, upload-time = "2023-05-18T09:28:10.323Z" }, @@ -6889,7 +6906,7 @@ wheels = [ [[package]] name = "uri-template" version = "1.3.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/31/c7/0336f2bd0bcbada6ccef7aaa25e443c118a704f828a0620c6fa0207c1b64/uri-template-1.3.0.tar.gz", hash = "sha256:0e00f8eb65e18c7de20d595a14336e9f337ead580c70934141624b6d1ffdacc7", size = 21678, upload-time = "2023-06-21T01:49:05.374Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/e7/00/3fca040d7cf8a32776d3d81a00c8ee7457e00f80c649f1e4a863c8321ae9/uri_template-1.3.0-py3-none-any.whl", hash = "sha256:a44a133ea12d44a0c0f06d7d42a52d71282e77e2f937d8abd5655b8d56fc1363", size = 11140, upload-time = "2023-06-21T01:49:03.467Z" }, @@ -6898,7 +6915,7 @@ wheels = [ [[package]] name = "urllib3" version = "2.7.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/53/0c/06f8b233b8fd13b9e5ee11424ef85419ba0d8ba0b3138bf360be2ff56953/urllib3-2.7.0.tar.gz", hash = "sha256:231e0ec3b63ceb14667c67be60f2f2c40a518cb38b03af60abc813da26505f4c", size = 433602, upload-time = "2026-05-07T16:13:18.596Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/7f/3e/5db95bcf282c52709639744ca2a8b149baccf648e39c8cc87553df9eae0c/urllib3-2.7.0-py3-none-any.whl", hash = "sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897", size = 131087, upload-time = "2026-05-07T16:13:17.151Z" }, @@ -6907,7 +6924,7 @@ wheels = [ [[package]] name = "uvicorn" version = "0.49.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "click" }, { name = "h11" }, @@ -6931,7 +6948,7 @@ standard = [ [[package]] name = "uvloop" version = "0.22.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/06/f0/18d39dbd1971d6d62c4629cc7fa67f74821b0dc1f5a77af43719de7936a7/uvloop-0.22.1.tar.gz", hash = "sha256:6c84bae345b9147082b17371e3dd5d42775bddce91f885499017f4607fdaf39f", size = 2443250, upload-time = "2025-10-16T22:17:19.342Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/3d/ff/7f72e8170be527b4977b033239a83a68d5c881cc4775fca255c677f7ac5d/uvloop-0.22.1-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:fe94b4564e865d968414598eea1a6de60adba0c040ba4ed05ac1300de402cd42", size = 1359936, upload-time = "2025-10-16T22:16:29.436Z" }, @@ -6963,7 +6980,7 @@ wheels = [ [[package]] name = "virtualenv" version = "21.5.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "distlib" }, { name = "filelock" }, @@ -6978,7 +6995,7 @@ wheels = [ [[package]] name = "wandb" version = "0.27.2" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "click" }, { name = "gitpython" }, @@ -7007,7 +7024,7 @@ wheels = [ [[package]] name = "watchfiles" version = "1.2.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "anyio" }, ] @@ -7093,7 +7110,7 @@ wheels = [ [[package]] name = "wcwidth" version = "0.8.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/49/b4/51fe890511f0f242d07cb1ebe6a5b6db417262b9d2568b460347c57d95cc/wcwidth-0.8.1.tar.gz", hash = "sha256:faf5b4a5366a72dc49cad48cdf21f52bdf63bdda995178e483ba247ff79089b9", size = 1466072, upload-time = "2026-06-08T05:57:23.146Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/bd/6e/95b0e537de1f4d4301f76f944642c6da50d1511cc7b3d64dc418a66c7509/wcwidth-0.8.1-py3-none-any.whl", hash = "sha256:f453740b1e4a4f3291faa37944c555d71056c4da08d59809b307ef4feba695c8", size = 323092, upload-time = "2026-06-08T05:57:21.413Z" }, @@ -7102,7 +7119,7 @@ wheels = [ [[package]] name = "webcolors" version = "25.10.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/1d/7a/eb316761ec35664ea5174709a68bbd3389de60d4a1ebab8808bfc264ed67/webcolors-25.10.0.tar.gz", hash = "sha256:62abae86504f66d0f6364c2a8520de4a0c47b80c03fc3a5f1815fedbef7c19bf", size = 53491, upload-time = "2025-10-31T07:51:03.977Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/e2/cc/e097523dd85c9cf5d354f78310927f1656c422bd7b2613b2db3e3f9a0f2c/webcolors-25.10.0-py3-none-any.whl", hash = "sha256:032c727334856fc0b968f63daa252a1ac93d33db2f5267756623c210e57a4f1d", size = 14905, upload-time = "2025-10-31T07:51:01.778Z" }, @@ -7111,7 +7128,7 @@ wheels = [ [[package]] name = "webencodings" version = "0.5.1" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/0b/02/ae6ceac1baeda530866a85075641cec12989bd8d31af6d5ab4a3e8c92f47/webencodings-0.5.1.tar.gz", hash = "sha256:b36a1c245f2d304965eb4e0a82848379241dc04b865afcc4aab16748587e1923", size = 9721, upload-time = "2017-04-05T20:21:34.189Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/f4/24/2a3e3df732393fed8b3ebf2ec078f05546de641fe1b667ee316ec1dcf3b7/webencodings-0.5.1-py2.py3-none-any.whl", hash = "sha256:a0af1213f3c2226497a97e2b3aa01a7e4bee4f403f95be16fc9acd2947514a78", size = 11774, upload-time = "2017-04-05T20:21:32.581Z" }, @@ -7120,7 +7137,7 @@ wheels = [ [[package]] name = "websocket-client" version = "1.9.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/2c/41/aa4bf9664e4cda14c3b39865b12251e8e7d239f4cd0e3cc1b6c2ccde25c1/websocket_client-1.9.0.tar.gz", hash = "sha256:9e813624b6eb619999a97dc7958469217c3176312b3a16a4bd1bc7e08a46ec98", size = 70576, upload-time = "2025-10-07T21:16:36.495Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/34/db/b10e48aa8fff7407e67470363eac595018441cf32d5e1001567a7aeba5d2/websocket_client-1.9.0-py3-none-any.whl", hash = "sha256:af248a825037ef591efbf6ed20cc5faa03d3b47b9e5a2230a529eeee1c1fc3ef", size = 82616, upload-time = "2025-10-07T21:16:34.951Z" }, @@ -7129,7 +7146,7 @@ wheels = [ [[package]] name = "websockets" version = "16.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/04/24/4b2031d72e840ce4c1ccb255f693b15c334757fc50023e4db9537080b8c4/websockets-16.0.tar.gz", hash = "sha256:5f6261a5e56e8d5c42a4497b364ea24d94d9563e8fbd44e78ac40879c60179b5", size = 179346, upload-time = "2026-01-10T09:23:47.181Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/84/7b/bac442e6b96c9d25092695578dda82403c77936104b5682307bd4deb1ad4/websockets-16.0-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:71c989cbf3254fbd5e84d3bff31e4da39c43f884e64f2551d14bb3c186230f00", size = 177365, upload-time = "2026-01-10T09:22:46.787Z" }, @@ -7174,7 +7191,7 @@ wheels = [ [[package]] name = "werkzeug" version = "3.1.8" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "markupsafe", marker = "sys_platform == 'linux'" }, ] @@ -7186,7 +7203,7 @@ wheels = [ [[package]] name = "widgetsnbextension" version = "4.0.15" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/bd/f4/c67440c7fb409a71b7404b7aefcd7569a9c0d6bd071299bf4198ae7a5d95/widgetsnbextension-4.0.15.tar.gz", hash = "sha256:de8610639996f1567952d763a5a41af8af37f2575a41f9852a38f947eb82a3b9", size = 1097402, upload-time = "2025-11-01T21:15:55.178Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/3f/0e/fa3b193432cfc60c93b42f3be03365f5f909d2b3ea410295cf36df739e31/widgetsnbextension-4.0.15-py3-none-any.whl", hash = "sha256:8156704e4346a571d9ce73b84bee86a29906c9abfd7223b7228a28899ccf3366", size = 2196503, upload-time = "2025-11-01T21:15:53.565Z" }, @@ -7195,7 +7212,7 @@ wheels = [ [[package]] name = "wrapt" version = "1.17.3" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/95/8f/aeb76c5b46e273670962298c23e7ddde79916cb74db802131d49a85e4b7d/wrapt-1.17.3.tar.gz", hash = "sha256:f66eb08feaa410fe4eebd17f2a2c8e2e46d3476e9f8c783daa8e09e0faa666d0", size = 55547, upload-time = "2025-08-12T05:53:21.714Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/9f/41/cad1aba93e752f1f9268c77270da3c469883d56e2798e7df6240dcb2287b/wrapt-1.17.3-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:ab232e7fdb44cdfbf55fc3afa31bcdb0d8980b9b95c38b6405df2acb672af0e0", size = 53998, upload-time = "2025-08-12T05:51:47.138Z" }, @@ -7244,7 +7261,7 @@ wheels = [ [[package]] name = "xxhash" version = "3.7.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/24/2f/e183a1b407002f5af81822bee18b61cdb94b8670208ef34734d8d2b8ebe9/xxhash-3.7.0.tar.gz", hash = "sha256:6cc4eefbb542a5d6ffd6d70ea9c502957c925e800f998c5630ecc809d6702bae", size = 82022, upload-time = "2026-04-25T11:10:32.553Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/f2/8a/51a14cdef4728c6c2337db8a7d8704422cc65676d9199d77215464c880af/xxhash-3.7.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:082c87bfdd2b9f457606c7a4a53457f4c4b48b0cdc48de0277f4349d79bb3d7a", size = 33357, upload-time = "2026-04-25T11:06:20.44Z" }, @@ -7357,7 +7374,7 @@ wheels = [ [[package]] name = "yarl" version = "1.24.2" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "idna" }, { name = "multidict" }, @@ -7439,7 +7456,7 @@ wheels = [ [[package]] name = "zipp" version = "4.1.0" -source = { registry = "https://pypi.registries.huggingface.tech/" } +source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/b9/d8/eab98a517c14134c0b2eb4e2387bc5f457334293ec5d2dd3857ec2966802/zipp-4.1.0.tar.gz", hash = "sha256:4cb57381f544315db7688e976e922a2b18cdb513d21cc194eb42232ba2a3e602", size = 26214, upload-time = "2026-05-18T20:08:57.967Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/3a/13/547360d81e6d88d58492968ffda9f9542854f11310ee556fef14260cc886/zipp-4.1.0-py3-none-any.whl", hash = "sha256:25ad4e16390cd314347dd8f1de67a2ac538ae658ed4ab9db16029c07c188e97f", size = 10238, upload-time = "2026-05-18T20:08:57.045Z" },