Compare commits

4 Commits

Author SHA1 Message Date
chenlin 456f202685 Merge pull request 'Cen' (#4) from Cen into main
web-platform-ci / TypeScript、Lint、Unit、Build (push) Has been cancelled
web-platform-ci / Playwright E2E (push) Has been cancelled
Reviewed-on: #4
2026-08-28 14:12:31 +08:00
chenlin fe4d98b4a5 Merge pull request 'Cen' (#3) from Cen into main
build / setup (compute matrix) (push) Has been cancelled
build / macos-15-arm64-studio (push) Has been cancelled
build / ubuntu-24.04-clang-18-studio (push) Has been cancelled
build / ubuntu-24.04-gcc-14-studio (push) Has been cancelled
build / windows-2025-ninja-studio (push) Has been cancelled
build / ubuntu-24.04-clang-18-wasm (push) Has been cancelled
build / ubuntu-24.04-clang-18-mjx (push) Has been cancelled
build / ${{ matrix.label }} (push) Has been cancelled
Reviewed-on: #3
2026-08-24 17:24:21 +08:00
chenlin 1263d274dd Merge pull request 'Cen' (#2) from Cen into main
build / setup (compute matrix) (push) Has been cancelled
build / macos-15-arm64-studio (push) Has been cancelled
build / ubuntu-24.04-clang-18-studio (push) Has been cancelled
build / ubuntu-24.04-gcc-14-studio (push) Has been cancelled
build / windows-2025-ninja-studio (push) Has been cancelled
build / ubuntu-24.04-clang-18-wasm (push) Has been cancelled
build / ubuntu-24.04-clang-18-mjx (push) Has been cancelled
build / ${{ matrix.label }} (push) Has been cancelled
Reviewed-on: #2
2026-08-21 16:55:40 +08:00
chenlin 2a50fab918 Merge pull request 'feat: add browser-based MuJoCo simulation platform' (#1) from Cen into main
build / setup (compute matrix) (push) Has been cancelled
build / ${{ matrix.label }} (push) Has been cancelled
build / macos-15-arm64-studio (push) Has been cancelled
build / ubuntu-24.04-clang-18-studio (push) Has been cancelled
build / ubuntu-24.04-gcc-14-studio (push) Has been cancelled
build / windows-2025-ninja-studio (push) Has been cancelled
build / ubuntu-24.04-clang-18-wasm (push) Has been cancelled
build / ubuntu-24.04-clang-18-mjx (push) Has been cancelled
Reviewed-on: #1
2026-08-20 15:22:18 +08:00
297 changed files with 2941 additions and 798932 deletions
-14
View File
@@ -1,14 +0,0 @@
root = true
[*]
charset = utf-8
end_of_line = lf
insert_final_newline = true
indent_style = space
indent_size = 2
[*.py]
indent_size = 4
[*.md]
trim_trailing_whitespace = false
+3 -12
View File
@@ -9,22 +9,15 @@ permissions:
jobs:
quality:
name: TypeScript, lint, unit, build
name: TypeScript、Lint、Unit、Build
runs-on: ubuntu-24.04
steps:
- uses: actions/checkout@v6
- uses: actions/setup-node@v6
with:
node-version-file: .nvmrc
node-version: 24
cache: npm
- uses: actions/setup-python@v6
with:
python-version: '3.12'
cache: pip
- run: npm install --global npm@11.17.0
- run: npm ci
- run: python -m pip install -r requirements-dev.txt
- run: npm run lint:python
- run: npm run check
e2e:
@@ -34,10 +27,8 @@ jobs:
- uses: actions/checkout@v6
- uses: actions/setup-node@v6
with:
node-version-file: .nvmrc
node-version: 24
cache: npm
- run: npm install --global npm@11.17.0
- run: npm ci
- run: npx playwright install --with-deps chromium
- run: npm run build
- run: npm run test:e2e
-46
View File
@@ -1,46 +0,0 @@
name: web-platform-release
on:
push:
tags:
- 'V*'
permissions:
contents: write
jobs:
release:
name: Build and publish release
runs-on: ubuntu-24.04
steps:
- uses: actions/checkout@v6
- uses: actions/setup-node@v6
with:
node-version-file: .nvmrc
cache: npm
- uses: actions/setup-python@v6
with:
python-version: '3.12'
cache: pip
- run: npm install --global npm@11.17.0
- run: npm ci
- run: python -m pip install -r requirements-dev.txt
- run: npm run lint:python
- run: npm run check
- run: npx playwright install --with-deps chromium
- run: npm run test:e2e
- name: Package static site
shell: bash
run: |
archive="mujoco-web-platform-${GITHUB_REF_NAME}.tar.gz"
tar -C web-platform-dist -czf "$archive" .
sha256sum "$archive" > SHA256SUMS
- name: Publish GitHub Release
env:
GH_TOKEN: ${{ github.token }}
run: |
gh release create "$GITHUB_REF_NAME" \
"mujoco-web-platform-${GITHUB_REF_NAME}.tar.gz" \
SHA256SUMS \
--generate-notes \
--title "$GITHUB_REF_NAME"
+1 -7
View File
@@ -7,7 +7,6 @@ playwright-report/
web_platform/test-results/
web_platform/playwright-report/
web_platform/node_modules/.vite/
.playwright-cli/
*.tsbuildinfo
# Python
@@ -16,10 +15,9 @@ __pycache__/
*.py[cod]
*.egg-info/
# Local build, cache, and planning directories
# Local build and cache directories
build/
.cache/
/plans/
# Editors and operating systems
.vscode/
@@ -29,7 +27,3 @@ build/
# Runtime logs
*.log
MUJOCO_LOG.TXT
training_server/rl/logs/
training_server/rl/wandb/
training_server/rl/outputs/
training_server/tuning-data/
-1
View File
@@ -1 +0,0 @@
engine-strict=true
-1
View File
@@ -1 +0,0 @@
24.19.0
-15
View File
@@ -1,15 +0,0 @@
AGENTS.md
context.md
plans/
.git/
.venv/
build/
node_modules/
web-platform-dist/
coverage/
playwright-report/
test-results/
.playwright-cli/
web_platform/fixtures/
web_platform/public/
package-lock.json
-7
View File
@@ -1,7 +0,0 @@
{
"singleQuote": true,
"semi": true,
"trailingComma": "all",
"printWidth": 100,
"proseWrap": "preserve"
}
-3
View File
@@ -1,3 +0,0 @@
Please always speak chinese.
python 虚拟环境路径在/home/cen/Embodied_Workspace/Mujoco_Projects/mujoco/.venv/bin/activate
系统是Ubuntu 24.04 LTS
-64
View File
@@ -1,64 +0,0 @@
# 更新日志
本项目的重要变更记录在此文件中,版本标签沿用仓库现有的 `V主版本.次版本[.修订版本]` 格式。
## [0.7.3] - 2026-09-01
### 新增
- 将默认 `Unitree-Go2-Flat` 训练所需的 PPO 环境、奖励/观测配置、ONNX 导出逻辑和 Go2 模型资产内置到 `training_server/rl/`。
### 变更
- 训练服务默认使用仓库内置训练器,不再要求单独克隆 `unitree_rl_mjlab`;仍可通过 `--trainer-root` 覆盖。
## [0.7.2] - 2026-09-01
### 变更
- 前端改为“应用壳 + 业务域 + 共享组件”的目录结构,并补充架构与依赖边界文档。
- 拆分工程侧栏、模型控制侧栏、执行器控制和通用侧栏容器,降低单文件职责和跨域耦合。
- 将地图、控制器、强化学习、本地训练和遥测面板及测试与各自业务域共置。
- 将遥测模块从仿真调度目录独立到 `telemetry/`,保留通过显式接口与 `SimulationSession` 集成。
## [0.7.1] - 2026-09-01
### 新增
- 增加浏览器内仿真遥测记录模块,可配置记录 Body、采样频率和样本上限。
- 记录位置、移动速度、机身侧倾/俯仰/偏航角、角速度、累计里程、接触与驱动指标。
- 增加实时数据与摘要面板,并支持导出稳定列结构的 CSV 和 Schema V1 JSON。
- 提供 `TelemetrySource`、`registerDataChannel` 及记录生命周期接口,便于扩展业务指标和其他仿真后端。
- 仿真重置使用数据分段,避免跨重置计算错误速度;达到样本上限时自动停止。
## [0.6.1] - 2026-08-28
### 新增
- 固定并强制使用 Node.js、npm、Prettier 和 Ruff 开发工具版本。
- 为核心导入、仿真数学、状态管理和训练客户端增加覆盖率门槛。
- Git 标签触发的自动构建、校验和与 GitHub Release 工作流。
- 本地训练服务随机 Bearer Token 鉴权、Host 校验和任务历史上限。
- 训练进程启动/取消竞态回归测试和 SIGTERM 受控退出。
### 变更
- Playwright CI 改用版本固定的 Chromium。
- TypeScript、TSX、配置和文档统一使用 Prettier 格式化。
- Python 训练服务统一使用 Ruff 检查和格式化。
## [0.6.0] - 2026-08-28
### 变更
- 将仓库重构为以 `web_platform/` 为核心的 Web 应用仓库。
- MuJoCo 运行时改为依赖官方 `@mujoco/mujoco` npm 包。
- 移除原生 C++、Python、MJX、Unity、桌面模拟器、CMake 和上游测试镜像。
- 将训练桥接服务和 Python 控制器示例提升到仓库根目录。
## [0.5.2] - 2026-08-28
### 变更
- 优化响应式工作区、可访问性、首屏加载、纹理兼容性和视口交互。
- 增加碰撞体、坐标系、关节轴、质心和惯量辅助可视化。
+14 -32
View File
@@ -10,22 +10,17 @@
- ROS `package://`、常见 URDF 兼容转换及 DAE 降级处理
- Three.js 模型、碰撞体、坐标系、关节轴、质心和惯量可视化
- 播放、暂停、单步、重置、变速、关节拖动与外力交互
- 统一地图资产库:工程地图、认证资产、内置地图及 9 类系统参数化地形共用点击/拖放、场景树、轻量预览、实例变换、放弃和一次编译事务
- 工程地图包:静态 MJCF/OBJ/STL/高度场碰撞层、GLB 视觉层和机器人出生点;同源实例拥有独立位姿,物理/视觉/出生点同步变换
- V3 地图创作层:认证资产可在程序地形上自动重力落位,共用右侧属性和视口操纵器,支持源内容编辑、空白取消选择及导出地图 ZIP
- 浏览器内 Python 控制器(Pyodide)
- ONNX 强化学习策略推理(ONNX Runtime Web)
- 内置 Go2 PPO 任务的本机 mjlab 训练桥接服务,以及 DeepSeek 驱动的奖励函数自调参、固定评估和 TensorBoard 风格独立工作台
- 可选的本机 mjlab 训练桥接服务
- 响应式工作区、源码编辑、性能监控和中文诊断
## 快速开始
环境要求:Node.js 24(版本见 `.nvmrc`)和 npm 11.17;仅使用训练桥接服务或执行 Python 检查时需要 Python 3.12。
环境要求:Node.js 24+;仅使用训练桥接服务时需要 Python 3。
```bash
nvm use
npm install --global npm@11.17.0
npm ci
npm install
npm run dev
```
@@ -39,16 +34,10 @@ npm run build # 生产构建到 web-platform-dist/
npm run preview # 预览生产构建
npm run typecheck # TypeScript 检查
npm run lint # ESLint
npm run check:format # Prettier 格式检查
npm test # Vitest 单元测试
npm run test:coverage # 核心模块覆盖率检查
npm run test:e2e # Playwright Chromium 浏览器测试
npm run test:e2e # Playwright 浏览器测试
npm run test:training-server # Python 训练桥接服务测试
npm run check # 除 E2E 和 Ruff 外的完整检查
# 修改 training_server/ 时额外执行
python3 -m pip install -r requirements-dev.txt
npm run lint:python
npm run check # 除 E2E 外的完整检查
```
## 仓库结构
@@ -56,33 +45,26 @@ npm run lint:python
```text
.
├── web_platform/ React、TypeScript、Three.js 前端与测试夹具
├── training_server/ 本机训练桥接服务及内置 Go2 PPO 训练代码/资产
├── training_server/ 可选的本机强化学习训练桥接服务
├── controllers/ 可导入平台的 Python 控制器示例
├── package.json 根目录统一开发入口
└── LICENSE Apache License 2.0
```
核心前端采用“应用壳 + 业务域 + 共享组件”分层:
核心前端分层:
- `web_platform/src/app/`:工作区编排、侧栏外壳、工具栏和全局对话框
- `web_platform/src/components/ui/`:无业务依赖的共享 UI 基础组件
- `web_platform/src/project/`:工程导入、路径规范化、项目树和 MEMFS
- `web_platform/src/simulation/`:MuJoCo WASM 会话、仿真调度及执行器控制 UI
- `web_platform/src/app/`:工作区、工具栏、面板和对话框
- `web_platform/src/project/`:工程导入、路径规范化和 MEMFS
- `web_platform/src/simulation/`:MuJoCo WASM 会话与仿真调度
- `web_platform/src/viewer/`:Three.js 渲染、选择和物理交互
- `web_platform/src/map/`:地图加载、创作、校验及对应 UI
- `web_platform/src/controller/`:Pyodide 控制器运行时及控制面板
- `web_platform/src/rl/`:ONNX 策略运行时、任务绑定及策略面板
- `web_platform/src/training/`:本地训练服务客户端、类型及训练面板
- `web_platform/src/telemetry/`:遥测数据源、记录器、导出与数据面板
- `web_platform/src/stores/`:跨业务域的应用状态
- `web_platform/src/controller/`:Pyodide 控制器运行时
- `web_platform/src/rl/`:ONNX 策略运行时
组件和测试优先与所属业务域共置;`app/components/` 仅保留跨域编排和应用壳组件,避免业务逻辑继续集中到单一侧栏文件。
详细使用说明见 [`web_platform/README.md`](web_platform/README.md),前端依赖边界见 [`web_platform/ARCHITECTURE.md`](web_platform/ARCHITECTURE.md),训练服务说明见 [`training_server/README.md`](training_server/README.md)。
详细使用说明见 [`web_platform/README.md`](web_platform/README.md),训练服务说明见 [`training_server/README.md`](training_server/README.md)。
## 数据与安全边界
模型、资源、Python 控制器和 ONNX 策略默认只在当前浏览器会话中处理,不上传到服务器。训练桥接服务只监听本机回环地址,并仅执行服务端允许列表中的任务。启用自调参时,DeepSeek 只接收脱敏后的奖励参数、曲线摘要和评估数值;API key 仅存在训练服务环境中,不进入浏览器、URL、SQLite 或训练日志。
模型、资源、Python 控制器和 ONNX 策略默认只在当前浏览器会话中处理,不上传到服务器。训练桥接服务只监听本机回环地址,并仅执行服务端允许列表中的任务。
## 上游与许可证
-116
View File
@@ -1,116 +0,0 @@
# Code Context
## Files Retrieved
1. `wasm/web_platform/src/simulation/SimulationSession.ts`(lines 1-107)- 仿真步进、控制器生命周期、模型名称解析和 `mj_step` 前控制写入的核心集成点。
2. `wasm/web_platform/src/simulation/PhysicsAdapter.ts`(lines 12-110)- UI 与 session 的稳定门面;模型切换、资源释放和 Go2-W URDF 增强发生于此。
3. `wasm/web_platform/src/app/App.tsx`(lines 37-78, 100-105)- 工程导入、控制器状态、错误反馈、右侧栏 wiring。
4. `wasm/web_platform/src/app/components/SidebarPanel.tsx`(lines 19-35)- 当前“控制”tab 仅嵌入 Python 控制器,是 ONNX 面板的最小 UI 插槽。
5. `wasm/web_platform/src/controller/PythonControllerRuntime.ts`(lines 30-120)- 可复用的控制器生命周期/状态设计;同时揭示现有控制回调必须同步。
6. `wasm/web_platform/src/controller/types.ts`(lines 1-42)- 当前状态被写死为 `language:'python'`,且 step API 可读关节、sensor、body quaternion。
7. `wasm/web_platform/src/project/importer.ts`(lines 1-48, 184-223)- 所有二进制文件已原样进入 manifest,ONNX 无需进入模型 entry 发现逻辑。
8. `wasm/web_platform/src/project/types.ts`(lines 1-41)- ProjectFile/大小上限;单文件允许 128 MiB、工程 512 MiB。
9. `wasm/package.json`(scripts/dependencies)- 已有 Vitest/Playwright,尚无 `onnxruntime-web`。
10. `wasm/web_platform/vite.config.ts`(lines 1-27)- Pyodide 本地资产复制是 ORT wasm 离线资产处理的直接范例。
11. `wasm/src/go2_w_balance.py`(lines 1-约480)- 仓库已有 Go2-W 手写 200 Hz 平衡/轮速控制,可作模型命名、IMU 和安全回退参考,但不是 RL 策略。
12. `/home/cen/Embodied_Workspace/unitree_rl_mjlab/deploy/robots/go2/config/policy/velocity/v0/params/deploy.yaml`(lines 1-56)- 部署时序、动作变换及精确 observation 顺序。
13. `/home/cen/Embodied_Workspace/unitree_rl_mjlab/deploy/robots/go2/src/State_RLBase.cpp`(lines 6-31)- 部署约定为 `params/deploy.yaml` + `exported/policy.onnx`,输出经 action manager 处理后按 joint map 下发。
14. `/home/cen/Embodied_Workspace/unitree_rl_mjlab/src/tasks/velocity/config/go2/env_cfgs.py`(lines 1-约160)- flat 任务移除 height scan,Go2 base 为 `base_link`。
## Key Code
### 已有正确插入顺序
`SimulationSession.singleStep/advance` 均执行:控制器 → 外力 → `mj_step`(`SimulationSession.ts:46-57`)。ONNX 动作必须沿用该位置,而不能在 React frame snapshot 后直接写 ctrl,否则会引入一物理步延迟和不稳定的渲染帧时序。
### 参考部署契约
Go2 deploy YAML 定义:
- `step_dt: 0.02`,即策略 50 Hz;仿真步可保持 0.002 s,动作 hold 10 步。
- observation 串接顺序和维数:`base_ang_vel(3), projected_gravity(3), velocity_commands(3), gait_phase(2), joint_pos_rel(12), joint_vel_rel(12), last_action(12)`,总计 **47 float32**。
- action 为 12 维位置目标:`processed = raw * 0.25 + default_joint_pos`;默认位姿见 YAML lines 3-5/18-20。
- joint 顺序不能取 MuJoCo ID 偶然顺序;部署 `joint_ids_map=[3,4,5,0,1,2,9,10,11,6,7,8]` 表示 RL 顺序与硬件/模型顺序存在显式重排。
- ONNX 与参数是配套件:`params/deploy.yaml` + `exported/policy.onnx`(`State_RLBase.cpp:12-16`)。仅导入 `.onnx` 无法可靠解释 observation/action。
### 最小文件级实现方案(建议,不修改)
1. **新增独立策略层**:`src/policy/types.ts`, `OnnxPolicyRuntime.ts`, `Go2VelocityTask.ts`。
- `OnnxPolicyRuntime` 只负责 ORT session、输入/输出名称与 shape 校验、generation cancellation、dispose。
- `Go2VelocityTask` 固化首个 task 的 47→12 契约、命令/phase/lastAction、名称到 joint/actuator 的显式映射、PD 目标或 position actuator 输出。
- 不应把 ONNX 逻辑塞入 `PythonControllerRuntime`;其 Pyodide proxy 与同步函数假设完全不同。
2. **`SimulationSession.ts`**:将单一 `pythonController` 最小泛化为互斥 controller slot,或新增 `onnxPolicy` 并在同一个 `runController()` 中互斥调度;load/remove/reset/dispose 复用 generation 防模型切换竞态。添加只读 observation helpers(按 joint/body/sensor 名称解析),避免 UI 从 snapshot 拼向量。禁用/异常时必须 `ctrl.fill(0)`、暂停并保留错误。
3. **异步风险处理**:`onnxruntime-web` 的 `session.run()` 返回 Promise,不能直接 await 于当前同步 `advance()` while loop。最小侵入方案是“最新完成动作 hold”:到 50 Hz 时复制 observation 发起一次推理(仅允许一个 in-flight),后续物理步继续持有上次动作,完成后原子替换;首帧保持安全默认位姿而非零力矩。状态须显示 inference latency/overrun。若要求训练等价的严格每 20 ms 动作,则需把物理 loop 改为异步/worker,已非最小侵入。
4. **`PhysicsAdapter.ts`**:接口新增 `loadOnnxPolicy(policyBytes, config)`, `setPolicyEnabled`, `setVelocityCommand`, `removePolicy`;MainThread adapter 只转发,模型 `releaseCurrent()` 自动 dispose。不要把 policy 写入 MuJoCo MEMFS,ORT 可直接吃 `Uint8Array`。
5. **项目导入**:`importer.ts` 已保留任意扩展名,ZIP/目录中的 `.onnx/.yaml` 会存在 `manifest.files`;无需修改 `discoverEntries`(它只应发现模型)。在 App 中按配套目录约定寻找 `exported/policy.onnx` 及 `params/deploy.yaml`,首版更稳妥的是新增浏览器可直接解析的 JSON policy manifest;不要为 YAML 引入重量解析器后猜测任意训练配置。独立单文件导入应沿用 Python 的 upsert 流程,但 accept 改 `.onnx` 并施加专用上限。
6. **依赖/构建**:`wasm/package.json` 加 `onnxruntime-web`;`vite.config.ts` 仿照 Pyodide 将 ORT wasm 文件复制到固定相对目录并设置 `ort.env.wasm.wasmPaths`,否则离线 `base:'./'` 构建容易 404。首版使用 wasm CPU 单线程,避免 WebGPU/COOP-COEP 扩大范围。
7. **UI**:新增 `PolicyPanel.tsx`,插在 `SidebarPanel.tsx:32` 的控制 tab,与 Python 控制器并列;展示模型/config、47×12 shape、50 Hz、启停、速度 x/y/yaw、推理耗时/错误。`App.tsx:40,55,72-77,104` 按现有 controller 状态模式增加 policy 状态及 callbacks;加载另一控制器时明确移除/禁用前一个,不能二者同时写 ctrl。
8. **首个 Go2-W 任务边界**:参考仓库提供的是 **Go2 12-DOF** velocity policy,不是 Go2-W policy;Go2-W 多四个轮关节。首版可把 12 个腿关节按该 policy 控制、轮 actuator 固定阻尼/零速,仅称“Go2-W 腿式站立/速度实验”,不能宣称与训练分布一致。真正轮式平衡/速度需要 Go2-W 专属训练导出(observation/action 很可能含轮速和 16 维动作)。
## Architecture
浏览器文件 → `ProjectManifest.files`(二进制 ONNX 已保留)→ `App` 选择配套 policy/config → `PhysicsAdapter` → `SimulationSession` 拥有 policy runtime。每个 MuJoCo 子步前 task 从 `MjData` 组 observation;每 0.02 s 异步提交 ONNX,完成输出经 scale/offset、显式 joint mapping 和限幅后成为 held action;物理 loop 每步应用 held action。snapshot 仅向 UI报告状态,不作为 observation 数据源。
## 观测/动作风险与严重度
- **blocker**:参考资产无 Go2-W ONNX,且仓库参考目录下 Go2 velocity 的 `exported/policy.onnx` 实际不存在;不能验证 input/output 名或数值一致性。
- **blocker**:Go2 12-DOF policy 与 Go2-W 16 actuator 拓扑不匹配。若按 `model.njnt/nactuator` 顺序拼接会 shape 错误或静默错控。
- **high**:`projected_gravity` 必须是重力单位向量由世界系旋到 body/IMU 局部系。MuJoCo `xquat` 为 wxyz;符号、quat inverse、IMU 安装姿态任一错误都会使策略立刻摔倒。加速度计并不等价于无噪声 projected gravity。
- **high**:`base_ang_vel` 需 body/local frame;直接使用 free joint `qvel[3:6]` 或 gyro sensor前必须确认 MuJoCo sensor frame与训练定义一致。
- **high**:`joint_pos_rel = q-default`,joint_vel 和 action 都必须是 YAML 的 RL 顺序;四足 FL/FR/RL/RR 次序在当前手写 Go2-W controller 与 deploy map 中并不天然一致。
- **high**:输出是 normalized action,不是 torque;必须 scale+offset 后交给 position actuator/PD。当前 `setActuator` 只写 ctrl,URDF 自动生成的是 motor(`PhysicsAdapter.ts:68`),因此若直接写位置值会被当 N·m。需要 task 内 PD torque,或为该 task 构造 position actuator。
- **medium**:gait phase 是 0.6 s 周期的 sin/cos,但应确认零点、sin/cos 顺序以及站立零命令时是否冻结;YAML 没编码完整函数语义。
- **medium**:异步 ORT latency 若超过 20 ms 会跳过策略 tick;必须统计 dropped/late inference,且避免并发堆积。
- **medium**:reset 应同时清 lastAction、phase、in-flight generation 和 held action;否则旧 Promise 可污染新 episode。
- **medium**:导入不可信 ONNX 可能造成大内存/计算消耗;现有 128 MiB 文件上限不足以约束 tensor shape/运行时内存。
## 测试建议
1. `src/policy/Go2VelocityTask.test.ts`:固定 qpos/qvel/quaternion/command,逐元素断言 47 维顺序、dtype、projected gravity、phase、joint reorder;固定 raw action 断言 `*0.25+offset` 和 actuator 映射。
2. `src/policy/OnnxPolicyRuntime.test.ts`:mock ORT,覆盖 input/output shape/name 错误、NaN/Inf、single-flight、超时/overrun、generation cancellation、dispose;用极小 identity ONNX fixture 做一次真实 wasm smoke test。
3. 新增 `SimulationSession.test.ts`(当前该核心类无测试):mock MainModule 验证 controller 在 `mj_step` 前执行、50 Hz decimation、held action、reset/disable/error 清 ctrl、Python/ONNX 互斥、模型切换后旧推理无效。
4. `PhysicsAdapter` contract test:load/remove/dispose 转发及失败后 workspace/runtime 都释放。
5. `importer.test.ts`:ZIP/目录保留 `.onnx` 与 config、路径规范化、重复/超限拒绝;entry discovery 不把 ONNX 当模型。
6. `PolicyPanel.test.tsx` + `SidebarPanel` 测试:导入、启停、命令范围、shape/错误展示、另一控制器启用时互斥。
7. Playwright:导入 Go2-W 工程+fixture policy,启用后推进若干秒,断言页面不崩、policy 状态与命令变化;数值验收应另做离线 golden trajectory(相同初态前 N 次 observation/action 与 Python/ORT reference 对齐),容差逐元素约 `1e-5`,并检查 base height/倾角安全阈值。
8. 常规命令:`npm run typecheck:platform`, `npm run lint:platform`, `npm run test:platform`, `npm run build:platform`, 最后 `npm run test:e2e:platform`。
## Start Here
先打开 `wasm/web_platform/src/simulation/SimulationSession.ts:44-93`:这是唯一能保证 observation 采样、控制输出和 `mj_step` 顺序正确,并统一 reset/error/dispose 语义的位置。
## Residual Risks
- 未获得 Go2-W 专属 ONNX 与导出 metadata,无法完成真实 shape/name/golden 验证。
- unitree_rl_mjlab 的 YAML 说明字段顺序与缩放,但 gait phase、frame convention 等函数语义仍需用实际导出/运行器确认。
- 主线程异步 inference 的 held-action 方案是最小侵入折衷,不提供硬实时或训练环境严格等价。
```acceptance-report
{
"criteriaSatisfied": [
{
"id": "criterion-1",
"status": "satisfied",
"evidence": "已给出 SimulationSession、PhysicsAdapter、App/Sidebar、导入与测试结构的具体文件/行范围方案,并按 blocker/high/medium 列出观测动作风险。"
}
],
"changedFiles": [],
"testsAddedOrUpdated": [],
"commandsRun": [
{
"command": "find/grep/read/nl(只读仓库与 unitree_rl_mjlab 参考配置)",
"result": "passed",
"summary": "确认现有控制步序、项目导入行为、测试布局及 Go2 deploy 的 47维观测/12维动作约定。"
}
],
"validationOutput": [
"参考 deploy.yaml: step_dt=0.02,观测 3+3+3+2+12+12+12=47,动作 12,scale=0.25。",
"参考部署约定为 params/deploy.yaml + exported/policy.onnx;当前 Go2 目录未发现实际 ONNX。",
"当前仓库 package dependencies 未包含 onnxruntime-web。"
],
"residualRisks": [
"缺少 Go2-W 专属 ONNX/metadata,参考仅为 Go2 12-DOF。",
"异步浏览器 ORT 与同步物理 loop 存在时序折衷。",
"projected gravity、角速度 frame、joint mapping 必须做 golden 对齐。"
],
"noStagedFiles": true,
"diffSummary": "只读分析;未修改仓库文件,仅写入指定 context.md 报告。",
"reviewFindings": [
"blocker: /home/cen/Embodied_Workspace/unitree_rl_mjlab/deploy/robots/go2 - 未发现 Go2 velocity policy.onnx,无法验证模型 I/O。",
"blocker: deploy.yaml:1-56 - 参考策略为 12-DOF Go2,与 Go2-W 16 actuator 不匹配。",
"high: wasm/web_platform/src/simulation/SimulationSession.ts:49-57 - 同步物理 while loop 无法直接 await ORT Promise。",
"high: wasm/web_platform/src/simulation/PhysicsAdapter.ts:68-70 - URDF 自动生成 motor,normalized position action 不可直接写 ctrl。"
],
"manualNotes": "报告已写入权威路径;除 context.md 外未改动项目。"
}
```
+1 -15
View File
@@ -3,18 +3,4 @@ import globals from 'globals';
import tseslint from 'typescript-eslint';
import reactHooks from 'eslint-plugin-react-hooks';
import reactRefresh from 'eslint-plugin-react-refresh';
export default tseslint.config(
{ ignores: ['web-platform-dist', 'node_modules'] },
js.configs.recommended,
...tseslint.configs.recommended,
{
files: ['web_platform/**/*.{ts,tsx}'],
languageOptions: { globals: { ...globals.browser, ...globals.node } },
plugins: { 'react-hooks': reactHooks, 'react-refresh': reactRefresh },
rules: {
...reactHooks.configs.recommended.rules,
'react-refresh/only-export-components': ['warn', { allowConstantExport: true }],
'@typescript-eslint/no-explicit-any': 'off',
},
},
);
export default tseslint.config({ignores:['web-platform-dist','node_modules']},js.configs.recommended,...tseslint.configs.recommended,{files:['web_platform/**/*.{ts,tsx}'],languageOptions:{globals:{...globals.browser,...globals.node}},plugins:{'react-hooks':reactHooks,'react-refresh':reactRefresh},rules:{...reactHooks.configs.recommended.rules,'react-refresh/only-export-components':['warn',{allowConstantExport:true}],'@typescript-eslint/no-explicit-any':'off'}});
+35 -212
View File
@@ -1,12 +1,12 @@
{
"name": "mujoco-web-platform",
"version": "0.8.3",
"version": "0.6.0",
"lockfileVersion": 3,
"requires": true,
"packages": {
"": {
"name": "mujoco-web-platform",
"version": "0.8.3",
"version": "0.6.0",
"license": "Apache-2.0",
"dependencies": {
"@monaco-editor/react": "^4.7.0",
@@ -18,7 +18,6 @@
"pyodide": "^0.29.4",
"react": "^19.2.8",
"react-dom": "^19.2.8",
"uplot": "^1.6.32",
"zustand": "^5.0.15"
},
"devDependencies": {
@@ -31,7 +30,6 @@
"@types/react-dom": "^19.2.4",
"@types/three": "^0.185.4",
"@vitejs/plugin-react": "^6.1.0",
"@vitest/coverage-v8": "4.1.11",
"autoprefixer": "^10.5.4",
"eslint": "^10.8.1",
"eslint-plugin-react-hooks": "^7.1.1",
@@ -39,17 +37,12 @@
"globals": "^17.11.0",
"jsdom": "^30.0.1",
"postcss": "^8.5.26",
"prettier": "3.9.6",
"tailwindcss": "^3.4.17",
"three": "^0.178.0",
"typescript": "5.8.2",
"typescript-eslint": "^8.67.0",
"vite": "^8.0.16",
"vitest": "^4.1.11"
},
"engines": {
"node": ">=24 <25",
"npm": "11.17.0"
}
},
"node_modules/@adobe/css-tools": {
@@ -198,6 +191,17 @@
"node": ">=6.9.0"
}
},
"node_modules/@babel/generator/node_modules/@jridgewell/trace-mapping": {
"version": "0.3.31",
"resolved": "https://registry.npmjs.org/@jridgewell/trace-mapping/-/trace-mapping-0.3.31.tgz",
"integrity": "sha512-zzNR+SdQSDJzc8joaeP8QQoCQr8NuYx2dIIytl1QeBEZHJ9uW6hebsrYgbz8hJwUQao3TWCMtmfV8Nu1twOLAw==",
"dev": true,
"license": "MIT",
"dependencies": {
"@jridgewell/resolve-uri": "^3.1.0",
"@jridgewell/sourcemap-codec": "^1.4.14"
}
},
"node_modules/@babel/helper-compilation-targets": {
"version": "7.29.7",
"resolved": "https://registry.npmjs.org/@babel/helper-compilation-targets/-/helper-compilation-targets-7.29.7.tgz",
@@ -385,16 +389,6 @@
"node": ">=6.9.0"
}
},
"node_modules/@bcoe/v8-coverage": {
"version": "1.0.2",
"resolved": "https://registry.npmjs.org/@bcoe/v8-coverage/-/v8-coverage-1.0.2.tgz",
"integrity": "sha512-6zABk/ECA/QYSCQ1NGiVwwbQerUCZ+TQbp64Q3AgmfNvurHH0j8TtXa1qbShXA6qqkpAj4V5W8pP6mLe1mcMqA==",
"dev": true,
"license": "MIT",
"engines": {
"node": ">=18"
}
},
"node_modules/@bramus/specificity": {
"version": "2.4.2",
"resolved": "https://registry.npmjs.org/@bramus/specificity/-/specificity-2.4.2.tgz",
@@ -851,6 +845,17 @@
"@jridgewell/trace-mapping": "^0.3.24"
}
},
"node_modules/@jridgewell/gen-mapping/node_modules/@jridgewell/trace-mapping": {
"version": "0.3.31",
"resolved": "https://registry.npmjs.org/@jridgewell/trace-mapping/-/trace-mapping-0.3.31.tgz",
"integrity": "sha512-zzNR+SdQSDJzc8joaeP8QQoCQr8NuYx2dIIytl1QeBEZHJ9uW6hebsrYgbz8hJwUQao3TWCMtmfV8Nu1twOLAw==",
"dev": true,
"license": "MIT",
"dependencies": {
"@jridgewell/resolve-uri": "^3.1.0",
"@jridgewell/sourcemap-codec": "^1.4.14"
}
},
"node_modules/@jridgewell/remapping": {
"version": "2.3.5",
"resolved": "https://registry.npmjs.org/@jridgewell/remapping/-/remapping-2.3.5.tgz",
@@ -862,6 +867,17 @@
"@jridgewell/trace-mapping": "^0.3.24"
}
},
"node_modules/@jridgewell/remapping/node_modules/@jridgewell/trace-mapping": {
"version": "0.3.31",
"resolved": "https://registry.npmjs.org/@jridgewell/trace-mapping/-/trace-mapping-0.3.31.tgz",
"integrity": "sha512-zzNR+SdQSDJzc8joaeP8QQoCQr8NuYx2dIIytl1QeBEZHJ9uW6hebsrYgbz8hJwUQao3TWCMtmfV8Nu1twOLAw==",
"dev": true,
"license": "MIT",
"dependencies": {
"@jridgewell/resolve-uri": "^3.1.0",
"@jridgewell/sourcemap-codec": "^1.4.14"
}
},
"node_modules/@jridgewell/resolve-uri": {
"version": "3.1.2",
"resolved": "https://registry.npmjs.org/@jridgewell/resolve-uri/-/resolve-uri-3.1.2.tgz",
@@ -877,17 +893,6 @@
"integrity": "sha512-cYQ9310grqxueWbl+WuIUIaiUaDcj7WOq5fVhEljNVgRfOUhY9fy2zTvfoqWsnebh8Sl70VScFbICvJnLKB0Og==",
"dev": true
},
"node_modules/@jridgewell/trace-mapping": {
"version": "0.3.31",
"resolved": "https://registry.npmjs.org/@jridgewell/trace-mapping/-/trace-mapping-0.3.31.tgz",
"integrity": "sha512-zzNR+SdQSDJzc8joaeP8QQoCQr8NuYx2dIIytl1QeBEZHJ9uW6hebsrYgbz8hJwUQao3TWCMtmfV8Nu1twOLAw==",
"dev": true,
"license": "MIT",
"dependencies": {
"@jridgewell/resolve-uri": "^3.1.0",
"@jridgewell/sourcemap-codec": "^1.4.14"
}
},
"node_modules/@monaco-editor/loader": {
"version": "1.7.0",
"resolved": "https://registry.npmjs.org/@monaco-editor/loader/-/loader-1.7.0.tgz",
@@ -1878,37 +1883,6 @@
}
}
},
"node_modules/@vitest/coverage-v8": {
"version": "4.1.11",
"resolved": "https://registry.npmjs.org/@vitest/coverage-v8/-/coverage-v8-4.1.11.tgz",
"integrity": "sha512-8MVGEFnJIcdGjcbfKmeq8z0pZHH0JlVtoVZH9Q/qwUp6wyFnEJUBMrw9DCaj+ra3vShGmhavjalMIhPNxZAUcw==",
"dev": true,
"license": "MIT",
"dependencies": {
"@bcoe/v8-coverage": "^1.0.2",
"@vitest/utils": "4.1.11",
"ast-v8-to-istanbul": "^1.0.0",
"istanbul-lib-coverage": "^3.2.2",
"istanbul-lib-report": "^3.0.1",
"istanbul-reports": "^3.2.0",
"magicast": "^0.5.2",
"obug": "^2.1.1",
"std-env": "^4.0.0-rc.1",
"tinyrainbow": "^3.1.0"
},
"funding": {
"url": "https://opencollective.com/vitest"
},
"peerDependencies": {
"@vitest/browser": "4.1.11",
"vitest": "4.1.11"
},
"peerDependenciesMeta": {
"@vitest/browser": {
"optional": true
}
}
},
"node_modules/@vitest/expect": {
"version": "4.1.11",
"resolved": "https://registry.npmjs.org/@vitest/expect/-/expect-4.1.11.tgz",
@@ -2116,25 +2090,6 @@
"node": ">=12"
}
},
"node_modules/ast-v8-to-istanbul": {
"version": "1.0.5",
"resolved": "https://registry.npmjs.org/ast-v8-to-istanbul/-/ast-v8-to-istanbul-1.0.5.tgz",
"integrity": "sha512-UPAgKJFSEGMWSDr3LX4tqnAb4f7KGT8O40Tyx8wbYmmZ/yn58lNCm8h3svs3eXgiGd5AXxz8NDOvXWvicq+rJA==",
"dev": true,
"license": "MIT",
"dependencies": {
"@jridgewell/trace-mapping": "^0.3.31",
"estree-walker": "^3.0.3",
"js-tokens": "^10.0.0"
}
},
"node_modules/ast-v8-to-istanbul/node_modules/js-tokens": {
"version": "10.0.0",
"resolved": "https://registry.npmjs.org/js-tokens/-/js-tokens-10.0.0.tgz",
"integrity": "sha512-lM/UBzQmfJRo9ABXbPWemivdCW8V2G8FHaHdypQaIy523snUjog0W71ayWXTjiR+ixeMyVHN2XcpnTd/liPg/Q==",
"dev": true,
"license": "MIT"
},
"node_modules/autoprefixer": {
"version": "10.5.4",
"resolved": "https://registry.npmjs.org/autoprefixer/-/autoprefixer-10.5.4.tgz",
@@ -3058,16 +3013,6 @@
"integrity": "sha512-Y8T4vYhEfwJOTbouREvG+3XDsjr8E3kIr7uf+JZ0BYloFsttiHU0WfvANVsR7TxNUJa/WpCnw/Ino/p+DeBhBQ==",
"license": "ISC"
},
"node_modules/has-flag": {
"version": "4.0.0",
"resolved": "https://registry.npmjs.org/has-flag/-/has-flag-4.0.0.tgz",
"integrity": "sha512-EykJT/Q1KjTWctppgIAgfSO0tKVuZUjhgMr17kqTumMl6Afv3EISleU7qZUzoXDFTAHTDC4NOoG/ZxU3EvlMPQ==",
"dev": true,
"license": "MIT",
"engines": {
"node": ">=8"
}
},
"node_modules/hasown": {
"version": "2.0.4",
"resolved": "https://registry.npmjs.org/hasown/-/hasown-2.0.4.tgz",
@@ -3111,13 +3056,6 @@
"node": "^20.19.0 || ^22.12.0 || >=24.0.0"
}
},
"node_modules/html-escaper": {
"version": "2.0.2",
"resolved": "https://registry.npmjs.org/html-escaper/-/html-escaper-2.0.2.tgz",
"integrity": "sha512-H2iMtd0I4Mt5eYiapRdIDjp+XzelXQ0tFE4JS7YFwFevXXMmOp9myNrUvCg0D6ws8iqkRPBfKHgbwig1SmlLfg==",
"dev": true,
"license": "MIT"
},
"node_modules/ignore": {
"version": "5.3.2",
"resolved": "https://registry.npmjs.org/ignore/-/ignore-5.3.2.tgz",
@@ -3223,45 +3161,6 @@
"integrity": "sha512-RHxMLp9lnKHGHRng9QFhRCMbYAcVpn69smSGcq3f36xjgVVWThj4qqLbTLlq7Ssj8B+fIQ1EuCEGI2lKsyQeIw==",
"dev": true
},
"node_modules/istanbul-lib-coverage": {
"version": "3.2.2",
"resolved": "https://registry.npmjs.org/istanbul-lib-coverage/-/istanbul-lib-coverage-3.2.2.tgz",
"integrity": "sha512-O8dpsF+r0WV/8MNRKfnmrtCWhuKjxrq2w+jpzBL5UZKTi2LeVWnWOmWRxFlesJONmc+wLAGvKQZEOanko0LFTg==",
"dev": true,
"license": "BSD-3-Clause",
"engines": {
"node": ">=8"
}
},
"node_modules/istanbul-lib-report": {
"version": "3.0.1",
"resolved": "https://registry.npmjs.org/istanbul-lib-report/-/istanbul-lib-report-3.0.1.tgz",
"integrity": "sha512-GCfE1mtsHGOELCU8e/Z7YWzpmybrx/+dSTfLrvY8qRmaY6zXTKWn6WQIjaAFw069icm6GVMNkgu0NzI4iPZUNw==",
"dev": true,
"license": "BSD-3-Clause",
"dependencies": {
"istanbul-lib-coverage": "^3.0.0",
"make-dir": "^4.0.0",
"supports-color": "^7.1.0"
},
"engines": {
"node": ">=10"
}
},
"node_modules/istanbul-reports": {
"version": "3.2.0",
"resolved": "https://registry.npmjs.org/istanbul-reports/-/istanbul-reports-3.2.0.tgz",
"integrity": "sha512-HGYWWS/ehqTV3xN10i23tkPkpH46MLCIMFNCaaKNavAXTF1RkqxawEPtnjnGZ6XKSInBKkiOA5BKS+aZiY3AvA==",
"dev": true,
"license": "BSD-3-Clause",
"dependencies": {
"html-escaper": "^2.0.0",
"istanbul-lib-report": "^3.0.0"
},
"engines": {
"node": ">=8"
}
},
"node_modules/jiti": {
"version": "1.21.7",
"resolved": "https://registry.npmjs.org/jiti/-/jiti-1.21.7.tgz",
@@ -3746,47 +3645,6 @@
"@jridgewell/sourcemap-codec": "^1.5.5"
}
},
"node_modules/magicast": {
"version": "0.5.4",
"resolved": "https://registry.npmjs.org/magicast/-/magicast-0.5.4.tgz",
"integrity": "sha512-llBEhWm1SacoRwgHUoQJYtwp4PBLF4faQi5TCpIGyGs9n4y5+juI0tDgyKIfpqxckRHaHzouUEph3THklWh03w==",
"dev": true,
"license": "MIT",
"dependencies": {
"@babel/parser": "^7.29.7",
"@babel/types": "^7.29.7",
"source-map-js": "^1.2.1"
}
},
"node_modules/make-dir": {
"version": "4.0.0",
"resolved": "https://registry.npmjs.org/make-dir/-/make-dir-4.0.0.tgz",
"integrity": "sha512-hXdUTZYIVOt1Ex//jAQi+wTZZpUpwBj/0QsOzqegb3rGMMeJiSEu5xLHnYfBrRV4RH2+OCSOO95Is/7x1WJ4bw==",
"dev": true,
"license": "MIT",
"dependencies": {
"semver": "^7.5.3"
},
"engines": {
"node": ">=10"
},
"funding": {
"url": "https://github.com/sponsors/sindresorhus"
}
},
"node_modules/make-dir/node_modules/semver": {
"version": "7.8.5",
"resolved": "https://registry.npmjs.org/semver/-/semver-7.8.5.tgz",
"integrity": "sha512-Y7/KDsb8LjooZpwaqGyulO6DQlksgCncchHGk+sZIY4SBvUocMBEFH5Ur1fI4dV+Jvl0w6cjvucaIi40puRioA==",
"dev": true,
"license": "ISC",
"bin": {
"semver": "bin/semver.js"
},
"engines": {
"node": ">=10"
}
},
"node_modules/marked": {
"version": "14.0.0",
"resolved": "https://registry.npmjs.org/marked/-/marked-14.0.0.tgz",
@@ -4334,22 +4192,6 @@
"node": ">= 0.8.0"
}
},
"node_modules/prettier": {
"version": "3.9.6",
"resolved": "https://registry.npmjs.org/prettier/-/prettier-3.9.6.tgz",
"integrity": "sha512-OpN0zzVdiaiAhxpuuj5efpIS4sY9j7bY6uR5mnj5yPzGkdkjNKSJeUThPb60Jw29QuAZgA4o+/iB49kFiaBX6g==",
"dev": true,
"license": "MIT",
"bin": {
"prettier": "bin/prettier.cjs"
},
"engines": {
"node": ">=14"
},
"funding": {
"url": "https://github.com/prettier/prettier?sponsor=1"
}
},
"node_modules/pretty-format": {
"version": "27.5.1",
"resolved": "https://registry.npmjs.org/pretty-format/-/pretty-format-27.5.1.tgz",
@@ -4758,19 +4600,6 @@
"node": ">=16 || 14 >=14.17"
}
},
"node_modules/supports-color": {
"version": "7.2.0",
"resolved": "https://registry.npmjs.org/supports-color/-/supports-color-7.2.0.tgz",
"integrity": "sha512-qpCAvRl9stuOHveKsn7HncJRvv501qIacKzQlO/+Lwxc9+0q2wLyv4Dfvt80/DPn2pqOBsJdDiogXGR9+OvwRw==",
"dev": true,
"license": "MIT",
"dependencies": {
"has-flag": "^4.0.0"
},
"engines": {
"node": ">=8"
}
},
"node_modules/supports-preserve-symlinks-flag": {
"version": "1.0.0",
"resolved": "https://registry.npmjs.org/supports-preserve-symlinks-flag/-/supports-preserve-symlinks-flag-1.0.0.tgz",
@@ -5092,12 +4921,6 @@
"browserslist": ">= 4.21.0"
}
},
"node_modules/uplot": {
"version": "1.6.32",
"resolved": "https://registry.npmjs.org/uplot/-/uplot-1.6.32.tgz",
"integrity": "sha512-KIMVnG68zvu5XXUbC4LQEPnhwOxBuLyW1AHtpm6IKTXImkbLgkMy+jabjLgSLMasNuGGzQm/ep3tOkyTxpiQIw==",
"license": "MIT"
},
"node_modules/uri-js": {
"version": "4.4.1",
"resolved": "https://registry.npmjs.org/uri-js/-/uri-js-4.4.1.tgz",
+3 -16
View File
@@ -1,6 +1,6 @@
{
"name": "mujoco-web-platform",
"version": "0.8.3",
"version": "0.6.0",
"description": "基于 MuJoCo WebAssembly 的本地机器人仿真与控制平台",
"private": true,
"type": "module",
@@ -14,12 +14,7 @@
"test:e2e": "playwright test -c web_platform/playwright.config.ts",
"training-server": "python3 training_server/server.py",
"test:training-server": "python3 -m unittest discover -s training_server/tests",
"check": "npm run typecheck && npm run lint && npm run check:format && npm run test:coverage && npm run test:training-server && npm run build",
"format": "prettier --write .",
"check:format": "prettier --check .",
"test:coverage": "vitest run --coverage --config web_platform/vite.config.ts",
"lint:python": "python3 -m ruff check training_server",
"format:python": "python3 -m ruff format training_server"
"check": "npm run typecheck && npm run lint && npm run test && npm run test:training-server && npm run build"
},
"license": "Apache-2.0",
"devDependencies": {
@@ -32,7 +27,6 @@
"@types/react-dom": "^19.2.4",
"@types/three": "^0.185.4",
"@vitejs/plugin-react": "^6.1.0",
"@vitest/coverage-v8": "4.1.11",
"autoprefixer": "^10.5.4",
"eslint": "^10.8.1",
"eslint-plugin-react-hooks": "^7.1.1",
@@ -40,7 +34,6 @@
"globals": "^17.11.0",
"jsdom": "^30.0.1",
"postcss": "^8.5.26",
"prettier": "3.9.6",
"tailwindcss": "^3.4.17",
"three": "^0.178.0",
"typescript": "5.8.2",
@@ -61,12 +54,6 @@
"pyodide": "^0.29.4",
"react": "^19.2.8",
"react-dom": "^19.2.8",
"uplot": "^1.6.32",
"zustand": "^5.0.15"
},
"engines": {
"node": ">=24 <25",
"npm": "11.17.0"
},
"packageManager": "npm@11.17.0"
}
}
-11
View File
@@ -1,11 +0,0 @@
[tool.ruff]
target-version = "py312"
line-length = 100
extend-exclude = [".venv", "build", "training_server/rl"]
[tool.ruff.lint]
select = ["E", "F", "I", "UP", "B", "SIM"]
[tool.ruff.format]
quote-style = "double"
indent-style = "space"
-2
View File
@@ -1,2 +0,0 @@
-r training_server/requirements.txt
ruff==0.16.5
+10 -59
View File
@@ -1,65 +1,30 @@
# 本地强化学习训练服务
该服务把 Web 平台发出的受限训练请求转换为本机训练子进程,并提供状态轮询、停止任务和 `policy.onnx` 下载接口。它还提供 `Unitree-Go2-Flat` 奖励函数自调参:DeepSeek Agent 根据训练曲线与固定评估指标提出受限参数 patch,系统支持全自动或逐轮审批、successive-halving、TensorBoard scalar 查询、最佳 preset 与 ONNX 导出。服务只绑定 `127.0.0.1`,不执行前端传入的任意命令或 Agent 生成的代码。
仓库已在 [`rl/`](rl/) 内置 `Unitree-Go2-Flat` 所需的 PPO 训练代码、Go2 模型资产和 ONNX 导出逻辑,不再要求另外克隆 `unitree_rl_mjlab`。`mjlab`、PyTorch 等大型运行依赖仍需安装在本机训练环境中。
## 准备训练环境
使用仓库已有的 `.venv` 项目虚拟环境:
```bash
source .venv/bin/activate
python -m pip install -r training_server/rl/requirements.txt
python -m pip install -r training_server/requirements.txt
```
该服务把 Web 平台发出的受限训练请求转换为本机 `unitree_rl_mjlab/scripts/train.py` 子进程,并提供状态轮询、停止任务和 `policy.onnx` 下载接口。服务只绑定 `127.0.0.1`,不执行前端传入的任意命令。
## 启动
使用已安装训练依赖的 Python 解释器启动服务:
必须使用已经安装 `mjlab`、PyTorch 和训练工程依赖的 Python 解释器:
```bash
npm run training-server -- \
--trainer-python "$PWD/.venv/bin/python"
```
服务启动时会在终端显示一个随机访问令牌。将该令牌填入前端“访问令牌”字段后再连接。令牌只保存在当前浏览器标签页的 `sessionStorage` 中。自动化启动时可固定令牌:
```bash
MUJOCO_TRAINING_TOKEN='至少十六个字符的随机令牌' npm run training-server -- \
--trainer-root /path/to/unitree_rl_mjlab \
--trainer-python /path/to/training-env/bin/python
```
启用云端自调参 Agent 时,在**服务端环境变量**中配置 DeepSeek;不要把 key 填到浏览器、URL 或命令行参数:
如:
```bash
export DEEPSEEK_API_KEY='你的 DeepSeek API key'
export MUJOCO_TUNING_AGENT_BASE_URL='https://api.deepseek.com' # 可省略
export MUJOCO_TUNING_AGENT_MODEL='deepseek-v4-flash' # 可省略
npm run training-server -- --trainer-python "$PWD/.venv/bin/python"
npm run training-server -- --trainer-root /home/cen/Embodied_Workspace/unitree_rl_mjlab --trainer-python /home/cen/miniconda3/envs/unitree_rl_mjlab/bin/python
```
普通训练不要求 DeepSeek key。未配置时健康接口会把 tuning 标记为不可用;只有创建 session 时明确勾选 fallback,才允许 Agent 失败后使用 Optuna 候选,不会静默降级。
默认训练工程是仓库内的 `training_server/rl`。如需使用包含其他已注册任务的外部训练工程,仍可通过 `--trainer-root /path/to/trainer` 或 `UNITREE_RL_MJLAB_ROOT` 覆盖。默认端口是 `8765`。如果前端不是从 `localhost` 或 `127.0.0.1` 提供,可显式添加来源:
也可用 `UNITREE_RL_MJLAB_ROOT` 指定工程目录。默认端口是 `8765`。如果前端不是从 `localhost` 或 `127.0.0.1` 提供,可显式添加来源:
```bash
python training_server/server.py \
--trainer-python /path/to/training-env/bin/python \
--trainer-root /path/to/unitree_rl_mjlab \
--allow-origin http://192.168.1.10:5173
```
普通训练与自调参共享同一个计算资源锁,任何时刻只允许一个训练/评估子进程占用 GPU。普通任务最多保留 20 个内存状态和每个任务 200 行最近日志;调参 session、trial、proposal、审计和 scalar 写入 SQLite/WAL,默认保存在 `training_server/rl/logs/auto_tuning/`。服务重启后等待审批/暂停状态可恢复,正在训练或评估的 trial 标记为 interrupted,只能从已完整保存的 checkpoint 显式恢复。停止服务或点击“停止”会终止整个进程组。
普通训练的 W&B 模式默认为 `offline`;调参 trial 强制使用本地 TensorBoard writer,DeepSeek 只接收最多 12 个 trial 的脱敏数值摘要和降采样曲线,不接收源代码、机器人资产、checkpoint、服务 token 或本地路径。所有 HTTP API 请求都必须携带启动时生成的 Bearer Token。
## 自调参流程
在主工作台连接训练服务后,点击“打开自调参 Agent 工作台”。默认预算为 12 个唯一配置:所有配置先训练 300 iterations,前 4 名续训到 900,前 2 名续训到 2000;默认使用 GPU 0 和 4096 个并行环境。首次使用建议先降低为 256–512 environments 做 smoke test。
固定评估使用站立、前进/侧移、转向和组合命令以及 3 个固定 seed。最终分数不直接使用可被权重放大的总 reward,而由速度跟踪 35%、动作平滑 20%、姿态稳定 15%、减少跌倒 15%、足端滑移 10%、能耗 5% 的权重无关指标组成。跌倒率高于基线 2% 或速度误差恶化超过 5% 的 trial 不晋级。逐轮审批模式会自动运行基线,之后每条 Agent 建议都等待批准、修改后批准或拒绝反馈;运行中的 session 也可在逐轮审批和全自动之间切换,切到全自动时会批准当前待处理建议。
最佳结果保存为不可变 preset,可在普通训练面板的“奖励配置”中选择,也可导出 JSON;不会覆盖仓库里的 Python 默认奖励配置。
服务一次只运行一个训练任务。停止服务或在界面点击“停止训练”会向整个训练进程组发送终止信号。训练请求的 W&B 模式默认为 `offline`,保留本地指标但不登录;也可以在界面选择完全禁用或在线模式。
## 接口
@@ -67,26 +32,12 @@ python training_server/server.py \
- `POST /api/training/jobs`:发起训练;
- `GET /api/training/jobs/{id}`:状态、迭代进度和最近日志;
- `DELETE /api/training/jobs/{id}`:停止训练;
- `GET /api/training/jobs/{id}/artifacts/policy.onnx`:下载本次生成的策略;
- `GET /api/tuning/capabilities`、`POST /api/tuning/agent/test`:检查/测试 Agent;
- `GET|POST /api/tuning/sessions`、`GET|DELETE /api/tuning/sessions/{id}`:列出、创建、查询、停止 session;
- `POST /api/tuning/sessions/{id}/pause|resume`:暂停后续调度或恢复;
- `POST /api/tuning/sessions/{id}/mode`:运行时切换 `automatic`/`approval` 模式;
- `PUT /api/tuning/sessions/{id}/constraints`:以 revision CAS 保存参数固定值/工程上下限,服务端在 Agent、fallback 与人工修改三条路径统一强制;
- `POST /api/tuning/sessions/{id}/step`:发放且只消费一个 Trial 调度令牌,完成训练与固定评估后重新暂停;
- `POST /api/tuning/sessions/{id}/rollback`:把同 Session 内已完成且通过安全门槛的 Trial 设为非破坏性后续基准,可同时验证其 checkpoint;
- `POST /api/tuning/sessions/{id}/proposals/{proposalId}/approve|reject`:审批、修改或拒绝建议;
- `GET /api/tuning/sessions/{id}/trials/{trialId}/metrics?afterStep=N`:查询降采样或增量 scalar;
- `GET /api/tuning/sessions/{id}/artifacts/best/policy.onnx`:下载最佳策略;
- `GET /api/tuning/presets`:列出可供普通训练复用的最佳奖励 preset。
- `GET /api/training/jobs/{id}/artifacts/policy.onnx`:下载本次生成的策略。
普通任务状态在服务重启后丢失,但日志、checkpoint 和 ONNX 保留在 `training_server/rl/logs/rsl_rl/`;调参状态、调度令牌、参数护栏、回滚基准及产物持久化在 `logs/auto_tuning/`。候选配置数可在 1–100 间设置(包含基线配置,仍受连续无提升早停约束)。API 只接收 32 位资源 ID,不接收客户端文件路径;奖励 patch 受到名称、符号、上下界、Session 护栏、每轮最多 4 项及 `0.5×–2×` 变化率校验。
任务保存在服务内存中,服务重启后历史任务状态会丢失,但 mjlab 日志和 checkpoint 仍保留在训练工程中。
## 测试
```bash
source .venv/bin/activate
python -m pip install -r requirements-dev.txt
npm run lint:python
npm run test:training-server
```
-5
View File
@@ -1,5 +0,0 @@
# 本地自调参服务(Python 3.12)
pydantic-ai-slim[openai]==2.37.0
httpx2[socks]==2.12.0
optuna==4.9.0
tensorboard==2.21.0
-201
View File
@@ -1,201 +0,0 @@
Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
1. Definitions.
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
"Legal Entity" shall mean the union of the acting entity and all
other entities that control, are controlled by, or are under common
control with that entity. For the purposes of this definition,
"control" means (i) the power, direct or indirect, to cause the
direction or management of such entity, whether by contract or
otherwise, or (ii) ownership of fifty percent (50%) or more of the
outstanding shares, or (iii) beneficial ownership of such entity.
"You" (or "Your") shall mean an individual or Legal Entity
exercising permissions granted by this License.
"Source" form shall mean the preferred form for making modifications,
including but not limited to software source code, documentation
source, and configuration files.
"Object" form shall mean any form resulting from mechanical
transformation or translation of a Source form, including but
not limited to compiled object code, generated documentation,
and conversions to other media types.
"Work" shall mean the work of authorship, whether in Source or
Object form, made available under the License, as indicated by a
copyright notice that is included in or attached to the work
(an example is provided in the Appendix below).
"Derivative Works" shall mean any work, whether in Source or Object
form, that is based on (or derived from) the Work and for which the
editorial revisions, annotations, elaborations, or other modifications
represent, as a whole, an original work of authorship. For the purposes
of this License, Derivative Works shall not include works that remain
separable from, or merely link (or bind by name) to the interfaces of,
the Work and Derivative Works thereof.
"Contribution" shall mean any work of authorship, including
the original version of the Work and any modifications or additions
to that Work or Derivative Works thereof, that is intentionally
submitted to Licensor for inclusion in the Work by the copyright owner
or by an individual or Legal Entity authorized to submit on behalf of
the copyright owner. For the purposes of this definition, "submitted"
means any form of electronic, verbal, or written communication sent
to the Licensor or its representatives, including but not limited to
communication on electronic mailing lists, source code control systems,
and issue tracking systems that are managed by, or on behalf of, the
Licensor for the purpose of discussing and improving the Work, but
excluding communication that is conspicuously marked or otherwise
designated in writing by the copyright owner as "Not a Contribution."
"Contributor" shall mean Licensor and any individual or Legal Entity
on behalf of whom a Contribution has been received by Licensor and
subsequently incorporated within the Work.
2. Grant of Copyright License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
copyright license to reproduce, prepare Derivative Works of,
publicly display, publicly perform, sublicense, and distribute the
Work and such Derivative Works in Source or Object form.
3. Grant of Patent License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
(except as stated in this section) patent license to make, have made,
use, offer to sell, sell, import, and otherwise transfer the Work,
where such license applies only to those patent claims licensable
by such Contributor that are necessarily infringed by their
Contribution(s) alone or by combination of their Contribution(s)
with the Work to which such Contribution(s) was submitted. If You
institute patent litigation against any entity (including a
cross-claim or counterclaim in a lawsuit) alleging that the Work
or a Contribution incorporated within the Work constitutes direct
or contributory patent infringement, then any patent licenses
granted to You under this License for that Work shall terminate
as of the date such litigation is filed.
4. Redistribution. You may reproduce and distribute copies of the
Work or Derivative Works thereof in any medium, with or without
modifications, and in Source or Object form, provided that You
meet the following conditions:
(a) You must give any other recipients of the Work or
Derivative Works a copy of this License; and
(b) You must cause any modified files to carry prominent notices
stating that You changed the files; and
(c) You must retain, in the Source form of any Derivative Works
that You distribute, all copyright, patent, trademark, and
attribution notices from the Source form of the Work,
excluding those notices that do not pertain to any part of
the Derivative Works; and
(d) If the Work includes a "NOTICE" text file as part of its
distribution, then any Derivative Works that You distribute must
include a readable copy of the attribution notices contained
within such NOTICE file, excluding those notices that do not
pertain to any part of the Derivative Works, in at least one
of the following places: within a NOTICE text file distributed
as part of the Derivative Works; within the Source form or
documentation, if provided along with the Derivative Works; or,
within a display generated by the Derivative Works, if and
wherever such third-party notices normally appear. The contents
of the NOTICE file are for informational purposes only and
do not modify the License. You may add Your own attribution
notices within Derivative Works that You distribute, alongside
or as an addendum to the NOTICE text from the Work, provided
that such additional attribution notices cannot be construed
as modifying the License.
You may add Your own copyright statement to Your modifications and
may provide additional or different license terms and conditions
for use, reproduction, or distribution of Your modifications, or
for any such Derivative Works as a whole, provided Your use,
reproduction, and distribution of the Work otherwise complies with
the conditions stated in this License.
5. Submission of Contributions. Unless You explicitly state otherwise,
any Contribution intentionally submitted for inclusion in the Work
by You to the Licensor shall be under the terms and conditions of
this License, without any additional terms or conditions.
Notwithstanding the above, nothing herein shall supersede or modify
the terms of any separate license agreement you may have executed
with Licensor regarding such Contributions.
6. Trademarks. This License does not grant permission to use the trade
names, trademarks, service marks, or product names of the Licensor,
except as required for reasonable and customary use in describing the
origin of the Work and reproducing the content of the NOTICE file.
7. Disclaimer of Warranty. Unless required by applicable law or
agreed to in writing, Licensor provides the Work (and each
Contributor provides its Contributions) on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, including, without limitation, any warranties or conditions
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
PARTICULAR PURPOSE. You are solely responsible for determining the
appropriateness of using or redistributing the Work and assume any
risks associated with Your exercise of permissions under this License.
8. Limitation of Liability. In no event and under no legal theory,
whether in tort (including negligence), contract, or otherwise,
unless required by applicable law (such as deliberate and grossly
negligent acts) or agreed to in writing, shall any Contributor be
liable to You for damages, including any direct, indirect, special,
incidental, or consequential damages of any character arising as a
result of this License or out of the use or inability to use the
Work (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all
other commercial damages or losses), even if such Contributor
has been advised of the possibility of such damages.
9. Accepting Warranty or Additional Liability. While redistributing
the Work or Derivative Works thereof, You may choose to offer,
and charge a fee for, acceptance of support, warranty, indemnity,
or other liability obligations and/or rights consistent with this
License. However, in accepting such obligations, You may act only
on Your own behalf and on Your sole responsibility, not on behalf
of any other Contributor, and only if You agree to indemnify,
defend, and hold each Contributor harmless for any liability
incurred by, or claims asserted against, such Contributor by reason
of your accepting any such warranty or additional liability.
END OF TERMS AND CONDITIONS
APPENDIX: How to apply the Apache License to your work.
To apply the Apache License to your work, attach the following
boilerplate notice, with the fields enclosed by brackets "[]"
replaced with your own identifying information. (Don't include
the brackets!) The text should be enclosed in the appropriate
comment syntax for the file format. We also recommend that a
file or class name and description of purpose be included on the
same "printed page" as the copyright notice for easier
identification within third-party archives.
Copyright 2024 The Isaac Lab Project Developers
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.
-28
View File
@@ -1,28 +0,0 @@
# 内置 Go2 强化学习训练器
这里包含 `training_server` 默认任务 `Unitree-Go2-Flat` 所需的训练代码与资产:
- RSL-RL/PPO 训练入口;
- Go2 速度跟踪环境、观测、奖励、终止条件和课程配置;
- Go2 MJCF 与网格资产;
- checkpoint 保存及 `policy.onnx` 导出 runner。
这是 [unitreerobotics/unitree_rl_mjlab](https://github.com/unitreerobotics/unitree_rl_mjlab) 的最小内置子集,来源提交为 `1425b15f73bd4095f0df53709d7c389c3eb9e790`,按上游 Apache License 2.0 许可证分发,见 [`LICENSE`](LICENSE)。为避免复制与默认训练任务无关的机器人、真机部署和运动模仿代码,只保留了 Go2 速度任务依赖。
## 环境
从仓库根目录使用已有的 `.venv` 项目虚拟环境:
```bash
source .venv/bin/activate
python -m pip install -r training_server/rl/requirements.txt
```
训练服务会直接从本目录启动 `scripts/train.py`,不需要安装本目录自身。训练产物写入本目录的 `logs/rsl_rl/`(已被仓库忽略)。
可以脱离服务验证任务是否成功注册:
```bash
cd training_server/rl
python scripts/train.py Unitree-Go2-Flat --help
```
-9
View File
@@ -1,9 +0,0 @@
# 与内置 unitree_rl_mjlab 子集验证一致的训练栈。
# mjlab 1.2.0 会使用 MuJoCo 3.5 的 mjENBL_MULTICCD API,不能放任
# pip 将 mujoco/warp-lang 升级到较新的不兼容版本。
mjlab==1.2.0
mujoco==3.5.0
mujoco-warp==3.5.0
warp-lang==1.15.0
# RSL-RL 5.0.1 仍向 wandb.Settings 传递 start_method,0.29 已删除该字段。
wandb==0.28.2
-222
View File
@@ -1,222 +0,0 @@
"""Deterministic, headless evaluation for Unitree Go2 velocity policies."""
from __future__ import annotations
import hashlib
import json
import os
import sys
from dataclasses import asdict, dataclass, field
from pathlib import Path
from statistics import fmean, pstdev
from typing import Literal
TRAINER_ROOT = Path(__file__).resolve().parents[1]
SERVICE_ROOT = TRAINER_ROOT.parent
for source_root in (TRAINER_ROOT, SERVICE_ROOT):
if str(source_root) not in sys.path:
sys.path.insert(0, str(source_root))
import torch
import tyro
import warp as wp
if not hasattr(wp, "context"):
from warp._src import context as warp_context
wp.context = warp_context # type: ignore[attr-defined]
from mjlab.envs import ManagerBasedRlEnv
from mjlab.rl import MjlabOnPolicyRunner, RslRlVecEnvWrapper
from mjlab.tasks.registry import list_tasks, load_env_cfg, load_rl_cfg, load_runner_cls
from mjlab.tasks.velocity.mdp import UniformVelocityCommandCfg
from mjlab.utils.torch import configure_torch_backends
from torch.utils.tensorboard import SummaryWriter
from tuning.schema import apply_reward_configuration, validate_configuration
SCENARIOS = (
(0.0, 0.0, 0.0),
(0.5, 0.0, 0.0),
(1.0, 0.0, 0.0),
(1.5, 0.0, 0.0),
(0.0, 0.5, 0.0),
(0.0, -0.5, 0.0),
(0.0, 0.0, 0.5),
(0.0, 0.0, -0.5),
(0.8, 0.25, 0.35),
)
METRIC_NAMES = (
"linear_velocity_rmse",
"angular_velocity_rmse",
"mean_action_acc",
"orientation_error",
"slip_velocity",
"mechanical_power",
)
@dataclass(frozen=True)
class EvaluateConfig:
checkpoint: str
output: str
reward_config: str | None = None
num_envs: int = 256
steps_per_seed: int = 1000
seeds: tuple[int, ...] = field(default_factory=lambda: (101, 202, 303))
device: str | None = None
gpu_ids: list[int] | Literal["all"] | None = field(default_factory=lambda: [0])
def _sha256(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as stream:
for chunk in iter(lambda: stream.read(1024 * 1024), b""):
digest.update(chunk)
return digest.hexdigest()
def _fixed_commands(command: torch.Tensor) -> torch.Tensor:
values = torch.as_tensor(SCENARIOS, device=command.device, dtype=command.dtype)
indexes = torch.arange(command.shape[0], device=command.device) % values.shape[0]
command[:] = values[indexes]
return command
def _evaluate_seed(task_id: str, cfg: EvaluateConfig, seed: int) -> dict[str, float]:
torch.manual_seed(seed)
env_cfg = load_env_cfg(task_id, play=False)
agent_cfg = load_rl_cfg(task_id)
env_cfg.seed = seed
env_cfg.scene.num_envs = cfg.num_envs
env_cfg.curriculum = {}
env_cfg.observations["actor"].enable_corruption = False
env_cfg.events.pop("push_robot", None)
twist_cfg = env_cfg.commands["twist"]
assert isinstance(twist_cfg, UniformVelocityCommandCfg)
twist_cfg.heading_command = False
twist_cfg.ranges.heading = None
twist_cfg.rel_heading_envs = 0.0
twist_cfg.rel_standing_envs = 0.0
twist_cfg.resampling_time_range = (1.0e9, 1.0e9)
if cfg.reward_config:
reward_path = Path(cfg.reward_config).expanduser().resolve(strict=True)
with reward_path.open(encoding="utf-8") as stream:
apply_reward_configuration(env_cfg, validate_configuration(json.load(stream)))
env = ManagerBasedRlEnv(cfg=env_cfg, device=cfg.device or "cuda:0")
wrapped = RslRlVecEnvWrapper(env, clip_actions=agent_cfg.clip_actions)
try:
runner_cls = load_runner_cls(task_id) or MjlabOnPolicyRunner
runner = runner_cls(wrapped, asdict(agent_cfg), log_dir=None, device=wrapped.device)
runner.load(
str(Path(cfg.checkpoint).expanduser().resolve(strict=True)),
load_cfg={"actor": True},
strict=True,
map_location=str(wrapped.device),
)
policy = runner.get_inference_policy(device=str(wrapped.device))
twist = wrapped.unwrapped.command_manager.get_term("twist")
_fixed_commands(twist.command)
obs = wrapped.get_observations()
sums = {name: 0.0 for name in METRIC_NAMES}
samples = 0
terminations = 0
completions = 0
with torch.inference_mode():
for _ in range(cfg.steps_per_seed):
_fixed_commands(twist.command)
obs = wrapped.get_observations()
actions = policy(obs)
obs, _rewards, _dones, _extras = wrapped.step(actions)
manager = wrapped.unwrapped.metrics_manager
for index, name in enumerate(manager.active_terms):
if name in sums:
sums[name] += float(torch.sum(manager._step_values[:, index]).item())
samples += wrapped.num_envs
terminated = wrapped.unwrapped.reset_terminated
timed_out = wrapped.unwrapped.reset_time_outs
terminations += int(torch.count_nonzero(terminated).item())
completions += int(torch.count_nonzero(terminated | timed_out).item())
result = {name: sums[name] / max(samples, 1) for name in METRIC_NAMES}
result["fall_rate"] = terminations / max(completions, wrapped.num_envs)
return result
finally:
wrapped.close()
def run_evaluation(task_id: str, cfg: EvaluateConfig) -> dict:
configure_torch_backends()
selected = cfg.gpu_ids
if selected is None:
os.environ["CUDA_VISIBLE_DEVICES"] = ""
device = "cpu"
else:
if selected == "all":
selected = [0]
os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(map(str, selected))
device = cfg.device or "cuda:0"
os.environ["MUJOCO_GL"] = "egl"
cfg = EvaluateConfig(**{**asdict(cfg), "device": device})
checkpoint = Path(cfg.checkpoint).expanduser().resolve(strict=True)
output = Path(cfg.output).expanduser().resolve()
output.parent.mkdir(parents=True, exist_ok=True)
per_seed = [_evaluate_seed(task_id, cfg, seed) for seed in cfg.seeds]
metrics = {
key: fmean(seed_metrics[key] for seed_metrics in per_seed)
for key in (*METRIC_NAMES, "fall_rate")
}
deviations = {
key: pstdev(seed_metrics[key] for seed_metrics in per_seed)
for key in (*METRIC_NAMES, "fall_rate")
}
result = {
"protocolVersion": 1,
"taskId": task_id,
"checkpoint": checkpoint.name,
"checkpointSha256": _sha256(checkpoint),
"seeds": list(cfg.seeds),
"numEnvs": cfg.num_envs,
"stepsPerSeed": cfg.steps_per_seed,
"scenarios": [list(value) for value in SCENARIOS],
"metrics": metrics,
"metricStd": deviations,
"seedMetrics": [
{"seed": seed, "metrics": values}
for seed, values in zip(cfg.seeds, per_seed, strict=True)
],
}
output.write_text(json.dumps(result, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
writer = SummaryWriter(log_dir=str(output.parent / "evaluation-events"))
try:
for name, value in metrics.items():
writer.add_scalar(f"Evaluation/{name}", value, 0)
finally:
writer.close()
print("MUJOCO_EVALUATION " + json.dumps({"output": str(output), "metrics": metrics}))
return result
def main() -> None:
import mjlab.tasks # noqa: F401
import src.tasks # noqa: F401
chosen_task, remaining = tyro.cli(
tyro.extras.literal_type_from_choices(list_tasks()),
add_help=False,
return_unknown_args=True,
config=mjlab.TYRO_FLAGS,
)
args = tyro.cli(
EvaluateConfig,
args=remaining,
prog=sys.argv[0] + f" {chosen_task}",
config=mjlab.TYRO_FLAGS,
)
run_evaluation(chosen_task, args)
if __name__ == "__main__":
main()
-286
View File
@@ -1,286 +0,0 @@
"""Script to train RL agent with RSL-RL."""
import json
import logging
import os
import sys
from dataclasses import asdict, dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Literal, cast
# 训练器作为仓库内置子集直接从 scripts/ 启动,不要求额外执行 pip install -e。
TRAINER_ROOT = Path(__file__).resolve().parents[1]
SERVICE_ROOT = TRAINER_ROOT.parent
for source_root in (TRAINER_ROOT, SERVICE_ROOT):
if str(source_root) not in sys.path:
sys.path.insert(0, str(source_root))
import tyro
import warp as wp
# mjlab 1.2.0 的 GPU CUDA-graph 检查仍访问旧公开路径 wp.context;
# Warp 1.15 已把实现移到 warp._src.context,但保留了相同 runtime 契约。
if not hasattr(wp, "context"):
from warp._src import context as warp_context
wp.context = warp_context # type: ignore[attr-defined]
from mjlab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg
from mjlab.rl import MjlabOnPolicyRunner, RslRlBaseRunnerCfg, RslRlVecEnvWrapper
from mjlab.tasks.registry import list_tasks, load_env_cfg, load_rl_cfg, load_runner_cls
from mjlab.tasks.tracking.mdp import MotionCommandCfg
from mjlab.utils.gpu import select_gpus
from mjlab.utils.os import dump_yaml, get_checkpoint_path
from mjlab.utils.torch import configure_torch_backends
from mjlab.utils.wrappers import VideoRecorder
from tuning.schema import apply_reward_configuration, validate_configuration
@dataclass(frozen=True)
class TrainConfig:
env: ManagerBasedRlEnvCfg
agent: RslRlBaseRunnerCfg
motion_file: str | None = None
video: bool = False
video_length: int = 200
video_interval: int = 2000
enable_nan_guard: bool = False
torchrunx_log_dir: str | None = None
gpu_ids: list[int] | Literal["all"] | None = field(default_factory=lambda: [0])
output_dir: str | None = None
resume_checkpoint: str | None = None
reward_config: str | None = None
reward_config_json: str | None = None
@staticmethod
def from_task(task_id: str) -> "TrainConfig":
env_cfg = load_env_cfg(task_id)
agent_cfg = load_rl_cfg(task_id)
return TrainConfig(env=env_cfg, agent=agent_cfg)
def _load_reward_config(path: str | None, inline: str | None) -> dict | None:
if path is not None and inline is not None:
raise ValueError("Use only one of reward_config and reward_config_json")
if inline is not None:
if len(inline.encode("utf-8")) > 64 * 1024:
raise ValueError("Reward configuration is larger than 64 KiB")
return validate_configuration(json.loads(inline))
if path is None:
return None
source = Path(path).expanduser().resolve(strict=True)
if source.stat().st_size > 64 * 1024:
raise ValueError("Reward configuration is larger than 64 KiB")
with source.open(encoding="utf-8") as stream:
return validate_configuration(json.load(stream))
def run_train(task_id: str, cfg: TrainConfig, log_dir: Path) -> None:
reward_config = _load_reward_config(cfg.reward_config, cfg.reward_config_json)
if reward_config is not None:
apply_reward_configuration(cfg.env, reward_config)
cuda_visible = os.environ.get("CUDA_VISIBLE_DEVICES", "")
if cuda_visible == "":
device = "cpu"
seed = cfg.agent.seed
rank = 0
else:
local_rank = int(os.environ.get("LOCAL_RANK", "0"))
rank = int(os.environ.get("RANK", "0"))
# Set EGL device to match the CUDA device.
os.environ["MUJOCO_EGL_DEVICE_ID"] = str(local_rank)
device = f"cuda:{local_rank}"
# Set seed to have diversity in different processes.
seed = cfg.agent.seed + local_rank
configure_torch_backends()
cfg.agent.seed = seed
cfg.env.seed = seed
print(f"[INFO] Training with: device={device}, seed={seed}, rank={rank}")
# Check if this is a tracking task by checking for motion command.
is_tracking_task = "motion" in cfg.env.commands and isinstance(
cfg.env.commands["motion"], MotionCommandCfg
)
if is_tracking_task:
if not cfg.motion_file:
raise ValueError("For tracking tasks, --motion-file must be set ...")
motion_path = Path(cfg.motion_file).expanduser().resolve()
if not motion_path.exists():
raise FileNotFoundError(f"Motion file not found: {motion_path}")
motion_cmd = cfg.env.commands["motion"]
assert isinstance(motion_cmd, MotionCommandCfg)
motion_cmd.motion_file = str(motion_path)
print(f"[INFO] Using motion file: {motion_cmd.motion_file}")
# Check if motion_file is already set (e.g., via CLI --env.commands.motion.motion-file).
if motion_cmd.motion_file and Path(motion_cmd.motion_file).exists():
print(f"[INFO] Using local motion file: {motion_cmd.motion_file}")
# Enable NaN guard if requested.
if cfg.enable_nan_guard:
cfg.env.sim.nan_guard.enabled = True
print(f"[INFO] NaN guard enabled, output dir: {cfg.env.sim.nan_guard.output_dir}")
if rank == 0:
print(f"[INFO] Logging experiment in directory: {log_dir}")
env = ManagerBasedRlEnv(
cfg=cfg.env, device=device, render_mode="rgb_array" if cfg.video else None
)
log_root_path = log_dir.parent # Go up from specific run dir to experiment dir.
resume_path: Path | None = None
explicit_resume = cfg.resume_checkpoint is not None
if explicit_resume:
resume_path = Path(cfg.resume_checkpoint).expanduser().resolve(strict=True)
elif cfg.agent.resume:
# Load checkpoint from local filesystem.
resume_path = get_checkpoint_path(
log_root_path, cfg.agent.load_run, cfg.agent.load_checkpoint
)
# Only record videos on rank 0 to avoid multiple workers writing to the same files.
if cfg.video and rank == 0:
env = VideoRecorder(
env,
video_folder=Path(log_dir) / "videos" / "train",
step_trigger=lambda step: step % cfg.video_interval == 0,
video_length=cfg.video_length,
disable_logger=True,
)
print("[INFO] Recording videos during training.")
env = RslRlVecEnvWrapper(env, clip_actions=cfg.agent.clip_actions)
agent_cfg = asdict(cfg.agent)
env_cfg = asdict(cfg.env)
runner_cls = load_runner_cls(task_id)
if runner_cls is None:
runner_cls = MjlabOnPolicyRunner
runner_kwargs = {}
runner = runner_cls(env, agent_cfg, str(log_dir), device, **runner_kwargs)
runner.add_git_repo_to_log(__file__)
if resume_path is not None:
print(f"[INFO]: Loading model checkpoint from: {resume_path}")
runner.load(str(resume_path), map_location=device)
if explicit_resume:
# RSL-RL stores the last completed zero-based iteration and otherwise
# repeats it after load. Explicit tuning promotion uses an absolute target.
runner.current_learning_iteration += 1
# Only write config files from rank 0 to avoid race conditions.
if rank == 0:
dump_yaml(log_dir / "params" / "env.yaml", env_cfg)
dump_yaml(log_dir / "params" / "agent.yaml", agent_cfg)
if reward_config is not None:
reward_snapshot = log_dir / "params" / "reward_config.json"
reward_snapshot.parent.mkdir(parents=True, exist_ok=True)
reward_snapshot.write_text(
json.dumps(reward_config, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
)
iterations = cfg.agent.max_iterations
if explicit_resume:
iterations = max(0, cfg.agent.max_iterations - runner.current_learning_iteration)
print(
f"[INFO] Learning target: current={runner.current_learning_iteration}, "
f"additional={iterations}, target={cfg.agent.max_iterations}",
flush=True,
)
runner.learn(num_learning_iterations=iterations, init_at_random_ep_len=True)
env.close()
def launch_training(task_id: str, args: TrainConfig | None = None):
args = args or TrainConfig.from_task(task_id)
# Create log directory once before launching workers.
if args.output_dir:
log_dir = Path(args.output_dir).expanduser().resolve()
else:
log_root_path = (Path("logs") / "rsl_rl" / args.agent.experiment_name).resolve()
log_dir_name = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
if args.agent.run_name:
log_dir_name += f"_{args.agent.run_name}"
log_dir = log_root_path / log_dir_name
log_dir.mkdir(parents=True, exist_ok=True)
# Select GPUs based on CUDA_VISIBLE_DEVICES and user specification.
selected_gpus, num_gpus = select_gpus(args.gpu_ids)
# Set environment variables for all modes.
if selected_gpus is None:
os.environ["CUDA_VISIBLE_DEVICES"] = ""
else:
os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(map(str, selected_gpus))
os.environ["MUJOCO_GL"] = "egl"
if num_gpus <= 1:
# CPU or single GPU: run directly without torchrunx.
run_train(task_id, args, log_dir)
else:
# Multi-GPU: use torchrunx.
import torchrunx
# torchrunx redirects stdout to logging.
logging.basicConfig(level=logging.INFO)
# Configure torchrunx logging directory.
# Priority: 1) existing env var, 2) user flag, 3) default to {log_dir}/torchrunx.
if "TORCHRUNX_LOG_DIR" not in os.environ:
if args.torchrunx_log_dir is not None:
# User specified a value via flag (could be "" to disable).
os.environ["TORCHRUNX_LOG_DIR"] = args.torchrunx_log_dir
else:
# Default: put logs in training directory.
os.environ["TORCHRUNX_LOG_DIR"] = str(log_dir / "torchrunx")
print(f"[INFO] Launching training with {num_gpus} GPUs", flush=True)
torchrunx.Launcher(
hostnames=["localhost"],
workers_per_host=num_gpus,
backend=None, # Let rsl_rl handle process group initialization.
copy_env_vars=torchrunx.DEFAULT_ENV_VARS_FOR_COPY + ("MUJOCO*",),
).run(run_train, task_id, args, log_dir)
def main():
# Parse first argument to choose the task.
# Import tasks to populate the registry.
import mjlab.tasks # noqa: F401
import src.tasks
all_tasks = list_tasks()
chosen_task, remaining_args = tyro.cli(
tyro.extras.literal_type_from_choices(all_tasks),
add_help=False,
return_unknown_args=True,
config=mjlab.TYRO_FLAGS,
)
args = tyro.cli(
TrainConfig,
args=remaining_args,
default=TrainConfig.from_task(chosen_task),
prog=sys.argv[0] + f" {chosen_task}",
config=mjlab.TYRO_FLAGS,
)
del remaining_args
launch_training(task_id=chosen_task, args=args)
if __name__ == "__main__":
main()
-4
View File
@@ -1,4 +0,0 @@
from pathlib import Path
SRC_PATH: Path = Path(__file__).parent
@@ -1,3 +0,0 @@
"""仓库内置训练器支持的机器人资产。"""
from .unitree_go2.go2_constants import get_go2_robot_cfg as get_go2_robot_cfg
@@ -1 +0,0 @@
"""Unitree Go2 quadruped."""
@@ -1,148 +0,0 @@
"""Unitree Go2 constants."""
from pathlib import Path
import mujoco
from src import SRC_PATH
from mjlab.actuator import BuiltinPositionActuatorCfg
from mjlab.entity import EntityArticulationInfoCfg, EntityCfg
from mjlab.utils.actuator import ElectricActuator, reflected_inertia
from mjlab.utils.os import update_assets
from mjlab.utils.spec_config import CollisionCfg
##
# MJCF and assets.
##
GO2_XML: Path = (
SRC_PATH / "assets" / "robots" / "unitree_go2" / "xmls" / "go2.xml"
)
assert GO2_XML.exists()
def get_assets(meshdir: str) -> dict[str, bytes]:
assets: dict[str, bytes] = {}
update_assets(assets, GO2_XML.parent / "assets", meshdir)
return assets
def get_spec() -> mujoco.MjSpec:
spec = mujoco.MjSpec.from_file(str(GO2_XML))
spec.assets = get_assets(spec.meshdir)
return spec
##
# Actuator config.
##
GO2_ACTUATOR_HIP = BuiltinPositionActuatorCfg(
target_names_expr=(
".*hip_.*",
),
stiffness=20.0,
damping=1.0,
effort_limit=23.5,
armature=0.01,
)
GO2_ACTUATOR_THIGH = BuiltinPositionActuatorCfg(
target_names_expr=(
".*thigh_.*",
),
stiffness=20.0,
damping=1.0,
effort_limit=23.5,
armature=0.01,
)
GO2_ACTUATOR_CALF = BuiltinPositionActuatorCfg(
target_names_expr=(
".*calf_.*",
),
stiffness=40.0,
damping=2.0,
effort_limit=45,
armature=0.02,
)
##
# Keyframes.
##
INIT_STATE = EntityCfg.InitialStateCfg(
pos=(0.0, 0.0, 0.32),
joint_pos={
".*thigh_joint": 0.9,
".*calf_joint": -1.8,
".*R_hip_joint": 0.1,
".*L_hip_joint": -0.1,
},
joint_vel={".*": 0.0},
)
##
# Collision config.
##
_foot_regex = "^[FR][LR]_foot_collision$"
# This disables all collisions except the feet.
# Furthermore, feet self collisions are disabled.
FEET_ONLY_COLLISION = CollisionCfg(
geom_names_expr=(_foot_regex,),
contype=0,
conaffinity=1,
condim=3,
priority=1,
friction=(0.6,),
solimp=(0.9, 0.95, 0.023),
)
# This enables all collisions, excluding self collisions.
# Foot collisions are given custom condim, friction and solimp.
FULL_COLLISION = CollisionCfg(
geom_names_expr=(".*_collision",),
condim={_foot_regex: 3, ".*_collision": 1},
priority={_foot_regex: 1},
friction={_foot_regex: (0.6,)},
solimp={_foot_regex: (0.9, 0.95, 0.023)},
contype=1,
conaffinity=0,
)
##
# Final config.
##
GO2_ARTICULATION = EntityArticulationInfoCfg(
actuators=(
GO2_ACTUATOR_HIP,
GO2_ACTUATOR_THIGH,
GO2_ACTUATOR_CALF,
),
soft_joint_pos_limit_factor=0.9,
)
def get_go2_robot_cfg() -> EntityCfg:
"""Get a fresh Go2 robot configuration instance.
Returns a new EntityCfg instance each time to avoid mutation issues when
the config is shared across multiple places.
"""
return EntityCfg(
init_state=INIT_STATE,
collisions=(FULL_COLLISION,),
spec_fn=get_spec,
articulation=GO2_ARTICULATION,
)
if __name__ == "__main__":
import mujoco.viewer as viewer
from mjlab.entity.entity import Entity
robot = Entity(get_go2_robot_cfg())
viewer.launch(robot.spec.compile())
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
@@ -1,161 +0,0 @@
<mujoco model="go2">
<compiler angle="radian" meshdir="assets" autolimits="true" />
<default>
<default class="go2">
<default class="visual">
<geom type="mesh" contype="0" conaffinity="0" density="0" group="2"/>
</default>
<default class="collision">
<geom priority="1" condim="6" group="3"/>
</default>
<site rgba="1 0 0 1" group="5"/>
</default>
</default>
<asset>
<material name="metal" rgba=".9 .95 .95 1" />
<material name="black" rgba="0 0 0 1" />
<material name="white" rgba="1 1 1 1" />
<material name="gray" rgba="0.671705 0.692426 0.774270 1" />
<mesh file="base_0.obj" />
<mesh file="base_1.obj" />
<mesh file="base_2.obj" />
<mesh file="base_3.obj" />
<mesh file="base_4.obj" />
<mesh file="hip_0.obj" />
<mesh file="hip_1.obj" />
<mesh file="thigh_0.obj" />
<mesh file="thigh_1.obj" />
<mesh file="thigh_mirror_0.obj" />
<mesh file="thigh_mirror_1.obj" />
<mesh file="calf_0.obj" />
<mesh file="calf_1.obj" />
<mesh file="calf_mirror_0.obj" />
<mesh file="calf_mirror_1.obj" />
<mesh file="foot.obj" />
</asset>
<worldbody>
<body name="base_link" pos="0 0 0.445" childclass="go2">
<inertial pos="0.021112 0 -0.005366" quat="-0.000543471 0.713435 -0.00173769 0.700719" mass="6.921" diaginertia="0.107027 0.0980771 0.0244531" />
<freejoint name="floating_base_joint"/>
<geom mesh="base_0" material="black" class="visual" />
<geom mesh="base_1" material="black" class="visual" />
<geom mesh="base_2" material="black" class="visual" />
<geom mesh="base_3" material="white" class="visual" />
<geom mesh="base_4" material="gray" class="visual" />
<geom name="base1_collision" size="0.1881 0.04675 0.057" type="box" class="collision" />
<geom name="base2_collision" size="0.05 0.045" pos="0.285 0 0.01" type="cylinder" class="collision" />
<geom name="base3_collision" size="0.047" pos="0.293 0 -0.06" type="sphere" class="collision" />
<site name="imu" pos="-0.02557 0 0.04232" group="5"/>
<body name="FL_hip" pos="0.1934 0.0465 0">
<inertial pos="-0.0054 0.00194 -0.000105" quat="0.497014 0.499245 0.505462 0.498237" mass="0.678" diaginertia="0.00088403 0.000596003 0.000479967" />
<joint name="FL_hip_joint" axis="1 0 0" range="-1.0472 1.0472"/>
<geom mesh="hip_0" material="metal" class="visual" />
<geom mesh="hip_1" material="gray" class="visual" />
<geom name="FL_hip_collision" size="0.046 0.02" pos="0 0.08 0" quat="1 1 0 0" type="cylinder" class="collision" />
<body name="FL_thigh" pos="0 0.0955 0">
<inertial pos="-0.00374 -0.0223 -0.0327" quat="0.829533 0.0847635 -0.0200632 0.551623" mass="1.152" diaginertia="0.00594973 0.00584149 0.000878787" />
<joint name="FL_thigh_joint" axis="0 1 0" range="-1.5708 3.4907"/>
<geom mesh="thigh_0" material="metal" class="visual" />
<geom mesh="thigh_1" material="gray" class="visual" />
<geom name="FL_thigh_collision" size="0.1065 0.01225 0.017" pos="0 0 -0.1065" quat="0.707107 0 0.707107 0" type="box" class="collision" />
<body name="FL_calf" pos="0 0 -0.213">
<inertial pos="0.00629595 -0.000622121 -0.141417" quat="0.710672 0.00154099 -0.00450087 0.703508" mass="0.241352" diaginertia="0.0014901 0.00146356 5.31397e-05" />
<joint name="FL_calf_joint" axis="0 1 0" range="-2.7227 -0.83776"/>
<geom mesh="calf_0" material="gray" class="visual" />
<geom mesh="calf_1" material="black" class="visual" />
<geom name="FL_calf1_collision" size="0.012 0.06" pos="0.008 0 -0.06" quat="0.994493 0 -0.104807 0" type="cylinder" class="collision" />
<geom name="FL_calf2_collision" size="0.011 0.0325" pos="0.02 0 -0.148" quat="0.999688 0 0.0249974 0" type="cylinder" class="collision" />
<geom pos="0 0 -0.213" mesh="foot" class="visual" material="black" />
<geom name="FL_foot_collision" pos="0 0 -0.213" type="sphere" size="0.022" class="collision"/>
<site name="FL" pos="0 0 -0.213" type="sphere" size="0.022"/>
</body>
</body>
</body>
<body name="FR_hip" pos="0.1934 -0.0465 0">
<inertial pos="-0.0054 -0.00194 -0.000105" quat="0.498237 0.505462 0.499245 0.497014" mass="0.678" diaginertia="0.00088403 0.000596003 0.000479967" />
<joint name="FR_hip_joint" axis="1 0 0" range="-1.0472 1.0472"/>
<geom mesh="hip_0" material="metal" class="visual" quat="4.63268e-05 1 0 0" />
<geom mesh="hip_1" material="gray" class="visual" quat="4.63268e-05 1 0 0" />
<geom name="FR_hip_collision" size="0.046 0.02" pos="0 -0.08 0" quat="0.707107 0.707107 0 0" type="cylinder" class="collision" />
<body name="FR_thigh" pos="0 -0.0955 0">
<inertial pos="-0.00374 0.0223 -0.0327" quat="0.551623 -0.0200632 0.0847635 0.829533" mass="1.152" diaginertia="0.00594973 0.00584149 0.000878787" />
<joint name="FR_thigh_joint" axis="0 1 0" range="-1.5708 3.4907"/>
<geom mesh="thigh_mirror_0" material="metal" class="visual" />
<geom mesh="thigh_mirror_1" material="gray" class="visual" />
<geom name="FR_thigh_collision" size="0.1065 0.01225 0.017" pos="0 0 -0.1065" quat="0.707107 0 0.707107 0" type="box" class="collision" />
<body name="FR_calf" pos="0 0 -0.213">
<inertial pos="0.00629595 0.000622121 -0.141417" quat="0.703508 -0.00450087 0.00154099 0.710672" mass="0.241352" diaginertia="0.0014901 0.00146356 5.31397e-05" />
<joint name="FR_calf_joint" axis="0 1 0" range="-2.7227 -0.83776"/>
<geom mesh="calf_mirror_0" material="gray" class="visual" />
<geom mesh="calf_mirror_1" material="black" class="visual" />
<geom name="FR_calf1_collision" size="0.013 0.06" pos="0.01 0 -0.06" quat="0.995004 0 -0.0998334 0" type="cylinder" class="collision" />
<geom name="FR_calf2_collision" size="0.011 0.0325" pos="0.02 0 -0.148" quat="0.999688 0 0.0249974 0" type="cylinder" class="collision" />
<geom pos="0 0 -0.213" mesh="foot" class="visual" material="black" />
<geom name="FR_foot_collision" pos="0 0 -0.213" type="sphere" size="0.022" class="collision"/>
<site name="FR" pos="0 0 -0.213" type="sphere" size="0.022"/>
</body>
</body>
</body>
<body name="RL_hip" pos="-0.1934 0.0465 0">
<inertial pos="0.0054 0.00194 -0.000105" quat="0.505462 0.498237 0.497014 0.499245" mass="0.678" diaginertia="0.00088403 0.000596003 0.000479967" />
<joint name="RL_hip_joint" axis="1 0 0" range="-1.0472 1.0472"/>
<geom mesh="hip_0" material="metal" class="visual" quat="4.63268e-05 0 1 0" />
<geom mesh="hip_1" material="gray" class="visual" quat="4.63268e-05 0 1 0" />
<geom name="RL_hip_collision" size="0.046 0.02" pos="0 0.08 0" quat="0.707107 0.707107 0 0" type="cylinder" class="collision" />
<body name="RL_thigh" pos="0 0.0955 0">
<inertial pos="-0.00374 -0.0223 -0.0327" quat="0.829533 0.0847635 -0.0200632 0.551623" mass="1.152" diaginertia="0.00594973 0.00584149 0.000878787" />
<joint name="RL_thigh_joint" axis="0 1 0" range="-0.5236 4.5379"/>
<geom mesh="thigh_0" material="metal" class="visual" />
<geom mesh="thigh_1" material="gray" class="visual" />
<geom name="RL_thigh_collision" size="0.1065 0.01225 0.017" pos="0 0 -0.1065" quat="0.707107 0 0.707107 0" type="box" class="collision" />
<body name="RL_calf" pos="0 0 -0.213">
<inertial pos="0.00629595 -0.000622121 -0.141417" quat="0.710672 0.00154099 -0.00450087 0.703508" mass="0.241352" diaginertia="0.0014901 0.00146356 5.31397e-05" />
<joint name="RL_calf_joint" axis="0 1 0" range="-2.7227 -0.83776"/>
<geom mesh="calf_0" material="gray" class="visual" />
<geom mesh="calf_1" material="black" class="visual" />
<geom name="RL_calf1_collision" size="0.013 0.06" pos="0.01 0 -0.06" quat="0.995004 0 -0.0998334 0" type="cylinder" class="collision" />
<geom name="RL_calf2_collision" size="0.011 0.0325" pos="0.02 0 -0.148" quat="0.999688 0 0.0249974 0" type="cylinder" class="collision" />
<geom pos="0 0 -0.213" mesh="foot" class="visual" material="black" />
<geom name="RL_foot_collision" pos="0 0 -0.213" type="sphere" size="0.022" class="collision"/>
<site name="RL" pos="0 0 -0.213" type="sphere" size="0.022"/>
</body>
</body>
</body>
<body name="RR_hip" pos="-0.1934 -0.0465 0">
<inertial pos="0.0054 -0.00194 -0.000105" quat="0.499245 0.497014 0.498237 0.505462" mass="0.678" diaginertia="0.00088403 0.000596003 0.000479967" />
<joint name="RR_hip_joint" axis="1 0 0" range="-1.0472 1.0472"/>
<geom mesh="hip_0" material="metal" class="visual" quat="2.14617e-09 4.63268e-05 4.63268e-05 -1" />
<geom mesh="hip_1" material="gray" class="visual" quat="2.14617e-09 4.63268e-05 4.63268e-05 -1" />
<geom name="RR_hip_collision" size="0.046 0.02" pos="0 -0.08 0" quat="0.707107 0.707107 0 0" type="cylinder" class="collision" />
<body name="RR_thigh" pos="0 -0.0955 0">
<inertial pos="-0.00374 0.0223 -0.0327" quat="0.551623 -0.0200632 0.0847635 0.829533" mass="1.152" diaginertia="0.00594973 0.00584149 0.000878787" />
<joint name="RR_thigh_joint" axis="0 1 0" range="-0.5236 4.5379"/>
<geom mesh="thigh_mirror_0" material="metal" class="visual" />
<geom mesh="thigh_mirror_1" material="gray" class="visual" />
<geom name="RR_thigh_collision" size="0.1065 0.01225 0.017" pos="0 0 -0.1065" quat="0.707107 0 0.707107 0" type="box" class="collision" />
<body name="RR_calf" pos="0 0 -0.213">
<inertial pos="0.00629595 0.000622121 -0.141417" quat="0.703508 -0.00450087 0.00154099 0.710672" mass="0.241352" diaginertia="0.0014901 0.00146356 5.31397e-05" />
<joint name="RR_calf_joint" axis="0 1 0" range="-2.7227 -0.83776"/>
<geom mesh="calf_mirror_0" material="gray" class="visual" />
<geom mesh="calf_mirror_1" material="black" class="visual" />
<geom name="RR_calf1_collision" size="0.013 0.06" pos="0.01 0 -0.06" quat="0.995004 0 -0.0998334 0" type="cylinder" class="collision" />
<geom name="RR_calf2_collision" size="0.011 0.0325" pos="0.02 0 -0.148" quat="0.999688 0 0.0249974 0" type="cylinder" class="collision" />
<geom pos="0 0 -0.213" mesh="foot" class="visual" material="black" />
<geom name="RR_foot_collision" pos="0 0 -0.213" type="sphere" size="0.022" class="collision"/>
<site name="RR" pos="0 0 -0.213" type="sphere" size="0.022"/>
</body>
</body>
</body>
</body>
</worldbody>
<sensor>
<gyro name="imu_ang_vel" site="imu"/>
<velocimeter name="imu_lin_vel" site="imu"/>
<accelerometer name="imu_lin_acc" site="imu"/>
<subtreeangmom name="root_angmom" body="base_link"/>
</sensor>
</mujoco>
@@ -1,306 +0,0 @@
<mujoco model="scene_go2">
<compiler angle="radian" meshdir="assets" autolimits="true" />
<option cone="elliptic" impratio="100" />
<default>
<default class="go2">
<geom friction="0.4" margin="0.001" condim="1"/>
<joint axis="0 1 0" damping="0.1" armature="0.01" frictionloss="0.2"/>
<motor ctrlrange="-23.7 23.7"/>
<default class="abduction">
<joint axis="1 0 0" range="-1.0472 1.0472"/>
</default>
<default class="hip">
<default class="front_hip">
<joint range="-1.5708 3.4907"/>
</default>
<default class="back_hip">
<joint range="-0.5236 4.5379"/>
</default>
</default>
<default class="knee">
<joint range="-2.7227 -0.83776"/>
<motor ctrlrange="-45.43 45.43"/>
</default>
<default class="visual">
<geom type="mesh" contype="0" conaffinity="0" group="2"/>
</default>
<default class="collision">
<geom group="3"/>
<default class="foot">
<geom size="0.022" pos="-0.002 0 -0.213" priority="1" condim="6"
friction="0.4 0.02 0.01"/>
</default>
</default>
</default>
</default>
<asset>
<material name="metal" rgba=".9 .95 .95 1" />
<material name="black" rgba="0 0 0 1" />
<material name="white" rgba="1 1 1 1" />
<material name="gray" rgba="0.671705 0.692426 0.774270 1" />
<mesh file="base_0.obj" />
<mesh file="base_1.obj" />
<mesh file="base_2.obj" />
<mesh file="base_3.obj" />
<mesh file="base_4.obj" />
<mesh file="hip_0.obj" />
<mesh file="hip_1.obj" />
<mesh file="thigh_0.obj" />
<mesh file="thigh_1.obj" />
<mesh file="thigh_mirror_0.obj" />
<mesh file="thigh_mirror_1.obj" />
<mesh file="calf_0.obj" />
<mesh file="calf_1.obj" />
<mesh file="calf_mirror_0.obj" />
<mesh file="calf_mirror_1.obj" />
<mesh file="foot.obj" />
</asset>
<worldbody>
<body name="base_link" pos="0 0 0.445" childclass="go2">
<inertial pos="0.021112 0 -0.005366" quat="-0.000543471 0.713435 -0.00173769 0.700719"
mass="6.921"
diaginertia="0.107027 0.0980771 0.0244531" />
<freejoint />
<geom mesh="base_0" material="black" class="visual" />
<geom mesh="base_1" material="black" class="visual" />
<geom mesh="base_2" material="black" class="visual" />
<geom mesh="base_3" material="white" class="visual" />
<geom mesh="base_4" material="gray" class="visual" />
<geom size="0.1881 0.04675 0.057" type="box" class="collision" />
<geom size="0.05 0.045" pos="0.285 0 0.01" type="cylinder" class="collision" />
<geom size="0.047" pos="0.293 0 -0.06" class="collision" />
<site name="imu" pos="-0.02557 0 0.04232" />
<body name="FL_hip" pos="0.1934 0.0465 0">
<inertial pos="-0.0054 0.00194 -0.000105" quat="0.497014 0.499245 0.505462 0.498237"
mass="0.678"
diaginertia="0.00088403 0.000596003 0.000479967" />
<joint name="FL_hip_joint" class="abduction" />
<geom mesh="hip_0" material="metal" class="visual" />
<geom mesh="hip_1" material="gray" class="visual" />
<geom size="0.046 0.02" pos="0 0.08 0" quat="1 1 0 0" type="cylinder" class="collision" />
<body name="FL_thigh" pos="0 0.0955 0">
<inertial pos="-0.00374 -0.0223 -0.0327" quat="0.829533 0.0847635 -0.0200632 0.551623"
mass="1.152"
diaginertia="0.00594973 0.00584149 0.000878787" />
<joint name="FL_thigh_joint" class="front_hip" />
<geom mesh="thigh_0" material="metal" class="visual" />
<geom mesh="thigh_1" material="gray" class="visual" />
<geom size="0.1065 0.01225 0.017" pos="0 0 -0.1065" quat="0.707107 0 0.707107 0"
type="box" class="collision" />
<body name="FL_calf" pos="0 0 -0.213">
<inertial pos="0.00629595 -0.000622121 -0.141417"
quat="0.710672 0.00154099 -0.00450087 0.703508"
mass="0.241352" diaginertia="0.0014901 0.00146356 5.31397e-05" />
<joint name="FL_calf_joint" class="knee" />
<geom mesh="calf_0" material="gray" class="visual" />
<geom mesh="calf_1" material="black" class="visual" />
<geom size="0.012 0.06" pos="0.008 0 -0.06" quat="0.994493 0 -0.104807 0"
type="cylinder" class="collision" />
<geom size="0.011 0.0325" pos="0.02 0 -0.148" quat="0.999688 0 0.0249974 0"
type="cylinder" class="collision" />
<geom pos="0 0 -0.213" mesh="foot" class="visual" material="black" />
<geom name="FL" class="foot" />
<body name="FL_foot" pos="0 0 -0.213" />
</body>
</body>
</body>
<body name="FR_hip" pos="0.1934 -0.0465 0">
<inertial pos="-0.0054 -0.00194 -0.000105" quat="0.498237 0.505462 0.499245 0.497014"
mass="0.678"
diaginertia="0.00088403 0.000596003 0.000479967" />
<joint name="FR_hip_joint" class="abduction" />
<geom mesh="hip_0" material="metal" class="visual" quat="4.63268e-05 1 0 0" />
<geom mesh="hip_1" material="gray" class="visual" quat="4.63268e-05 1 0 0" />
<geom size="0.046 0.02" pos="0 -0.08 0" quat="0.707107 0.707107 0 0" type="cylinder"
class="collision" />
<body name="FR_thigh" pos="0 -0.0955 0">
<inertial pos="-0.00374 0.0223 -0.0327" quat="0.551623 -0.0200632 0.0847635 0.829533"
mass="1.152"
diaginertia="0.00594973 0.00584149 0.000878787" />
<joint name="FR_thigh_joint" class="front_hip" />
<geom mesh="thigh_mirror_0" material="metal" class="visual" />
<geom mesh="thigh_mirror_1" material="gray" class="visual" />
<geom size="0.1065 0.01225 0.017" pos="0 0 -0.1065" quat="0.707107 0 0.707107 0"
type="box" class="collision" />
<body name="FR_calf" pos="0 0 -0.213">
<inertial pos="0.00629595 0.000622121 -0.141417"
quat="0.703508 -0.00450087 0.00154099 0.710672"
mass="0.241352" diaginertia="0.0014901 0.00146356 5.31397e-05" />
<joint name="FR_calf_joint" class="knee" />
<geom mesh="calf_mirror_0" material="gray" class="visual" />
<geom mesh="calf_mirror_1" material="black" class="visual" />
<geom size="0.013 0.06" pos="0.01 0 -0.06" quat="0.995004 0 -0.0998334 0"
type="cylinder" class="collision" />
<geom size="0.011 0.0325" pos="0.02 0 -0.148" quat="0.999688 0 0.0249974 0"
type="cylinder" class="collision" />
<geom pos="0 0 -0.213" mesh="foot" class="visual" material="black" />
<geom name="FR" class="foot" />
<body name="FR_foot" pos="0 0 -0.213" />
</body>
</body>
</body>
<body name="RL_hip" pos="-0.1934 0.0465 0">
<inertial pos="0.0054 0.00194 -0.000105" quat="0.505462 0.498237 0.497014 0.499245"
mass="0.678"
diaginertia="0.00088403 0.000596003 0.000479967" />
<joint name="RL_hip_joint" class="abduction" />
<geom mesh="hip_0" material="metal" class="visual" quat="4.63268e-05 0 1 0" />
<geom mesh="hip_1" material="gray" class="visual" quat="4.63268e-05 0 1 0" />
<geom size="0.046 0.02" pos="0 0.08 0" quat="0.707107 0.707107 0 0" type="cylinder"
class="collision" />
<body name="RL_thigh" pos="0 0.0955 0">
<inertial pos="-0.00374 -0.0223 -0.0327" quat="0.829533 0.0847635 -0.0200632 0.551623"
mass="1.152"
diaginertia="0.00594973 0.00584149 0.000878787" />
<joint name="RL_thigh_joint" class="back_hip" />
<geom mesh="thigh_0" material="metal" class="visual" />
<geom mesh="thigh_1" material="gray" class="visual" />
<geom size="0.1065 0.01225 0.017" pos="0 0 -0.1065" quat="0.707107 0 0.707107 0"
type="box" class="collision" />
<body name="RL_calf" pos="0 0 -0.213">
<inertial pos="0.00629595 -0.000622121 -0.141417"
quat="0.710672 0.00154099 -0.00450087 0.703508"
mass="0.241352" diaginertia="0.0014901 0.00146356 5.31397e-05" />
<joint name="RL_calf_joint" class="knee" />
<geom mesh="calf_0" material="gray" class="visual" />
<geom mesh="calf_1" material="black" class="visual" />
<geom size="0.013 0.06" pos="0.01 0 -0.06" quat="0.995004 0 -0.0998334 0"
type="cylinder" class="collision" />
<geom size="0.011 0.0325" pos="0.02 0 -0.148" quat="0.999688 0 0.0249974 0"
type="cylinder" class="collision" />
<geom pos="0 0 -0.213" mesh="foot" class="visual" material="black" />
<geom name="RL" class="foot" />
<body name="RL_foot" pos="0 0 -0.213" />
</body>
</body>
</body>
<body name="RR_hip" pos="-0.1934 -0.0465 0">
<inertial pos="0.0054 -0.00194 -0.000105" quat="0.499245 0.497014 0.498237 0.505462"
mass="0.678"
diaginertia="0.00088403 0.000596003 0.000479967" />
<joint name="RR_hip_joint" class="abduction" />
<geom mesh="hip_0" material="metal" class="visual"
quat="2.14617e-09 4.63268e-05 4.63268e-05 -1" />
<geom mesh="hip_1" material="gray" class="visual"
quat="2.14617e-09 4.63268e-05 4.63268e-05 -1" />
<geom size="0.046 0.02" pos="0 -0.08 0" quat="0.707107 0.707107 0 0" type="cylinder"
class="collision" />
<body name="RR_thigh" pos="0 -0.0955 0">
<inertial pos="-0.00374 0.0223 -0.0327" quat="0.551623 -0.0200632 0.0847635 0.829533"
mass="1.152"
diaginertia="0.00594973 0.00584149 0.000878787" />
<joint name="RR_thigh_joint" class="back_hip" />
<geom mesh="thigh_mirror_0" material="metal" class="visual" />
<geom mesh="thigh_mirror_1" material="gray" class="visual" />
<geom size="0.1065 0.01225 0.017" pos="0 0 -0.1065" quat="0.707107 0 0.707107 0"
type="box" class="collision" />
<body name="RR_calf" pos="0 0 -0.213">
<inertial pos="0.00629595 0.000622121 -0.141417"
quat="0.703508 -0.00450087 0.00154099 0.710672"
mass="0.241352" diaginertia="0.0014901 0.00146356 5.31397e-05" />
<joint name="RR_calf_joint" class="knee" />
<geom mesh="calf_mirror_0" material="gray" class="visual" />
<geom mesh="calf_mirror_1" material="black" class="visual" />
<geom size="0.013 0.06" pos="0.01 0 -0.06" quat="0.995004 0 -0.0998334 0"
type="cylinder" class="collision" />
<geom size="0.011 0.0325" pos="0.02 0 -0.148" quat="0.999688 0 0.0249974 0"
type="cylinder" class="collision" />
<geom pos="0 0 -0.213" mesh="foot" class="visual" material="black" />
<geom name="RR" class="foot" />
<body name="RR_foot" pos="0 0 -0.213" />
</body>
</body>
</body>
</body>
</worldbody>
<actuator>
<motor class="abduction" name="FR_hip" joint="FR_hip_joint" />
<motor class="hip" name="FR_thigh" joint="FR_thigh_joint" />
<motor class="knee" name="FR_calf" joint="FR_calf_joint" />
<motor class="abduction" name="FL_hip" joint="FL_hip_joint" />
<motor class="hip" name="FL_thigh" joint="FL_thigh_joint" />
<motor class="knee" name="FL_calf" joint="FL_calf_joint" />
<motor class="abduction" name="RR_hip" joint="RR_hip_joint" />
<motor class="hip" name="RR_thigh" joint="RR_thigh_joint" />
<motor class="knee" name="RR_calf" joint="RR_calf_joint" />
<motor class="abduction" name="RL_hip" joint="RL_hip_joint" />
<motor class="hip" name="RL_thigh" joint="RL_thigh_joint" />
<motor class="knee" name="RL_calf" joint="RL_calf_joint" />
</actuator>
<sensor>
<jointpos name="FR_hip_pos" joint="FR_hip_joint" />
<jointpos name="FR_thigh_pos" joint="FR_thigh_joint" />
<jointpos name="FR_calf_pos" joint="FR_calf_joint" />
<jointpos name="FL_hip_pos" joint="FL_hip_joint" />
<jointpos name="FL_thigh_pos" joint="FL_thigh_joint" />
<jointpos name="FL_calf_pos" joint="FL_calf_joint" />
<jointpos name="RR_hip_pos" joint="RR_hip_joint" />
<jointpos name="RR_thigh_pos" joint="RR_thigh_joint" />
<jointpos name="RR_calf_pos" joint="RR_calf_joint" />
<jointpos name="RL_hip_pos" joint="RL_hip_joint" />
<jointpos name="RL_thigh_pos" joint="RL_thigh_joint" />
<jointpos name="RL_calf_pos" joint="RL_calf_joint" />
<jointvel name="FR_hip_vel" joint="FR_hip_joint" />
<jointvel name="FR_thigh_vel" joint="FR_thigh_joint" />
<jointvel name="FR_calf_vel" joint="FR_calf_joint" />
<jointvel name="FL_hip_vel" joint="FL_hip_joint" />
<jointvel name="FL_thigh_vel" joint="FL_thigh_joint" />
<jointvel name="FL_calf_vel" joint="FL_calf_joint" />
<jointvel name="RR_hip_vel" joint="RR_hip_joint" />
<jointvel name="RR_thigh_vel" joint="RR_thigh_joint" />
<jointvel name="RR_calf_vel" joint="RR_calf_joint" />
<jointvel name="RL_hip_vel" joint="RL_hip_joint" />
<jointvel name="RL_thigh_vel" joint="RL_thigh_joint" />
<jointvel name="RL_calf_vel" joint="RL_calf_joint" />
<jointactuatorfrc name="FR_hip_torque" joint="FR_hip_joint" noise="0.01" />
<jointactuatorfrc name="FR_thigh_torque" joint="FR_thigh_joint" noise="0.01" />
<jointactuatorfrc name="FR_calf_torque" joint="FR_calf_joint" noise="0.01" />
<jointactuatorfrc name="FL_hip_torque" joint="FL_hip_joint" noise="0.01" />
<jointactuatorfrc name="FL_thigh_torque" joint="FL_thigh_joint" noise="0.01" />
<jointactuatorfrc name="FL_calf_torque" joint="FL_calf_joint" noise="0.01" />
<jointactuatorfrc name="RR_hip_torque" joint="RR_hip_joint" noise="0.01" />
<jointactuatorfrc name="RR_thigh_torque" joint="RR_thigh_joint" noise="0.01" />
<jointactuatorfrc name="RR_calf_torque" joint="RR_calf_joint" noise="0.01" />
<jointactuatorfrc name="RL_hip_torque" joint="RL_hip_joint" noise="0.01" />
<jointactuatorfrc name="RL_thigh_torque" joint="RL_thigh_joint" noise="0.01" />
<jointactuatorfrc name="RL_calf_torque" joint="RL_calf_joint" noise="0.01" />
<framequat name="imu_quat" objtype="site" objname="imu" />
<gyro name="imu_gyro" site="imu" />
<accelerometer name="imu_acc" site="imu" />
<framepos name="frame_pos" objtype="site" objname="imu" />
<framelinvel name="frame_vel" objtype="site" objname="imu" />
</sensor>
<keyframe>
<key name="home" qpos="0 0 0.27 1 0 0 0 0 0.9 -1.8 0 0.9 -1.8 0 0.9 -1.8 0 0.9 -1.8"
ctrl="0 0.9 -1.8 0 0.9 -1.8 0 0.9 -1.8 0 0.9 -1.8" />
</keyframe>
<!-- setup scene -->
<statistic center="1.0 1.0 1.0" extent="0.8"/>
<visual>
<headlight diffuse="0.6 0.6 0.6" ambient="0.1 0.1 0.1" specular="0.9 0.9 0.9"/>
<rgba haze="0.15 0.25 0.35 1"/>
<global azimuth="-140" elevation="-20"/>
</visual>
<asset>
<texture type="skybox" builtin="flat" rgb1="0 0 0" rgb2="0 0 0" width="512" height="3072"/>
<texture type="2d" name="groundplane" builtin="checker" mark="edge" rgb1="0.2 0.3 0.4" rgb2="0.1 0.2 0.3" markrgb="0.8 0.8 0.8" width="300" height="300"/>
<material name="groundplane" texture="groundplane" texuniform="true" texrepeat="5 5" reflectance="0.2"/>
</asset>
<worldbody>
<light pos="1 0 3.5" dir="0 0 -1" directional="true"/>
<geom name="floor" size="0 0 0.05" type="plane" material="groundplane"/>
</worldbody>
</mujoco>
-5
View File
@@ -1,5 +0,0 @@
from mjlab.utils.lab_api.tasks.importer import import_packages
_BLACKLIST_PKGS = ["utils", ".mdp"]
import_packages(__name__, _BLACKLIST_PKGS)
@@ -1 +0,0 @@
"""Velocity tracking environments for legged robots."""
@@ -1,24 +0,0 @@
from mjlab.tasks.registry import register_mjlab_task
from src.tasks.velocity.rl import VelocityOnPolicyRunner
from .env_cfgs import (
unitree_go2_flat_env_cfg,
unitree_go2_rough_env_cfg,
)
from .rl_cfg import unitree_go2_ppo_runner_cfg
register_mjlab_task(
task_id="Unitree-Go2-Rough",
env_cfg=unitree_go2_rough_env_cfg(),
play_env_cfg=unitree_go2_rough_env_cfg(play=True),
rl_cfg=unitree_go2_ppo_runner_cfg(),
runner_cls=VelocityOnPolicyRunner,
)
register_mjlab_task(
task_id="Unitree-Go2-Flat",
env_cfg=unitree_go2_flat_env_cfg(),
play_env_cfg=unitree_go2_flat_env_cfg(play=True),
rl_cfg=unitree_go2_ppo_runner_cfg(),
runner_cls=VelocityOnPolicyRunner,
)
@@ -1,171 +0,0 @@
"""Unitree Go2 velocity environment configurations."""
from typing import Literal
from src.assets.robots import (
get_go2_robot_cfg,
)
from mjlab.envs import ManagerBasedRlEnvCfg
from mjlab.envs import mdp as envs_mdp
from mjlab.envs.mdp.actions import JointPositionActionCfg
from mjlab.managers import TerminationTermCfg
from mjlab.managers.event_manager import EventTermCfg
from mjlab.sensor import ContactMatch, ContactSensorCfg, RayCastSensorCfg
from mjlab.tasks.velocity import mdp
from mjlab.tasks.velocity.mdp import UniformVelocityCommandCfg
from src.tasks.velocity.velocity_env_cfg import make_velocity_env_cfg
TerrainType = Literal["rough", "obstacles"]
def unitree_go2_rough_env_cfg(
play: bool = False,
) -> ManagerBasedRlEnvCfg:
"""Create Unitree Go2 rough terrain velocity configuration."""
cfg = make_velocity_env_cfg()
cfg.sim.mujoco.ccd_iterations = 500
cfg.sim.contact_sensor_maxmatch = 500
cfg.scene.entities = {"robot": get_go2_robot_cfg()}
# Set raycast sensor frame to Go2 base_link.
for sensor in cfg.scene.sensors or ():
if sensor.name == "terrain_scan":
assert isinstance(sensor, RayCastSensorCfg)
sensor.frame.name = "base_link"
foot_names = ("FR", "FL", "RR", "RL")
site_names = ("FR", "FL", "RR", "RL")
geom_names = tuple(f"{name}_foot_collision" for name in foot_names)
feet_ground_cfg = ContactSensorCfg(
name="feet_ground_contact",
primary=ContactMatch(mode="geom", pattern=geom_names, entity="robot"),
secondary=ContactMatch(mode="body", pattern="terrain"),
fields=("found", "force"),
reduce="netforce",
num_slots=1,
track_air_time=True,
)
nonfoot_ground_cfg = ContactSensorCfg(
name="nonfoot_ground_touch",
primary=ContactMatch(
mode="geom",
entity="robot",
# Grab all collision geoms...
pattern=r".*_collision\d*$",
# Except for the foot geoms.
exclude=tuple(geom_names),
),
secondary=ContactMatch(mode="body", pattern="terrain"),
fields=("found", "force"),
reduce="none",
num_slots=1,
history_length=4,
)
cfg.scene.sensors = (cfg.scene.sensors or ()) + (
feet_ground_cfg,
nonfoot_ground_cfg,
)
if cfg.scene.terrain is not None and cfg.scene.terrain.terrain_generator is not None:
cfg.scene.terrain.terrain_generator.curriculum = True
joint_pos_action = cfg.actions["joint_pos"]
assert isinstance(joint_pos_action, JointPositionActionCfg)
cfg.viewer.body_name = "base_link"
cfg.viewer.distance = 1.5
cfg.viewer.elevation = -10.0
cfg.observations["critic"].terms["foot_height"].params["asset_cfg"].site_names = site_names
cfg.events["foot_friction"].params["asset_cfg"].geom_names = geom_names
cfg.events["base_com"].params["asset_cfg"].body_names = ("base_link",)
cfg.rewards["pose"].params["std_standing"] = {
r".*(FR|FL|RR|RL)_hip_joint.*": 0.05,
r".*(FR|FL|RR|RL)_thigh_joint.*": 0.1,
r".*(FR|FL|RR|RL)_calf_joint.*": 0.15,
}
cfg.rewards["pose"].params["std_walking"] = {
r".*(FR|FL|RR|RL)_hip_joint.*": 0.15,
r".*(FR|FL|RR|RL)_thigh_joint.*": 0.35,
r".*(FR|FL|RR|RL)_calf_joint.*": 0.5,
}
cfg.rewards["pose"].params["std_running"] = {
r".*(FR|FL|RR|RL)_hip_joint.*": 0.15,
r".*(FR|FL|RR|RL)_thigh_joint.*": 0.35,
r".*(FR|FL|RR|RL)_calf_joint.*": 0.5,
}
cfg.rewards["foot_gait"].params["offset"] = [0.0, 0.5, 0.5, 0.0]
cfg.rewards["body_orientation_l2"].params["asset_cfg"].body_names = ("base_link",)
cfg.rewards["body_ang_vel"].params["asset_cfg"].body_names = ("base_link",)
cfg.rewards["foot_clearance"].params["asset_cfg"].site_names = site_names
cfg.rewards["foot_slip"].params["asset_cfg"].site_names = site_names
cfg.metrics["slip_velocity"].params["asset_cfg"].site_names = site_names
cfg.terminations["illegal_contact"] = TerminationTermCfg(
func=mdp.illegal_contact,
params={"sensor_name": nonfoot_ground_cfg.name, "force_threshold": 10.0},
)
# Apply play mode overrides.
if play:
# Effectively infinite episode length.
cfg.episode_length_s = int(1e9)
cfg.observations["actor"].enable_corruption = False
cfg.events.pop("push_robot", None)
cfg.curriculum = {}
cfg.events["randomize_terrain"] = EventTermCfg(
func=envs_mdp.randomize_terrain,
mode="reset",
params={},
)
if cfg.scene.terrain is not None:
if cfg.scene.terrain.terrain_generator is not None:
cfg.scene.terrain.terrain_generator.curriculum = False
cfg.scene.terrain.terrain_generator.num_cols = 5
cfg.scene.terrain.terrain_generator.num_rows = 5
cfg.scene.terrain.terrain_generator.border_width = 10.0
return cfg
def unitree_go2_flat_env_cfg(play: bool = False) -> ManagerBasedRlEnvCfg:
"""Create Unitree Go2 flat terrain velocity configuration."""
cfg = unitree_go2_rough_env_cfg(play=play)
cfg.sim.njmax = 300
cfg.sim.mujoco.ccd_iterations = 50
cfg.sim.contact_sensor_maxmatch = 64
cfg.sim.nconmax = None
# Switch to flat terrain.
assert cfg.scene.terrain is not None
cfg.scene.terrain.terrain_type = "plane"
cfg.scene.terrain.terrain_generator = None
# Remove raycast sensor and height scan (no terrain to scan).
cfg.scene.sensors = tuple(
s for s in (cfg.scene.sensors or ()) if s.name != "terrain_scan"
)
del cfg.observations["actor"].terms["height_scan"]
del cfg.observations["critic"].terms["height_scan"]
# Disable terrain curriculum (not present in play mode since rough clears all).
cfg.curriculum.pop("terrain_levels", None)
if play:
twist_cmd = cfg.commands["twist"]
assert isinstance(twist_cmd, UniformVelocityCommandCfg)
twist_cmd.ranges.lin_vel_x = (-0.5, 1.0)
twist_cmd.ranges.lin_vel_y = (-0.5, 0.5)
twist_cmd.ranges.ang_vel_z = (-0.5, 0.5)
return cfg
@@ -1,46 +0,0 @@
"""RL configuration for Unitree Go2 velocity task."""
from mjlab.rl import (
RslRlModelCfg,
RslRlOnPolicyRunnerCfg,
RslRlPpoAlgorithmCfg,
)
def unitree_go2_ppo_runner_cfg() -> RslRlOnPolicyRunnerCfg:
"""Create RL runner configuration for Unitree Go2 velocity task."""
return RslRlOnPolicyRunnerCfg(
actor=RslRlModelCfg(
hidden_dims=(512, 256, 128),
activation="elu",
obs_normalization=True,
distribution_cfg={
"class_name": "GaussianDistribution",
"init_std": 1.0,
"std_type": "scalar",
},
),
critic=RslRlModelCfg(
hidden_dims=(512, 256, 128),
activation="elu",
obs_normalization=True,
),
algorithm=RslRlPpoAlgorithmCfg(
value_loss_coef=1.0,
use_clipped_value_loss=True,
clip_param=0.2,
entropy_coef=0.01,
num_learning_epochs=5,
num_mini_batches=4,
learning_rate=1.0e-3,
schedule="adaptive",
gamma=0.99,
lam=0.95,
desired_kl=0.01,
max_grad_norm=1.0,
),
experiment_name="go2_velocity",
save_interval=100,
num_steps_per_env=24,
max_iterations=10001,
)
@@ -1,8 +0,0 @@
from mjlab.envs.mdp import * # noqa: F401, F403
from .curriculums import * # noqa: F403
from .metrics import * # noqa: F403
from .observations import * # noqa: F403
from .rewards import * # noqa: F403
from .terminations import * # noqa: F403
from .velocity_command import * # noqa: F403
@@ -1,107 +0,0 @@
from __future__ import annotations
from typing import TYPE_CHECKING, TypedDict, cast
import torch
from mjlab.entity import Entity
from mjlab.managers.scene_entity_config import SceneEntityCfg
from .velocity_command import UniformVelocityCommandCfg
if TYPE_CHECKING:
from mjlab.envs import ManagerBasedRlEnv
_DEFAULT_SCENE_CFG = SceneEntityCfg("robot")
class VelocityStage(TypedDict):
step: int
lin_vel_x: tuple[float, float] | None
lin_vel_y: tuple[float, float] | None
ang_vel_z: tuple[float, float] | None
class RewardWeightStage(TypedDict):
step: int
weight: float
def terrain_levels_vel(
env: ManagerBasedRlEnv,
env_ids: torch.Tensor,
command_name: str,
asset_cfg: SceneEntityCfg = _DEFAULT_SCENE_CFG,
) -> torch.Tensor:
asset: Entity = env.scene[asset_cfg.name]
terrain = env.scene.terrain
assert terrain is not None
terrain_generator = terrain.cfg.terrain_generator
assert terrain_generator is not None
command = env.command_manager.get_command(command_name)
assert command is not None
# Compute the distance the robot walked.
distance = torch.norm(
asset.data.root_link_pos_w[env_ids, :2] - env.scene.env_origins[env_ids, :2], dim=1
)
# Robots that walked far enough progress to harder terrains.
move_up = distance > terrain_generator.size[0] / 2
# Robots that walked less than half of their required distance go to simpler
# terrains.
move_down = (
distance < torch.norm(command[env_ids, :2], dim=1) * env.max_episode_length_s * 0.5
)
move_down *= ~move_up
# Update terrain levels.
terrain.update_env_origins(env_ids, move_up, move_down)
return torch.mean(terrain.terrain_levels.float())
def commands_vel(
env: ManagerBasedRlEnv,
env_ids: torch.Tensor,
command_name: str,
velocity_stages: list[VelocityStage],
) -> dict[str, torch.Tensor]:
del env_ids # Unused.
command_term = env.command_manager.get_term(command_name)
assert command_term is not None
cfg = cast(UniformVelocityCommandCfg, command_term.cfg)
for stage in velocity_stages:
if env.common_step_counter > stage["step"]:
if "lin_vel_x" in stage and stage["lin_vel_x"] is not None:
cfg.ranges.lin_vel_x = stage["lin_vel_x"]
if "lin_vel_y" in stage and stage["lin_vel_y"] is not None:
cfg.ranges.lin_vel_y = stage["lin_vel_y"]
if "ang_vel_z" in stage and stage["ang_vel_z"] is not None:
cfg.ranges.ang_vel_z = stage["ang_vel_z"]
return {
# "lin_vel_x_min": torch.tensor(cfg.ranges.lin_vel_x[0]),
# "lin_vel_x_max": torch.tensor(cfg.ranges.lin_vel_x[1]),
# "lin_vel_y_min": torch.tensor(cfg.ranges.lin_vel_y[0]),
# "lin_vel_y_max": torch.tensor(cfg.ranges.lin_vel_y[1]),
# "ang_vel_z_min": torch.tensor(cfg.ranges.ang_vel_z[0]),
# "ang_vel_z_max": torch.tensor(cfg.ranges.ang_vel_z[1]),
}
def reward_weight(
env: ManagerBasedRlEnv,
env_ids: torch.Tensor,
reward_name: str,
weight_stages: list[RewardWeightStage],
) -> torch.Tensor:
"""Update a reward term's weight based on training step stages."""
del env_ids # Unused.
reward_term_cfg = env.reward_manager.get_term_cfg(reward_name)
for stage in weight_stages:
if env.common_step_counter > stage["step"]:
reward_term_cfg.weight = stage["weight"]
return torch.tensor([reward_term_cfg.weight])
@@ -1,88 +0,0 @@
"""Reward-weight-independent quality metrics for velocity tasks."""
from __future__ import annotations
from typing import TYPE_CHECKING
import torch
from mjlab.entity import Entity
from mjlab.managers.scene_entity_config import SceneEntityCfg
from mjlab.sensor import ContactSensor
if TYPE_CHECKING:
from mjlab.envs import ManagerBasedRlEnv
_DEFAULT_ASSET_CFG = SceneEntityCfg("robot")
def linear_velocity_rmse(
env: ManagerBasedRlEnv,
command_name: str,
asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG,
) -> torch.Tensor:
"""Per-step commanded-vs-actual base linear velocity RMSE in body frame."""
asset: Entity = env.scene[asset_cfg.name]
command = env.command_manager.get_command(command_name)
assert command is not None
actual = asset.data.root_link_lin_vel_b
error = torch.cat((command[:, :2] - actual[:, :2], -actual[:, 2:3]), dim=1)
return torch.sqrt(torch.mean(torch.square(error), dim=1))
def angular_velocity_rmse(
env: ManagerBasedRlEnv,
command_name: str,
asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG,
) -> torch.Tensor:
"""Per-step commanded-vs-actual base angular velocity RMSE in body frame."""
asset: Entity = env.scene[asset_cfg.name]
command = env.command_manager.get_command(command_name)
assert command is not None
actual = asset.data.root_link_ang_vel_b
desired = torch.zeros_like(actual)
desired[:, 2] = command[:, 2]
return torch.sqrt(torch.mean(torch.square(desired - actual), dim=1))
def orientation_error(
env: ManagerBasedRlEnv,
asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG,
) -> torch.Tensor:
"""Magnitude of projected gravity in the base x/y plane; zero is upright."""
asset: Entity = env.scene[asset_cfg.name]
return torch.linalg.vector_norm(asset.data.projected_gravity_b[:, :2], dim=1)
def fall_indicator(env: ManagerBasedRlEnv) -> torch.Tensor:
"""One on non-timeout terminal steps, otherwise zero."""
return env.termination_manager.terminated.float()
def slip_velocity(
env: ManagerBasedRlEnv,
sensor_name: str,
asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG,
) -> torch.Tensor:
"""Mean x/y velocity of feet currently touching the ground."""
asset: Entity = env.scene[asset_cfg.name]
sensor: ContactSensor = env.scene[sensor_name]
assert sensor.data.found is not None
contact = (sensor.data.found > 0).float()
speed = torch.linalg.vector_norm(
asset.data.site_lin_vel_w[:, asset_cfg.site_ids, :2], dim=-1
)
count = torch.clamp(torch.sum(contact, dim=1), min=1.0)
return torch.sum(speed * contact, dim=1) / count
def mechanical_power(
env: ManagerBasedRlEnv,
asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG,
) -> torch.Tensor:
"""Positive actuator mechanical power in watts (regeneration is ignored)."""
asset: Entity = env.scene[asset_cfg.name]
torque = asset.data.actuator_force[:, asset_cfg.actuator_ids]
velocity = asset.data.joint_vel[:, asset_cfg.joint_ids]
count = min(torque.shape[1], velocity.shape[1])
return torch.sum(torch.clamp(torque[:, :count] * velocity[:, :count], min=0.0), dim=1)
@@ -1,54 +0,0 @@
from __future__ import annotations
from typing import TYPE_CHECKING
import torch
from mjlab.entity import Entity
from mjlab.managers.scene_entity_config import SceneEntityCfg
from mjlab.sensor import ContactSensor
if TYPE_CHECKING:
from mjlab.envs import ManagerBasedRlEnv
_DEFAULT_ASSET_CFG = SceneEntityCfg("robot")
def foot_height(
env: ManagerBasedRlEnv, asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG
) -> torch.Tensor:
asset: Entity = env.scene[asset_cfg.name]
return asset.data.site_pos_w[:, asset_cfg.site_ids, 2] # (num_envs, num_sites)
def foot_air_time(env: ManagerBasedRlEnv, sensor_name: str) -> torch.Tensor:
sensor: ContactSensor = env.scene[sensor_name]
sensor_data = sensor.data
current_air_time = sensor_data.current_air_time
assert current_air_time is not None
return current_air_time
def foot_contact(env: ManagerBasedRlEnv, sensor_name: str) -> torch.Tensor:
sensor: ContactSensor = env.scene[sensor_name]
sensor_data = sensor.data
assert sensor_data.found is not None
return (sensor_data.found > 0).float()
def foot_contact_forces(env: ManagerBasedRlEnv, sensor_name: str) -> torch.Tensor:
sensor: ContactSensor = env.scene[sensor_name]
sensor_data = sensor.data
assert sensor_data.force is not None
forces_flat = sensor_data.force.flatten(start_dim=1) # [B, N*3]
return torch.sign(forces_flat) * torch.log1p(torch.abs(forces_flat))
def phase(env: ManagerBasedRlEnv, period: float, command_name: str) -> torch.Tensor:
global_phase = (env.episode_length_buf * env.step_dt) % period / period
phase = torch.zeros(env.num_envs, 2, device=env.device)
phase[:, 0] = torch.sin(global_phase * torch.pi * 2.0)
phase[:, 1] = torch.cos(global_phase * torch.pi * 2.0)
stand_mask = torch.linalg.norm(env.command_manager.get_command(command_name), dim=1) < 0.1
phase = torch.where(stand_mask.unsqueeze(1), torch.zeros_like(phase), phase)
return phase
@@ -1,427 +0,0 @@
from __future__ import annotations
from typing import TYPE_CHECKING
import torch
from mjlab.entity import Entity
from mjlab.managers.reward_manager import RewardTermCfg
from mjlab.managers.scene_entity_config import SceneEntityCfg
from mjlab.sensor import BuiltinSensor, ContactSensor
from mjlab.utils.lab_api.math import quat_apply_inverse
from mjlab.utils.lab_api.string import (
resolve_matching_names_values,
)
if TYPE_CHECKING:
from mjlab.envs import ManagerBasedRlEnv
_DEFAULT_ASSET_CFG = SceneEntityCfg("robot")
def track_linear_velocity(
env: ManagerBasedRlEnv,
std: float,
command_name: str,
asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG,
) -> torch.Tensor:
"""Reward for tracking the commanded base linear velocity.
The commanded z velocity is assumed to be zero.
"""
asset: Entity = env.scene[asset_cfg.name]
command = env.command_manager.get_command(command_name)
assert command is not None, f"Command '{command_name}' not found."
actual = asset.data.root_link_lin_vel_b
xy_error = torch.sum(torch.square(command[:, :2] - actual[:, :2]), dim=1)
z_error = torch.square(actual[:, 2])
lin_vel_error = xy_error + (2 * z_error)
return torch.exp(-lin_vel_error / std**2)
def track_angular_velocity(
env: ManagerBasedRlEnv,
std: float,
command_name: str,
asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG,
) -> torch.Tensor:
"""Reward heading error for heading-controlled envs, angular velocity for others.
The commanded xy angular velocities are assumed to be zero.
"""
asset: Entity = env.scene[asset_cfg.name]
command = env.command_manager.get_command(command_name)
assert command is not None, f"Command '{command_name}' not found."
actual = asset.data.root_link_ang_vel_b
z_error = torch.square(command[:, 2] - actual[:, 2])
xy_error = torch.sum(torch.square(actual[:, :2]), dim=1)
ang_vel_error = z_error + (0.05 * xy_error)
return torch.exp(-ang_vel_error / std**2)
def body_orientation_l2(
env: ManagerBasedRlEnv,
asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG,
) -> torch.Tensor:
"""Reward flat base orientation (robot being upright).
If asset_cfg has body_ids specified, computes the projected gravity
for that specific body. Otherwise, uses the root link projected gravity.
"""
asset: Entity = env.scene[asset_cfg.name]
# If body_ids are specified, compute projected gravity for that body.
if asset_cfg.body_ids:
body_quat_w = asset.data.body_link_quat_w[:, asset_cfg.body_ids, :] # [B, N, 4]
body_quat_w = body_quat_w.squeeze(1) # [B, 4]
gravity_w = asset.data.gravity_vec_w # [3]
projected_gravity_b = quat_apply_inverse(body_quat_w, gravity_w) # [B, 3]
xy_squared = torch.sum(torch.square(projected_gravity_b[:, :2]), dim=1)
else:
# Use root link projected gravity.
xy_squared = torch.sum(torch.square(asset.data.projected_gravity_b[:, :2]), dim=1)
return xy_squared
def self_collision_cost(
env: ManagerBasedRlEnv,
sensor_name: str,
force_threshold: float = 10.0,
) -> torch.Tensor:
"""Penalize self-collisions.
When the sensor provides force history (from ``history_length > 0``),
counts substeps where any contact force exceeds *force_threshold*.
Falls back to the instantaneous ``found`` count otherwise.
"""
sensor: ContactSensor = env.scene[sensor_name]
data = sensor.data
if data.force_history is not None:
# force_history: [B, N, H, 3]
force_mag = torch.norm(data.force_history, dim=-1) # [B, N, H]
hit = (force_mag > force_threshold).any(dim=1) # [B, H]
return hit.sum(dim=-1).float() # [B]
assert data.found is not None
return data.found.squeeze(-1)
def body_angular_velocity_penalty(
env: ManagerBasedRlEnv,
asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG,
) -> torch.Tensor:
"""Penalize excessive body angular velocities."""
asset: Entity = env.scene[asset_cfg.name]
ang_vel = asset.data.body_link_ang_vel_w[:, asset_cfg.body_ids, :]
ang_vel = ang_vel.squeeze(1)
ang_vel_xy = ang_vel[:, :2] # Don't penalize z-angular velocity.
return torch.sum(torch.square(ang_vel_xy), dim=1)
def angular_momentum_penalty(
env: ManagerBasedRlEnv,
sensor_name: str,
) -> torch.Tensor:
"""Penalize whole-body angular momentum to encourage natural arm swing."""
angmom_sensor: BuiltinSensor = env.scene[sensor_name]
angmom = angmom_sensor.data
angmom_magnitude_sq = torch.sum(torch.square(angmom), dim=-1)
angmom_magnitude = torch.sqrt(angmom_magnitude_sq)
env.extras["log"]["Metrics/angular_momentum_mean"] = torch.mean(angmom_magnitude)
return angmom_magnitude_sq
def feet_air_time(
env: ManagerBasedRlEnv,
sensor_name: str,
threshold: float = 0.4,
command_name: str | None = None,
command_threshold: float = 0.1,
) -> torch.Tensor:
"""Reward feet air time."""
sensor: ContactSensor = env.scene[sensor_name]
sensor_data = sensor.data
air_time = sensor_data.current_air_time
contact_time = sensor_data.current_contact_time
in_contact = contact_time > 0.0
in_mode_time = torch.where(in_contact, contact_time, air_time)
single_stance = torch.mean(in_contact.float(), dim=1) == 0.5
mode_time = torch.min(torch.where(single_stance.unsqueeze(-1), in_mode_time, 0.0), dim=1)[0]
error = torch.abs(mode_time - threshold)
reward = torch.clamp(threshold - error, min=0.0)
if command_name is not None:
command = env.command_manager.get_command(command_name)
if command is not None:
linear_norm = torch.norm(command[:, :2], dim=1)
angular_norm = torch.abs(command[:, 2])
total_command = linear_norm + angular_norm
scale = (total_command > command_threshold).float()
reward *= scale
return reward
def feet_clearance(
env: ManagerBasedRlEnv,
target_height: float,
command_name: str | None = None,
command_threshold: float = 0.1,
asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG,
) -> torch.Tensor:
"""Penalize deviation from target clearance height, weighted by foot velocity."""
asset: Entity = env.scene[asset_cfg.name]
foot_z = asset.data.site_pos_w[:, asset_cfg.site_ids, 2] # [B, N]
foot_vel_xy = asset.data.site_lin_vel_w[:, asset_cfg.site_ids, :2] # [B, N, 2]
vel_norm = torch.norm(foot_vel_xy, dim=-1) # [B, N]
delta = torch.abs(foot_z - target_height) # [B, N]
cost = torch.sum(delta * vel_norm, dim=1) # [B]
if command_name is not None:
command = env.command_manager.get_command(command_name)
if command is not None:
linear_norm = torch.norm(command[:, :2], dim=1)
angular_norm = torch.abs(command[:, 2])
total_command = linear_norm + angular_norm
active = (total_command > command_threshold).float()
cost = cost * active
return cost
def feet_gait(
env: ManagerBasedRlEnv,
period: float,
offset: list[float],
threshold: float,
command_threshold: float,
command_name: str,
sensor_name: str,
) -> torch.Tensor:
sensor: ContactSensor = env.scene[sensor_name]
is_contact = sensor.data.current_contact_time > 0
global_phase = ((env.episode_length_buf * env.step_dt) / period).unsqueeze(1)
offsets = torch.as_tensor(offset, device=env.device, dtype=global_phase.dtype).view(1, -1)
leg_phase = (global_phase + offsets) % 1.0
is_stance = (leg_phase < threshold)
reward = (is_stance == is_contact).float().mean(dim=1)
if command_name is not None:
command = env.command_manager.get_command(command_name)
if command is not None:
linear_norm = torch.norm(command[:, :2], dim=1)
angular_norm = torch.abs(command[:, 2])
total_command = linear_norm + angular_norm
scale = (total_command > command_threshold).float()
reward *= scale
return reward
class feet_swing_height:
"""Penalize deviation from target swing height, evaluated at landing."""
def __init__(self, cfg: RewardTermCfg, env: ManagerBasedRlEnv):
self.sensor_name = cfg.params["sensor_name"]
self.site_names = cfg.params["asset_cfg"].site_names
self.peak_heights = torch.zeros(
(env.num_envs, len(self.site_names)), device=env.device, dtype=torch.float32
)
self.step_dt = env.step_dt
def __call__(
self,
env: ManagerBasedRlEnv,
sensor_name: str,
target_height: float,
command_name: str,
command_threshold: float,
asset_cfg: SceneEntityCfg,
) -> torch.Tensor:
asset: Entity = env.scene[asset_cfg.name]
contact_sensor: ContactSensor = env.scene[sensor_name]
command = env.command_manager.get_command(command_name)
assert command is not None
foot_heights = asset.data.site_pos_w[:, asset_cfg.site_ids, 2]
in_air = contact_sensor.data.found == 0
self.peak_heights = torch.where(
in_air,
torch.maximum(self.peak_heights, foot_heights),
self.peak_heights,
)
first_contact = contact_sensor.compute_first_contact(dt=self.step_dt)
linear_norm = torch.norm(command[:, :2], dim=1)
angular_norm = torch.abs(command[:, 2])
total_command = linear_norm + angular_norm
active = (total_command > command_threshold).float()
error = self.peak_heights / target_height - 1.0
cost = torch.sum(torch.square(error) * first_contact.float(), dim=1) * active
num_landings = torch.sum(first_contact.float())
peak_heights_at_landing = self.peak_heights * first_contact.float()
mean_peak_height = torch.sum(peak_heights_at_landing) / torch.clamp(
num_landings, min=1
)
env.extras["log"]["Metrics/peak_height_mean"] = mean_peak_height
self.peak_heights = torch.where(
first_contact,
torch.zeros_like(self.peak_heights),
self.peak_heights,
)
return cost
def feet_slip(
env: ManagerBasedRlEnv,
sensor_name: str,
command_name: str,
command_threshold: float = 0.01,
asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG,
) -> torch.Tensor:
"""Penalize foot sliding (xy velocity while in contact)."""
asset: Entity = env.scene[asset_cfg.name]
contact_sensor: ContactSensor = env.scene[sensor_name]
command = env.command_manager.get_command(command_name)
assert command is not None
linear_norm = torch.norm(command[:, :2], dim=1)
angular_norm = torch.abs(command[:, 2])
total_command = linear_norm + angular_norm
active = (total_command > command_threshold).float()
assert contact_sensor.data.found is not None
in_contact = (contact_sensor.data.found > 0).float() # [B, N]
foot_vel_xy = asset.data.site_lin_vel_w[:, asset_cfg.site_ids, :2] # [B, N, 2]
vel_xy_norm = torch.norm(foot_vel_xy, dim=-1) # [B, N]
vel_xy_norm_sq = torch.square(vel_xy_norm) # [B, N]
cost = torch.sum(vel_xy_norm_sq * in_contact, dim=1) * active
num_in_contact = torch.sum(in_contact)
mean_slip_vel = torch.sum(vel_xy_norm * in_contact) / torch.clamp(
num_in_contact, min=1
)
env.extras["log"]["Metrics/slip_velocity_mean"] = mean_slip_vel
return cost
def soft_landing(
env: ManagerBasedRlEnv,
sensor_name: str,
command_name: str | None = None,
command_threshold: float = 0.05,
) -> torch.Tensor:
"""Penalize high impact forces at landing to encourage soft footfalls."""
contact_sensor: ContactSensor = env.scene[sensor_name]
sensor_data = contact_sensor.data
assert sensor_data.force is not None
forces = sensor_data.force # [B, N, 3]
force_magnitude = torch.norm(forces, dim=-1) # [B, N]
first_contact = contact_sensor.compute_first_contact(dt=env.step_dt) # [B, N]
landing_impact = force_magnitude * first_contact.float() # [B, N]
cost = torch.sum(landing_impact, dim=1) # [B]
num_landings = torch.sum(first_contact.float())
mean_landing_force = torch.sum(landing_impact) / torch.clamp(num_landings, min=1)
env.extras["log"]["Metrics/landing_force_mean"] = mean_landing_force
if command_name is not None:
command = env.command_manager.get_command(command_name)
if command is not None:
linear_norm = torch.norm(command[:, :2], dim=1)
angular_norm = torch.abs(command[:, 2])
total_command = linear_norm + angular_norm
active = (total_command > command_threshold).float()
cost = cost * active
return cost
class variable_posture:
"""Penalize deviation from default pose with speed-dependent tolerance.
Uses per-joint standard deviations to control how much each joint can deviate
from default pose. Smaller std = stricter (less deviation allowed), larger
std = more forgiving. The reward is: exp(-mean(error² / std²))
Three speed regimes (based on linear + angular command velocity):
- std_standing (speed < walking_threshold): Tight tolerance for holding pose.
- std_walking (walking_threshold <= speed < running_threshold): Moderate.
- std_running (speed >= running_threshold): Loose tolerance for large motion.
Tune std values per joint based on how much motion that joint needs at each
speed. Map joint name patterns to std values, e.g. {".*knee.*": 0.35}.
"""
def __init__(self, cfg: RewardTermCfg, env: ManagerBasedRlEnv):
asset: Entity = env.scene[cfg.params["asset_cfg"].name]
default_joint_pos = asset.data.default_joint_pos
assert default_joint_pos is not None
self.default_joint_pos = default_joint_pos
_, joint_names = asset.find_joints(cfg.params["asset_cfg"].joint_names)
_, _, std_standing = resolve_matching_names_values(
data=cfg.params["std_standing"],
list_of_strings=joint_names,
)
self.std_standing = torch.tensor(
std_standing, device=env.device, dtype=torch.float32
)
_, _, std_walking = resolve_matching_names_values(
data=cfg.params["std_walking"],
list_of_strings=joint_names,
)
self.std_walking = torch.tensor(std_walking, device=env.device, dtype=torch.float32)
_, _, std_running = resolve_matching_names_values(
data=cfg.params["std_running"],
list_of_strings=joint_names,
)
self.std_running = torch.tensor(std_running, device=env.device, dtype=torch.float32)
def __call__(
self,
env: ManagerBasedRlEnv,
std_standing,
std_walking,
std_running,
asset_cfg: SceneEntityCfg,
command_name: str,
walking_threshold: float = 0.5,
running_threshold: float = 1.5,
) -> torch.Tensor:
del std_standing, std_walking, std_running # Unused.
asset: Entity = env.scene[asset_cfg.name]
command = env.command_manager.get_command(command_name)
assert command is not None
linear_speed = torch.norm(command[:, :2], dim=1)
angular_speed = torch.abs(command[:, 2])
total_speed = linear_speed + angular_speed
standing_mask = (total_speed < walking_threshold).float()
walking_mask = (
(total_speed >= walking_threshold) & (total_speed < running_threshold)
).float()
running_mask = (total_speed >= running_threshold).float()
std = (
self.std_standing * standing_mask.unsqueeze(1)
+ self.std_walking * walking_mask.unsqueeze(1)
+ self.std_running * running_mask.unsqueeze(1)
)
current_joint_pos = asset.data.joint_pos[:, asset_cfg.joint_ids]
desired_joint_pos = self.default_joint_pos[:, asset_cfg.joint_ids]
error_squared = torch.square(current_joint_pos - desired_joint_pos)
return torch.exp(-torch.mean(error_squared / (std**2), dim=1))
def stand_still(
env: ManagerBasedRlEnv,
command_name: str,
command_threshold: float = 0.1,
asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG
) -> torch.Tensor:
asset: Entity = env.scene[asset_cfg.name]
diff_angle = asset.data.joint_pos[:, asset_cfg.joint_ids] - asset.data.default_joint_pos[:, asset_cfg.joint_ids]
reward = torch.sum(torch.square(diff_angle), dim=1)
if command_name is not None:
command = env.command_manager.get_command(command_name)
if command is not None:
linear_norm = torch.norm(command[:, :2], dim=1)
angular_norm = torch.abs(command[:, 2])
total_command = linear_norm + angular_norm
scale = (total_command <= command_threshold).float()
reward *= scale
return reward
@@ -1,25 +0,0 @@
from __future__ import annotations
from typing import TYPE_CHECKING
import torch
from mjlab.sensor import ContactSensor
if TYPE_CHECKING:
from mjlab.envs import ManagerBasedRlEnv
def illegal_contact(
env: ManagerBasedRlEnv,
sensor_name: str,
force_threshold: float = 10.0,
) -> torch.Tensor:
sensor: ContactSensor = env.scene[sensor_name]
data = sensor.data
if data.force_history is not None:
# force_history: [B, N, H, 3]
force_mag = torch.norm(data.force_history, dim=-1) # [B, N, H]
return (force_mag > force_threshold).any(dim=-1).any(dim=-1) # [B]
assert data.found is not None
return torch.any(data.found, dim=-1)
@@ -1,282 +0,0 @@
from __future__ import annotations
from collections.abc import Callable
from dataclasses import dataclass, field
from typing import TYPE_CHECKING
import numpy as np
import torch
from mjlab.entity import Entity
from mjlab.managers.command_manager import CommandTerm, CommandTermCfg
from mjlab.utils.lab_api.math import (
matrix_from_quat,
quat_apply,
wrap_to_pi,
)
if TYPE_CHECKING:
import viser
from mjlab.envs.manager_based_rl_env import ManagerBasedRlEnv
from mjlab.viewer.debug_visualizer import DebugVisualizer
class UniformVelocityCommand(CommandTerm):
cfg: UniformVelocityCommandCfg
def __init__(self, cfg: UniformVelocityCommandCfg, env: ManagerBasedRlEnv):
super().__init__(cfg, env)
if self.cfg.heading_command and self.cfg.ranges.heading is None:
raise ValueError("heading_command=True but ranges.heading is set to None.")
if self.cfg.ranges.heading and not self.cfg.heading_command:
raise ValueError("ranges.heading is set but heading_command=False.")
self.robot: Entity = env.scene[cfg.entity_name]
self.vel_command_b = torch.zeros(self.num_envs, 3, device=self.device)
self.heading_target = torch.zeros(self.num_envs, device=self.device)
self.heading_error = torch.zeros(self.num_envs, device=self.device)
self.is_heading_env = torch.zeros(
self.num_envs, dtype=torch.bool, device=self.device
)
self.is_standing_env = torch.zeros_like(self.is_heading_env)
self.metrics["error_vel_xy"] = torch.zeros(self.num_envs, device=self.device)
self.metrics["error_vel_yaw"] = torch.zeros(self.num_envs, device=self.device)
# Set by create_gui() when the viewer is active.
self._joystick_enabled: viser.GuiCheckboxHandle | None = None
self._joystick_sliders: list[viser.GuiSliderHandle] = []
self._joystick_get_env_idx: Callable[[], int] | None = None
@property
def command(self) -> torch.Tensor:
return self.vel_command_b
def _update_metrics(self) -> None:
max_command_time = self.cfg.resampling_time_range[1]
max_command_step = max_command_time / self._env.step_dt
self.metrics["error_vel_xy"] += (
torch.norm(
self.vel_command_b[:, :2] - self.robot.data.root_link_lin_vel_b[:, :2], dim=-1
)
/ max_command_step
)
self.metrics["error_vel_yaw"] += (
torch.abs(self.vel_command_b[:, 2] - self.robot.data.root_link_ang_vel_b[:, 2])
/ max_command_step
)
def _resample_command(self, env_ids: torch.Tensor) -> None:
r = torch.empty(len(env_ids), device=self.device)
self.vel_command_b[env_ids, 0] = r.uniform_(*self.cfg.ranges.lin_vel_x)
self.vel_command_b[env_ids, 1] = r.uniform_(*self.cfg.ranges.lin_vel_y)
self.vel_command_b[env_ids, 2] = r.uniform_(*self.cfg.ranges.ang_vel_z)
self.vel_command_b[env_ids, :] *= (torch.norm(self.vel_command_b[env_ids, :], dim=1) > 0.1).unsqueeze(1)
if self.cfg.heading_command:
assert self.cfg.ranges.heading is not None
self.heading_target[env_ids] = r.uniform_(*self.cfg.ranges.heading)
self.is_heading_env[env_ids] = r.uniform_(0.0, 1.0) <= self.cfg.rel_heading_envs
self.is_standing_env[env_ids] = r.uniform_(0.0, 1.0) <= self.cfg.rel_standing_envs
init_vel_mask = r.uniform_(0.0, 1.0) < self.cfg.init_velocity_prob
init_vel_env_ids = env_ids[init_vel_mask]
if len(init_vel_env_ids) > 0:
root_pos = self.robot.data.root_link_pos_w[init_vel_env_ids]
root_quat = self.robot.data.root_link_quat_w[init_vel_env_ids]
lin_vel_b = self.robot.data.root_link_lin_vel_b[init_vel_env_ids]
lin_vel_b[:, :2] = self.vel_command_b[init_vel_env_ids, :2]
root_lin_vel_w = quat_apply(root_quat, lin_vel_b)
root_ang_vel_b = self.robot.data.root_link_ang_vel_b[init_vel_env_ids]
root_ang_vel_b[:, 2] = self.vel_command_b[init_vel_env_ids, 2]
root_state = torch.cat(
[root_pos, root_quat, root_lin_vel_w, root_ang_vel_b], dim=-1
)
self.robot.write_root_state_to_sim(root_state, init_vel_env_ids)
def _update_command(self) -> None:
if self.cfg.heading_command:
self.heading_error = wrap_to_pi(self.heading_target - self.robot.data.heading_w)
env_ids = self.is_heading_env.nonzero(as_tuple=False).flatten()
self.vel_command_b[env_ids, 2] = torch.clip(
self.cfg.heading_control_stiffness * self.heading_error[env_ids],
min=self.cfg.ranges.ang_vel_z[0],
max=self.cfg.ranges.ang_vel_z[1],
)
standing_env_ids = self.is_standing_env.nonzero(as_tuple=False).flatten()
self.vel_command_b[standing_env_ids, :] = 0.0
# GUI.
def create_gui(
self,
name: str,
server: "viser.ViserServer",
get_env_idx: Callable[[], int],
) -> None:
"""Create velocity joystick sliders in the Viser viewer."""
from viser import Icon
ranges = self.cfg.ranges
axes = [
("lin_vel_x", ranges.lin_vel_x[1]),
("lin_vel_y", ranges.lin_vel_y[1]),
("ang_vel_z", ranges.ang_vel_z[1]),
]
sliders: list = []
with server.gui.add_folder(name.capitalize()):
enabled = server.gui.add_checkbox("Enable", initial_value=False)
for label, max_val in axes:
max_input = server.gui.add_slider(
f"Max {label}",
initial_value=max_val,
step=0.1,
min=0.1,
max=10.0,
)
slider = server.gui.add_slider(
label,
min=-max_val,
max=max_val,
step=0.05,
initial_value=0.0,
)
@max_input.on_update
def _(_ev, _s=slider, _m=max_input) -> None:
_s.min = -_m.value
_s.max = _m.value
sliders.append(slider)
zero_btn = server.gui.add_button("Zero", icon=Icon.SQUARE_X)
@zero_btn.on_click
def _(_) -> None:
for s in sliders:
s.value = 0.0
# Store GUI state for compute() override.
self._joystick_enabled = enabled
self._joystick_sliders = sliders
self._joystick_get_env_idx = get_env_idx
def compute(self, dt: float) -> None:
super().compute(dt)
if self._joystick_enabled is not None and self._joystick_enabled.value:
assert self._joystick_get_env_idx is not None
idx = self._joystick_get_env_idx()
for i, s in enumerate(self._joystick_sliders):
self.vel_command_b[idx, i] = s.value
# Visualization.
def _debug_vis_impl(self, visualizer: "DebugVisualizer") -> None:
"""Draw velocity command and actual velocity arrows."""
env_indices = visualizer.get_env_indices(self.num_envs)
if not env_indices:
return
cmds = self.command.cpu().numpy()
base_pos_ws = self.robot.data.root_link_pos_w.cpu().numpy()
base_quat_w = self.robot.data.root_link_quat_w
base_mat_ws = matrix_from_quat(base_quat_w).cpu().numpy()
lin_vel_bs = self.robot.data.root_link_lin_vel_b.cpu().numpy()
ang_vel_bs = self.robot.data.root_link_ang_vel_b.cpu().numpy()
scale = self.cfg.viz.scale
z_offset = self.cfg.viz.z_offset
for batch in env_indices:
base_pos_w = base_pos_ws[batch]
base_mat_w = base_mat_ws[batch]
cmd = cmds[batch]
lin_vel_b = lin_vel_bs[batch]
ang_vel_b = ang_vel_bs[batch]
# Skip if robot appears uninitialized (at origin).
if np.linalg.norm(base_pos_w) < 1e-6:
continue
# Helper to transform local to world coordinates.
def local_to_world(
vec: np.ndarray, pos: np.ndarray = base_pos_w, mat: np.ndarray = base_mat_w
) -> np.ndarray:
return pos + mat @ vec
# Command linear velocity arrow (blue).
cmd_lin_from = local_to_world(np.array([0, 0, z_offset]) * scale)
cmd_lin_to = local_to_world(
(np.array([0, 0, z_offset]) + np.array([cmd[0], cmd[1], 0])) * scale
)
visualizer.add_arrow(
cmd_lin_from, cmd_lin_to, color=(0.2, 0.2, 0.6, 0.6), width=0.015
)
# Command angular velocity arrow (green).
cmd_ang_from = cmd_lin_from
cmd_ang_to = local_to_world(
(np.array([0, 0, z_offset]) + np.array([0, 0, cmd[2]])) * scale
)
visualizer.add_arrow(
cmd_ang_from, cmd_ang_to, color=(0.2, 0.6, 0.2, 0.6), width=0.015
)
# Actual linear velocity arrow (cyan).
act_lin_from = local_to_world(np.array([0, 0, z_offset]) * scale)
act_lin_to = local_to_world(
(np.array([0, 0, z_offset]) + np.array([lin_vel_b[0], lin_vel_b[1], 0])) * scale
)
visualizer.add_arrow(
act_lin_from, act_lin_to, color=(0.0, 0.6, 1.0, 0.7), width=0.015
)
# Actual angular velocity arrow (light green).
act_ang_from = act_lin_from
act_ang_to = local_to_world(
(np.array([0, 0, z_offset]) + np.array([0, 0, ang_vel_b[2]])) * scale
)
visualizer.add_arrow(
act_ang_from, act_ang_to, color=(0.0, 1.0, 0.4, 0.7), width=0.015
)
@dataclass(kw_only=True)
class UniformVelocityCommandCfg(CommandTermCfg):
entity_name: str
heading_command: bool = False
heading_control_stiffness: float = 1.0
rel_standing_envs: float = 0.0
rel_heading_envs: float = 1.0
init_velocity_prob: float = 0.0
@dataclass
class Ranges:
lin_vel_x: tuple[float, float]
lin_vel_y: tuple[float, float]
ang_vel_z: tuple[float, float]
heading: tuple[float, float] | None = None
ranges: Ranges
@dataclass
class VizCfg:
z_offset: float = 0.2
scale: float = 0.5
viz: VizCfg = field(default_factory=VizCfg)
def build(self, env: ManagerBasedRlEnv) -> UniformVelocityCommand:
return UniformVelocityCommand(self, env)
def __post_init__(self):
if self.heading_command and self.ranges.heading is None:
raise ValueError(
"The velocity command has heading commands active (heading_command=True) but "
"the `ranges.heading` parameter is set to None."
)
@@ -1,3 +0,0 @@
from .runner import (
VelocityOnPolicyRunner as VelocityOnPolicyRunner,
)
@@ -1,28 +0,0 @@
import os
import wandb
from mjlab.rl import RslRlVecEnvWrapper
from mjlab.rl.exporter_utils import (
attach_metadata_to_onnx,
get_base_metadata,
)
from mjlab.rl.runner import MjlabOnPolicyRunner
class VelocityOnPolicyRunner(MjlabOnPolicyRunner):
env: RslRlVecEnvWrapper
def save(self, path: str, infos=None):
super().save(path, infos)
policy_path = path.split("model")[0]
filename = "policy.onnx"
self.export_policy_to_onnx(policy_path, filename)
run_name: str = (
wandb.run.name if self.logger.logger_type == "wandb" and wandb.run else "local"
) # type: ignore[assignment]
onnx_path = os.path.join(policy_path, filename)
metadata = get_base_metadata(self.env.unwrapped, run_name)
attach_metadata_to_onnx(onnx_path, metadata)
if self.logger.logger_type in ["wandb"]:
wandb.save(policy_path + filename, base_path=os.path.dirname(policy_path))
@@ -1,455 +0,0 @@
"""Velocity task configuration.
This module provides a factory function to create a base velocity task config.
Robot-specific configurations call the factory and customize as needed.
"""
import math
from dataclasses import replace
from mjlab.envs import ManagerBasedRlEnvCfg
from mjlab.envs import mdp as envs_mdp
from mjlab.envs.mdp import dr
from mjlab.envs.mdp.actions import JointPositionActionCfg
from mjlab.managers.action_manager import ActionTermCfg
from mjlab.managers.command_manager import CommandTermCfg
from mjlab.managers.curriculum_manager import CurriculumTermCfg
from mjlab.managers.event_manager import EventTermCfg
from mjlab.managers.metrics_manager import MetricsTermCfg
from mjlab.managers.observation_manager import ObservationGroupCfg, ObservationTermCfg
from mjlab.managers.reward_manager import RewardTermCfg
from mjlab.managers.scene_entity_config import SceneEntityCfg
from mjlab.managers.termination_manager import TerminationTermCfg
from mjlab.scene import SceneCfg
from mjlab.sensor import GridPatternCfg, ObjRef, RayCastSensorCfg
from mjlab.sim import MujocoCfg, SimulationCfg
from mjlab.tasks.velocity import mdp
from mjlab.tasks.velocity.mdp import UniformVelocityCommandCfg
from mjlab.terrains import TerrainEntityCfg
from mjlab.terrains.config import ROUGH_TERRAINS_CFG
from mjlab.utils.noise import UniformNoiseCfg as Unoise
from mjlab.viewer import ViewerConfig
import src.tasks.velocity.mdp as mdp
def make_velocity_env_cfg() -> ManagerBasedRlEnvCfg:
"""Create base velocity tracking task configuration."""
##
# Sensors
##
terrain_scan = RayCastSensorCfg(
name="terrain_scan",
frame=ObjRef(type="body", name="", entity="robot"), # Set per-robot.
ray_alignment="yaw",
pattern=GridPatternCfg(size=(1.6, 1.0), resolution=0.1),
max_distance=5.0,
exclude_parent_body=True,
debug_vis=True,
viz=RayCastSensorCfg.VizCfg(show_normals=True),
)
##
# Observations
##
actor_terms = {
"base_ang_vel": ObservationTermCfg(
func=mdp.builtin_sensor,
params={"sensor_name": "robot/imu_ang_vel"},
noise=Unoise(n_min=-0.2, n_max=0.2),
),
"projected_gravity": ObservationTermCfg(
func=mdp.projected_gravity,
noise=Unoise(n_min=-0.05, n_max=0.05),
),
"command": ObservationTermCfg(
func=mdp.generated_commands,
params={"command_name": "twist"},
),
"phase": ObservationTermCfg(
func=mdp.phase,
params={"period": 0.6, "command_name": "twist"},
),
"joint_pos": ObservationTermCfg(
func=mdp.joint_pos_rel,
noise=Unoise(n_min=-0.01, n_max=0.01),
),
"joint_vel": ObservationTermCfg(
func=mdp.joint_vel_rel,
noise=Unoise(n_min=-1.5, n_max=1.5),
),
"actions": ObservationTermCfg(func=mdp.last_action),
"height_scan": ObservationTermCfg(
func=envs_mdp.height_scan,
params={"sensor_name": "terrain_scan"},
noise=Unoise(n_min=-0.1, n_max=0.1),
scale=1 / terrain_scan.max_distance,
),
}
critic_terms = {
**actor_terms,
"base_lin_vel": ObservationTermCfg(
func=mdp.builtin_sensor,
params={"sensor_name": "robot/imu_lin_vel"},
noise=Unoise(n_min=-0.5, n_max=0.5),
),
"height_scan": ObservationTermCfg(
func=envs_mdp.height_scan,
params={"sensor_name": "terrain_scan"},
scale=1 / terrain_scan.max_distance,
),
"foot_height": ObservationTermCfg(
func=mdp.foot_height,
params={"asset_cfg": SceneEntityCfg("robot", site_names=())}, # Set per-robot.
),
"foot_air_time": ObservationTermCfg(
func=mdp.foot_air_time,
params={"sensor_name": "feet_ground_contact"},
),
"foot_contact": ObservationTermCfg(
func=mdp.foot_contact,
params={"sensor_name": "feet_ground_contact"},
),
"foot_contact_forces": ObservationTermCfg(
func=mdp.foot_contact_forces,
params={"sensor_name": "feet_ground_contact"},
),
}
observations = {
"actor": ObservationGroupCfg(
terms=actor_terms,
concatenate_terms=True,
enable_corruption=True,
history_length=1,
),
"critic": ObservationGroupCfg(
terms=critic_terms,
concatenate_terms=True,
enable_corruption=False,
history_length=1,
),
}
##
# Metrics
##
metrics = {
"linear_velocity_rmse": MetricsTermCfg(
func=mdp.linear_velocity_rmse,
params={"command_name": "twist"},
),
"angular_velocity_rmse": MetricsTermCfg(
func=mdp.angular_velocity_rmse,
params={"command_name": "twist"},
),
"mean_action_acc": MetricsTermCfg(func=mdp.mean_action_acc),
"orientation_error": MetricsTermCfg(func=mdp.orientation_error),
"fall_indicator": MetricsTermCfg(func=mdp.fall_indicator),
"slip_velocity": MetricsTermCfg(
func=mdp.slip_velocity,
params={
"sensor_name": "feet_ground_contact",
"asset_cfg": SceneEntityCfg("robot", site_names=()), # Set per-robot.
},
),
"mechanical_power": MetricsTermCfg(
func=mdp.mechanical_power,
params={"asset_cfg": SceneEntityCfg("robot", joint_names=(".*",))},
),
}
##
# Actions
##
actions: dict[str, ActionTermCfg] = {
"joint_pos": JointPositionActionCfg(
entity_name="robot",
actuator_names=(".*",),
scale=0.25, # Override per-robot.
use_default_offset=True,
)
}
##
# Commands
##
commands: dict[str, CommandTermCfg] = {
"twist": UniformVelocityCommandCfg(
entity_name="robot",
resampling_time_range=(3.0, 8.0),
rel_standing_envs=0.05,
heading_command=True,
heading_control_stiffness=0.5,
debug_vis=True,
ranges=UniformVelocityCommandCfg.Ranges(
lin_vel_x=(-1.0, 2.0),
lin_vel_y=(-1.0, 1.0),
ang_vel_z=(-1.0, 1.0),
heading=(-math.pi, math.pi),
),
)
}
##
# Events
##
events = {
"reset_base": EventTermCfg(
func=mdp.reset_root_state_uniform,
mode="reset",
params={
"pose_range": {
"x": (-0.5, 0.5),
"y": (-0.5, 0.5),
"z": (0.0, 0.0),
"yaw": (-3.14, 3.14),
},
"velocity_range": {},
},
),
"reset_robot_joints": EventTermCfg(
func=mdp.reset_joints_by_offset,
mode="reset",
params={
"position_range": (-0.0, 0.0),
"velocity_range": (-0.0, 0.0),
"asset_cfg": SceneEntityCfg("robot", joint_names=(".*",)),
},
),
"push_robot": EventTermCfg(
func=mdp.push_by_setting_velocity,
mode="interval",
interval_range_s=(5.0, 6.0),
params={
"velocity_range": {
"x": (-0.5, 0.5),
"y": (-0.5, 0.5),
"z": (-0.4, 0.4),
"roll": (-0.52, 0.52),
"pitch": (-0.52, 0.52),
"yaw": (-0.78, 0.78),
},
},
),
"foot_friction": EventTermCfg(
mode="startup",
func=dr.geom_friction,
params={
"asset_cfg": SceneEntityCfg("robot", geom_names=()), # Set per-robot.
"operation": "abs",
"ranges": (0.3, 1.6),
"shared_random": True, # All foot geoms share the same friction.
},
),
"encoder_bias": EventTermCfg(
mode="startup",
func=dr.encoder_bias,
params={
"asset_cfg": SceneEntityCfg("robot"),
"bias_range": (-0.015, 0.015),
},
),
"base_com": EventTermCfg(
mode="startup",
func=dr.body_com_offset,
params={
"asset_cfg": SceneEntityCfg("robot", body_names=()), # Set per-robot.
"operation": "add",
"ranges": {
0: (-0.05, 0.05),
1: (-0.05, 0.05),
2: (-0.05, 0.05),
},
},
),
}
##
# Rewards
##
rewards = {
"track_linear_velocity": RewardTermCfg(
func=mdp.track_linear_velocity,
weight=1.0,
params={"command_name": "twist", "std": math.sqrt(0.25)},
),
"track_angular_velocity": RewardTermCfg(
func=mdp.track_angular_velocity,
weight=1.0,
params={"command_name": "twist", "std": math.sqrt(0.5)},
),
"body_orientation_l2": RewardTermCfg(
func=mdp.body_orientation_l2,
weight=-1.0,
params={"asset_cfg": SceneEntityCfg("robot", body_names=())}, # Set per-robot.
),
"pose": RewardTermCfg(
func=mdp.variable_posture,
weight=1.0,
params={
"asset_cfg": SceneEntityCfg("robot", joint_names=".*"),
"command_name": "twist",
"std_standing": {}, # Set per-robot.
"std_walking": {}, # Set per-robot.
"std_running": {}, # Set per-robot.
"walking_threshold": 0.1,
"running_threshold": 1.5,
},
),
"body_ang_vel": RewardTermCfg(
func=mdp.body_angular_velocity_penalty,
weight=-0.05, # Override per-robot
params={"asset_cfg": SceneEntityCfg("robot", body_names=())}, # Set per-robot.
),
"angular_momentum": RewardTermCfg(
func=mdp.angular_momentum_penalty,
weight=-0.025, # Override per-robot
params={"sensor_name": "robot/root_angmom"},
),
"is_terminated": RewardTermCfg(func=mdp.is_terminated, weight=-200.0),
"electrical_power": RewardTermCfg(
func=mdp.electrical_power_cost,
weight=0.0,
params={"asset_cfg": SceneEntityCfg("robot", joint_names=(".*",))},
),
"joint_acc_l2": RewardTermCfg(func=mdp.joint_acc_l2, weight=-2.5e-7),
"joint_pos_limits": RewardTermCfg(func=mdp.joint_pos_limits, weight=-10.0),
"action_rate_l2": RewardTermCfg(func=mdp.action_rate_l2, weight=-0.05),
"foot_gait": RewardTermCfg(
func=mdp.feet_gait,
weight=0.5,
params={
"period": 0.6,
"offset": [0.0, 0.5],
"threshold": 0.56,
"command_threshold": 0.1,
"command_name": "twist",
"sensor_name": "feet_ground_contact",
}
),
"foot_clearance": RewardTermCfg(
func=mdp.feet_clearance,
weight=-1.0,
params={
"target_height": 0.10,
"command_name": "twist",
"command_threshold": 0.1,
"asset_cfg": SceneEntityCfg("robot", site_names=()), # Set per-robot.
},
),
"foot_slip": RewardTermCfg(
func=mdp.feet_slip,
weight=-0.25,
params={
"sensor_name": "feet_ground_contact",
"command_name": "twist",
"command_threshold": 0.1,
"asset_cfg": SceneEntityCfg("robot", site_names=()), # Set per-robot.
},
),
"soft_landing": RewardTermCfg(
func=mdp.soft_landing,
weight=-1e-3,
params={
"sensor_name": "feet_ground_contact",
"command_name": "twist",
"command_threshold": 0.1,
},
),
"stand_still": RewardTermCfg(
func=mdp.stand_still,
weight=-1.0,
params={
"command_name": "twist",
"command_threshold": 0.1,
"asset_cfg": SceneEntityCfg("robot", joint_names=".*"),
},
),
}
##
# Terminations
##
terminations = {
"time_out": TerminationTermCfg(func=mdp.time_out, time_out=True),
"fell_over": TerminationTermCfg(
func=mdp.bad_orientation,
params={"limit_angle": math.radians(70.0)},
),
}
##
# Curriculum
##
curriculum = {
"terrain_levels": CurriculumTermCfg(
func=mdp.terrain_levels_vel,
params={"command_name": "twist"},
),
"command_vel": CurriculumTermCfg(
func=mdp.commands_vel,
params={
"command_name": "twist",
"velocity_stages": [
{"step": 0, "lin_vel_x": (-0.5, 1.0), "lin_vel_y": (-0.5, 0.5), "ang_vel_z": (-1.0, 1.0)},
{"step": 5000 * 24, "lin_vel_x": (-1.0, 2.0), "lin_vel_y": (-1.0, 1.0)},
],
},
),
}
##
# Assemble and return
##
return ManagerBasedRlEnvCfg(
scene=SceneCfg(
terrain=TerrainEntityCfg(
terrain_type="generator",
terrain_generator=replace(ROUGH_TERRAINS_CFG),
max_init_terrain_level=5,
),
sensors=(terrain_scan,),
num_envs=1,
extent=2.0,
),
observations=observations,
actions=actions,
commands=commands,
events=events,
rewards=rewards,
terminations=terminations,
curriculum=curriculum,
metrics=metrics,
viewer=ViewerConfig(
origin_type=ViewerConfig.OriginType.ASSET_BODY,
entity_name="robot",
body_name="", # Set per-robot.
distance=3.0,
elevation=-5.0,
azimuth=90.0,
),
sim=SimulationCfg(
nconmax=35,
njmax=1500,
mujoco=MujocoCfg(
timestep=0.005,
iterations=10,
ls_iterations=20,
),
),
decimation=4,
episode_length_s=20.0,
)
+393 -698
View File
File diff suppressed because it is too large Load Diff
+54 -189
View File
@@ -1,32 +1,20 @@
import json
import subprocess
import sys
import tempfile
import threading
import time
import unittest
from pathlib import Path
from unittest.mock import patch
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from server import ( # noqa: E402
MAX_JOBS,
ApiError,
TrainingJob,
TrainingManager,
TrainingRequestHandler,
default_trainer_root,
termination_signal_handler,
)
from server import ApiError, TrainingManager # noqa: E402
class TrainingManagerTest(unittest.TestCase):
def setUp(self):
self.temporary = tempfile.TemporaryDirectory()
self.root = Path(self.temporary.name)
(self.root / "scripts").mkdir()
(self.root / "scripts" / "train.py").write_text(
"""import os, pathlib, time
def setUp(self):
self.temporary = tempfile.TemporaryDirectory()
self.root = Path(self.temporary.name)
(self.root / "scripts").mkdir()
(self.root / "scripts" / "train.py").write_text(
"""import os, pathlib, time
print('WANDB_MODE=' + os.environ.get('WANDB_MODE', ''), flush=True)
print('Learning iteration 1 / 2', flush=True)
time.sleep(0.02)
@@ -35,181 +23,58 @@ out=pathlib.Path('logs/rsl_rl/test/run/policy.onnx')
out.parent.mkdir(parents=True, exist_ok=True)
out.write_bytes(b'onnx')
""",
encoding="utf-8",
)
self.manager = TrainingManager(
self.root, sys.executable, ("Unitree-Go2-Flat",), check_environment=False
)
encoding="utf-8",
)
self.manager = TrainingManager(self.root, sys.executable, ("Unitree-Go2-Flat",), check_environment=False)
def tearDown(self):
self.temporary.cleanup()
def tearDown(self):
self.temporary.cleanup()
@staticmethod
def payload(**patch):
value = {
"taskId": "Unitree-Go2-Flat",
"numEnvs": 16,
"maxIterations": 2,
"seed": 42,
"runName": "browser-test",
"device": "cpu",
"gpuIds": [],
"wandbMode": "offline",
}
value.update(patch)
return value
@staticmethod
def payload(**patch):
value = {
"taskId": "Unitree-Go2-Flat",
"numEnvs": 16,
"maxIterations": 2,
"seed": 42,
"runName": "browser-test",
"device": "cpu",
"gpuIds": [],
"wandbMode": "offline",
}
value.update(patch)
return value
def test_default_trainer_is_bundled_with_go2_assets(self):
trainer_root = default_trainer_root()
self.assertEqual(trainer_root, Path(__file__).resolve().parents[1] / "rl")
self.assertTrue((trainer_root / "scripts" / "train.py").is_file())
self.assertTrue(
(
trainer_root / "src" / "assets" / "robots" / "unitree_go2" / "xmls" / "go2.xml"
).is_file()
)
def test_validates_allowlist_and_limits(self):
with self.assertRaises(ApiError):
self.manager.parse_config(self.payload(taskId="shell injection"))
with self.assertRaises(ApiError):
self.manager.parse_config(self.payload(numEnvs=0))
with self.assertRaises(ApiError):
self.manager.parse_config(self.payload(runName="bad name"))
with self.assertRaises(ApiError):
self.manager.parse_config(self.payload(wandbMode="login"))
def test_validates_allowlist_and_limits(self):
with self.assertRaises(ApiError):
self.manager.parse_config(self.payload(taskId="shell injection"))
with self.assertRaises(ApiError):
self.manager.parse_config(self.payload(numEnvs=0))
with self.assertRaises(ApiError):
self.manager.parse_config(self.payload(runName="bad name"))
with self.assertRaises(ApiError):
self.manager.parse_config(self.payload(wandbMode="login"))
def test_builds_argument_array_without_shell(self):
config = self.manager.parse_config(self.payload(device="gpu", gpuIds=[0, 2]))
command = self.manager.command_for(config)
self.assertEqual(command[:4], [sys.executable, "-u", "scripts/train.py", "Unitree-Go2-Flat"])
self.assertEqual(command[-2:], ["--gpu-ids", "[0,2]"])
def test_builds_argument_array_without_shell(self):
config = self.manager.parse_config(self.payload(device="gpu", gpuIds=[0, 2]))
command = self.manager.command_for(config)
self.assertEqual(
command[:4], [sys.executable, "-u", "scripts/train.py", "Unitree-Go2-Flat"]
)
self.assertEqual(command[-2:], ["--gpu-ids", "[0,2]"])
def test_resolves_reward_preset_to_inline_validated_trainer_argument(self):
preset_id = "f" * 32
reward_config = {"weights": {"pose": 1.2}, "params": {}}
self.manager.preset_resolver = lambda value: reward_config if value == preset_id else None
config = self.manager.parse_config(self.payload(rewardPresetId=preset_id))
command = self.manager.command_for(config)
index = command.index("--reward-config-json")
self.assertEqual(json.loads(command[index + 1]), reward_config)
def test_requires_local_host_origin_and_bearer_token(self):
handler = object.__new__(TrainingRequestHandler)
handler.access_token = "secret-token-1234"
handler.allowed_origins = ()
handler.headers = {
"Host": "127.0.0.1:8765",
"Origin": "http://localhost:5173",
"Authorization": "Bearer secret-token-1234",
}
handler._ensure_request()
handler.headers["Authorization"] = "Bearer wrong-token"
with self.assertRaises(ApiError) as error:
handler._ensure_request()
self.assertEqual(error.exception.status, 401)
handler.headers["Authorization"] = "Bearer secret-token-1234"
handler.headers["Host"] = "attacker.example"
with self.assertRaises(ApiError) as error:
handler._ensure_request()
self.assertEqual(error.exception.status, 403)
def test_caps_completed_job_history(self):
config = self.manager.parse_config(self.payload())
for index in range(MAX_JOBS):
job_id = f"{index:032x}"
self.manager.jobs[job_id] = TrainingJob(
id=job_id,
config=config,
state="succeeded",
)
with patch("server.threading.Thread") as thread:
created = self.manager.start(self.payload())
self.assertEqual(len(self.manager.jobs), MAX_JOBS)
self.assertNotIn(f"{0:032x}", self.manager.jobs)
self.assertIn(created["id"], self.manager.jobs)
thread.return_value.start.assert_called_once()
def test_cancel_waits_until_starting_process_is_registered(self):
entered_popen = threading.Event()
release_popen = threading.Event()
terminated = threading.Event()
class FakeStdout:
def __iter__(self):
terminated.wait(2)
return iter(())
def close(self):
pass
class FakeProcess:
pid = 1234
stdout = FakeStdout()
@staticmethod
def poll():
return -15 if terminated.is_set() else None
@staticmethod
def wait(timeout=None):
if not terminated.wait(timeout):
raise subprocess.TimeoutExpired("fake-training", timeout)
return -15
def create_process(*_args, **_kwargs):
entered_popen.set()
self.assertTrue(release_popen.wait(2))
return FakeProcess()
config = self.manager.parse_config(self.payload())
job = TrainingJob(id="a" * 32, config=config)
self.manager.jobs[job.id] = job
runner = threading.Thread(target=self.manager._run, args=(job,))
cancel_done = threading.Event()
def cancel():
self.manager.cancel(job.id)
cancel_done.set()
with (
patch("server.subprocess.Popen", side_effect=create_process),
patch("server.os.killpg", side_effect=lambda *_args: terminated.set()) as killpg,
):
runner.start()
self.assertTrue(entered_popen.wait(2))
canceller = threading.Thread(target=cancel)
canceller.start()
self.assertFalse(cancel_done.wait(0.05))
release_popen.set()
canceller.join(2)
runner.join(2)
self.assertFalse(runner.is_alive())
self.assertFalse(canceller.is_alive())
killpg.assert_called_once_with(FakeProcess.pid, 15)
self.assertEqual(self.manager.get(job.id)["state"], "cancelled")
def test_sigterm_enters_controlled_shutdown(self):
with self.assertRaises(KeyboardInterrupt):
termination_signal_handler(15, None)
def test_runs_job_and_exposes_new_onnx_artifact(self):
job = self.manager.start(self.payload())
deadline = time.monotonic() + 5
while time.monotonic() < deadline:
job = self.manager.get(job["id"])
if job["state"] not in ("queued", "running"):
break
time.sleep(0.02)
self.assertEqual(job["state"], "succeeded")
self.assertEqual(job["iteration"], 2)
self.assertIn("WANDB_MODE=offline", job["logs"])
self.assertTrue(job["artifactReady"])
self.assertEqual(self.manager.artifact(job["id"]).read_bytes(), b"onnx")
def test_runs_job_and_exposes_new_onnx_artifact(self):
job = self.manager.start(self.payload())
deadline = time.monotonic() + 5
while time.monotonic() < deadline:
job = self.manager.get(job["id"])
if job["state"] not in ("queued", "running"):
break
time.sleep(0.02)
self.assertEqual(job["state"], "succeeded")
self.assertEqual(job["iteration"], 2)
self.assertIn("WANDB_MODE=offline", job["logs"])
self.assertTrue(job["artifactReady"])
self.assertEqual(self.manager.artifact(job["id"]).read_bytes(), b"onnx")
if __name__ == "__main__":
unittest.main()
unittest.main()
-239
View File
@@ -1,239 +0,0 @@
import math
import sys
import tempfile
import unittest
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from tuning.advisor import AdvisorConfig, DeepSeekAdvisor # noqa: E402
from tuning.schema import ( # noqa: E402
BASE_REWARD_CONFIGURATION,
RewardConfigError,
merge_proposal,
validate_configuration,
validate_constraints,
validate_proposal,
)
from tuning.scoring import ( # noqa: E402
DEFAULT_OBJECTIVE_WEIGHTS,
EvaluationError,
score_evaluation,
)
from tuning.storage import StorageConflict, TuningStorage # noqa: E402
class RewardSchemaTest(unittest.TestCase):
def test_baseline_is_complete_and_energy_is_disabled(self):
config = validate_configuration(BASE_REWARD_CONFIGURATION)
self.assertEqual(len(config["weights"]), 16)
self.assertEqual(config["weights"]["electrical_power"], 0.0)
def test_sparse_proposal_constraints(self):
patch = validate_proposal(
{
"weights": {"track_linear_velocity": 1.2, "foot_slip": -0.3},
"params": {"foot_gait.period": 0.65},
},
BASE_REWARD_CONFIGURATION,
)
merged = merge_proposal(BASE_REWARD_CONFIGURATION, patch)
self.assertEqual(merged["weights"]["track_linear_velocity"], 1.2)
with self.assertRaises(RewardConfigError):
validate_proposal({"weights": {"track_linear_velocity": 0}}, BASE_REWARD_CONFIGURATION)
with self.assertRaises(RewardConfigError):
validate_proposal({"weights": {"foot_slip": 0.2}}, BASE_REWARD_CONFIGURATION)
with self.assertRaises(RewardConfigError):
validate_proposal(
{"weights": {"track_linear_velocity": 3.0}}, BASE_REWARD_CONFIGURATION
)
with self.assertRaises(RewardConfigError):
validate_proposal({"params": {"foot_gait.period": math.nan}}, BASE_REWARD_CONFIGURATION)
with self.assertRaises(RewardConfigError):
validate_proposal(
{
"weights": {
"pose": 1.1,
"foot_gait": 0.6,
"foot_slip": -0.3,
"soft_landing": -0.002,
},
"params": {"foot_gait.period": 0.65},
},
BASE_REWARD_CONFIGURATION,
)
def test_cross_parameter_order(self):
with self.assertRaises(RewardConfigError):
validate_proposal(
{"params": {"pose.walking_threshold": 0.5, "pose.running_threshold": 0.4}},
BASE_REWARD_CONFIGURATION,
)
def test_session_constraints_reject_unknown_out_of_range_and_fixed_changes(self):
constraints = validate_constraints(
{
"weights.track_linear_velocity": {"kind": "fixed", "value": 1.0},
"params.foot_gait.period": {"kind": "range", "min": 0.5, "max": 0.7},
}
)
validate_proposal(
{"params": {"foot_gait.period": 0.65}},
BASE_REWARD_CONFIGURATION,
constraints,
)
with self.assertRaisesRegex(RewardConfigError, "已固定"):
validate_proposal(
{"weights": {"track_linear_velocity": 1.1}},
BASE_REWARD_CONFIGURATION,
constraints,
)
with self.assertRaisesRegex(RewardConfigError, "工程锁定范围"):
validate_proposal(
{"params": {"foot_gait.period": 0.75}},
BASE_REWARD_CONFIGURATION,
constraints,
)
with self.assertRaisesRegex(RewardConfigError, "未知参数约束"):
validate_constraints({"weights.not_allowed": {"kind": "fixed", "value": 1.0}})
class ScoringTest(unittest.TestCase):
baseline = {
"linear_velocity_rmse": 0.3,
"angular_velocity_rmse": 0.2,
"mean_action_acc": 0.1,
"orientation_error": 0.2,
"fall_rate": 0.1,
"slip_velocity": 0.2,
"mechanical_power": 100.0,
}
def test_improvement_and_safety_gate(self):
better = {key: value * 0.8 for key, value in self.baseline.items()}
scored = score_evaluation(self.baseline, better, DEFAULT_OBJECTIVE_WEIGHTS)
self.assertTrue(scored["eligible"])
self.assertGreater(scored["score"], 0)
unsafe = dict(better, fall_rate=0.2)
scored = score_evaluation(self.baseline, unsafe)
self.assertFalse(scored["eligible"])
self.assertEqual(scored["score"], -1.0)
def test_rejects_missing_and_nonfinite_metrics(self):
with self.assertRaises(EvaluationError):
score_evaluation(self.baseline, {"fall_rate": 0.1})
bad = dict(self.baseline, mechanical_power=math.inf)
with self.assertRaises(EvaluationError):
score_evaluation(self.baseline, bad)
class AdvisorTest(unittest.TestCase):
class Output:
weights = {"pose": 1.1}
params = {}
rationale = "improve posture"
expected_impact = {"posture": "better"}
confidence = 0.7
class Result:
output = None
@staticmethod
def usage():
return type("Usage", (), {"requests": 1, "input_tokens": 10, "output_tokens": 5})()
class Agent:
def __init__(self, output):
self.output = output
def run_sync(self, _prompt):
result = AdvisorTest.Result()
result.output = self.output
return result
def test_structured_result_is_revalidated_locally(self):
advisor = DeepSeekAdvisor(AdvisorConfig("fake"))
advisor._cached_agent = self.Agent(self.Output())
proposal = advisor.propose({"trials": []}, BASE_REWARD_CONFIGURATION)
self.assertEqual(proposal["patch"]["weights"]["pose"], 1.1)
self.assertEqual(proposal["usage"]["input_tokens"], 10)
def test_invalid_model_patch_is_rejected(self):
output = self.Output()
output.weights = {"track_linear_velocity": -1.0}
advisor = DeepSeekAdvisor(AdvisorConfig("fake"))
advisor._cached_agent = self.Agent(output)
with self.assertRaises(RewardConfigError):
advisor.propose({"trials": []}, BASE_REWARD_CONFIGURATION)
class StorageTest(unittest.TestCase):
def setUp(self):
self.temporary = tempfile.TemporaryDirectory()
self.storage = TuningStorage(Path(self.temporary.name) / "state.sqlite3")
def tearDown(self):
self.temporary.cleanup()
def test_persists_session_trial_proposal_and_metrics(self):
session = self.storage.create_session(
"approval", {"taskId": "Unitree-Go2-Flat"}, DEFAULT_OBJECTIVE_WEIGHTS, False
)
trial = self.storage.create_trial(
session["id"], 0, 0, 10, BASE_REWARD_CONFIGURATION, None, "trial-000-rung-0"
)
proposal = self.storage.create_proposal(
session["id"],
trial["id"],
{"weights": {"pose": 1.1}, "params": {}},
"test",
{},
0.8,
)
self.assertTrue(self.storage.decide_proposal(proposal["id"], "approved", None))
self.assertFalse(self.storage.decide_proposal(proposal["id"], "approved", None))
points = [
("Train/reward", step, float(step), 50.0 if step == 50 else float(step % 7))
for step in range(100)
]
self.storage.insert_metrics(trial["id"], points)
sampled = self.storage.metrics(trial["id"], max_points=10)[0]["points"]
self.assertEqual(len(sampled), 10)
self.assertEqual(sampled[0]["step"], 0)
self.assertEqual(sampled[-1]["step"], 99)
self.assertIn(50.0, [point["value"] for point in sampled])
control = self.storage.replace_constraints(
session["id"],
0,
{"weights.pose": {"kind": "range", "min": 0.5, "max": 1.5}},
)
self.assertEqual(control["constraintsRevision"], 1)
with self.assertRaises(StorageConflict):
self.storage.replace_constraints(session["id"], 0, {})
self.storage.grant_dispatch_token(session["id"])
with self.assertRaises(StorageConflict):
self.storage.grant_dispatch_token(session["id"])
self.assertTrue(self.storage.use_dispatch_token(session["id"]))
self.assertFalse(self.storage.use_dispatch_token(session["id"]))
incremental = self.storage.metrics(trial["id"], max_points=100, after_step=90)[0]
self.assertEqual(incremental["points"][0]["step"], 91)
reopened = TuningStorage(self.storage.path)
self.assertEqual(reopened.get_session(session["id"])["mode"], "approval")
self.assertEqual(reopened.get_control(session["id"])["constraintsRevision"], 1)
def test_recovery_marks_inflight_records(self):
session = self.storage.create_session(
"automatic", {"taskId": "Unitree-Go2-Flat"}, DEFAULT_OBJECTIVE_WEIGHTS, True
)
trial = self.storage.create_trial(
session["id"], 0, 0, 10, BASE_REWARD_CONFIGURATION, None, "trial-000-rung-0"
)
self.storage.update_session(session["id"], state="running")
self.storage.update_trial(trial["id"], state="training")
self.storage.recover_interrupted()
self.assertEqual(self.storage.get_session(session["id"])["state"], "interrupted")
self.assertEqual(self.storage.get_trial(trial["id"])["state"], "interrupted")
if __name__ == "__main__":
unittest.main()
@@ -1,350 +0,0 @@
import sys
import tempfile
import time
import unittest
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from tuning.manager import TuningManager # noqa: E402
from tuning.process import GpuLease # noqa: E402
from tuning.schema import BASE_REWARD_CONFIGURATION # noqa: E402
from tuning.scoring import score_evaluation # noqa: E402
from tuning.storage import now_iso # noqa: E402
BASE_METRICS = {
"linear_velocity_rmse": 0.3,
"angular_velocity_rmse": 0.2,
"mean_action_acc": 0.1,
"orientation_error": 0.2,
"fall_rate": 0.1,
"slip_velocity": 0.2,
"mechanical_power": 100.0,
}
class FakeAdvisor:
def capability(self):
return {
"configured": True,
"apiKeyConfigured": True,
"frameworkInstalled": True,
"model": "fake",
"baseUrl": "https://example.invalid",
}
def propose(self, _context, previous):
value = min(2.4, previous["weights"]["pose"] * 1.05)
return {
"patch": {"weights": {"pose": value}, "params": {}},
"rationale": "fake",
"expectedImpact": {},
"confidence": 0.8,
"promptHash": "abc",
"usage": {},
"model": "fake",
}
def test_connection(self):
return {"ok": True, "model": "fake", "outputType": "fake"}
class FakeTuningManager(TuningManager):
def _execute_trial(self, session, trial, resume_checkpoint=None):
del resume_checkpoint
factor = max(0.5, 1.0 - 0.03 * trial["number"] - 0.01 * trial["rung"])
metrics = {key: value * factor for key, value in BASE_METRICS.items()}
trials = self.storage.list_trials(session["id"])
baseline = next(
(item for item in trials if item["number"] == 0 and item["rung"] == 0), None
)
if baseline and baseline["evaluation"]:
scored = score_evaluation(
baseline["evaluation"]["metrics"], metrics, session["objectiveWeights"]
)
else:
scored = {"score": 0.0, "eligible": True, "components": {}}
root = self._session_root(session["id"])
run = root / trial["runDir"]
run.mkdir(parents=True, exist_ok=True)
(run / "model_1.pt").write_bytes(b"checkpoint")
(run / "policy.onnx").write_bytes(b"onnx")
evaluation = {"metrics": metrics, "score": scored}
self.storage.update_trial(
trial["id"],
state="completed",
started_at=now_iso(),
ended_at=now_iso(),
message="fake complete",
checkpoint_path=str((run / "model_1.pt").relative_to(root)),
policy_path=str((run / "policy.onnx").relative_to(root)),
evaluation=evaluation,
score=scored["score"],
eligible=scored["eligible"],
)
return self.storage.get_trial(trial["id"])
class TuningManagerTest(unittest.TestCase):
def setUp(self):
self.temporary = tempfile.TemporaryDirectory()
self.root = Path(self.temporary.name)
(self.root / "trainer" / "scripts").mkdir(parents=True)
(self.root / "trainer" / "scripts" / "evaluate.py").write_text("", encoding="utf-8")
self.manager = FakeTuningManager(
self.root / "trainer",
sys.executable,
self.root / "data",
GpuLease(),
advisor=FakeAdvisor(),
)
def tearDown(self):
self.manager.shutdown()
self.temporary.cleanup()
@staticmethod
def payload(mode="automatic"):
return {
"taskId": "Unitree-Go2-Flat",
"mode": mode,
"runName": "test",
"numEnvs": 16,
"gpuIds": [0],
"trialCount": 4,
"initialIterations": 1,
"middleIterations": 2,
"finalIterations": 3,
"evalNumEnvs": 8,
"evalSteps": 10,
}
def wait_terminal(self, session_id, timeout=5):
deadline = time.monotonic() + timeout
while time.monotonic() < deadline:
session = self.manager.detail(session_id)
if session["state"] in {"succeeded", "failed", "cancelled"}:
return session
time.sleep(0.01)
self.fail("session did not finish")
def test_automatic_session_runs_rungs_and_persists_best_artifact(self):
session = self.manager.create(self.payload())
completed = self.wait_terminal(session["id"])
self.assertEqual(completed["state"], "succeeded", completed["message"])
self.assertGreaterEqual(len(completed["trials"]), 7)
self.assertTrue(self.manager.best_artifact(session["id"]).is_file())
self.assertEqual(len(self.manager.storage.list_presets()), 1)
def test_approval_session_waits_and_accepts_modified_patch(self):
session = self.manager.create(self.payload("approval"))
deadline = time.monotonic() + 3
while time.monotonic() < deadline:
detail = self.manager.detail(session["id"])
if detail["state"] == "awaiting_approval":
break
time.sleep(0.01)
else:
self.fail("session did not wait for approval")
proposal = detail["proposals"][-1]
patch = {"weights": {"pose": 1.1}, "params": {}}
with self.assertRaisesRegex(Exception, "feedback"):
self.manager.approve(session["id"], proposal["id"], {"feedback": {}})
approved = self.manager.approve(
session["id"], proposal["id"], {"feedback": "ok", "patch": patch}
)
self.assertEqual(approved["proposals"][-1]["state"], "approved")
self.manager.cancel(session["id"])
self.assertEqual(self.wait_terminal(session["id"])["state"], "cancelled")
def test_runtime_mode_switch_auto_approves_pending_proposal(self):
session = self.manager.create(self.payload("approval"))
deadline = time.monotonic() + 3
while time.monotonic() < deadline:
detail = self.manager.detail(session["id"])
if detail["state"] == "awaiting_approval":
break
time.sleep(0.01)
else:
self.fail("session did not wait for approval")
changed = self.manager.set_mode(session["id"], {"mode": "automatic"})
self.assertEqual(changed["mode"], "automatic")
self.assertEqual(changed["proposals"][-1]["state"], "approved")
completed = self.wait_terminal(session["id"])
self.assertEqual(completed["state"], "succeeded", completed["message"])
def test_trial_count_is_user_configurable(self):
payload = self.payload()
payload["trialCount"] = 1
_, config, _, _ = self.manager.parse_create(payload)
self.assertEqual(config["trialCount"], 1)
payload["trialCount"] = 100
_, config, _, _ = self.manager.parse_create(payload)
self.assertEqual(config["trialCount"], 100)
payload["trialCount"] = 101
with self.assertRaisesRegex(Exception, "trialCount"):
self.manager.parse_create(payload)
def test_resume_discards_only_interrupted_trial_and_continues(self):
mode, config, objective, fallback = self.manager.parse_create(self.payload())
session = self.manager.storage.create_session(mode, config, objective, fallback)
baseline = self.manager.storage.create_trial(
session["id"],
0,
0,
1,
BASE_REWARD_CONFIGURATION,
None,
"trial-000-rung-0",
)
self.manager._execute_trial(self.manager.storage.get_session(session["id"]), baseline)
interrupted = self.manager.storage.create_trial(
session["id"],
1,
0,
1,
baseline["rewardConfig"],
None,
"trial-001-rung-0",
)
self.manager.storage.update_trial(interrupted["id"], state="interrupted")
self.manager.storage.reset_step_gate(session["id"])
self.manager.storage.update_session(session["id"], state="interrupted")
self.manager.resume(session["id"])
self.assertEqual(self.manager.storage.get_control(session["id"])["runPolicy"], "continuous")
completed = self.wait_terminal(session["id"])
self.assertEqual(completed["state"], "succeeded", completed["message"])
self.assertNotIn(interrupted["id"], [trial["id"] for trial in completed["trials"]])
def test_create_validation_and_agent_capability(self):
self.assertTrue(self.manager.capability()["configured"])
with self.assertRaisesRegex(Exception, "只支持"):
self.manager.parse_create({"taskId": "Other"})
self.assertEqual(self.manager.test_agent()["model"], "fake")
def test_pause_closes_persistent_dispatch_gate_at_trial_boundary(self):
mode, config, objective, fallback = self.manager.parse_create(self.payload())
session = self.manager.storage.create_session(mode, config, objective, fallback)
self.manager.storage.update_session(session["id"], state="running")
self.manager.storage.grant_dispatch_token(session["id"])
paused = self.manager.pause(session["id"])
self.assertEqual(paused["state"], "paused")
self.assertEqual(paused["control"]["runPolicy"], "step")
self.assertEqual(paused["control"]["dispatchTokens"], 0)
def test_step_token_executes_exactly_one_trial_then_pauses(self):
session = self.manager.create(self.payload("approval"))
deadline = time.monotonic() + 3
while time.monotonic() < deadline:
detail = self.manager.detail(session["id"])
if detail["state"] == "awaiting_approval":
break
time.sleep(0.01)
else:
self.fail("session did not wait for approval")
baseline_count = len([trial for trial in detail["trials"] if trial["state"] == "completed"])
stepped = self.manager.step(session["id"], {"count": 1})
self.assertEqual(stepped["control"]["runPolicy"], "step")
self.assertEqual(stepped["control"]["dispatchTokens"], 1)
proposal = stepped["proposals"][-1]
self.manager.approve(session["id"], proposal["id"], {})
deadline = time.monotonic() + 3
while time.monotonic() < deadline:
detail = self.manager.detail(session["id"])
completed_count = len(
[trial for trial in detail["trials"] if trial["state"] == "completed"]
)
if detail["state"] == "paused" and completed_count == baseline_count + 1:
break
time.sleep(0.01)
else:
self.fail("single-step trial did not pause at the next boundary")
time.sleep(0.05)
self.assertEqual(
len(
[
trial
for trial in self.manager.detail(session["id"])["trials"]
if trial["state"] == "completed"
]
),
baseline_count + 1,
)
self.manager.cancel(session["id"])
self.assertEqual(self.wait_terminal(session["id"])["state"], "cancelled")
def test_constraints_are_revisioned_and_enforced_during_approval(self):
session = self.manager.create(self.payload("approval"))
deadline = time.monotonic() + 3
while time.monotonic() < deadline:
detail = self.manager.detail(session["id"])
if detail["state"] == "awaiting_approval":
break
time.sleep(0.01)
else:
self.fail("session did not wait for approval")
proposal = detail["proposals"][-1]
constrained = self.manager.set_constraints(
session["id"],
{
"revision": 0,
"constraints": {"weights.track_linear_velocity": {"kind": "fixed", "value": 1.0}},
},
)
self.assertEqual(constrained["control"]["constraintsRevision"], 1)
with self.assertRaisesRegex(Exception, "已固定"):
self.manager.approve(
session["id"],
proposal["id"],
{"patch": {"weights": {"track_linear_velocity": 1.1}, "params": {}}},
)
with self.assertRaisesRegex(Exception, "revision"):
self.manager.set_constraints(session["id"], {"revision": 0, "constraints": {}})
self.manager.approve(session["id"], proposal["id"], {})
self.manager.cancel(session["id"])
self.assertEqual(self.wait_terminal(session["id"])["state"], "cancelled")
def test_rollback_uses_safe_completed_trial_as_next_proposal_base(self):
session = self.manager.create(self.payload("approval"))
deadline = time.monotonic() + 3
while time.monotonic() < deadline:
detail = self.manager.detail(session["id"])
if detail["state"] == "awaiting_approval":
break
time.sleep(0.01)
else:
self.fail("session did not wait for approval")
baseline = detail["trials"][0]
old_proposal = detail["proposals"][-1]
rolled_back = self.manager.rollback(
session["id"], {"trialId": baseline["id"], "checkpoint": True}
)
self.assertEqual(rolled_back["state"], "paused")
self.assertEqual(rolled_back["control"]["activeBaseTrialId"], baseline["id"])
self.assertEqual(
next(item for item in rolled_back["proposals"] if item["id"] == old_proposal["id"])[
"state"
],
"rejected",
)
self.manager.step(session["id"], {"count": 1})
deadline = time.monotonic() + 3
while time.monotonic() < deadline:
detail = self.manager.detail(session["id"])
if (
detail["state"] == "awaiting_approval"
and detail["proposals"][-1]["id"] != old_proposal["id"]
):
break
time.sleep(0.01)
else:
self.fail("rollback base did not produce a replacement proposal")
self.assertEqual(detail["proposals"][-1]["baseTrialId"], baseline["id"])
self.manager.cancel(session["id"])
self.assertEqual(self.wait_terminal(session["id"])["state"], "cancelled")
if __name__ == "__main__":
unittest.main()
-26
View File
@@ -1,26 +0,0 @@
"""Reward auto-tuning support for the local training service."""
from .schema import (
BASE_REWARD_CONFIGURATION,
PARAMETER_SPECS,
WEIGHT_SPECS,
RewardConfigError,
apply_reward_configuration,
merge_proposal,
validate_configuration,
validate_proposal,
)
from .scoring import DEFAULT_OBJECTIVE_WEIGHTS, score_evaluation
__all__ = [
"BASE_REWARD_CONFIGURATION",
"DEFAULT_OBJECTIVE_WEIGHTS",
"PARAMETER_SPECS",
"WEIGHT_SPECS",
"RewardConfigError",
"apply_reward_configuration",
"merge_proposal",
"score_evaluation",
"validate_configuration",
"validate_proposal",
]
-134
View File
@@ -1,134 +0,0 @@
"""PydanticAI adapter for the DeepSeek reward-tuning advisor."""
from __future__ import annotations
import hashlib
import json
import os
from dataclasses import dataclass
from typing import Any
from .schema import validate_proposal
SYSTEM_PROMPT = """你是 Unitree Go2 强化学习奖励调参专家。
只根据提供的数值配置、训练曲线摘要和固定评估结果提出下一轮稀疏修改。
必须优先保持速度跟踪与跌倒安全门槛;每轮最多修改四个白名单标量,不得改变符号、函数、传感器或结构。
不要建议 Python 代码、命令、文件路径或白名单外参数。输出必须符合 RewardProposal schema。
"""
class AdvisorUnavailable(RuntimeError):
pass
@dataclass(frozen=True)
class AdvisorConfig:
api_key: str | None
base_url: str = "https://api.deepseek.com"
model: str = "deepseek-v4-flash"
@classmethod
def from_environment(cls) -> AdvisorConfig:
return cls(
api_key=os.environ.get("DEEPSEEK_API_KEY"),
base_url=os.environ.get("MUJOCO_TUNING_AGENT_BASE_URL", "https://api.deepseek.com"),
model=os.environ.get("MUJOCO_TUNING_AGENT_MODEL", "deepseek-v4-flash"),
)
class DeepSeekAdvisor:
def __init__(self, config: AdvisorConfig | None = None):
self.config = config or AdvisorConfig.from_environment()
self._cached_agent = None
def capability(self) -> dict[str, Any]:
try:
import pydantic_ai # noqa: F401
except ImportError:
installed = False
else:
installed = True
return {
"configured": bool(self.config.api_key) and installed,
"apiKeyConfigured": bool(self.config.api_key),
"frameworkInstalled": installed,
"model": self.config.model,
"baseUrl": self.config.base_url,
}
def _agent(self):
if self._cached_agent is not None:
return self._cached_agent
if not self.config.api_key:
raise AdvisorUnavailable("未配置 DEEPSEEK_API_KEY")
try:
import httpx2
from pydantic import BaseModel, Field
from pydantic_ai import Agent, PromptedOutput
from pydantic_ai.models.openai import OpenAIChatModel
from pydantic_ai.providers.openai import OpenAIProvider
except ImportError as error:
raise AdvisorUnavailable(
"缺少 PydanticAI,请安装 training_server/requirements.txt"
) from error
class RewardProposalOutput(BaseModel):
weights: dict[str, float] = Field(default_factory=dict)
params: dict[str, float] = Field(default_factory=dict)
rationale: str = Field(min_length=1, max_length=2000)
expected_impact: dict[str, str] = Field(default_factory=dict)
confidence: float = Field(ge=0.0, le=1.0)
proxy = os.environ.get("HTTPS_PROXY") or os.environ.get("ALL_PROXY")
if proxy and proxy.startswith("socks://"):
proxy = "socks5://" + proxy.removeprefix("socks://")
http_client = httpx2.AsyncClient(proxy=proxy, trust_env=False, timeout=60.0)
provider = OpenAIProvider(
base_url=self.config.base_url, api_key=self.config.api_key, http_client=http_client
)
model = OpenAIChatModel(self.config.model, provider=provider) # type: ignore[arg-type]
self._cached_agent = Agent(
model,
output_type=PromptedOutput(RewardProposalOutput),
system_prompt=SYSTEM_PROMPT,
retries=2,
model_settings={"temperature": 0.2},
)
return self._cached_agent
def propose(self, context: dict[str, Any], previous: dict) -> dict[str, Any]:
prompt = json.dumps(context, ensure_ascii=False, separators=(",", ":"), allow_nan=False)
result = self._agent().run_sync(prompt)
output = result.output
patch = validate_proposal(
{"weights": dict(output.weights), "params": dict(output.params)}, previous
)
try:
usage = result.usage()
usage_value = {
key: getattr(usage, key)
for key in ("requests", "input_tokens", "output_tokens", "total_tokens")
if getattr(usage, key, None) is not None
}
except (AttributeError, TypeError):
usage_value = {}
return {
"patch": patch,
"rationale": output.rationale,
"expectedImpact": dict(output.expected_impact),
"confidence": float(output.confidence),
"promptHash": hashlib.sha256(prompt.encode()).hexdigest(),
"usage": usage_value,
"model": self.config.model,
}
def test_connection(self) -> dict[str, Any]:
base = {
"weights": {"track_linear_velocity": 1.0},
"params": {},
"instruction": "仅返回一个合法示例:把 track_linear_velocity 改为 1.1。",
}
# A minimal full previous config is supplied by callers for actual proposals;
# connectivity probing only verifies the provider and structured response path.
result = self._agent().run_sync(json.dumps(base, ensure_ascii=False))
return {"ok": True, "model": self.config.model, "outputType": type(result.output).__name__}
-999
View File
@@ -1,999 +0,0 @@
"""Persistent reward tuning session orchestrator."""
from __future__ import annotations
import json
import os
import re
import shutil
import subprocess
import threading
from contextlib import suppress
from copy import deepcopy
from pathlib import Path
from typing import Any
from .advisor import AdvisorUnavailable, DeepSeekAdvisor
from .process import GpuLease, ResourceBusyError, terminate_process
from .schema import (
BASE_REWARD_CONFIGURATION,
merge_proposal,
validate_configuration_constraints,
validate_constraints,
validate_proposal,
)
from .scoring import DEFAULT_OBJECTIVE_WEIGHTS, score_evaluation, validate_objective_weights
from .storage import StorageConflict, TuningStorage, now_iso
from .study import OptunaStudies
from .tensorboard import ingest_scalars
ACTIVE_SESSION_STATES = {
"queued",
"running",
"evaluating",
"awaiting_approval",
"paused",
"interrupted",
}
RUNNING_STATES = {"queued", "running", "evaluating"}
RUN_NAME = re.compile(r"^[A-Za-z0-9_.-]{1,64}$")
class TuningError(RuntimeError):
pass
class TuningManager:
def __init__(
self,
trainer_root: Path,
python: str,
data_root: Path,
lease: GpuLease,
advisor: Any | None = None,
):
self.trainer_root = trainer_root.expanduser().resolve()
self.python = python
self.data_root = data_root.expanduser().resolve()
self.data_root.mkdir(parents=True, exist_ok=True)
self.storage = TuningStorage(self.data_root / "tuning.sqlite3")
self.storage.recover_interrupted()
self.studies = OptunaStudies(self.data_root)
self.lease = lease
self.advisor = advisor or DeepSeekAdvisor()
self.lock = threading.RLock()
self.condition = threading.Condition(self.lock)
self.workers: dict[str, threading.Thread] = {}
self.processes: dict[str, subprocess.Popen[str]] = {}
self.cancel_events: dict[str, threading.Event] = {}
def capability(self) -> dict[str, Any]:
capability = self.advisor.capability()
capability.update({"ready": (self.trainer_root / "scripts" / "evaluate.py").is_file()})
return capability
@staticmethod
def _integer(payload: dict, name: str, default: int, minimum: int, maximum: int) -> int:
value = payload.get(name, default)
if isinstance(value, bool) or not isinstance(value, int) or not minimum <= value <= maximum:
raise TuningError(f"{name} 必须在 {minimum}–{maximum} 之间")
return value
def parse_create(self, payload: Any) -> tuple[str, dict, dict, bool]:
if not isinstance(payload, dict):
raise TuningError("请求体必须是 JSON 对象")
mode = payload.get("mode", "automatic")
if mode not in ("automatic", "approval"):
raise TuningError("mode 必须是 automatic 或 approval")
if payload.get("taskId", "Unitree-Go2-Flat") != "Unitree-Go2-Flat":
raise TuningError("第一版只支持 Unitree-Go2-Flat")
run_name = payload.get("runName", "auto-tune")
if not isinstance(run_name, str) or not RUN_NAME.fullmatch(run_name):
raise TuningError("runName 格式无效")
gpu_ids = payload.get("gpuIds", [0])
if (
not isinstance(gpu_ids, list)
or not gpu_ids
or any(
isinstance(value, bool) or not isinstance(value, int) or not 0 <= value <= 255
for value in gpu_ids
)
):
raise TuningError("gpuIds 必须是非空非负整数数组")
trial_count = self._integer(payload, "trialCount", 12, 1, 100)
rung0 = self._integer(payload, "initialIterations", 300, 1, 1000000)
rung1 = self._integer(payload, "middleIterations", 900, rung0, 1000000)
rung2 = self._integer(payload, "finalIterations", 2000, rung1, 1000000)
config = {
"taskId": "Unitree-Go2-Flat",
"numEnvs": self._integer(payload, "numEnvs", 4096, 1, 16384),
"seed": self._integer(payload, "seed", 42, 0, 2147483647),
"runName": run_name,
"gpuIds": gpu_ids,
"trialCount": trial_count,
"rungs": [rung0, rung1, rung2],
"promote": [trial_count, min(4, trial_count), min(2, trial_count)],
"evalNumEnvs": self._integer(payload, "evalNumEnvs", 256, 1, 4096),
"evalSteps": self._integer(payload, "evalSteps", 1000, 10, 100000),
"earlyStopPatience": self._integer(payload, "earlyStopPatience", 4, 1, 20),
}
objective = validate_objective_weights(
payload.get("objectiveWeights", DEFAULT_OBJECTIVE_WEIGHTS)
)
fallback = payload.get("fallbackEnabled", False)
if not isinstance(fallback, bool):
raise TuningError("fallbackEnabled 必须是布尔值")
return mode, config, objective, fallback
def create(self, payload: Any) -> dict:
mode, config, objective, fallback = self.parse_create(payload)
capability = self.capability()
if not capability["ready"]:
raise TuningError("评估入口未就绪")
if not capability["configured"] and not fallback:
raise TuningError("DeepSeek Agent 未配置;设置 DEEPSEEK_API_KEY 或显式启用 fallback")
with self.lock:
active = [
session
for session in self.storage.list_sessions()
if session["state"] in ACTIVE_SESSION_STATES
]
if active:
raise ResourceBusyError("已有调参 session 未结束")
session = self.storage.create_session(mode, config, objective, fallback)
self._start_worker(session["id"], resume=False)
return self.detail(session["id"])
def _start_worker(self, session_id: str, resume: bool) -> None:
cancel = threading.Event()
self.cancel_events[session_id] = cancel
worker = threading.Thread(
target=self._run_session,
args=(session_id, resume, cancel),
name=f"tuning-{session_id[:8]}",
daemon=True,
)
self.workers[session_id] = worker
worker.start()
def detail(self, session_id: str) -> dict:
session = self.storage.get_session(session_id)
session["trials"] = self.storage.list_trials(session_id)
session["proposals"] = self.storage.list_proposals(session_id)
session["audit"] = self.storage.audit_events(session_id)
control = self.storage.get_control(session_id)
current = next(
(trial for trial in session["trials"] if trial["id"] == session["currentTrialId"]),
None,
)
control["effectiveAfterCurrent"] = bool(
current and current["state"] in {"training", "evaluating"}
)
session["control"] = control
return session
def list(self) -> list[dict]:
return self.storage.list_sessions()
def _session_root(self, session_id: str) -> Path:
root = (self.data_root / "sessions" / session_id).resolve()
if not root.is_relative_to(self.data_root):
raise TuningError("非法 session 路径")
root.mkdir(parents=True, exist_ok=True)
return root
def _run_command(
self,
session_id: str,
command: list[str],
cwd: Path,
environment: dict[str, str],
log_path: Path,
) -> int:
owner = f"tuning:{session_id}"
self.lease.acquire(owner)
try:
process = subprocess.Popen(
command,
cwd=cwd,
env=environment,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
text=True,
encoding="utf-8",
errors="replace",
bufsize=1,
start_new_session=True,
)
with self.lock:
self.processes[session_id] = process
assert process.stdout is not None
with log_path.open("a", encoding="utf-8") as log:
for line in process.stdout:
log.write(line)
log.flush()
if self.cancel_events[session_id].is_set():
terminate_process(process)
break
return process.wait()
finally:
with self.lock:
self.processes.pop(session_id, None)
self.lease.release(owner)
@staticmethod
def _latest_checkpoint(run_dir: Path) -> Path | None:
values = []
for path in run_dir.glob("model_*.pt"):
match = re.fullmatch(r"model_(\d+)\.pt", path.name)
if match:
values.append((int(match.group(1)), path))
return max(values, default=(0, None), key=lambda value: value[0])[1]
def _execute_trial(
self, session: dict, trial: dict, resume_checkpoint: Path | None = None
) -> dict:
session_id, trial_id = session["id"], trial["id"]
root = self._session_root(session_id)
run_dir = (root / trial["runDir"]).resolve()
if not run_dir.is_relative_to(root):
raise TuningError("trial 目录越界")
run_dir.mkdir(parents=True, exist_ok=True)
reward_path = run_dir / "reward_config.json"
reward_path.write_text(
json.dumps(trial["rewardConfig"], ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
)
config = session["config"]
command = [
self.python,
"-u",
"scripts/train.py",
config["taskId"],
f"--env.scene.num-envs={config['numEnvs']}",
f"--agent.max-iterations={trial['targetIterations']}",
f"--agent.seed={config['seed']}",
f"--agent.run-name={config['runName']}-t{trial['number']}-r{trial['rung']}",
"--agent.logger=tensorboard",
"--agent.upload-model=False",
"--gpu-ids",
json.dumps(config["gpuIds"], separators=(",", ":")),
"--output-dir",
str(run_dir),
"--reward-config",
str(reward_path),
]
if resume_checkpoint is not None:
command.extend(("--resume-checkpoint", str(resume_checkpoint)))
environment = os.environ.copy()
environment["WANDB_MODE"] = "disabled"
environment["WANDB_SILENT"] = "true"
self.storage.update_trial(
trial_id, state="training", started_at=now_iso(), message="正在训练"
)
self.storage.update_session(
session_id,
state="running",
message=f"正在训练 trial {trial['number']} / rung {trial['rung']}",
)
return_code = self._run_command(
session_id, command, self.trainer_root, environment, run_dir / "train.log"
)
ingest_scalars(self.storage, trial_id, run_dir)
if self.cancel_events[session_id].is_set():
raise TuningError("session 已取消")
checkpoint = self._latest_checkpoint(run_dir)
policy = run_dir / "policy.onnx"
if return_code != 0 or checkpoint is None or not policy.is_file():
raise TuningError(f"训练失败(返回码 {return_code})或缺少 checkpoint/policy.onnx")
# 暂停只关闭后续 Trial 调度门;已经开始的 Trial 必须连同固定评估一起
# 完成,避免把一次单步令牌错误地消耗在半个 Trial 上。
eval_output = run_dir / "evaluation.json"
eval_command = [
self.python,
"-u",
"scripts/evaluate.py",
config["taskId"],
"--checkpoint",
str(checkpoint),
"--output",
str(eval_output),
"--reward-config",
str(reward_path),
f"--num-envs={config['evalNumEnvs']}",
f"--steps-per-seed={config['evalSteps']}",
"--gpu-ids",
json.dumps(config["gpuIds"], separators=(",", ":")),
]
self.storage.update_trial(trial_id, state="evaluating", message="正在固定协议评估")
self.storage.update_session(
session_id, state="evaluating", message=f"正在评估 trial {trial['number']}"
)
return_code = self._run_command(
session_id, eval_command, self.trainer_root, environment, run_dir / "evaluate.log"
)
ingest_scalars(self.storage, trial_id, run_dir / "evaluation-events")
if return_code != 0 or not eval_output.is_file():
raise TuningError(f"评估失败(返回码 {return_code})")
evaluation = json.loads(eval_output.read_text(encoding="utf-8"))
baseline_trial = self.storage.list_trials(session_id)[0]
if baseline_trial["evaluation"] is None:
scored = {
"eligible": True,
"score": 0.0,
"components": {},
"metrics": evaluation["metrics"],
}
else:
scored = score_evaluation(
baseline_trial["evaluation"]["metrics"],
evaluation["metrics"],
session["objectiveWeights"],
)
evaluation["score"] = scored
rel_checkpoint = str(checkpoint.relative_to(root))
rel_policy = str(policy.relative_to(root))
self.storage.update_trial(
trial_id,
state="completed",
ended_at=now_iso(),
message="训练与评估完成",
checkpoint_path=rel_checkpoint,
policy_path=rel_policy,
evaluation=evaluation,
score=scored["score"],
eligible=scored["eligible"],
)
best = self._best_highest_rung(session_id)
if best is not None:
self.storage.update_session(session_id, best_trial_id=best["id"])
try:
optuna_number = self.studies.record(
session_id,
trial["rewardConfig"],
scored["score"],
scored["eligible"],
trial["rung"],
)
self.storage.audit(
session_id, "optuna_trial_recorded", {"trialId": trial_id, "number": optuna_number}
)
except Exception as error:
self.storage.audit(
session_id, "optuna_record_failed", {"trialId": trial_id, "error": str(error)[:500]}
)
return self.storage.get_trial(trial_id)
def _best(self, session_id: str, rung: int | None = None) -> dict | None:
trials = [
trial
for trial in self.storage.list_trials(session_id)
if trial["state"] == "completed" and trial["eligible"]
]
if rung is not None:
trials = [trial for trial in trials if trial["rung"] == rung]
return max(
trials,
key=lambda trial: trial["score"] if trial["score"] is not None else -999,
default=None,
)
def _best_highest_rung(self, session_id: str) -> dict | None:
for rung in (2, 1, 0):
best = self._best(session_id, rung=rung)
if best is not None:
return best
return None
def _proposal_context(self, session: dict) -> dict:
trials = self.storage.list_trials(session["id"])[-12:]
rejected_feedback = [
proposal["feedback"]
for proposal in self.storage.list_proposals(session["id"])
if proposal["state"] == "rejected" and proposal["feedback"]
][-4:]
return {
"task": session["config"]["taskId"],
"objectiveWeights": session["objectiveWeights"],
"allowlist": "服务端将验证固定 schema;最多四项修改",
"parameterConstraints": self.storage.get_control(session["id"])["constraints"],
"rejectedFeedback": rejected_feedback,
"trials": [
{
"number": t["number"],
"rung": t["rung"],
"score": t["score"],
"eligible": t["eligible"],
"rewardConfig": t["rewardConfig"],
"evaluation": t["evaluation"] and t["evaluation"].get("metrics"),
}
for t in trials
],
}
def _fallback_patch(self, previous: dict, index: int) -> dict:
names = ("track_linear_velocity", "action_rate_l2", "body_orientation_l2", "foot_slip")
name = names[index % len(names)]
old = previous["weights"][name]
factor = 1.1 if index % 2 == 0 else 0.9
return validate_proposal({"weights": {name: old * factor}}, previous)
def _request_proposal(
self, session: dict, previous: dict, base_trial_id: str, index: int
) -> dict:
try:
result = self.advisor.propose(self._proposal_context(session), previous)
source = "agent"
except Exception as error:
if not session["fallbackEnabled"]:
raise AdvisorUnavailable(str(error)) from error
result = {
"patch": self._fallback_patch(previous, index),
"rationale": f"Agent 不可用,显式 fallback:{error}",
"expectedImpact": {},
"confidence": 0.2,
"promptHash": None,
"usage": {},
"model": "optuna-fallback",
}
source = "fallback"
constraints = self.storage.get_control(session["id"])["constraints"]
result["patch"] = validate_proposal(result["patch"], previous, constraints)
proposal = self.storage.create_proposal(
session["id"],
base_trial_id,
result["patch"],
result["rationale"],
result.get("expectedImpact", {}),
result["confidence"],
source,
)
self.storage.audit(
session["id"],
"proposal_created",
{
"proposalId": proposal["id"],
"source": source,
"promptHash": result.get("promptHash"),
"usage": result.get("usage", {}),
"model": result.get("model"),
},
)
return proposal
def _wait_for_approval(
self, session_id: str, proposal_id: str, cancel: threading.Event
) -> dict:
with self.condition:
while not cancel.is_set():
proposal = self.storage.get_proposal(proposal_id)
if proposal["state"] != "pending":
return proposal
self.condition.wait(timeout=1.0)
raise TuningError("session 已取消")
def _run_session(self, session_id: str, resume: bool, cancel: threading.Event) -> None:
try:
session = self.storage.get_session(session_id)
trials = self.storage.list_trials(session_id)
if resume:
root = self._session_root(session_id)
for interrupted in [trial for trial in trials if trial["state"] == "interrupted"]:
run_dir = (root / interrupted["runDir"]).resolve()
if run_dir.is_relative_to(root):
shutil.rmtree(run_dir, ignore_errors=True)
self.storage.delete_trial(interrupted["id"])
self.storage.audit(
session_id,
"session_resumed",
{
"discardedInterruptedTrials": [
t["id"] for t in trials if t["state"] == "interrupted"
]
},
)
trials = self.storage.list_trials(session_id)
if not trials:
baseline_dir = "trial-000-rung-0"
trial = self.storage.create_trial(
session_id,
0,
0,
session["config"]["rungs"][0],
deepcopy(BASE_REWARD_CONFIGURATION),
None,
baseline_dir,
)
self._execute_trial(session, trial)
session = self.storage.get_session(session_id)
completed_rung0 = [
t
for t in self.storage.list_trials(session_id)
if t["rung"] == 0 and t["state"] == "completed"
]
next_number = len({t["number"] for t in completed_rung0})
best_score = max(
(t["score"] if t["score"] is not None else float("-inf") for t in completed_rung0),
default=0.0,
)
no_improve = 0
while (
next_number < session["config"]["trialCount"]
and no_improve < session["config"]["earlyStopPatience"]
):
if cancel.is_set():
raise TuningError("session 已取消")
self._wait_for_dispatch(session_id, cancel)
control = self.storage.get_control(session_id)
base = (
self.storage.get_trial(control["activeBaseTrialId"])
if control["activeBaseTrialId"]
else self._best(session_id, rung=0) or completed_rung0[0]
)
proposal = self._request_proposal(
session, base["rewardConfig"], base["id"], next_number
)
# 模式允许在 session 运行期间切换,因此每次决策都读取最新持久化值。
current_mode = self.storage.get_session(session_id)["mode"]
if current_mode == "approval":
self.storage.update_session(
session_id, state="awaiting_approval", message="等待批准 Agent 建议"
)
proposal = self._wait_for_approval(session_id, proposal["id"], cancel)
if proposal["state"] == "rejected":
self.storage.audit(
session_id,
"proposal_rejected",
{"proposalId": proposal["id"], "feedback": proposal["feedback"]},
)
continue
else:
self.storage.decide_proposal(proposal["id"], "approved", "自动模式")
proposal = self.storage.get_proposal(proposal["id"])
self._claim_dispatch(session_id, cancel)
constraints = self.storage.get_control(session_id)["constraints"]
reward_config = merge_proposal(base["rewardConfig"], proposal["patch"], constraints)
trial = self.storage.create_trial(
session_id,
next_number,
0,
session["config"]["rungs"][0],
reward_config,
proposal["id"],
f"trial-{next_number:03d}-rung-0",
)
result = self._execute_trial(session, trial)
if control["activeBaseTrialId"]:
self.storage.set_active_base(session_id, None)
self._pause_after_step(session_id)
result_score = result["score"]
if (
result["eligible"]
and result_score is not None
and result_score > best_score + 0.01
):
best_score = result_score
no_improve = 0
else:
no_improve += 1
self.storage.update_session(session_id, consecutive_no_improve=no_improve)
next_number += 1
# Promote top configurations; each new rung resumes its own previous checkpoint.
for rung in (1, 2):
previous = [
t
for t in self.storage.list_trials(session_id)
if t["rung"] == rung - 1 and t["state"] == "completed" and t["eligible"]
]
previous.sort(
key=lambda trial: (
trial["score"] if trial["score"] is not None else float("-inf")
),
reverse=True,
)
promoted_numbers = {
trial["number"]
for trial in self.storage.list_trials(session_id)
if trial["rung"] == rung and trial["state"] == "completed"
}
for parent in previous[: session["config"]["promote"][rung]]:
if parent["number"] in promoted_numbers:
continue
if cancel.is_set():
raise TuningError("session 已取消")
self._claim_dispatch(session_id, cancel)
root = self._session_root(session_id)
checkpoint = root / parent["checkpointPath"]
trial = self.storage.create_trial(
session_id,
parent["number"],
rung,
session["config"]["rungs"][rung],
parent["rewardConfig"],
parent["proposalId"],
f"trial-{parent['number']:03d}-rung-{rung}",
)
self._execute_trial(session, trial, checkpoint)
self._pause_after_step(session_id)
best = self._best_highest_rung(session_id)
if best is None:
raise TuningError("没有通过安全门槛的 trial")
preset_name = f"{session['config']['runName']}-{session_id[:8]}"
self.storage.save_preset(preset_name, session_id, best["id"], best["rewardConfig"])
self.storage.update_session(
session_id,
state="succeeded",
best_trial_id=best["id"],
current_trial_id=None,
message="调参完成",
)
self.storage.audit(
session_id, "session_completed", {"bestTrialId": best["id"], "preset": preset_name}
)
except Exception as error:
state = self.storage.get_session(session_id)["state"]
if cancel.is_set() or state == "cancelled":
self.storage.update_session(
session_id, state="cancelled", message="调参已取消", current_trial_id=None
)
else:
self.storage.update_session(
session_id, state="failed", message=str(error), current_trial_id=None
)
self.storage.audit(session_id, "session_failed", {"error": str(error)[:1000]})
finally:
with self.lock:
self.workers.pop(session_id, None)
self.processes.pop(session_id, None)
def _wait_for_dispatch(self, session_id: str, cancel: threading.Event) -> None:
"""Wait at a scheduler boundary without consuming a one-shot token."""
with self.condition:
while not cancel.is_set():
session = self.storage.get_session(session_id)
control = self.storage.get_control(session_id)
if session["state"] == "paused":
self.condition.wait(timeout=1.0)
continue
if control["runPolicy"] == "continuous" or control["dispatchTokens"] > 0:
return
self.storage.update_session(
session_id,
state="paused",
message="单步 Trial 已完成;等待下一个调度令牌",
)
self.storage.audit(session_id, "step_gate_waiting", {})
self.condition.wait(timeout=1.0)
raise TuningError("session 已取消")
def _claim_dispatch(self, session_id: str, cancel: threading.Event) -> None:
while not cancel.is_set():
self._wait_for_dispatch(session_id, cancel)
if self.storage.use_dispatch_token(session_id):
return
raise TuningError("session 已取消")
def _pause_after_step(self, session_id: str) -> None:
control = self.storage.get_control(session_id)
if control["runPolicy"] == "step" and control["dispatchTokens"] == 0:
self.storage.update_session(
session_id,
state="paused",
message="单步 Trial 已完成;后续调度已暂停",
)
self.storage.audit(session_id, "step_trial_completed", {})
def approve(self, session_id: str, proposal_id: str, payload: Any) -> dict:
proposal = self.storage.get_proposal(proposal_id)
if proposal["sessionId"] != session_id:
raise TuningError("proposal 不属于该 session")
patch = proposal["patch"]
feedback = None
base = self.storage.get_trial(proposal["baseTrialId"])
constraints = self.storage.get_control(session_id)["constraints"]
if isinstance(payload, dict):
feedback = payload.get("feedback")
if "patch" in payload:
patch = payload["patch"]
if feedback is not None and (not isinstance(feedback, str) or len(feedback) > 2000):
raise TuningError("feedback 无效")
patch = validate_proposal(patch, base["rewardConfig"], constraints)
if not self.storage.decide_proposal(proposal_id, "approved", feedback, patch):
raise TuningError("proposal 已处理")
self.storage.audit(
session_id,
"proposal_approved",
{"proposalId": proposal_id, "modified": patch != proposal["patch"]},
)
with self.condition:
self.condition.notify_all()
return self.detail(session_id)
def reject(self, session_id: str, proposal_id: str, payload: Any) -> dict:
feedback = payload.get("feedback", "") if isinstance(payload, dict) else ""
if not isinstance(feedback, str) or len(feedback) > 2000:
raise TuningError("feedback 无效")
proposal = self.storage.get_proposal(proposal_id)
if proposal["sessionId"] != session_id:
raise TuningError("proposal 不属于该 session")
if not self.storage.decide_proposal(proposal_id, "rejected", feedback):
raise TuningError("proposal 已处理")
with self.condition:
self.condition.notify_all()
return self.detail(session_id)
def set_mode(self, session_id: str, payload: Any) -> dict:
if not isinstance(payload, dict) or payload.get("mode") not in ("automatic", "approval"):
raise TuningError("mode 必须是 automatic 或 approval")
mode = payload["mode"]
with self.condition:
session = self.storage.get_session(session_id)
if session["state"] not in RUNNING_STATES | {"awaiting_approval", "paused"}:
raise TuningError("当前状态不能切换运行模式")
previous = session["mode"]
if previous == mode:
return self.detail(session_id)
self.storage.update_session(session_id, mode=mode)
approved_ids = []
if mode == "automatic":
constraints = self.storage.get_control(session_id)["constraints"]
for proposal in self.storage.list_proposals(session_id):
if proposal["state"] != "pending":
continue
base = self.storage.get_trial(proposal["baseTrialId"])
try:
validate_proposal(proposal["patch"], base["rewardConfig"], constraints)
except Exception as error:
self.storage.decide_proposal(
proposal["id"], "rejected", f"参数护栏已变化:{error}"
)
continue
if self.storage.decide_proposal(
proposal["id"], "approved", "运行时切换为全自动模式"
):
approved_ids.append(proposal["id"])
if session["state"] == "awaiting_approval":
self.storage.update_session(
session_id, state="running", message="已切换为全自动模式,继续调参"
)
self.storage.audit(
session_id,
"session_mode_changed",
{"from": previous, "to": mode, "autoApprovedProposalIds": approved_ids},
)
self.condition.notify_all()
return self.detail(session_id)
def set_constraints(self, session_id: str, payload: Any) -> dict:
if not isinstance(payload, dict) or set(payload) != {"revision", "constraints"}:
raise TuningError("参数护栏请求必须包含 revision 与 constraints")
revision = payload["revision"]
if isinstance(revision, bool) or not isinstance(revision, int) or revision < 0:
raise TuningError("constraints revision 必须是非负整数")
constraints = validate_constraints(payload["constraints"])
session = self.storage.get_session(session_id)
if session["state"] not in ACTIVE_SESSION_STATES:
raise TuningError("终态 session 不能修改参数护栏")
control = self.storage.get_control(session_id)
base = None
if control["activeBaseTrialId"]:
base = self.storage.get_trial(control["activeBaseTrialId"])
elif session["currentTrialId"]:
base = self.storage.get_trial(session["currentTrialId"])
else:
base = self._best_highest_rung(session_id)
if base is not None:
validate_configuration_constraints(base["rewardConfig"], constraints)
try:
updated = self.storage.replace_constraints(session_id, revision, constraints)
except StorageConflict as error:
raise ResourceBusyError(str(error)) from error
rejected = []
for proposal in self.storage.list_proposals(session_id):
if proposal["state"] != "pending":
continue
proposal_base = self.storage.get_trial(proposal["baseTrialId"])
try:
validate_proposal(proposal["patch"], proposal_base["rewardConfig"], constraints)
except Exception as error:
message = f"参数护栏 revision {updated['constraintsRevision']}:{error}"
if self.storage.decide_proposal(proposal["id"], "rejected", message):
rejected.append(proposal["id"])
self.storage.audit(
session_id,
"constraints_updated",
{
"revision": updated["constraintsRevision"],
"paths": sorted(constraints),
"rejectedProposalIds": rejected,
},
)
if rejected:
with self.condition:
self.condition.notify_all()
return self.detail(session_id)
def step(self, session_id: str, payload: Any) -> dict:
if not isinstance(payload, dict) or payload.get("count", 1) != 1:
raise TuningError("单步调度一次只能发放 1 个 Trial 令牌")
session = self.storage.get_session(session_id)
if session["state"] not in {"paused", "awaiting_approval"}:
raise TuningError("请先暂停或等待 Proposal 审批,再执行单步 Trial")
control = self.storage.get_control(session_id)
if control["dispatchTokens"] > 0:
raise TuningError("已有未消费的单步 Trial 令牌")
try:
self.storage.grant_dispatch_token(session_id)
except StorageConflict as error:
raise ResourceBusyError(str(error)) from error
if session["state"] == "paused":
has_pending = any(
proposal["state"] == "pending"
for proposal in self.storage.list_proposals(session_id)
)
self.storage.update_session(
session_id,
state="awaiting_approval" if has_pending else "running",
message="单步令牌已就绪;请审批 Proposal"
if has_pending
else "已授权执行一个 Trial",
)
self.storage.audit(session_id, "step_token_granted", {"count": 1})
with self.condition:
self.condition.notify_all()
return self.detail(session_id)
def rollback(self, session_id: str, payload: Any) -> dict:
if not isinstance(payload, dict):
raise TuningError("rollback 请求体必须是对象")
session = self.storage.get_session(session_id)
if session["state"] not in {"paused", "awaiting_approval"}:
raise TuningError("回滚只能在安全暂停或等待审批时执行")
target_id = payload.get("trialId")
if payload.get("target") == "best":
target_id = session["bestTrialId"] or (self._best_highest_rung(session_id) or {}).get(
"id"
)
if not isinstance(target_id, str):
raise TuningError("rollback 必须指定 trialId 或 target=best")
target = self.storage.get_trial(target_id)
if target["sessionId"] != session_id:
raise TuningError("rollback Trial 不属于该 session")
if target["state"] != "completed" or not target["eligible"]:
raise TuningError("只能回滚到已完成且通过安全门槛的 Trial")
checkpoint = payload.get("checkpoint", False)
if not isinstance(checkpoint, bool):
raise TuningError("checkpoint 必须是布尔值")
if checkpoint:
path_value = target.get("checkpointPath")
if not path_value:
raise TuningError("目标 Trial 没有可用 checkpoint")
root = self._session_root(session_id)
path = (root / path_value).resolve()
if not path.is_relative_to(root) or not path.is_file():
raise TuningError("目标 checkpoint 不存在或路径非法")
superseded = []
for proposal in self.storage.list_proposals(session_id):
if proposal["state"] == "pending" and self.storage.decide_proposal(
proposal["id"], "rejected", f"由回滚到 Trial {target_id[:8]} 取代"
):
superseded.append(proposal["id"])
self.storage.set_active_base(session_id, target_id)
self.storage.reset_step_gate(session_id)
self.storage.update_session(
session_id,
state="paused",
message=f"已回滚到 Trial {target['number']} / Rung {target['rung']};等待单步或继续",
)
self.storage.audit(
session_id,
"rollback_selected",
{
"trialId": target_id,
"checkpoint": checkpoint,
"supersededProposalIds": superseded,
},
)
with self.condition:
self.condition.notify_all()
return self.detail(session_id)
def pause(self, session_id: str) -> dict:
session = self.storage.get_session(session_id)
if session["state"] not in RUNNING_STATES | {"awaiting_approval"}:
raise TuningError("当前状态不能暂停")
# 以持久化调度门记录暂停意图,而不是只依赖 session.state。当前 Trial
# 的训练/评估会继续完成,下一次 _claim_dispatch 必须等待显式继续或单步。
self.storage.reset_step_gate(session_id)
self.storage.update_session(
session_id, state="paused", message="已暂停后续调度;在途 Trial(如有)将完整结束"
)
return self.detail(session_id)
def resume(self, session_id: str) -> dict:
session = self.storage.get_session(session_id)
if session["state"] == "paused":
self.storage.set_run_policy(session_id, "continuous")
has_pending = any(
proposal["state"] == "pending"
for proposal in self.storage.list_proposals(session_id)
)
self.storage.update_session(
session_id,
state="awaiting_approval" if has_pending else "running",
message="请审批待处理 Proposal" if has_pending else "连续调参已恢复",
)
with self.condition:
self.condition.notify_all()
elif session["state"] == "interrupted":
self.storage.set_run_policy(session_id, "continuous")
self.storage.update_session(session_id, state="queued", message="从最近完整结果恢复")
self._start_worker(session_id, resume=True)
else:
raise TuningError("当前状态不能恢复")
return self.detail(session_id)
def cancel(self, session_id: str) -> dict:
session = self.storage.get_session(session_id)
if session["state"] not in ACTIVE_SESSION_STATES:
return self.detail(session_id)
self.storage.update_session(session_id, state="cancelled", message="正在取消")
event = self.cancel_events.get(session_id)
if event:
event.set()
process = self.processes.get(session_id)
if process:
terminate_process(process)
with self.condition:
self.condition.notify_all()
return self.detail(session_id)
def metrics(
self,
session_id: str,
trial_id: str,
tags: list[str] | None,
max_points: int,
after_step: int | None = None,
) -> dict:
trial = self.storage.get_trial(trial_id)
if trial["sessionId"] != session_id:
raise TuningError("trial 不属于该 session")
series = self.storage.metrics(trial_id, tags, max_points, after_step)
next_step = max(
(point["step"] for item in series for point in item["points"]),
default=after_step,
)
return {"trialId": trial_id, "series": series, "nextStep": next_step}
def best_artifact(self, session_id: str) -> Path:
session = self.storage.get_session(session_id)
if not session["bestTrialId"]:
raise TuningError("尚无最佳策略")
trial = self.storage.get_trial(session["bestTrialId"])
if not trial["policyPath"]:
raise TuningError("最佳策略文件不存在")
root = self._session_root(session_id)
path = (root / trial["policyPath"]).resolve()
if not path.is_relative_to(root) or not path.is_file():
raise TuningError("最佳策略文件不存在")
return path
def preset_config(self, preset_id: str) -> dict:
return self.storage.get_preset(preset_id)["rewardConfig"]
def test_agent(self) -> dict:
try:
return self.advisor.test_connection()
except Exception as error:
raise TuningError(f"Agent 连接测试失败:{error}") from error
def shutdown(self) -> None:
for session_id in list(self.workers):
with suppress(KeyError, TuningError):
self.cancel(session_id)
for worker in list(self.workers.values()):
worker.join(timeout=7)
-46
View File
@@ -1,46 +0,0 @@
"""Shared GPU lease and process-group lifecycle helpers."""
from __future__ import annotations
import os
import signal
import subprocess
import threading
from contextlib import suppress
class ResourceBusyError(RuntimeError):
pass
class GpuLease:
def __init__(self):
self.lock = threading.RLock()
self.owner: str | None = None
def acquire(self, owner: str) -> None:
with self.lock:
if self.owner is not None and self.owner != owner:
raise ResourceBusyError(f"计算资源正由 {self.owner} 使用")
self.owner = owner
def release(self, owner: str) -> None:
with self.lock:
if self.owner == owner:
self.owner = None
def public(self) -> str | None:
with self.lock:
return self.owner
def terminate_process(process: subprocess.Popen[str], grace_seconds: float = 5.0) -> None:
if process.poll() is not None:
return
with suppress(ProcessLookupError):
os.killpg(process.pid, signal.SIGTERM)
try:
process.wait(timeout=grace_seconds)
except subprocess.TimeoutExpired:
with suppress(ProcessLookupError):
os.killpg(process.pid, signal.SIGKILL)
-281
View File
@@ -1,281 +0,0 @@
"""Pure-Python reward tuning schema shared by the service and trainer."""
from __future__ import annotations
import math
from collections.abc import Mapping
from copy import deepcopy
from dataclasses import dataclass
from typing import Any
MAX_PROPOSAL_CHANGES = 4
MIN_CHANGE_RATIO = 0.5
MAX_CHANGE_RATIO = 2.0
class RewardConfigError(ValueError):
"""A reward configuration or proposal violated the allowlist."""
@dataclass(frozen=True)
class NumericSpec:
minimum: float
maximum: float
default: float
allow_zero: bool = True
WEIGHT_SPECS: dict[str, NumericSpec] = {
"track_linear_velocity": NumericSpec(0.5, 3.0, 1.0, False),
"track_angular_velocity": NumericSpec(0.25, 2.0, 1.0, False),
"body_orientation_l2": NumericSpec(-3.0, -0.1, -1.0, False),
"pose": NumericSpec(0.0, 2.5, 1.0),
"body_ang_vel": NumericSpec(-0.2, 0.0, -0.05),
"angular_momentum": NumericSpec(-0.1, 0.0, -0.025),
"is_terminated": NumericSpec(-400.0, -50.0, -200.0, False),
"joint_acc_l2": NumericSpec(-2.0e-6, 0.0, -2.5e-7),
"joint_pos_limits": NumericSpec(-30.0, -2.0, -10.0, False),
"action_rate_l2": NumericSpec(-0.2, -0.005, -0.05, False),
"foot_gait": NumericSpec(0.0, 1.5, 0.5),
"foot_clearance": NumericSpec(-3.0, 0.0, -1.0),
"foot_slip": NumericSpec(-1.0, 0.0, -0.25),
"soft_landing": NumericSpec(-5.0e-3, 0.0, -1.0e-3),
"stand_still": NumericSpec(-3.0, 0.0, -1.0),
"electrical_power": NumericSpec(-5.0e-3, 0.0, 0.0),
}
PARAMETER_SPECS: dict[str, NumericSpec] = {
"track_linear_velocity.std": NumericSpec(0.25, 1.0, math.sqrt(0.25), False),
"track_angular_velocity.std": NumericSpec(0.35, 1.2, math.sqrt(0.5), False),
"pose.std_standing_scale": NumericSpec(0.5, 2.0, 1.0, False),
"pose.std_walking_scale": NumericSpec(0.5, 2.0, 1.0, False),
"pose.std_running_scale": NumericSpec(0.5, 2.0, 1.0, False),
"pose.walking_threshold": NumericSpec(0.05, 0.5, 0.1, False),
"pose.running_threshold": NumericSpec(1.0, 2.5, 1.5, False),
"foot_gait.period": NumericSpec(0.4, 0.8, 0.6, False),
"foot_gait.threshold": NumericSpec(0.45, 0.65, 0.56, False),
"foot_gait.command_threshold": NumericSpec(0.02, 0.3, 0.1, False),
"foot_clearance.target_height": NumericSpec(0.06, 0.16, 0.1, False),
"foot_clearance.command_threshold": NumericSpec(0.02, 0.3, 0.1, False),
"foot_slip.command_threshold": NumericSpec(0.02, 0.3, 0.1, False),
"soft_landing.command_threshold": NumericSpec(0.02, 0.3, 0.1, False),
"stand_still.command_threshold": NumericSpec(0.02, 0.3, 0.1, False),
}
BASE_REWARD_CONFIGURATION: dict[str, dict[str, float]] = {
"weights": {name: spec.default for name, spec in WEIGHT_SPECS.items()},
"params": {name: spec.default for name, spec in PARAMETER_SPECS.items()},
}
def _number(name: str, value: Any, spec: NumericSpec) -> float:
if isinstance(value, bool) or not isinstance(value, (int, float)):
raise RewardConfigError(f"{name} 必须是数值")
result = float(value)
if not math.isfinite(result):
raise RewardConfigError(f"{name} 必须是有限数值")
if result == 0.0 and not spec.allow_zero:
raise RewardConfigError(f"{name} 不允许关闭")
if result < spec.minimum or result > spec.maximum:
raise RewardConfigError(f"{name} 必须在 {spec.minimum}–{spec.maximum} 之间")
return result
def _mapping(value: Any, name: str) -> Mapping[str, Any]:
if not isinstance(value, Mapping):
raise RewardConfigError(f"{name} 必须是对象")
return value
def _cross_validate(config: Mapping[str, Mapping[str, float]]) -> None:
params = config["params"]
if params["pose.walking_threshold"] >= params["pose.running_threshold"]:
raise RewardConfigError("pose.walking_threshold 必须小于 pose.running_threshold")
def _path_spec(path: str) -> tuple[str, str, NumericSpec]:
if path.startswith("weights."):
section, name = "weights", path.removeprefix("weights.")
spec = WEIGHT_SPECS.get(name)
elif path.startswith("params."):
section, name = "params", path.removeprefix("params.")
spec = PARAMETER_SPECS.get(name)
else:
section, name, spec = "", "", None
if spec is None:
raise RewardConfigError(f"未知参数约束:{path}")
return section, name, spec
def validate_constraints(value: Any) -> dict[str, dict[str, float | str]]:
"""Validate sparse per-session range/fixed safety constraints."""
root = _mapping(value, "constraints")
if len(root) > len(WEIGHT_SPECS) + len(PARAMETER_SPECS):
raise RewardConfigError("constraints 数量超过白名单参数总数")
result: dict[str, dict[str, float | str]] = {}
for raw_path, raw_constraint in root.items():
if not isinstance(raw_path, str):
raise RewardConfigError("constraint path 必须是字符串")
_, _, spec = _path_spec(raw_path)
constraint = _mapping(raw_constraint, raw_path)
kind = constraint.get("kind")
if kind == "fixed":
if set(constraint) != {"kind", "value"}:
raise RewardConfigError(f"{raw_path} fixed 约束只能包含 kind/value")
fixed = _number(f"{raw_path}.value", constraint["value"], spec)
result[raw_path] = {"kind": "fixed", "value": fixed}
elif kind == "range":
if set(constraint) != {"kind", "min", "max"}:
raise RewardConfigError(f"{raw_path} range 约束只能包含 kind/min/max")
minimum = _number(f"{raw_path}.min", constraint["min"], spec)
maximum = _number(f"{raw_path}.max", constraint["max"], spec)
if minimum > maximum:
raise RewardConfigError(f"{raw_path} 下限不能大于上限")
result[raw_path] = {"kind": "range", "min": minimum, "max": maximum}
else:
raise RewardConfigError(f"{raw_path}.kind 必须是 range 或 fixed")
return result
def validate_configuration_constraints(value: Any, constraints: Any) -> None:
"""Ensure a complete reward configuration satisfies every session constraint."""
config = validate_configuration(value)
checked = validate_constraints(constraints)
for path, constraint in checked.items():
section, name, _ = _path_spec(path)
current = config[section][name]
if constraint["kind"] == "fixed":
if current != constraint["value"]:
raise RewardConfigError(
f"{path} 已固定为 {constraint['value']},不能设为 {current}"
)
elif current < constraint["min"] or current > constraint["max"]:
raise RewardConfigError(
f"{path}={current} 超出工程锁定范围 {constraint['min']}–{constraint['max']}"
)
def validate_configuration(value: Any) -> dict[str, dict[str, float]]:
"""Validate a complete configuration and reject missing/unknown fields."""
root = _mapping(value, "rewardConfig")
if set(root) != {"weights", "params"}:
raise RewardConfigError("rewardConfig 只能包含 weights 和 params")
raw_weights = _mapping(root["weights"], "weights")
raw_params = _mapping(root["params"], "params")
if set(raw_weights) != set(WEIGHT_SPECS):
raise RewardConfigError("weights 必须完整且不能包含未知奖励项")
if set(raw_params) != set(PARAMETER_SPECS):
raise RewardConfigError("params 必须完整且不能包含未知参数")
config = {
"weights": {
name: _number(f"weights.{name}", raw_weights[name], spec)
for name, spec in WEIGHT_SPECS.items()
},
"params": {
name: _number(f"params.{name}", raw_params[name], spec)
for name, spec in PARAMETER_SPECS.items()
},
}
_cross_validate(config)
return config
def validate_proposal(
value: Any, previous: Any, constraints: Any | None = None
) -> dict[str, dict[str, float]]:
"""Validate a sparse Agent patch relative to a complete previous config and guardrails."""
current = validate_configuration(previous)
root = _mapping(value, "proposal")
if not set(root).issubset({"weights", "params"}):
raise RewardConfigError("proposal 只能包含 weights 和 params")
raw_weights = _mapping(root.get("weights", {}), "weights")
raw_params = _mapping(root.get("params", {}), "params")
if len(raw_weights) + len(raw_params) == 0:
raise RewardConfigError("proposal 至少需要一项修改")
if len(raw_weights) + len(raw_params) > MAX_PROPOSAL_CHANGES:
raise RewardConfigError(f"proposal 每轮最多修改 {MAX_PROPOSAL_CHANGES} 项")
unknown_weights = set(raw_weights) - set(WEIGHT_SPECS)
unknown_params = set(raw_params) - set(PARAMETER_SPECS)
if unknown_weights:
raise RewardConfigError(f"未知奖励项:{', '.join(sorted(unknown_weights))}")
if unknown_params:
raise RewardConfigError(f"未知奖励参数:{', '.join(sorted(unknown_params))}")
patch: dict[str, dict[str, float]] = {"weights": {}, "params": {}}
for name, raw in raw_weights.items():
value_number = _number(f"weights.{name}", raw, WEIGHT_SPECS[name])
old = current["weights"][name]
if old != 0.0 and value_number != 0.0:
ratio = abs(value_number / old)
if ratio < MIN_CHANGE_RATIO or ratio > MAX_CHANGE_RATIO:
raise RewardConfigError(
f"weights.{name} 单轮变化必须在旧值幅度的 "
f"{MIN_CHANGE_RATIO}×–{MAX_CHANGE_RATIO}×"
)
if value_number == old:
raise RewardConfigError(f"weights.{name} 没有发生变化")
patch["weights"][name] = value_number
for name, raw in raw_params.items():
value_number = _number(f"params.{name}", raw, PARAMETER_SPECS[name])
old = current["params"][name]
ratio = abs(value_number / old)
if ratio < MIN_CHANGE_RATIO or ratio > MAX_CHANGE_RATIO:
raise RewardConfigError(
f"params.{name} 单轮变化必须在旧值的 {MIN_CHANGE_RATIO}×–{MAX_CHANGE_RATIO}×"
)
if value_number == old:
raise RewardConfigError(f"params.{name} 没有发生变化")
patch["params"][name] = value_number
candidate = deepcopy(current)
candidate["weights"].update(patch["weights"])
candidate["params"].update(patch["params"])
_cross_validate(candidate)
if constraints is not None:
validate_configuration_constraints(candidate, constraints)
return patch
def merge_proposal(
previous: Any, proposal: Any, constraints: Any | None = None
) -> dict[str, dict[str, float]]:
current = validate_configuration(previous)
patch = validate_proposal(proposal, current, constraints)
merged = deepcopy(current)
merged["weights"].update(patch["weights"])
merged["params"].update(patch["params"])
return validate_configuration(merged)
def apply_reward_configuration(env_cfg: Any, value: Any) -> None:
"""Apply a validated full config to a fresh mjlab environment config."""
config = validate_configuration(value)
for name, weight in config["weights"].items():
if name not in env_cfg.rewards:
raise RewardConfigError(f"环境缺少奖励项:{name}")
env_cfg.rewards[name].weight = weight
params = config["params"]
direct = {
"track_linear_velocity.std": ("track_linear_velocity", "std"),
"track_angular_velocity.std": ("track_angular_velocity", "std"),
"pose.walking_threshold": ("pose", "walking_threshold"),
"pose.running_threshold": ("pose", "running_threshold"),
"foot_gait.period": ("foot_gait", "period"),
"foot_gait.threshold": ("foot_gait", "threshold"),
"foot_gait.command_threshold": ("foot_gait", "command_threshold"),
"foot_clearance.target_height": ("foot_clearance", "target_height"),
"foot_clearance.command_threshold": ("foot_clearance", "command_threshold"),
"foot_slip.command_threshold": ("foot_slip", "command_threshold"),
"soft_landing.command_threshold": ("soft_landing", "command_threshold"),
"stand_still.command_threshold": ("stand_still", "command_threshold"),
}
for path, (term, parameter) in direct.items():
env_cfg.rewards[term].params[parameter] = params[path]
for regime in ("standing", "walking", "running"):
key = f"std_{regime}"
scale = params[f"pose.{key}_scale"]
baseline = env_cfg.rewards["pose"].params[key]
env_cfg.rewards["pose"].params[key] = {
pattern: float(std) * scale for pattern, std in baseline.items()
}
-118
View File
@@ -1,118 +0,0 @@
"""Stable, reward-weight-independent evaluation scoring."""
from __future__ import annotations
import math
from collections.abc import Mapping
from typing import Any
DEFAULT_OBJECTIVE_WEIGHTS = {
"velocity_tracking": 0.35,
"action_smoothness": 0.20,
"posture_stability": 0.15,
"fall_avoidance": 0.15,
"foot_slip": 0.10,
"energy": 0.05,
}
REQUIRED_METRICS = {
"linear_velocity_rmse",
"angular_velocity_rmse",
"mean_action_acc",
"orientation_error",
"fall_rate",
"slip_velocity",
"mechanical_power",
}
PHYSICAL_FLOORS = {
"linear_velocity_rmse": 0.10,
"angular_velocity_rmse": 0.10,
"mean_action_acc": 0.01,
"orientation_error": 0.05,
"fall_rate": 0.02,
"slip_velocity": 0.05,
"mechanical_power": 10.0,
}
class EvaluationError(ValueError):
pass
def validate_objective_weights(value: Any) -> dict[str, float]:
if not isinstance(value, Mapping) or set(value) != set(DEFAULT_OBJECTIVE_WEIGHTS):
raise EvaluationError("objectiveWeights 必须完整包含六个目标")
result: dict[str, float] = {}
for key in DEFAULT_OBJECTIVE_WEIGHTS:
raw = value[key]
if isinstance(raw, bool) or not isinstance(raw, (int, float)):
raise EvaluationError(f"objectiveWeights.{key} 必须是数值")
number = float(raw)
if not math.isfinite(number) or number < 0.0 or number > 1.0:
raise EvaluationError(f"objectiveWeights.{key} 必须在 0–1 之间")
result[key] = number
if not math.isclose(sum(result.values()), 1.0, abs_tol=1.0e-6):
raise EvaluationError("objectiveWeights 总和必须为 1")
return result
def validate_metrics(value: Any) -> dict[str, float]:
if not isinstance(value, Mapping):
raise EvaluationError("metrics 必须是对象")
missing = REQUIRED_METRICS - set(value)
if missing:
raise EvaluationError(f"metrics 缺少:{', '.join(sorted(missing))}")
result: dict[str, float] = {}
for key in REQUIRED_METRICS:
raw = value[key]
if isinstance(raw, bool) or not isinstance(raw, (int, float)):
raise EvaluationError(f"metrics.{key} 必须是数值")
number = float(raw)
if not math.isfinite(number) or number < 0.0:
raise EvaluationError(f"metrics.{key} 必须是非负有限数值")
result[key] = number
if result["fall_rate"] > 1.0:
raise EvaluationError("metrics.fall_rate 必须在 0–1 之间")
return result
def _improvement(baseline: Mapping[str, float], current: Mapping[str, float], key: str) -> float:
scale = max(abs(baseline[key]), PHYSICAL_FLOORS[key])
return max(-1.0, min(1.0, (baseline[key] - current[key]) / scale))
def score_evaluation(
baseline_value: Any,
current_value: Any,
objective_weights: Any = DEFAULT_OBJECTIVE_WEIGHTS,
) -> dict[str, Any]:
baseline = validate_metrics(baseline_value)
current = validate_metrics(current_value)
weights = validate_objective_weights(objective_weights)
components = {
"velocity_tracking": 0.8 * _improvement(baseline, current, "linear_velocity_rmse")
+ 0.2 * _improvement(baseline, current, "angular_velocity_rmse"),
"action_smoothness": _improvement(baseline, current, "mean_action_acc"),
"posture_stability": _improvement(baseline, current, "orientation_error"),
"fall_avoidance": _improvement(baseline, current, "fall_rate"),
"foot_slip": _improvement(baseline, current, "slip_velocity"),
"energy": _improvement(baseline, current, "mechanical_power"),
}
tracking_limit = max(
baseline["linear_velocity_rmse"] * 1.05, baseline["linear_velocity_rmse"] + 1.0e-6
)
eligible = (
current["fall_rate"] <= baseline["fall_rate"] + 0.02
and current["linear_velocity_rmse"] <= tracking_limit
)
total = sum(weights[key] * components[key] for key in weights)
if not eligible:
total = min(total, -1.0)
return {
"eligible": eligible,
"score": total,
"components": components,
"metrics": current,
"baselineMetrics": baseline,
}
-689
View File
@@ -1,689 +0,0 @@
"""SQLite persistence for tuning sessions, trials, proposals and scalar data."""
from __future__ import annotations
import json
import sqlite3
import threading
import uuid
from collections.abc import Iterator
from contextlib import contextmanager
from datetime import UTC, datetime
from pathlib import Path
from typing import Any
SCHEMA_VERSION = 2
class StorageConflict(RuntimeError):
"""Optimistic-concurrency revision mismatch."""
def now_iso() -> str:
return datetime.now(UTC).isoformat()
def _json(value: Any) -> str:
return json.dumps(value, ensure_ascii=False, separators=(",", ":"), allow_nan=False)
def _decode(value: str | None) -> Any:
return json.loads(value) if value else None
def _lttb(points: list[dict], threshold: int) -> list[dict]:
"""Largest-Triangle-Three-Buckets downsampling preserving peaks and endpoints."""
if threshold >= len(points) or threshold < 3:
return points[:threshold]
sampled = [points[0]]
bucket_width = (len(points) - 2) / (threshold - 2)
anchor_index = 0
for bucket in range(threshold - 2):
average_start = int((bucket + 1) * bucket_width) + 1
average_end = min(int((bucket + 2) * bucket_width) + 1, len(points))
average_bucket = points[average_start:average_end] or [points[-1]]
average_x = sum(point["step"] for point in average_bucket) / len(average_bucket)
average_y = sum(point["value"] for point in average_bucket) / len(average_bucket)
range_start = int(bucket * bucket_width) + 1
range_end = min(int((bucket + 1) * bucket_width) + 1, len(points) - 1)
anchor = points[anchor_index]
selected_index = range_start
maximum_area = -1.0
for index in range(range_start, max(range_start + 1, range_end)):
point = points[index]
area = abs(
(anchor["step"] - average_x) * (point["value"] - anchor["value"])
- (anchor["step"] - point["step"]) * (average_y - anchor["value"])
)
if area > maximum_area:
maximum_area = area
selected_index = index
sampled.append(points[selected_index])
anchor_index = selected_index
sampled.append(points[-1])
return sampled
class TuningStorage:
def __init__(self, path: Path):
self.path = path.expanduser().resolve()
self.path.parent.mkdir(parents=True, exist_ok=True)
self.local = threading.local()
self._migrate()
def connection(self) -> sqlite3.Connection:
connection = getattr(self.local, "connection", None)
if connection is None:
connection = sqlite3.connect(self.path, timeout=10, isolation_level=None)
connection.row_factory = sqlite3.Row
connection.execute("PRAGMA foreign_keys=ON")
connection.execute("PRAGMA journal_mode=WAL")
connection.execute("PRAGMA busy_timeout=10000")
self.local.connection = connection
return connection
@contextmanager
def transaction(self) -> Iterator[sqlite3.Connection]:
connection = self.connection()
connection.execute("BEGIN IMMEDIATE")
try:
yield connection
connection.execute("COMMIT")
except Exception:
connection.execute("ROLLBACK")
raise
def _migrate(self) -> None:
connection = self.connection()
connection.executescript(
"""
CREATE TABLE IF NOT EXISTS schema_migrations(version INTEGER PRIMARY KEY);
CREATE TABLE IF NOT EXISTS sessions(
id TEXT PRIMARY KEY, state TEXT NOT NULL, mode TEXT NOT NULL,
created_at TEXT NOT NULL, updated_at TEXT NOT NULL,
config_json TEXT NOT NULL, objective_json TEXT NOT NULL,
message TEXT NOT NULL, current_trial_id TEXT, best_trial_id TEXT,
consecutive_no_improve INTEGER NOT NULL DEFAULT 0,
fallback_enabled INTEGER NOT NULL DEFAULT 0
);
CREATE TABLE IF NOT EXISTS trials(
id TEXT PRIMARY KEY,
session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE,
number INTEGER NOT NULL, state TEXT NOT NULL, rung INTEGER NOT NULL,
target_iterations INTEGER NOT NULL, reward_config_json TEXT NOT NULL,
proposal_id TEXT, run_dir TEXT NOT NULL, checkpoint_path TEXT,
policy_path TEXT, evaluation_json TEXT, score REAL, eligible INTEGER,
created_at TEXT NOT NULL, started_at TEXT, ended_at TEXT, message TEXT NOT NULL,
UNIQUE(session_id, number, rung)
);
CREATE TABLE IF NOT EXISTS proposals(
id TEXT PRIMARY KEY,
session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE,
base_trial_id TEXT, state TEXT NOT NULL, source TEXT NOT NULL,
patch_json TEXT NOT NULL, rationale TEXT NOT NULL,
expected_json TEXT, confidence REAL NOT NULL,
created_at TEXT NOT NULL, decided_at TEXT, feedback TEXT
);
CREATE TABLE IF NOT EXISTS metric_points(
trial_id TEXT NOT NULL REFERENCES trials(id) ON DELETE CASCADE,
tag TEXT NOT NULL, step INTEGER NOT NULL,
wall_time REAL NOT NULL, value REAL NOT NULL,
PRIMARY KEY(trial_id, tag, step)
);
CREATE TABLE IF NOT EXISTS audit_events(
id INTEGER PRIMARY KEY AUTOINCREMENT,
session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE,
event_type TEXT NOT NULL, payload_json TEXT NOT NULL, created_at TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS presets(
id TEXT PRIMARY KEY, name TEXT NOT NULL UNIQUE, session_id TEXT NOT NULL,
trial_id TEXT NOT NULL, reward_config_json TEXT NOT NULL, created_at TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS session_controls(
session_id TEXT PRIMARY KEY REFERENCES sessions(id) ON DELETE CASCADE,
run_policy TEXT NOT NULL DEFAULT 'continuous'
CHECK(run_policy IN ('continuous','step')),
dispatch_tokens INTEGER NOT NULL DEFAULT 0 CHECK(dispatch_tokens >= 0),
active_base_trial_id TEXT,
revision INTEGER NOT NULL DEFAULT 0
);
CREATE TABLE IF NOT EXISTS session_constraints(
session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE,
path TEXT NOT NULL, kind TEXT NOT NULL CHECK(kind IN ('range','fixed')),
min_value REAL, max_value REAL, fixed_value REAL,
updated_at TEXT NOT NULL,
PRIMARY KEY(session_id, path)
);
CREATE INDEX IF NOT EXISTS idx_trials_session ON trials(session_id, number, rung);
CREATE INDEX IF NOT EXISTS idx_proposals_session ON proposals(session_id, created_at);
CREATE INDEX IF NOT EXISTS idx_metrics_trial_tag ON metric_points(trial_id, tag, step);
"""
)
connection.execute(
"INSERT OR IGNORE INTO session_controls(session_id) SELECT id FROM sessions"
)
connection.execute(
"INSERT OR IGNORE INTO schema_migrations(version) VALUES (?)", (SCHEMA_VERSION,)
)
def recover_interrupted(self) -> None:
at = now_iso()
with self.transaction() as connection:
connection.execute(
"UPDATE trials SET state='interrupted', ended_at=?, "
"message='服务重启中断,等待显式恢复' "
"WHERE state IN ('training','evaluating')",
(at,),
)
connection.execute(
"UPDATE sessions SET state='interrupted', updated_at=?, "
"message='服务重启中断,可从完整 checkpoint 恢复' "
"WHERE state IN ('running','evaluating')",
(at,),
)
def create_session(self, mode: str, config: dict, objective: dict, fallback: bool) -> dict:
session_id, at = uuid.uuid4().hex, now_iso()
with self.transaction() as connection:
connection.execute(
"INSERT INTO sessions("
"id,state,mode,created_at,updated_at,config_json,objective_json,"
"message,fallback_enabled) "
"VALUES (?, 'queued', ?, ?, ?, ?, ?, '等待基线训练', ?)",
(session_id, mode, at, at, _json(config), _json(objective), int(fallback)),
)
connection.execute("INSERT INTO session_controls(session_id) VALUES (?)", (session_id,))
connection.execute(
"INSERT INTO audit_events(session_id,event_type,payload_json,created_at) "
"VALUES (?,?,?,?)",
(session_id, "session_created", _json({"mode": mode}), at),
)
return self.get_session(session_id)
def _session(self, row: sqlite3.Row) -> dict:
return {
"id": row["id"],
"state": row["state"],
"mode": row["mode"],
"createdAt": row["created_at"],
"updatedAt": row["updated_at"],
"config": _decode(row["config_json"]),
"objectiveWeights": _decode(row["objective_json"]),
"message": row["message"],
"currentTrialId": row["current_trial_id"],
"bestTrialId": row["best_trial_id"],
"consecutiveNoImprove": row["consecutive_no_improve"],
"fallbackEnabled": bool(row["fallback_enabled"]),
}
def get_session(self, session_id: str) -> dict:
row = (
self.connection().execute("SELECT * FROM sessions WHERE id=?", (session_id,)).fetchone()
)
if row is None:
raise KeyError(session_id)
return self._session(row)
def list_sessions(self, limit: int = 50) -> list[dict]:
rows = (
self.connection()
.execute("SELECT * FROM sessions ORDER BY created_at DESC LIMIT ?", (limit,))
.fetchall()
)
return [self._session(row) for row in rows]
def update_session(self, session_id: str, **changes: Any) -> bool:
columns = {
"state": "state",
"mode": "mode",
"message": "message",
"current_trial_id": "current_trial_id",
"best_trial_id": "best_trial_id",
"consecutive_no_improve": "consecutive_no_improve",
}
values, assignments = [], []
for key, value in changes.items():
if key not in columns:
raise ValueError(key)
assignments.append(f"{columns[key]}=?")
values.append(value)
assignments.append("updated_at=?")
values.extend((now_iso(), session_id))
cursor = self.connection().execute(
f"UPDATE sessions SET {', '.join(assignments)} WHERE id=?", values
)
return cursor.rowcount == 1
def get_control(self, session_id: str) -> dict:
row = (
self.connection()
.execute("SELECT * FROM session_controls WHERE session_id=?", (session_id,))
.fetchone()
)
if row is None:
raise KeyError(session_id)
constraint_rows = (
self.connection()
.execute(
"SELECT * FROM session_constraints WHERE session_id=? ORDER BY path", (session_id,)
)
.fetchall()
)
constraints = {}
for constraint in constraint_rows:
if constraint["kind"] == "fixed":
value = {"kind": "fixed", "value": constraint["fixed_value"]}
else:
value = {
"kind": "range",
"min": constraint["min_value"],
"max": constraint["max_value"],
}
constraints[constraint["path"]] = value
return {
"runPolicy": row["run_policy"],
"dispatchTokens": row["dispatch_tokens"],
"constraintsRevision": row["revision"],
"constraints": constraints,
"activeBaseTrialId": row["active_base_trial_id"],
}
def replace_constraints(
self, session_id: str, expected_revision: int, constraints: dict
) -> dict:
at = now_iso()
with self.transaction() as connection:
row = connection.execute(
"SELECT revision FROM session_controls WHERE session_id=?", (session_id,)
).fetchone()
if row is None:
raise KeyError(session_id)
if row["revision"] != expected_revision:
raise StorageConflict(
f"参数护栏 revision 已变化(当前 {row['revision']},请求 {expected_revision})"
)
connection.execute("DELETE FROM session_constraints WHERE session_id=?", (session_id,))
for path, constraint in constraints.items():
connection.execute(
"INSERT INTO session_constraints("
"session_id,path,kind,min_value,max_value,fixed_value,updated_at) "
"VALUES (?,?,?,?,?,?,?)",
(
session_id,
path,
constraint["kind"],
constraint.get("min"),
constraint.get("max"),
constraint.get("value"),
at,
),
)
connection.execute(
"UPDATE session_controls SET revision=revision+1 WHERE session_id=?",
(session_id,),
)
return self.get_control(session_id)
def grant_dispatch_token(self, session_id: str) -> dict:
"""Atomically grant the sole outstanding one-Trial token."""
with self.transaction() as connection:
cursor = connection.execute(
"UPDATE session_controls SET run_policy='step',dispatch_tokens=1 "
"WHERE session_id=? AND dispatch_tokens=0",
(session_id,),
)
if cursor.rowcount != 1:
exists = connection.execute(
"SELECT 1 FROM session_controls WHERE session_id=?", (session_id,)
).fetchone()
if exists is None:
raise KeyError(session_id)
raise StorageConflict("已有未消费的单步 Trial 令牌")
return self.get_control(session_id)
def use_dispatch_token(self, session_id: str) -> bool:
"""Atomically consume one step token; continuous mode never needs a token."""
with self.transaction() as connection:
row = connection.execute(
"SELECT run_policy,dispatch_tokens FROM session_controls WHERE session_id=?",
(session_id,),
).fetchone()
if row is None:
raise KeyError(session_id)
if row["run_policy"] == "continuous":
return True
if row["dispatch_tokens"] <= 0:
return False
connection.execute(
"UPDATE session_controls SET dispatch_tokens=dispatch_tokens-1 WHERE session_id=?",
(session_id,),
)
return True
def set_run_policy(self, session_id: str, policy: str) -> dict:
if policy not in {"continuous", "step"}:
raise ValueError(policy)
cursor = self.connection().execute(
"UPDATE session_controls SET run_policy=?,"
"dispatch_tokens=CASE WHEN ?='continuous' THEN 0 ELSE dispatch_tokens END "
"WHERE session_id=?",
(policy, policy, session_id),
)
if cursor.rowcount != 1:
raise KeyError(session_id)
return self.get_control(session_id)
def reset_step_gate(self, session_id: str) -> dict:
cursor = self.connection().execute(
"UPDATE session_controls SET run_policy='step',dispatch_tokens=0 WHERE session_id=?",
(session_id,),
)
if cursor.rowcount != 1:
raise KeyError(session_id)
return self.get_control(session_id)
def set_active_base(self, session_id: str, trial_id: str | None) -> dict:
cursor = self.connection().execute(
"UPDATE session_controls SET active_base_trial_id=? WHERE session_id=?",
(trial_id, session_id),
)
if cursor.rowcount != 1:
raise KeyError(session_id)
return self.get_control(session_id)
def create_trial(
self,
session_id: str,
number: int,
rung: int,
target: int,
reward_config: dict,
proposal_id: str | None,
run_dir: str,
) -> dict:
trial_id, at = uuid.uuid4().hex, now_iso()
with self.transaction() as connection:
connection.execute(
"INSERT INTO trials("
"id,session_id,number,state,rung,target_iterations,reward_config_json,"
"proposal_id,run_dir,created_at,message) "
"VALUES (?,?,?,'queued',?,?,?,?,?,?,'等待训练')",
(
trial_id,
session_id,
number,
rung,
target,
_json(reward_config),
proposal_id,
run_dir,
at,
),
)
connection.execute(
"UPDATE sessions SET current_trial_id=?,updated_at=? WHERE id=?",
(trial_id, at, session_id),
)
return self.get_trial(trial_id)
def _trial(self, row: sqlite3.Row) -> dict:
return {
"id": row["id"],
"sessionId": row["session_id"],
"number": row["number"],
"state": row["state"],
"rung": row["rung"],
"targetIterations": row["target_iterations"],
"rewardConfig": _decode(row["reward_config_json"]),
"proposalId": row["proposal_id"],
"runDir": row["run_dir"],
"checkpointPath": row["checkpoint_path"],
"policyPath": row["policy_path"],
"evaluation": _decode(row["evaluation_json"]),
"score": row["score"],
"eligible": None if row["eligible"] is None else bool(row["eligible"]),
"createdAt": row["created_at"],
"startedAt": row["started_at"],
"endedAt": row["ended_at"],
"message": row["message"],
}
def get_trial(self, trial_id: str) -> dict:
row = self.connection().execute("SELECT * FROM trials WHERE id=?", (trial_id,)).fetchone()
if row is None:
raise KeyError(trial_id)
return self._trial(row)
def list_trials(self, session_id: str) -> list[dict]:
rows = (
self.connection()
.execute("SELECT * FROM trials WHERE session_id=? ORDER BY number,rung", (session_id,))
.fetchall()
)
return [self._trial(row) for row in rows]
def delete_trial(self, trial_id: str) -> bool:
cursor = self.connection().execute(
"DELETE FROM trials WHERE id=? AND state='interrupted'", (trial_id,)
)
return cursor.rowcount == 1
def update_trial(self, trial_id: str, **changes: Any) -> bool:
columns = {
"state": "state",
"message": "message",
"checkpoint_path": "checkpoint_path",
"policy_path": "policy_path",
"score": "score",
"eligible": "eligible",
"started_at": "started_at",
"ended_at": "ended_at",
"evaluation": "evaluation_json",
}
values, assignments = [], []
for key, value in changes.items():
if key not in columns:
raise ValueError(key)
if key == "evaluation":
value = _json(value)
if key == "eligible":
value = int(value)
assignments.append(f"{columns[key]}=?")
values.append(value)
values.append(trial_id)
cursor = self.connection().execute(
f"UPDATE trials SET {', '.join(assignments)} WHERE id=?", values
)
return cursor.rowcount == 1
def create_proposal(
self,
session_id: str,
base_trial_id: str | None,
patch: dict,
rationale: str,
expected: Any,
confidence: float,
source: str = "agent",
) -> dict:
proposal_id, at = uuid.uuid4().hex, now_iso()
self.connection().execute(
"INSERT INTO proposals("
"id,session_id,base_trial_id,state,source,patch_json,rationale,"
"expected_json,confidence,created_at) "
"VALUES (?,?,?,'pending',?,?,?,?,?,?)",
(
proposal_id,
session_id,
base_trial_id,
source,
_json(patch),
rationale,
_json(expected),
confidence,
at,
),
)
return self.get_proposal(proposal_id)
def _proposal(self, row: sqlite3.Row) -> dict:
return {
"id": row["id"],
"sessionId": row["session_id"],
"baseTrialId": row["base_trial_id"],
"state": row["state"],
"source": row["source"],
"patch": _decode(row["patch_json"]),
"rationale": row["rationale"],
"expectedImpact": _decode(row["expected_json"]),
"confidence": row["confidence"],
"createdAt": row["created_at"],
"decidedAt": row["decided_at"],
"feedback": row["feedback"],
}
def get_proposal(self, proposal_id: str) -> dict:
row = (
self.connection()
.execute("SELECT * FROM proposals WHERE id=?", (proposal_id,))
.fetchone()
)
if row is None:
raise KeyError(proposal_id)
return self._proposal(row)
def list_proposals(self, session_id: str) -> list[dict]:
rows = (
self.connection()
.execute(
"SELECT * FROM proposals WHERE session_id=? ORDER BY created_at", (session_id,)
)
.fetchall()
)
return [self._proposal(row) for row in rows]
def decide_proposal(
self, proposal_id: str, state: str, feedback: str | None, patch: dict | None = None
) -> bool:
at = now_iso()
assignments, values = ["state=?", "feedback=?", "decided_at=?"], [state, feedback, at]
if patch is not None:
assignments.append("patch_json=?")
values.append(_json(patch))
values.extend((proposal_id,))
cursor = self.connection().execute(
f"UPDATE proposals SET {', '.join(assignments)} WHERE id=? AND state='pending'", values
)
return cursor.rowcount == 1
def insert_metrics(self, trial_id: str, points: list[tuple[str, int, float, float]]) -> None:
self.connection().executemany(
"INSERT INTO metric_points(trial_id,tag,step,wall_time,value) "
"VALUES (?,?,?,?,?) ON CONFLICT(trial_id,tag,step) DO UPDATE SET "
"wall_time=excluded.wall_time,value=excluded.value",
[(trial_id, *point) for point in points],
)
def metrics(
self,
trial_id: str,
tags: list[str] | None = None,
max_points: int = 1000,
after_step: int | None = None,
) -> list[dict]:
parameters: list[Any] = [trial_id]
clause = "trial_id=?"
if tags:
clause += f" AND tag IN ({','.join('?' for _ in tags)})"
parameters.extend(tags)
if after_step is not None:
clause += " AND step>?"
parameters.append(after_step)
rows = (
self.connection()
.execute(
f"SELECT tag,step,wall_time,value FROM metric_points "
f"WHERE {clause} ORDER BY tag,step",
parameters,
)
.fetchall()
)
grouped: dict[str, list[dict]] = {}
for row in rows:
grouped.setdefault(row["tag"], []).append(
{"step": row["step"], "wallTime": row["wall_time"], "value": row["value"]}
)
series = []
for tag, values in grouped.items():
if len(values) > max_points:
values = _lttb(values, max_points)
series.append({"tag": tag, "points": values})
return series
def audit(self, session_id: str, event_type: str, payload: Any) -> None:
self.connection().execute(
"INSERT INTO audit_events(session_id,event_type,payload_json,created_at) "
"VALUES (?,?,?,?)",
(session_id, event_type, _json(payload), now_iso()),
)
def audit_events(self, session_id: str) -> list[dict]:
rows = (
self.connection()
.execute("SELECT * FROM audit_events WHERE session_id=? ORDER BY id", (session_id,))
.fetchall()
)
return [
{
"id": row["id"],
"type": row["event_type"],
"payload": _decode(row["payload_json"]),
"createdAt": row["created_at"],
}
for row in rows
]
def save_preset(self, name: str, session_id: str, trial_id: str, reward_config: dict) -> dict:
preset_id, at = uuid.uuid4().hex, now_iso()
self.connection().execute(
"INSERT INTO presets(id,name,session_id,trial_id,reward_config_json,created_at) "
"VALUES (?,?,?,?,?,?)",
(preset_id, name, session_id, trial_id, _json(reward_config), at),
)
return {
"id": preset_id,
"name": name,
"sessionId": session_id,
"trialId": trial_id,
"rewardConfig": reward_config,
"createdAt": at,
}
def get_preset(self, preset_id: str) -> dict:
row = self.connection().execute("SELECT * FROM presets WHERE id=?", (preset_id,)).fetchone()
if row is None:
raise KeyError(preset_id)
return {
"id": row["id"],
"name": row["name"],
"sessionId": row["session_id"],
"trialId": row["trial_id"],
"rewardConfig": _decode(row["reward_config_json"]),
"createdAt": row["created_at"],
}
def list_presets(self) -> list[dict]:
rows = (
self.connection().execute("SELECT * FROM presets ORDER BY created_at DESC").fetchall()
)
return [
{
"id": row["id"],
"name": row["name"],
"sessionId": row["session_id"],
"trialId": row["trial_id"],
"rewardConfig": _decode(row["reward_config_json"]),
"createdAt": row["created_at"],
}
for row in rows
]
-55
View File
@@ -1,55 +0,0 @@
"""Lazy Optuna study integration used for durable trial history and pruning metadata."""
from __future__ import annotations
from pathlib import Path
from typing import Any
class OptunaStudies:
def __init__(self, root: Path):
self.path = (root / "optuna.sqlite3").resolve()
self.url = f"sqlite:///{self.path}"
def _study(self, session_id: str):
import optuna
optuna.logging.set_verbosity(optuna.logging.WARNING)
return optuna.create_study(
study_name=f"reward-tuning-{session_id}",
storage=self.url,
direction="maximize",
load_if_exists=True,
pruner=optuna.pruners.SuccessiveHalvingPruner(
min_resource=300, reduction_factor=3, min_early_stopping_rate=0
),
)
def record(
self,
session_id: str,
reward_config: dict[str, Any],
score: float,
eligible: bool,
rung: int,
) -> int:
"""Record an externally proposed Agent config through Optuna ask/tell."""
import optuna
study = self._study(session_id)
trial = study.ask()
trial.set_user_attr("reward_config", reward_config)
trial.set_user_attr("eligible", eligible)
trial.set_user_attr("rung", rung)
state = optuna.trial.TrialState.COMPLETE if eligible else optuna.trial.TrialState.PRUNED
study.tell(trial, score if eligible else None, state=state)
return trial.number
def summary(self, session_id: str) -> dict[str, Any]:
study = self._study(session_id)
completed = [trial for trial in study.trials if trial.value is not None]
return {
"studyName": study.study_name,
"trialCount": len(study.trials),
"bestValue": max((trial.value for trial in completed), default=None),
}
-35
View File
@@ -1,35 +0,0 @@
"""TensorBoard scalar ingestion with optional dependency isolation."""
from __future__ import annotations
from pathlib import Path
from .storage import TuningStorage
class TensorboardUnavailable(RuntimeError):
pass
def ingest_scalars(storage: TuningStorage, trial_id: str, log_dir: Path) -> int:
"""Reload all scalar events and idempotently upsert them into SQLite."""
try:
from tensorboard.backend.event_processing.event_accumulator import EventAccumulator
except ImportError as error:
raise TensorboardUnavailable(
"缺少 tensorboard,请安装 training_server/requirements.txt"
) from error
if not log_dir.is_dir():
return 0
accumulator = EventAccumulator(str(log_dir), size_guidance={"scalars": 0})
try:
accumulator.Reload()
except (OSError, ValueError):
return 0
points: list[tuple[str, int, float, float]] = []
for tag in accumulator.Tags().get("scalars", []):
for event in accumulator.Scalars(tag):
points.append((tag, int(event.step), float(event.wall_time), float(event.value)))
if points:
storage.insert_metrics(trial_id, points)
return len(points)
-44
View File
@@ -1,44 +0,0 @@
# 前端架构
`src/` 采用“应用壳 + 业务域 + 共享基础设施”的组织方式。
## 目录职责
```text
src/
├── app/ 应用入口、跨域编排、工作区壳和全局反馈
├── components/ui/ 无业务依赖的共享 UI 基础组件
├── controller/ Python 控制器运行时、类型和面板
├── map/ 地图加载、生成、编辑和面板
├── project/ 工程导入、转换、文件树和工作区
├── rl/ ONNX 策略运行时、任务绑定、类型和面板
├── simulation/ MuJoCo 会话、物理适配器和仿真控制组件
├── telemetry/ 数据源抽象、记录器、导出和数据面板
├── training/ 本地训练/调参客户端、类型、共享连接和面板
├── tuning/ 独立 tuning.html 的 Agent dashboard 与 scalar 图表
├── viewer/ Three.js 场景、渲染和交互
├── stores/ 跨域应用状态
└── test/ 全局测试初始化
```
## 依赖规则
1. `components/ui` 不依赖任何业务域。
2. 业务组件、类型和测试与所属业务域共置。
3. `app` 负责组合业务域,不承载可复用的业务算法或运行时。
4. `simulation` 不依赖 React 应用状态;遥测通过 `TelemetrySource` 接口读取仿真数据。
5. 业务域之间通过显式类型和门面交互,避免从其他域读取内部状态。
6. 新文件优先放入已有业务域;只有跨域复用且不含业务语义的组件才进入 `components/ui`。
## 主要运行链路
```text
项目文件 → project → PhysicsAdapter → SimulationSession → MuJoCo WASM
├→ controller / rl
└→ telemetry
SimulationSession snapshot → app → viewer / 各业务面板
```
主工作台由 `index.html → src/main.tsx` 启动;自调参工作台由 Vite MPA 入口 `tuning.html → src/tuning/main.tsx` 启动,避免把 MuJoCo/Three.js 主应用依赖打入监控页面。两页仅通过训练 HTTP API和严格同源的短消息交接训练服务凭据/策略导入请求,不在 URL 中传 token。
测试文件使用 `*.test.ts(x)` 与被测模块共置;端到端测试统一保存在 `e2e/`。
+5 -49
View File
@@ -12,11 +12,9 @@
- Three.js primitive、mesh、材质/贴图显示与对象选择
- 播放、暂停、单步、重置、0.25×–4× 速度
- actuator 滑杆、hinge/slide 关节拖动、动态 body 外力拖拽
- 内置平地、坡道、楼梯、可复现随机障碍物及 9 类系统参数化地形,可配置尺寸、摩擦、难度、种子与高度场采样精度
- 导入单文件 `.py` 控制器,通过本地 Pyodide 在 `mj_step` 前按仿真时间同步执行
- 导入 mjlab 导出的 `policy.onnx`,在浏览器本地执行 Go2-W 平衡/速度策略推理
- 从图形界面向本机训练桥接服务发起 mjlab 强化学习训练、查看进度/日志、停止任务并导入训练生成的 ONNX;可在独立 TensorBoard 风格页面运行 DeepSeek 奖励函数自调参
- 可配置仿真遥测记录,实时查看速度、机身姿态、位置、驱动力等指标并导出 CSV/JSON
- 从图形界面向本机训练桥接服务发起 mjlab 强化学习训练、查看进度/日志、停止任务并导入训练生成的 ONNX
- FPS、物理耗时和主线程步进预算提示
## 开发
@@ -63,47 +61,6 @@ python3 -m http.server 8080 --directory web-platform-dist
- 默认限制:2000 个文件、单文件 128 MiB、总解压大小 512 MiB、ZIP 文件 128 MiB。
- 文件夹或 ZIP 中的 `.py` 会显示在“控制 → Python 控制器”;也可以在加载模型后单独导入不超过 1 MiB 的 `.py`。
## 地图模块
模型加载后打开右侧“地图”标签,即可使用统一地图资产库。工程地图、认证资产、内置程序地图和系统参数化地形共用一条“点击/拖放 → 轻量预览 → 场景树管理 → 一次应用”链路;任何来源的增删或参数修改都先进入同一个场景草稿,不会隐式触发 MuJoCo 编译。参数化地形包括离散障碍、沟壑、倒金字塔阶梯、深坑、金字塔阶梯、轨道、随机粗糙、踏石和波浪地形;相同参数与随机种子会确定性生成相同碰撞层。卡片可直接拖到三维画布,落点按 0.1 m 吸附;所有内置程序地图都可在视口选择、移动和绕 Z 旋转。认证资产的“自动落位(重力)”会查询程序地图真实表面,可落到坡道、楼梯、坑底或高度场,而不是固定假设 `z=0`。最后点击“一次编译应用”才组合 MJCF;编译失败时保留上一个可用仿真会话和完整草稿,普通模型重载也不会提前提交草稿。
工程地图由 `map.json`、静态 MJCF 碰撞层和可选的自包含 GLB 视觉层组成:
```text
maps/warehouse/
├── map.json
├── physics/world.xml
├── physics/meshes/*.obj
├── visuals/scene.glb
└── authoring/map.scene.json # V3 可选创作层
```
最小描述示例:
```json
{
"schemaVersion": 1,
"id": "warehouse",
"name": "仓库",
"coordinateSystem": { "units": "m", "up": "Z", "forward": "+X" },
"physics": { "source": "physics/world.xml" },
"visual": { "source": "visuals/scene.glb" },
"spawnPoints": [{ "id": "main", "name": "主入口", "position": [0, 0, 0.35], "yawDeg": 0 }]
}
```
物理地图仅允许静态 `worldbody` 以及 mesh、heightfield、texture、material 等基础 asset,不允许 joint、mocap body、actuator、sensor、include 或 default class。OBJ/STL 应使用简化碰撞模型;高精度模型只放入 GLB。GLB 必须是 2.0 自包含文件,外部 URI 会被拒绝。地图统一使用米制、Z-up、+X 前向坐标系。
V3 可编辑地图使用 `schemaVersion: 2`,并增加 `"authoring": { "source": "authoring/map.scene.json" }`。创作层支持方盒、圆柱、胶囊、坡道、楼梯和出生点。“场景 · 资产库”中的认证资产可点击添加或拖到画布落位;没有可编辑地图时,首个资产会创建临时 Schema V2 场景,但不触发 MuJoCo 重编译,“放弃场景更改”会同时移除临时实例和三份临时工程文件。新增对象支持自动贴地、沿世界 `-Z` 自动重力落位及锁定位姿;属性面板可编辑位置、绕 Z 旋转、原语尺寸、摩擦、颜色和透明度。工程地图源与场景实例分离:同一来源可以重复放置,实例拥有独立 XY/绕 Z 位姿,物理层、GLB 视觉层和出生点同步变换;编辑 authoring 源内容会明确同步到同源实例。视口支持吸附、拾取、空白取消选择、复制、删除及对齐地面;`W`/`E`/`S` 切换移动、旋转和缩放工具,`Delete` 删除,`Ctrl+Z`/`Ctrl+Y` 撤销重做。浏览器不会直接写回原目录,可使用“导出地图 ZIP”下载当前已提交地图包。没有 `authoring.source` 的 V1/V2 地图默认只读;仅由 `box`、`cylinder`、`capsule` 构成且不含 asset、材质、碰撞过滤或隐藏姿态语义的静态 MJCF,可通过“创建可编辑副本”显式升级;任何不可逆语义都会导致整体拒绝,不会静默丢失内容。
物理地图超过 2000 个 geom 会产生性能警告,超过 10000 个会被拒绝;GLB 超过 100 万三角面会警告,超过 300 万会被拒绝。当前原生 URDF 模式不支持地图,请切换到“转换为 MJCF”。
## 数据记录
加载模型后打开右侧“数据”标签,可以选择需要跟踪的 Body、采样频率和样本上限。内置通道包括世界系位置/速度、水平与三维速度、机身侧倾/俯仰/偏航角及角速度、累计里程、接触数、控制输入 RMS、驱动力 RMS、绝对驱动功率和广义速度 RMS。仿真重置不会删除已有数据,而是创建新分段,避免跨重置计算出错误速度;切换记录 Body 或采样配置会清空不兼容的旧数据。
记录只保留在当前浏览器会话中,达到样本上限后自动停止,可导出带稳定列名的 CSV 或包含通道元数据、摘要和样本的 Schema V1 JSON。`SimulationSession`/`PhysicsAdapter` 保留 `configureDataRecorder`、`startDataRecording`、`stopDataRecording`、`clearDataRecording`、`exportDataRecording` 接口;还可通过 `registerDataChannel({ key, label, unit, read })` 在开始记录前注册业务自定义标量通道。数据源通过 `TelemetrySource` 抽象与 MuJoCo 解耦,后续可复用于 Worker 或远端仿真。
## Python 控制器
Python 控制器是可信的单文件脚本,必须同步定义 `step(ctx, state)`;可选定义 `NAME`、`CONTROL_HZ`(限制为 1–500 Hz)、`init(api)`、`command(name, state)`、`reset(state)` 和 `dispose(state)`。`init` 可用 `api.joint(name)`、`api.actuator(name)`、`api.sensor(name)`、`api.body(name)` 预解析 ID;`step` 可用 `ctx.qpos(id)`、`ctx.qvel(id)`、`ctx.sensor(id)`、`ctx.body_quat(id)`、`ctx.body_position(id)` 读取状态,并用 `ctx.set_control(id, value)` 写入经过有限值检查和 actuator 限幅的控制量。定义 `command` 后,界面会显示停止、前进、后退、左转、右转和起跳按钮,并分别传入 `stop`、`forward`、`backward`、`turn_left`、`turn_right`、`jump`。所有回调都必须同步;异常会自动停止控制器或显示诊断,运行期异常还会暂停仿真并清零 `ctrl`。
@@ -112,18 +69,17 @@ Python 控制器是可信的单文件脚本,必须同步定义 `step(ctx, stat
## 本地强化学习训练
训练仍由本机 Python/mjlab 进程执行,但可以从右侧“控制 → 本地强化学习训练”直接发起和管理。仓库已经内置默认 Go2 任务的训练代码与资产;先按 [`training_server/README.md`](../training_server/README.md) 安装训练依赖,再使用对应 Python 启动本地桥接服务:
训练仍由本机 Python/mjlab 进程执行,但可以从右侧“控制 → 本地强化学习训练”直接发起和管理。先使用安装了 mjlab、PyTorch 及训练依赖的 Python 启动本地桥接服务:
```bash
npm run training-server -- \
--trainer-root /path/to/unitree_rl_mjlab \
--trainer-python /path/to/training-env/bin/python
```
服务启动时会在终端输出一个随机访问令牌;在界面中填写该令牌后连接。令牌仅保存在当前标签页的 `sessionStorage`。界面默认连接 `http://127.0.0.1:8765`,可选择服务端允许的任务、并行环境数、训练迭代、随机种子、CPU/GPU、GPU 编号和实验记录方式。W&B 默认为本地离线模式,无需登录或 API Key;也可完全禁用,只有明确选择在线模式时才会联网登录。训练期间页面轮询迭代进度与最近日志,可以停止任务;训练成功后点击“导入策略”,生成的 `policy.onnx` 会进入现有 ONNX 加载流程。普通训练还可以选择自调参产生的命名 reward preset,而不会改写仓库默认配置。
界面默认连接 `http://127.0.0.1:8765`,可选择服务端允许的任务、并行环境数、训练迭代、随机种子、CPU/GPU、GPU 编号和实验记录方式。W&B 默认为本地离线模式,无需登录或 API Key;也可完全禁用,只有明确选择在线模式时才会联网登录。训练期间页面轮询迭代进度与最近日志,可以停止任务;训练成功后点击“导入策略”,生成的 `policy.onnx` 会进入现有 ONNX 加载流程。
连接服务后点击“打开自调参 Agent 工作台”会打开独立 `tuning.html`。该页面采用 Cyber-Industrial 三栏控制台:左侧展示 Session/ASHA 晋级树和 Trial 对比选择,中间使用 uPlot 叠加多 Trial 增量收敛曲线(金线标记历史最优)及六维物理评分,右侧以可折叠因果时间线展示 rationale、expected impact、置信度和评估结果。工具栏支持自动/逐轮审批切换、一次一 Trial 的调度令牌、服务端参数范围/固定值护栏、回滚历史最优或复现任意安全 Trial;Reward Merge Patch 与回滚差异由按需加载的 Monaco Diff 审查。scalar 以 1 Hz 非重入方式增量轮询,进入固定容量环形缓冲并由 `requestAnimationFrame` 合批后调用 `uPlot.setData`,不会在每次轮询时重建图表。新标签页 URL 不包含 token;同源 opener 会一次性交接凭据,直接打开页面时也可手工输入。DeepSeek key 始终由本地 Python 服务的 `DEEPSEEK_API_KEY` 环境变量读取,浏览器不会接触该 key。
桥接服务只监听本机回环地址,并检查 Host、Origin 和 Bearer Token;仅接受允许列表中的任务和经过范围校验的参数,不执行前端提供的 Shell 命令;一次只运行一个训练进程。默认任务使用仓库内置的 Go2 机器人资产与环境配置,**不会自动把浏览器中临时编辑的 MJCF/URDF 作为训练环境**。自定义浏览器模型训练需要在兼容的外部训练工程中注册 task,并通过服务的 `--trainer-root` 指定该工程。服务配置、接口和安全边界见 [`../training_server/README.md`](../training_server/README.md)。
桥接服务只监听本机回环地址、仅接受允许列表中的任务和经过范围校验的参数,不执行前端提供的 Shell 命令;一次只运行一个训练进程。当前任务使用 `unitree_rl_mjlab` 自带的机器人资产与环境配置,**不会自动把浏览器中临时编辑的 MJCF/URDF 作为训练环境**。自定义浏览器模型训练需要先在 mjlab 中注册对应 task。服务配置、接口和安全边界见 [`../training_server/README.md`](../training_server/README.md)。
## ONNX 强化学习策略
File diff suppressed because it is too large Load Diff
+1 -72
View File
@@ -1,72 +1 @@
<!doctype html>
<html lang="zh-CN">
<head>
<meta charset="UTF-8" />
<meta name="viewport" content="width=device-width,initial-scale=1" />
<meta name="theme-color" content="#09111e" />
<meta
name="description"
content="基于 MuJoCo WebAssembly 的本地机器人仿真、控制与地图编辑工作台"
/>
<link rel="icon" href="data:," />
<title>MuJoCo Web 仿真平台</title>
<style>
html,
body,
#root {
height: 100%;
margin: 0;
}
body {
background: #09111e;
}
.boot-screen {
display: grid;
height: 100%;
place-items: center;
color: #f1f5f9;
font:
13px Inter,
system-ui,
sans-serif;
}
.boot-mark {
display: grid;
width: 42px;
height: 42px;
margin: 0 auto 14px;
place-items: center;
border: 1px solid #2b604f;
border-radius: 13px;
background: #123b31;
color: #38d39f;
font-weight: 800;
box-shadow: 0 16px 48px rgb(0 0 0 / 35%);
animation: boot-pulse 1.4s ease-in-out infinite;
}
.boot-caption {
color: #8fa0b5;
font-size: 11px;
letter-spacing: 0.04em;
text-align: center;
}
@keyframes boot-pulse {
50% {
transform: translateY(-2px);
box-shadow: 0 18px 54px rgb(56 211 159 / 15%);
}
}
</style>
</head>
<body>
<div id="root">
<div class="boot-screen" role="status" aria-label="正在启动仿真工作台">
<div>
<div class="boot-mark">M</div>
<div class="boot-caption">正在启动本地仿真工作台…</div>
</div>
</div>
</div>
<script type="module" src="/src/main.tsx"></script>
</body>
</html>
<!doctype html><html lang="zh-CN"><head><meta charset="UTF-8"/><meta name="viewport" content="width=device-width,initial-scale=1"/><meta name="theme-color" content="#101827"/><link rel="icon" href="data:,"/><title>MuJoCo Web 仿真平台</title></head><body><div id="root"></div><script type="module" src="/src/main.tsx"></script></body></html>
+2 -11
View File
@@ -1,11 +1,2 @@
import { defineConfig } from '@playwright/test';
export default defineConfig({
testDir: './e2e',
timeout: 120_000,
use: { baseURL: 'http://127.0.0.1:4173' },
webServer: {
command: 'npm run preview --prefix .. -- --host 127.0.0.1',
url: 'http://127.0.0.1:4173',
reuseExistingServer: true,
},
});
import {defineConfig} from '@playwright/test';
export default defineConfig({testDir:'./e2e', timeout:120_000, use:{baseURL:'http://127.0.0.1:4173',channel:'chrome'}, webServer:{command:'npm run preview --prefix .. -- --host 127.0.0.1',url:'http://127.0.0.1:4173',reuseExistingServer:true}});
+1 -3
View File
@@ -1,3 +1 @@
module.exports = {
plugins: { tailwindcss: { config: './web_platform/tailwind.config.cjs' }, autoprefixer: {} },
};
module.exports = {plugins: {tailwindcss: {config: './web_platform/tailwind.config.cjs'}, autoprefixer: {}}};
+113 -2760
View File
File diff suppressed because it is too large Load Diff
+3 -28
View File
@@ -1,28 +1,3 @@
import { Component, type ErrorInfo, type ReactNode } from 'react';
import { Button } from '../components/ui';
export class ErrorBoundary extends Component<{ children: ReactNode }, { error?: Error }> {
state: { error?: Error } = {};
static getDerivedStateFromError(error: Error) {
return { error };
}
componentDidCatch(error: Error, info: ErrorInfo) {
console.error('React fatal error', error, info);
}
render() {
return this.state.error ? (
<main className="grid h-screen place-items-center bg-app text-text-primary">
<section className="max-w-xl rounded-xl border border-danger-border bg-panel p-6 shadow-xl">
<h1 className="text-xl font-semibold">界面发生致命错误</h1>
<pre className="mt-3 whitespace-pre-wrap text-sm text-danger">
{this.state.error.message}
</pre>
<Button variant="danger" className="mt-4" onClick={() => location.reload()}>
重新加载
</Button>
</section>
</main>
) : (
this.props.children
);
}
}
import {Component,type ErrorInfo,type ReactNode} from 'react';
import {Button} from '../components/ui';
export class ErrorBoundary extends Component<{children:ReactNode},{error?:Error}>{state:{error?:Error}={};static getDerivedStateFromError(error:Error){return {error};}componentDidCatch(error:Error,info:ErrorInfo){console.error('React fatal error',error,info);}render(){return this.state.error?<main className="grid h-screen place-items-center bg-app text-text-primary"><section className="max-w-xl rounded-xl border border-danger-border bg-panel p-6 shadow-xl"><h1 className="text-xl font-semibold">界面发生致命错误</h1><pre className="mt-3 whitespace-pre-wrap text-sm text-danger">{this.state.error.message}</pre><Button variant="danger" className="mt-4" onClick={()=>location.reload()}>重新加载</Button></section></main>:this.props.children;}}
@@ -0,0 +1,41 @@
import {fireEvent,render,screen} from '@testing-library/react';
import {ActuatorControl} from './SidebarPanel';
import type {ActuatorInfo} from '../../simulation/SimulationSession';
const actuator:ActuatorInfo={id:0,name:'shoulder_motor',value:.5,min:-1,max:1,limited:true,jointId:0,jointName:'shoulder',jointType:3,unit:'N·m',kind:'motor',controlCount:1,gear:2,gain:1,kp:0,kv:0,ctrlLimited:true,ctrlMin:-1,ctrlMax:1,forceLimited:true,forceMin:-20,forceMax:20};
describe('ActuatorControl',()=>{
it('显示对应关节和常用力矩单位',()=>{
render(<ActuatorControl actuator={actuator} onControl={()=>{}} onParameters={()=>{}}/>);
expect(screen.getByText('shoulder_motor')).toBeVisible();
expect(screen.getByText('关节:shoulder')).toBeVisible();
expect(screen.getByText('1.000 N·m')).toBeVisible();
});
it('内部按 gear 换算输出,但参数面板只开放 kp、kv 等业务参数',()=>{
const onControl=vi.fn(),onParameters=vi.fn();
render(<ActuatorControl actuator={actuator} onControl={onControl} onParameters={onParameters}/>);
fireEvent.change(screen.getByRole('slider'),{target:{value:'2'}});
expect(onControl).toHaveBeenCalledWith(1);
fireEvent.click(screen.getByText('常用参数'));
expect(screen.queryByLabelText('传动比 gear')).not.toBeInTheDocument();
expect(screen.queryByLabelText('固定增益 gain')).not.toBeInTheDocument();
const kp=screen.getByLabelText(/kp(MJCF stiffness/);fireEvent.change(kp,{target:{value:'3'}});fireEvent.blur(kp);
expect(onParameters).toHaveBeenCalledWith(expect.objectContaining({kp:3,ctrlLimited:true,forceLimited:true}));
});
it('position 伺服使用角度目标并开放 kp、kv',()=>{
const onParameters=vi.fn();
render(<ActuatorControl actuator={{...actuator,name:'shoulder_servo',kind:'position',unit:'°',value:Math.PI/2,min:-Math.PI,max:Math.PI,gear:1,kp:100,kv:10,gain:100}} onControl={()=>{}} onParameters={onParameters}/>);
expect(screen.getByText('90.000 °')).toBeVisible();fireEvent.click(screen.getByText('常用参数'));
const kp=screen.getByLabelText(/位置增益 kp/);fireEvent.change(kp,{target:{value:'150'}});fireEvent.blur(kp);
expect(onParameters).toHaveBeenCalledWith(expect.objectContaining({kp:150,kv:10}));
});
it('非 motor 驱动器保持原始控制单位且不开放通用参数编辑',()=>{
render(<ActuatorControl actuator={{...actuator,name:'custom',kind:'other',unit:'',value:.25}} onControl={()=>{}} onParameters={()=>{}}/>);
expect(screen.getByText('0.250')).toBeVisible();
expect(screen.queryByText('常用参数')).not.toBeInTheDocument();
expect(screen.getByText(/不是可直接编辑的 motor\/position/)).toBeVisible();
});
});
@@ -1,120 +1,12 @@
import { useEffect, useId, useMemo, useRef, useState, type ReactNode } from 'react';
import { Search } from 'lucide-react';
import { Dialog, EmptySearchState, Kbd } from '../../components/ui';
export interface WorkbenchCommand {
id: string;
label: string;
group: string;
icon?: ReactNode;
shortcut?: string;
disabled?: boolean;
run: () => void;
}
export function CommandPalette({
open,
onClose,
commands,
}: {
open: boolean;
onClose: () => void;
commands: WorkbenchCommand[];
}) {
const [query, setQuery] = useState(''),
[active, setActive] = useState(0),
input = useRef<HTMLInputElement>(null),
listId = useId();
const filtered = useMemo(() => {
const needle = query.trim().toLocaleLowerCase();
return commands.filter(
(command) =>
!needle || `${command.label} ${command.group}`.toLocaleLowerCase().includes(needle),
);
}, [commands, query]);
const enabled = filtered.flatMap((command, index) => (command.disabled ? [] : [index])),
highlighted = filtered[active] && !filtered[active].disabled ? active : (enabled[0] ?? -1);
useEffect(() => {
if (open) requestAnimationFrame(() => input.current?.focus());
}, [open]);
const close = () => {
setQuery('');
setActive(0);
onClose();
},
execute = (command?: WorkbenchCommand) => {
if (!command || command.disabled) return;
command.run();
close();
};
return (
<Dialog open={open} onClose={close} title="命令面板" className="max-w-xl">
<div className="relative -m-4 mb-2 border-b border-border">
<Search className="pointer-events-none absolute left-4 top-1/2 h-4 w-4 -translate-y-1/2 text-text-tertiary" />
<input
ref={input}
role="combobox"
aria-label="搜索命令"
aria-autocomplete="list"
aria-expanded="true"
aria-controls={listId}
aria-activedescendant={
highlighted >= 0 ? `${listId}-${filtered[highlighted].id}` : undefined
}
value={query}
onChange={(event) => {
setQuery(event.target.value);
setActive(0);
}}
onKeyDown={(event) => {
if (!enabled.length) return;
const current = Math.max(0, enabled.indexOf(highlighted));
if (event.key === 'ArrowDown') {
event.preventDefault();
setActive(enabled[(current + 1) % enabled.length]);
} else if (event.key === 'ArrowUp') {
event.preventDefault();
setActive(enabled[(current - 1 + enabled.length) % enabled.length]);
} else if (event.key === 'Enter') {
event.preventDefault();
execute(filtered[highlighted]);
}
}}
placeholder="输入命令名称…"
className="h-11 w-full bg-input pl-11 pr-4 text-sm text-text-primary placeholder:text-text-tertiary focus:outline-none"
/>
</div>
<div
id={listId}
role="listbox"
aria-label="可用命令"
className="max-h-80 space-y-1 overflow-auto pt-1"
>
{filtered.length ? (
filtered.map((command, index) => (
<button
key={command.id}
id={`${listId}-${command.id}`}
type="button"
role="option"
aria-selected={index === highlighted}
disabled={command.disabled}
onMouseEnter={() => {
if (!command.disabled) setActive(index);
}}
onClick={() => execute(command)}
className={`flex w-full items-center gap-3 rounded-md px-3 py-2 text-left text-xs outline-none ${index === highlighted ? 'bg-accent-soft text-accent' : 'text-text-secondary hover:bg-element-hover'} disabled:opacity-40`}
>
<span className="flex h-5 w-5 items-center justify-center">{command.icon}</span>
<span className="min-w-0 flex-1">
<span className="block truncate font-medium">{command.label}</span>
<span className="block text-[10px] text-text-tertiary">{command.group}</span>
</span>
{command.shortcut && <Kbd>{command.shortcut}</Kbd>}
</button>
))
) : (
<EmptySearchState label="没有匹配的命令" />
)}
</div>
</Dialog>
);
import {useEffect,useId,useMemo,useRef,useState,type ReactNode} from 'react';
import {Search} from 'lucide-react';
import {Dialog,EmptySearchState,Kbd} from '../../components/ui';
export interface WorkbenchCommand{id:string;label:string;group:string;icon?:ReactNode;shortcut?:string;disabled?:boolean;run:()=>void;}
export function CommandPalette({open,onClose,commands}:{open:boolean;onClose:()=>void;commands:WorkbenchCommand[]}){
const [query,setQuery]=useState(''),[active,setActive]=useState(0),input=useRef<HTMLInputElement>(null),listId=useId();
const filtered=useMemo(()=>{const needle=query.trim().toLocaleLowerCase();return commands.filter(command=>!needle||`${command.label} ${command.group}`.toLocaleLowerCase().includes(needle));},[commands,query]);
const enabled=filtered.flatMap((command,index)=>command.disabled?[]:[index]),highlighted=filtered[active]&&!filtered[active].disabled?active:(enabled[0]??-1);
useEffect(()=>{if(open)requestAnimationFrame(()=>input.current?.focus());},[open]);
const close=()=>{setQuery('');setActive(0);onClose();},execute=(command?:WorkbenchCommand)=>{if(!command||command.disabled)return;command.run();close();};
return <Dialog open={open} onClose={close} title="命令面板" className="max-w-xl"><div className="relative -m-4 mb-2 border-b border-border"><Search className="pointer-events-none absolute left-4 top-1/2 h-4 w-4 -translate-y-1/2 text-text-tertiary"/><input ref={input} role="combobox" aria-label="搜索命令" aria-autocomplete="list" aria-expanded="true" aria-controls={listId} aria-activedescendant={highlighted>=0?`${listId}-${filtered[highlighted].id}`:undefined} value={query} onChange={event=>{setQuery(event.target.value);setActive(0);}} onKeyDown={event=>{if(!enabled.length)return;const current=Math.max(0,enabled.indexOf(highlighted));if(event.key==='ArrowDown'){event.preventDefault();setActive(enabled[(current+1)%enabled.length]);}else if(event.key==='ArrowUp'){event.preventDefault();setActive(enabled[(current-1+enabled.length)%enabled.length]);}else if(event.key==='Enter'){event.preventDefault();execute(filtered[highlighted]);}}} placeholder="输入命令名称…" className="h-11 w-full bg-input pl-11 pr-4 text-sm text-text-primary placeholder:text-text-tertiary focus:outline-none"/></div><div id={listId} role="listbox" aria-label="可用命令" className="max-h-80 space-y-1 overflow-auto pt-1">{filtered.length?filtered.map((command,index)=><button key={command.id} id={`${listId}-${command.id}`} type="button" role="option" aria-selected={index===highlighted} disabled={command.disabled} onMouseEnter={()=>{if(!command.disabled)setActive(index);}} onClick={()=>execute(command)} className={`flex w-full items-center gap-3 rounded-md px-3 py-2 text-left text-xs outline-none ${index===highlighted?'bg-accent-soft text-accent':'text-text-secondary hover:bg-element-hover'} disabled:opacity-40`}><span className="flex h-5 w-5 items-center justify-center">{command.icon}</span><span className="min-w-0 flex-1"><span className="block truncate font-medium">{command.label}</span><span className="block text-[10px] text-text-tertiary">{command.group}</span></span>{command.shortcut&&<Kbd>{command.shortcut}</Kbd>}</button>):<EmptySearchState label="没有匹配的命令"/>}</div></Dialog>;
}
@@ -1,50 +1,5 @@
import { useState } from 'react';
import { ChevronDown, TriangleAlert, X } from 'lucide-react';
import type { AppDiagnostic } from '../../stores/useAppStore';
import { IconButton } from '../../components/ui';
export function DiagnosticNotice({
value,
onClose,
}: {
value: AppDiagnostic;
onClose: () => void;
}) {
const [expanded, setExpanded] = useState(false);
return (
<section
role="alert"
className="absolute bottom-4 left-1/2 z-30 w-[min(42rem,calc(100%-2rem))] -translate-x-1/2 overflow-hidden rounded-xl border border-danger-border bg-panel shadow-2xl"
>
<div className="flex items-start gap-3 p-3">
<span className="mt-0.5 grid h-7 w-7 shrink-0 place-items-center rounded-full bg-danger-soft text-danger">
<TriangleAlert className="h-4 w-4" />
</span>
<div className="min-w-0 flex-1">
<h2 className="text-sm font-semibold text-text-primary">{value.summary}</h2>
{value.path && (
<p className="mt-0.5 truncate text-xs text-text-tertiary" title={value.path}>
路径:{value.path}
</p>
)}
<button
type="button"
aria-expanded={expanded}
className="mt-1 flex items-center gap-1 text-xs text-danger hover:underline"
onClick={() => setExpanded((v) => !v)}
>
技术详情
<ChevronDown className={`h-3 w-3 ${expanded ? 'rotate-180' : ''}`} />
</button>
</div>
<IconButton aria-label="关闭错误" tooltip="关闭" onClick={onClose}>
<X className="h-4 w-4" />
</IconButton>
</div>
{expanded && (
<pre className="max-h-36 overflow-auto border-t border-danger-border bg-danger-soft p-3 text-xs text-danger">
{value.detail}
</pre>
)}
</section>
);
}
import {useState} from 'react';
import {ChevronDown,TriangleAlert,X} from 'lucide-react';
import type {AppDiagnostic} from '../../stores/useAppStore';
import {IconButton} from '../../components/ui';
export function DiagnosticNotice({value,onClose}:{value:AppDiagnostic;onClose:()=>void}){const [expanded,setExpanded]=useState(false);return <section role="alert" className="absolute bottom-4 left-1/2 z-30 w-[min(42rem,calc(100%-2rem))] -translate-x-1/2 overflow-hidden rounded-xl border border-danger-border bg-panel shadow-2xl"><div className="flex items-start gap-3 p-3"><span className="mt-0.5 grid h-7 w-7 shrink-0 place-items-center rounded-full bg-danger-soft text-danger"><TriangleAlert className="h-4 w-4"/></span><div className="min-w-0 flex-1"><h2 className="text-sm font-semibold text-text-primary">{value.summary}</h2>{value.path&&<p className="mt-0.5 truncate text-xs text-text-tertiary" title={value.path}>路径:{value.path}</p>}<button type="button" aria-expanded={expanded} className="mt-1 flex items-center gap-1 text-xs text-danger hover:underline" onClick={()=>setExpanded(v=>!v)}>技术详情<ChevronDown className={`h-3 w-3 ${expanded?'rotate-180':''}`}/></button></div><IconButton aria-label="关闭错误" tooltip="关闭" onClick={onClose}><X className="h-4 w-4"/></IconButton></div>{expanded&&<pre className="max-h-36 overflow-auto border-t border-danger-border bg-danger-soft p-3 text-xs text-danger">{value.detail}</pre>}</section>;}
@@ -1,97 +1,6 @@
import { useState } from 'react';
import { CheckCircle2, Info, TriangleAlert, XCircle } from 'lucide-react';
import { Button, CopyButton, Dialog, Tabs } from '../../components/ui';
import type { WorkbenchNotification } from './NotificationCenter';
type Filter = 'all' | 'warning' | 'danger';
export function DiagnosticsDrawer({
open,
items,
onClose,
onClear,
}: {
open: boolean;
items: WorkbenchNotification[];
onClose: () => void;
onClear: () => void;
}) {
const [filter, setFilter] = useState<Filter>('all');
const content = (value: Filter) => {
const filtered = items.filter((item) => value === 'all' || item.tone === value);
return (
<div className="space-y-2">
{filtered.length ? (
filtered.map((item) => {
const Icon =
item.tone === 'danger'
? XCircle
: item.tone === 'warning'
? TriangleAlert
: item.tone === 'success'
? CheckCircle2
: Info;
return (
<article key={item.id} className="rounded-lg border border-border bg-surface p-3">
<div className="flex items-start gap-2">
<Icon
className={`mt-0.5 h-4 w-4 ${item.tone === 'danger' ? 'text-danger' : item.tone === 'warning' ? 'text-warning' : 'text-success'}`}
/>
<div className="min-w-0 flex-1">
<h3 className="text-xs font-semibold">{item.title}</h3>
<time className="text-[10px] text-text-tertiary">
{new Date(item.at).toLocaleString('zh-CN')}
</time>
{item.detail && (
<pre className="mt-2 whitespace-pre-wrap text-[10px] leading-4 text-text-secondary">
{item.detail}
</pre>
)}
</div>
{item.detail && (
<CopyButton value={`${item.title}\n${item.detail}`} label="复制事件详情" />
)}
</div>
</article>
);
})
) : (
<p className="p-8 text-center text-xs text-text-tertiary">没有符合条件的事件</p>
)}
</div>
);
};
return (
<Dialog
open={open}
onClose={onClose}
title="诊断与事件日志"
className="max-w-2xl"
footer={
<div className="flex justify-end">
<Button variant="danger" disabled={!items.length} onClick={onClear}>
清空事件
</Button>
</div>
}
>
<Tabs
label="事件筛选"
value={filter}
onValueChange={setFilter}
keepMounted={false}
items={[
{ value: 'all', label: `全部 ${items.length}`, content: content('all') },
{
value: 'warning',
label: `警告 ${items.filter((item) => item.tone === 'warning').length}`,
content: content('warning'),
},
{
value: 'danger',
label: `错误 ${items.filter((item) => item.tone === 'danger').length}`,
content: content('danger'),
},
]}
/>
</Dialog>
);
}
import {useState} from 'react';
import {CheckCircle2,Info,TriangleAlert,XCircle} from 'lucide-react';
import {Button,CopyButton,Dialog,Tabs} from '../../components/ui';
import type {WorkbenchNotification} from './NotificationCenter';
type Filter='all'|'warning'|'danger';
export function DiagnosticsDrawer({open,items,onClose,onClear}:{open:boolean;items:WorkbenchNotification[];onClose:()=>void;onClear:()=>void}){const [filter,setFilter]=useState<Filter>('all');const content=(value:Filter)=>{const filtered=items.filter(item=>value==='all'||item.tone===value);return <div className="space-y-2">{filtered.length?filtered.map(item=>{const Icon=item.tone==='danger'?XCircle:item.tone==='warning'?TriangleAlert:item.tone==='success'?CheckCircle2:Info;return <article key={item.id} className="rounded-lg border border-border bg-surface p-3"><div className="flex items-start gap-2"><Icon className={`mt-0.5 h-4 w-4 ${item.tone==='danger'?'text-danger':item.tone==='warning'?'text-warning':'text-success'}`}/><div className="min-w-0 flex-1"><h3 className="text-xs font-semibold">{item.title}</h3><time className="text-[10px] text-text-tertiary">{new Date(item.at).toLocaleString('zh-CN')}</time>{item.detail&&<pre className="mt-2 whitespace-pre-wrap text-[10px] leading-4 text-text-secondary">{item.detail}</pre>}</div>{item.detail&&<CopyButton value={`${item.title}\n${item.detail}`} label="复制事件详情"/>}</div></article>}):<p className="p-8 text-center text-xs text-text-tertiary">没有符合条件的事件</p>}</div>;};return <Dialog open={open} onClose={onClose} title="诊断与事件日志" className="max-w-2xl" footer={<div className="flex justify-end"><Button variant="danger" disabled={!items.length} onClick={onClear}>清空事件</Button></div>}><Tabs label="事件筛选" value={filter} onValueChange={setFilter} keepMounted={false} items={[{value:'all',label:`全部 ${items.length}`,content:content('all')},{value:'warning',label:`警告 ${items.filter(item=>item.tone==='warning').length}`,content:content('warning')},{value:'danger',label:`错误 ${items.filter(item=>item.tone==='danger').length}`,content:content('danger')}]}/></Dialog>;}
@@ -1,18 +1,4 @@
import { render, screen } from '@testing-library/react';
import { EntrySelectionDialog } from './EntrySelectionDialog';
import {render,screen} from '@testing-library/react';
import {EntrySelectionDialog} from './EntrySelectionDialog';
describe('EntrySelectionDialog', () => {
it('父组件重渲染时不抢走入口按钮焦点,且不暴露无效关闭动作', () => {
const entries = [
{ path: 'a.xml', label: '模型 A' },
{ path: 'b.xml', label: '模型 B' },
],
select = vi.fn();
const { rerender } = render(<EntrySelectionDialog entries={entries} onSelect={select} />);
const entry = screen.getByRole('button', { name: '模型 A' });
entry.focus();
rerender(<EntrySelectionDialog entries={[...entries]} onSelect={select} />);
expect(entry).toHaveFocus();
expect(screen.queryByRole('button', { name: '关闭' })).not.toBeInTheDocument();
});
});
describe('EntrySelectionDialog',()=>{it('父组件重渲染时不抢走入口按钮焦点,且不暴露无效关闭动作',()=>{const entries=[{path:'a.xml',label:'模型 A'},{path:'b.xml',label:'模型 B'}],select=vi.fn();const {rerender}=render(<EntrySelectionDialog entries={entries} onSelect={select}/>);const entry=screen.getByRole('button',{name:'模型 A'});entry.focus();rerender(<EntrySelectionDialog entries={[...entries]} onSelect={select}/>);expect(entry).toHaveFocus();expect(screen.queryByRole('button',{name:'关闭'})).not.toBeInTheDocument();});});
@@ -1,28 +1,4 @@
import { FileCode2 } from 'lucide-react';
import { Button, Dialog } from '../../components/ui';
const noop = () => {};
export function EntrySelectionDialog({
entries,
onSelect,
}: {
entries: { path: string; label: string }[];
onSelect: (path: string) => void;
}) {
return (
<Dialog open={entries.length > 0} onClose={noop} closable={false} title="选择模型入口">
<p className="mb-4 text-sm text-text-secondary">工程包含多个可加载模型,请选择一个。</p>
<div className="space-y-2">
{entries.map((entry) => (
<Button
key={entry.path}
className="w-full justify-start overflow-hidden"
onClick={() => onSelect(entry.path)}
icon={<FileCode2 className="h-4 w-4" />}
>
<span className="truncate">{entry.label}</span>
</Button>
))}
</div>
</Dialog>
);
}
import {FileCode2} from 'lucide-react';
import {Button,Dialog} from '../../components/ui';
const noop=()=>{};
export function EntrySelectionDialog({entries,onSelect}:{entries:{path:string;label:string}[];onSelect:(path:string)=>void}){return <Dialog open={entries.length>0} onClose={noop} closable={false} title="选择模型入口"><p className="mb-4 text-sm text-text-secondary">工程包含多个可加载模型,请选择一个。</p><div className="space-y-2">{entries.map(entry=><Button key={entry.path} className="w-full justify-start overflow-hidden" onClick={()=>onSelect(entry.path)} icon={<FileCode2 className="h-4 w-4"/>}><span className="truncate">{entry.label}</span></Button>)}</div></Dialog>;}
@@ -1,68 +1,5 @@
import { useState } from 'react';
import { ChevronDown, FolderTree, RefreshCw, TriangleAlert, X } from 'lucide-react';
import type { AppDiagnostic } from '../../stores/useAppStore';
import { Button, CopyButton, IconButton } from '../../components/ui';
export function ErrorRecoveryPanel({
value,
onClose,
onRetry,
onOpenProject,
}: {
value: AppDiagnostic;
onClose: () => void;
onRetry?: () => void;
onOpenProject: () => void;
}) {
const [expanded, setExpanded] = useState(false);
return (
<section
role="alert"
className="absolute bottom-4 left-1/2 z-30 w-[min(42rem,calc(100%-2rem))] -translate-x-1/2 overflow-hidden rounded-xl border border-danger-border bg-panel shadow-2xl"
>
<div className="flex items-start gap-3 p-3">
<span className="mt-0.5 grid h-7 w-7 shrink-0 place-items-center rounded-full bg-danger-soft text-danger">
<TriangleAlert className="h-4 w-4" />
</span>
<div className="min-w-0 flex-1">
<h2 className="text-sm font-semibold">{value.summary}</h2>
{value.path && <p className="truncate text-xs text-text-tertiary">路径:{value.path}</p>}
<div className="mt-2 flex flex-wrap gap-2">
{onRetry && (
<Button
variant="danger"
onClick={onRetry}
icon={<RefreshCw className="h-3.5 w-3.5" />}
>
重试当前入口
</Button>
)}
<Button onClick={onOpenProject} icon={<FolderTree className="h-3.5 w-3.5" />}>
返回工程树
</Button>
<CopyButton
value={`${value.summary}\n${value.path ?? ''}\n${value.detail}`}
label="复制错误详情"
/>
</div>
<button
type="button"
aria-expanded={expanded}
className="mt-2 flex items-center gap-1 text-xs text-danger"
onClick={() => setExpanded((v) => !v)}
>
技术详情
<ChevronDown className={`h-3 w-3 ${expanded ? 'rotate-180' : ''}`} />
</button>
</div>
<IconButton aria-label="关闭错误" tooltip="关闭" onClick={onClose}>
<X className="h-4 w-4" />
</IconButton>
</div>
{expanded && (
<pre className="max-h-36 overflow-auto border-t border-danger-border bg-danger-soft p-3 text-xs text-danger">
{value.detail}
</pre>
)}
</section>
);
}
import {useState} from 'react';
import {ChevronDown,FolderTree,RefreshCw,TriangleAlert,X} from 'lucide-react';
import type {AppDiagnostic} from '../../stores/useAppStore';
import {Button,CopyButton,IconButton} from '../../components/ui';
export function ErrorRecoveryPanel({value,onClose,onRetry,onOpenProject}:{value:AppDiagnostic;onClose:()=>void;onRetry?:()=>void;onOpenProject:()=>void}){const [expanded,setExpanded]=useState(false);return <section role="alert" className="absolute bottom-4 left-1/2 z-30 w-[min(42rem,calc(100%-2rem))] -translate-x-1/2 overflow-hidden rounded-xl border border-danger-border bg-panel shadow-2xl"><div className="flex items-start gap-3 p-3"><span className="mt-0.5 grid h-7 w-7 shrink-0 place-items-center rounded-full bg-danger-soft text-danger"><TriangleAlert className="h-4 w-4"/></span><div className="min-w-0 flex-1"><h2 className="text-sm font-semibold">{value.summary}</h2>{value.path&&<p className="truncate text-xs text-text-tertiary">路径:{value.path}</p>}<div className="mt-2 flex flex-wrap gap-2">{onRetry&&<Button variant="danger" onClick={onRetry} icon={<RefreshCw className="h-3.5 w-3.5"/>}>重试当前入口</Button>}<Button onClick={onOpenProject} icon={<FolderTree className="h-3.5 w-3.5"/>}>返回工程树</Button><CopyButton value={`${value.summary}\n${value.path??''}\n${value.detail}`} label="复制错误详情"/></div><button type="button" aria-expanded={expanded} className="mt-2 flex items-center gap-1 text-xs text-danger" onClick={()=>setExpanded(v=>!v)}>技术详情<ChevronDown className={`h-3 w-3 ${expanded?'rotate-180':''}`}/></button></div><IconButton aria-label="关闭错误" tooltip="关闭" onClick={onClose}><X className="h-4 w-4"/></IconButton></div>{expanded&&<pre className="max-h-36 overflow-auto border-t border-danger-border bg-danger-soft p-3 text-xs text-danger">{value.detail}</pre>}</section>;}
@@ -1,57 +1,5 @@
import { fireEvent, render, screen } from '@testing-library/react';
import { DiagnosticNotice } from './DiagnosticNotice';
import { WorkspaceOverlays } from './WorkspaceOverlays';
import { StatusBar } from './StatusBar';
describe('工作台反馈组件', () => {
it('诊断详情可展开并关闭', () => {
const close = vi.fn();
render(
<DiagnosticNotice
value={{
category: '模型编译',
summary: '模型编译失败',
detail: 'bad xml',
path: 'robot.xml',
at: 1,
}}
onClose={close}
/>,
);
expect(screen.queryByText('bad xml')).not.toBeInTheDocument();
fireEvent.click(screen.getByRole('button', { name: '技术详情' }));
expect(screen.getByText('bad xml')).toBeVisible();
fireEvent.click(screen.getByRole('button', { name: '关闭错误' }));
expect(close).toHaveBeenCalledTimes(1);
});
it('加载态与空态互斥,并展示真实阶段进度', () => {
const { rerender } = render(<WorkspaceOverlays loading={false} hasSnapshot={false} />);
expect(screen.getByText('拖放模型工程到此处')).toBeVisible();
rerender(
<WorkspaceOverlays
loading
hasSnapshot={false}
progress={{
title: '正在导入工程',
label: '读取工程文件',
detail: 'robot.zip',
value: 0.42,
}}
/>,
);
expect(screen.queryByText('拖放模型工程到此处')).not.toBeInTheDocument();
expect(screen.getByRole('status')).toHaveTextContent('正在导入工程');
expect(screen.getByRole('progressbar', { name: '读取工程文件' })).toHaveAttribute(
'aria-valuenow',
'42',
);
});
it('拖入文件时显示明确落点反馈', () => {
render(<WorkspaceOverlays loading={false} hasSnapshot={false} dragActive />);
expect(screen.getByText('松开即可导入工程')).toBeVisible();
});
it('展示格式化状态数据', () => {
render(<StatusBar time={1.25} fps={60} stepMs={0.5} memoryMb={10} loaded overBudget={false} />);
expect(screen.getByText(/时间 1.250 s/)).toBeVisible();
expect(screen.getByText(/WASM 已加载/)).toBeVisible();
});
});
import {fireEvent,render,screen} from '@testing-library/react';
import {DiagnosticNotice} from './DiagnosticNotice';
import {WorkspaceOverlays} from './WorkspaceOverlays';
import {StatusBar} from './StatusBar';
describe('工作台反馈组件',()=>{it('诊断详情可展开并关闭',()=>{const close=vi.fn();render(<DiagnosticNotice value={{category:'模型编译',summary:'模型编译失败',detail:'bad xml',path:'robot.xml',at:1}} onClose={close}/>);expect(screen.queryByText('bad xml')).not.toBeInTheDocument();fireEvent.click(screen.getByRole('button',{name:'技术详情'}));expect(screen.getByText('bad xml')).toBeVisible();fireEvent.click(screen.getByRole('button',{name:'关闭错误'}));expect(close).toHaveBeenCalledTimes(1);});it('加载态与空态互斥',()=>{const {rerender}=render(<WorkspaceOverlays loading={false} hasSnapshot={false}/>);expect(screen.getByText('拖放模型工程到此处')).toBeVisible();rerender(<WorkspaceOverlays loading hasSnapshot={false}/>);expect(screen.queryByText('拖放模型工程到此处')).not.toBeInTheDocument();expect(screen.getByRole('status')).toBeVisible();});it('展示格式化状态数据',()=>{render(<StatusBar time={1.25} fps={60} stepMs={0.5} memoryMb={10} loaded overBudget={false}/>);expect(screen.getByText(/时间 1.250 s/)).toBeVisible();expect(screen.getByText(/WASM 已加载/)).toBeVisible();});});
@@ -1,66 +1,12 @@
import { fireEvent, render, screen, within } from '@testing-library/react';
import { DiagnosticsDrawer } from './DiagnosticsDrawer';
import { ErrorRecoveryPanel } from './ErrorRecoveryPanel';
import { ToolbarOverflowMenu } from './ToolbarOverflowMenu';
import { WorkspaceOverlays } from './WorkspaceOverlays';
const event = { id: 1, title: '编译失败', detail: 'bad xml', tone: 'danger' as const, at: 0 };
describe('第五批工作台组件', () => {
it('事件日志支持分类和清空', () => {
const clear = vi.fn();
render(<DiagnosticsDrawer open items={[event]} onClose={() => {}} onClear={clear} />);
expect(
within(screen.getByRole('tabpanel', { name: /全部/ })).getByText('bad xml'),
).toBeVisible();
expect(screen.getAllByText('bad xml')).toHaveLength(1);
fireEvent.click(screen.getByRole('button', { name: '清空事件' }));
expect(clear).toHaveBeenCalled();
});
it('错误恢复面板透传重试与工程树动作', () => {
const retry = vi.fn(),
project = vi.fn();
render(
<ErrorRecoveryPanel
value={{ category: '模型编译', summary: '失败', detail: 'bad', path: 'a.xml', at: 1 }}
onClose={() => {}}
onRetry={retry}
onOpenProject={project}
/>,
);
fireEvent.click(screen.getByRole('button', { name: '重试当前入口' }));
fireEvent.click(screen.getByRole('button', { name: '返回工程树' }));
expect(retry).toHaveBeenCalled();
expect(project).toHaveBeenCalled();
});
it('导入叠层显示阶段进度', () => {
render(
<div className="relative">
<WorkspaceOverlays
loading
hasSnapshot={false}
progress={{ label: '处理模型资源', value: 0.4 }}
/>
</div>,
);
expect(screen.getByRole('progressbar', { name: '处理模型资源' })).toHaveAttribute(
'aria-valuenow',
'40',
);
});
it('工具栏更多菜单提供窄桌面动作', () => {
const settings = vi.fn();
render(
<ToolbarOverflowMenu
fullscreen={false}
onCommands={() => {}}
onLayout={() => {}}
onSettings={settings}
onFullscreen={() => {}}
onHelp={() => {}}
onTheme={() => {}}
/>,
);
fireEvent.click(screen.getByRole('button', { name: '更多工作台操作' }));
fireEvent.click(screen.getByRole('menuitem', { name: '工作台设置' }));
expect(settings).toHaveBeenCalled();
});
import {fireEvent,render,screen,within} from '@testing-library/react';
import {DiagnosticsDrawer} from './DiagnosticsDrawer';
import {ErrorRecoveryPanel} from './ErrorRecoveryPanel';
import {ToolbarOverflowMenu} from './ToolbarOverflowMenu';
import {WorkspaceOverlays} from './WorkspaceOverlays';
const event={id:1,title:'编译失败',detail:'bad xml',tone:'danger' as const,at:0};
describe('第五批工作台组件',()=>{
it('事件日志支持分类和清空',()=>{const clear=vi.fn();render(<DiagnosticsDrawer open items={[event]} onClose={()=>{}} onClear={clear}/>);expect(within(screen.getByRole('tabpanel',{name:/全部/})).getByText('bad xml')).toBeVisible();expect(screen.getAllByText('bad xml')).toHaveLength(1);fireEvent.click(screen.getByRole('button',{name:'清空事件'}));expect(clear).toHaveBeenCalled();});
it('错误恢复面板透传重试与工程树动作',()=>{const retry=vi.fn(),project=vi.fn();render(<ErrorRecoveryPanel value={{category:'模型编译',summary:'失败',detail:'bad',path:'a.xml',at:1}} onClose={()=>{}} onRetry={retry} onOpenProject={project}/>);fireEvent.click(screen.getByRole('button',{name:'重试当前入口'}));fireEvent.click(screen.getByRole('button',{name:'返回工程树'}));expect(retry).toHaveBeenCalled();expect(project).toHaveBeenCalled();});
it('导入叠层显示阶段进度',()=>{render(<div className="relative"><WorkspaceOverlays loading hasSnapshot={false} progress={{label:'处理模型资源',value:.4}}/></div>);expect(screen.getByRole('progressbar',{name:'处理模型资源'})).toHaveAttribute('aria-valuenow','40');});
it('工具栏更多菜单提供窄桌面动作',()=>{const settings=vi.fn();render(<ToolbarOverflowMenu fullscreen={false} onCommands={()=>{}} onLayout={()=>{}} onSettings={settings} onFullscreen={()=>{}} onHelp={()=>{}} onTheme={()=>{}}/>);fireEvent.click(screen.getByRole('button',{name:'更多工作台操作'}));fireEvent.click(screen.getByRole('menuitem',{name:'工作台设置'}));expect(settings).toHaveBeenCalled();});
});
@@ -1,103 +1,13 @@
import { act, fireEvent, render, screen } from '@testing-library/react';
import {
NotificationCenter,
ToastViewport,
type WorkbenchNotification,
} from './NotificationCenter';
import { ProjectBreadcrumb } from './ProjectBreadcrumb';
import { SettingsDialog } from './SettingsDialog';
import { LayoutSettingsDialog } from './LayoutSettingsDialog';
const item: WorkbenchNotification = {
id: 1,
title: '模型加载完成',
detail: '完成',
tone: 'success',
at: 0,
};
describe('第四批工作台组件', () => {
it('通知中心展示、移除并清空消息', () => {
const dismiss = vi.fn(),
clear = vi.fn();
render(<NotificationCenter items={[item]} onDismiss={dismiss} onClear={clear} />);
fireEvent.click(screen.getByRole('button', { name: '通知中心' }));
expect(screen.getByRole('dialog', { name: '通知中心' })).toHaveTextContent('模型加载完成');
fireEvent.click(screen.getByRole('button', { name: '移除通知:模型加载完成' }));
expect(dismiss).toHaveBeenCalledWith(1);
fireEvent.click(screen.getByText('清空'));
expect(clear).toHaveBeenCalled();
});
it('Toast 自动关闭', () => {
vi.useFakeTimers();
const close = vi.fn();
render(<ToastViewport item={item} onDismiss={close} />);
act(() => vi.advanceTimersByTime(4000));
expect(close).toHaveBeenCalledWith(1);
vi.useRealTimers();
});
it('工程面包屑可切换多入口', () => {
const select = vi.fn();
render(
<ProjectBreadcrumb
projectName="robot"
selectedEntry="models/a.xml"
entries={[
{ path: 'models/a.xml', label: 'A', format: 'mjcf' },
{ path: 'models/b.xml', label: 'B', format: 'mjcf' },
]}
onSelect={select}
/>,
);
fireEvent.click(screen.getByRole('button', { name: '切换模型入口' }));
fireEvent.click(screen.getByRole('option', { name: /B/ }));
expect(select).toHaveBeenCalledWith('models/b.xml');
});
it('模型加载期间禁用入口切换', () => {
render(
<ProjectBreadcrumb
projectName="robot"
loading
entries={[
{ path: 'a.xml', label: 'A', format: 'mjcf' },
{ path: 'b.xml', label: 'B', format: 'mjcf' },
]}
selectedEntry="a.xml"
onSelect={() => {}}
/>,
);
expect(screen.getByRole('button', { name: '切换模型入口' })).toBeDisabled();
});
it('设置和布局弹窗透传现有设置动作', () => {
const theme = vi.fn(),
preset = vi.fn();
render(
<>
<SettingsDialog
open
onClose={() => {}}
theme="dark"
angleUnit="rad"
showCollision={false}
jointAdvanced={false}
forceScale={50}
onTheme={theme}
onAngleUnit={() => {}}
onShowCollision={() => {}}
onJointAdvanced={() => {}}
onForceScale={() => {}}
/>
<LayoutSettingsDialog
open={false}
onClose={() => {}}
leftOpen
rightOpen
onLeftOpen={() => {}}
onRightOpen={() => {}}
onPreset={preset}
onReset={() => {}}
/>
</>,
);
fireEvent.change(screen.getByLabelText('设置主题'), { target: { value: 'light' } });
expect(theme).toHaveBeenCalledWith('light');
});
import {act,fireEvent,render,screen} from '@testing-library/react';
import {NotificationCenter,ToastViewport,type WorkbenchNotification} from './NotificationCenter';
import {ProjectBreadcrumb} from './ProjectBreadcrumb';
import {SettingsDialog} from './SettingsDialog';
import {LayoutSettingsDialog} from './LayoutSettingsDialog';
const item:WorkbenchNotification={id:1,title:'模型加载完成',detail:'完成',tone:'success',at:0};
describe('第四批工作台组件',()=>{
it('通知中心展示、移除并清空消息',()=>{const dismiss=vi.fn(),clear=vi.fn();render(<NotificationCenter items={[item]} onDismiss={dismiss} onClear={clear}/>);fireEvent.click(screen.getByRole('button',{name:'通知中心'}));expect(screen.getByRole('dialog',{name:'通知中心'})).toHaveTextContent('模型加载完成');fireEvent.click(screen.getByRole('button',{name:'移除通知:模型加载完成'}));expect(dismiss).toHaveBeenCalledWith(1);fireEvent.click(screen.getByText('清空'));expect(clear).toHaveBeenCalled();});
it('Toast 自动关闭',()=>{vi.useFakeTimers();const close=vi.fn();render(<ToastViewport item={item} onDismiss={close}/>);act(()=>vi.advanceTimersByTime(4000));expect(close).toHaveBeenCalledWith(1);vi.useRealTimers();});
it('工程面包屑可切换多入口',()=>{const select=vi.fn();render(<ProjectBreadcrumb projectName="robot" selectedEntry="models/a.xml" entries={[{path:'models/a.xml',label:'A',format:'mjcf'},{path:'models/b.xml',label:'B',format:'mjcf'}]} onSelect={select}/>);fireEvent.click(screen.getByRole('button',{name:'切换模型入口'}));fireEvent.click(screen.getByRole('option',{name:/B/}));expect(select).toHaveBeenCalledWith('models/b.xml');});
it('模型加载期间禁用入口切换',()=>{render(<ProjectBreadcrumb projectName="robot" loading entries={[{path:'a.xml',label:'A',format:'mjcf'},{path:'b.xml',label:'B',format:'mjcf'}]} selectedEntry="a.xml" onSelect={()=>{}}/>);expect(screen.getByRole('button',{name:'切换模型入口'})).toBeDisabled();});
it('设置和布局弹窗透传现有设置动作',()=>{const theme=vi.fn(),preset=vi.fn();render(<><SettingsDialog open onClose={()=>{}} theme="dark" angleUnit="rad" showCollision={false} jointAdvanced={false} forceScale={50} onTheme={theme} onAngleUnit={()=>{}} onShowCollision={()=>{}} onJointAdvanced={()=>{}} onForceScale={()=>{}}/><LayoutSettingsDialog open={false} onClose={()=>{}} leftOpen rightOpen onLeftOpen={()=>{}} onRightOpen={()=>{}} onPreset={preset} onReset={()=>{}}/></>);fireEvent.change(screen.getByLabelText('设置主题'),{target:{value:'light'}});expect(theme).toHaveBeenCalledWith('light');});
});
@@ -1,78 +1,7 @@
import { Columns3, Focus, PanelLeft, PanelRight, RotateCcw } from 'lucide-react';
import { Button, Dialog } from '../../components/ui';
export type LayoutPreset = 'default' | 'viewport' | 'project' | 'control';
const presets = [
{ value: 'default' as const, label: '默认布局', detail: '左右面板均衡显示', icon: Columns3 },
{ value: 'viewport' as const, label: '宽视口', detail: '隐藏两侧面板', icon: Focus },
{ value: 'project' as const, label: '工程浏览', detail: '加宽工程面板', icon: PanelLeft },
{ value: 'control' as const, label: '控制调试', detail: '加宽控制面板', icon: PanelRight },
];
export function LayoutSettingsDialog({
open,
onClose,
leftOpen,
rightOpen,
onLeftOpen,
onRightOpen,
onPreset,
onReset,
}: {
open: boolean;
onClose: () => void;
leftOpen: boolean;
rightOpen: boolean;
onLeftOpen: (value: boolean) => void;
onRightOpen: (value: boolean) => void;
onPreset: (preset: LayoutPreset) => void;
onReset: () => void;
}) {
return (
<Dialog open={open} onClose={onClose} title="布局设置">
<div className="flex gap-2">
<Button
variant={leftOpen ? 'primary' : 'secondary'}
aria-pressed={leftOpen}
onClick={() => onLeftOpen(!leftOpen)}
icon={<PanelLeft className="h-3.5 w-3.5" />}
>
工程面板
</Button>
<Button
variant={rightOpen ? 'primary' : 'secondary'}
aria-pressed={rightOpen}
onClick={() => onRightOpen(!rightOpen)}
icon={<PanelRight className="h-3.5 w-3.5" />}
>
右侧面板
</Button>
</div>
<h3 className="mb-2 mt-4 text-xs font-semibold">布局预设</h3>
<div className="grid grid-cols-2 gap-2">
{presets.map((item) => (
<button
key={item.value}
onClick={() => onPreset(item.value)}
className="flex gap-2 rounded-lg border border-border bg-surface p-3 text-left hover:border-accent hover:bg-accent-soft focus-visible:ring-2 focus-visible:ring-accent/30"
>
<item.icon className="h-4 w-4 shrink-0 text-accent" />
<span>
<span className="block text-xs font-medium">{item.label}</span>
<span className="mt-0.5 block text-[10px] text-text-tertiary">{item.detail}</span>
</span>
</button>
))}
</div>
<Button
className="mt-4 w-full"
onClick={onReset}
icon={<RotateCcw className="h-3.5 w-3.5" />}
>
恢复默认布局与宽度
</Button>
</Dialog>
);
}
import {Columns3,Focus,PanelLeft,PanelRight,RotateCcw} from 'lucide-react';
import {Button,Dialog} from '../../components/ui';
export type LayoutPreset='default'|'viewport'|'project'|'control';
const presets=[{value:'default' as const,label:'默认布局',detail:'左右面板均衡显示',icon:Columns3},{value:'viewport' as const,label:'宽视口',detail:'隐藏两侧面板',icon:Focus},{value:'project' as const,label:'工程浏览',detail:'加宽工程面板',icon:PanelLeft},{value:'control' as const,label:'控制调试',detail:'加宽控制面板',icon:PanelRight}];
export function LayoutSettingsDialog({open,onClose,leftOpen,rightOpen,onLeftOpen,onRightOpen,onPreset,onReset}:{open:boolean;onClose:()=>void;leftOpen:boolean;rightOpen:boolean;onLeftOpen:(value:boolean)=>void;onRightOpen:(value:boolean)=>void;onPreset:(preset:LayoutPreset)=>void;onReset:()=>void}){return <Dialog open={open} onClose={onClose} title="布局设置"><div className="flex gap-2"><Button variant={leftOpen?'primary':'secondary'} aria-pressed={leftOpen} onClick={()=>onLeftOpen(!leftOpen)} icon={<PanelLeft className="h-3.5 w-3.5"/>}>工程面板</Button><Button variant={rightOpen?'primary':'secondary'} aria-pressed={rightOpen} onClick={()=>onRightOpen(!rightOpen)} icon={<PanelRight className="h-3.5 w-3.5"/>}>属性面板</Button></div><h3 className="mb-2 mt-4 text-xs font-semibold">布局预设</h3><div className="grid grid-cols-2 gap-2">{presets.map(item=><button key={item.value} onClick={()=>onPreset(item.value)} className="flex gap-2 rounded-lg border border-border bg-surface p-3 text-left hover:border-accent hover:bg-accent-soft focus-visible:ring-2 focus-visible:ring-accent/30"><item.icon className="h-4 w-4 shrink-0 text-accent"/><span><span className="block text-xs font-medium">{item.label}</span><span className="mt-0.5 block text-[10px] text-text-tertiary">{item.detail}</span></span></button>)}</div><Button className="mt-4 w-full" onClick={onReset} icon={<RotateCcw className="h-3.5 w-3.5"/>}>恢复默认布局与宽度</Button></Dialog>;}
// eslint-disable-next-line react-refresh/only-export-components
export function dispatchLayoutWidths(left: number, right: number) {
window.dispatchEvent(new CustomEvent('mujoco-layout-widths', { detail: { left, right } }));
}
export function dispatchLayoutWidths(left:number,right:number){window.dispatchEvent(new CustomEvent('mujoco-layout-widths',{detail:{left,right}}));}
@@ -0,0 +1,25 @@
import {fireEvent,render,screen,waitFor} from '@testing-library/react';
import {beforeEach,describe,expect,it,vi} from 'vitest';
import {LocalTrainingPanel} from './LocalTrainingPanel';
beforeEach(()=>{localStorage.clear();vi.unstubAllGlobals();});
describe('LocalTrainingPanel',()=>{
it('连接本地服务并从图形界面发起训练请求',async()=>{
const health={version:'0.1.0',ready:true,trainerRoot:'/opt/unitree_rl_mjlab',python:'/env/bin/python',tasks:['Unitree-Go2-Flat']};
const job={id:'a'.repeat(32),state:'queued',taskId:'Unitree-Go2-Flat',createdAt:'2025-01-01T00:00:00Z',iteration:0,maxIterations:2000,progress:0,message:'等待启动',logs:[],artifactReady:false};
const fetchMock=vi.fn()
.mockResolvedValueOnce(new Response(JSON.stringify(health),{status:200,headers:{'Content-Type':'application/json'}}))
.mockResolvedValueOnce(new Response(JSON.stringify(job),{status:202,headers:{'Content-Type':'application/json'}}));
vi.stubGlobal('fetch',fetchMock);
render(<LocalTrainingPanel onPolicyReady={vi.fn()}/>);
fireEvent.click(screen.getByRole('button',{name:'连接'}));
expect(await screen.findByText('/opt/unitree_rl_mjlab')).toBeInTheDocument();
fireEvent.change(screen.getByLabelText('并行环境'),{target:{value:'32'}});
fireEvent.click(screen.getByRole('button',{name:'发起本地训练'}));
await waitFor(()=>expect(fetchMock).toHaveBeenCalledTimes(2));
const request=fetchMock.mock.calls[1][1] as RequestInit;
expect(JSON.parse(String(request.body))).toMatchObject({taskId:'Unitree-Go2-Flat',numEnvs:32,device:'gpu',gpuIds:[0],wandbMode:'offline'});
expect(await screen.findByText('排队中')).toBeInTheDocument();
});
});
@@ -0,0 +1,77 @@
import {useEffect,useState,type ReactNode} from 'react';
import {Download,Link,Play,Server,Square} from 'lucide-react';
import {Badge,Button,ProgressBar,PropertyRow,Select} from '../../components/ui';
import {LocalTrainingClient} from '../../training/LocalTrainingClient';
import type {TrainingDevice,TrainingJob,TrainingServerInfo,WandbMode} from '../../training/types';
const ENDPOINT_KEY='mujoco-local-training-endpoint',JOB_KEY='mujoco-local-training-job';
const DEFAULT_ENDPOINT='http://127.0.0.1:8765';
const ACTIVE_STATES=new Set(['queued','running']);
function stored(key:string,fallback=''):string{try{return localStorage.getItem(key)??fallback;}catch{return fallback;}}
function errorText(error:unknown):string{return error instanceof Error?error.message:String(error);}
function stateLabel(state:TrainingJob['state']):string{return {queued:'排队中',running:'训练中',succeeded:'已完成',failed:'失败',cancelled:'已取消'}[state];}
export function LocalTrainingPanel({onPolicyReady}:{onPolicyReady(file:File):void}){
const [endpoint,setEndpoint]=useState(()=>stored(ENDPOINT_KEY,DEFAULT_ENDPOINT));
const [server,setServer]=useState<TrainingServerInfo>();
const [job,setJob]=useState<TrainingJob>();
const [busy,setBusy]=useState(false),[error,setError]=useState<string>();
const [taskId,setTaskId]=useState('Unitree-Go2-Flat'),[numEnvs,setNumEnvs]=useState(4096),[maxIterations,setMaxIterations]=useState(2000),[seed,setSeed]=useState(42),[runName,setRunName]=useState('web'),[device,setDevice]=useState<TrainingDevice>('gpu'),[gpuIds,setGpuIds]=useState('0'),[wandbMode,setWandbMode]=useState<WandbMode>('offline');
const connect=async()=>{
setBusy(true);setError(undefined);
try{
const client=new LocalTrainingClient(endpoint),info=await client.health();
setServer(info);try{localStorage.setItem(ENDPOINT_KEY,client.endpoint);}catch{/* 当前会话仍可连接 */}
if(info.tasks.length&&!info.tasks.includes(taskId))setTaskId(info.tasks[0]);
const remembered=info.activeJobId??stored(JOB_KEY);
if(remembered){try{setJob(await client.job(remembered));}catch{try{localStorage.removeItem(JOB_KEY);}catch{/* ignore */}}}
if(!info.ready)setError(info.error??'训练服务尚未就绪');
}catch(value){setServer(undefined);setError(errorText(value));}
finally{setBusy(false);}
};
const jobId=job?.id,jobState=job?.state;
useEffect(()=>{
if(!jobId||!jobState||!ACTIVE_STATES.has(jobState))return;
let disposed=false;
const refresh=async()=>{try{const next=await new LocalTrainingClient(endpoint).job(jobId);if(!disposed)setJob(next);}catch(value){if(!disposed)setError(errorText(value));}};
const timer=window.setInterval(()=>void refresh(),1500);return()=>{disposed=true;window.clearInterval(timer);};
},[endpoint,jobId,jobState]);
const start=async()=>{
setBusy(true);setError(undefined);
try{
const ids=device==='gpu'?gpuIds.split(/[\s,]+/).filter(Boolean).map(Number):[];
if(ids.some(id=>!Number.isInteger(id)||id<0))throw new Error('GPU 编号必须是非负整数');
const next=await new LocalTrainingClient(endpoint).start({taskId,numEnvs,maxIterations,seed,runName,device,gpuIds:ids,wandbMode});
setJob(next);try{localStorage.setItem(JOB_KEY,next.id);}catch{/* ignore */}
}catch(value){setError(errorText(value));}finally{setBusy(false);}
};
const cancel=async()=>{if(!job)return;setBusy(true);setError(undefined);try{setJob(await new LocalTrainingClient(endpoint).cancel(job.id));}catch(value){setError(errorText(value));}finally{setBusy(false);}};
const importResult=async()=>{if(!job)return;setBusy(true);setError(undefined);try{onPolicyReady(await new LocalTrainingClient(endpoint).downloadPolicy(job.id));}catch(value){setError(errorText(value));}finally{setBusy(false);}};
const active=Boolean(job&&ACTIVE_STATES.has(job.state));
return <div>
<label className="block text-xs text-text-secondary"><span className="mb-1 block">本地训练服务</span><div className="flex gap-2"><input aria-label="本地训练服务地址" className="field h-7 min-w-0 flex-1 px-2 text-xs text-text-primary" value={endpoint} disabled={active} onChange={event=>setEndpoint(event.target.value)}/><Button icon={<Link className="h-3.5 w-3.5"/>} disabled={busy||active} onClick={()=>void connect()}>连接</Button></div></label>
<div className="mt-2 flex items-center justify-between rounded-md border border-border bg-surface px-2 py-1.5 text-[10px] text-text-tertiary"><span className="flex min-w-0 items-center gap-1.5 truncate"><Server className="h-3.5 w-3.5"/>{server?.trainerRoot??'请先启动本地训练服务'}</span><Badge tone={server?.ready?'success':'warning'}>{server?.ready?'可用':'离线'}</Badge></div>
{server?.ready&&!job&&<div className="mt-3 space-y-2">
<Field label="训练任务"><Select aria-label="训练任务" className="w-full" value={taskId} onChange={event=>setTaskId(event.target.value)}>{server.tasks.map(task=><option key={task} value={task}>{task}</option>)}</Select></Field>
<div className="grid grid-cols-2 gap-2"><NumberField label="并行环境" value={numEnvs} min={1} max={16384} onChange={setNumEnvs}/><NumberField label="训练迭代" value={maxIterations} min={1} max={1000000} onChange={setMaxIterations}/><NumberField label="随机种子" value={seed} min={0} max={2147483647} onChange={setSeed}/><Field label="运行名称"><input aria-label="运行名称" className="field h-7 w-full px-2 text-xs text-text-primary" value={runName} onChange={event=>setRunName(event.target.value)}/></Field></div>
<div className="grid grid-cols-2 gap-2"><Field label="计算设备"><Select aria-label="计算设备" className="w-full" value={device} onChange={event=>setDevice(event.target.value as TrainingDevice)}><option value="gpu">GPU</option><option value="cpu">CPU</option></Select></Field><Field label="GPU 编号"><input aria-label="GPU 编号" className="field h-7 w-full px-2 text-xs text-text-primary disabled:opacity-40" value={gpuIds} disabled={device==='cpu'} onChange={event=>setGpuIds(event.target.value)}/></Field></div>
<Field label="实验记录"><Select aria-label="W&B 模式" className="w-full" value={wandbMode} onChange={event=>setWandbMode(event.target.value as WandbMode)}><option value="offline">本地离线(默认,无需登录)</option><option value="disabled">完全禁用 W&amp;B</option><option value="online">在线 W&amp;B(需要 API Key)</option></Select></Field>
<Button variant="primary" className="w-full" icon={<Play className="h-3.5 w-3.5"/>} disabled={busy} onClick={()=>void start()}>发起本地训练</Button>
<p className="text-[10px] leading-4 text-text-tertiary">训练使用本地 mjlab 任务资产,不会把浏览器中的模型上传到网络。服务一次只运行一个训练任务。</p>
</div>}
{job&&<div className="mt-3 rounded-lg border border-border bg-surface p-2.5">
<div className="mb-2 flex items-center justify-between gap-2"><span className="truncate text-xs font-medium text-text-primary" title={job.id}>{job.taskId}</span><Badge tone={job.state==='succeeded'?'success':job.state==='failed'||job.state==='cancelled'?'warning':'accent'}>{stateLabel(job.state)}</Badge></div>
<ProgressBar value={job.progress} label="训练进度"/><div className="mt-2"><PropertyRow label="迭代" value={`${job.iteration} / ${job.maxIterations}`}/><PropertyRow label="状态" value={job.message}/></div>
{job.logs.length>0&&<details className="mt-2"><summary className="cursor-pointer text-[10px] text-text-secondary">最近日志</summary><pre className="mt-1 max-h-36 overflow-auto whitespace-pre-wrap break-all rounded bg-app p-2 text-[9px] leading-4 text-text-tertiary">{job.logs.slice(-40).join('\n')}</pre></details>}
<div className="mt-3 grid grid-cols-2 gap-2">{active?<Button variant="danger" className="col-span-2" icon={<Square className="h-3.5 w-3.5"/>} disabled={busy} onClick={()=>void cancel()}>停止训练</Button>:<><Button disabled={busy||!job.artifactReady} icon={<Download className="h-3.5 w-3.5"/>} onClick={()=>void importResult()}>导入策略</Button><Button onClick={()=>{setJob(undefined);try{localStorage.removeItem(JOB_KEY);}catch{/* ignore */}}}>新建任务</Button></>}</div>
</div>}
{error&&<p role="alert" className="mt-2 break-words rounded bg-danger/10 p-2 text-[10px] leading-4 text-danger">{error}</p>}
</div>;
}
function Field({label,children}:{label:string;children:ReactNode}){return <label className="block text-[10px] text-text-tertiary"><span className="mb-1 block">{label}</span>{children}</label>;}
function NumberField({label,value,min,max,onChange}:{label:string;value:number;min:number;max:number;onChange(value:number):void}){return <Field label={label}><input aria-label={label} type="number" className="field h-7 w-full px-2 text-xs text-text-primary" value={value} min={min} max={max} onChange={event=>onChange(Number(event.target.value))}/></Field>;}
@@ -1,99 +0,0 @@
import { fireEvent, render, screen } from '@testing-library/react';
import { MapDraftStatusOverlay, MapViewportToolbar } from './MapViewportTools';
const noop = () => {};
describe('MapViewportToolbar', () => {
it('统一联动变换、吸附与贴地检测命令', () => {
const onModeChange = vi.fn();
const onSnappingChange = vi.fn();
const onPlacementModeChange = vi.fn();
const onAlignToSurface = vi.fn();
render(
<MapViewportToolbar
visible
interactionActive
mode="translate"
snapping
placementMode="auto_ground"
hasSelectedObject
allowScale
loading={false}
onActivate={noop}
onModeChange={onModeChange}
onSnappingChange={onSnappingChange}
onPlacementModeChange={onPlacementModeChange}
onAlignToSurface={onAlignToSurface}
/>,
);
fireEvent.click(screen.getByRole('button', { name: '旋转工具 E' }));
fireEvent.click(screen.getByRole('button', { name: '网格吸附' }));
fireEvent.change(screen.getByLabelText('贴地检测模式'), {
target: { value: 'gravity' },
});
fireEvent.click(screen.getByRole('button', { name: '立即贴合承载面' }));
expect(onModeChange).toHaveBeenCalledWith('rotate');
expect(onSnappingChange).toHaveBeenCalledWith(false);
expect(onPlacementModeChange).toHaveBeenCalledWith('gravity');
expect(onAlignToSurface).toHaveBeenCalledOnce();
});
it('参数地形或未选中对象时不提供缩放与贴合动作', () => {
render(
<MapViewportToolbar
visible
interactionActive
mode="translate"
snapping={false}
placementMode="auto_ground"
hasSelectedObject={false}
allowScale={false}
loading={false}
onActivate={noop}
onModeChange={noop}
onSnappingChange={noop}
onPlacementModeChange={noop}
onAlignToSurface={noop}
/>,
);
expect(screen.queryByRole('button', { name: '缩放工具 R' })).not.toBeInTheDocument();
expect(screen.getByRole('button', { name: '立即贴合承载面' })).toBeDisabled();
});
});
describe('MapDraftStatusOverlay', () => {
it('常驻显示未保存数量并提供提交与丢弃入口', () => {
const onCommit = vi.fn();
const onDiscard = vi.fn();
render(
<MapDraftStatusOverlay
visible
changeCount={3}
loading={false}
onCommit={onCommit}
onDiscard={onDiscard}
/>,
);
expect(screen.getByText('3 项未保存改动')).toBeVisible();
fireEvent.click(screen.getByRole('button', { name: '提交地图草稿' }));
fireEvent.click(screen.getByRole('button', { name: '丢弃地图草稿' }));
expect(onCommit).toHaveBeenCalledOnce();
expect(onDiscard).toHaveBeenCalledOnce();
});
it('无草稿时显示已同步并禁用动作', () => {
render(
<MapDraftStatusOverlay
visible
changeCount={0}
loading={false}
onCommit={noop}
onDiscard={noop}
/>,
);
expect(screen.getByText('地图草稿已同步')).toBeVisible();
expect(screen.getByRole('button', { name: '提交地图草稿' })).toBeDisabled();
});
});
@@ -1,202 +0,0 @@
import {
ArrowDownToLine,
Check,
Grid3X3,
LockKeyhole,
MapPinned,
Move3d,
Rotate3d,
Scaling,
Trash2,
} from 'lucide-react';
import {
Button,
IconButton,
Select,
ToolbarToggleGroup,
type ToolbarItem,
} from '../../components/ui';
import {
MAP_OBJECT_PLACEMENT_LABELS,
type MapEditorTransformMode,
type MapObjectPlacementMode,
} from '../../map/editor/types';
const TRANSFORM_TOOLS: ToolbarItem<MapEditorTransformMode>[] = [
{ value: 'translate', label: '移动工具 W', icon: Move3d },
{ value: 'rotate', label: '旋转工具 E', icon: Rotate3d },
{ value: 'scale', label: '缩放工具 R', icon: Scaling },
];
export function MapViewportToolbar({
visible,
interactionActive,
mode,
snapping,
placementMode,
hasSelectedObject,
allowScale,
loading,
onActivate,
onModeChange,
onSnappingChange,
onPlacementModeChange,
onAlignToSurface,
}: {
visible: boolean;
interactionActive: boolean;
mode: MapEditorTransformMode;
snapping: boolean;
placementMode: MapObjectPlacementMode;
hasSelectedObject: boolean;
allowScale: boolean;
loading: boolean;
onActivate: () => void;
onModeChange: (mode: MapEditorTransformMode) => void;
onSnappingChange: (value: boolean) => void;
onPlacementModeChange: (mode: MapObjectPlacementMode) => void;
onAlignToSurface: () => void;
}) {
if (!visible) return null;
return (
<div
aria-label="地图视口工具"
className="engineering-glass map-tool-enter absolute left-1/2 top-3 z-20 flex max-w-[calc(100%-24px)] -translate-x-1/2 items-center gap-1.5 rounded-xl border p-1.5"
>
<button
type="button"
aria-label="激活地图编辑"
aria-pressed={interactionActive}
disabled={loading}
onClick={onActivate}
className={`hidden h-8 items-center gap-1.5 rounded-lg px-2 text-[10px] font-semibold transition-colors sm:flex ${interactionActive ? 'bg-accent-soft text-accent' : 'bg-surface/80 text-text-secondary hover:bg-element-hover'}`}
>
<MapPinned className="h-3.5 w-3.5" aria-hidden="true" />
地图编辑
</button>
<ToolbarToggleGroup
items={TRANSFORM_TOOLS.filter((item) => allowScale || item.value !== 'scale')}
value={mode}
onChange={onModeChange}
label="地图变换模式"
/>
<IconButton
active={snapping}
tooltip={snapping ? '关闭网格吸附(当前 0.1 m / 5°)' : '开启网格吸附'}
aria-label="网格吸附"
aria-pressed={snapping}
disabled={loading}
onClick={() => onSnappingChange(!snapping)}
>
<Grid3X3 className="h-3.5 w-3.5" />
</IconButton>
<span aria-hidden="true" className="mx-0.5 h-5 w-px bg-border" />
<label className="flex items-center gap-1 text-[10px] text-text-tertiary">
<span className="hidden lg:inline">贴地</span>
<Select
aria-label="贴地检测模式"
className="w-[112px]"
value={placementMode}
disabled={loading}
onChange={(event) => onPlacementModeChange(event.target.value as MapObjectPlacementMode)}
>
{(Object.keys(MAP_OBJECT_PLACEMENT_LABELS) as MapObjectPlacementMode[]).map((value) => (
<option key={value} value={value}>
{value === 'auto_ground' ? '地面 z=0' : value === 'gravity' ? '表面检测' : '锁定位姿'}
</option>
))}
</Select>
</label>
<IconButton
tooltip={hasSelectedObject ? '重新检测并贴合承载面' : '选中地图物体后可贴合表面'}
aria-label="立即贴合承载面"
disabled={loading || !hasSelectedObject || placementMode === 'locked'}
onClick={onAlignToSurface}
>
{placementMode === 'locked' ? (
<LockKeyhole className="h-3.5 w-3.5" />
) : (
<ArrowDownToLine className="h-3.5 w-3.5" />
)}
</IconButton>
</div>
);
}
export function MapDraftStatusOverlay({
visible,
changeCount,
loading,
onCommit,
onDiscard,
}: {
visible: boolean;
changeCount: number;
loading: boolean;
onCommit: () => void;
onDiscard: () => void;
}) {
if (!visible) return null;
const dirty = changeCount > 0;
return (
<div
aria-label="地图草稿状态"
role="status"
className="engineering-glass absolute bottom-3 left-1/2 z-20 flex max-w-[calc(100%-24px)] -translate-x-1/2 items-center gap-2 rounded-xl border px-2 py-1.5"
>
<span
aria-hidden="true"
className={`h-2 w-2 shrink-0 rounded-full ${dirty ? 'draft-dirty-dot bg-warning' : 'bg-success'}`}
/>
<span className="min-w-0 whitespace-nowrap text-[10px] font-medium text-text-secondary sm:text-[11px]">
{dirty ? `${changeCount} 项未保存改动` : '地图草稿已同步'}
</span>
<Button
variant="ghost"
aria-label="丢弃地图草稿"
disabled={loading || !dirty}
icon={<Trash2 className="h-3 w-3" />}
onClick={onDiscard}
>
<span className="hidden sm:inline">丢弃</span>
</Button>
<Button
variant="primary"
aria-label="提交地图草稿"
disabled={loading || !dirty}
icon={<Check className="h-3 w-3" />}
onClick={onCommit}
>
<span className="hidden sm:inline">提交</span>
</Button>
</div>
);
}
export interface MapAssetDropTarget {
left: number;
top: number;
position: readonly [number, number, number] | null;
}
export function MapAssetDropIndicator({ target }: { target?: MapAssetDropTarget }) {
if (!target) return null;
const valid = Boolean(target.position);
return (
<div className="pointer-events-none absolute inset-0 z-30" aria-label="地图资产放置预览">
<div
className={`absolute -translate-x-1/2 -translate-y-1/2 rounded-full border-2 p-1 shadow-2xl ${valid ? 'border-accent bg-accent/25 text-accent' : 'border-danger bg-danger/20 text-danger'}`}
style={{ left: target.left, top: target.top }}
>
<span className="grid h-8 w-8 place-items-center rounded-full border border-current bg-panel/85 backdrop-blur">
<MapPinned className="h-4 w-4" />
</span>
<span className="absolute left-1/2 top-full mt-2 -translate-x-1/2 whitespace-nowrap rounded-lg border border-border-strong bg-panel/90 px-2 py-1 text-[10px] font-medium shadow-xl backdrop-blur">
{valid
? `释放放置 · ${target.position![0].toFixed(1)}, ${target.position![1].toFixed(1)}`
: '请拖到 3D 地面'}
</span>
</div>
</div>
);
}
@@ -1,162 +0,0 @@
import type { ComponentProps } from 'react';
import { fireEvent, render, screen } from '@testing-library/react';
import { DEFAULT_PHYSICAL_MAP_CONFIG } from '../../map/types';
import type { SimulationSnapshot } from '../../simulation/SimulationSession';
import { ModelControlsSidebar } from './ModelControlsSidebar';
const snapshot = {
bodies: [
{ id: 0, name: 'world', parentId: 0 },
{ id: 1, name: 'base', parentId: 0 },
{ id: 2, name: 'arm', parentId: 1 },
],
joints: [
{
id: 7,
name: 'arm_joint',
bodyId: 2,
type: 3,
value: 0.25,
min: -1,
max: 1,
limitMin: -1,
limitMax: 1,
limited: true,
limitsIgnored: false,
editable: true,
axis: [0, 0, 1],
},
],
actuators: [],
model: { nbody: 3, njnt: 1, ngeom: 2, ncam: 0, nactuator: 0, nu: 0, nq: 1, nv: 1 },
} as unknown as SimulationSnapshot;
const noop = () => {};
const asyncTrue = async () => true;
function props(
patch: Partial<ComponentProps<typeof ModelControlsSidebar>> = {},
): ComponentProps<typeof ModelControlsSidebar> {
return {
snapshot,
selection: null,
viewerSelection: null,
loading: false,
urdfMode: 'mjcf',
baseMode: 'floating',
showCollision: false,
ignoreJointLimits: false,
jointAdvanced: false,
angleUnit: 'rad',
mapSelection: { kind: 'none' },
mapSceneDirty: false,
maps: [],
showVisualMap: true,
showMapCollision: false,
editorDocument: null,
workspaceTool: null,
onWorkspaceToolChange: noop,
onSelectJoint: noop,
onUrdfMode: noop,
onBaseMode: noop,
onShowCollision: noop,
onResetJoints: noop,
onToggleJointLimits: noop,
onToggleAdvanced: noop,
onToggleAngleUnit: noop,
onActuator: noop,
onActuatorParameters: noop,
onJoint: noop,
onApplyMap: noop,
onMapDraft: noop,
onEditorPreview: noop,
onEditorDraftChange: noop,
onEditorApply: asyncTrue,
onEditorExport: noop,
onEditorConvert: asyncTrue,
onEditorBindInteraction: noop,
onEditorSelect: noop,
onEditorSessionStateChange: noop,
onMapDisplay: noop,
...patch,
};
}
describe('ModelControlsSidebar', () => {
it('没有选择时显示场景摘要,并保留基础地图快速入口', () => {
render(<ModelControlsSidebar {...props()} />);
expect(screen.getByText('未选择对象')).toBeVisible();
expect(screen.getByText('模型摘要')).toBeVisible();
expect(screen.getByText('未选择地图实例')).toBeVisible();
});
it('按照统一选择自动路由 Body 与 Joint 检查器', () => {
const view = render(
<ModelControlsSidebar {...props({ selection: { kind: 'body', bodyId: 2 } })} />,
);
expect(screen.getByText('Robot / Body')).toBeVisible();
expect(screen.getByText('arm', { selector: '[title="arm"]' })).toBeVisible();
view.rerender(
<ModelControlsSidebar {...props({ selection: { kind: 'joint', jointId: 7, bodyId: 2 } })} />,
);
expect(screen.getByText('Robot / Joint')).toBeVisible();
expect(screen.getByText('Hinge · arm')).toBeVisible();
expect(screen.getByText('关联 Actuator')).toBeVisible();
});
it('控制台与数据录制在右侧面板内切换,不创建模态窗口', () => {
const onWorkspaceToolChange = vi.fn();
render(
<ModelControlsSidebar
{...props({
selection: { kind: 'body', bodyId: 2 },
workspaceTool: 'controls',
workspaceTools: (
<div role="tabpanel" aria-label="控制台">
内嵌控制工具
</div>
),
onWorkspaceToolChange,
})}
/>,
);
expect(screen.queryByRole('dialog')).not.toBeInTheDocument();
expect(screen.getByRole('tabpanel', { name: '控制台' })).toBeVisible();
expect(screen.getByText('内嵌控制工具')).toBeVisible();
const inspector = document.querySelector('[role="tabpanel"][aria-label="检查器"]');
expect(inspector).not.toBeNull();
expect(inspector).not.toBeVisible();
const controlsTab = screen.getByRole('tab', { name: '控制台' });
expect(controlsTab).toHaveAttribute('aria-selected', 'true');
fireEvent.keyDown(controlsTab, { key: 'ArrowRight' });
expect(onWorkspaceToolChange).toHaveBeenCalledWith('data');
fireEvent.click(screen.getByRole('tab', { name: '检查器' }));
expect(onWorkspaceToolChange).toHaveBeenCalledWith(null);
});
it('选择地图实例后自动展示地图检查器,并把全局工具移到右侧入口', () => {
const onWorkspaceToolChange = vi.fn();
render(
<ModelControlsSidebar
{...props({
selection: { kind: 'map', mapAssetId: 'terrain-a' },
activeMapAssetId: 'terrain-a',
activeMapAssetName: '随机粗糙地形',
mapSelection: {
kind: 'builtin',
config: { ...DEFAULT_PHYSICAL_MAP_CONFIG, preset: 'rough' },
},
onWorkspaceToolChange,
})}
/>,
);
expect(screen.getByText('Map / Instance')).toBeVisible();
expect(screen.getByText('随机粗糙地形', { selector: '[title="随机粗糙地形"]' })).toBeVisible();
fireEvent.click(screen.getByRole('tab', { name: '控制台' }));
fireEvent.click(screen.getByRole('tab', { name: '数据录制' }));
expect(onWorkspaceToolChange.mock.calls).toEqual([['controls'], ['data']]);
});
});

Some files were not shown because too many files have changed in this diff Show More