Validate RTX 3080 LeRobot environment
This commit is contained in:
ylhan
@@ -1,7 +1,7 @@
|
||||
# LeRobot 双臂移动操作项目方案
|
||||
|
||||
> 决策与阶段验收以本文件为准。更新：2026-09-08，设计评审修订。
|
||||
> 当前 P0：已有规划与环境脚本，尚无 GPU 环境、训练效果或真机验收结果。
|
||||
> 当前 P0：3080 基础环境与合成数据 ACT 管线通过，见 [实测记录](docs/p0-3080-20260908.md)；3090 和真机未验收。
|
||||
|
||||
## 1. 目标与边界
|
||||
|
||||
@@ -59,9 +59,9 @@ P1 起验证停止手段、限位、限速与人工接管。P3 前定义最大
|
||||
### P0：环境与采购确认
|
||||
|
||||
- [ ] 两台补齐 OS、驱动、GPU、Python、uv 信息（[机器记录](env/machines.md)）。
|
||||
- [ ] 3080 用 `env/setup_lerobot.sh --resolve` 生成依赖锁并通过软件冒烟检查，保留报告。
|
||||
- [x] 3080 用 `env/setup_lerobot.sh --resolve` 生成依赖锁并通过软件冒烟检查，保留报告。
|
||||
- [ ] 将锁文件和输入指纹同步到 3090，以默认模式安装，核对两台锁的 SHA256。
|
||||
- [ ] 完成短视频编码/解码、小样本训练和模型加载；解决 CUDA/FFmpeg/TorchCodec 兼容性后再认可该锁。
|
||||
- [x] 完成短视频编码/解码、小样本训练和模型加载；解决 CUDA/FFmpeg/TorchCodec 兼容性后再认可该锁。
|
||||
- [ ] 验证过的锁、指纹与精简环境记录入 git，不提交凭据或原始日志。
|
||||
- [ ] 确认实际套件的 leader/follower 电机、供电规格，先买一对臂、相机和夹具。
|
||||
|
||||
@@ -119,7 +119,7 @@ P1 起验证停止手段、限位、限速与人工接管。P3 前定义最大
|
||||
|
||||
## 6. 参考与下一步
|
||||
|
||||
下一步：3080 完成 P0 首次解析和报告，再同步 3090；同时确认 P1 型号。未经测试的项保持未勾选。
|
||||
下一步：将已验证的锁同步至 3090 并复验；机械臂到货后完成 P1 装配、校准及少量真实示范检查。未经测试的项保持未勾选。
|
||||
|
||||
- [v0.6.1 依赖](https://github.com/huggingface/lerobot/blob/v0.6.1/pyproject.toml)
|
||||
- [v0.6.1 双臂源码](https://github.com/huggingface/lerobot/blob/v0.6.1/src/lerobot/robots/bi_so_follower/bi_so_follower.py)
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
|
||||
基于 LeRobot，建立人工遥操示范、采集、训练与真机执行闭环。先完成单臂抓放，再扩展桌面双臂，最后实现“移动到位 → 停稳 → 操作”的分阶段任务。
|
||||
|
||||
当前处于 P0：规划与环境脚本已准备，GPU、数据管线和机器人效果尚未验收。
|
||||
当前处于 P0：3080 基础环境、视频读写、合成数据 ACT 短训练及模型重载已通过；3090 与真机尚未验收。详见 [本机验收记录](docs/p0-3080-20260908.md) 和 [到货前复验步骤](docs/before_hardware.md)。
|
||||
|
||||
## 导览
|
||||
|
||||
@@ -33,7 +33,7 @@
|
||||
```bash
|
||||
mkdir -p ~/robot
|
||||
git clone https://git.hyonglin.top/ylhan/lerobot_learn.git ~/robot/proj
|
||||
bash ~/robot/proj/env/setup_lerobot.sh --resolve
|
||||
bash ~/robot/proj/env/setup_lerobot.sh
|
||||
```
|
||||
|
||||
第一次在目标主机解析完整依赖并生成锁；通过 P0 验证后，把锁和输入指纹同步至 3090，再执行不带 `--resolve` 的脚本。当前未附经过 GPU 验证的完整锁，不能声称两机已可复现。脚本不会自动拉取上游最新代码，也不会操作机器人。
|
||||
仓库已附本机验证的依赖锁及输入指纹，默认按锁安装；不要再次使用 `--resolve` 覆盖锁。3090 仍需同锁独立验收。已有本仓库时无需重复 clone，可直接在项目目录运行 `bash env/setup_lerobot.sh`。脚本不会自动拉取上游最新代码，也不会操作机器人。
|
||||
|
||||
@@ -0,0 +1,40 @@
|
||||
# 机械臂到货前的软件准备
|
||||
|
||||
本机项目位于 `/home/ubuntu/lerobot_learn`，运行环境位于 `~/robot/.venv-lerobot-0.6.1`。
|
||||
|
||||
## 复验命令
|
||||
|
||||
```bash
|
||||
cd /home/ubuntu/lerobot_learn
|
||||
export PATH="$HOME/.local/bin:$PATH"
|
||||
bash env/setup_lerobot.sh
|
||||
source "$HOME/robot/.venv-lerobot-0.6.1/bin/activate"
|
||||
python env/smoke_pipeline.py --output "$HOME/robot/reports/pipeline-$(date +%Y%m%dT%H%M%S)"
|
||||
```
|
||||
|
||||
每次使用新输出目录；脚本拒绝覆盖。安装检查使用已有锁，不重新解析。原始安装日志位于 `~/robot/reports/`。
|
||||
|
||||
管线检查生成 3 段共 72 帧合成数据（6 维状态/动作，单相机 128×96），编码 H.264，分别通过 PyAV 和 TorchCodec 读取，再用 batch=2、chunk=8 跑 10 步 ACT。关闭预训练视觉权重和 Hub 联网，保存 checkpoint 与处理器，再重载并检查有限值动作。输出训练日志、训练峰值显存和推理延迟报告。
|
||||
|
||||
这验证的是软件通路；合成数据模型不能用于机械臂。小图像、短 chunk 的显存不能作为正式 640×480 采集训练的预算，真实配置需再次测量。SmolVLA 当前只做模块导入检查，预训练权重下载、微调和推理另行验证。
|
||||
|
||||
也可指定接近正式 ACT 的参数：
|
||||
|
||||
```bash
|
||||
python env/smoke_pipeline.py \
|
||||
--output "$HOME/robot/reports/pipeline-640x480-$(date +%Y%m%dT%H%M%S)" \
|
||||
--height 480 --width 640 --episodes 2 --frames-per-episode 120 \
|
||||
--batch-size 8 --steps 10 --chunk-size 100 --pretrained-backbone
|
||||
```
|
||||
|
||||
该命令仍使用合成画面，但会验证 640×480 输入、ACT 默认动作块和预训练 ResNet18 的下载与加载。
|
||||
|
||||
## 到货后顺序
|
||||
|
||||
1. 按实际套件确认 leader/follower 电机和供电，组装固定并确认停止手段。
|
||||
2. 在本机发现端口和相机，回填 `docs/cli_reference.md` 中的设备参数。
|
||||
3. 按官方装配流程逐个设置新电机，校准两条臂并备份校准文件。
|
||||
4. 验证遥操与相机视野，先录 2–3 段真实示范；检查回放、状态动作和时间对齐。
|
||||
5. 正式采约 50 条同任务示范，再同步至 3090 训练 ACT，回传完整模型与处理器。
|
||||
|
||||
双臂、底盘与远程推理在后续阶段开展。3090 仍需使用相同锁独立验证，不能凭 3080 的结果勾选两机验收。
|
||||
@@ -0,0 +1,44 @@
|
||||
# 3080 软件验收记录 · 2026-09-08
|
||||
|
||||
状态：本机基础环境及离线合成数据管线通过。机械臂未到货，3090 未验收，整个 P0 尚未完成。
|
||||
|
||||
## 环境
|
||||
|
||||
- Ubuntu 24.04.3 x86_64；RTX 3080 10240 MiB；驱动 550.163.01。
|
||||
- Python 3.12.3；uv 0.12.10；FFmpeg 6.1.1-3ubuntu5；已安装 build-essential/python3-dev。
|
||||
- LeRobot 0.6.1；PyTorch 2.7.1+cu126；TorchVision 0.22.1；TorchCodec 0.4.0；PyAV 15.1.0。
|
||||
- 112 个依赖包；`uv pip check` 通过；按哈希同步安装通过。
|
||||
- 锁 SHA256：`d3876e8ff74ba8356618bf6f6059fc32be33e3ed6301a50d164b668884405e18`。
|
||||
|
||||
原始输入会选择 PyTorch 2.11 / CUDA 13，因此约束为 CUDA 12 的版本组合。配对依据：[TorchCodec 官方兼容表](https://github.com/pytorch/torchcodec/tree/v0.5.0#installing-torchcodec)、[NVIDIA CUDA minor compatibility](https://docs.nvidia.com/deploy/cuda-compatibility/minor-version-compatibility.html)；本机 CUDA 运算及 ACT 训练已实际通过。未升级系统驱动。
|
||||
|
||||
## 通过的检查
|
||||
|
||||
- CUDA 矩阵运算；ACT、SmolVLA、SO follower/leader 和视频库导入。
|
||||
- 电机设置、校准、遥操、录制、训练、rollout CLI 帮助加载；端口/相机发现 CLI 存在，未连接硬件。
|
||||
- 3 段共 72 帧合成 LeRobot 数据写入、H.264 编码、PyAV/TorchCodec 读取及跨 episode 边界取样。
|
||||
- ACT batch=2、chunk=8、128×96 单相机、6 维状态/动作、seed=42、无预训练视觉权重，10 步 CUDA 训练完成。
|
||||
- 保存及重载完整模型/处理器；输出形状为 (1, 6)，无 NaN/Inf。
|
||||
|
||||
## 测量
|
||||
|
||||
- 训练子进程含初始化和保存：4.55 秒；逐步日志峰值 allocated 显存约 1.38 GiB（上游日志精度 0.01 GiB）。
|
||||
- 重载推理：预热 5 次后测 20 次，每次清空动作缓存并重新生成 chunk；包括预处理和后处理，不包括图像解码/硬件采集。
|
||||
- 推理 P50/P95/P99：2.19 / 3.48 / 3.56 ms；推理峰值 allocated 显存 214.49 MiB。
|
||||
- 最后一步显存另列于 JSON；上游逐步重置峰值计数，不能拿末步值替代全程峰值。
|
||||
|
||||
2026-09-09 又完成一轮更接近正式配置的单相机基准：640×480、batch=8、chunk=100、预训练 ResNet18、10 步训练。训练子进程含初始化和保存共 8.26 秒，逐步日志峰值 allocated 显存约 2.17 GiB；重载推理 P50/P95/P99 为 5.17 / 5.19 / 5.22 ms，推理峰值 allocated 显存 265.21 MiB。结果见 [640×480 基准 JSON](p0-3080-benchmark-640x480.json)。
|
||||
|
||||
上述数字只适用于合成小图像配置，不能推断真实任务成功率、640×480 多相机显存或端到端控制延迟。SmolVLA 只通过导入，尚未下载基座或微调。
|
||||
|
||||
## 本机证据路径
|
||||
|
||||
- 基础环境：`/home/ubuntu/robot/reports/20260908T090146Z-51776/`
|
||||
- 最终管线：`/home/ubuntu/robot/reports/pipeline-20260908-final/`
|
||||
- 管线内含 `result.json`、`train.log`、合成数据及 `training/checkpoints/000010/pretrained_model/`。
|
||||
- 可移植精简结果：[p0-3080-result.json](p0-3080-result.json)。复验命令：[到货前准备](before_hardware.md)。
|
||||
- 640×480 基准：`/home/ubuntu/robot/reports/pipeline-640x480-b8/`。
|
||||
|
||||
## 剩余
|
||||
|
||||
3090 同锁安装及验收；真实相机和示范数据检查；机械臂组装、校准、遥操和停止验证；真实 ACT 训练及 20 次评测。合成 checkpoint 不用于机械臂。
|
||||
@@ -0,0 +1,31 @@
|
||||
{
|
||||
"status": "PASS",
|
||||
"scope": "synthetic software pipeline only; not robot/task validation",
|
||||
"frames": 240,
|
||||
"episodes": 2,
|
||||
"train_steps": 10,
|
||||
"batch_size": 8,
|
||||
"chunk_size": 100,
|
||||
"image_shape": [
|
||||
480,
|
||||
640,
|
||||
3
|
||||
],
|
||||
"pretrained_backbone": true,
|
||||
"training_wall_seconds": 8.26456517999759,
|
||||
"training_peak_allocated_gib_from_step_logs": 2.17,
|
||||
"training_last_step_memory": {
|
||||
"allocated_mib": 2144.20751953125,
|
||||
"reserved_mib": 2588.0
|
||||
},
|
||||
"checkpoint": "/home/ubuntu/robot/reports/pipeline-640x480-b8/training/checkpoints/000010/pretrained_model",
|
||||
"inference_ms_p50_p95_p99": [
|
||||
5.1704970028367825,
|
||||
5.194654352089856,
|
||||
5.221867668733466
|
||||
],
|
||||
"inference_peak_allocated_mib": 265.20556640625,
|
||||
"torch": "2.7.1+cu126",
|
||||
"cuda": "12.6",
|
||||
"gpu": "NVIDIA GeForce RTX 3080"
|
||||
}
|
||||
@@ -0,0 +1,30 @@
|
||||
{
|
||||
"status": "PASS",
|
||||
"scope": "synthetic software pipeline only; not robot/task validation",
|
||||
"frames": 72,
|
||||
"episodes": 3,
|
||||
"train_steps": 10,
|
||||
"batch_size": 2,
|
||||
"image_shape": [
|
||||
96,
|
||||
128,
|
||||
3
|
||||
],
|
||||
"pretrained_backbone": false,
|
||||
"training_wall_seconds": 4.55074907900007,
|
||||
"training_peak_allocated_gib_from_step_logs": 1.38,
|
||||
"training_last_step_memory": {
|
||||
"allocated_mib": 959.85107421875,
|
||||
"reserved_mib": 1040.0
|
||||
},
|
||||
"checkpoint": "/home/ubuntu/robot/reports/pipeline-20260908-final/training/checkpoints/000010/pretrained_model",
|
||||
"inference_ms_p50_p95_p99": [
|
||||
2.192092500081344,
|
||||
3.4814052502724735,
|
||||
3.5563298500483143
|
||||
],
|
||||
"inference_peak_allocated_mib": 214.486328125,
|
||||
"torch": "2.7.1+cu126",
|
||||
"cuda": "12.6",
|
||||
"gpu": "NVIDIA GeForce RTX 3080"
|
||||
}
|
||||
Vendored
+6
-6
@@ -1,13 +1,13 @@
|
||||
# GPU 环境与机器记录
|
||||
|
||||
基线：LeRobot 0.6.1 / Python 3.12。脚本适用 Ubuntu x86_64 GPU 主机；不用于 Windows 或 RPi。当前 GPU 环境与模型训练均未验收。
|
||||
基线：LeRobot 0.6.1 / Python 3.12。脚本适用 Ubuntu x86_64 GPU 主机；不用于 Windows 或 RPi。3080 的 CUDA、视频和合成 ACT 短训练已验收，详见 [记录](../docs/p0-3080-20260908.md)；3090 与真实任务未验收。
|
||||
|
||||
## 分工与路径
|
||||
|
||||
| 机器 | 用途 | 环境记录 |
|
||||
|---|---|---|
|
||||
| Windows | 编辑、静态检查、git 同步 | 本仓库 `D:\robot\lerobot` |
|
||||
| 3080 10GB | USB 遥操、采集、推理，首次解析依赖 | OS/驱动/Python/uv/锁 SHA256：待填 |
|
||||
| 3080 10GB | USB 遥操、采集、推理，首次解析依赖 | Ubuntu 24.04.3 / 550.163.01 / Python 3.12.3 / uv 0.12.10；锁见实测记录 |
|
||||
| 3090 24GB | 训练，同局域网 | OS/驱动/Python/uv/锁 SHA256：待填 |
|
||||
| RPi 5 | P3 起硬件驱动、采图、执行和故障处理 | ARM64 环境另行验证，不复用 x86_64 锁 |
|
||||
| DRIVE Orin | 可选实验 | 不阻塞主线 |
|
||||
@@ -16,9 +16,9 @@ GPU 机统一路径：项目 `~/robot/proj`；专用环境 `~/robot/.venv-lerobo
|
||||
|
||||
## 首次安装与依赖锁
|
||||
|
||||
前置：安装 NVIDIA 驱动、git、uv、FFmpeg。脚本用 uv 创建 Python 3.12 环境（必要时 uv 下载 Python），不自动 sudo 安装系统组件。缺少前置即退出。
|
||||
前置：安装 NVIDIA 驱动、git、uv、FFmpeg，以及构建遥操依赖 evdev 所需的 `build-essential`、`python3-dev`。Ubuntu 可执行 `sudo apt-get install -y ffmpeg build-essential python3-dev`。脚本用 uv 创建 Python 3.12 环境（必要时 uv 下载 Python），不自动 sudo 安装系统组件。缺少前置即退出。
|
||||
|
||||
先在 3080：
|
||||
首次解析流程（本机已完成；新机器有仓库附带锁时直接使用下方默认安装命令）：
|
||||
|
||||
```bash
|
||||
mkdir -p ~/robot
|
||||
@@ -35,7 +35,7 @@ bash ~/robot/proj/env/setup_lerobot.sh
|
||||
source ~/robot/.venv-lerobot-0.6.1/bin/activate
|
||||
```
|
||||
|
||||
默认安装严格使用锁与包哈希，不重新解析。两台比较报告中的锁 SHA256；验收通过后把锁、输入指纹与精简环境信息提交 git。**当前仓库尚未包含经过 Ubuntu/GPU 验证的传递依赖锁，默认安装会明确拒绝继续。**
|
||||
默认安装严格使用锁与包哈希，不重新解析。两台比较报告中的锁 SHA256；验收通过后把锁、输入指纹与精简环境信息提交 git。**当前仓库已包含 3080 验证的完整依赖锁和指纹，默认安装可直接使用。不要对现有锁执行 `--resolve`。**
|
||||
|
||||
两台尽量统一 Ubuntu 版本，记录 Python 补丁版本、uv 版本与 NVIDIA 驱动；同一 Python 包锁不能锁定系统库或驱动。脚本 `uv pip sync` 仅作用于新命名的项目专用环境，会使其中包与锁一致，请勿把它用于其他项目。
|
||||
|
||||
@@ -53,7 +53,7 @@ rsync -avP ~/robot/data/ USER@TRAIN_HOST:~/robot/data/
|
||||
rsync -avP USER@TRAIN_HOST:~/robot/outputs/ ~/robot/outputs/
|
||||
```
|
||||
|
||||
每个实验保留数据版本/校验值、模型与处理器、配置、校准备份和评测记录。模型快照记录来源 revision，不自动随 Hub main 更新。同步代码使用 git；当前未配置 GPU 主机远程访问，因此 GPU 执行结果需从目标机器回传。
|
||||
每个实验保留数据版本/校验值、模型与处理器、配置、校准备份和评测记录。模型快照记录来源 revision，不自动随 Hub main 更新。同步代码使用 git；3080 为当前本机，3090 的远程访问尚未验证。
|
||||
|
||||
## 验收报告
|
||||
|
||||
|
||||
Vendored
+2456
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1 @@
|
||||
e236f8a3986eecd6bafb490108f52dd8adc1c9f76fd85d30c49a67ea00e3fd80 requirements.in
|
||||
Vendored
+5
@@ -2,3 +2,8 @@
|
||||
# Resolve once on 3080; share the lock and its input fingerprint with 3090.
|
||||
# ACT uses the base dependencies; v0.6.1 has no 'act' extra.
|
||||
lerobot[core_scripts,training,feetech,smolvla]==0.6.1
|
||||
# Keep CUDA 12 wheels for the 3080 host's 550 driver; unconstrained resolution
|
||||
# selects torch 2.11 / CUDA 13. TorchCodec must match PyTorch's ABI.
|
||||
torch==2.7.1
|
||||
torchvision==0.22.1
|
||||
torchcodec==0.4.0
|
||||
|
||||
Vendored
+8
-1
@@ -18,7 +18,7 @@ esac
|
||||
[[ "$(uname -s)" == Linux && "$(uname -m)" == x86_64 ]] || {
|
||||
echo "Only Linux x86_64 GPU hosts are supported; RPi needs a separate environment." >&2; exit 1;
|
||||
}
|
||||
for cmd in uv nvidia-smi ffmpeg sha256sum; do
|
||||
for cmd in uv nvidia-smi ffmpeg sha256sum gcc; do
|
||||
command -v "$cmd" >/dev/null || { echo "Missing prerequisite: $cmd" >&2; exit 1; }
|
||||
done
|
||||
if ! $RESOLVE && [[ ! -s "$LOCK" ]]; then
|
||||
@@ -34,12 +34,19 @@ exec > >(tee "$REPORT/setup.log") 2>&1
|
||||
echo "[1/5] Prerequisites"
|
||||
nvidia-smi --query-gpu=name,memory.total,driver_version --format=csv,noheader
|
||||
uv --version
|
||||
cat /etc/os-release > "$REPORT/os-release.txt"
|
||||
uname -srmo > "$REPORT/kernel.txt"
|
||||
nvidia-smi > "$REPORT/nvidia-smi.txt"
|
||||
ffmpeg -version > "$REPORT/ffmpeg.txt"
|
||||
# Set UV_DEFAULT_INDEX explicitly if a mirror is needed. No automatic Hub/token routing.
|
||||
echo "[2/5] Dedicated Python 3.12 environment"
|
||||
if [[ ! -d "$VENV" ]]; then uv venv --python 3.12 "$VENV"; fi
|
||||
PYTHON="${VENV}/bin/python"
|
||||
"$PYTHON" -c 'import sys; assert sys.version_info[:2] == (3, 12), sys.version'
|
||||
PYTHON_INCLUDE="$("$PYTHON" -c 'import sysconfig; print(sysconfig.get_path("include"))')"
|
||||
[[ -f "$PYTHON_INCLUDE/Python.h" ]] || {
|
||||
echo "Missing Python development headers: $PYTHON_INCLUDE/Python.h" >&2; exit 1;
|
||||
}
|
||||
if $RESOLVE; then
|
||||
TEMP_LOCK="$(mktemp "${LOCK}.tmp.XXXXXX")"
|
||||
trap 'rm -f -- "$TEMP_LOCK"' EXIT
|
||||
|
||||
Vendored
+145
@@ -0,0 +1,145 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Offline synthetic-data acceptance. No robot access; no task-quality claim."""
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
from pathlib import Path
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--output", type=Path, required=True)
|
||||
parser.add_argument("--height", type=int, default=96)
|
||||
parser.add_argument("--width", type=int, default=128)
|
||||
parser.add_argument("--episodes", type=int, default=3)
|
||||
parser.add_argument("--frames-per-episode", type=int, default=24)
|
||||
parser.add_argument("--batch-size", type=int, default=2)
|
||||
parser.add_argument("--steps", type=int, default=10)
|
||||
parser.add_argument("--chunk-size", type=int, default=8)
|
||||
parser.add_argument("--pretrained-backbone", action="store_true")
|
||||
args = parser.parse_args()
|
||||
root = args.output.resolve()
|
||||
root.mkdir(parents=True, exist_ok=False)
|
||||
os.environ["HF_HUB_OFFLINE"] = "1"
|
||||
os.environ["WANDB_MODE"] = "disabled"
|
||||
import numpy as np
|
||||
import torch
|
||||
from lerobot.datasets.lerobot_dataset import LeRobotDataset
|
||||
from lerobot.configs.video import RGBEncoderConfig
|
||||
|
||||
torch.set_num_threads(4)
|
||||
assert torch.cuda.is_available()
|
||||
features = {
|
||||
"observation.state": {"dtype": "float32", "shape": (6,), "names": [f"joint_{i}" for i in range(6)]},
|
||||
"action": {"dtype": "float32", "shape": (6,), "names": [f"joint_{i}" for i in range(6)]},
|
||||
"observation.images.front": {"dtype": "video", "shape": (args.height, args.width, 3), "names": ["height", "width", "channels"]},
|
||||
}
|
||||
dataset = LeRobotDataset.create(
|
||||
repo_id="local/p0_synthetic", root=root / "dataset", fps=10,
|
||||
features=features, robot_type="synthetic", use_videos=True,
|
||||
rgb_encoder=RGBEncoderConfig(vcodec="h264"), encoder_threads=2,
|
||||
)
|
||||
for episode in range(args.episodes):
|
||||
for frame in range(args.frames_per_episode):
|
||||
phase = frame / 10 + episode
|
||||
state = np.sin(phase + np.arange(6)).astype(np.float32)
|
||||
image = np.zeros((args.height, args.width, 3), dtype=np.uint8)
|
||||
image[:, :, 1] = 40 + episode * 50
|
||||
box = max(8, min(args.height, args.width) // 6)
|
||||
x = frame * max(1, (args.width - box) // max(1, args.frames_per_episode - 1))
|
||||
y = (args.height - box) // 2
|
||||
image[y:y+box, x:x+box, 0] = 220
|
||||
dataset.add_frame({"observation.state": state,
|
||||
"action": np.sin(phase + 0.1 + np.arange(6)).astype(np.float32),
|
||||
"observation.images.front": image, "task": "synthetic pipeline check"})
|
||||
dataset.save_episode()
|
||||
dataset.finalize()
|
||||
from torchcodec.decoders import VideoDecoder
|
||||
videos = list((root / "dataset").rglob("*.mp4"))
|
||||
assert videos
|
||||
for video in videos:
|
||||
decoded = VideoDecoder(str(video))
|
||||
assert len(decoded) > 0 and decoded[0].shape == (3, args.height, args.width)
|
||||
frame_count = args.episodes * args.frames_per_episode
|
||||
for backend in ("pyav", "torchcodec"):
|
||||
loaded = LeRobotDataset("local/p0_synthetic", root=root / "dataset", video_backend=backend)
|
||||
assert len(loaded) == frame_count
|
||||
boundary_indices = sorted({0, frame_count - 1} | {
|
||||
i for episode in range(1, args.episodes)
|
||||
for i in (episode * args.frames_per_episode - 1, episode * args.frames_per_episode)
|
||||
})
|
||||
for index in boundary_indices:
|
||||
item = loaded[index]
|
||||
assert item["observation.images.front"].shape == (3, args.height, args.width)
|
||||
assert torch.isfinite(item["action"]).all()
|
||||
|
||||
train_wrapper = (
|
||||
"import json,torch; from pathlib import Path; "
|
||||
"from lerobot.scripts.lerobot_train import main; "
|
||||
"torch.cuda.reset_peak_memory_stats(); main(); "
|
||||
f"Path({str(root / 'training_memory.json')!r}).write_text(json.dumps("
|
||||
"{'allocated_mib':torch.cuda.max_memory_allocated()/2**20,"
|
||||
"'reserved_mib':torch.cuda.max_memory_reserved()/2**20}))"
|
||||
)
|
||||
command = [sys.executable, "-c", train_wrapper,
|
||||
"--policy.type=act", "--policy.device=cuda", "--policy.push_to_hub=false",
|
||||
f"--policy.chunk_size={args.chunk_size}", f"--policy.n_action_steps={args.chunk_size}",
|
||||
"--dataset.repo_id=local/p0_synthetic",
|
||||
f"--dataset.root={root / 'dataset'}", "--dataset.video_backend=torchcodec",
|
||||
f"--batch_size={args.batch_size}", f"--steps={args.steps}", "--num_workers=0", "--seed=42",
|
||||
"--log_freq=1", "--save_freq=10", "--wandb.enable=false",
|
||||
f"--output_dir={root / 'training'}", "--job_name=p0_synthetic"]
|
||||
if not args.pretrained_backbone:
|
||||
command.append("--policy.pretrained_backbone_weights=null")
|
||||
env = dict(os.environ, OMP_NUM_THREADS="4", MKL_NUM_THREADS="4")
|
||||
start = time.perf_counter()
|
||||
with (root / "train.log").open("w") as log:
|
||||
subprocess.run(command, env=env, stdout=log, stderr=subprocess.STDOUT, check=True)
|
||||
elapsed = time.perf_counter() - start
|
||||
checkpoints = sorted((root / "training" / "checkpoints").glob("[0-9]*/pretrained_model"))
|
||||
assert checkpoints, "No checkpoint saved"
|
||||
from lerobot.policies.act.modeling_act import ACTPolicy
|
||||
from lerobot.policies.factory import make_pre_post_processors
|
||||
checkpoint = checkpoints[-1]
|
||||
policy = ACTPolicy.from_pretrained(checkpoint).to("cuda").eval()
|
||||
pre, post = make_pre_post_processors(policy.config, pretrained_path=checkpoint,
|
||||
preprocessor_overrides={"device_processor": {"device": "cuda"}})
|
||||
observation = {k: v.unsqueeze(0) for k, v in loaded[0].items()
|
||||
if k.startswith("observation.")}
|
||||
timings = []
|
||||
torch.cuda.reset_peak_memory_stats()
|
||||
with torch.inference_mode():
|
||||
for i in range(25):
|
||||
policy.reset()
|
||||
torch.cuda.synchronize()
|
||||
begin = time.perf_counter()
|
||||
action = post(policy.select_action(pre(observation)))
|
||||
torch.cuda.synchronize()
|
||||
assert action.shape == (1, 6) and torch.isfinite(action).all()
|
||||
if i >= 5:
|
||||
timings.append((time.perf_counter() - begin) * 1000)
|
||||
report = {"status": "PASS", "scope": "synthetic software pipeline only; not robot/task validation",
|
||||
"frames": frame_count, "episodes": args.episodes, "train_steps": args.steps,
|
||||
"batch_size": args.batch_size, "chunk_size": args.chunk_size,
|
||||
"image_shape": [args.height, args.width, 3],
|
||||
"pretrained_backbone": args.pretrained_backbone,
|
||||
"training_wall_seconds": elapsed,
|
||||
# Upstream resets CUDA peak counters each step. Its per-step log is rounded to 0.01 GiB.
|
||||
"training_peak_allocated_gib_from_step_logs": max(float(v) for v in re.findall(
|
||||
r"mem_gb:([0-9.]+)", (root / "train.log").read_text())),
|
||||
"training_last_step_memory": json.loads((root / "training_memory.json").read_text()),
|
||||
"checkpoint": str(checkpoint),
|
||||
"inference_ms_p50_p95_p99": np.percentile(timings, [50, 95, 99]).tolist(),
|
||||
"inference_peak_allocated_mib": torch.cuda.max_memory_allocated() / 2**20,
|
||||
"torch": torch.__version__, "cuda": torch.version.cuda,
|
||||
"gpu": torch.cuda.get_device_name(0)}
|
||||
(root / "result.json").write_text(json.dumps(report, indent=2) + "\n")
|
||||
print(json.dumps(report, indent=2))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user