Initial commit
This commit is contained in:
+42
@@ -0,0 +1,42 @@
|
||||
# Python
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
*.egg-info/
|
||||
.eggs/
|
||||
build/
|
||||
dist/
|
||||
.pytest_cache/
|
||||
.mypy_cache/
|
||||
.ruff_cache/
|
||||
.coverage
|
||||
htmlcov/
|
||||
|
||||
# Environments
|
||||
.conda/
|
||||
.venv/
|
||||
venv/
|
||||
|
||||
# Editors / OS
|
||||
.idea/
|
||||
.vscode/
|
||||
.DS_Store
|
||||
*.swp
|
||||
*.swo
|
||||
|
||||
# Training outputs: publish only the final 50k artifacts of the two released runs.
|
||||
checkpoints/
|
||||
logs/*
|
||||
!logs/s10_terrain_v4_50k_full/
|
||||
!logs/s10_terrain_v7_50k_20261101_225518/
|
||||
logs/s10_terrain_v4_50k_full/*
|
||||
logs/s10_terrain_v7_50k_20261101_225518/*
|
||||
!logs/s10_terrain_v4_50k_full/model_50000.pt
|
||||
!logs/s10_terrain_v4_50k_full/model_50000.onnx
|
||||
!logs/s10_terrain_v4_50k_full/model_50000.json
|
||||
!logs/s10_terrain_v4_50k_full/train.log
|
||||
!logs/s10_terrain_v4_50k_full/metrics.jsonl
|
||||
!logs/s10_terrain_v7_50k_20261101_225518/model_50000.pt
|
||||
!logs/s10_terrain_v7_50k_20261101_225518/model_50000.onnx
|
||||
!logs/s10_terrain_v7_50k_20261101_225518/model_50000.json
|
||||
!logs/s10_terrain_v7_50k_20261101_225518/train.log
|
||||
!logs/s10_terrain_v7_50k_20261101_225518/metrics.jsonl
|
||||
@@ -0,0 +1,109 @@
|
||||
# S10 SuperMimic
|
||||
|
||||
基于 mjlab 和 MuJoCo Warp 的 S10 轮足机器人强化学习项目。通过仿真数据训练世界模型，再利用世界模型生成想象轨迹来优化速度控制策略，支持平地、楼梯、平台、斜坡和碎石地形。
|
||||
|
||||
演示使用v7_50k策略，在仓库中附带提交权重和模型
|
||||
|
||||
<video src="docs/WorldModel_demo1.mp4" controls width="75%"></video>
|
||||
|
||||
<video src="docs/Mujoco_demo2.mp4" controls width="75%"></video>
|
||||
|
||||
## Requirements
|
||||
|
||||
Python 3.10–3.13、PyTorch ≥ 2.7、mjlab 1.6.0、MuJoCo Warp；参考配置需要 16 GB 显存，如果使用更高显存训练，可以提高 `--imag-batch`。依赖安装：`pip install -r requirements.txt`。
|
||||
|
||||
## Tasks
|
||||
|
||||
- `s10-v4`：平地激进策略。
|
||||
- `s10-v7`：平地/台阶保守策略，可上 20 cm 台阶。
|
||||
- `s10-v4-upward`：倒地起身训练，开发中。
|
||||
|
||||
## Train
|
||||
|
||||
```bash
|
||||
export PYTHONPATH="$PWD"
|
||||
python scripts/train.py --task s10-v4 --output logs/s10_v4
|
||||
python scripts/train.py --task s10-v7 --output logs/s10_v7
|
||||
python scripts/train.py --task s10-v4-upward --output logs/s10_v4_upward
|
||||
```
|
||||
|
||||
`--task` 加载任务默认配置，命令行参数可覆盖。训练输出 `model_<step>.pt` 和 `metrics.jsonl`；`logs/` 保留 v4、v7 的 50k 最终权重。
|
||||
|
||||
## Play and Export
|
||||
|
||||
```bash
|
||||
python scripts/export_onnx.py --checkpoint <pt> --output <onnx>
|
||||
```
|
||||
|
||||
ONNX 只包含确定性 actor。`./scripts/play_gui.sh --checkpoint <pt>` 提供 GUI 交互演示，可可视化世界模型的未来预测，但不能替代 sim2sim；sim2sim 建议使用导出的 ONNX。
|
||||
|
||||
## World model
|
||||
|
||||
世界模型是带残差预测的前馈多层感知机（MLP），默认使用三层 512 维隐藏层和 SiLU 激活函数。它根据当前状态和动作预测下一时刻的状态变化，以及机身平面位移和偏航角变化。输入和输出使用运行统计量进行归一化。
|
||||
|
||||
| 输入 | 维度 | 内容 |
|
||||
| --- | ---: | --- |
|
||||
| 机器人状态 | 38 | 高度 1、投影重力 3、机身线速度 3、角速度 3、腿关节位置 12、腿关节速度 12、轮速 4 |
|
||||
| 当前动作 | 16 | 12 维腿部位置控制和 4 维轮速控制 |
|
||||
| 上一动作 | 16 | 上一控制周期的动作 |
|
||||
| 物理参数 | 7 | 四个轮子的摩擦系数和机身质心偏移 |
|
||||
| 地形扫描 | 273 | 复杂地形模型使用的 21 × 13 高度采样网格 |
|
||||
|
||||
平地世界模型输入为 **77 维**；复杂地形世界模型加入高度扫描后为 **350 维**。两者的动力学输出均为 **41 维**：38 维状态变化和 3 维平面位姿变化。复杂地形模型（revision ≥ 4）带有接触概率与接触力预测头，地形 v7 还带有非轮部位碰撞预测头。世界模型仍使用真实轮速作为训练时的特权状态，这与策略观测不同。
|
||||
|
||||
训练时从仿真回放中采样连续片段，世界模型按时间自回归预测并拟合实际轨迹；默认训练窗口为 8 步。随后使用世界模型展开想象轨迹，并结合奖励和 critic 更新策略。
|
||||
|
||||
## Actor
|
||||
|
||||
Actor 根据机器人自身观测和速度指令输出控制动作。当前 S10 Actor（v4 / v4-upward / v7）每帧接收 **53 维**观测：
|
||||
|
||||
| 观测 | 每帧维度 |
|
||||
| --- | ---: |
|
||||
| 机身角速度 | 3 |
|
||||
| 投影重力 | 3 |
|
||||
| 腿关节位置 | 12 |
|
||||
| 腿关节速度 | 12 |
|
||||
| 轮速 | 4 |
|
||||
| 上一动作 | 16 |
|
||||
| 速度指令（前后、横移、转向） | 3 |
|
||||
|
||||
默认保留 **5 帧**历史，因此 Actor 输入为 **265 维**；各观测项分别按时间从旧到新排列。控制周期为 0.02 秒，5 帧覆盖 5 个控制时刻。
|
||||
|
||||
Actor 是默认隐藏层为 **256 → 128 → 64** 的 MLP，使用输入归一化和 ELU 激活，最后通过 tanh 输出 **16 维动作**：12 维腿部位置控制和 4 维轮速控制。Actor 不接收地形扫描，实现盲走控制；演示时使用 Actor 生成动作。
|
||||
|
||||
## Critic
|
||||
|
||||
Critic 在训练中估计当前状态的未来回报，输出一个标量价值。它使用训练时可获得的特权信息，输入不需要堆叠 Actor 的历史帧：
|
||||
|
||||
| 输入 | 维度 | 内容 |
|
||||
| --- | ---: | --- |
|
||||
| 机器人状态 | 38 | 包含机身高度、线速度、角速度、关节状态和轮速 |
|
||||
| 速度指令 | 3 | 前后、横移和转向指令 |
|
||||
| 物理参数 | 7 | 四个轮子摩擦系数和机身质心偏移 |
|
||||
| 上一动作 | 16 | 上一控制周期的动作 |
|
||||
| 地形扫描 | 273 | 复杂地形的 21 × 13 高度采样网格 |
|
||||
|
||||
平地 Critic 输入为 **64 维**，复杂地形 Critic 输入为 **337 维**。Critic 默认同样使用 **256 → 128 → 64** 的 MLP、输入归一化和 ELU 激活，输出 **1 维**价值估计。Critic 和地形扫描只用于训练；运行策略时不需要 Critic。
|
||||
|
||||
训练时，世界模型生成想象轨迹，Actor 通过轨迹中的奖励和目标 Critic 的末端价值优化动作；Critic 使用想象轨迹计算的 TD(λ) 回报作为目标进行拟合。Actor 的 5 帧历史与世界模型默认的 8 步训练窗口是不同的参数，分别由 `--history` 和 `--wm-window` 设置。
|
||||
|
||||
## Structure
|
||||
|
||||
```text
|
||||
s10_supermimic/
|
||||
├── core/ # 世界模型、策略、想象训练、经验回放与演示界面
|
||||
├── env/ # 平地及复杂地形环境
|
||||
├── tasks/ # 任务注册与默认参数
|
||||
└── robot/ # S10 机器人配置与模型资产
|
||||
scripts/ # 训练、导出与演示入口
|
||||
docs/ # 演示素材
|
||||
tests/ # 模型、任务与奖励检查
|
||||
logs/ # 保留的最终权重与训练日志
|
||||
```
|
||||
|
||||
## References
|
||||
|
||||
- 仿真基于：https://github.com/mujocolab/mjlab
|
||||
- 算法基于：
|
||||
- https://github.com/zita-ch/supermimic_vibed
|
||||
- https://github.com/uwvwko-zzz/go2w-supermimic
|
||||
Binary file not shown.
Binary file not shown.
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,131 @@
|
||||
{
|
||||
"format": "s10-terrain-policy-v1",
|
||||
"step": 50000,
|
||||
"history": 5,
|
||||
"frame_dim": 53,
|
||||
"observation_dim": 265,
|
||||
"action_dim": 16,
|
||||
"dt": 0.02,
|
||||
"actor_hidden": [
|
||||
256,
|
||||
128,
|
||||
64
|
||||
],
|
||||
"actor_params": 110288,
|
||||
"frame_terms": [
|
||||
{
|
||||
"name": "base_ang_vel",
|
||||
"dim": 3,
|
||||
"scale": 1.0
|
||||
},
|
||||
{
|
||||
"name": "projected_gravity",
|
||||
"dim": 3,
|
||||
"scale": 1.0
|
||||
},
|
||||
{
|
||||
"name": "joint_pos_rel",
|
||||
"dim": 12,
|
||||
"scale": 1.0
|
||||
},
|
||||
{
|
||||
"name": "joint_vel_rel",
|
||||
"dim": 12,
|
||||
"scale": 0.05
|
||||
},
|
||||
{
|
||||
"name": "wheel_vel_rel",
|
||||
"dim": 4,
|
||||
"scale": 0.05
|
||||
},
|
||||
{
|
||||
"name": "last_action",
|
||||
"dim": 16,
|
||||
"scale": 1.0
|
||||
},
|
||||
{
|
||||
"name": "command",
|
||||
"dim": 3,
|
||||
"scale": 1.0
|
||||
}
|
||||
],
|
||||
"history_layout": "term-wise blocks, oldest to newest (mjlab flattening)",
|
||||
"env_action_order": [
|
||||
0,
|
||||
1,
|
||||
2,
|
||||
3,
|
||||
4,
|
||||
5,
|
||||
6,
|
||||
7,
|
||||
8,
|
||||
9,
|
||||
10,
|
||||
11,
|
||||
12,
|
||||
13,
|
||||
14,
|
||||
15
|
||||
],
|
||||
"leg_scale": [
|
||||
0.25,
|
||||
0.6,
|
||||
0.8,
|
||||
0.25,
|
||||
0.6,
|
||||
0.8,
|
||||
0.25,
|
||||
0.6,
|
||||
0.8,
|
||||
0.25,
|
||||
0.6,
|
||||
0.8
|
||||
],
|
||||
"wheel_scale": 35.0,
|
||||
"wheel_clip": 35.0,
|
||||
"default_joint_pos": [
|
||||
0.1,
|
||||
-0.35,
|
||||
0.65,
|
||||
-0.1,
|
||||
-0.35,
|
||||
0.65,
|
||||
0.1,
|
||||
0.35,
|
||||
-0.65,
|
||||
-0.1,
|
||||
0.35,
|
||||
-0.65,
|
||||
0.0,
|
||||
0.0,
|
||||
0.0,
|
||||
0.0
|
||||
],
|
||||
"leg_joint_names": [
|
||||
"fl_hipx_joint",
|
||||
"fl_hipy_joint",
|
||||
"fl_knee_joint",
|
||||
"fr_hipx_joint",
|
||||
"fr_hipy_joint",
|
||||
"fr_knee_joint",
|
||||
"hl_hipx_joint",
|
||||
"hl_hipy_joint",
|
||||
"hl_knee_joint",
|
||||
"hr_hipx_joint",
|
||||
"hr_hipy_joint",
|
||||
"hr_knee_joint"
|
||||
],
|
||||
"wheel_joint_names": [
|
||||
"fl_wheel_joint",
|
||||
"fr_wheel_joint",
|
||||
"hl_wheel_joint",
|
||||
"hr_wheel_joint"
|
||||
],
|
||||
"onnx": {
|
||||
"input": "obs",
|
||||
"output": "actions",
|
||||
"opset": 17,
|
||||
"max_abs_error_vs_torch": 7.487833499908447e-07
|
||||
}
|
||||
}
|
||||
Binary file not shown.
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,133 @@
|
||||
{
|
||||
"format": "s10-terrain-policy-v1",
|
||||
"step": 50000,
|
||||
"history": 5,
|
||||
"actor_schema": "s10-v7-53",
|
||||
"actor_wheel_vel": true,
|
||||
"frame_dim": 53,
|
||||
"observation_dim": 265,
|
||||
"action_dim": 16,
|
||||
"dt": 0.02,
|
||||
"actor_hidden": [
|
||||
256,
|
||||
128,
|
||||
64
|
||||
],
|
||||
"actor_params": 110288,
|
||||
"frame_terms": [
|
||||
{
|
||||
"name": "base_ang_vel",
|
||||
"dim": 3,
|
||||
"scale": 1.0
|
||||
},
|
||||
{
|
||||
"name": "projected_gravity",
|
||||
"dim": 3,
|
||||
"scale": 1.0
|
||||
},
|
||||
{
|
||||
"name": "joint_pos_rel",
|
||||
"dim": 12,
|
||||
"scale": 1.0
|
||||
},
|
||||
{
|
||||
"name": "joint_vel_rel",
|
||||
"dim": 12,
|
||||
"scale": 0.05
|
||||
},
|
||||
{
|
||||
"name": "wheel_vel_rel",
|
||||
"dim": 4,
|
||||
"scale": 0.05
|
||||
},
|
||||
{
|
||||
"name": "last_action",
|
||||
"dim": 16,
|
||||
"scale": 1.0
|
||||
},
|
||||
{
|
||||
"name": "command",
|
||||
"dim": 3,
|
||||
"scale": 1.0
|
||||
}
|
||||
],
|
||||
"history_layout": "term-wise blocks, oldest to newest (mjlab flattening)",
|
||||
"env_action_order": [
|
||||
0,
|
||||
1,
|
||||
2,
|
||||
3,
|
||||
4,
|
||||
5,
|
||||
6,
|
||||
7,
|
||||
8,
|
||||
9,
|
||||
10,
|
||||
11,
|
||||
12,
|
||||
13,
|
||||
14,
|
||||
15
|
||||
],
|
||||
"leg_scale": [
|
||||
0.25,
|
||||
0.6,
|
||||
0.8,
|
||||
0.25,
|
||||
0.6,
|
||||
0.8,
|
||||
0.25,
|
||||
0.6,
|
||||
0.8,
|
||||
0.25,
|
||||
0.6,
|
||||
0.8
|
||||
],
|
||||
"wheel_scale": 35.0,
|
||||
"wheel_clip": 35.0,
|
||||
"default_joint_pos": [
|
||||
0.05,
|
||||
-0.4363323129985824,
|
||||
1.0,
|
||||
-0.05,
|
||||
-0.4363323129985824,
|
||||
1.0,
|
||||
0.05,
|
||||
0.4363323129985824,
|
||||
-1.0,
|
||||
-0.05,
|
||||
0.4363323129985824,
|
||||
-1.0,
|
||||
0.0,
|
||||
0.0,
|
||||
0.0,
|
||||
0.0
|
||||
],
|
||||
"leg_joint_names": [
|
||||
"fl_hipx_joint",
|
||||
"fl_hipy_joint",
|
||||
"fl_knee_joint",
|
||||
"fr_hipx_joint",
|
||||
"fr_hipy_joint",
|
||||
"fr_knee_joint",
|
||||
"hl_hipx_joint",
|
||||
"hl_hipy_joint",
|
||||
"hl_knee_joint",
|
||||
"hr_hipx_joint",
|
||||
"hr_hipy_joint",
|
||||
"hr_knee_joint"
|
||||
],
|
||||
"wheel_joint_names": [
|
||||
"fl_wheel_joint",
|
||||
"fr_wheel_joint",
|
||||
"hl_wheel_joint",
|
||||
"hr_wheel_joint"
|
||||
],
|
||||
"onnx": {
|
||||
"input": "obs",
|
||||
"output": "actions",
|
||||
"opset": 17,
|
||||
"max_abs_error_vs_torch": 7.897615432739258e-07
|
||||
}
|
||||
}
|
||||
Binary file not shown.
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,34 @@
|
||||
[build-system]
|
||||
requires = ["setuptools>=68", "wheel"]
|
||||
build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "s10-supermimic"
|
||||
version = "0.1.0"
|
||||
description = "Blind wheel-legged locomotion for Deep Robotics S10 with a learned world model"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10,<3.14"
|
||||
dependencies = [
|
||||
"mjlab==1.6.0",
|
||||
"torch>=2.7.0",
|
||||
"numpy",
|
||||
]
|
||||
|
||||
[tool.setuptools]
|
||||
include-package-data = true
|
||||
packages = [
|
||||
"s10_supermimic",
|
||||
"s10_supermimic.core",
|
||||
"s10_supermimic.env",
|
||||
"s10_supermimic.tasks",
|
||||
"s10_supermimic.robot",
|
||||
]
|
||||
|
||||
[tool.setuptools.package-data]
|
||||
s10_supermimic = [
|
||||
"core/gui.html",
|
||||
"robot/LICENSE_s10",
|
||||
"robot/model/*.xml",
|
||||
"robot/model/*.md",
|
||||
"robot/model/meshes/*.STL",
|
||||
]
|
||||
@@ -0,0 +1,12 @@
|
||||
# Core runtime
|
||||
torch>=2.7.0
|
||||
mjlab==1.6.0
|
||||
mujoco~=3.11.0
|
||||
mujoco-warp~=3.11.0
|
||||
warp-lang>=1.14.0
|
||||
numpy
|
||||
|
||||
# Video recording and ONNX export
|
||||
mediapy>=1.2.7
|
||||
onnx
|
||||
onnxruntime
|
||||
@@ -0,0 +1 @@
|
||||
"""S10 wheel-legged locomotion with model-based RL in mjlab."""
|
||||
@@ -0,0 +1 @@
|
||||
"""S10 policies, world models, replay, imagined rollouts and control UI."""
|
||||
@@ -0,0 +1,88 @@
|
||||
"""Explicit name-based mapping between mjlab and canonical S10 model tensors."""
|
||||
|
||||
import torch
|
||||
|
||||
from s10_supermimic.robot.config import (
|
||||
LEG_JOINT_NAMES, WHEEL_BODY_NAMES, WHEEL_GEOM_NAMES, WHEEL_JOINT_NAMES,
|
||||
)
|
||||
from .world_model import STATE_DIM, ACTION_DIM
|
||||
|
||||
|
||||
class S10Bridge:
|
||||
def __init__(self, env):
|
||||
self.env = env
|
||||
self.robot = env.scene["robot"]
|
||||
self.legs, _ = self.robot.find_joints(LEG_JOINT_NAMES, preserve_order=True)
|
||||
self.wheels, _ = self.robot.find_joints(WHEEL_JOINT_NAMES, preserve_order=True)
|
||||
if len(self.legs) != 12 or len(self.wheels) != 4:
|
||||
raise ValueError("S10 requires 12 leg joints and 4 wheel joints")
|
||||
pos = env.action_manager.get_term("joint_pos")
|
||||
vel = env.action_manager.get_term("wheel_vel")
|
||||
if list(env.cfg.actions) != ["joint_pos", "wheel_vel"]:
|
||||
raise ValueError("Expected joint_pos then wheel_vel action terms")
|
||||
names = list(pos.target_names) + list(vel.target_names)
|
||||
canonical = LEG_JOINT_NAMES + WHEEL_JOINT_NAMES
|
||||
if len(names) != ACTION_DIM or set(names) != set(canonical):
|
||||
raise ValueError("Unexpected S10 action targets")
|
||||
self.to_canonical = [names.index(n) for n in canonical]
|
||||
self.to_env = [canonical.index(n) for n in names]
|
||||
geoms, _ = self.robot.find_geoms(WHEEL_GEOM_NAMES, preserve_order=True)
|
||||
self.geoms = self.robot.indexing.geom_ids[geoms]
|
||||
body, _ = self.robot.find_bodies("base_link")
|
||||
self.base = self.robot.indexing.body_ids[body[0]]
|
||||
# XML reference, not an already-randomized simulator value after reset.
|
||||
from s10_supermimic.robot.config import get_spec
|
||||
self.nominal_com = torch.as_tensor(get_spec().body("base_link").ipos.copy(),
|
||||
device=env.device, dtype=torch.float32)
|
||||
|
||||
def canonical_action(self, action):
|
||||
return action[:, self.to_canonical]
|
||||
|
||||
def env_action(self, action):
|
||||
return action[:, self.to_env]
|
||||
|
||||
def state(self):
|
||||
d = self.robot.data
|
||||
s = torch.cat((d.root_link_pos_w[:, 2:3], d.projected_gravity_b,
|
||||
d.root_link_lin_vel_b, d.root_link_ang_vel_b,
|
||||
(d.joint_pos - d.default_joint_pos)[:, self.legs],
|
||||
d.joint_vel[:, self.legs], d.joint_vel[:, self.wheels]), -1)
|
||||
assert s.shape[-1] == STATE_DIM
|
||||
return s
|
||||
|
||||
def pose(self):
|
||||
d = self.robot.data
|
||||
return torch.cat((d.root_link_pos_w[:, :2], d.heading_w[:, None]), -1)
|
||||
|
||||
def context(self):
|
||||
model = self.env.sim.model
|
||||
friction = model.geom_friction
|
||||
com = model.body_ipos
|
||||
mu = friction[:, self.geoms, 0] if friction.ndim == 3 else friction[self.geoms, 0].expand(self.env.num_envs, -1)
|
||||
offset = com[:, self.base] if com.ndim == 3 else com[self.base].expand(self.env.num_envs, -1)
|
||||
return torch.cat((mu, offset - self.nominal_com), -1)
|
||||
|
||||
def contacts(self):
|
||||
"""Endpoint targets: contact flags4, net wheel forces12 in base frame (N).
|
||||
|
||||
Explicit names prevent silent FL/FR swaps in sensor expansion order.
|
||||
Terminal/reset endpoints are excluded by replay validity in the learner.
|
||||
"""
|
||||
from mjlab.utils.lab_api.math import quat_apply_inverse
|
||||
sensor = self.env.scene["wheel_ground_contact"]
|
||||
names = [name.split('/')[-1] for name in sensor.primary_names]
|
||||
order = [names.index(name) for name in WHEEL_BODY_NAMES]
|
||||
found = (sensor.data.found[:,order] > 0).float().reshape(-1,4)
|
||||
force = sensor.data.force[:,order].reshape(-1,4,3)
|
||||
quat = self.robot.data.root_link_quat_w[:,None,:].expand(-1,4,-1)
|
||||
force = quat_apply_inverse(quat.reshape(-1,4),force.reshape(-1,3)).reshape(-1,12)
|
||||
return torch.cat((found,force),-1)
|
||||
|
||||
|
||||
def pose_delta(before, after):
|
||||
dxy = after[:, :2] - before[:, :2]
|
||||
c, s = before[:, 2].cos(), before[:, 2].sin()
|
||||
yaw = after[:, 2] - before[:, 2]
|
||||
return torch.stack((c * dxy[:, 0] + s * dxy[:, 1],
|
||||
-s * dxy[:, 0] + c * dxy[:, 1],
|
||||
torch.atan2(yaw.sin(), yaw.cos())), -1)
|
||||
@@ -0,0 +1,24 @@
|
||||
"""Non-wheel collision termination and pre-reset labels for imagined survival."""
|
||||
import torch
|
||||
|
||||
|
||||
def illegal_nonwheel_contact(env, force_threshold=10.):
|
||||
data = env.scene['nonwheel_ground_contact'].data
|
||||
force = data.force_history if data.force_history is not None else data.force
|
||||
finite = torch.isfinite(force).flatten(1).all(-1)
|
||||
magnitude = torch.nan_to_num(force, nan=0., posinf=0., neginf=0.).norm(dim=-1)
|
||||
peak = magnitude.flatten(1).amax(-1)
|
||||
hit = (peak > force_threshold) & finite
|
||||
# Scene/termination buffers are cleared by automatic reset. These copies
|
||||
# survive it, so the learner retains the transition that caused termination.
|
||||
env._nonwheel_hit = hit.detach().clone()
|
||||
env._nonwheel_label_valid = finite.detach().clone()
|
||||
env._nonwheel_peak = peak.detach().clone()
|
||||
return hit
|
||||
|
||||
|
||||
def collision_events(batch):
|
||||
valid = batch['valid'].squeeze(-1).bool()
|
||||
prefix = torch.cat((torch.ones_like(valid[:,:1]), valid.cumprod(1).bool()[:,:-1]),1)
|
||||
return ((batch['collision'].squeeze(-1) > .5)
|
||||
& batch['collision_valid'].squeeze(-1).bool() & prefix).any(-1)
|
||||
@@ -0,0 +1,32 @@
|
||||
"""Velocity range presets. V4 is time-based; v5/v6/v7 are stair-gated.
|
||||
|
||||
S10 v7 uses ``V7_STAGES`` as the 1.5/0.8/1.2 cap, but stage advancement comes
|
||||
from ``stair_curriculum.next_stage`` (actual up-stair pass rate). Only v4
|
||||
advances by iteration count alone.
|
||||
"""
|
||||
STAGES = (
|
||||
(0, (-.3,.8), (-.2,.2), (-.8,.8)),
|
||||
(1000, (-.5,1.), (-.4,.4), (-1.,1.)),
|
||||
(2500, (-.8,1.5), (-.6,.6), (-1.5,1.5)),
|
||||
(5000, (-1.5,1.8), (-.8,.8), (-2.,2.)),
|
||||
(7500, (-2.,2.), (-1.,1.), (-3.,3.)),
|
||||
)
|
||||
V7_STAGES = (
|
||||
(0, (-.3,.8), (-.2,.2), (-.5,.5)),
|
||||
(1000, (-.5,1.), (-.4,.4), (-.7,.7)),
|
||||
(2500, (-.8,1.2), (-.6,.6), (-.9,.9)),
|
||||
(5000, (-1.2,1.4), (-.7,.7), (-1.1,1.1)),
|
||||
(7500, (-1.5,1.5), (-.8,.8), (-1.2,1.2)),
|
||||
)
|
||||
|
||||
|
||||
def stages_for(revision):
|
||||
return V7_STAGES if revision >= 7 else STAGES
|
||||
|
||||
|
||||
def set_command_stage(cfg, iteration, stage=None, stages=STAGES):
|
||||
index = stage if stage is not None else max(
|
||||
i for i, item in enumerate(stages) if max(0, iteration) >= item[0])
|
||||
_, x, y, yaw = stages[index]
|
||||
cfg.ranges.lin_vel_x, cfg.ranges.lin_vel_y, cfg.ranges.ang_vel_z = x, y, yaw
|
||||
return index
|
||||
@@ -0,0 +1,114 @@
|
||||
<!doctype html>
|
||||
<html lang="zh-CN">
|
||||
<head>
|
||||
<meta charset="utf-8"><meta name="viewport" content="width=device-width,initial-scale=1">
|
||||
<title>S10 · 运行控制</title>
|
||||
<style>
|
||||
:root{color-scheme:dark;font-family:system-ui,-apple-system,sans-serif;background:#10141b;color:#e4eaf3}*{box-sizing:border-box}body{margin:0;padding:28px}header{display:flex;justify-content:space-between;align-items:center;margin-bottom:24px}h1{font-size:25px;margin:0}header small{color:#93a1b5;font-size:13px}.status{background:#202a36;padding:8px 14px;border-radius:20px;color:#a2b7cb;font-size:13px}main{display:grid;grid-template-columns:minmax(400px,1fr) 350px;gap:22px;max-width:1600px;margin:auto}.card{background:#19212d;border:1px solid #2c3748;border-radius:16px;overflow:hidden}.viewer{background:#080b10;display:flex;align-items:center;justify-content:center;min-height:400px}.viewer img{display:block;width:100%;aspect-ratio:4/3;object-fit:contain}.caption{padding:16px;color:#9cabc0;font-size:13px}.panel{padding:24px}h2{font-size:18px;margin:0 0 8px}.muted{color:#9cabc0;font-size:13px;line-height:1.7}.slider{margin:26px 0}.slider label{display:flex;justify-content:space-between;align-items:center;font-size:14px}output{font:600 21px ui-monospace,monospace;color:#64debf}input[type=range]{width:100%;accent-color:#58cfb1;margin:17px 0 5px;cursor:pointer}.bounds{display:flex;justify-content:space-between;color:#8292a7;font-size:12px}.actions{display:flex;gap:10px;margin:12px 0}button{cursor:pointer;border:1px solid #415067;border-radius:9px;background:#273448;color:#edf3fa;padding:11px 14px;font:inherit;font-size:14px;flex:1}button:hover{filter:brightness(1.2)}button.stop{background:#94404c;border-color:#c05b69;width:100%;font-weight:600}.telemetry{border-top:1px solid #344054;padding-top:20px;margin-top:24px}.metric{display:flex;justify-content:space-between;padding:8px 0;font-size:14px}.metric b{font-family:ui-monospace,monospace}#wm{font-size:11px;max-width:210px;text-align:right;white-space:normal;line-height:1.35}.dot{color:#64debf}.note{margin-top:20px;font-size:12px;color:#8a9ab0;line-height:1.7}@media(max-width:900px){body{padding:16px}main{grid-template-columns:1fr}.viewer{min-height:200px}}
|
||||
select{background:#273448;color:#edf3fa;padding:8px;border:1px solid #415067;border-radius:6px;margin:5px}.viewer img{cursor:grab;touch-action:none;user-select:none}.viewer img:active{cursor:grabbing}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<header><div><h1>S10 <span style="font-weight:400;color:#9cabc0">/ 运行控制</span></h1><small id="model">正在加载策略…</small></div><div class="status" id="connection">连接中</div></header>
|
||||
<main>
|
||||
<section class="card"><div class="viewer"><img id="frame" src="/frame.jpg" alt="S10 实时仿真画面"></div><div class="caption" id="caption">单环境 · CPU · 实时画面</div></section>
|
||||
<section class="card panel">
|
||||
<h2>地形与视角</h2>
|
||||
<div><select id="terrainKind" aria-label="地形"><option value="flat">平地</option><option value="up">上楼梯</option><option value="down">下楼梯</option><option value="platform">障碍平台</option><option value="slope">斜坡</option><option value="rubble">碎石地</option></select>
|
||||
<select id="terrainLevel" aria-label="难度"><option value="0">等级 0 · 2 cm / 3°</option><option value="1">等级 1 · 4 cm / 5.4°</option><option value="2">等级 2 · 8 cm / 7.8°</option><option value="3">等级 3 · 12 cm / 10.2°</option><option value="4">等级 4 · 16 cm / 12.6°</option><option value="5">等级 5 · 20 cm / 15°</option></select></div>
|
||||
<div class="actions"><button id="changeScene">切换地形</button><button id="resetCamera">恢复视角</button></div>
|
||||
<div class="muted">在画面上按住左键拖动旋转，滚轮缩放，双击恢复视角。镜头始终跟随机器人。切换地形会重置机器人。</div>
|
||||
|
||||
<h2>键盘控制 · 速度大小</h2><div class="muted">↑↓ 前进/后退 · ←→ 横移 · Q/E 转向<br>按住移动，松开停止；滑条只调速度大小。</div><div class="muted" id="requested">当前指令：0 / 0 / 0</div>
|
||||
<div class="slider"><label for="vx">前进 / 后退 · vx <output id="vxout">0.00</output></label><input id="vx" type="range" min="0" max="2" step="0.01" value="0.3"><div class="bounds"><span>0 m/s</span><span>+2.00 m/s</span></div></div>
|
||||
<div class="slider"><label for="vy">左移 / 右移 · vy <output id="vyout">0.00</output></label><input id="vy" type="range" min="0" max="1" step="0.01" value="0.2"><div class="bounds"><span>0 m/s</span><span>+1.00 m/s</span></div></div>
|
||||
<div class="slider"><label for="yaw">左转 / 右转 · yaw rate <output id="yawout">0.00</output></label><input id="yaw" type="range" min="0" max="3" step="0.01" value="0.4"><div class="bounds"><span>0 rad/s</span><span>+3.00 rad/s</span></div></div>
|
||||
<button class="stop" id="stop">停止 · 空格</button><div class="actions"><button id="pause">暂停仿真</button><button id="reset">重置机器人</button></div>
|
||||
<div class="telemetry"><h2>实际速度</h2><div class="metric"><span>vx</span><b id="actual0">—</b></div><div class="metric"><span>vy</span><b id="actual1">—</b></div><div class="metric"><span>yaw rate</span><b id="actual2">—</b></div><div class="metric"><span>WM 预测误差</span><b id="wm">—</b></div></div>
|
||||
<div class="note" id="note">黄→橙→红渐变残影是世界模型从当前状态想象的未来 32 步（默认每 4 步画一个）；WM 预测误差会随仿真推进与真实状态对齐后更新。切换标签页或关闭页面会归零。连接中断超过 2 秒也会归零。暂停会清空指令。</div>
|
||||
</section>
|
||||
</main>
|
||||
<script>
|
||||
const sliders=['vx','vy','yaw'].map(id=>document.getElementById(id));
|
||||
const connection=document.getElementById('connection');
|
||||
let dirty=false,busy=false,actionQueue=[],dragging=false;
|
||||
let cameraDelta=[0,0,0],pointer=null,sceneInitialized=false;
|
||||
const keys=new Set(),motionKeys=new Set(["ArrowUp","ArrowDown","ArrowLeft","ArrowRight","KeyQ","KeyE"]);
|
||||
function command(){
|
||||
const x=Number(keys.has("ArrowUp"))-Number(keys.has("ArrowDown"));
|
||||
const y=Number(keys.has("ArrowLeft"))-Number(keys.has("ArrowRight"));
|
||||
const z=Number(keys.has("KeyQ"))-Number(keys.has("KeyE"));
|
||||
const scale=Math.max(1,Math.hypot(x,y));
|
||||
return [Number(sliders[0].value)*x/scale,Number(sliders[1].value)*y/scale,Number(sliders[2].value)*z];
|
||||
}
|
||||
function labels(){sliders.forEach(s=>document.getElementById(s.id+'out').textContent=Number(s.value).toFixed(2))}
|
||||
function zero(){keys.clear();dirty=true}
|
||||
function action(name){if(name!=='heartbeat'){zero();actionQueue.push(name)}}
|
||||
sliders.forEach(s=>{s.addEventListener('input',()=>{labels();dirty=true});s.addEventListener('pointerdown',()=>dragging=true)});
|
||||
window.addEventListener('pointerup',()=>dragging=false);
|
||||
document.getElementById('stop').onclick=()=>action('stop');
|
||||
document.getElementById('pause').onclick=()=>action('pause');
|
||||
document.getElementById('reset').onclick=()=>action('reset');
|
||||
window.addEventListener('keydown',e=>{
|
||||
if(e.code==='Space'){e.preventDefault();action('stop');return}
|
||||
if(motionKeys.has(e.code)){e.preventDefault();if(!e.repeat){keys.add(e.code);dirty=true}}
|
||||
});
|
||||
window.addEventListener('keyup',e=>{
|
||||
if(motionKeys.has(e.code)){e.preventDefault();keys.delete(e.code);dirty=true}
|
||||
});
|
||||
window.addEventListener('blur',()=>{endDrag();action('stop');leave()});
|
||||
labels();
|
||||
function leave(){zero();navigator.sendBeacon('/api/control',JSON.stringify({action:'stop'}))}
|
||||
document.addEventListener('visibilitychange',()=>{if(document.hidden)leave()});
|
||||
window.addEventListener('pagehide',leave);
|
||||
function show(s){
|
||||
connection.textContent=s.loading?'正在切换地形…':s.ready?(s.paused?'已暂停':'● 已连接'):'仿真初始化中';
|
||||
if(!sceneInitialized && s.terrain_kind){
|
||||
document.getElementById('terrainKind').value=s.terrain_kind;
|
||||
document.getElementById('terrainLevel').value=String(s.terrain_level);sceneInitialized=true;
|
||||
}
|
||||
document.getElementById('changeScene').disabled=!!s.loading||!s.terrain;
|
||||
document.getElementById('pause').textContent=s.paused?'继续仿真':'暂停仿真';
|
||||
document.getElementById('requested').textContent='当前指令：vx '+s.command[0].toFixed(2)+' / vy '+s.command[1].toFixed(2)+' / yaw '+s.command[2].toFixed(2);
|
||||
(s.actual||[]).forEach((v,i)=>document.getElementById('actual'+i).textContent=Number(v).toFixed(3)+(i===2?' rad/s':' m/s'));
|
||||
document.getElementById('wm').textContent=s.wm_error||'—';
|
||||
if(s.checkpoint)document.getElementById('model').textContent=s.checkpoint+' · 训练步数 '+s.iteration;
|
||||
if(s.scene)document.getElementById('caption').textContent=s.scene+' · 单环境 · '+s.device+' · 渲染 '+Number(s.render_fps||0).toFixed(1)+' FPS · 实时倍率 '+Number(s.realtime_factor||0).toFixed(2)+'×';
|
||||
if(s.terrain)document.getElementById('note').textContent=(s.terrain_revision>=3?'新版训练指令范围：vx ±2 m/s，vy ±1 m/s，yaw ±3 rad/s。实际能力以评估为准。':'当前为前向地形策略。训练范围：vx 0.2–0.45 m/s，vy 0，yaw ±0.05 rad/s；反向和大幅转向尚未验证。')+' 黄→橙→红渐变残影是世界模型预测的未来 32 步；WM 误差在与真实状态对齐后更新。切换标签页或断开连接后指令归零。';
|
||||
}
|
||||
async function tick(){
|
||||
if(busy||document.hidden)return;busy=true;
|
||||
try{
|
||||
let body={action:'heartbeat'};
|
||||
if(actionQueue.length){const item=actionQueue.shift();body=typeof item==='string'?{action:item}:item;}
|
||||
else if(dirty){body={action:'command',values:command()};dirty=false}
|
||||
if(cameraDelta.some(v=>v!==0)){body.camera_delta=cameraDelta;cameraDelta=[0,0,0]}
|
||||
const response=await fetch('/api/control',{method:'POST',headers:{'Content-Type':'application/json'},body:JSON.stringify(body),signal:AbortSignal.timeout(1500)});
|
||||
if(!response.ok)throw Error(response.status);show(await response.json());
|
||||
}catch(e){connection.textContent='连接已断开';zero()}finally{busy=false}
|
||||
}
|
||||
setInterval(tick,33);tick();
|
||||
const image=document.getElementById('frame');
|
||||
image.draggable=false;
|
||||
image.addEventListener('pointerdown',e=>{
|
||||
if(e.button!==0)return;e.preventDefault();pointer={id:e.pointerId,x:e.clientX,y:e.clientY};image.setPointerCapture(e.pointerId);
|
||||
});
|
||||
image.addEventListener('pointermove',e=>{
|
||||
if(!pointer||pointer.id!==e.pointerId)return;
|
||||
cameraDelta[0]-=(e.clientX-pointer.x)*.4;cameraDelta[1]-=(e.clientY-pointer.y)*.3;
|
||||
pointer.x=e.clientX;pointer.y=e.clientY;
|
||||
});
|
||||
function endDrag(){pointer=null}
|
||||
image.addEventListener('pointerup',endDrag);image.addEventListener('pointercancel',endDrag);image.addEventListener('lostpointercapture',endDrag);
|
||||
|
||||
image.addEventListener('wheel',e=>{e.preventDefault();cameraDelta[2]+=Math.max(-1,Math.min(1,e.deltaY*(e.deltaMode===1?.03:e.deltaMode===2?.5:.001)));},{passive:false});
|
||||
function resetCamera(){cameraDelta=[0,0,0];actionQueue.push('camera_reset')}
|
||||
image.addEventListener('dblclick',resetCamera);
|
||||
document.getElementById('resetCamera').onclick=resetCamera;
|
||||
document.getElementById('changeScene').onclick=()=>{
|
||||
zero();actionQueue.push({action:'scene',kind:document.getElementById('terrainKind').value,level:Number(document.getElementById('terrainLevel').value)});
|
||||
};
|
||||
function nextFrame(){setTimeout(()=>{if(!document.hidden)image.src='/frame.jpg?t='+Date.now();else nextFrame()},16)}
|
||||
image.onload=nextFrame;image.onerror=nextFrame;if(image.complete)nextFrame();
|
||||
</script>
|
||||
</body></html>
|
||||
@@ -0,0 +1,171 @@
|
||||
"""Local browser controls. Latest-frame storage only; simulation owns MuJoCo."""
|
||||
import errno
|
||||
import json
|
||||
import math
|
||||
from pathlib import Path
|
||||
import threading
|
||||
import time
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
|
||||
|
||||
class ControlState:
|
||||
def __init__(self, timeout=2.):
|
||||
self.speed, self.lateral_speed, self.turn_speed = 2., 1., 3.
|
||||
self.lock = threading.RLock()
|
||||
self.values = (0., 0., 0.)
|
||||
self.last_seen = time.monotonic()
|
||||
self.timeout = timeout
|
||||
self.paused = False
|
||||
self.reset_requested = False
|
||||
self.camera = [140., -20., 2.5]
|
||||
self.scene_requested = None
|
||||
self.loading = False
|
||||
self.jpeg = b""
|
||||
self.status = {"ready": False, "actual": [0., 0., 0.]}
|
||||
|
||||
def clear(self):
|
||||
with self.lock:
|
||||
self.values = (0., 0., 0.)
|
||||
|
||||
def command(self):
|
||||
with self.lock:
|
||||
if time.monotonic() - self.last_seen > self.timeout:
|
||||
self.clear()
|
||||
return (0., 0., 0.) if self.paused or self.loading else self.values
|
||||
|
||||
def update(self, request):
|
||||
if not isinstance(request, dict):
|
||||
raise ValueError("Expected an object")
|
||||
action = request.get("action", "command")
|
||||
if action not in ("command", "heartbeat", "stop", "pause", "reset", "camera_reset", "scene"):
|
||||
raise ValueError("Unknown action")
|
||||
values = None
|
||||
if action == "command":
|
||||
values = request.get("values")
|
||||
if not isinstance(values, list) or len(values) != 3:
|
||||
raise ValueError("Expected vx, vy, yaw_rate")
|
||||
if any(isinstance(v, bool) or not isinstance(v, (int, float)) or not math.isfinite(v) for v in values):
|
||||
raise ValueError("Commands must be finite numbers")
|
||||
values = tuple(max(-limit, min(limit, value)) for value, limit in
|
||||
zip(values, (self.speed, self.lateral_speed, self.turn_speed)))
|
||||
camera_delta = request.get("camera_delta", [0., 0., 0.])
|
||||
if (not isinstance(camera_delta, list) or len(camera_delta) != 3
|
||||
or any(isinstance(v, bool) or not isinstance(v, (int, float)) or not math.isfinite(v) for v in camera_delta)):
|
||||
raise ValueError("Invalid camera delta")
|
||||
scene = None
|
||||
if action == "scene":
|
||||
kind, level = request.get("kind"), request.get("level")
|
||||
if kind not in ("flat", "up", "down", "platform", "slope", "rubble") or type(level) is not int or level not in range(6):
|
||||
raise ValueError("Invalid terrain selection")
|
||||
if not self.status.get("terrain", False):
|
||||
raise ValueError("Terrain selection requires a terrain policy")
|
||||
scene = (kind, level)
|
||||
with self.lock:
|
||||
self.camera[0] = (self.camera[0] + camera_delta[0]) % 360
|
||||
self.camera[1] = max(-89., min(0., self.camera[1] + camera_delta[1]))
|
||||
self.camera[2] = max(.5, min(12., self.camera[2] * math.exp(max(-2., min(2., camera_delta[2])))))
|
||||
if action == "camera_reset":
|
||||
self.camera = [140., -20., 2.5]
|
||||
if scene is not None:
|
||||
self.scene_requested = scene
|
||||
self.loading = True
|
||||
self.clear()
|
||||
# Expire stale commands before a late heartbeat can revive them.
|
||||
self.command()
|
||||
self.last_seen = time.monotonic()
|
||||
if values is not None and not self.loading:
|
||||
self.values = values
|
||||
elif action in ("stop", "reset", "pause"):
|
||||
self.clear()
|
||||
if action == "pause":
|
||||
self.paused = not self.paused
|
||||
elif action == "reset":
|
||||
self.reset_requested = True
|
||||
return self.snapshot()
|
||||
|
||||
def snapshot(self):
|
||||
with self.lock:
|
||||
return dict(self.status, command=list(self.command()), paused=self.paused, loading=self.loading, camera=list(self.camera))
|
||||
|
||||
def take_scene(self):
|
||||
with self.lock:
|
||||
scene, self.scene_requested = self.scene_requested, None
|
||||
return scene
|
||||
|
||||
def finish_scene(self):
|
||||
with self.lock:
|
||||
self.loading = self.scene_requested is not None
|
||||
self.clear()
|
||||
|
||||
def take_reset(self):
|
||||
with self.lock:
|
||||
reset, self.reset_requested = self.reset_requested, False
|
||||
return reset
|
||||
|
||||
def publish(self, actual, step, **metadata):
|
||||
with self.lock:
|
||||
self.status.update(ready=True, actual=actual, step=step, **metadata)
|
||||
|
||||
|
||||
def make_server(state, port=8765):
|
||||
html = Path(__file__).with_name("gui.html").read_bytes()
|
||||
|
||||
class Handler(BaseHTTPRequestHandler):
|
||||
def log_message(self, *args):
|
||||
pass
|
||||
|
||||
def send(self, status, data, kind):
|
||||
self.send_response(status)
|
||||
self.send_header("Content-Type", kind)
|
||||
self.send_header("Content-Length", str(len(data)))
|
||||
self.send_header("Cache-Control", "no-store")
|
||||
self.end_headers()
|
||||
try:
|
||||
self.wfile.write(data)
|
||||
except (BrokenPipeError, ConnectionResetError):
|
||||
pass
|
||||
|
||||
def do_GET(self):
|
||||
route = self.path.split("?", 1)[0]
|
||||
if route == "/":
|
||||
self.send(200, html, "text/html; charset=utf-8")
|
||||
elif route == "/api/state":
|
||||
self.send(200, json.dumps(state.snapshot()).encode(), "application/json")
|
||||
elif route == "/frame.jpg":
|
||||
with state.lock:
|
||||
jpeg = state.jpeg
|
||||
self.send(200 if jpeg else 503, jpeg, "image/jpeg")
|
||||
else:
|
||||
self.send(404, b"Not found", "text/plain")
|
||||
|
||||
def do_POST(self):
|
||||
if self.path != "/api/control":
|
||||
self.send(404, b"Not found", "text/plain")
|
||||
return
|
||||
# Local-only API: browser requests must come from this page's origin.
|
||||
origin = self.headers.get("Origin")
|
||||
if origin and origin != "http://" + self.headers.get("Host", ""):
|
||||
self.send(403, b"Wrong origin", "text/plain")
|
||||
return
|
||||
try:
|
||||
length = int(self.headers.get("Content-Length", "0"))
|
||||
if not 0 < length <= 2048:
|
||||
raise ValueError("Invalid body length")
|
||||
result = state.update(json.loads(self.rfile.read(length)))
|
||||
self.send(200, json.dumps(result).encode(), "application/json")
|
||||
except (ValueError, TypeError) as error:
|
||||
self.send(400, json.dumps({"error": str(error)}).encode(), "application/json")
|
||||
|
||||
requested_port = port
|
||||
while True:
|
||||
try:
|
||||
server = ThreadingHTTPServer(("127.0.0.1", port), Handler)
|
||||
break
|
||||
except OSError as error:
|
||||
if error.errno != errno.EADDRINUSE or port >= 65535:
|
||||
raise
|
||||
port += 1
|
||||
if port != requested_port:
|
||||
print(f"Port {requested_port} is occupied; using http://127.0.0.1:{server.server_port}/", flush=True)
|
||||
server.daemon_threads = True
|
||||
return server
|
||||
@@ -0,0 +1,198 @@
|
||||
"""Differentiable rewards, full-return actor objective and TD(lambda) critic targets."""
|
||||
|
||||
import math
|
||||
from contextlib import contextmanager
|
||||
import torch
|
||||
from .world_model import GRAVITY, LIN_VEL, ANG_VEL, LEG_POS, LEG_VEL, WHEEL_VEL
|
||||
from .policy import frame, push_history, critic_input
|
||||
|
||||
|
||||
@contextmanager
|
||||
def frozen(*modules):
|
||||
"""Freeze weights without disabling gradients through inputs/actions."""
|
||||
parameters = [(p, p.requires_grad) for m in modules for p in m.parameters()]
|
||||
try:
|
||||
for p, _ in parameters:
|
||||
p.requires_grad_(False)
|
||||
yield
|
||||
finally:
|
||||
for p, enabled in parameters:
|
||||
p.requires_grad_(enabled)
|
||||
|
||||
|
||||
def survival(state):
|
||||
gravity = torch.nn.functional.normalize(state[:, GRAVITY], dim=-1, eps=1e-6)
|
||||
# Smooth proxy for the environment's 70-degree fall cutoff, width ~5 deg.
|
||||
return torch.sigmoid((-gravity[:, 2] - math.cos(math.radians(70))) / 0.08)
|
||||
|
||||
|
||||
def continuation_value(state, constants):
|
||||
"""Imagined continuation probability for the current task.
|
||||
|
||||
The v4 proxy treats a 70-degree tilt as an immediate fall. The upward task
|
||||
instead terminates only after 4 s of bad orientation, which is 200 control
|
||||
steps, far beyond the 32/64-step imagination horizon, and the counter is not
|
||||
part of the state. Applying the v4 proxy there truncates the return after
|
||||
the first step whenever the robot is inverted, so the actor cannot learn to
|
||||
flip back.
|
||||
"""
|
||||
if constants.get("upright_mode") == "upward":
|
||||
return torch.ones(state.shape[0], dtype=state.dtype, device=state.device)
|
||||
return survival(state)
|
||||
|
||||
|
||||
def reward_terms(state, action, previous, command, constants, survival_prob, scan=None):
|
||||
v, w = state[:, LIN_VEL], state[:, ANG_VEL]
|
||||
q, qd, wd = state[:, LEG_POS], state[:, LEG_VEL], state[:, WHEEL_VEL]
|
||||
standing = (command.square().sum(-1) < 1e-6).float()
|
||||
# Endpoint PD torque surrogate: actual MuJoCo forces also reflect implicit
|
||||
# integration over substeps. The low-weight effort term is approximate.
|
||||
leg_torque = (constants["kp"] * (constants.get("leg_action_scale", 0.25) * action[:, :12] - q)
|
||||
- constants["kd"] * qd).clamp(-constants["effort"], constants["effort"])
|
||||
wheel_limit = constants.get("wheel_velocity_limit", 35.)
|
||||
wheel_effort = constants.get("wheel_effort", 45.)
|
||||
wheel_torque = (constants.get("wheel_kd", 2.) * ((constants.get("wheel_action_scale", 15.) * action[:, 12:]).clamp(-wheel_limit, wheel_limit) - wd)).clamp(-wheel_effort, wheel_effort)
|
||||
absolute_q = q + constants["default_q"]
|
||||
limits = ((constants["limits"][..., 0] - absolute_q).clamp_min(0)
|
||||
+ (absolute_q - constants["limits"][..., 1]).clamp_min(0)).sum(-1)
|
||||
posture = q.square().sum(-1)
|
||||
if constants.get("terrain_posture_v2", False) and not constants.get("wheel_local_rewards", False):
|
||||
if scan is None:
|
||||
raise ValueError("Terrain v2 posture reward requires the privileged scan")
|
||||
from .terrain_posture import posture_cost
|
||||
posture = posture_cost(q, scan, constants.get("posture_relief_floor", 0.125))
|
||||
terrain = constants.get("terrain_rewards", False)
|
||||
gravity = state[:, GRAVITY]
|
||||
if constants.get("upright_mode") == "upward":
|
||||
# Sign-aware recovery reward: 1 upright, 0.5 sideways, 0 upside down.
|
||||
upright = (0.5 * (1. - gravity[:, 2])).clamp(0., 1.)
|
||||
else:
|
||||
upright = torch.exp(
|
||||
-gravity[:, :2].square().sum(-1) / constants.get("upright_std", 0.45)**2)
|
||||
terms = {
|
||||
"track_linear_velocity": torch.exp(-((v[:, :2] - command[:, :2]).square().sum(-1)
|
||||
+ (0. if terrain else v[:, 2].square())) / 0.25),
|
||||
"track_angular_velocity": torch.exp(-((w[:, 2] - command[:, 2]).square()
|
||||
+ (0. if terrain else w[:, :2].square().sum(-1))) / 0.5),
|
||||
"upright": upright,
|
||||
"leg_posture": posture,
|
||||
"standstill": standing * (v.square().sum(-1) + w.square().sum(-1)
|
||||
+ 0.01 * wd.square().sum(-1)),
|
||||
"action_rate_l2": (action - previous).square().sum(-1),
|
||||
"joint_torques_l2": leg_torque.square().sum(-1) + wheel_torque.square().sum(-1),
|
||||
"dof_pos_limits": limits,
|
||||
"terminated": 1. - survival_prob,
|
||||
}
|
||||
|
||||
if constants.get("rolling_rewards", False) and not constants.get("wheel_local_rewards", False):
|
||||
if scan is None:
|
||||
raise ValueError("Rolling rewards require a terrain scan")
|
||||
from .rolling_rewards import rolling_costs
|
||||
terms.update(rolling_costs(state, command, scan, constants["default_q"]))
|
||||
if constants.get("obstacle_progress", False):
|
||||
from .obstacle_rewards import obstacle_progress
|
||||
terms["obstacle_progress"] = obstacle_progress(state, command, scan)
|
||||
if constants.get("wheel_body_intrusion", False):
|
||||
from .wheel_intrusion import wheel_body_intrusion
|
||||
terms["wheel_body_intrusion"] = wheel_body_intrusion(
|
||||
state, constants["default_q"])
|
||||
if constants.get("wheel_local_rewards", False):
|
||||
from .wheel_terrain_rewards import wheel_terms
|
||||
terms.update(wheel_terms(state, command, scan, constants["default_q"]))
|
||||
if constants.get("bad_orientation_penalty", False):
|
||||
from .upward import bad_orientation_value
|
||||
terms["bad_orientation_penalty"] = bad_orientation_value(gravity)
|
||||
from .upward import apply_reward_clamps
|
||||
apply_reward_clamps(terms, constants)
|
||||
return terms
|
||||
|
||||
|
||||
def lambda_returns(rewards, continuation, next_values, gamma=0.99, lam=0.95):
|
||||
carry = next_values[-1]
|
||||
result = []
|
||||
for t in reversed(range(len(rewards))):
|
||||
carry = rewards[t] + gamma * continuation[t] * ((1 - lam) * next_values[t] + lam * carry)
|
||||
result.append(carry)
|
||||
return torch.stack(result[::-1])
|
||||
|
||||
|
||||
def discounted_objective(rewards, continuation, terminal_value, gamma=.99):
|
||||
alive = torch.ones_like(terminal_value)
|
||||
total = torch.zeros_like(terminal_value)
|
||||
discount = 1.
|
||||
for reward, cont in zip(rewards, continuation):
|
||||
total = total + discount * alive * reward
|
||||
alive = alive * cont
|
||||
discount *= gamma
|
||||
return total + discount * alive * terminal_value
|
||||
|
||||
|
||||
def make_imagined_step(model, actor, constants, weights, env_action_order,
|
||||
dt=.02, noise=True, field=None):
|
||||
"""Pure single-step graph, shared by eager and compiled execution."""
|
||||
def step(state, previous, history, pose, scan, command, context, bias, timers):
|
||||
action = actor(history)
|
||||
model_context = context if field is None else torch.cat((context, scan), -1)
|
||||
collision_prob = None
|
||||
if constants.get('nonwheel_collision', False):
|
||||
nxt, delta, _, collision_logit = model.forward_all(state, action, previous, model_context)
|
||||
collision_prob = collision_logit.squeeze(-1).sigmoid()
|
||||
else:
|
||||
nxt, delta = model(state, action, previous, model_context)
|
||||
if field is not None:
|
||||
from .terrain import advance_pose
|
||||
pose = advance_pose(pose, delta)
|
||||
next_scan = None if field is None else field.scan(nxt, pose)
|
||||
cont = continuation_value(nxt, constants)
|
||||
if collision_prob is not None:
|
||||
cont = cont * (1. - collision_prob)
|
||||
terms = reward_terms(nxt, action, previous, command, constants, cont,
|
||||
scan=next_scan)
|
||||
if constants.get('pair_body_intrusion', False):
|
||||
from .wheel_intrusion import pair_body_intrusion
|
||||
terms['pair_body_intrusion'], timers = pair_body_intrusion(
|
||||
nxt, constants['default_q'], timers, dt)
|
||||
reward = dt * sum(weights[name] * value for name, value in terms.items())
|
||||
latest = frame(nxt, action[:, env_action_order], command, bias, noise=noise,
|
||||
wheel_vel=actor.actor_wheel_vel)
|
||||
history = push_history(history, latest, actor.history, actor.terms)
|
||||
return nxt, action, history, pose, next_scan, reward, cont, timers
|
||||
return step
|
||||
|
||||
|
||||
def rollout(model, actor, target_critic, seed, constants, weights, env_action_order,
|
||||
horizon=16, dt=0.02, gamma=0.99, lam=0.95, noise=True, field=None,
|
||||
step_fn=None):
|
||||
"""Actor: full discounted return + terminal value. Critic: detached TD(lambda).
|
||||
|
||||
Caller freezes dynamics and target critic weights, preserving state gradients.
|
||||
Reuse one terrain scan at each state, including the rollout's initial state.
|
||||
"""
|
||||
state, previous, history = seed["state"], seed["previous"], seed["history"]
|
||||
command, context, bias = seed["command"], seed["context"], seed["bias"]
|
||||
pose = seed.get("world_pose")
|
||||
timers = seed.get("pair_timer")
|
||||
if timers is None:
|
||||
timers = torch.zeros(state.shape[0], 4, device=state.device)
|
||||
def value_input(s, prev, scan):
|
||||
base = critic_input(s, command, context, prev)
|
||||
return base if field is None else torch.cat((base, scan), -1)
|
||||
if step_fn is None:
|
||||
step_fn = make_imagined_step(model, actor, constants, weights, env_action_order,
|
||||
dt, noise, field)
|
||||
rewards, continuations, inputs, next_inputs = [], [], [], []
|
||||
scan = None if field is None else field.scan(state, pose)
|
||||
for _ in range(horizon):
|
||||
inputs.append(value_input(state, previous, scan).detach())
|
||||
state, previous, history, pose, scan, reward, cont, timers = step_fn(
|
||||
state, previous, history, pose, scan, command, context, bias, timers)
|
||||
rewards.append(reward)
|
||||
continuations.append(cont)
|
||||
next_inputs.append(value_input(state, previous, scan).detach())
|
||||
final_value = target_critic(value_input(state, previous, scan)).squeeze(-1)
|
||||
r, c = torch.stack(rewards), torch.stack(continuations)
|
||||
objective = discounted_objective(r, c, final_value, gamma)
|
||||
with torch.no_grad():
|
||||
next_values = target_critic(torch.stack(next_inputs)).squeeze(-1)
|
||||
targets = lambda_returns(r.detach(), c.detach(), next_values, gamma, lam)
|
||||
return -objective.mean(), torch.stack(inputs), targets
|
||||
@@ -0,0 +1,52 @@
|
||||
"""Explicit straight/side/yaw command modes, with more stair traversal samples.
|
||||
|
||||
``deadzone`` zeroes per-axis commands below the threshold; v7 uses 0.05 so tiny
|
||||
creep commands become standing instead of ambiguous near-zero motion.
|
||||
"""
|
||||
from dataclasses import dataclass
|
||||
import torch
|
||||
from mjlab.tasks.velocity.mdp.velocity_command import UniformVelocityCommand, UniformVelocityCommandCfg
|
||||
|
||||
|
||||
def mix_commands(values, stairs):
|
||||
result = values.clone()
|
||||
mode = torch.rand(len(values),device=values.device)
|
||||
# General terrain: 40% pure vx, 20% pure vy, 20% pure yaw, 20% combined.
|
||||
result[mode<.4,1:] = 0.
|
||||
result[(mode>=.4)&(mode<.6),0] = 0.
|
||||
result[(mode>=.4)&(mode<.6),2] = 0.
|
||||
result[(mode>=.6)&(mode<.8),:2] = 0.
|
||||
# Stairs: 80% low-speed pure forward/backward, 20% full-range commands.
|
||||
result[stairs] = values[stairs]
|
||||
traverse = stairs & (mode<.8)
|
||||
magnitude = .2 + .4*torch.rand(len(values),device=values.device)
|
||||
result[traverse,0] = torch.where(values[traverse,0]<0,-magnitude[traverse],magnitude[traverse])
|
||||
result[traverse,1:] = 0.
|
||||
return result
|
||||
|
||||
|
||||
class MixedVelocityCommand(UniformVelocityCommand):
|
||||
def _resample_command(self, env_ids):
|
||||
super()._resample_command(env_ids)
|
||||
gen = self._env.cfg.scene.terrain.terrain_generator
|
||||
stair_columns = torch.tensor([c.kind in ('up','down') for c in gen.sub_terrains.values()],device=self.device)
|
||||
stairs = stair_columns[self._env.scene.terrain.terrain_types[env_ids]]
|
||||
mixed = mix_commands(self.vel_command_b[env_ids],stairs)
|
||||
# Stair magnitudes must obey the current curriculum bounds, including reverse.
|
||||
for i, bounds in enumerate((self.cfg.ranges.lin_vel_x, self.cfg.ranges.lin_vel_y, self.cfg.ranges.ang_vel_z)):
|
||||
mixed[:,i].clamp_(*bounds)
|
||||
# Per-axis deadband: tiny commands are standing, not an ambiguous creep.
|
||||
if self.cfg.deadzone > 0.:
|
||||
mixed = torch.where(mixed.abs() < self.cfg.deadzone,
|
||||
torch.zeros_like(mixed), mixed)
|
||||
self.vel_command_b[env_ids] = mixed
|
||||
if self.cfg.hold_stair_command:
|
||||
self.time_left[env_ids[stairs]] = self._env.cfg.episode_length_s + 1.
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class MixedVelocityCommandCfg(UniformVelocityCommandCfg):
|
||||
hold_stair_command: bool = False
|
||||
deadzone: float = 0.
|
||||
def build(self, env):
|
||||
return MixedVelocityCommand(self,env)
|
||||
@@ -0,0 +1,34 @@
|
||||
"""Bounded directional motion near uneven ground, shared by real and imagined steps."""
|
||||
import torch
|
||||
from .world_model import LIN_VEL
|
||||
|
||||
|
||||
def obstacle_gate(scan):
|
||||
local = scan.reshape(-1, 21, 13)[:, 5:16, 3:10].flatten(1)
|
||||
return ((local.amax(-1) - local.amin(-1) - .02) / .08).clamp(0., 1.)
|
||||
|
||||
|
||||
def obstacle_progress(state, command, scan):
|
||||
speed = command[:, :2].norm(dim=-1)
|
||||
along = (state[:, LIN_VEL][:, :2] * command[:, :2]).sum(-1) / speed.clamp_min(.1)
|
||||
# Signed and capped: no benefit for wheel spin, pure yaw, lifting or overspeed.
|
||||
return obstacle_gate(scan) * (speed > .1) * (along / speed.clamp_min(.2)).clamp(-1., 1.)
|
||||
|
||||
|
||||
def env_obstacle_progress(env):
|
||||
from .terrain import TerrainField, TerrainBridge
|
||||
from .stair_curriculum import buffers
|
||||
bridge = getattr(env, '_s10_posture_bridge', None)
|
||||
if bridge is None:
|
||||
bridge = TerrainBridge(env, TerrainField.from_env(env))
|
||||
env._s10_posture_bridge = bridge
|
||||
state = bridge.state()
|
||||
scan = bridge.field.scan(state, bridge.pose())
|
||||
command = env.command_manager.get_command('twist')
|
||||
buffers(env)
|
||||
stalled = ((command[:, :2].norm(dim=-1) > .1)
|
||||
& (state[:, LIN_VEL][:, :2].norm(dim=-1) < .05)
|
||||
& (obstacle_gate(scan) > .5))
|
||||
env._stair_stall_time = torch.where(stalled, env._stair_stall_time + env.step_dt, 0.)
|
||||
env._stair_stuck |= env._stair_stall_time >= 3.
|
||||
return obstacle_progress(state, command, scan)
|
||||
@@ -0,0 +1,149 @@
|
||||
"""S10 actor history and actor/critic networks, independent of the quadruped learner.
|
||||
|
||||
Two Actor observation schemas coexist:
|
||||
|
||||
* v4 / v4-upward / new v7: 53 values/frame, including measured wheel speed;
|
||||
* legacy v7: 49 values/frame, no measured wheel speed (the last action still
|
||||
contains the four wheel commands).
|
||||
|
||||
The learner, replay and deployment exporter select the schema explicitly. Old
|
||||
checkpoints are recognized from their first-layer width.
|
||||
"""
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
from .world_model import RunningNorm, ANG_VEL, GRAVITY, LEG_POS, LEG_VEL, WHEEL_VEL
|
||||
|
||||
LEGACY_TERMS = (3, 3, 12, 12, 4, 16, 3)
|
||||
TERMS = (3, 3, 12, 12, 16, 3)
|
||||
LEGACY_FRAME_DIM = sum(LEGACY_TERMS)
|
||||
FRAME_DIM = sum(TERMS)
|
||||
CRITIC_DIM = 38 + 3 + 7 + 16
|
||||
POLICY_FORMAT = "s10-flat-policy-v1"
|
||||
|
||||
|
||||
def frame(state, last_env_action, command, bias, noise=False, wheel_vel=True):
|
||||
"""Exactly the term order/scaling in env.flat (noise before scaling)."""
|
||||
values = [state[:, ANG_VEL], state[:, GRAVITY], state[:, LEG_POS] + bias,
|
||||
state[:, LEG_VEL]]
|
||||
widths = [0.2, 0.05, 0.01, 1.5]
|
||||
if wheel_vel:
|
||||
values.append(state[:, WHEEL_VEL])
|
||||
widths.append(1.5)
|
||||
values.extend((last_env_action, command))
|
||||
widths.extend((0., 0.))
|
||||
if noise:
|
||||
values = [v + (torch.rand_like(v) * 2 - 1) * w if w else v
|
||||
for v, w in zip(values, widths)]
|
||||
# Leg velocity is always index 3; wheel velocity, when present, is index 4.
|
||||
velocity_indices = (3, 4) if wheel_vel else (3,)
|
||||
return torch.cat([v * (0.05 if i in velocity_indices else 1.)
|
||||
for i, v in enumerate(values)], -1)
|
||||
|
||||
|
||||
def push_history(history, latest, length, terms):
|
||||
"""Term-wise histories, oldest to newest, matching mjlab flattening."""
|
||||
result, hidx, fidx = [], 0, 0
|
||||
for dim in terms:
|
||||
result.append(torch.cat((history[:, hidx + dim:hidx + dim * length],
|
||||
latest[:, fidx:fidx + dim]), -1))
|
||||
hidx += dim * length
|
||||
fidx += dim
|
||||
return torch.cat(result, -1)
|
||||
|
||||
|
||||
def critic_input(state, command, context, previous):
|
||||
return torch.cat((state, command, context, previous), -1)
|
||||
|
||||
|
||||
class Network(nn.Module):
|
||||
def __init__(self, inputs, outputs, hidden):
|
||||
super().__init__()
|
||||
self.norm = RunningNorm(inputs)
|
||||
layers = []
|
||||
for n in hidden:
|
||||
layers.extend((nn.Linear(inputs, n), nn.ELU()))
|
||||
inputs = n
|
||||
layers.append(nn.Linear(inputs, outputs))
|
||||
self.net = nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x):
|
||||
return self.net(self.norm(x))
|
||||
|
||||
|
||||
class Actor(Network):
|
||||
def __init__(self, history=5, hidden=(256, 128, 64), actor_wheel_vel=True):
|
||||
self.actor_wheel_vel = bool(actor_wheel_vel)
|
||||
self.terms = LEGACY_TERMS if self.actor_wheel_vel else TERMS
|
||||
self.frame_dim = sum(self.terms)
|
||||
super().__init__(self.frame_dim * history, 16, hidden)
|
||||
self.history, self.hidden = history, tuple(hidden)
|
||||
nn.init.uniform_(self.net[-1].weight, -0.01, 0.01)
|
||||
nn.init.zeros_(self.net[-1].bias)
|
||||
|
||||
def forward(self, obs):
|
||||
# Same bounded action distribution during collection and imagination.
|
||||
return super().forward(obs).tanh()
|
||||
|
||||
|
||||
def checkpoint_wheel_vel(checkpoint):
|
||||
"""Infer legacy observation schema from weights; validate explicit metadata."""
|
||||
flag = checkpoint.get("actor_wheel_vel")
|
||||
if "actor" in checkpoint:
|
||||
dim = checkpoint["actor"]["net.0.weight"].shape[1]
|
||||
history = checkpoint["history"]
|
||||
if dim not in (FRAME_DIM * history, LEGACY_FRAME_DIM * history):
|
||||
raise ValueError("Unsupported actor observation dimension")
|
||||
inferred = dim == LEGACY_FRAME_DIM * history
|
||||
if flag is not None and bool(flag) != inferred:
|
||||
raise ValueError("Actor observation metadata disagrees with weights")
|
||||
return inferred
|
||||
return bool(flag) if flag is not None else False
|
||||
|
||||
|
||||
def load_actor(checkpoint, device="cpu"):
|
||||
if checkpoint.get("format") not in (POLICY_FORMAT, "s10-terrain-policy-v1"):
|
||||
raise ValueError("Expected a S10 policy checkpoint, not other-robot or dynamics-only weights")
|
||||
actor = Actor(checkpoint["history"], checkpoint["actor_hidden"],
|
||||
checkpoint_wheel_vel(checkpoint)).to(device)
|
||||
actor.load_state_dict(checkpoint["actor"])
|
||||
return actor.eval()
|
||||
|
||||
|
||||
def evaluation_env_cfg(checkpoint, level=0, kind=None):
|
||||
"""Select the matching task; terrain remains privileged at deployment."""
|
||||
from s10_supermimic.tasks import get_task
|
||||
actor_wheel_vel = checkpoint_wheel_vel(checkpoint)
|
||||
if checkpoint.get("format") == "s10-terrain-policy-v1":
|
||||
checkpoint_task = checkpoint.get("config", {}).get("task")
|
||||
task = get_task(checkpoint_task) if checkpoint_task else None
|
||||
if task is not None:
|
||||
if task.terrain_revision != checkpoint["config"].get("terrain_revision", 1):
|
||||
raise ValueError("Checkpoint task revision differs from its config")
|
||||
if task.actor.wheel_vel != actor_wheel_vel:
|
||||
raise ValueError("Checkpoint task Actor schema differs from its weights")
|
||||
cfg = task.build_env_cfg(play=True, history=checkpoint["history"],
|
||||
seed=checkpoint["config"]["seed"])
|
||||
else:
|
||||
from s10_supermimic.env.terrain import s10_velocity_terrain_env_cfg
|
||||
cfg = s10_velocity_terrain_env_cfg(
|
||||
play=True, history=checkpoint["history"],
|
||||
seed=checkpoint["config"]["seed"],
|
||||
revision=checkpoint["config"].get("terrain_revision", 1),
|
||||
actor_wheel_vel=actor_wheel_vel)
|
||||
# A single staircase can overlap the robot's pre-reset origin at compile
|
||||
# time. Reserve contacts for that transient pose as well as normal rolling.
|
||||
cfg.sim.nconmax = max(cfg.sim.nconmax or 0, 512)
|
||||
cfg.sim.njmax = max(cfg.sim.njmax or 0, 2048)
|
||||
if level not in range(6):
|
||||
raise ValueError("Terrain level must be 0..5")
|
||||
gen = cfg.scene.terrain.terrain_generator
|
||||
# Evaluation fixes difficulty; one row, same seeded generator.
|
||||
gen.num_rows = 1
|
||||
gen.difficulty_range = (level/5,level/5)
|
||||
if kind is not None:
|
||||
gen.sub_terrains = {kind: gen.sub_terrains[kind]}
|
||||
return cfg
|
||||
from s10_supermimic.env.flat import s10_velocity_flat_env_cfg
|
||||
return s10_velocity_flat_env_cfg(play=True, history=checkpoint["history"],
|
||||
actor_wheel_vel=actor_wheel_vel)
|
||||
@@ -0,0 +1,21 @@
|
||||
"""Isolate invalid simulation rows before terrain lookup or replay insertion."""
|
||||
import torch
|
||||
|
||||
|
||||
def invalid_rows(env, obs):
|
||||
d = env.scene["robot"].data
|
||||
tensors = (d.root_link_pos_w, d.projected_gravity_b, d.root_link_lin_vel_b,
|
||||
d.root_link_ang_vel_b, d.joint_pos, d.joint_vel, obs["actor"])
|
||||
return ~torch.stack([torch.isfinite(t).flatten(1).all(-1) for t in tensors]).all(0)
|
||||
|
||||
|
||||
def recover_invalid_envs(env, obs, reward):
|
||||
bad = invalid_rows(env, obs) | ~torch.isfinite(reward)
|
||||
ids = bad.nonzero(as_tuple=False).flatten()
|
||||
if ids.numel():
|
||||
print(f"RECOVERY reset {ids.numel()} invalid environments; first IDs={ids[:16].tolist()}", flush=True)
|
||||
obs, _ = env.reset(env_ids=ids)
|
||||
# A failed reset must be handled by the supervisor, never added to replay.
|
||||
if invalid_rows(env, obs).any():
|
||||
raise RuntimeError("Environment reset failed to restore finite state")
|
||||
return obs, bad
|
||||
@@ -0,0 +1,157 @@
|
||||
"""Bounded S10 transition storage with chronological multi-step sampling."""
|
||||
|
||||
import torch
|
||||
from .world_model import STATE_DIM, ACTION_DIM, CONTEXT_DIM
|
||||
|
||||
|
||||
class Replay:
|
||||
def __init__(self, length, num_envs, device, history_dim=0, terrain=False,
|
||||
contacts=False, courses=False, collisions=False, pairs=False):
|
||||
if min(length, num_envs) < 1:
|
||||
raise ValueError("Replay length and environment count must be positive")
|
||||
self.length, self.size, self.cursor = length, 0, 0
|
||||
self.archive = None
|
||||
self.success_archive = None
|
||||
dims = dict(state=STATE_DIM, action=ACTION_DIM, previous=ACTION_DIM,
|
||||
context=CONTEXT_DIM, next=STATE_DIM, pose=3, valid=1)
|
||||
if terrain:
|
||||
dims["world_pose"] = 3
|
||||
if contacts:
|
||||
dims["next_contacts"] = 16
|
||||
if collisions:
|
||||
dims.update(collision=1, collision_valid=1)
|
||||
if courses:
|
||||
dims["course"] = 4 # type, level, uphill ground transition, top crossing
|
||||
if pairs:
|
||||
dims["pair_timer"] = 4 # per-wheel continuous 5/6-region intrusion time
|
||||
if history_dim:
|
||||
dims.update(history=history_dim, command=3, bias=12)
|
||||
self.data = {k: torch.zeros(length, num_envs, dim, device=device) for k, dim in dims.items()}
|
||||
|
||||
def add(self, **items):
|
||||
if items.keys() != self.data.keys():
|
||||
raise ValueError("Replay transition fields do not match configured fields")
|
||||
for k, value in items.items():
|
||||
self.data[k][self.cursor].copy_(value.detach())
|
||||
self.cursor = (self.cursor + 1) % self.length
|
||||
self.size = min(self.size + 1, self.length)
|
||||
|
||||
def sample(self, batch, window, keys=None):
|
||||
if self.size < window or window < 1:
|
||||
raise ValueError("Not enough transitions for a window")
|
||||
device = self.data["state"].device
|
||||
start = torch.randint(self.size - window + 1, (batch,), device=device)
|
||||
oldest = self.cursor if self.size == self.length else 0
|
||||
times = (oldest + start[:, None] + torch.arange(window, device=device)) % self.length
|
||||
envs = torch.randint(self.data["state"].shape[1], (batch, 1), device=device)
|
||||
return {k: v[times, envs] for k, v in self.data.items() if keys is None or k in keys}
|
||||
|
||||
|
||||
def sample_world_model(self, batch, window, prioritized=.25, archived=.25,
|
||||
capacity=100000, keep_prob=.02):
|
||||
keys = ('state','next','action','previous','context','pose','valid',
|
||||
'world_pose','next_contacts')
|
||||
if 'next_contacts' not in self.data:
|
||||
return self.sample(batch, window, keys)
|
||||
if self.archive is None:
|
||||
self.archive = WindowArchive(capacity)
|
||||
if 'course' in self.data:
|
||||
return self._sample_courses(batch, window, capacity, keep_prob)
|
||||
n_priority = int(batch * prioritized)
|
||||
n_archive = int(batch * archived) if self.archive.size else 0
|
||||
parts = [self.sample(batch-n_priority-n_archive, window, keys)]
|
||||
candidates = self.sample(4*max(n_priority,1), window, keys)
|
||||
score, rare = event_scores(candidates)
|
||||
selected = rare & (torch.rand_like(score) < keep_prob)
|
||||
self.archive.add({k:v[selected] for k,v in candidates.items()})
|
||||
if n_priority:
|
||||
indices = (score + .001*torch.rand_like(score)).topk(n_priority).indices
|
||||
parts.append({k:v[indices] for k,v in candidates.items()})
|
||||
if n_archive:
|
||||
parts.append(self.archive.sample(n_archive))
|
||||
return {k:torch.cat([p[k] for p in parts]) for k in parts[0]}
|
||||
|
||||
|
||||
def _sample_courses(self, batch, window, capacity, keep_prob):
|
||||
# Total archive capacity remains bounded by --archive-size, split evenly.
|
||||
if self.success_archive is None:
|
||||
self.archive = WindowArchive(max(1, capacity // 2))
|
||||
self.success_archive = WindowArchive(max(1, capacity - capacity // 2))
|
||||
keys = tuple(k for k in self.data if k not in ('history', 'command', 'bias'))
|
||||
n_priority = batch // 4
|
||||
candidates = self.sample(max(1, n_priority * 4), window, keys)
|
||||
score, rare = event_scores(candidates)
|
||||
valid = candidates['valid'].squeeze(-1).bool().cumprod(1).bool()
|
||||
tag = candidates['course']
|
||||
moving = (candidates['pose'][:, :, :2].norm(dim=-1) > .001) & valid
|
||||
up = tag[:, :, 0] == 1
|
||||
useful = ((tag[:, :, 2:] > .5).any(-1) & valid & up).any(-1)
|
||||
# Do not fill the general archive with repeated stationary stair impacts.
|
||||
rare &= (~up.any(-1) | moving.any(-1))
|
||||
if 'collision' in candidates:
|
||||
from .collision import collision_events
|
||||
rare |= collision_events(candidates)
|
||||
for archive, chosen in ((self.archive, rare), (self.success_archive, useful)):
|
||||
chosen = chosen & (torch.rand_like(score) < keep_prob)
|
||||
archive.add({k:v[chosen] for k,v in candidates.items()})
|
||||
n_success = batch // 8 if self.success_archive.size else 0
|
||||
n_rare = batch // 8 if self.archive.size else 0
|
||||
# At least half the batch is always uniform replay, including when sparse.
|
||||
parts = [self.sample(batch - n_priority - n_success - n_rare, window, keys)]
|
||||
if n_priority:
|
||||
rank = score + 8 * useful.float() + 2 * (moving & up).any(-1).float()
|
||||
selected = (rank + .001 * torch.rand_like(rank)).topk(n_priority).indices
|
||||
parts.append({k:v[selected] for k,v in candidates.items()})
|
||||
if n_success: parts.append(self.success_archive.sample(n_success))
|
||||
if n_rare: parts.append(self.archive.sample(n_rare))
|
||||
return {k:torch.cat([part[k] for part in parts]) for k in keys}
|
||||
|
||||
|
||||
def event_scores(w):
|
||||
"""Measured wheel contact transitions; ignore all labels after a reset.
|
||||
|
||||
Force impacts, sustained flight, hard vertical motion and episode ends enter
|
||||
the archive. Contact changes also include normal wheel touchdown/liftoff.
|
||||
"""
|
||||
valid = w['valid'].squeeze(-1).bool()
|
||||
mask = valid.cumprod(1).bool()
|
||||
contacts = w['next_contacts'][:,:,:4] > .5
|
||||
pairs = mask[:,1:] & mask[:,:-1]
|
||||
changes = ((contacts[:,1:] != contacts[:,:-1]).sum(-1)*pairs).sum(-1)
|
||||
flight = ((~contacts.any(-1)) & mask).sum(-1)
|
||||
force = torch.nan_to_num(w['next_contacts'][:,:,4:],nan=0.).reshape(*mask.shape,4,3).norm(dim=-1)
|
||||
impact = ((force.amax(-1)>200.) & mask).any(-1)
|
||||
hard_vz = ((w['state'][:,:,6].abs()>.6) & mask).any(-1)
|
||||
# A terminal endpoint is not a valid dynamics target, but its preceding
|
||||
# window is still useful. Windows starting after termination are excluded.
|
||||
ended = ((~valid) & torch.cat((torch.ones_like(mask[:,:1]),mask[:,:-1]),1)).any(-1)
|
||||
score = changes.float() + flight.float() + 4.*ended + 2.*impact
|
||||
rare = ((flight>=2) | hard_vz | impact | ended) & mask[:,0]
|
||||
score = score * mask[:,0]
|
||||
if 'collision' in w:
|
||||
from .collision import collision_events
|
||||
collision = collision_events(w)
|
||||
score = score + 8.*collision
|
||||
rare |= collision
|
||||
return score, rare
|
||||
|
||||
|
||||
class WindowArchive:
|
||||
"""Bounded GPU ring storing only world-model fields, not actor histories."""
|
||||
def __init__(self, capacity):
|
||||
self.capacity, self.size, self.cursor = capacity, 0, 0
|
||||
self.data = None
|
||||
|
||||
def add(self, windows):
|
||||
n = min(len(windows['state']),self.capacity)
|
||||
if not n:return
|
||||
if self.data is None:
|
||||
self.data = {k:v.new_empty((self.capacity,*v.shape[1:])) for k,v in windows.items()}
|
||||
ids = (torch.arange(n,device=windows['state'].device)+self.cursor)%self.capacity
|
||||
for k,v in windows.items():self.data[k][ids] = v[-n:].detach()
|
||||
self.cursor = (self.cursor+n)%self.capacity
|
||||
self.size = min(self.capacity,self.size+n)
|
||||
|
||||
def sample(self, batch):
|
||||
ids = torch.randint(self.size,(batch,),device=self.data['state'].device)
|
||||
return {k:v[ids] for k,v in self.data.items()}
|
||||
@@ -0,0 +1,77 @@
|
||||
"""Terrain-gated rolling preference, shared by real and imagined rewards."""
|
||||
import torch
|
||||
from .world_model import GRAVITY, LEG_POS, LEG_VEL
|
||||
|
||||
|
||||
def rolling_gate(scan, command):
|
||||
local = scan.reshape(-1, 21, 13)[:, 5:16, 3:10].flatten(1)
|
||||
relief = local.amax(-1) - local.amin(-1)
|
||||
# Flat (<2 cm): favor rolling; near obstacles (>=6 cm): allow lifting.
|
||||
flat = (1. - (relief - .02) / .04).clamp(0., 1.)
|
||||
# Non-steerable wheels may require stepping for lateral/yaw maneuvers.
|
||||
straight = torch.exp(-(command[:,1]/.25).square() - (command[:,2]/1.).square())
|
||||
return flat * straight
|
||||
|
||||
|
||||
def wheel_geometry(state, default_q):
|
||||
"""S10 wheel centres in body coordinates and vertical cylinder half-extents.
|
||||
|
||||
Joint order is hipx (abduction), hipy (flexion), knee. Geometry follows
|
||||
``S10_description/s10_mjcf/S10.xml``: hipx at x=±0.2277/y=±0.0574, hipy
|
||||
offset y=±0.0779, knee at z=-0.18, wheel at y=±0.045891/z=-0.18,
|
||||
wheel radius 0.081 m and half-width 0.0185 m. The analytic result was
|
||||
checked against MuJoCo ``mj_forward`` to machine precision.
|
||||
|
||||
The cylinder is symmetric about its axle, so wheel phase is unnecessary.
|
||||
"""
|
||||
q = (state[:,LEG_POS] + default_q).reshape(-1,4,3)
|
||||
hipx, hipy, knee = q.unbind(-1)
|
||||
side = state.new_tensor([1.,-1.,1.,-1.])
|
||||
# Leg sagittal plane: knee link (0.18 m) + wheel link (0.18 m).
|
||||
x = .18*knee.sin()
|
||||
z = -.18 - .18*knee.cos()
|
||||
x3 = hipy.cos()*x - hipy.sin()*z
|
||||
z3 = hipy.sin()*x + hipy.cos()*z
|
||||
y3 = side*(.0779+.045891)
|
||||
y4 = hipx.cos()*y3 + hipx.sin()*z3
|
||||
z4 = -hipx.sin()*y3 + hipx.cos()*z3
|
||||
point = torch.stack((state.new_tensor([.2277,.2277,-.2277,-.2277]) + x3,
|
||||
state.new_tensor([.0574,-.0574,.0574,-.0574]) + y4,
|
||||
z4),-1)
|
||||
up = -torch.nn.functional.normalize(state[:,GRAVITY],dim=-1,eps=1e-6)
|
||||
# Wheel axle is local -y, rotated only by hipx; sign is irrelevant after abs.
|
||||
axle_z = -up[:,1:2]*hipx.cos()+up[:,2:3]*hipx.sin()
|
||||
radius_z = .081*(1.-axle_z.square()).clamp_min(1e-8).sqrt()+.0185*axle_z.abs()
|
||||
return point, radius_z
|
||||
|
||||
|
||||
def wheel_bottom_height(state, default_q):
|
||||
"""Wheel bottoms relative to the ground below the base.
|
||||
|
||||
Only used on near-flat ground by rolling_gate, not as obstacle clearance.
|
||||
"""
|
||||
point, radius_z = wheel_geometry(state, default_q)
|
||||
up = -torch.nn.functional.normalize(state[:,GRAVITY],dim=-1,eps=1e-6)
|
||||
return state[:,:1] + (point*up[:,None,:]).sum(-1) - radius_z
|
||||
|
||||
|
||||
def rolling_costs(state, command, scan, default_q):
|
||||
gate = rolling_gate(scan, command)
|
||||
# Tolerate 5 mm of clearance; no positive reward for lifting anywhere.
|
||||
lift = ((wheel_bottom_height(state,default_q)-.005).clamp_min(0)/.05).square().sum(-1)
|
||||
# Reference mjlab task uses mean squared speed and relaxes during recovery.
|
||||
tilt = (state[:,GRAVITY][:,:2].norm(dim=-1)-.08).div(.22).clamp(0.,1.)
|
||||
motion = state[:,LEG_VEL].square().mean(-1) * (1.-.8*tilt)
|
||||
return {'unnecessary_wheel_lift':gate*lift, 'rolling_leg_motion':gate*motion}
|
||||
|
||||
|
||||
def env_rolling_cost(env, term):
|
||||
from .terrain import TerrainField, TerrainBridge
|
||||
bridge = getattr(env,'_s10_posture_bridge',None)
|
||||
if bridge is None:
|
||||
bridge = TerrainBridge(env,TerrainField.from_env(env))
|
||||
env._s10_posture_bridge = bridge
|
||||
state = bridge.state()
|
||||
return rolling_costs(state, env.command_manager.get_command('twist'),
|
||||
bridge.field.scan(state,bridge.pose()),
|
||||
bridge.robot.data.default_joint_pos[0,bridge.legs])[term]
|
||||
@@ -0,0 +1,75 @@
|
||||
"""V5 episode outcomes and performance-gated velocity curriculum.
|
||||
|
||||
All buffers are training-only metadata, never actor observations.
|
||||
"""
|
||||
import torch
|
||||
|
||||
KINDS = ('flat', 'up', 'down', 'platform', 'slope', 'rubble')
|
||||
MIN_ITERATIONS = (0, 3000, 8000, 15000, 22000)
|
||||
|
||||
|
||||
def buffers(env):
|
||||
if not hasattr(env, '_stair_stats'):
|
||||
env._stair_stats = torch.zeros(6, 6, 4, device=env.device) # trials, pass, fall, stuck
|
||||
env._stair_stall_time = torch.zeros(env.num_envs, device=env.device)
|
||||
env._stair_stuck = torch.zeros(env.num_envs, dtype=torch.bool, device=env.device)
|
||||
return env._stair_stats
|
||||
|
||||
|
||||
def stair_pass(offset, gravity, level, up):
|
||||
"""Require body/rear clearance of the last riser, correct landing height and stability."""
|
||||
heights = offset.new_tensor([.02, .04, .08, .12, .16, .20])[level]
|
||||
expected = torch.where(up, 5 * heights, -5 * heights)
|
||||
return ((offset[:, 0].abs() > 2.8) & (offset[:, 1].abs() < 2.5)
|
||||
& (offset[:, 2] > expected + .20) & (offset[:, 2] < expected + .60)
|
||||
& (gravity[:, :2].norm(dim=-1) < .4))
|
||||
|
||||
|
||||
def curriculum(env, env_ids):
|
||||
stats = buffers(env)
|
||||
terrain = env.scene.terrain
|
||||
d = env.scene['robot'].data
|
||||
offset = d.root_link_pos_w[env_ids] - env.scene.env_origins[env_ids]
|
||||
levels = terrain.terrain_levels[env_ids].clone()
|
||||
types = terrain.terrain_types[env_ids]
|
||||
# Six fixed columns, generated in KINDS order.
|
||||
stairs = (types == 1) | (types == 2)
|
||||
active = env.episode_length_buf[env_ids] > 0
|
||||
fallen = env.termination_manager.terminated[env_ids]
|
||||
success = (offset[:, 0].abs() > 5.4) & (offset[:, 1].abs() < 6.5)
|
||||
top = stair_pass(offset, d.projected_gravity_b[env_ids], levels, types == 1)
|
||||
# Landing contacts rule out simply jumping through the clearance test.
|
||||
contact = env.scene['wheel_ground_contact'].data.found[env_ids].reshape(len(env_ids), -1).gt(0).sum(-1) >= 2
|
||||
success = torch.where(stairs, top & contact, success) & active & ~fallen
|
||||
command = env.command_manager.get_command('twist')[env_ids]
|
||||
traversal = (command[:, 0].abs() > .1) & (command[:, 1:].abs().amax(-1) < .05)
|
||||
trials = (~stairs | traversal) & active & ((env.episode_length_buf[env_ids] * env.step_dt >= 5.) | fallen)
|
||||
values = torch.stack((trials, success & trials, fallen & trials,
|
||||
env._stair_stuck[env_ids] & trials), -1).float()
|
||||
stats.view(-1, 4).index_add_(0, types * 6 + levels, values)
|
||||
terrain.terrain_levels[env_ids] = (levels + success.long() - (active & fallen).long()).clamp(0, 5)
|
||||
terrain.env_origins[env_ids] = terrain.terrain_origins[terrain.terrain_levels[env_ids], types]
|
||||
env._stair_stall_time[env_ids] = 0
|
||||
env._stair_stuck[env_ids] = False
|
||||
return terrain.terrain_levels.float().mean()
|
||||
|
||||
|
||||
def next_stage(stage, iteration, stats, velocity_mae):
|
||||
"""Minimum age AND sufficient actual stair passes, never time alone."""
|
||||
if stage >= 4 or iteration < MIN_ITERATIONS[stage + 1]:
|
||||
return stage
|
||||
# Stage increases require success at 4, 8, 12, then 16 cm respectively.
|
||||
trials, passed = stats[1, stage + 1, :2]
|
||||
accurate = max(velocity_mae[:2]) < .18 and velocity_mae[2] < .30
|
||||
return stage + int(trials >= 64 and passed / max(trials, 1) >= .60 and accurate)
|
||||
|
||||
|
||||
def summary(env):
|
||||
stats = buffers(env).detach().cpu().tolist()
|
||||
levels, types = env.scene.terrain.terrain_levels, env.scene.terrain.terrain_types
|
||||
result = {}
|
||||
for i, name in enumerate(KINDS):
|
||||
selected = levels[types == i]
|
||||
result[name] = dict(level_counts=torch.bincount(selected, minlength=6).cpu().tolist(),
|
||||
episodes_by_level=stats[i])
|
||||
return result
|
||||
@@ -0,0 +1,95 @@
|
||||
"""Shared static height atlas and privileged, differentiable yaw-frame scans."""
|
||||
import numpy as np
|
||||
import torch
|
||||
from .bridge import S10Bridge
|
||||
|
||||
SCAN_DIM = 21 * 13
|
||||
|
||||
class TerrainField:
|
||||
def __init__(self, heights, origin, resolution=.04):
|
||||
self.h, self.origin, self.res = heights, origin, resolution
|
||||
x,y = torch.meshgrid(torch.linspace(-1,1,21,device=heights.device),
|
||||
torch.linspace(-.6,.6,13,device=heights.device), indexing="ij")
|
||||
self.offsets = torch.stack((x.flatten(),y.flatten()),-1)
|
||||
|
||||
def sample(self, xy):
|
||||
# Clamp *integer* neighbour indices too. NaN survives float clamp and
|
||||
# casts to INT64_MIN; on CUDA that would poison the entire device context.
|
||||
finite = torch.isfinite(xy).all(-1)
|
||||
clean = torch.where(torch.isfinite(xy), xy, torch.zeros_like(xy))
|
||||
gx = ((clean[...,0]-self.origin[0])/self.res).clamp(0,self.h.shape[1]-1)
|
||||
gy = ((clean[...,1]-self.origin[1])/self.res).clamp(0,self.h.shape[0]-1)
|
||||
x = gx.floor().long().clamp(0,self.h.shape[1]-2)
|
||||
y = gy.floor().long().clamp(0,self.h.shape[0]-2)
|
||||
a,b = gx-x,gy-y
|
||||
value = ((1-a)*self.h[y,x]+a*self.h[y,x+1])*(1-b)+((1-a)*self.h[y+1,x]+a*self.h[y+1,x+1])*b
|
||||
# Do not disguise corrupt states as valid ground: the learner will discard
|
||||
# this update, while real invalid environments are reset before querying.
|
||||
return torch.where(finite, value, torch.full_like(value, float("nan")))
|
||||
|
||||
def scan(self, state, pose):
|
||||
c,s = pose[:,2:3].cos(),pose[:,2:3].sin()
|
||||
dx,dy = self.offsets[:,0],self.offsets[:,1]
|
||||
xy = pose[:,None,:2]+torch.stack((c*dx-s*dy,s*dx+c*dy),-1)
|
||||
return (state[:,:1]+self.sample(pose[:,:2])[:,None]-self.sample(xy)).clamp(-1.5,1.5)
|
||||
|
||||
@classmethod
|
||||
def from_env(cls, env, resolution=.04):
|
||||
"""Rasterize exact oriented boxes in small bounding rectangles, no per-pixel rays.
|
||||
|
||||
Shared CPU build is ~6 MB at 4 cm for a 48 x 48 m atlas. No robot geometry.
|
||||
Only these static box courses are supported; fail if new shapes are introduced.
|
||||
"""
|
||||
import mujoco
|
||||
model = env.sim.mj_model
|
||||
data = mujoco.MjData(model)
|
||||
mujoco.mj_forward(model,data)
|
||||
robot_ids = set(env.scene["robot"].indexing.geom_ids.cpu().tolist())
|
||||
ids = [i for i in range(model.ngeom) if i not in robot_ids and model.geom_contype[i]]
|
||||
corners = np.array([[a,b,c] for a in (-1,1) for b in (-1,1) for c in (-1,1)])
|
||||
boxes=[]
|
||||
for i in ids:
|
||||
if model.geom_type[i] != mujoco.mjtGeom.mjGEOM_BOX:
|
||||
raise ValueError(f"Unsupported terrain geometry: {model.geom(i).name}")
|
||||
center, rot, size = data.geom_xpos[i],data.geom_xmat[i].reshape(3,3),model.geom_size[i]
|
||||
points = (corners*size) @ rot.T+center
|
||||
boxes.append((center.copy(),rot.copy(),size.copy(),points.min(0),points.max(0)))
|
||||
lo=np.min([b[3][:2] for b in boxes],axis=0)-resolution
|
||||
hi=np.max([b[4][:2] for b in boxes],axis=0)+resolution
|
||||
shape=np.ceil((hi-lo)/resolution).astype(int)+1
|
||||
heights=np.zeros((shape[1],shape[0]),dtype=np.float32)
|
||||
for center,rot,size,low,high in boxes:
|
||||
start=np.maximum(0,np.floor((low[:2]-lo)/resolution).astype(int))
|
||||
end=np.minimum(shape,np.ceil((high[:2]-lo)/resolution).astype(int)+1)
|
||||
xs=lo[0]+np.arange(start[0],end[0])*resolution
|
||||
ys=lo[1]+np.arange(start[1],end[1])*resolution
|
||||
x,y=np.meshgrid(xs,ys)
|
||||
# Intersect vertical line p=(x,y,z) with all three oriented-box slabs.
|
||||
lower=np.full_like(x,-np.inf); upper=np.full_like(x,np.inf); valid=np.ones_like(x,dtype=bool)
|
||||
for k in range(3):
|
||||
base=rot[0,k]*(x-center[0])+rot[1,k]*(y-center[1])-rot[2,k]*center[2]
|
||||
slope=rot[2,k]
|
||||
if abs(slope)<1e-10:
|
||||
valid &= np.abs(base)<=size[k]+1e-9
|
||||
else:
|
||||
z1=(-size[k]-base)/slope; z2=(size[k]-base)/slope
|
||||
lower=np.maximum(lower,np.minimum(z1,z2)); upper=np.minimum(upper,np.maximum(z1,z2))
|
||||
surface=np.where(valid & (lower<=upper),upper,0)
|
||||
patch=heights[start[1]:end[1],start[0]:end[0]]
|
||||
np.maximum(patch,surface,out=patch)
|
||||
return cls(torch.as_tensor(heights,device=env.device),tuple(lo),resolution)
|
||||
|
||||
|
||||
def advance_pose(pose, delta):
|
||||
c,s=pose[:,2].cos(),pose[:,2].sin()
|
||||
return pose+torch.stack((c*delta[:,0]-s*delta[:,1],s*delta[:,0]+c*delta[:,1],delta[:,2]),-1)
|
||||
|
||||
class TerrainBridge(S10Bridge):
|
||||
def __init__(self, env, field):
|
||||
super().__init__(env)
|
||||
self.field=field
|
||||
|
||||
def state(self):
|
||||
state=super().state()
|
||||
state[:,0]-=self.field.sample(self.robot.data.root_link_pos_w[:,:2])
|
||||
return state
|
||||
@@ -0,0 +1,33 @@
|
||||
"""Shared real/imagined posture objective for the opt-in terrain v2 task.
|
||||
|
||||
S10 joint order is hipx (abduction), hipy (flexion), knee. Hipx keeps a 2x
|
||||
weight while hipy/knee relax on rough terrain. v7 sets
|
||||
``relief_floor=0.02`` so stairs leave knee/hipy deviation almost unpenalised
|
||||
and the policy can flex instead of holding a stiff leg.
|
||||
"""
|
||||
import torch
|
||||
|
||||
LEG_SCALES = (0.25, 0.6, 0.8) * 4
|
||||
|
||||
|
||||
def posture_cost(q, scan, relief_floor=0.125):
|
||||
# Only the supporting/approaching ground near the body should relax posture.
|
||||
# scan is body height minus ground; its range is independent of body height.
|
||||
local = scan.reshape(-1, 21, 13)[:, 5:16, 3:10].flatten(1)
|
||||
relief = ((local.amax(-1) - local.amin(-1)) / 0.12).clamp(0., 1.)
|
||||
multiplier = (1. - relief).clamp_min(relief_floor)
|
||||
legs = q.reshape(-1, 4, 3).square()
|
||||
# Hip abduction stays constrained even on steps; flexion is freer there.
|
||||
return 2. * legs[:, :, 0].sum(-1) + multiplier * legs[:, :, 1:].sum((1, 2))
|
||||
|
||||
|
||||
def env_posture_cost(env, asset_cfg, relief_floor=0.125):
|
||||
from .terrain import TerrainField, TerrainBridge
|
||||
bridge = getattr(env, "_s10_posture_bridge", None)
|
||||
if bridge is None:
|
||||
bridge = TerrainBridge(env, TerrainField.from_env(env))
|
||||
env._s10_posture_bridge = bridge
|
||||
d = env.scene[asset_cfg.name].data
|
||||
q = d.joint_pos[:, asset_cfg.joint_ids] - d.default_joint_pos[:, asset_cfg.joint_ids]
|
||||
return posture_cost(q, bridge.field.scan(bridge.state(), bridge.pose()),
|
||||
relief_floor)
|
||||
@@ -0,0 +1,120 @@
|
||||
"""S10 terrain dynamics. Flat model checkpoints remain a separate format."""
|
||||
import torch
|
||||
from torch import nn
|
||||
from .world_model import S10WorldModel, RunningNorm
|
||||
from .terrain import SCAN_DIM, advance_pose
|
||||
|
||||
FORMAT = "s10-terrain-dynamics-v1"
|
||||
POLICY_FORMAT = "s10-terrain-policy-v1"
|
||||
|
||||
class TerrainWorldModel(S10WorldModel):
|
||||
def __init__(self, hidden=(512,512,512), contacts=False, collisions=False):
|
||||
super().__init__(hidden)
|
||||
self.net[0] = nn.Linear(77+SCAN_DIM,hidden[0])
|
||||
self.input_norm = RunningNorm(77+SCAN_DIM)
|
||||
self.predict_contacts = contacts
|
||||
if contacts:
|
||||
self.contact_head = nn.Linear(hidden[-1],16)
|
||||
self.predict_collisions = collisions
|
||||
if collisions:
|
||||
self.collision_head = nn.Linear(hidden[-1],1)
|
||||
nn.init.zeros_(self.collision_head.weight)
|
||||
nn.init.constant_(self.collision_head.bias, -4.6)
|
||||
self.last_contact_metrics = {}
|
||||
|
||||
def forward_all(self, state, action, previous_action, context):
|
||||
"""Predict endpoint state, pose, optional contacts and non-wheel collision logit."""
|
||||
x = self.input_norm(self.inputs(state,action,previous_action,context))
|
||||
for layer in list(self.net.children())[:-1]:
|
||||
x = layer(x)
|
||||
delta = self.net[-1](x)*self.output_norm.scale+self.output_norm.mean
|
||||
contacts = self.contact_head(x) if self.predict_contacts else None
|
||||
collision = self.collision_head(x) if self.predict_collisions else None
|
||||
return state+delta[:,:38],delta[:,38:],contacts,collision
|
||||
|
||||
def forward_with_contacts(self, state, action, previous_action, context):
|
||||
"""Predict endpoint state, pose, and optional [logits4, asinh(force/100)12]."""
|
||||
pred,delta,contacts,_ = self.forward_all(state,action,previous_action,context)
|
||||
return pred,delta,contacts
|
||||
|
||||
def forward(self, state, action, previous_action, context):
|
||||
pred,delta,_ = self.forward_with_contacts(state,action,previous_action,context)
|
||||
return pred,delta
|
||||
|
||||
@staticmethod
|
||||
def decode_contacts(prediction):
|
||||
"""Contact probabilities and signed base-frame force vectors in newtons."""
|
||||
return prediction[:,:4].sigmoid(),100.*prediction[:,4:].clamp(-10,10).sinh().reshape(-1,4,3)
|
||||
|
||||
@staticmethod
|
||||
def inputs(state, action, previous_action, context):
|
||||
if context.shape[-1] != 7+SCAN_DIM:
|
||||
raise ValueError("Terrain dynamics require context7 plus height scan273")
|
||||
return torch.cat((S10WorldModel.inputs(state,action,previous_action,context[:,:7]),context[:,7:]),-1)
|
||||
|
||||
def terrain_loss(self, b, field):
|
||||
state,pose=b["state"][:,0],b["world_pose"][:,0]
|
||||
mask=b["valid"].squeeze(-1).bool().cumprod(1).bool()
|
||||
# Include the first terminal transition for collision classification only.
|
||||
prefix=torch.cat((torch.ones_like(mask[:,:1]),mask[:,:-1]),1)
|
||||
total=next(self.parameters()).sum()*0.
|
||||
contact_total=total.detach().clone(); force_total=total.detach().clone()
|
||||
collision_total=total.clone(); collision_count=total.detach().clone()
|
||||
positives=total.detach().clone(); predicted=total.detach().clone()
|
||||
for t in range(b["action"].shape[1]):
|
||||
keep=prefix[:,t] if self.predict_collisions else mask[:,t]
|
||||
if not keep.any(): break
|
||||
context=torch.cat((b["context"][keep,0],field.scan(state[keep],pose[keep])),-1)
|
||||
pred,delta,contact,collision=self.forward_all(state[keep],b["action"][keep,t],b["previous"][keep,t],context)
|
||||
current=mask[keep,t]
|
||||
# Never regress the reset state as the endpoint of a collision.
|
||||
if current.any():
|
||||
error=torch.cat((pred[current]-b["next"][keep,t][current],delta[current]-b["pose"][keep,t][current]),-1)
|
||||
total=total+(error/self.output_norm.scale).square().mean(-1).sum()
|
||||
if self.predict_contacts:
|
||||
target=b["next_contacts"][keep,t][current]
|
||||
finite=torch.isfinite(target).all(-1)
|
||||
# Missing/non-finite sensor labels must not poison dynamics gradients.
|
||||
prediction=contact[current][finite]; target=target[finite]
|
||||
classification=torch.nn.functional.binary_cross_entropy_with_logits(
|
||||
prediction[:,:4],target[:,:4],reduction='none').mean(-1)
|
||||
force=torch.nn.functional.smooth_l1_loss(prediction[:,4:],
|
||||
torch.asinh(target[:,4:]/100.),reduction='none').mean(-1)
|
||||
total=total+.1*(classification+force).sum()
|
||||
contact_total=contact_total+classification.detach().sum()
|
||||
force_total=force_total+force.detach().sum()
|
||||
if self.predict_collisions:
|
||||
label=b['collision'][keep,t]
|
||||
usable=b['collision_valid'][keep,t].squeeze(-1).bool() & torch.isfinite(label).all(-1)
|
||||
logits=collision[usable]; labels=label[usable]
|
||||
loss=torch.nn.functional.binary_cross_entropy_with_logits(logits,labels,reduction='sum')
|
||||
collision_total=collision_total+loss
|
||||
collision_count=collision_count+usable.sum()
|
||||
positives=positives+labels.detach().sum()
|
||||
predicted=predicted+logits.detach().sigmoid().sum()
|
||||
next_pose=advance_pose(pose[keep],delta)
|
||||
state=state.clone(); pose=pose.clone()
|
||||
state[keep]=pred; pose[keep]=next_pose
|
||||
count=mask.sum().clamp_min(1)
|
||||
if self.predict_contacts:
|
||||
self.last_contact_metrics={'contact_bce':contact_total/count,'contact_force_loss':force_total/count}
|
||||
result=total/count
|
||||
if self.predict_collisions:
|
||||
denom=collision_count.clamp_min(1)
|
||||
result=result+collision_total/denom
|
||||
self.last_contact_metrics.update(collision_bce=collision_total.detach()/denom,
|
||||
collision_label_rate=positives/denom,collision_prediction_mean=predicted/denom,
|
||||
collision_label_count=collision_count)
|
||||
return result
|
||||
|
||||
def checkpoint(self):
|
||||
return {"format":FORMAT,"hidden":self.hidden,"contacts":self.predict_contacts,
|
||||
"collisions":self.predict_collisions,"model":self.state_dict()}
|
||||
|
||||
@classmethod
|
||||
def from_checkpoint(cls, checkpoint, device="cpu"):
|
||||
if checkpoint.get("format") != FORMAT: raise ValueError("Expected terrain dynamics checkpoint")
|
||||
model=cls(checkpoint["hidden"],contacts=checkpoint.get("contacts",False),
|
||||
collisions=checkpoint.get("collisions",False)).to(device)
|
||||
model.load_state_dict(checkpoint["model"])
|
||||
return model
|
||||
@@ -0,0 +1,782 @@
|
||||
"""Online S10 model-based learning: collect, fit dynamics, imagine, optimize."""
|
||||
|
||||
import argparse
|
||||
import copy
|
||||
import json
|
||||
import math
|
||||
import time
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from zoneinfo import ZoneInfo
|
||||
|
||||
import torch
|
||||
|
||||
from .bridge import S10Bridge, pose_delta
|
||||
from .world_model import S10WorldModel
|
||||
from .policy import (Actor, Network, CRITIC_DIM, POLICY_FORMAT,
|
||||
checkpoint_wheel_vel, critic_input)
|
||||
from .replay import Replay
|
||||
from .imagine import frozen, rollout, make_imagined_step
|
||||
from .recovery import recover_invalid_envs
|
||||
from .upward import (TILT_STAGES_DEG, init_upward_state,
|
||||
upward_curriculum_stats, upward_video_tilt_limit)
|
||||
from s10_supermimic.tasks import get_task, list_tasks, task_names
|
||||
|
||||
|
||||
def _format_value(value):
|
||||
if value is None:
|
||||
return "n/a"
|
||||
if isinstance(value, float):
|
||||
return f"{value:+.4f}" if abs(value) >= 1e-3 else f"{value:.3e}"
|
||||
if isinstance(value, (list, tuple)):
|
||||
return "[" + ", ".join(_format_value(v) for v in value) + "]"
|
||||
return str(value)
|
||||
|
||||
|
||||
def format_training_log(row, args, reward_terms, termination_rates):
|
||||
"""IsaacLab-style multi-line training summary; the JSONL stays machine-readable."""
|
||||
def line(label, value):
|
||||
return f"{label:>42s}: {value}"
|
||||
total = row.get("total_transitions", 0)
|
||||
elapsed = max(row.get("elapsed_s", 0.0), 1e-9)
|
||||
lines = [
|
||||
"",
|
||||
"=" * 78,
|
||||
f"Learning iteration {row.get('step', 0)}/{args.iterations}",
|
||||
line("Total steps", f"{total:,}"),
|
||||
line("Steps per second", f"{total / elapsed:,.0f}"),
|
||||
line("Iteration time", f"{row.get('iteration_time_s', 0.0):.3f}s"),
|
||||
line("Collection time", f"{row.get('collection_time_s', 0.0):.3f}s"),
|
||||
line("Learning time", f"{row.get('learning_time_s', 0.0):.3f}s"),
|
||||
line("Mean reward", f"{row.get('reward', 0.0):+.4f}"),
|
||||
line("Mean episode length", f"{row.get('mean_episode_length', 0.0):.1f}"),
|
||||
line("Actor loss", _format_value(row.get("actor_loss"))),
|
||||
line("Critic loss", _format_value(row.get("critic_loss"))),
|
||||
line("World model loss", _format_value(row.get("wm_loss"))),
|
||||
line("Mean action noise", f"{args.action_noise:.3f}"),
|
||||
line("Skipped updates", row.get("skipped_updates", 0)),
|
||||
"-" * 78,
|
||||
]
|
||||
for name, value in reward_terms.items():
|
||||
lines.append(line(f"Episode_Reward/{name}", f"{value:+.4f}"))
|
||||
for key in ("terrain_level_mean", "terrain_level_frac",
|
||||
"tilt_level_mean", "tilt_level_q75", "tilt_level_frac",
|
||||
"tilt_limit_mean_deg", "recovery_window_rate",
|
||||
"traverse_window_rate", "recovery_window_ready"):
|
||||
if key in row:
|
||||
lines.append(line(f"Curriculum/{key}", _format_value(row[key])))
|
||||
for key in ("command_stage", "command_mode", "command_ranges"):
|
||||
if key in row:
|
||||
lines.append(line(f"Command/{key}", _format_value(row[key])))
|
||||
vmae = row.get("velocity_mae")
|
||||
if vmae:
|
||||
lines.append(line("Metrics/base_velocity/error_vel_xy",
|
||||
f"{(vmae[0] + vmae[1]) / 2:.4f}"))
|
||||
lines.append(line("Metrics/base_velocity/error_vel_yaw", f"{vmae[2]:.4f}"))
|
||||
for name, value in termination_rates.items():
|
||||
lines.append(line(f"Episode_Termination/{name}", f"{value:.4f}"))
|
||||
lines.append("=" * 78)
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def constants_from_bridge(bridge):
|
||||
"""Fixed gains/limits of this S10 task; no per-environment gain randomization."""
|
||||
d, dev = bridge.robot.data, bridge.env.device
|
||||
from .terrain_posture import env_posture_cost, LEG_SCALES
|
||||
from .upward import UPWARD_REWARD_CLAMPS, is_upward_task
|
||||
from s10_supermimic.robot.config import LEG_JOINT_NAMES
|
||||
upward = is_upward_task(bridge.env)
|
||||
revised = (bridge.env.cfg.rewards["leg_posture"].func is env_posture_cost
|
||||
or "wheel_obstacle_clearance" in bridge.env.cfg.rewards
|
||||
or upward)
|
||||
# Read the actual leg actuator gains so the imagined torque surrogate matches
|
||||
# v4 (kp=80) and the softer v7 (kp=20) without duplicating constants.
|
||||
robot_cfg = bridge.env.cfg.scene.entities["robot"]
|
||||
leg_act = next(
|
||||
(a for a in robot_cfg.articulation.actuators
|
||||
if set(LEG_JOINT_NAMES).issubset(set(getattr(a, "target_names_expr", ())))),
|
||||
None)
|
||||
leg_kp = float(getattr(leg_act, "stiffness", 80.0))
|
||||
leg_kd = float(getattr(leg_act, "damping", 2.0))
|
||||
leg_effort = float(getattr(leg_act, "effort_limit", 50.0))
|
||||
return {
|
||||
"terrain_posture_v2": revised,
|
||||
"nonwheel_collision": "illegal_nonwheel_contact" in bridge.env.cfg.terminations,
|
||||
"wheel_local_rewards": "wheel_obstacle_clearance" in bridge.env.cfg.rewards,
|
||||
"rolling_rewards": "unnecessary_wheel_lift" in bridge.env.cfg.rewards,
|
||||
"obstacle_progress": "obstacle_progress" in bridge.env.cfg.rewards,
|
||||
"wheel_body_intrusion": "wheel_body_intrusion" in bridge.env.cfg.rewards,
|
||||
"pair_body_intrusion": "pair_body_intrusion" in bridge.env.cfg.rewards,
|
||||
"posture_relief_floor": bridge.env.cfg.rewards["leg_posture"].params.get(
|
||||
"relief_floor", 0.125),
|
||||
"upright_mode": "upward" if upward else "exp",
|
||||
"bad_orientation_penalty": upward,
|
||||
"reward_clamps": dict(UPWARD_REWARD_CLAMPS) if upward else {},
|
||||
"wheel_action_scale": bridge.env.cfg.actions["wheel_vel"].scale,
|
||||
"leg_action_scale": torch.tensor(LEG_SCALES if revised else (.25,) * 12, device=dev),
|
||||
"terrain_rewards": hasattr(bridge, "field"),
|
||||
"upright_std": bridge.env.cfg.rewards["upright"].params.get("std", .6),
|
||||
"default_q": d.default_joint_pos[0, bridge.legs].detach().clone(),
|
||||
"limits": d.soft_joint_pos_limits[0, bridge.legs].detach().clone(),
|
||||
# S10 leg actuators: v4 kp=80/kd=2, v7 kp=20/kd=2; effort=50 N·m.
|
||||
# TODO(s10): the -2e-5 joint_torques_l2 weight was inherited from the
|
||||
# initial recipe; retune if S10 training shows torque suppression or chatter.
|
||||
"kp": torch.full((12,), leg_kp, device=dev),
|
||||
"kd": torch.full((12,), leg_kd, device=dev),
|
||||
"effort": torch.full((12,), leg_effort, device=dev),
|
||||
# S10 wheel actuator: velocity damping 0.8, effort 14 N·m, 35 rad/s clip.
|
||||
"wheel_kd": 0.8,
|
||||
"wheel_effort": 14.0,
|
||||
"wheel_velocity_limit": 35.0,
|
||||
}
|
||||
|
||||
|
||||
def actor_learning_rate(args, iteration):
|
||||
age = max(0, iteration - args.actor_schedule_start)
|
||||
progress = min(1., max(0, age - args.actor_hold_iterations) / max(1, args.lr_decay_iterations))
|
||||
decay = args.lr_final_frac + (1.-args.lr_final_frac)*.5*(1.+math.cos(math.pi*progress))
|
||||
ramp = min(1., (age+1) / max(1, args.actor_ramp_iterations))
|
||||
return args.lr * args.actor_lr_scale * decay * ramp
|
||||
|
||||
|
||||
def update(optimizer, loss, parameters, element_clip=None):
|
||||
parameters = list(parameters)
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
if not torch.isfinite(loss):
|
||||
print("RECOVERY skipped non-finite loss; optimizer unchanged", flush=True)
|
||||
return False
|
||||
loss.backward()
|
||||
norm = torch.nn.utils.clip_grad_norm_(parameters, 1., error_if_nonfinite=False)
|
||||
if not torch.isfinite(norm):
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
print("RECOVERY skipped non-finite gradient; optimizer unchanged", flush=True)
|
||||
return False
|
||||
if element_clip is not None:
|
||||
torch.nn.utils.clip_grad_value_(parameters, element_clip)
|
||||
optimizer.step()
|
||||
return True
|
||||
|
||||
|
||||
def record_policy_video(actor, args, output, iteration, env_order,
|
||||
upward_tilt_limit=None, command_ranges=None):
|
||||
"""Record a short MP4 with the current actor on a fresh single-env scene.
|
||||
|
||||
This mirrors Isaac Lab's ``--video``/``--video_length``/``--video_interval``
|
||||
workflow. A separate one-env scene is used so recording does not perturb the
|
||||
training environments. Recording is skipped if the optional mediapy/EGL
|
||||
dependencies or the extra scene cannot be created.
|
||||
"""
|
||||
import random
|
||||
import numpy as np
|
||||
|
||||
state = (
|
||||
random.getstate(),
|
||||
np.random.get_state(),
|
||||
torch.get_rng_state(),
|
||||
torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None,
|
||||
)
|
||||
video_env = None
|
||||
try:
|
||||
import mediapy
|
||||
from mjlab.envs import ManagerBasedRlEnv
|
||||
from .bridge import S10Bridge
|
||||
|
||||
task = get_task(args.task) if getattr(args, "task", None) else None
|
||||
if task is not None:
|
||||
if task.actor.wheel_vel != actor.actor_wheel_vel:
|
||||
raise ValueError("Task Actor schema differs from the training Actor")
|
||||
cfg = task.build_env_cfg(play=True, history=args.history, seed=args.seed)
|
||||
elif args.terrain:
|
||||
from s10_supermimic.env.terrain import s10_velocity_terrain_env_cfg
|
||||
cfg = s10_velocity_terrain_env_cfg(
|
||||
play=True, history=args.history, seed=args.seed,
|
||||
revision=args.terrain_revision, actor_wheel_vel=actor.actor_wheel_vel)
|
||||
else:
|
||||
from s10_supermimic.env.flat import s10_velocity_flat_env_cfg
|
||||
cfg = s10_velocity_flat_env_cfg(
|
||||
play=True, history=args.history,
|
||||
actor_wheel_vel=actor.actor_wheel_vel)
|
||||
cfg.scene.num_envs = 1
|
||||
cfg.rewards = {}
|
||||
cfg.metrics = {}
|
||||
cfg.observations.pop("critic", None)
|
||||
cfg.episode_length_s = 1e9
|
||||
cfg.viewer.width = args.video_width
|
||||
cfg.viewer.height = args.video_height
|
||||
cfg.viewer.distance = 2.5
|
||||
cfg.viewer.elevation = -20.
|
||||
cfg.viewer.azimuth = 140.
|
||||
video_env = ManagerBasedRlEnv(
|
||||
cfg=cfg, device=args.device, render_mode="rgb_array")
|
||||
if command_ranges is not None:
|
||||
ranges = video_env.command_manager.get_term("twist").cfg.ranges
|
||||
ranges.lin_vel_x, ranges.lin_vel_y, ranges.ang_vel_z = command_ranges
|
||||
bridge = S10Bridge(video_env)
|
||||
if list(bridge.to_env) != list(env_order):
|
||||
raise ValueError("Video env action order differs from the training env")
|
||||
if upward_tilt_limit is not None:
|
||||
video_env._upward_tilt_limits = torch.full(
|
||||
(video_env.num_envs,), float(upward_tilt_limit), device=video_env.device)
|
||||
obs, _ = video_env.reset()
|
||||
path = output / "videos" / f"iteration_{iteration:06d}.mp4"
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
was_training = actor.training
|
||||
actor.eval()
|
||||
with torch.no_grad(), mediapy.VideoWriter(
|
||||
str(path), shape=(args.video_height, args.video_width),
|
||||
fps=args.video_fps) as writer:
|
||||
for _ in range(args.video_length):
|
||||
action = actor(obs["actor"])
|
||||
obs, _, terminated, truncated, _ = video_env.step(
|
||||
bridge.env_action(action))
|
||||
writer.add_image(video_env.render())
|
||||
if bool(terminated.any() or truncated.any()):
|
||||
obs, _ = video_env.reset()
|
||||
actor.train(was_training)
|
||||
print(f"S10 video saved: {path}", flush=True)
|
||||
finally:
|
||||
if video_env is not None:
|
||||
video_env.close()
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
random.setstate(state[0])
|
||||
np.random.set_state(state[1])
|
||||
torch.set_rng_state(state[2])
|
||||
if state[3] is not None:
|
||||
torch.cuda.set_rng_state_all(state[3])
|
||||
|
||||
|
||||
class Learner:
|
||||
"""Pure-Torch learner; also usable in small CPU tests without mjlab."""
|
||||
def __init__(self, args, field=None):
|
||||
self.args, self.field = args, field
|
||||
self.format = POLICY_FORMAT
|
||||
self.model_class = S10WorldModel
|
||||
extra = 0
|
||||
if field is not None:
|
||||
from .terrain_world_model import TerrainWorldModel, POLICY_FORMAT as TERRAIN_FORMAT
|
||||
from .terrain import SCAN_DIM
|
||||
self.model_class, self.format, extra = TerrainWorldModel, TERRAIN_FORMAT, SCAN_DIM
|
||||
self.model = self.model_class(tuple(args.wm_hidden), **(
|
||||
{"contacts":args.terrain_revision >= 4,
|
||||
"collisions":args.terrain_revision >= 7} if field is not None else {})).to(args.device)
|
||||
self.actor = Actor(args.history, tuple(args.actor_hidden),
|
||||
args.actor_wheel_vel).to(args.device)
|
||||
self.critic = Network(CRITIC_DIM + extra, 1, tuple(args.critic_hidden)).to(args.device)
|
||||
self.target = copy.deepcopy(self.critic)
|
||||
self.optimizers = [torch.optim.Adam(module.parameters(), lr=args.lr)
|
||||
for module in (self.model, self.actor, self.critic)]
|
||||
self.policy_updates = 0
|
||||
self.skipped_updates = 0
|
||||
self.terrain_state = None
|
||||
self.curriculum_state = None
|
||||
self.env_steps = 0
|
||||
self.imagined_step = None
|
||||
|
||||
def critic_input(self, state, command, context, previous, pose):
|
||||
x = critic_input(state, command, context, previous)
|
||||
return x if self.field is None else torch.cat((x, self.field.scan(state, pose)), -1)
|
||||
|
||||
def fit(self, b):
|
||||
if self.field is not None:
|
||||
loss = self.model.terrain_loss(b, self.field)
|
||||
if not update(self.optimizers[0], loss, self.model.parameters()):
|
||||
self.skipped_updates += 1
|
||||
return None
|
||||
return loss.item()
|
||||
loss = self.model.loss(b["state"][:, 0], b["action"], b["previous"], b["context"][:, 0],
|
||||
b["next"], b["pose"], b["valid"].squeeze(-1))
|
||||
if not update(self.optimizers[0], loss, self.model.parameters()):
|
||||
self.skipped_updates += 1
|
||||
return None
|
||||
return loss.item()
|
||||
|
||||
def improve(self, seed, constants, weights, env_order, dt):
|
||||
args = self.args
|
||||
if args.compile_imagination and self.imagined_step is None:
|
||||
step = make_imagined_step(self.model,self.actor,constants,weights,env_order,
|
||||
dt=dt,field=self.field)
|
||||
self.imagined_step = torch.compile(step, dynamic=False, fullgraph=True)
|
||||
with frozen(self.model, self.target):
|
||||
actor_loss, inputs, targets = rollout(
|
||||
self.model, self.actor, self.target, seed, constants, weights, env_order,
|
||||
horizon=args.horizon, dt=dt, gamma=args.gamma, lam=args.lam, field=self.field,
|
||||
step_fn=self.imagined_step)
|
||||
if not update(self.optimizers[1], actor_loss, self.actor.parameters(), args.actor_grad_clip_element):
|
||||
self.skipped_updates += 1
|
||||
return None, None
|
||||
x, y = inputs.flatten(0, 1), targets.flatten()
|
||||
critic_losses = []
|
||||
for _ in range(args.critic_epochs):
|
||||
order = torch.randperm(len(x),device=x.device)
|
||||
for indices in order.tensor_split(min(args.critic_minibatches,len(x))):
|
||||
value = self.critic(x[indices]).squeeze(-1)
|
||||
critic_loss = torch.nn.functional.mse_loss(value, y[indices])
|
||||
if not update(self.optimizers[2], critic_loss, self.critic.parameters()):
|
||||
self.skipped_updates += 1
|
||||
return actor_loss.item(), None
|
||||
critic_losses.append(critic_loss.detach())
|
||||
with torch.no_grad():
|
||||
for target, source in zip(self.target.parameters(), self.critic.parameters()):
|
||||
target.lerp_(source, 0.02)
|
||||
# Running normalization is a coordinate system, not an EMA weight.
|
||||
self.target.norm.load_state_dict(self.critic.norm.state_dict())
|
||||
self.policy_updates += 1
|
||||
return actor_loss.item(), torch.stack(critic_losses).mean().item()
|
||||
|
||||
def save(self, path, step, env_order, dt):
|
||||
checkpoint = {
|
||||
"training_algorithm": "s10-full-return-events-v1",
|
||||
"format": self.format, "step": step, "history": self.actor.history,
|
||||
"actor_wheel_vel": self.actor.actor_wheel_vel,
|
||||
"actor_hidden": self.actor.hidden, "actor": self.actor.state_dict(),
|
||||
"world_model": self.model.checkpoint(), "critic": self.critic.state_dict(),
|
||||
"target_critic": self.target.state_dict(), "config": vars(self.args),
|
||||
"optimizers": [o.state_dict() for o in self.optimizers],
|
||||
"terrain_state": self.terrain_state,
|
||||
"curriculum_state": self.curriculum_state,
|
||||
"env_steps": self.env_steps,
|
||||
"policy_updates": self.policy_updates, "env_action_order": list(env_order), "dt": dt,
|
||||
}
|
||||
path = Path(path)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
tmp = path.with_suffix(".tmp")
|
||||
torch.save(checkpoint, tmp)
|
||||
tmp.replace(path)
|
||||
|
||||
def load(self, path, env_order, dt):
|
||||
ckpt = torch.load(path, map_location=self.args.device, weights_only=True)
|
||||
if ckpt.get("format") != self.format:
|
||||
raise ValueError("Resume requires a S10 policy checkpoint")
|
||||
if ckpt["env_action_order"] != list(env_order) or abs(ckpt["dt"] - dt) > 1e-8:
|
||||
raise ValueError("Checkpoint action order/control timestep differs from this environment")
|
||||
if ckpt["history"] != self.args.history:
|
||||
raise ValueError("Checkpoint observation history differs")
|
||||
if self.field is not None and ckpt["config"]["seed"] != self.args.seed:
|
||||
raise ValueError("Terrain resume requires the same terrain seed")
|
||||
if ckpt["config"].get("terrain_revision", 1) != getattr(self.args, "terrain_revision", 1):
|
||||
raise ValueError("Terrain revision/action scale differs; start a new run for a different revision")
|
||||
if checkpoint_wheel_vel(ckpt) != self.actor.actor_wheel_vel:
|
||||
raise ValueError("Actor observation schema differs; choose the matching task or start a new run")
|
||||
self.actor.load_state_dict(ckpt["actor"])
|
||||
self.model.load_state_dict(self.model_class.from_checkpoint(ckpt["world_model"]).state_dict())
|
||||
self.critic.load_state_dict(ckpt["critic"])
|
||||
self.target.load_state_dict(ckpt["target_critic"])
|
||||
for optimizer, data in zip(self.optimizers, ckpt["optimizers"]):
|
||||
optimizer.load_state_dict(data)
|
||||
self.policy_updates = ckpt["policy_updates"]
|
||||
self.terrain_state = ckpt.get("terrain_state")
|
||||
self.curriculum_state = ckpt.get("curriculum_state")
|
||||
self.env_steps = ckpt.get("env_steps", ckpt["step"] * ckpt["config"].get("num_steps_per_env", 1))
|
||||
return ckpt["step"]
|
||||
|
||||
|
||||
def parser():
|
||||
p = argparse.ArgumentParser(description=__doc__)
|
||||
p.add_argument("--task", choices=task_names(), default=None,
|
||||
help="Select a registered S10 task; its defaults are applied before CLI overrides")
|
||||
p.add_argument("--list-tasks", action="store_true", help="List registered tasks and exit")
|
||||
for name, default in (("num-envs", 64), ("iterations", 10000), ("buffer-len", 512),
|
||||
("history", 5), ("wm-window", 8), ("wm-batch", 128),
|
||||
("wm-updates", 2), ("imag-batch", 128), ("horizon", 16),
|
||||
("warmup", 300), ("critic-epochs", 5), ("save-interval", 1000),
|
||||
("log-interval", 50), ("seed", 0), ("num-steps-per-env", 1), ("critic-minibatches", 4),
|
||||
("archive-size", 100000), ("actor-ramp-iterations", 1000),
|
||||
("lr-decay-iterations", 30000), ("actor-hold-iterations", 0),
|
||||
("actor-schedule-start", 0)):
|
||||
p.add_argument("--" + name, type=int, default=default)
|
||||
for name, default in (("lr", 2e-4), ("action-noise", 0.15), ("gamma", 0.99), ("lam", 0.95)):
|
||||
p.add_argument("--" + name, type=float, default=default)
|
||||
p.add_argument("--compile-imagination", action="store_true")
|
||||
p.add_argument("--no-compile-imagination", dest="compile_imagination",
|
||||
action="store_false", default=False)
|
||||
p.add_argument("--prioritized-replay", action="store_true")
|
||||
p.add_argument("--actor-lr-scale", type=float, default=.3)
|
||||
p.add_argument("--lr-final-frac", type=float, default=.03)
|
||||
p.add_argument("--actor-grad-clip-element", type=float, default=.1)
|
||||
p.add_argument("--wm-hidden", type=int, nargs="+", default=[512, 512, 512])
|
||||
p.add_argument("--actor-hidden", type=int, nargs="+", default=[256, 128, 64])
|
||||
p.add_argument("--critic-hidden", type=int, nargs="+", default=[256, 128, 64])
|
||||
p.add_argument("--device", default="cuda:0")
|
||||
p.add_argument("--terrain", action="store_true", help="Blind static rough-terrain task")
|
||||
p.add_argument("--terrain-revision", type=int, choices=(1, 2, 3, 4, 5, 6, 7), default=1,
|
||||
help="2: larger leg actions; 3: bidirectional commands; 4: rolling rewards; 5: performance-gated stair curriculum; 6: wheel-local clearance; 7: safe v7 (53-value Actor, time-ramped 1.5/0.8/1.2 commands, ±45° init, wheel/body intrusion, collision); new run required")
|
||||
p.add_argument("--command-curriculum", action="store_true", help="Progressively widen velocity ranges")
|
||||
p.add_argument("--command-curriculum-start", type=int, default=0, help="Learner iteration where the velocity schedule starts; preserve on resume")
|
||||
p.add_argument("--init-actor", default="", help="Initialize only the actor from a flat policy; dynamics/critic start fresh")
|
||||
p.add_argument("--resume", default="")
|
||||
p.add_argument("--output", default="")
|
||||
p.add_argument("--video", action="store_true", help="Record periodic MP4 videos with the current actor")
|
||||
p.add_argument("--video-length", type=int, default=500, help="Frames per video")
|
||||
p.add_argument("--video-interval", type=int, default=15000, help="Record every N learner iterations")
|
||||
p.add_argument("--video-fps", type=int, default=50, help="Video frame rate")
|
||||
p.add_argument("--video-width", type=int, default=640)
|
||||
p.add_argument("--video-height", type=int, default=480)
|
||||
return p
|
||||
|
||||
|
||||
def main():
|
||||
p = parser()
|
||||
pre, _ = p.parse_known_args()
|
||||
if pre.list_tasks:
|
||||
for spec in list_tasks():
|
||||
print(f"{spec.name}: {spec.description}")
|
||||
return
|
||||
task = get_task(pre.task) if pre.task else None
|
||||
if task is not None:
|
||||
p.set_defaults(**task.defaults)
|
||||
args = p.parse_args()
|
||||
if task is not None and args.terrain_revision != task.terrain_revision:
|
||||
p.error(f"Task {task.name} requires terrain-revision {task.terrain_revision}")
|
||||
if task is not None:
|
||||
args.actor_wheel_vel = task.actor.wheel_vel
|
||||
elif args.resume:
|
||||
resumed = torch.load(args.resume, map_location="cpu", weights_only=True)
|
||||
args.actor_wheel_vel = checkpoint_wheel_vel(resumed)
|
||||
else:
|
||||
args.actor_wheel_vel = args.terrain_revision <= 4 or args.terrain_revision >= 7
|
||||
positive = (args.archive_size, args.critic_minibatches, args.lr_decay_iterations, args.num_steps_per_env, args.num_envs, args.iterations, args.buffer_len, args.history, args.wm_window,
|
||||
args.wm_batch, args.wm_updates, args.imag_batch, args.horizon, args.critic_epochs,
|
||||
args.save_interval, args.log_interval, *args.wm_hidden, *args.actor_hidden,
|
||||
*args.critic_hidden)
|
||||
if min(positive) < 1 or args.warmup < 0 or args.buffer_len < args.wm_window:
|
||||
p.error("Positive sizes required; buffer-len must be >= wm-window; warmup >= 0")
|
||||
if not (0 < args.gamma <= 1 and 0 <= args.lam <= 1 and args.lr > 0 and args.action_noise >= 0):
|
||||
p.error("Invalid optimizer, noise or return parameters")
|
||||
if args.actor_lr_scale <= 0 or not 0 < args.lr_final_frac <= 1 or args.actor_ramp_iterations < 0 or args.actor_grad_clip_element <= 0:
|
||||
p.error("Invalid actor learning-rate/gradient protection")
|
||||
if args.actor_hold_iterations < 0:
|
||||
p.error("actor-hold-iterations must be nonnegative")
|
||||
if args.terrain_revision >= 5 and args.archive_size < 2:
|
||||
p.error("terrain v5+ requires archive-size >=2")
|
||||
if 5 <= args.terrain_revision < 7 and not args.command_curriculum:
|
||||
p.error("terrain v5/v6 requires command-curriculum; v7 uses a time-based cap")
|
||||
if args.video and (args.video_length < 1 or args.video_interval < 1
|
||||
or not 1 <= args.video_fps <= 50
|
||||
or args.video_width < 64 or args.video_height < 64
|
||||
or args.video_width % 2 or args.video_height % 2):
|
||||
p.error("Video requires length/interval >=1, fps 1..50, even width/height >=64")
|
||||
if args.prioritized_replay and (not args.terrain or args.terrain_revision < 4):
|
||||
p.error("prioritized-replay requires measured contacts from terrain revision >=4")
|
||||
from mjlab.envs import ManagerBasedRlEnv
|
||||
from s10_supermimic.env.flat import s10_velocity_flat_env_cfg
|
||||
|
||||
torch.manual_seed(args.seed)
|
||||
if args.resume and args.init_actor:
|
||||
p.error("Choose resume or init-actor")
|
||||
if args.terrain_revision >= 2 and (not args.terrain or args.init_actor):
|
||||
p.error("terrain-revision >=2 requires --terrain and cannot initialize a flat actor")
|
||||
if task is not None:
|
||||
if task.actor.wheel_vel != args.actor_wheel_vel:
|
||||
p.error("Task Actor schema differs from the requested Actor schema")
|
||||
cfg = task.build_env_cfg(play=False, history=args.history, seed=args.seed)
|
||||
elif args.terrain:
|
||||
from s10_supermimic.env.terrain import s10_velocity_terrain_env_cfg
|
||||
cfg = s10_velocity_terrain_env_cfg(
|
||||
history=args.history, seed=args.seed, revision=args.terrain_revision,
|
||||
actor_wheel_vel=args.actor_wheel_vel)
|
||||
else:
|
||||
cfg = s10_velocity_flat_env_cfg(
|
||||
history=args.history, actor_wheel_vel=args.actor_wheel_vel)
|
||||
if args.command_curriculum:
|
||||
if not args.terrain or args.terrain_revision < 4:
|
||||
p.error("command-curriculum requires terrain revision >=4")
|
||||
from .command_curriculum import set_command_stage, stages_for
|
||||
set_command_stage(cfg.commands["twist"], 0,
|
||||
stages=stages_for(args.terrain_revision))
|
||||
cfg.seed = args.seed
|
||||
cfg.scene.num_envs = args.num_envs
|
||||
env = ManagerBasedRlEnv(cfg=cfg, device=args.device)
|
||||
try:
|
||||
if args.task == "s10-v4-upward":
|
||||
init_upward_state(env)
|
||||
obs, _ = env.reset()
|
||||
field = None
|
||||
if args.terrain:
|
||||
from .terrain import TerrainField, TerrainBridge
|
||||
field = TerrainField.from_env(env)
|
||||
bridge = TerrainBridge(env, field)
|
||||
env._s10_posture_bridge = bridge
|
||||
print(f"Shared terrain atlas: {field.h.numel()*4/2**20:.1f} MiB", flush=True)
|
||||
else:
|
||||
bridge = S10Bridge(env)
|
||||
learner = Learner(args, field)
|
||||
if args.init_actor:
|
||||
ckpt = torch.load(args.init_actor, map_location=args.device, weights_only=True)
|
||||
if ckpt.get("format") != POLICY_FORMAT or ckpt["history"] != args.history:
|
||||
raise ValueError("init-actor requires a matching flat S10 actor")
|
||||
if ckpt["env_action_order"] != list(bridge.to_env) or abs(ckpt["dt"]-env.step_dt)>1e-8:
|
||||
raise ValueError("Warm-start action contract differs")
|
||||
if checkpoint_wheel_vel(ckpt) != learner.actor.actor_wheel_vel:
|
||||
raise ValueError("Warm-start actor observation schema differs")
|
||||
learner.actor.load_state_dict(ckpt["actor"])
|
||||
constants = constants_from_bridge(bridge)
|
||||
weights = {name: term.weight for name, term in cfg.rewards.items()}
|
||||
contact_targets = args.terrain and args.terrain_revision >= 4
|
||||
replay = Replay(args.buffer_len, args.num_envs, args.device,
|
||||
learner.actor.frame_dim * args.history,
|
||||
terrain=args.terrain, contacts=contact_targets,
|
||||
courses=args.terrain_revision >= 5,
|
||||
collisions=args.terrain_revision >= 7,
|
||||
pairs=args.terrain_revision >= 7)
|
||||
start = learner.load(args.resume, bridge.to_env, env.step_dt) if args.resume else 0
|
||||
command_stage = 0
|
||||
if args.command_curriculum:
|
||||
saved_stage = None
|
||||
if args.terrain_revision >= 5:
|
||||
saved_stage = (int(learner.curriculum_state['stage'])
|
||||
if learner.curriculum_state else 0)
|
||||
command_stage = set_command_stage(
|
||||
env.command_manager.get_term("twist").cfg,
|
||||
start - args.command_curriculum_start,
|
||||
stage=saved_stage, stages=stages_for(args.terrain_revision))
|
||||
if args.terrain and learner.terrain_state is not None:
|
||||
terrain = env.scene.terrain
|
||||
levels = learner.terrain_state["levels"].to(args.device)
|
||||
types = learner.terrain_state["types"].to(args.device)
|
||||
if levels.shape != terrain.terrain_levels.shape:
|
||||
raise ValueError("Terrain resume requires the same num-envs")
|
||||
terrain.terrain_levels.copy_(levels)
|
||||
terrain.terrain_types.copy_(types)
|
||||
terrain.env_origins.copy_(terrain.terrain_origins[levels, types])
|
||||
obs, _ = env.reset()
|
||||
if args.terrain_revision >= 5:
|
||||
from .stair_curriculum import buffers, next_stage, summary
|
||||
stats = buffers(env)
|
||||
saved = learner.curriculum_state
|
||||
velocity_ema = saved['velocity_ema'] if saved else [1., 1., 1.]
|
||||
if saved:
|
||||
stats.copy_(saved['stats'].to(env.device))
|
||||
if args.terrain_revision >= 5:
|
||||
command_stage = int(saved['stage']) if saved else 0
|
||||
if args.command_curriculum:
|
||||
set_command_stage(env.command_manager.get_term('twist').cfg, 0,
|
||||
stage=command_stage,
|
||||
stages=stages_for(args.terrain_revision))
|
||||
if start >= args.iterations:
|
||||
raise ValueError("iterations must exceed the resumed checkpoint step")
|
||||
if obs["actor"].shape[-1] != learner.actor.frame_dim * args.history:
|
||||
raise ValueError("S10 actor observation dimension mismatch")
|
||||
output = Path(args.output or "logs/s10/" + datetime.now(ZoneInfo("Asia/Shanghai")).strftime("%Y-%m-%d_%H-%M-%S"))
|
||||
output.mkdir(parents=True, exist_ok=True)
|
||||
previous = torch.zeros(args.num_envs, 16, device=args.device)
|
||||
ou = torch.zeros_like(previous)
|
||||
totals = torch.zeros(6, device=args.device)
|
||||
reward_totals = torch.zeros(len(env.reward_manager.active_terms), device=args.device)
|
||||
interval_steps = 0
|
||||
nonwheel_terminations = 0
|
||||
invalid_env_resets = 0
|
||||
episode_length_sum = 0.0
|
||||
episode_count = 0
|
||||
termination_counts = {name: 0 for name in env.termination_manager.active_terms}
|
||||
collection_time_sum = 0.0
|
||||
learning_time_sum = 0.0
|
||||
began = time.monotonic()
|
||||
print(f"S10 replay storage: {sum(t.numel() * t.element_size() for t in replay.data.values()) / 2**20:.1f} MiB", flush=True)
|
||||
if args.video:
|
||||
try:
|
||||
ranges = env.command_manager.get_term("twist").cfg.ranges
|
||||
record_policy_video(
|
||||
learner.actor, args, output, start, bridge.to_env,
|
||||
upward_tilt_limit=(upward_video_tilt_limit(env)
|
||||
if args.task == "s10-v4-upward" else None),
|
||||
command_ranges=(ranges.lin_vel_x, ranges.lin_vel_y, ranges.ang_vel_z))
|
||||
except Exception as exc:
|
||||
print(f"RECOVERY skipped initial video: {exc}", flush=True)
|
||||
with (output / "metrics.jsonl").open("a") as log:
|
||||
for step in range(start, args.iterations):
|
||||
iter_start = time.monotonic()
|
||||
for group in learner.optimizers[1].param_groups:
|
||||
group["lr"] = actor_learning_rate(args, step)
|
||||
if args.command_curriculum and args.terrain_revision < 5:
|
||||
command_stage = set_command_stage(
|
||||
env.command_manager.get_term("twist").cfg,
|
||||
step - args.command_curriculum_start,
|
||||
stages=stages_for(args.terrain_revision))
|
||||
for _ in range(args.num_steps_per_env):
|
||||
with torch.no_grad():
|
||||
history = obs["actor"].clone()
|
||||
state, pose, context = bridge.state(), bridge.pose(), bridge.context()
|
||||
command = env.command_manager.get_command("twist").clone()
|
||||
bias = bridge.robot.data.encoder_bias[:, bridge.legs].clone()
|
||||
learner.actor.norm.update(history)
|
||||
learner.critic.norm.update(learner.critic_input(state, command, context, previous, pose))
|
||||
learner.target.norm.load_state_dict(learner.critic.norm.state_dict())
|
||||
ou = 0.95 * ou + (1 - 0.95**2)**0.5 * torch.randn_like(ou)
|
||||
action = (learner.actor(history) + args.action_noise * ou).clamp(-1., 1.)
|
||||
if args.terrain_revision >= 5:
|
||||
course_types = env.scene.terrain.terrain_types.clone()
|
||||
course_levels = env.scene.terrain.terrain_levels.clone()
|
||||
course_origins = env.scene.env_origins.clone()
|
||||
prev_episode_length = env.episode_length_buf.clone()
|
||||
obs, real_reward, terminated, truncated, _ = env.step(bridge.env_action(action))
|
||||
terminated, truncated = terminated.clone(), truncated.clone()
|
||||
obs, bad = recover_invalid_envs(env, obs, real_reward)
|
||||
invalid_env_resets += int(bad.sum())
|
||||
valid = ~(terminated | truncated | bad)
|
||||
done = (terminated | truncated) & ~bad
|
||||
if done.any():
|
||||
episode_length_sum += float(prev_episode_length[done].sum().item())
|
||||
episode_count += int(done.sum())
|
||||
for name in termination_counts:
|
||||
termination_counts[name] += int(
|
||||
(env.termination_manager.get_term(name) & done).sum().item())
|
||||
real_reward = torch.where(bad, torch.zeros_like(real_reward), real_reward)
|
||||
nxt, delta = bridge.state(), pose_delta(pose, bridge.pose())
|
||||
measured = torch.cat((nxt[:, 4:6], nxt[:, 9:10]), -1)
|
||||
reward_totals += env.reward_manager._step_reward[~bad].sum(0)
|
||||
totals[0] += real_reward.sum()
|
||||
totals[1] += terminated.sum()
|
||||
totals[2] += valid.sum()
|
||||
totals[3:] += (measured[valid] - command[valid]).abs().sum(0)
|
||||
interval_steps += 1
|
||||
course_data = {}
|
||||
if args.terrain_revision >= 7:
|
||||
label_valid = env._nonwheel_label_valid & ~bad & torch.isfinite(state).all(-1)
|
||||
course_data.update(collision=env._nonwheel_hit[:,None].float(),
|
||||
collision_valid=label_valid[:,None].float())
|
||||
nonwheel_terminations += int((env._nonwheel_hit & label_valid).sum())
|
||||
if args.terrain_revision >= 5:
|
||||
new_pose = bridge.pose()
|
||||
ground_rise = field.sample(new_pose[:,:2]) - field.sample(pose[:,:2])
|
||||
forward_motion = (delta[:,:2] * command[:,:2]).sum(-1) > 0
|
||||
offset = bridge.robot.data.root_link_pos_w - course_origins
|
||||
from .stair_curriculum import stair_pass
|
||||
reached = stair_pass(offset, bridge.robot.data.projected_gravity_b,
|
||||
course_levels, course_types == 1)
|
||||
crossing = (pose[:,0] - course_origins[:,0]).abs() <= 2.8
|
||||
useful = (ground_rise > .005) & forward_motion & valid
|
||||
course_data['course'] = torch.stack((course_types.float(), course_levels.float(),
|
||||
useful.float(), (reached & crossing & valid).float()), -1)
|
||||
replay.add(state=state, action=action, previous=previous, context=context,
|
||||
next=nxt, pose=delta, valid=valid[:, None], history=history,
|
||||
command=command, bias=bias, **({"world_pose": pose} if args.terrain else {}),
|
||||
**({"next_contacts":bridge.contacts()} if contact_targets else {}),
|
||||
**({"pair_timer": env._s10_pair_timer} if args.terrain_revision >= 7 else {}),
|
||||
**course_data)
|
||||
if valid.any():
|
||||
model_context = context if field is None else torch.cat((context, field.scan(state, pose)), -1)
|
||||
learner.model.update_stats(state[valid], action[valid], previous[valid], model_context[valid],
|
||||
nxt[valid], delta[valid])
|
||||
previous = action.clone()
|
||||
previous[~valid], ou[~valid] = 0., 0.
|
||||
learner.env_steps += 1
|
||||
collect_end = time.monotonic()
|
||||
collection_time_sum += collect_end - iter_start
|
||||
wm_loss = actor_loss = critic_loss = None
|
||||
if replay.size >= args.wm_window:
|
||||
for _ in range(args.wm_updates):
|
||||
batch = (replay.sample_world_model(args.wm_batch, args.wm_window, capacity=args.archive_size)
|
||||
if args.prioritized_replay else replay.sample(args.wm_batch, args.wm_window))
|
||||
wm_loss = learner.fit(batch)
|
||||
# Replay is not checkpointed; refit it during warmup after every resume.
|
||||
if step - start >= args.warmup:
|
||||
seed = {k: v[:, 0] for k, v in replay.sample(args.imag_batch, 1).items()}
|
||||
actor_loss, critic_loss = learner.improve(seed, constants, weights, bridge.to_env, env.step_dt)
|
||||
learning_time_sum += time.monotonic() - collect_end
|
||||
if (step + 1) % args.log_interval == 0 or step + 1 == args.iterations:
|
||||
mean_episode_length = episode_length_sum / max(episode_count, 1)
|
||||
# RewardManager._step_reward is raw*weight; the env reward also
|
||||
# multiplies by dt. Scale before turning rates into episode sums.
|
||||
weighted_terms = (reward_totals / max(interval_steps * args.num_envs, 1)).cpu()
|
||||
terms_per_step = weighted_terms * env.step_dt
|
||||
episode_reward_terms = dict(zip(
|
||||
env.reward_manager.active_terms,
|
||||
(terms_per_step * mean_episode_length).tolist()))
|
||||
termination_rates = {
|
||||
name: count / max(episode_count, 1)
|
||||
for name, count in termination_counts.items()}
|
||||
now = time.monotonic()
|
||||
row = dict(step=step + 1, actor_lr=learner.optimizers[1].param_groups[0]["lr"],
|
||||
archive_windows=0 if replay.archive is None else replay.archive.size, env_steps=learner.env_steps,
|
||||
total_transitions=learner.env_steps * args.num_envs,
|
||||
num_steps_per_env=args.num_steps_per_env, wm_loss=wm_loss, actor_loss=actor_loss,
|
||||
critic_loss=critic_loss, reward=(totals[0] / (interval_steps * args.num_envs)).item(),
|
||||
falls=int(totals[1]), interval_steps=interval_steps,
|
||||
velocity_mae=(totals[3:] / totals[2].clamp_min(1)).tolist(),
|
||||
mean_episode_length=mean_episode_length,
|
||||
episode_reward_terms=episode_reward_terms,
|
||||
weighted_reward_terms=dict(zip(env.reward_manager.active_terms,
|
||||
weighted_terms.tolist())),
|
||||
episode_termination_rates=termination_rates,
|
||||
collection_time_s=collection_time_sum / max(interval_steps, 1),
|
||||
learning_time_s=learning_time_sum / max(interval_steps, 1),
|
||||
iteration_time_s=(now - began) / max(step + 1 - start, 1),
|
||||
policy_updates=learner.policy_updates,
|
||||
skipped_updates=learner.skipped_updates, invalid_env_resets=invalid_env_resets,
|
||||
elapsed_s=now - began)
|
||||
if args.task == "s10-v4-upward":
|
||||
stats = upward_curriculum_stats(env)
|
||||
row.update(stats)
|
||||
row.update(upward_tilt_stage=stats["tilt_level_q75"],
|
||||
upward_tilt_deg=TILT_STAGES_DEG[stats["tilt_level_q75"]],
|
||||
upward_recovery_rate=stats["recovery_window_rate"],
|
||||
upward_traverse_rate=stats["traverse_window_rate"],
|
||||
upward_success_rate=stats["traverse_window_rate"])
|
||||
if args.terrain_revision >= 7:
|
||||
row['nonwheel_terminations'] = nonwheel_terminations
|
||||
if args.terrain:
|
||||
row.update({k:v.item() for k,v in learner.model.last_contact_metrics.items()})
|
||||
row["terrain_level_mean"] = env.scene.terrain.terrain_levels.float().mean().item()
|
||||
row["terrain_level_max"] = env.scene.terrain.terrain_levels.max().item()
|
||||
if args.terrain_revision >= 5:
|
||||
velocity_ema = [.9*a + .1*b for a,b in zip(velocity_ema, row['velocity_mae'])]
|
||||
row['terrain_details'] = summary(env)
|
||||
row['success_archive_windows'] = 0 if replay.success_archive is None else replay.success_archive.size
|
||||
promoted = command_stage
|
||||
if args.command_curriculum and args.terrain_revision >= 5:
|
||||
promoted = next_stage(command_stage, step + 1 - args.command_curriculum_start,
|
||||
env._stair_stats.cpu().numpy(), velocity_ema)
|
||||
row['next_command_stage'] = promoted
|
||||
if args.command_curriculum:
|
||||
ranges = env.command_manager.get_term("twist").cfg.ranges
|
||||
row.update(command_stage=command_stage, command_ranges=[ranges.lin_vel_x, ranges.lin_vel_y, ranges.ang_vel_z])
|
||||
else:
|
||||
ranges = env.command_manager.get_term("twist").cfg.ranges
|
||||
row.update(command_stage=0, command_mode="fixed",
|
||||
command_ranges=[ranges.lin_vel_x, ranges.lin_vel_y, ranges.ang_vel_z])
|
||||
if args.terrain_revision >= 3:
|
||||
row["command_negative_vx_fraction"] = (command[:,0] < 0).float().mean().item()
|
||||
row["command_abs_max"] = command.abs().amax(0).tolist()
|
||||
log.write(json.dumps(row) + "\n")
|
||||
log.flush()
|
||||
print(format_training_log(row, args, episode_reward_terms, termination_rates),
|
||||
flush=True)
|
||||
if (args.terrain_revision >= 5 and args.command_curriculum
|
||||
and promoted != command_stage):
|
||||
command_stage = promoted
|
||||
env._stair_stats.zero_() # Fresh performance window for the next stage.
|
||||
set_command_stage(env.command_manager.get_term('twist').cfg, 0,
|
||||
stage=command_stage,
|
||||
stages=stages_for(args.terrain_revision))
|
||||
reward_totals.zero_()
|
||||
totals.zero_()
|
||||
interval_steps = 0
|
||||
episode_length_sum = 0.0
|
||||
episode_count = 0
|
||||
for name in termination_counts:
|
||||
termination_counts[name] = 0
|
||||
collection_time_sum = 0.0
|
||||
learning_time_sum = 0.0
|
||||
if (step + 1) % args.save_interval == 0 or step + 1 == args.iterations:
|
||||
if args.terrain:
|
||||
learner.terrain_state = {"levels": env.scene.terrain.terrain_levels.cpu().clone(),
|
||||
"types": env.scene.terrain.terrain_types.cpu().clone()}
|
||||
if args.terrain_revision >= 5:
|
||||
learner.curriculum_state = dict(stage=command_stage, velocity_ema=velocity_ema,
|
||||
stats=env._stair_stats.detach().cpu().clone())
|
||||
learner.save(output / f"model_{step + 1}.pt", step + 1, bridge.to_env, env.step_dt)
|
||||
if args.video and ((step + 1) % args.video_interval == 0
|
||||
or step + 1 == args.iterations):
|
||||
try:
|
||||
ranges = env.command_manager.get_term("twist").cfg.ranges
|
||||
record_policy_video(
|
||||
learner.actor, args, output, step + 1, bridge.to_env,
|
||||
upward_tilt_limit=(upward_video_tilt_limit(env)
|
||||
if args.task == "s10-v4-upward" else None),
|
||||
command_ranges=(ranges.lin_vel_x, ranges.lin_vel_y, ranges.ang_vel_z))
|
||||
except Exception as exc:
|
||||
print(f"RECOVERY skipped video at {step + 1}: {exc}", flush=True)
|
||||
print(f"S10 training complete: {output}")
|
||||
finally:
|
||||
env.close()
|
||||
@@ -0,0 +1,336 @@
|
||||
"""V4-upward helpers: sign-aware upright reward, reward clamps and sustained fall.
|
||||
|
||||
This module is used only by the ``s10-v4-upward`` task. V4 itself keeps its
|
||||
baseline exp-upright reward and immediate 70-degree fall termination.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
|
||||
import torch
|
||||
from mjlab.envs import mdp
|
||||
from mjlab.managers import ManagerTermBase
|
||||
from mjlab.managers.scene_entity_config import SceneEntityCfg
|
||||
|
||||
from .world_model import GRAVITY
|
||||
|
||||
# Raw-term caps (before weight and dt) for the recovery task.
|
||||
# Chosen from v4 random-orientation samples: p99/max were approximately
|
||||
# leg_posture .78/.86, standstill 5.37/7.90, joint_torques 4772/5016,
|
||||
# unnecessary_wheel_lift 750/826, rolling_leg_motion 2.38/4.47.
|
||||
UPWARD_REWARD_CLAMPS = {
|
||||
"leg_posture": 2.0,
|
||||
"standstill": 4.0,
|
||||
"action_rate_l2": 4.0,
|
||||
"joint_torques_l2": 5000.0,
|
||||
"dof_pos_limits": 0.5,
|
||||
"unnecessary_wheel_lift": 2.0,
|
||||
"rolling_leg_motion": 5.0,
|
||||
}
|
||||
|
||||
# IsaacLab upward: upward max 4 (raw 4 x weight 1), bad-orientation -80,
|
||||
# i.e. the bad penalty is 20x the upward maximum. Here upward max is 2
|
||||
# (raw 1 x weight 2), so the same ratio is -40.
|
||||
BAD_ORIENTATION_WEIGHT = -40.0
|
||||
BAD_ORIENTATION_XY = 0.7
|
||||
# The old max(gz, |gx|-0.7, |gy|-0.7) severity saturated for gz >= 0.3 and had
|
||||
# a local maximum near 105 deg, where the |gx|-0.7 term crossed gz. The reward
|
||||
# severity is now a function of the tilt angle only: 0 at 45 deg, 1 at 180 deg,
|
||||
# monotonic in between. The correction term steepens the near-threshold part
|
||||
# while keeping a positive slope at 180 deg and a finite derivative at 45 deg.
|
||||
BAD_ORIENTATION_START_DEG = 45.0
|
||||
BAD_ORIENTATION_CORRECTION = 2.5
|
||||
|
||||
# Per-environment recovery curriculum, IsaacLab terrain-curriculum style.
|
||||
# Each env keeps the outcomes of its last HISTORY_LENGTH episodes and advances
|
||||
# when at least PROMOTE_SUCCESSES of them recovered (16/20 = 80% > 75%).
|
||||
# There is no global batch gate: terrain and tilt levels are per environment.
|
||||
TILT_STAGES_DEG = (60.0, 135.0, 180.0)
|
||||
HISTORY_LENGTH = 20
|
||||
PROMOTE_SUCCESSES = 16
|
||||
DEMOTE_SUCCESSES = 5
|
||||
|
||||
|
||||
def upright_upward_value(gravity: torch.Tensor) -> torch.Tensor:
|
||||
"""Normalized upward reward in [0, 1].
|
||||
|
||||
``projected_gravity_b[:, 2]`` is -1 upright, 0 sideways and +1 upside down.
|
||||
The old v4 exp-upright term is sign-blind: upside-down also scores 1. This
|
||||
linear form is monotonic through the whole range and has a constant gradient
|
||||
with respect to gravity_z, which is what the recovery task needs.
|
||||
"""
|
||||
return (0.5 * (1.0 - gravity[:, 2])).clamp(0.0, 1.0)
|
||||
|
||||
|
||||
def upright_upward(env, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
|
||||
"""Environment reward term replacing v4 ``upright`` for this task."""
|
||||
gravity = env.scene[asset_cfg.name].data.projected_gravity_b
|
||||
return upright_upward_value(gravity)
|
||||
|
||||
|
||||
def bad_orientation_value(gravity: torch.Tensor) -> torch.Tensor:
|
||||
"""Monotonic tilt-based bad-orientation severity in [0, 1].
|
||||
|
||||
``projected_gravity_b[:, 2]`` is -1 upright and +1 upside down, so
|
||||
``tilt = acos(-gz)`` goes from 0 to pi. Severity is 0 below
|
||||
``BAD_ORIENTATION_START_DEG``, then increases monotonically to 1 at 180 deg:
|
||||
|
||||
x = clamp((tilt - 45 deg) / 135 deg, 0, 1)
|
||||
severity = x + 2.5 * x * (1 - x)^2
|
||||
|
||||
Unlike the previous max(gz, |gx|-0.7, |gy|-0.7) formula, this has no
|
||||
saturation plateau for deep inversion and no local maximum near 105 deg.
|
||||
The polynomial correction is used instead of a fractional power because
|
||||
``x ** (1/3)`` has an infinite derivative at x=0 (upright), which would
|
||||
produce non-finite gradients during imagination.
|
||||
"""
|
||||
gz = gravity[:, 2].clamp(-1.0 + 1e-4, 1.0 - 1e-4)
|
||||
tilt = torch.acos(-gz)
|
||||
start = math.radians(BAD_ORIENTATION_START_DEG)
|
||||
x = ((tilt - start) / (math.pi - start)).clamp(0.0, 1.0)
|
||||
return x + BAD_ORIENTATION_CORRECTION * x * (1.0 - x).square()
|
||||
|
||||
|
||||
def bad_orientation_mask(gravity: torch.Tensor) -> torch.Tensor:
|
||||
"""Binary bad set used by the 4-second termination.
|
||||
|
||||
This intentionally keeps the IsaacLab-style boundary: upside-down
|
||||
hemisphere or |gx|/|gy| beyond 0.7. The reward severity above is a smooth
|
||||
tilt-based shaping term and is not used to decide termination.
|
||||
"""
|
||||
return ((gravity[:, 2] > 0.0)
|
||||
| (gravity[:, :2].abs() > BAD_ORIENTATION_XY).any(-1))
|
||||
|
||||
|
||||
def bad_orientation_penalty(env, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
|
||||
"""Continuous per-step bad-orientation penalty used by the environment."""
|
||||
gravity = env.scene[asset_cfg.name].data.projected_gravity_b
|
||||
return bad_orientation_value(gravity)
|
||||
|
||||
|
||||
def is_upward_task(env) -> bool:
|
||||
"""Identify the task from its upright reward function, not from a new flag."""
|
||||
rewards = getattr(env.cfg, "rewards", {})
|
||||
term = rewards.get("upright") if hasattr(rewards, "get") else None
|
||||
return term is not None and getattr(term, "func", None) is upright_upward
|
||||
|
||||
|
||||
def _sample_offsets(range_dict, count: int, device):
|
||||
"""Sample x/y/z/roll/pitch/yaw offsets; missing keys are zero."""
|
||||
offsets = []
|
||||
for key in ("x", "y", "z", "roll", "pitch", "yaw"):
|
||||
lo, hi = (range_dict or {}).get(key, (0.0, 0.0))
|
||||
lo, hi = float(lo), float(hi)
|
||||
if hi <= lo:
|
||||
offsets.append(torch.full((count,), lo, device=device))
|
||||
else:
|
||||
offsets.append(torch.empty(count, device=device).uniform_(lo, hi))
|
||||
return torch.stack(offsets, -1)
|
||||
|
||||
|
||||
def reset_root_state_upward(env, env_ids, pose_range, velocity_range=None,
|
||||
asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")):
|
||||
"""Reset with a per-environment roll/pitch limit.
|
||||
|
||||
``env._upward_tilt_limits`` is a per-env tensor updated by the curriculum.
|
||||
Play/evaluation has no such tensor and falls back to the full range in
|
||||
``pose_range`` (the task sets +/-pi there).
|
||||
"""
|
||||
env_ids = mdp.resolve_env_ids(env, env_ids)
|
||||
asset = env.scene[asset_cfg.name]
|
||||
count = len(env_ids)
|
||||
pose = _sample_offsets(pose_range, count, env.device)
|
||||
limits = getattr(env, "_upward_tilt_limits", None)
|
||||
if limits is None:
|
||||
roll = pose_range.get("roll", (-math.pi, math.pi))
|
||||
limit = max(abs(float(roll[0])), abs(float(roll[1])))
|
||||
limits = torch.full((count,), limit, device=env.device)
|
||||
else:
|
||||
limits = limits[env_ids]
|
||||
pose[:, 3] = (2.0 * torch.rand(count, device=env.device) - 1.0) * limits
|
||||
pose[:, 4] = (2.0 * torch.rand(count, device=env.device) - 1.0) * limits
|
||||
default = asset.data.default_root_state[env_ids].clone()
|
||||
positions = default[:, 0:3] + pose[:, 0:3] + env.scene.env_origins[env_ids]
|
||||
delta = mdp.quat_from_euler_xyz(pose[:, 3], pose[:, 4], pose[:, 5])
|
||||
orientations = mdp.quat_mul(default[:, 3:7], delta)
|
||||
velocities = default[:, 7:13] + _sample_offsets(velocity_range, count, env.device)
|
||||
asset.write_root_link_pose_to_sim(
|
||||
torch.cat((positions, orientations), dim=-1), env_ids=env_ids)
|
||||
asset.write_root_link_velocity_to_sim(velocities, env_ids=env_ids)
|
||||
|
||||
|
||||
def init_upward_state(env):
|
||||
"""Allocate the per-env recovery ring and tilt levels once."""
|
||||
if getattr(env, "_upward_history", None) is not None:
|
||||
return
|
||||
n, device = env.num_envs, env.device
|
||||
env._upward_history = torch.zeros(
|
||||
n, HISTORY_LENGTH, dtype=torch.bool, device=device)
|
||||
env._upward_traverse_history = torch.zeros_like(env._upward_history)
|
||||
env._upward_history_cursor = torch.zeros(n, dtype=torch.long, device=device)
|
||||
env._upward_history_filled = torch.zeros(n, dtype=torch.long, device=device)
|
||||
env._upward_tilt_levels = torch.zeros(n, dtype=torch.long, device=device)
|
||||
env._upward_tilt_limit_table = torch.tensor(
|
||||
[math.radians(d) for d in TILT_STAGES_DEG], device=device)
|
||||
env._upward_tilt_limits = env._upward_tilt_limit_table[0].expand(n).clone()
|
||||
|
||||
|
||||
def upward_curriculum(env, env_ids):
|
||||
"""Per-env rolling recovery curriculum (IsaacLab terrain-curriculum style).
|
||||
|
||||
Each environment stores the outcomes of its last ``HISTORY_LENGTH``
|
||||
episodes. ``recovered`` means the episode did not end in a bad-orientation
|
||||
termination and the final projected gravity is not in the bad set. An env
|
||||
advances when at least ``PROMOTE_SUCCESSES`` of the window recovered and
|
||||
drops when at most ``DEMOTE_SUCCESSES`` did. Terrain and tilt levels are
|
||||
updated independently per environment; no global batch rate is involved.
|
||||
"""
|
||||
init_upward_state(env)
|
||||
env_ids = mdp.resolve_env_ids(env, env_ids)
|
||||
terrain = env.scene.terrain
|
||||
data = env.scene["robot"].data
|
||||
active = env.episode_length_buf[env_ids] > 0
|
||||
fallen = env.termination_manager.terminated[env_ids]
|
||||
gravity = data.projected_gravity_b[env_ids]
|
||||
offset = data.root_link_pos_w[env_ids] - env.scene.env_origins[env_ids]
|
||||
recovered = active & ~fallen & ~bad_orientation_mask(gravity)
|
||||
traversed = (active & ~fallen & (offset[:, 0].abs() > 5.4)
|
||||
& (offset[:, 1].abs() < 6.5))
|
||||
ids = env_ids[active]
|
||||
if len(ids):
|
||||
cursor = env._upward_history_cursor[ids]
|
||||
env._upward_history[ids, cursor] = recovered[active]
|
||||
env._upward_traverse_history[ids, cursor] = traversed[active]
|
||||
env._upward_history_cursor[ids] = (cursor + 1) % HISTORY_LENGTH
|
||||
env._upward_history_filled[ids] = (
|
||||
env._upward_history_filled[ids] + 1).clamp_max(HISTORY_LENGTH)
|
||||
ready = env._upward_history_filled[env_ids] >= HISTORY_LENGTH
|
||||
successes = env._upward_history[env_ids].sum(-1)
|
||||
promote = ready & (successes >= PROMOTE_SUCCESSES)
|
||||
demote = ready & (successes <= DEMOTE_SUCCESSES)
|
||||
delta = promote.long() - demote.long()
|
||||
old_levels = terrain.terrain_levels[env_ids].clone()
|
||||
levels = (old_levels + delta).clamp(0, terrain.terrain_origins.shape[0] - 1)
|
||||
old_tilt_levels = env._upward_tilt_levels[env_ids].clone()
|
||||
tilt_levels = (old_tilt_levels + delta).clamp(0, len(TILT_STAGES_DEG) - 1)
|
||||
# A changed environment must prove itself for a fresh 20-episode window at
|
||||
# the new level; otherwise a long success streak could chain promotions in
|
||||
# consecutive episodes. A clamped promote/demote at the min/max level is not
|
||||
# a change and must keep its window so the logged rate stays meaningful.
|
||||
changed = (levels != old_levels) | (tilt_levels != old_tilt_levels)
|
||||
if changed.any():
|
||||
changed_ids = env_ids[changed]
|
||||
env._upward_history[changed_ids] = False
|
||||
env._upward_traverse_history[changed_ids] = False
|
||||
env._upward_history_filled[changed_ids] = 0
|
||||
env._upward_history_cursor[changed_ids] = 0
|
||||
terrain.terrain_levels[env_ids] = levels
|
||||
terrain.env_origins[env_ids] = terrain.terrain_origins[
|
||||
levels, terrain.terrain_types[env_ids]]
|
||||
env._upward_tilt_levels[env_ids] = tilt_levels
|
||||
env._upward_tilt_limits[env_ids] = env._upward_tilt_limit_table[tilt_levels]
|
||||
return terrain.terrain_levels.float().mean()
|
||||
|
||||
|
||||
def upward_curriculum_stats(env):
|
||||
"""Summary for logging and for choosing the training-video tilt level."""
|
||||
init_upward_state(env)
|
||||
levels = env._upward_tilt_levels
|
||||
terrain_levels = env.scene.terrain.terrain_levels
|
||||
ready = env._upward_history_filled >= HISTORY_LENGTH
|
||||
recovery = env._upward_history.sum(-1)
|
||||
traverse = env._upward_traverse_history.sum(-1)
|
||||
ready_count = int(ready.sum())
|
||||
recovery_rate = (float(recovery[ready].float().mean().item()) / HISTORY_LENGTH
|
||||
if ready_count else 0.0)
|
||||
traverse_rate = (float(traverse[ready].float().mean().item()) / HISTORY_LENGTH
|
||||
if ready_count else 0.0)
|
||||
def fractions(values, count):
|
||||
return [float((values == i).float().mean().item()) for i in range(count)]
|
||||
q75 = int(torch.quantile(levels.float(), 0.75).round().item())
|
||||
return {
|
||||
"tilt_level_mean": float(levels.float().mean().item()),
|
||||
"tilt_level_max": int(levels.max().item()),
|
||||
"tilt_level_q75": q75,
|
||||
"tilt_level_frac": fractions(levels, len(TILT_STAGES_DEG)),
|
||||
"tilt_limit_mean_deg": float(
|
||||
env._upward_tilt_limits.mean().item() * 180.0 / math.pi),
|
||||
"terrain_level_frac": fractions(
|
||||
terrain_levels, int(env.scene.terrain.terrain_origins.shape[0])),
|
||||
"recovery_window_rate": recovery_rate,
|
||||
"traverse_window_rate": traverse_rate,
|
||||
"recovery_window_ready": float(ready.float().mean().item()),
|
||||
}
|
||||
|
||||
|
||||
def upward_video_tilt_limit(env):
|
||||
"""Tilt limit reached by at least 75% of environments, for the video scene."""
|
||||
q75 = upward_curriculum_stats(env)["tilt_level_q75"]
|
||||
return math.radians(TILT_STAGES_DEG[q75])
|
||||
|
||||
|
||||
def _cap(value: torch.Tensor, max_value: float) -> torch.Tensor:
|
||||
return value.clamp(max=float(max_value))
|
||||
|
||||
|
||||
def clamped_leg_posture(env, asset_cfg: SceneEntityCfg, max_value: float) -> torch.Tensor:
|
||||
from .terrain_posture import env_posture_cost
|
||||
return _cap(env_posture_cost(env, asset_cfg), max_value)
|
||||
|
||||
|
||||
def clamped_standstill(env, wheel_cfg: SceneEntityCfg, max_value: float) -> torch.Tensor:
|
||||
from s10_supermimic.env.flat import standstill_cost
|
||||
return _cap(standstill_cost(env, wheel_cfg), max_value)
|
||||
|
||||
|
||||
def clamped_action_rate(env, max_value: float) -> torch.Tensor:
|
||||
return _cap(mdp.action_rate_l2(env), max_value)
|
||||
|
||||
|
||||
def clamped_joint_torques(env, max_value: float) -> torch.Tensor:
|
||||
return _cap(mdp.joint_torques_l2(env), max_value)
|
||||
|
||||
|
||||
def clamped_dof_pos_limits(env, asset_cfg: SceneEntityCfg, max_value: float) -> torch.Tensor:
|
||||
return _cap(mdp.joint_pos_limits(env, asset_cfg), max_value)
|
||||
|
||||
|
||||
def clamped_rolling_cost(env, term: str, max_value: float) -> torch.Tensor:
|
||||
from .rolling_rewards import env_rolling_cost
|
||||
return _cap(env_rolling_cost(env, term), max_value)
|
||||
|
||||
|
||||
def apply_reward_clamps(terms: dict[str, torch.Tensor], constants: dict) -> None:
|
||||
"""Apply the same raw caps to imagined rewards used by the env task."""
|
||||
for name, max_value in constants.get("reward_clamps", {}).items():
|
||||
if name in terms:
|
||||
terms[name] = terms[name].clamp(max=float(max_value))
|
||||
|
||||
|
||||
class BadOrientationSustained(ManagerTermBase):
|
||||
"""Terminate only after bad orientation persists for ``duration_s``.
|
||||
|
||||
This is the recovery-task state machine: a random roll/pitch spawn is allowed
|
||||
to be bad for a few seconds, but the episode ends if the robot never rights
|
||||
itself. Functionally it replaces v4's immediate 70-degree ``fell_over``.
|
||||
"""
|
||||
|
||||
def __init__(self, cfg, env):
|
||||
super().__init__(env)
|
||||
self.steps = torch.zeros(env.num_envs, dtype=torch.long, device=env.device)
|
||||
|
||||
def __call__(self, env, duration_s: float, asset_cfg: SceneEntityCfg) -> torch.Tensor:
|
||||
gravity = env.scene[asset_cfg.name].data.projected_gravity_b
|
||||
# Termination keeps the binary IsaacLab-style bad set; the reward severity
|
||||
# above is only a smooth shaping term.
|
||||
bad = bad_orientation_mask(gravity)
|
||||
self.steps = torch.where(bad, self.steps + 1, torch.zeros_like(self.steps))
|
||||
duration_steps = math.ceil(duration_s / env.step_dt)
|
||||
return self.steps >= duration_steps
|
||||
|
||||
def reset(self, env_ids=None) -> None:
|
||||
if env_ids is None:
|
||||
self.steps.zero_()
|
||||
else:
|
||||
self.steps[env_ids] = 0
|
||||
@@ -0,0 +1,101 @@
|
||||
"""Always-on wheel/body projection constraints for the safe v7 stance.
|
||||
|
||||
The body is 0.60 m long in ``base_link`` (x in [-0.30, 0.30]).
|
||||
|
||||
``wheel_body_intrusion`` keeps every individual wheel out of the opposite
|
||||
three quarters: a front wheel inside the rear three quarters (x < +0.15) or a
|
||||
rear wheel inside the front three quarters (x > -0.15) shortens the support
|
||||
polygon and makes the robot tip-prone. Its hinge is exactly zero outside the
|
||||
outer quarters and grows cubically inside, so a normal extended stance is
|
||||
untouched while a wheel tucked under the body is penalised quickly.
|
||||
|
||||
``pair_body_intrusion`` adds the wider five-sixths restriction requested for
|
||||
v7: both front wheels may not sit inside the rear five sixths (x < +0.20) at
|
||||
the same time, and both rear wheels may not sit inside the front five sixths
|
||||
(x > -0.20). A single leg crossing the line is allowed; a leg only counts as
|
||||
inside after it has stayed inside continuously for ``PAIR_GRACE_S`` (0.5 s).
|
||||
The pair penalty is therefore zero during a brief swing-through and applies
|
||||
immediately once both same-side legs are armed.
|
||||
|
||||
Both terms are kinematic and need no terrain scan. Wheels are ordered
|
||||
fl, fr, hl, hr.
|
||||
"""
|
||||
import torch
|
||||
|
||||
from .rolling_rewards import wheel_geometry
|
||||
|
||||
BODY_HALF_LENGTH = 0.30
|
||||
# Each wheel must stay in the outer quarter of the body length: x > +0.15 for
|
||||
# front wheels and x < -0.15 for rear wheels.
|
||||
WHEEL_STRIP_BOUNDARY = BODY_HALF_LENGTH / 2.0
|
||||
# Rear/front 5/6 of a 0.60 m body => the boundary sits at x = +-0.20.
|
||||
PAIR_BOUNDARY = 2.0 * BODY_HALF_LENGTH / 3.0
|
||||
PAIR_GRACE_S = 0.5
|
||||
|
||||
|
||||
def wheel_body_intrusion(state, default_q, boundary=WHEEL_STRIP_BOUNDARY):
|
||||
point, _ = wheel_geometry(state, default_q)
|
||||
x = point[..., 0]
|
||||
# Front wheels must stay ahead of +L/4, rear wheels behind -L/4.
|
||||
violation = torch.cat(((boundary - x[:, :2]).clamp_min(0.),
|
||||
(x[:, 2:] + boundary).clamp_min(0.)), -1)
|
||||
normalized = violation / boundary
|
||||
return normalized.pow(3).mean(-1)
|
||||
|
||||
|
||||
def pair_region(state, default_q, boundary=PAIR_BOUNDARY):
|
||||
"""Per-wheel inside mask and normalised depth for the 5/6 restriction."""
|
||||
point, _ = wheel_geometry(state, default_q)
|
||||
x = point[..., 0]
|
||||
inside = torch.cat((x[:, :2] < boundary, x[:, 2:] > -boundary), -1)
|
||||
depth = torch.cat(((boundary - x[:, :2]).clamp_min(0.),
|
||||
(x[:, 2:] + boundary).clamp_min(0.)), -1) / boundary
|
||||
return inside, depth
|
||||
|
||||
|
||||
def pair_body_intrusion(state, default_q, timers, dt, grace=PAIR_GRACE_S):
|
||||
"""Penalise both front or both rear wheels staying inside the 5/6 region.
|
||||
|
||||
``timers`` holds the continuous inside time of each wheel. Returning it lets
|
||||
the caller keep the 0.5 s grace across environment steps. A wheel outside
|
||||
the boundary resets its own timer, so a single-leg pass-through never
|
||||
contributes.
|
||||
"""
|
||||
inside, depth = pair_region(state, default_q)
|
||||
timers = torch.where(inside, timers + dt, torch.zeros_like(timers))
|
||||
armed = timers >= grace - 1e-6
|
||||
front = armed[:, 0] & armed[:, 1]
|
||||
rear = armed[:, 2] & armed[:, 3]
|
||||
cost = (front.float() * depth[:, :2].mean(-1)
|
||||
+ rear.float() * depth[:, 2:].mean(-1))
|
||||
return cost, timers
|
||||
|
||||
|
||||
def _bridge(env):
|
||||
from .terrain import TerrainBridge, TerrainField
|
||||
bridge = getattr(env, '_s10_posture_bridge', None)
|
||||
if bridge is None:
|
||||
bridge = TerrainBridge(env, TerrainField.from_env(env))
|
||||
env._s10_posture_bridge = bridge
|
||||
return bridge
|
||||
|
||||
|
||||
def env_wheel_body_intrusion(env):
|
||||
bridge = _bridge(env)
|
||||
return wheel_body_intrusion(
|
||||
bridge.state(), bridge.robot.data.default_joint_pos[0, bridge.legs])
|
||||
|
||||
|
||||
def env_pair_body_intrusion(env):
|
||||
bridge = _bridge(env)
|
||||
timers = getattr(env, '_s10_pair_timer', None)
|
||||
if timers is None or timers.shape[0] != env.num_envs:
|
||||
timers = torch.zeros(env.num_envs, 4, device=env.device)
|
||||
# episode_length_buf is 1 on the first reward step after an auto-reset.
|
||||
timers = timers.clone()
|
||||
timers[env.episode_length_buf <= 1] = 0.
|
||||
cost, timers = pair_body_intrusion(
|
||||
bridge.state(), bridge.robot.data.default_joint_pos[0, bridge.legs],
|
||||
timers, env.step_dt)
|
||||
env._s10_pair_timer = timers
|
||||
return cost
|
||||
@@ -0,0 +1,81 @@
|
||||
"""V6 wheel-local terrain constraints; privileged scans never enter the actor.
|
||||
|
||||
Clearance is a geometric surrogate, not a measured persistent-contact detector.
|
||||
S10 joint geometry comes from ``rolling_rewards.wheel_geometry``.
|
||||
"""
|
||||
import torch
|
||||
from .world_model import GRAVITY, LEG_POS, LEG_VEL
|
||||
from .rolling_rewards import wheel_geometry
|
||||
|
||||
|
||||
def yaw_frame_points(state, points):
|
||||
# Remove roll/pitch, retaining the heading frame used by TerrainField.scan.
|
||||
up = -torch.nn.functional.normalize(state[:, GRAVITY], dim=-1, eps=1e-6)
|
||||
cp = (up[:, 1:].square().sum(-1)+1e-12).sqrt()
|
||||
sr, cr, sp = up[:, 1] / cp, up[:, 2] / cp, -up[:, 0]
|
||||
x, y, z = points.unbind(-1)
|
||||
yr = cr[:, None]*y - sr[:, None]*z
|
||||
zr = sr[:, None]*y + cr[:, None]*z
|
||||
return torch.stack((cp[:, None]*x + sp[:, None]*zr, yr,
|
||||
-sp[:, None]*x + cp[:, None]*zr), -1)
|
||||
|
||||
|
||||
def sample_scan(scan, xy):
|
||||
"""Bilinear samples on the existing 21x13, 10cm heading-frame grid."""
|
||||
shape = xy.shape[:-1]
|
||||
x = ((xy[..., 0]+1.)*10).clamp(0,20).reshape(len(scan),-1)
|
||||
y = ((xy[..., 1]+.6)*10).clamp(0,12).reshape(len(scan),-1)
|
||||
ix, iy = x.floor().long().clamp(0,19), y.floor().long().clamp(0,11)
|
||||
ax, ay = x-ix, y-iy
|
||||
def at(dx,dy): return scan.gather(1,(ix+dx)*13+iy+dy)
|
||||
return (((1-ax)*at(0,0)+ax*at(1,0))*(1-ay)
|
||||
+ ((1-ax)*at(0,1)+ax*at(1,1))*ay).reshape(shape)
|
||||
|
||||
|
||||
def wheel_terms(state, command, scan, default_q):
|
||||
point, radius = wheel_geometry(state, default_q)
|
||||
point = yaw_frame_points(state, point)
|
||||
xy = point[..., :2]
|
||||
offsets = state.new_tensor([[x,y] for x in (-.15,0.,.15) for y in (-.10,0.,.10)])
|
||||
patch_xy = xy[:,:,None,:]+offsets
|
||||
samples = sample_scan(scan, patch_xy)
|
||||
# Index-based extrema avoid division by the number of equal extrema in
|
||||
# compiled backward through interpolated, nearly-flat patches.
|
||||
relief = samples.max(dim=-1).values-samples.min(dim=-1).values
|
||||
# Outside the scan, never assume flatness and penalize a recovering leg.
|
||||
inside = ((patch_xy[...,0].abs()<=1.) & (patch_xy[...,1].abs()<=.6)).all(-1)
|
||||
flat = (1.-(relief-.02)/.04).clamp(0,1)*inside
|
||||
straight = torch.exp(-(command[:,1]/.25).square()-(command[:,2]).square())
|
||||
gate = flat*straight[:,None]
|
||||
below = sample_scan(scan,xy)
|
||||
bottom = point[...,2]+below-radius
|
||||
lift = ((bottom-.005).clamp_min(0)/.05).square()
|
||||
tilt = ((state[:,GRAVITY][:,:2].square().sum(-1)+1e-12).sqrt()-.08).div(.22).clamp(0,1)
|
||||
motion = state[:,LEG_VEL].reshape(-1,4,3).square().mean(-1)*(1.-.8*tilt[:,None])
|
||||
q = state[:,LEG_POS].reshape(-1,4,3).square()
|
||||
posture_multiplier = torch.where(inside, 1.-.875*(relief/.12).clamp(0,1), .125)
|
||||
posture = 2.*q[:,:,0].sum(-1)+(posture_multiplier*q[:,:,1:].sum(-1)).sum(-1)
|
||||
# Only approaching a higher surface: no reward for lifting on flat ground.
|
||||
speed = command[:,:2].norm(dim=-1)
|
||||
heading_velocity = yaw_frame_points(state, torch.cat((command[:,:2], torch.zeros_like(command[:,:1])), -1)[:,None,:])[:,0,:2]
|
||||
direction = heading_velocity/(heading_velocity.square().sum(-1,keepdim=True)+1e-12).sqrt().clamp_min(.1)
|
||||
ahead_xy = xy+.16*direction[:,None,:]
|
||||
ahead = sample_scan(scan,ahead_xy)
|
||||
rise = below-ahead
|
||||
approaching = ((rise-.04)/.08).clamp(0,1)*(speed>.1)[:,None]*inside
|
||||
deficit = (radius+.015-point[...,2]-ahead).clamp_min(0)/.20
|
||||
clearance = (approaching*deficit/(1.+deficit)).mean(-1)
|
||||
return dict(unnecessary_wheel_lift=(gate*lift).sum(-1),
|
||||
rolling_leg_motion=(gate*motion).mean(-1), leg_posture=posture,
|
||||
wheel_obstacle_clearance=clearance)
|
||||
|
||||
|
||||
def env_wheel_term(env, term):
|
||||
from .terrain import TerrainBridge,TerrainField
|
||||
bridge=getattr(env,'_s10_posture_bridge',None)
|
||||
if bridge is None:
|
||||
bridge=TerrainBridge(env,TerrainField.from_env(env));env._s10_posture_bridge=bridge
|
||||
state=bridge.state()
|
||||
return wheel_terms(state,env.command_manager.get_command('twist'),
|
||||
bridge.field.scan(state,bridge.pose()),
|
||||
bridge.robot.data.default_joint_pos[0,bridge.legs])[term]
|
||||
@@ -0,0 +1,118 @@
|
||||
"""S10 residual dynamics with position legs and velocity wheels.
|
||||
|
||||
Canonical order: fl, fr, hl, hr; hipx/hipy/knee per leg, then four wheels.
|
||||
Actions are normalized policy outputs: leg targets = default + scale*a[:12],
|
||||
wheel targets = clamp(35*a[12:], -35, 35) rad/s. Wheel phase is excluded.
|
||||
This module is self-contained and depends only on torch.
|
||||
"""
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
STATE_DIM = 38
|
||||
ACTION_DIM = 16
|
||||
CONTEXT_DIM = 7 # four wheel sliding frictions, base COM xyz offset (metres)
|
||||
POSE_DIM = 3
|
||||
Z, GRAVITY, LIN_VEL, ANG_VEL = slice(0, 1), slice(1, 4), slice(4, 7), slice(7, 10)
|
||||
LEG_POS, LEG_VEL, WHEEL_VEL = slice(10, 22), slice(22, 34), slice(34, 38)
|
||||
FORMAT = "s10-flat-dynamics-v1"
|
||||
|
||||
|
||||
class RunningNorm(nn.Module):
|
||||
def __init__(self, dim):
|
||||
super().__init__()
|
||||
self.register_buffer("mean", torch.zeros(dim))
|
||||
self.register_buffer("var", torch.ones(dim))
|
||||
self.register_buffer("count", torch.tensor(1e-4))
|
||||
|
||||
@torch.no_grad()
|
||||
def update(self, x):
|
||||
x = x.reshape(-1, self.mean.numel())
|
||||
if not len(x):
|
||||
return
|
||||
delta = x.mean(0) - self.mean
|
||||
total = self.count + len(x)
|
||||
self.var.copy_((self.var * self.count + x.var(0, unbiased=False) * len(x)
|
||||
+ delta.square() * self.count * len(x) / total) / total)
|
||||
self.mean.add_(delta * len(x) / total)
|
||||
self.count.copy_(total)
|
||||
|
||||
@property
|
||||
def scale(self):
|
||||
return self.var.clamp_min(1e-6).sqrt()
|
||||
|
||||
def forward(self, x):
|
||||
return ((x - self.mean) / self.scale).clamp(-10, 10)
|
||||
|
||||
|
||||
class S10WorldModel(nn.Module):
|
||||
"""(state38, action16, previous16, context7) -> (next state38, pose delta3).
|
||||
|
||||
Flat ground only. Linear and angular velocity use the base frame. Planar
|
||||
displacement uses the previous yaw frame. No target/chase or terrain scan.
|
||||
Policy history is separate from physical state and is not fed to this model.
|
||||
"""
|
||||
def __init__(self, hidden=(512, 512, 512)):
|
||||
super().__init__()
|
||||
self.hidden = tuple(hidden)
|
||||
dims = (STATE_DIM + 2 * ACTION_DIM + CONTEXT_DIM, *hidden, STATE_DIM + POSE_DIM)
|
||||
layers = []
|
||||
for i, (a, b) in enumerate(zip(dims[:-1], dims[1:])):
|
||||
layers.append(nn.Linear(a, b))
|
||||
if i < len(dims) - 2:
|
||||
layers.append(nn.SiLU())
|
||||
self.net = nn.Sequential(*layers)
|
||||
self.input_norm = RunningNorm(dims[0])
|
||||
self.output_norm = RunningNorm(dims[-1])
|
||||
|
||||
@staticmethod
|
||||
def inputs(state, action, previous_action, context):
|
||||
for name, x, dim in (("state", state, STATE_DIM), ("action", action, ACTION_DIM),
|
||||
("previous_action", previous_action, ACTION_DIM),
|
||||
("context", context, CONTEXT_DIM)):
|
||||
if x.ndim != 2 or x.shape != (state.shape[0], dim):
|
||||
raise ValueError(f"{name}: expected (batch, {dim}), got {tuple(x.shape)}")
|
||||
return torch.cat((state, action, previous_action, context), -1)
|
||||
|
||||
@torch.no_grad()
|
||||
def update_stats(self, state, action, previous_action, context, next_state, pose_delta):
|
||||
self.input_norm.update(self.inputs(state, action, previous_action, context))
|
||||
self.output_norm.update(torch.cat((next_state - state, pose_delta), -1))
|
||||
|
||||
def forward(self, state, action, previous_action, context):
|
||||
x = self.inputs(state, action, previous_action, context)
|
||||
delta = self.net(self.input_norm(x)) * self.output_norm.scale + self.output_norm.mean
|
||||
return state + delta[:, :STATE_DIM], delta[:, STATE_DIM:]
|
||||
|
||||
def loss(self, state, actions, previous_actions, context, targets, poses, valid):
|
||||
"""Autoregressive window loss; validity must be false at/reset after done.
|
||||
|
||||
Tensor shapes: state[B,38], actions/previous[B,K,16], context[B,7],
|
||||
targets[B,K,38], poses[B,K,3], valid[B,K]. A reset invalidates the rest
|
||||
of that window, even if later transitions are individually valid.
|
||||
"""
|
||||
mask = valid.bool().cumprod(1).bool()
|
||||
total = next(self.parameters()).sum() * 0.0
|
||||
for t in range(actions.shape[1]):
|
||||
# Do not execute invalid rows: NaN terminal/reset data must not leak
|
||||
# into either the loss or gradients through multiplication by zero.
|
||||
keep = mask[:, t]
|
||||
if not keep.any():
|
||||
break
|
||||
pred, pose = self(state[keep], actions[keep, t], previous_actions[keep, t], context[keep])
|
||||
error = torch.cat((pred - targets[keep, t], pose - poses[keep, t]), -1)
|
||||
total = total + (error / self.output_norm.scale).square().mean(-1).sum()
|
||||
state = state.clone()
|
||||
state[keep] = pred
|
||||
return total / mask.sum().clamp_min(1)
|
||||
|
||||
def checkpoint(self):
|
||||
return {"format": FORMAT, "hidden": self.hidden, "model": self.state_dict()}
|
||||
|
||||
@classmethod
|
||||
def from_checkpoint(cls, checkpoint, device="cpu"):
|
||||
if checkpoint.get("format") != FORMAT:
|
||||
raise ValueError("Expected a S10 dynamics checkpoint; other-robot checkpoints are incompatible")
|
||||
model = cls(checkpoint["hidden"]).to(device)
|
||||
model.load_state_dict(checkpoint["model"])
|
||||
return model
|
||||
@@ -0,0 +1,296 @@
|
||||
"""World-model future preview for the play scripts.
|
||||
|
||||
The actor checkpoint already contains the world model (``world_model``) and the
|
||||
training config. We use it directly; no separate ONNX export is required.
|
||||
|
||||
The preview has two independent outputs:
|
||||
|
||||
* a kinematic ghost trail: every predicted state is mapped back to MJCF qpos
|
||||
and transparently added to the offscreen renderer;
|
||||
* a prediction-error ledger: predictions are stored with their target step and
|
||||
compared with the real state when that step arrives (xy/z/yaw/gravity/leg q).
|
||||
|
||||
The ghost is only a visualization. It never writes into ``env.sim.data`` or the
|
||||
physics model.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
import math
|
||||
|
||||
import mujoco
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from .imagine import frame, push_history
|
||||
from .policy import Actor
|
||||
from .terrain import advance_pose
|
||||
from .terrain_world_model import TerrainWorldModel
|
||||
from .world_model import GRAVITY, LEG_POS, S10WorldModel
|
||||
|
||||
DEFAULT_HORIZON = 32
|
||||
DEFAULT_STRIDE = 4
|
||||
GHOST_NEAR = (1.0, 0.86, 0.08, 0.20)
|
||||
GHOST_FAR = (1.0, 0.24, 0.02, 0.055)
|
||||
|
||||
|
||||
@dataclass
|
||||
class GhostPrediction:
|
||||
horizon: int
|
||||
state: np.ndarray # 38
|
||||
pose: np.ndarray # x, y, yaw
|
||||
qpos: np.ndarray # MJCF qpos
|
||||
|
||||
|
||||
def load_world_model(checkpoint, device="cpu"):
|
||||
payload = checkpoint.get("world_model")
|
||||
if not isinstance(payload, dict):
|
||||
raise ValueError("Checkpoint has no world_model payload")
|
||||
fmt = payload.get("format")
|
||||
if fmt == "s10-terrain-dynamics-v1":
|
||||
model = TerrainWorldModel.from_checkpoint(payload, device)
|
||||
elif fmt == "s10-flat-dynamics-v1":
|
||||
model = S10WorldModel.from_checkpoint(payload, device)
|
||||
else:
|
||||
raise ValueError(f"Unsupported world-model format: {fmt}")
|
||||
return model.to(device).eval()
|
||||
|
||||
|
||||
def _quat_mul(a, b):
|
||||
w1, x1, y1, z1 = a
|
||||
w2, x2, y2, z2 = b
|
||||
return np.array([
|
||||
w1*w2 - x1*x2 - y1*y2 - z1*z2,
|
||||
w1*x2 + x1*w2 + y1*z2 - z1*y2,
|
||||
w1*y2 - x1*z2 + y1*w2 + z1*x2,
|
||||
w1*z2 + x1*y2 - y1*x2 + z1*w2,
|
||||
], dtype=np.float64)
|
||||
|
||||
|
||||
def _quat_from_up(up_b, yaw):
|
||||
"""Rotation mapping body-frame up to world +z, followed by world yaw."""
|
||||
up = np.asarray(up_b, dtype=np.float64)
|
||||
up /= max(np.linalg.norm(up), 1e-9)
|
||||
target = np.array([0., 0., 1.])
|
||||
dot = float(np.clip(up @ target, -1., 1.))
|
||||
if dot < -0.999999:
|
||||
align = np.array([0., 1., 0., 0.])
|
||||
else:
|
||||
axis = np.cross(up, target)
|
||||
align = np.concatenate(([1. + dot], axis))
|
||||
align /= max(np.linalg.norm(align), 1e-9)
|
||||
yaw_q = np.array([math.cos(yaw / 2), 0., 0., math.sin(yaw / 2)])
|
||||
return _quat_mul(yaw_q, align)
|
||||
|
||||
|
||||
class WorldModelPreview:
|
||||
def __init__(self, checkpoint, env, bridge, field, actor, horizon=DEFAULT_HORIZON,
|
||||
device="cpu", env_action_order=None):
|
||||
self.env = env
|
||||
self.bridge = bridge
|
||||
self.field = field
|
||||
self.actor = actor
|
||||
self.model = load_world_model(checkpoint, device)
|
||||
self.horizon = int(horizon)
|
||||
self.device = device
|
||||
self.env_action_order = list(env_action_order or checkpoint["env_action_order"])
|
||||
self.default_q = bridge.robot.data.default_joint_pos[0, bridge.legs].detach().cpu().numpy()
|
||||
model = env.sim.mj_model
|
||||
free = np.flatnonzero(model.jnt_type == mujoco.mjtJoint.mjJNT_FREE)
|
||||
if len(free) != 1:
|
||||
raise ValueError(f"Expected one free root joint, found {len(free)}")
|
||||
self.root_qadr = int(model.jnt_qposadr[int(free[0])])
|
||||
# ``bridge.legs``/``bridge.wheels`` are entity joint indices, not MuJoCo
|
||||
# joint ids (the free joint is not in the entity list). Use the entity's
|
||||
# own qpos-address table so the ghost writes the intended joints.
|
||||
qadr = bridge.robot.indexing.joint_q_adr.cpu().numpy()
|
||||
self.leg_qadr = [int(qadr[int(j)]) for j in bridge.legs]
|
||||
self.wheel_qadr = [int(qadr[int(j)]) for j in bridge.wheels]
|
||||
self.qpos_template = env.sim.data.qpos[0].detach().cpu().numpy().copy()
|
||||
# root_link_pos_w may be env-local while qpos is world. Store the constant
|
||||
# offset so predicted future positions land in the renderer's world frame.
|
||||
self.world_offset = (self.qpos_template[self.root_qadr:self.root_qadr + 2]
|
||||
- bridge.pose()[0, :2].detach().cpu().numpy())
|
||||
self.tracker = PredictionTracker(self.horizon)
|
||||
|
||||
def _ghost(self, state, pose, horizon):
|
||||
s = state[0].detach().cpu().numpy().astype(np.float64)
|
||||
p = pose[0].detach().cpu().numpy().astype(np.float64)
|
||||
qpos = self.qpos_template.copy()
|
||||
z = float(s[0])
|
||||
if self.field is not None:
|
||||
z += float(self.field.sample(pose[:, :2])[0].detach().cpu())
|
||||
quat = _quat_from_up(-s[1:4], float(p[2]))
|
||||
qpos[self.root_qadr:self.root_qadr + 3] = (
|
||||
p[0] + self.world_offset[0], p[1] + self.world_offset[1], z)
|
||||
qpos[self.root_qadr + 3:self.root_qadr + 7] = quat
|
||||
qpos[self.leg_qadr] = s[LEG_POS] + self.default_q
|
||||
qpos[self.wheel_qadr] = self.qpos_template[self.wheel_qadr]
|
||||
return GhostPrediction(horizon, s, p, qpos)
|
||||
|
||||
@torch.no_grad()
|
||||
def predict(self, history, previous, command, context, bias, state, pose):
|
||||
"""Roll the actor + world model forward from the current real state."""
|
||||
ghosts = []
|
||||
history = history.clone()
|
||||
previous = previous.clone()
|
||||
state = state.clone()
|
||||
pose = pose.clone()
|
||||
scan = None if self.field is None else self.field.scan(state, pose)
|
||||
for i in range(self.horizon):
|
||||
action = self.actor(history)
|
||||
model_context = context if self.field is None else torch.cat((context, scan), -1)
|
||||
if self.field is not None:
|
||||
nxt, delta, _, _ = self.model.forward_all(state, action, previous, model_context)
|
||||
else:
|
||||
nxt, delta = self.model(state, action, previous, model_context)
|
||||
pose = advance_pose(pose, delta)
|
||||
if self.field is not None:
|
||||
scan = self.field.scan(nxt, pose)
|
||||
latest = frame(nxt, action[:, self.env_action_order], command, bias,
|
||||
noise=False, wheel_vel=self.actor.actor_wheel_vel)
|
||||
history = push_history(history, latest, self.actor.history, self.actor.terms)
|
||||
previous = action
|
||||
state = nxt
|
||||
ghosts.append(self._ghost(state, pose, i + 1))
|
||||
return ghosts
|
||||
|
||||
def update_errors(self, step, state, pose):
|
||||
self.tracker.update(step, state[0].detach().cpu().numpy(),
|
||||
pose[0].detach().cpu().numpy())
|
||||
|
||||
def push_errors(self, step, ghosts):
|
||||
self.tracker.push(step, ghosts)
|
||||
|
||||
def summary(self):
|
||||
return self.tracker.summary()
|
||||
|
||||
def clear(self):
|
||||
self.tracker.clear()
|
||||
|
||||
def render(self, ghosts, stride=DEFAULT_STRIDE, rgba=None):
|
||||
"""Actual frame plus transparent ghost geoms, using the env's renderer."""
|
||||
return render_with_ghosts(self.env, ghosts, stride=stride, rgba=rgba)
|
||||
|
||||
|
||||
class PredictionTracker:
|
||||
"""Compare predictions with the state that actually arrives later."""
|
||||
def __init__(self, horizon, keep=200):
|
||||
self.horizon = int(horizon)
|
||||
self.keep = int(keep)
|
||||
self.pending = []
|
||||
self.errors = {h: [] for h in (1, 4, 8, 16, 32) if h <= self.horizon}
|
||||
|
||||
def clear(self):
|
||||
self.pending.clear()
|
||||
for values in self.errors.values():
|
||||
values.clear()
|
||||
|
||||
@staticmethod
|
||||
def _error(pred_state, pred_pose, state, pose):
|
||||
dxy = float(np.linalg.norm(pred_pose[:2] - pose[:2]))
|
||||
dz = float(abs(pred_state[0] - state[0]))
|
||||
dyaw = float(abs(math.atan2(math.sin(pred_pose[2] - pose[2]),
|
||||
math.cos(pred_pose[2] - pose[2]))))
|
||||
pg = pred_state[GRAVITY]
|
||||
ag = state[GRAVITY]
|
||||
pg = pg / max(np.linalg.norm(pg), 1e-9)
|
||||
ag = ag / max(np.linalg.norm(ag), 1e-9)
|
||||
dg = float(math.degrees(math.acos(float(np.clip(pg @ ag, -1., 1.)))))
|
||||
dq = float(np.sqrt(np.mean((pred_state[LEG_POS] - state[LEG_POS]) ** 2)))
|
||||
return dxy, dz, dyaw, dg, dq
|
||||
|
||||
def push(self, step, ghosts):
|
||||
for ghost in ghosts:
|
||||
if ghost.horizon in self.errors:
|
||||
self.pending.append((step + ghost.horizon, ghost.horizon,
|
||||
ghost.state, ghost.pose))
|
||||
|
||||
def update(self, step, state, pose):
|
||||
keep = []
|
||||
for target, horizon, pred_state, pred_pose in self.pending:
|
||||
if target == step:
|
||||
values = self.errors[horizon]
|
||||
values.append(self._error(pred_state, pred_pose, state, pose))
|
||||
if len(values) > self.keep:
|
||||
del values[:-self.keep]
|
||||
elif target > step:
|
||||
keep.append((target, horizon, pred_state, pred_pose))
|
||||
self.pending = keep
|
||||
|
||||
def summary(self):
|
||||
parts = []
|
||||
for horizon in sorted(self.errors):
|
||||
values = self.errors[horizon]
|
||||
if not values:
|
||||
continue
|
||||
mean = np.mean(values, axis=0)
|
||||
parts.append(f"h{horizon} xy {mean[0]:.3f}m yaw {math.degrees(mean[2]):.1f}deg "
|
||||
f"leg {mean[4]:.3f}")
|
||||
return " | ".join(parts) if parts else "collecting..."
|
||||
|
||||
|
||||
def annotate_rgb(rgb, lines, color=(255, 185, 90)):
|
||||
"""Small bottom-left text overlay for MP4 frames (Pillow optional)."""
|
||||
if not lines:
|
||||
return rgb
|
||||
try:
|
||||
from PIL import Image, ImageDraw, ImageFont
|
||||
except ImportError:
|
||||
return rgb
|
||||
image = Image.fromarray(rgb)
|
||||
draw = ImageDraw.Draw(image)
|
||||
font = ImageFont.load_default()
|
||||
y = image.height - 8 - 14 * len(lines)
|
||||
for line in lines:
|
||||
draw.rectangle((6, y - 2, 6 + 8 * len(line), y + 13), fill=(8, 10, 14))
|
||||
draw.text((10, y), line, fill=color, font=font)
|
||||
y += 14
|
||||
return np.asarray(image)
|
||||
|
||||
|
||||
def render_with_ghosts(env, ghosts, stride=DEFAULT_STRIDE, rgba=None):
|
||||
"""Add ghost robot geoms to the existing offscreen scene and render once."""
|
||||
rgb = env.render()
|
||||
renderer = env._offline_renderer
|
||||
if renderer is None or not ghosts:
|
||||
return rgb
|
||||
import mujoco
|
||||
|
||||
model, data = renderer._model, renderer._data
|
||||
scene = renderer.renderer.scene
|
||||
old_qpos, old_qvel = data.qpos.copy(), data.qvel.copy()
|
||||
old_group = renderer._opt.geomgroup.copy()
|
||||
# S10 visual meshes are group 2; group 0 is terrain and group 1 collision
|
||||
# geometry. Adding only group 2 gives robot ghosts without duplicating terrain.
|
||||
visual = np.flatnonzero(model.geom_group == 2)
|
||||
old_geom_rgba = model.geom_rgba[visual].copy()
|
||||
mat_ids = np.unique(model.geom_matid[visual])
|
||||
mat_ids = mat_ids[mat_ids >= 0]
|
||||
old_mat_rgba = {int(i): model.mat_rgba[int(i)].copy() for i in mat_ids}
|
||||
near = np.asarray(rgba if rgba is not None else GHOST_NEAR, dtype=np.float32)
|
||||
far = np.asarray(GHOST_FAR, dtype=np.float32)
|
||||
try:
|
||||
renderer._opt.geomgroup[:] = 0
|
||||
renderer._opt.geomgroup[2] = 1
|
||||
for ghost in ghosts[max(0, int(stride) - 1)::max(1, int(stride))]:
|
||||
# Near future is yellow and stronger; the far tail fades to red.
|
||||
t = ghost.horizon / max(1, ghosts[-1].horizon)
|
||||
tint = near * (1. - t) + far * t
|
||||
for mat_id in mat_ids:
|
||||
model.mat_rgba[int(mat_id)] = tint
|
||||
model.geom_rgba[visual] = tint
|
||||
data.qpos[:] = ghost.qpos
|
||||
data.qvel[:] = 0.
|
||||
mujoco.mj_forward(model, data)
|
||||
mujoco.mjv_addGeoms(model, data, renderer._opt, renderer._pert,
|
||||
renderer._catmask.value, scene)
|
||||
finally:
|
||||
renderer._opt.geomgroup[:] = old_group
|
||||
model.geom_rgba[visual] = old_geom_rgba
|
||||
for mat_id, value in old_mat_rgba.items():
|
||||
model.mat_rgba[mat_id] = value
|
||||
data.qpos[:] = old_qpos
|
||||
data.qvel[:] = old_qvel
|
||||
mujoco.mj_forward(model, data)
|
||||
return renderer.render()
|
||||
Vendored
+5
@@ -0,0 +1,5 @@
|
||||
"""S10 flat and terrain environment factories; no implicit task registration.
|
||||
|
||||
Use flat.s10_velocity_flat_env_cfg or
|
||||
terrain.s10_velocity_terrain_env_cfg to construct an environment.
|
||||
"""
|
||||
Vendored
+148
@@ -0,0 +1,148 @@
|
||||
"""Small flat-ground velocity task for the S10 mixed leg/wheel controller.
|
||||
|
||||
This environment is independent of the chase-specific SuperMimic learner.
|
||||
"""
|
||||
|
||||
from copy import deepcopy
|
||||
|
||||
import torch
|
||||
|
||||
from mjlab.envs import mdp
|
||||
from mjlab.envs.mdp.actions import JointPositionActionCfg, JointVelocityActionCfg
|
||||
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.tasks.velocity import mdp as velocity_mdp
|
||||
from mjlab.tasks.velocity.velocity_env_cfg import make_velocity_env_cfg
|
||||
from mjlab.utils.noise import UniformNoiseCfg
|
||||
|
||||
from s10_supermimic.robot.config import (
|
||||
LEG_JOINT_NAMES, WHEEL_JOINT_NAMES, WHEEL_GEOM_NAMES, get_s10_robot_cfg,
|
||||
)
|
||||
|
||||
ACTOR_HISTORY = 5
|
||||
|
||||
|
||||
def leg_posture_cost(env, asset_cfg: SceneEntityCfg) -> torch.Tensor:
|
||||
d = env.scene[asset_cfg.name].data
|
||||
return (d.joint_pos[:, asset_cfg.joint_ids]
|
||||
- d.default_joint_pos[:, asset_cfg.joint_ids]).square().sum(-1)
|
||||
|
||||
|
||||
def standstill_cost(env, wheel_cfg: SceneEntityCfg) -> torch.Tensor:
|
||||
"""Extra braking at zero command, including wheels spinning in place.
|
||||
|
||||
This suppresses velocity; it does not fix a world-frame position or heading.
|
||||
"""
|
||||
command = env.command_manager.get_command("twist")
|
||||
standing = command.square().sum(-1) < 1e-6
|
||||
d = env.scene["robot"].data
|
||||
cost = (d.root_link_lin_vel_b.square().sum(-1)
|
||||
+ d.root_link_ang_vel_b.square().sum(-1)
|
||||
+ 0.01 * d.joint_vel[:, wheel_cfg.joint_ids].square().sum(-1))
|
||||
return standing * cost
|
||||
|
||||
|
||||
def s10_velocity_flat_env_cfg(play: bool = False, history: int = ACTOR_HISTORY,
|
||||
actor_wheel_vel: bool = True,
|
||||
hipy: float | None = None,
|
||||
hipx: float | None = None,
|
||||
knee: float | None = None,
|
||||
leg_kp: float = 80.0,
|
||||
leg_kd: float = 2.0):
|
||||
if history < 1:
|
||||
raise ValueError("history must be at least 1")
|
||||
cfg = make_velocity_env_cfg()
|
||||
cfg.scene.entities = {"robot": get_s10_robot_cfg(
|
||||
hipy=hipy, hipx=hipx, knee=knee, leg_kp=leg_kp, leg_kd=leg_kd)}
|
||||
cfg.scene.terrain.terrain_type = "plane"
|
||||
cfg.scene.terrain.terrain_generator = None
|
||||
cfg.scene.sensors = ()
|
||||
cfg.scene.num_envs = 1 # Training entry point must opt into larger batches.
|
||||
cfg.sim.mujoco.cone = "elliptic"
|
||||
cfg.sim.mujoco.impratio = 10
|
||||
cfg.sim.mujoco.ccd_iterations = 50
|
||||
cfg.viewer.body_name = "base_link"
|
||||
cfg.viewer.distance = 2.0
|
||||
|
||||
legs = SceneEntityCfg("robot", joint_names=LEG_JOINT_NAMES, preserve_order=True)
|
||||
wheels = SceneEntityCfg("robot", joint_names=WHEEL_JOINT_NAMES, preserve_order=True)
|
||||
cfg.actions = {
|
||||
# TODO(s10): flat pretraining keeps the smaller flat scales (0.25 / wheel 15).
|
||||
# Terrain revision 4 overrides these to 0.25/0.6/0.8 and wheel 35. If you
|
||||
# warm-start from a flat actor, consider matching the rev4 scales here.
|
||||
"joint_pos": JointPositionActionCfg(
|
||||
entity_name="robot", actuator_names=LEG_JOINT_NAMES,
|
||||
scale=0.25, use_default_offset=True,
|
||||
),
|
||||
"wheel_vel": JointVelocityActionCfg(
|
||||
entity_name="robot", actuator_names=WHEEL_JOINT_NAMES,
|
||||
scale=15.0, use_default_offset=False, clip={".*": (-35.0, 35.0)},
|
||||
),
|
||||
}
|
||||
# 53 values/frame (v4 / new v7, with measured wheel speed) or legacy 49
|
||||
# values/frame (old v7). The critic and world model always keep wheel speed.
|
||||
terms = {
|
||||
"base_ang_vel": ObservationTermCfg(
|
||||
func=mdp.base_ang_vel, noise=UniformNoiseCfg(n_min=-0.2, n_max=0.2)),
|
||||
"projected_gravity": ObservationTermCfg(
|
||||
func=mdp.projected_gravity, noise=UniformNoiseCfg(n_min=-0.05, n_max=0.05)),
|
||||
"joint_pos": ObservationTermCfg(
|
||||
func=mdp.joint_pos_rel, params={"asset_cfg": legs, "biased": True},
|
||||
noise=UniformNoiseCfg(n_min=-0.01, n_max=0.01)),
|
||||
"joint_vel": ObservationTermCfg(
|
||||
func=mdp.joint_vel_rel, params={"asset_cfg": legs}, scale=0.05,
|
||||
noise=UniformNoiseCfg(n_min=-1.5, n_max=1.5)),
|
||||
"wheel_vel": ObservationTermCfg(
|
||||
func=mdp.joint_vel_rel, params={"asset_cfg": wheels}, scale=0.05,
|
||||
noise=UniformNoiseCfg(n_min=-1.5, n_max=1.5)),
|
||||
"actions": ObservationTermCfg(func=mdp.last_action),
|
||||
"command": ObservationTermCfg(
|
||||
func=mdp.generated_commands, params={"command_name": "twist"}),
|
||||
}
|
||||
critic = deepcopy(terms)
|
||||
critic["joint_pos"].params["biased"] = False
|
||||
critic["base_lin_vel"] = ObservationTermCfg(func=mdp.base_lin_vel)
|
||||
if not actor_wheel_vel:
|
||||
terms.pop("wheel_vel")
|
||||
cfg.observations = {
|
||||
"actor": ObservationGroupCfg(terms=terms, concatenate_terms=True,
|
||||
enable_corruption=not play, history_length=history),
|
||||
"critic": ObservationGroupCfg(terms=critic, concatenate_terms=True,
|
||||
enable_corruption=False),
|
||||
}
|
||||
command = cfg.commands["twist"]
|
||||
command.heading_command = False
|
||||
command.ranges.heading = None
|
||||
command.rel_heading_envs = 0.0
|
||||
command.rel_standing_envs = 0.2
|
||||
command.ranges.lin_vel_x = (-1.0, 1.0)
|
||||
command.ranges.lin_vel_y = (-0.2, 0.2)
|
||||
command.ranges.ang_vel_z = (-0.5, 0.5)
|
||||
cfg.curriculum = {}
|
||||
cfg.events.pop("push_robot", None)
|
||||
cfg.events["base_com"].params["asset_cfg"].body_names = ("base_link",)
|
||||
cfg.events["foot_friction"].params["asset_cfg"].geom_names = WHEEL_GEOM_NAMES
|
||||
cfg.events["encoder_bias"].params["asset_cfg"] = legs
|
||||
# Wheel rolling must not incur quadruped foot slip / swing height costs.
|
||||
cfg.rewards = {
|
||||
"track_linear_velocity": cfg.rewards["track_linear_velocity"],
|
||||
"track_angular_velocity": cfg.rewards["track_angular_velocity"],
|
||||
"upright": RewardTermCfg(func=velocity_mdp.upright, weight=1.0,
|
||||
params={"std": 0.45, "asset_cfg": SceneEntityCfg(
|
||||
"robot", body_names=("base_link",))}),
|
||||
"leg_posture": RewardTermCfg(func=leg_posture_cost, weight=-0.2,
|
||||
params={"asset_cfg": legs}),
|
||||
"standstill": RewardTermCfg(func=standstill_cost, weight=-2.0,
|
||||
params={"wheel_cfg": wheels}),
|
||||
"action_rate_l2": RewardTermCfg(func=mdp.action_rate_l2, weight=-0.05),
|
||||
"joint_torques_l2": RewardTermCfg(func=mdp.joint_torques_l2, weight=-2e-5),
|
||||
"dof_pos_limits": RewardTermCfg(func=mdp.joint_pos_limits, weight=-1.0,
|
||||
params={"asset_cfg": legs}),
|
||||
"terminated": RewardTermCfg(func=mdp.is_terminated, weight=-20.0),
|
||||
}
|
||||
cfg.terminations.pop("out_of_terrain_bounds", None)
|
||||
if play:
|
||||
for event in ("base_com", "encoder_bias", "foot_friction"):
|
||||
cfg.events.pop(event, None)
|
||||
return cfg
|
||||
Vendored
+271
@@ -0,0 +1,271 @@
|
||||
"""Static S10 terrain courses."""
|
||||
from dataclasses import dataclass
|
||||
import math
|
||||
import mujoco
|
||||
import numpy as np
|
||||
import torch
|
||||
from mjlab.terrains.terrain_generator import SubTerrainCfg, TerrainOutput, TerrainGeometry, TerrainGeneratorCfg
|
||||
from mjlab.managers import CurriculumTermCfg, TerminationTermCfg, RewardTermCfg
|
||||
from .flat import s10_velocity_flat_env_cfg
|
||||
from s10_supermimic.robot.config import (
|
||||
V7_HIPX, V7_HIPY, V7_KNEE, V7_LEG_KP,
|
||||
WHEEL_BODY_NAMES, WHEEL_GEOM_NAMES, KNEE_GEOM_NAMES)
|
||||
|
||||
STEP_HEIGHTS = (0.02, 0.04, 0.08, 0.12, 0.16, 0.20)
|
||||
|
||||
@dataclass
|
||||
class CourseCfg(SubTerrainCfg):
|
||||
kind: str = "rubble"
|
||||
bidirectional: bool = False
|
||||
stair_tread: float = .45
|
||||
|
||||
def function(self, difficulty, spec, rng):
|
||||
level = min(5, max(0, round(difficulty * 5)))
|
||||
h = STEP_HEIGHTS[level]
|
||||
sx, sy = self.size
|
||||
if self.bidirectional:
|
||||
sx = 8. # Mirror the base course into a 16 m square.
|
||||
body, geoms = spec.body("terrain"), []
|
||||
def box(x, y, z, a, b, c, yaw=0., pitch=0., shade=.45):
|
||||
# Trim the source half at its spawn plane before mirroring, avoiding
|
||||
# coincident contact surfaces in the central floor/down-stair landing.
|
||||
if self.bidirectional and yaw == 0. and pitch == 0.:
|
||||
left, right = max(1., x-a/2), x+a/2
|
||||
if right <= left:
|
||||
return
|
||||
x, a = (left+right)/2, right-left
|
||||
# Composition Rz(yaw) Ry(pitch); z is centre, dimensions are full lengths.
|
||||
cy, ss, cp, sp = math.cos(yaw/2), math.sin(yaw/2), math.cos(pitch/2), math.sin(pitch/2)
|
||||
placements = [(x, (cy*cp,-ss*sp,cy*sp,ss*cp))]
|
||||
if self.bidirectional:
|
||||
# Reflect about x=8: yaw and pitch change sign, roll stays unchanged.
|
||||
placements = [(x+7., (cy*cp,-ss*sp,cy*sp,ss*cp)),
|
||||
(9.-x, (cy*cp,-ss*sp,-cy*sp,-ss*cp))]
|
||||
for xx, quat in placements:
|
||||
geom = body.add_geom(type=mujoco.mjtGeom.mjGEOM_BOX, pos=(xx,y,z),
|
||||
size=(a/2,b/2,c/2), quat=quat)
|
||||
geoms.append(TerrainGeometry(geom=geom, color=(shade,shade,.9*shade,1.)))
|
||||
box(sx/2,sy/2,-.1,sx,sy,.2,shade=.6)
|
||||
origin = np.array([1., sy/2, 0.])
|
||||
if self.kind in ("up", "down"):
|
||||
n, tread = 5, self.stair_tread
|
||||
if self.kind == "up":
|
||||
# Overlapping foundation boxes make solid risers and a long top landing.
|
||||
for k in range(n):
|
||||
start = 2. + k*tread
|
||||
box((start+sx)/2,sy/2,(k+1)*h/2,sx-start,sy,(k+1)*h)
|
||||
else:
|
||||
origin[2] = n*h
|
||||
for k in range(n):
|
||||
end = 2. + k*tread
|
||||
height = (n-k)*h
|
||||
box(end/2,sy/2,height/2,end,sy,height)
|
||||
elif self.kind == "platform":
|
||||
for x in (2.6,4.6):
|
||||
box(x,sy/2,h/2,.9,sy,h)
|
||||
elif self.kind == "slope":
|
||||
# A continuous wedge made by a tilted slab with its low edge buried.
|
||||
angle = math.radians(3+12*difficulty)
|
||||
length = 3.
|
||||
thickness = .16
|
||||
center_z = length/2*math.sin(angle)-thickness/2*math.cos(angle)
|
||||
box(3.5,sy/2,center_z,length,sy,thickness,pitch=-angle)
|
||||
end = 3.5+length/2*math.cos(angle)
|
||||
top = length*math.sin(angle)
|
||||
box((end+sx)/2,sy/2,top/2,sx-end,sy,top)
|
||||
elif self.kind == "rubble":
|
||||
# Staggered rows cover the lane. Static overlapping rocks, never loose bodies.
|
||||
# Heights of all layers together are capped by h (2--20 cm).
|
||||
for ix, x in enumerate(np.arange(2.15,5.9,.48)):
|
||||
for y in np.arange(.35,sy-.2,.52):
|
||||
a,b = rng.uniform(.35,.65,2)
|
||||
height = h*rng.uniform(.35,.7)
|
||||
xx,yy = x+rng.uniform(-.12,.12),y+rng.uniform(-.12,.12)
|
||||
yaw = rng.uniform(-math.pi,math.pi)
|
||||
box(xx,yy,height/2,a,b,height,yaw=yaw,shade=rng.uniform(.3,.55))
|
||||
if rng.random() < .25+.35*difficulty:
|
||||
cap = h-height
|
||||
# Deliberately overlapping cap; bounded surface height even in stacks.
|
||||
pitch = rng.uniform(-.15,.15)*difficulty
|
||||
half_z = (cap*math.cos(pitch)+a*.65*abs(math.sin(pitch)))/2
|
||||
box(xx+.04,yy-.03,h-half_z,a*.65,b*.65,cap,
|
||||
yaw=yaw+.6,pitch=pitch,shade=.55)
|
||||
elif self.kind != "flat":
|
||||
raise ValueError(self.kind)
|
||||
if self.bidirectional:
|
||||
origin[0] = 8.
|
||||
return TerrainOutput(origin=origin, geometries=geoms)
|
||||
|
||||
|
||||
def terrain_linear_tracking(env):
|
||||
velocity = env.scene["robot"].data.root_link_lin_vel_b
|
||||
command = env.command_manager.get_command("twist")
|
||||
return torch.exp(-(velocity[:,:2]-command[:,:2]).square().sum(-1)/.25)
|
||||
|
||||
|
||||
def terrain_yaw_tracking(env):
|
||||
velocity = env.scene["robot"].data.root_link_ang_vel_b
|
||||
command = env.command_manager.get_command("twist")
|
||||
return torch.exp(-(velocity[:,2]-command[:,2]).square()/.5)
|
||||
|
||||
|
||||
def course_exit(env):
|
||||
offset = env.scene["robot"].data.root_link_pos_w - env.scene.env_origins
|
||||
return (offset[:,0] > 5.4) | (offset[:,0] < -.65) | (offset[:,1].abs() > 2.7)
|
||||
|
||||
|
||||
def course_curriculum(env, env_ids):
|
||||
terrain = env.scene.terrain
|
||||
offset = env.scene["robot"].data.root_link_pos_w[env_ids] - env.scene.env_origins[env_ids]
|
||||
active = env.episode_length_buf[env_ids] > 0
|
||||
fallen = env.termination_manager.terminated[env_ids]
|
||||
success = active & ~fallen & (offset[:,0] > 5.4) & (offset[:,1].abs() < 2.7)
|
||||
failure = active & (fallen | (offset[:,0] < -.65) | (offset[:,1].abs() > 2.7))
|
||||
terrain.terrain_levels[env_ids] = (terrain.terrain_levels[env_ids] + success.long() - failure.long()).clamp(0,5)
|
||||
terrain.env_origins[env_ids] = terrain.terrain_origins[terrain.terrain_levels[env_ids], terrain.terrain_types[env_ids]]
|
||||
return terrain.terrain_levels.float().mean()
|
||||
|
||||
|
||||
def bidirectional_exit(env):
|
||||
offset = env.scene["robot"].data.root_link_pos_w - env.scene.env_origins
|
||||
return (offset[:,0].abs() > 6.5) | (offset[:,1].abs() > 6.5)
|
||||
|
||||
|
||||
def bidirectional_curriculum(env, env_ids):
|
||||
terrain = env.scene.terrain
|
||||
offset = env.scene["robot"].data.root_link_pos_w[env_ids] - env.scene.env_origins[env_ids]
|
||||
active = env.episode_length_buf[env_ids] > 0
|
||||
fallen = env.termination_manager.terminated[env_ids]
|
||||
# Both course directions count. Lateral exits are neutral, not evidence
|
||||
# of clearing steps; negative commands no longer count as failure.
|
||||
success = active & ~fallen & (offset[:,0].abs() > 5.4) & (offset[:,1].abs() < 6.5)
|
||||
failure = active & fallen
|
||||
terrain.terrain_levels[env_ids] = (terrain.terrain_levels[env_ids] + success.long() - failure.long()).clamp(0,5)
|
||||
terrain.env_origins[env_ids] = terrain.terrain_origins[terrain.terrain_levels[env_ids], terrain.terrain_types[env_ids]]
|
||||
return terrain.terrain_levels.float().mean()
|
||||
|
||||
|
||||
def s10_velocity_terrain_env_cfg(play=False, history=5, seed=0, revision=1,
|
||||
actor_wheel_vel=None):
|
||||
if revision not in (1, 2, 3, 4, 5, 6, 7):
|
||||
raise ValueError("Terrain revision must be 1, 2, 3, 4, 5, 6 or 7")
|
||||
if actor_wheel_vel is None:
|
||||
actor_wheel_vel = revision <= 4 or revision >= 7
|
||||
cfg = s10_velocity_flat_env_cfg(
|
||||
play=play, history=history, actor_wheel_vel=actor_wheel_vel,
|
||||
hipy=(V7_HIPY if revision >= 7 else None),
|
||||
hipx=(V7_HIPX if revision >= 7 else None),
|
||||
knee=(V7_KNEE if revision >= 7 else None),
|
||||
leg_kp=(V7_LEG_KP if revision >= 7 else 80.0))
|
||||
cfg.sim.nconmax = 128
|
||||
cfg.sim.njmax = 512
|
||||
cfg.scene.terrain.terrain_type = "generator"
|
||||
cfg.scene.terrain.max_init_terrain_level = 0
|
||||
cfg.scene.terrain.terrain_generator = TerrainGeneratorCfg(
|
||||
seed=seed, size=(8.,8.), num_rows=6, curriculum=True, color_scheme="none",
|
||||
sub_terrains={kind: CourseCfg(kind=kind, proportion=weight) for kind,weight in
|
||||
(("flat",.10),("up",.25),("down",.20),("platform",.15),("slope",.10),("rubble",.20))})
|
||||
cfg.events["reset_base"].params["pose_range"].update(x=(-.1,.1),y=(-.1,.1),yaw=(-.08,.08))
|
||||
command = cfg.commands["twist"]
|
||||
command.ranges.lin_vel_x = (.2,.45)
|
||||
command.ranges.lin_vel_y = (0.,0.)
|
||||
command.ranges.ang_vel_z = (-.05,.05)
|
||||
command.rel_standing_envs = .1
|
||||
cfg.episode_length_s = 35.
|
||||
cfg.rewards["leg_posture"].weight = -.05
|
||||
cfg.rewards["track_linear_velocity"].func = terrain_linear_tracking
|
||||
cfg.rewards["track_linear_velocity"].params = {}
|
||||
cfg.rewards["track_angular_velocity"].func = terrain_yaw_tracking
|
||||
cfg.rewards["track_angular_velocity"].params = {}
|
||||
cfg.rewards["upright"].params["std"] = .6
|
||||
if revision >= 2:
|
||||
from s10_supermimic.core.terrain_posture import env_posture_cost
|
||||
cfg.actions["joint_pos"].scale = {".*_hipx_joint": .25, ".*_hipy_joint": .6, ".*_knee_joint": .8}
|
||||
cfg.rewards["leg_posture"].func = env_posture_cost
|
||||
cfg.rewards["leg_posture"].weight = -.4
|
||||
cfg.curriculum = {} if play else {"course": CurriculumTermCfg(func=course_curriculum)}
|
||||
cfg.terminations["course_exit"] = TerminationTermCfg(func=course_exit, time_out=True)
|
||||
if revision >= 3:
|
||||
command.ranges.lin_vel_x = (-2., 2.)
|
||||
command.ranges.lin_vel_y = (-1., 1.)
|
||||
command.ranges.ang_vel_z = (-3., 3.)
|
||||
command.rel_forward_envs = 0.
|
||||
command.rel_world_envs = 0.
|
||||
cfg.actions["wheel_vel"].scale = 35.
|
||||
gen = cfg.scene.terrain.terrain_generator
|
||||
gen.size = (16., 16.)
|
||||
for course in gen.sub_terrains.values():
|
||||
course.bidirectional = True
|
||||
cfg.curriculum = {} if play else {"course": CurriculumTermCfg(func=bidirectional_curriculum)}
|
||||
cfg.terminations["course_exit"] = TerminationTermCfg(func=bidirectional_exit, time_out=True)
|
||||
if revision >= 4:
|
||||
from mjlab.sensor import ContactMatch, ContactSensorCfg
|
||||
cfg.scene.sensors = tuple(cfg.scene.sensors or ()) + (ContactSensorCfg(
|
||||
name="wheel_ground_contact", primary=ContactMatch(mode="body",
|
||||
pattern=tuple(WHEEL_BODY_NAMES), entity="robot"),
|
||||
secondary=ContactMatch(mode="body", pattern="terrain"),
|
||||
fields=("found","force"), reduce="netforce", num_slots=1),)
|
||||
from s10_supermimic.core.rolling_rewards import env_rolling_cost
|
||||
for course in cfg.scene.terrain.terrain_generator.sub_terrains.values():
|
||||
course.stair_tread = .30
|
||||
for name, weight in (("unnecessary_wheel_lift", -.2), ("rolling_leg_motion", -.12)):
|
||||
cfg.rewards[name] = RewardTermCfg(func=env_rolling_cost, weight=weight, params={"term":name})
|
||||
if not play:
|
||||
from s10_supermimic.core.mixed_commands import MixedVelocityCommandCfg
|
||||
cfg.commands["twist"] = MixedVelocityCommandCfg(**vars(command))
|
||||
if revision >= 5:
|
||||
from s10_supermimic.core.obstacle_rewards import env_obstacle_progress
|
||||
from s10_supermimic.core.stair_curriculum import curriculum
|
||||
cfg.rewards['obstacle_progress'] = RewardTermCfg(func=env_obstacle_progress, weight=.5)
|
||||
if not play:
|
||||
cfg.curriculum = {'course': CurriculumTermCfg(func=curriculum)}
|
||||
cfg.commands['twist'].hold_stair_command = True
|
||||
for kind, weight in zip(('flat','up','down','platform','slope','rubble'), (.1,.35,.15,.1,.1,.2)):
|
||||
cfg.scene.terrain.terrain_generator.sub_terrains[kind].proportion = weight
|
||||
if revision == 6:
|
||||
from s10_supermimic.core.wheel_terrain_rewards import env_wheel_term
|
||||
for name in ('leg_posture', 'unnecessary_wheel_lift', 'rolling_leg_motion'):
|
||||
cfg.rewards[name].func = env_wheel_term
|
||||
cfg.rewards[name].params = {'term': name}
|
||||
cfg.rewards['wheel_obstacle_clearance'] = RewardTermCfg(
|
||||
func=env_wheel_term, weight=-.25, params={'term':'wheel_obstacle_clearance'})
|
||||
if revision >= 7:
|
||||
from s10_supermimic.core.collision import illegal_nonwheel_contact
|
||||
from s10_supermimic.core.mixed_commands import MixedVelocityCommandCfg
|
||||
from s10_supermimic.core.wheel_intrusion import (
|
||||
env_pair_body_intrusion, env_wheel_body_intrusion)
|
||||
# V7 = V5 + non-wheel collisions; experimental V6 shaping is independent.
|
||||
cfg.scene.sensors += (ContactSensorCfg(
|
||||
name='nonwheel_ground_contact',
|
||||
primary=ContactMatch(mode='geom', pattern='.*_collision', entity='robot',
|
||||
exclude=WHEEL_GEOM_NAMES + KNEE_GEOM_NAMES),
|
||||
secondary=ContactMatch(mode='body', pattern='terrain'),
|
||||
fields=('found','force'), reduce='netforce', num_slots=1,
|
||||
history_length=cfg.decimation),)
|
||||
cfg.terminations['illegal_nonwheel_contact'] = TerminationTermCfg(
|
||||
func=illegal_nonwheel_contact, params={'force_threshold':50.})
|
||||
# Safe low-speed all-terrain envelope. Fixed for training and play so the
|
||||
# validation videos use exactly the command distribution seen in training.
|
||||
command = cfg.commands['twist']
|
||||
command.ranges.lin_vel_x = (-1.5, 1.5)
|
||||
command.ranges.lin_vel_y = (-.8, .8)
|
||||
command.ranges.ang_vel_z = (-1.2, 1.2)
|
||||
if not isinstance(command, MixedVelocityCommandCfg):
|
||||
command = MixedVelocityCommandCfg(**vars(command))
|
||||
command.deadzone = .05
|
||||
command.hold_stair_command = not play
|
||||
cfg.commands['twist'] = command
|
||||
# V7 initialization: small x/y/yaw only, no random roll/pitch.
|
||||
# Always active, independent of terrain relief: keep front/rear wheels out
|
||||
# of the opposite two-thirds of the body projection.
|
||||
cfg.rewards['wheel_body_intrusion'] = RewardTermCfg(
|
||||
func=env_wheel_body_intrusion, weight=-1.0)
|
||||
# Wider 5/6 restriction: both same-side legs may not sit under the body
|
||||
# together for more than 0.5 s. One leg passing through is free.
|
||||
cfg.rewards['pair_body_intrusion'] = RewardTermCfg(
|
||||
func=env_pair_body_intrusion, weight=-1.0)
|
||||
# Stairs need deep flexion. Keep hipx constrained (2x in posture_cost) but
|
||||
# let hipy/knee deviate almost freely once local step relief is large.
|
||||
cfg.rewards['leg_posture'].params = {
|
||||
**cfg.rewards['leg_posture'].params, 'relief_floor': .02}
|
||||
return cfg
|
||||
@@ -0,0 +1,28 @@
|
||||
BSD 3-Clause License
|
||||
|
||||
Copyright (c) 2025, DEEPRoboticsLab
|
||||
|
||||
Redistribution and use in source and binary forms, with or without
|
||||
modification, are permitted provided that the following conditions are met:
|
||||
|
||||
1. Redistributions of source code must retain the above copyright notice, this
|
||||
list of conditions and the following disclaimer.
|
||||
|
||||
2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
this list of conditions and the following disclaimer in the documentation
|
||||
and/or other materials provided with the distribution.
|
||||
|
||||
3. Neither the name of the copyright holder nor the names of its
|
||||
contributors may be used to endorse or promote products derived from
|
||||
this software without specific prior written permission.
|
||||
|
||||
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
@@ -0,0 +1 @@
|
||||
"""Deep Robotics S10 model."""
|
||||
@@ -0,0 +1,115 @@
|
||||
"""S10 asset adapter: twelve position-controlled legs and four velocity wheels.
|
||||
|
||||
The MJCF and default pose come from the Deep Robotics S10 deployment SDK
|
||||
(``S10_description/s10_mjcf``). The copied XML has its torque-motor block
|
||||
removed, so mjlab can add native position/velocity actuators. v4 keeps
|
||||
the SDK/IsaacLab leg gains (kp=80, kd=2, effort=50 N·m); v7 uses a softer
|
||||
kp=20 for deployment safety. Wheels use velocity damping 0.8, effort=14 N·m.
|
||||
"""
|
||||
|
||||
import math
|
||||
from pathlib import Path
|
||||
|
||||
import mujoco
|
||||
|
||||
from mjlab.actuator import (
|
||||
BuiltinPositionActuatorCfg,
|
||||
BuiltinVelocityActuatorCfg,
|
||||
)
|
||||
from mjlab.entity import EntityArticulationInfoCfg, EntityCfg
|
||||
from mjlab.utils.spec_config import CollisionCfg
|
||||
|
||||
S10_XML = Path(__file__).parent / "model" / "s10.xml"
|
||||
LEGS = ("fl", "fr", "hl", "hr")
|
||||
LEG_JOINT_NAMES = tuple(
|
||||
f"{leg}_{joint}_joint" for leg in LEGS for joint in ("hipx", "hipy", "knee")
|
||||
)
|
||||
WHEEL_JOINT_NAMES = tuple(f"{leg}_wheel_joint" for leg in LEGS)
|
||||
WHEEL_BODY_NAMES = tuple(f"{leg}_wheel" for leg in LEGS)
|
||||
WHEEL_GEOM_NAMES = tuple(f"{leg}_wheel_collision" for leg in LEGS)
|
||||
# Knee-joint housing, shin link and knee hub; v7 excludes these from the
|
||||
# non-wheel contact termination so stair shin/knee grazing does not end runs.
|
||||
KNEE_GEOM_NAMES = tuple(
|
||||
f"{leg}_{part}_collision" for leg in LEGS
|
||||
for part in ("knee_joint", "shin", "knee_hub"))
|
||||
|
||||
# v4 keeps the SDK neutral pose. v7 uses a more crouched, stair-friendly
|
||||
# pose closer to the gait it needs (hipx 0.05, hipy ±25°, knee 1.0) and a
|
||||
# softer leg actuator (kp=20) for safer deployment. Knee 1.0 keeps the wheel
|
||||
# centres outside the body projection even with the ±25° splay.
|
||||
DEFAULT_HIPY = 0.35
|
||||
V7_HIPX = 0.05
|
||||
V7_HIPY = math.radians(25.0)
|
||||
V7_KNEE = 1.0
|
||||
V7_LEG_KP = 20.0
|
||||
|
||||
|
||||
def get_spec() -> mujoco.MjSpec:
|
||||
spec = mujoco.MjSpec.from_file(str(S10_XML))
|
||||
spec.modelname = "s10"
|
||||
# Scene solver options are owned by the task, not the attached robot.
|
||||
spec.option.cone = mujoco.mjtCone.mjCONE_PYRAMIDAL
|
||||
spec.option.impratio = 1.0
|
||||
# Name every collision so domain randomisation can select actual contacts.
|
||||
for i, geom in enumerate(spec.geoms):
|
||||
if geom.contype or geom.conaffinity:
|
||||
# Warp MULTICCD cannot handle nonzero margins on these cylinder/box pairs.
|
||||
geom.margin = 0.0
|
||||
if not geom.name:
|
||||
geom.name = f"body_{i}_collision"
|
||||
return spec
|
||||
|
||||
|
||||
def get_s10_robot_cfg(hipy: float | None = None, hipx: float | None = None,
|
||||
knee: float | None = None, leg_kp: float = 80.0,
|
||||
leg_kd: float = 2.0) -> EntityCfg:
|
||||
"""Return a fresh S10 config with mjlab native position/velocity actuators.
|
||||
|
||||
``hipx``, ``hipy`` and ``knee`` are absolute default joint angles. Omit them
|
||||
for the legacy v4 pose; v7 passes its crouched stair pose. ``leg_kp``/``kd``
|
||||
let v7 use a softer actuator without changing v4 checkpoints. The joint
|
||||
action offset and posture penalty both read ``default_joint_pos``, so these
|
||||
values move the action target and the penalised posture together.
|
||||
"""
|
||||
hipx = 0.1 if hipx is None else float(hipx)
|
||||
hipy = DEFAULT_HIPY if hipy is None else float(hipy)
|
||||
knee = 0.65 if knee is None else float(knee)
|
||||
return EntityCfg(
|
||||
spec_fn=get_spec,
|
||||
init_state=EntityCfg.InitialStateCfg(
|
||||
# Default leg pose puts the wheel bottom at ~0.434 m; 0.44 m gives a
|
||||
# small margin before the reset z offset is applied.
|
||||
pos=(0.0, 0.0, 0.44),
|
||||
joint_pos={
|
||||
"fl_hipx_joint": hipx, "hl_hipx_joint": hipx,
|
||||
"fr_hipx_joint": -hipx, "hr_hipx_joint": -hipx,
|
||||
"fl_hipy_joint": -hipy, "fr_hipy_joint": -hipy,
|
||||
"hl_hipy_joint": hipy, "hr_hipy_joint": hipy,
|
||||
"fl_knee_joint": knee, "fr_knee_joint": knee,
|
||||
"hl_knee_joint": -knee, "hr_knee_joint": -knee,
|
||||
".*_wheel_joint": 0.0,
|
||||
},
|
||||
joint_vel={".*": 0.0},
|
||||
),
|
||||
collisions=(CollisionCfg(
|
||||
geom_names_expr=(".*_collision",), contype=1, conaffinity=1,
|
||||
condim={".*_wheel_collision": 3, "body_.*_collision": 1, ".*": 1},
|
||||
priority={".*_wheel_collision": 1, "body_.*_collision": 0, ".*": 0},
|
||||
),),
|
||||
articulation=EntityArticulationInfoCfg(
|
||||
actuators=(
|
||||
# TODO(s10): the SDK/IsaacLab S10 config also randomizes actuator delay
|
||||
# (0-5 physics steps) and PD gains. v7 uses a softer kp=20 for safe
|
||||
# deployment; v4 keeps the SDK kp=80 recipe.
|
||||
BuiltinPositionActuatorCfg(
|
||||
target_names_expr=LEG_JOINT_NAMES,
|
||||
stiffness=float(leg_kp), damping=float(leg_kd), effort_limit=50.0,
|
||||
),
|
||||
BuiltinVelocityActuatorCfg(
|
||||
target_names_expr=WHEEL_JOINT_NAMES,
|
||||
damping=0.8, effort_limit=14.0, armature=0.0005,
|
||||
),
|
||||
),
|
||||
soft_joint_pos_limit_factor=0.9,
|
||||
),
|
||||
)
|
||||
@@ -0,0 +1,29 @@
|
||||
S10 model adapted from the Deep Robotics deployment SDK
|
||||
(`S10_description/s10_mjcf/S10.xml`). Only the 17 mesh files referenced by
|
||||
`s10.xml` are included. The source XML is BSD-3-Clause, see `../LICENSE_s10`.
|
||||
|
||||
Local changes:
|
||||
|
||||
- Removed the world floor and light so mjlab can attach the robot to its own
|
||||
terrain.
|
||||
- Removed the torque `<motor>` actuator block; mjlab adds native position
|
||||
actuators for the 12 leg joints and native velocity actuators for the 4
|
||||
wheels.
|
||||
- Added names to the four wheel collision cylinders
|
||||
(`fl/fr/hl/hr_wheel_collision`) so wheel friction randomisation and contact
|
||||
sensors can address them.
|
||||
- Removed the `<sensor>` block. The SDK XML has unnamed `accelerometer` and
|
||||
`gyro` elements, which mjlab tries to wrap as a sensor with an empty name and
|
||||
fails with `KeyError: Invalid name ''` during scene initialization. The
|
||||
training actor/world model read MuJoCo data directly and do not use these XML
|
||||
sensors.
|
||||
- `robot/config.py` names remaining unnamed collision geoms, clears collision
|
||||
margins for Warp MULTICCD compatibility, and delegates solver options to the
|
||||
task.
|
||||
- `meshdir` changed from `../meshes/` to `meshes/`.
|
||||
|
||||
Default pose, joint limits, link geometry and inertial values are unchanged
|
||||
from the SDK model. Actuator gains follow the SDK/IsaacLab S10 configuration:
|
||||
legs kp=80, kd=2, effort=50 N·m; wheels damping=0.8, effort=14 N·m.
|
||||
These values are taken from the provided deployment/training projects and are
|
||||
not independently hardware-calibrated here.
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,176 @@
|
||||
<?xml version='1.0' encoding='utf-8'?>
|
||||
<mujoco model="S10">
|
||||
<default>
|
||||
<default class="robot">
|
||||
<default class="motor">
|
||||
<joint />
|
||||
<motor />
|
||||
</default>
|
||||
<default class="visual">
|
||||
<geom material="default_material" contype="0" conaffinity="0" group="2" />
|
||||
</default>
|
||||
<default class="collision">
|
||||
<geom material="collision_material" condim="3" contype="1" conaffinity="1" solref="0.005 1" friction="1 0.01 0.001" density="0" group="1" />
|
||||
</default>
|
||||
</default>
|
||||
</default>
|
||||
<compiler angle="radian" meshdir="meshes/" />
|
||||
<compiler autolimits="true" />
|
||||
<asset>
|
||||
<material name="default_material" rgba="0.7 0.7 0.7 1" />
|
||||
<material name="collision_material" rgba="0.0 0.4 0.8 0.2" />
|
||||
<mesh name="base_link" file="base_link.STL" />
|
||||
<mesh name="fl_hipx" file="fl_hipx.STL" />
|
||||
<mesh name="fl_hipy" file="fl_hipy.STL" />
|
||||
<mesh name="fl_knee" file="fl_knee.STL" />
|
||||
<mesh name="fl_wheel" file="fl_wheel.STL" />
|
||||
<mesh name="fr_hipx" file="fr_hipx.STL" />
|
||||
<mesh name="fr_hipy" file="fr_hipy.STL" />
|
||||
<mesh name="fr_knee" file="fr_knee.STL" />
|
||||
<mesh name="fr_wheel" file="fr_wheel.STL" />
|
||||
<mesh name="hl_hipx" file="hl_hipx.STL" />
|
||||
<mesh name="hl_hipy" file="hl_hipy.STL" />
|
||||
<mesh name="hl_knee" file="hl_knee.STL" />
|
||||
<mesh name="hl_wheel" file="hl_wheel.STL" />
|
||||
<mesh name="hr_hipx" file="hr_hipx.STL" />
|
||||
<mesh name="hr_hipy" file="hr_hipy.STL" />
|
||||
<mesh name="hr_knee" file="hr_knee.STL" />
|
||||
<mesh name="hr_wheel" file="hr_wheel.STL" />
|
||||
<texture name="texplane" type="2d" builtin="checker" rgb1=".2 .3 .4" rgb2=".1 0.15 0.2" width="512" height="512" />
|
||||
<material name="MatPlane" reflectance="0.3" texture="texplane" texrepeat="1 1" texuniform="true" />
|
||||
</asset>
|
||||
<worldbody>
|
||||
<body name="base_link" pos="0 0.0 0.5" quat="1 0 0 0" childclass="robot">
|
||||
<freejoint name="floating_base" />
|
||||
<inertial pos="-0.000512 0.057317 0.001182" quat="1 0 0 0" mass="8.023087" diaginertia="0.115559 0.171329 0.262028" />
|
||||
<site name="imu_site" pos="0.0075 -0.0373 0.0556" size="0.005" rgba="1 0 0 1" />
|
||||
<site name="base_site" pos="0 0 0" quat="1 0 0 0" />
|
||||
<geom type="mesh" class="visual" mesh="base_link" />
|
||||
<geom size="0.08 0.3" pos="0 0.024 0" quat="0.707107 0 0.707107 0" type="cylinder" class="collision" />
|
||||
<geom size="0.08 0.3" pos="0 -0.024 0" quat="0.707107 0 0.707107 0" type="cylinder" class="collision" />
|
||||
<geom size="0.08 0.024 0.3" quat="0.707107 0 0.707107 0" type="box" class="collision" />
|
||||
<body name="fl_hipx" pos="0.2277 0.0574 0">
|
||||
<inertial pos="-0.005034 0.023524 4.6e-05" quat="0.25609 0.665262 -0.228203 0.663151" mass="0.627405" diaginertia="0.000415056 0.000391723 0.000223221" />
|
||||
<joint name="fl_hipx_joint" pos="0 0 0" axis="-1 0 0" range="-0.6109 0.6109" actuatorfrcrange="-50 50" />
|
||||
<geom type="mesh" class="visual" mesh="fl_hipx" />
|
||||
<geom size="0.04 0.05" pos="0 0.054 0" quat="0.707107 0.707107 0 0" type="cylinder" class="collision" />
|
||||
<body name="fl_hipy" pos="0 0.0779 0">
|
||||
<inertial pos="7.4e-05 0.018439 -0.129675" quat="0.701304 -0.0510428 0.0516391 0.709155" mass="0.885874" diaginertia="0.00454086 0.00445199 0.00030515" />
|
||||
<joint name="fl_hipy_joint" pos="0 0 0" axis="0 -1 0" range="-2.5307 2.5307" actuatorfrcrange="-50 50" />
|
||||
<geom type="mesh" class="visual" mesh="fl_hipy" />
|
||||
<geom size="0.032 0.012 0.09" pos="0 0.012 -0.09" type="box" class="collision" />
|
||||
<geom size="0.04 0.04" pos="0 0.018 -0.18" quat="0.707107 0.707107 0 0" type="cylinder" class="collision" name="fl_knee_joint_collision" />
|
||||
<body name="fl_knee" pos="0 0 -0.18">
|
||||
<inertial pos="-3.2e-05 0.013504 -0.156308" quat="0.703786 -0.068447 0.068447 0.703786" mass="0.640283" diaginertia="0.0022296 0.002221 0.000193398" />
|
||||
<joint name="fl_knee_joint" pos="0 0 0" axis="0 -1 0" range="-2.7227 2.7227" actuatorfrcrange="-50 50" />
|
||||
<geom type="mesh" class="visual" mesh="fl_knee" />
|
||||
<geom size="0.029 0.01 0.09" pos="0 -0.012 -0.09" type="box" class="collision" name="fl_shin_collision" />
|
||||
<geom size="0.04 0.025" pos="0 0.002 -0.18" quat="0.707107 0.707107 0 0" type="cylinder" class="collision" name="fl_knee_hub_collision" />
|
||||
<body name="fl_wheel" pos="0 0.045891 -0.18">
|
||||
<inertial pos="-8e-06 0.001694 -5e-06" quat="0.707107 0 0 0.707107" mass="0.587952" diaginertia="0.002003 0.001059 0.001059" />
|
||||
<joint name="fl_wheel_joint" pos="0 0 0" axis="0 -1 0" actuatorfrcrange="-14 14" />
|
||||
<geom type="mesh" class="visual" mesh="fl_wheel" />
|
||||
<geom size="0.081 0.0185" quat="0.707107 0.707107 0 0" type="cylinder" class="collision" name="fl_wheel_collision" />
|
||||
</body>
|
||||
</body>
|
||||
</body>
|
||||
</body>
|
||||
<body name="fr_hipx" pos="0.2277 -0.0574 0">
|
||||
<inertial pos="-0.004831 -0.02349 0.000169" quat="-0.278522 0.675499 0.205833 0.650968" mass="0.626546" diaginertia="0.000415238 0.000395047 0.000218715" />
|
||||
<joint name="fr_hipx_joint" pos="0 0 0" axis="-1 0 0" range="-0.6109 0.6109" actuatorfrcrange="-50 50" />
|
||||
<geom type="mesh" class="visual" mesh="fr_hipx" />
|
||||
<geom size="0.04 0.05" pos="0 -0.054 0" quat="0.707107 0.707107 0 0" type="cylinder" class="collision" />
|
||||
<body name="fr_hipy" pos="0 -0.0779 0">
|
||||
<inertial pos="-7.4e-05 -0.018439 -0.129675" quat="0.701304 0.0510428 -0.0516391 0.709155" mass="0.885874" diaginertia="0.00454086 0.00445199 0.00030515" />
|
||||
<joint name="fr_hipy_joint" pos="0 0 0" axis="0 -1 0" range="-2.5307 2.5307" actuatorfrcrange="-50 50" />
|
||||
<geom type="mesh" class="visual" mesh="fr_hipy" />
|
||||
<geom size="0.032 0.012 0.09" pos="0 -0.012 -0.09" type="box" class="collision" />
|
||||
<geom size="0.04 0.04" pos="0 -0.018 -0.18" quat="0.707107 0.707107 0 0" type="cylinder" class="collision" name="fr_knee_joint_collision" />
|
||||
<body name="fr_knee" pos="0 0 -0.18">
|
||||
<inertial pos="3.2e-05 -0.013504 -0.156308" quat="0.703786 0.068447 -0.068447 0.703786" mass="0.640283" diaginertia="0.0022296 0.002221 0.000193398" />
|
||||
<joint name="fr_knee_joint" pos="0 0 0" axis="0 -1 0" range="-2.7227 2.7227" actuatorfrcrange="-50 50" />
|
||||
<geom type="mesh" class="visual" mesh="fr_knee" />
|
||||
<geom size="0.029 0.01 0.09" pos="0 0.012 -0.09" type="box" class="collision" name="fr_shin_collision" />
|
||||
<geom size="0.04 0.025" pos="0 -0.002 -0.18" quat="0.707107 0.707107 0 0" type="cylinder" class="collision" name="fr_knee_hub_collision" />
|
||||
<body name="fr_wheel" pos="0 -0.045891 -0.18">
|
||||
<inertial pos="8e-06 -0.001694 5e-06" quat="0.707107 0 0 0.707107" mass="0.587952" diaginertia="0.002003 0.001059 0.001059" />
|
||||
<joint name="fr_wheel_joint" pos="0 0 0" axis="0 -1 0" actuatorfrcrange="-14 14" />
|
||||
<geom type="mesh" class="visual" mesh="fr_wheel" />
|
||||
<geom size="0.081 0.0185" quat="0.707107 0.707107 0 0" type="cylinder" class="collision" name="fr_wheel_collision" />
|
||||
</body>
|
||||
</body>
|
||||
</body>
|
||||
</body>
|
||||
<body name="hl_hipx" pos="-0.2277 0.0574 0">
|
||||
<inertial pos="0.004831 0.02349 0.000169" quat="-0.205833 0.650968 0.278522 0.675499" mass="0.626546" diaginertia="0.000415238 0.000395047 0.000218715" />
|
||||
<joint name="hl_hipx_joint" pos="0 0 0" axis="-1 0 0" range="-0.6109 0.6109" actuatorfrcrange="-50 50" />
|
||||
<geom type="mesh" class="visual" mesh="hl_hipx" />
|
||||
<geom size="0.04 0.05" pos="0 0.054 0" quat="0.707107 0.707107 0 0" type="cylinder" class="collision" />
|
||||
<body name="hl_hipy" pos="0 0.0779 0">
|
||||
<inertial pos="7.4e-05 0.018439 -0.129675" quat="0.701304 -0.0510428 0.0516391 0.709155" mass="0.885874" diaginertia="0.00454086 0.00445199 0.00030515" />
|
||||
<joint name="hl_hipy_joint" pos="0 0 0" axis="0 -1 0" range="-2.5307 2.5307" actuatorfrcrange="-50 50" />
|
||||
<geom type="mesh" class="visual" mesh="hl_hipy" />
|
||||
<geom size="0.032 0.012 0.09" pos="0 0.012 -0.09" type="box" class="collision" />
|
||||
<geom size="0.04 0.04" pos="0 0.018 -0.18" quat="0.707107 0.707107 0 0" type="cylinder" class="collision" name="hl_knee_joint_collision" />
|
||||
<body name="hl_knee" pos="0 0 -0.18">
|
||||
<inertial pos="-3.2e-05 0.013504 -0.156308" quat="0.703786 -0.068447 0.068447 0.703786" mass="0.640283" diaginertia="0.0022296 0.002221 0.000193398" />
|
||||
<joint name="hl_knee_joint" pos="0 0 0" axis="0 -1 0" range="-2.7227 2.7227" actuatorfrcrange="-50 50" />
|
||||
<geom type="mesh" class="visual" mesh="hl_knee" />
|
||||
<geom size="0.029 0.01 0.09" pos="0 -0.012 -0.09" type="box" class="collision" name="hl_shin_collision" />
|
||||
<geom size="0.04 0.025" pos="0 0.002 -0.18" quat="0.707107 0.707107 0 0" type="cylinder" class="collision" name="hl_knee_hub_collision" />
|
||||
<body name="hl_wheel" pos="0 0.045891 -0.18">
|
||||
<inertial pos="-8e-06 0.001694 -5e-06" quat="0.707107 0 0 0.707107" mass="0.587952" diaginertia="0.002003 0.001059 0.001059" />
|
||||
<joint name="hl_wheel_joint" pos="0 0 0" axis="0 -1 0" actuatorfrcrange="-14 14" />
|
||||
<geom type="mesh" class="visual" mesh="hl_wheel" />
|
||||
<geom size="0.081 0.0185" quat="0.707107 0.707107 0 0" type="cylinder" class="collision" name="hl_wheel_collision" />
|
||||
</body>
|
||||
</body>
|
||||
</body>
|
||||
</body>
|
||||
<body name="hr_hipx" pos="-0.2277 -0.0574 0">
|
||||
<inertial pos="0.005034 -0.023524 4.6e-05" quat="0.228203 0.663151 -0.25609 0.665262" mass="0.627405" diaginertia="0.000415056 0.000391723 0.000223221" />
|
||||
<joint name="hr_hipx_joint" pos="0 0 0" axis="-1 0 0" range="-0.6109 0.6109" actuatorfrcrange="-50 50" />
|
||||
<geom type="mesh" class="visual" mesh="hr_hipx" />
|
||||
<geom size="0.04 0.05" pos="0 -0.054 0" quat="0.707107 0.707107 0 0" type="cylinder" class="collision" />
|
||||
<body name="hr_hipy" pos="0 -0.0779 0">
|
||||
<inertial pos="-7.4e-05 -0.018439 -0.129675" quat="0.701304 0.0510428 -0.0516391 0.709155" mass="0.885874" diaginertia="0.00454086 0.00445199 0.00030515" />
|
||||
<joint name="hr_hipy_joint" pos="0 0 0" axis="0 -1 0" range="-2.5307 2.5307" actuatorfrcrange="-50 50" />
|
||||
<geom type="mesh" class="visual" mesh="hr_hipy" />
|
||||
<geom size="0.032 0.012 0.09" pos="0 -0.012 -0.09" type="box" class="collision" />
|
||||
<geom size="0.04 0.04" pos="0 -0.018 -0.18" quat="0.707107 0.707107 0 0" type="cylinder" class="collision" name="hr_knee_joint_collision" />
|
||||
<body name="hr_knee" pos="0 0 -0.18">
|
||||
<inertial pos="3.2e-05 -0.013504 -0.156308" quat="0.703786 0.068447 -0.068447 0.703786" mass="0.640283" diaginertia="0.0022296 0.002221 0.000193398" />
|
||||
<joint name="hr_knee_joint" pos="0 0 0" axis="0 -1 0" range="-2.7227 2.7227" actuatorfrcrange="-50 50" />
|
||||
<geom type="mesh" class="visual" mesh="hr_knee" />
|
||||
<geom size="0.029 0.01 0.09" pos="0 0.012 -0.09" type="box" class="collision" name="hr_shin_collision" />
|
||||
<geom size="0.04 0.025" pos="0 -0.002 -0.18" quat="0.707107 0.707107 0 0" type="cylinder" class="collision" name="hr_knee_hub_collision" />
|
||||
<body name="hr_wheel" pos="0 -0.045891 -0.18">
|
||||
<inertial pos="8e-06 -0.001694 5e-06" quat="0.707107 0 0 0.707107" mass="0.587952" diaginertia="0.002003 0.001059 0.001059" />
|
||||
<joint name="hr_wheel_joint" pos="0 0 0" axis="0 -1 0" actuatorfrcrange="-14 14" />
|
||||
<geom type="mesh" class="visual" mesh="hr_wheel" />
|
||||
<geom size="0.081 0.0185" quat="0.707107 0.707107 0 0" type="cylinder" class="collision" name="hr_wheel_collision" />
|
||||
</body>
|
||||
</body>
|
||||
</body>
|
||||
</body>
|
||||
</body>
|
||||
</worldbody>
|
||||
<contact>
|
||||
<exclude body1="base_link" body2="fl_hipx" />
|
||||
<exclude body1="fl_hipx" body2="fl_hipy" />
|
||||
<exclude body1="fl_hipy" body2="fl_knee" />
|
||||
<exclude body1="fl_knee" body2="fl_wheel" />
|
||||
<exclude body1="base_link" body2="fr_hipx" />
|
||||
<exclude body1="fr_hipx" body2="fr_hipy" />
|
||||
<exclude body1="fr_hipy" body2="fr_knee" />
|
||||
<exclude body1="fr_knee" body2="fr_wheel" />
|
||||
<exclude body1="base_link" body2="hl_hipx" />
|
||||
<exclude body1="hl_hipx" body2="hl_hipy" />
|
||||
<exclude body1="hl_hipy" body2="hl_knee" />
|
||||
<exclude body1="hl_knee" body2="hl_wheel" />
|
||||
<exclude body1="base_link" body2="hr_hipx" />
|
||||
<exclude body1="hr_hipx" body2="hr_hipy" />
|
||||
<exclude body1="hr_hipy" body2="hr_knee" />
|
||||
<exclude body1="hr_knee" body2="hr_wheel" />
|
||||
</contact>
|
||||
</mujoco>
|
||||
@@ -0,0 +1,12 @@
|
||||
"""Selectable S10 training tasks.
|
||||
|
||||
Import a task module here to register it. The training entry point accepts
|
||||
``--task <name>`` and applies that task's defaults before explicit CLI flags.
|
||||
Shared algorithm code stays revision-agnostic; only the observation contract
|
||||
and environment wiring differ between tasks.
|
||||
"""
|
||||
|
||||
from .base import ActorSpec, TaskSpec, get_task, list_tasks, register, task_names
|
||||
from . import task_v4, task_v4_upward, task_v7
|
||||
|
||||
__all__ = ("ActorSpec", "TaskSpec", "get_task", "list_tasks", "register", "task_names")
|
||||
@@ -0,0 +1,95 @@
|
||||
"""Parallel S10 training tasks: architecture, environment and recipe.
|
||||
|
||||
A task owns the observation contract (the part that cannot be changed by a
|
||||
resume) and the training defaults that go with its terrain revision. Shared
|
||||
learner, reward and world-model code lives outside this package, so adding a
|
||||
new task means adding one module and registering it in ``tasks/__init__.py``.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Callable
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ActorSpec:
|
||||
"""Actor observation contract used by training, export and playback."""
|
||||
name: str
|
||||
terms: tuple[int, ...]
|
||||
frame_terms: tuple[dict[str, Any], ...]
|
||||
wheel_vel: bool
|
||||
|
||||
@property
|
||||
def frame_dim(self) -> int:
|
||||
return sum(self.terms)
|
||||
|
||||
def history_dim(self, history: int) -> int:
|
||||
if history < 1:
|
||||
raise ValueError("history must be positive")
|
||||
return self.frame_dim * history
|
||||
|
||||
def metadata(self, history: int) -> dict[str, Any]:
|
||||
return {
|
||||
"actor_schema": self.name,
|
||||
"actor_wheel_vel": self.wheel_vel,
|
||||
"frame_dim": self.frame_dim,
|
||||
"observation_dim": self.history_dim(history),
|
||||
"frame_terms": [dict(term) for term in self.frame_terms],
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TaskSpec:
|
||||
"""One selectable training task.
|
||||
|
||||
``defaults`` are argparse destinations applied before parsing, so explicit
|
||||
command-line flags still override them. That keeps a task reproducible while
|
||||
allowing a one-off run to change length or batch size.
|
||||
"""
|
||||
name: str
|
||||
description: str
|
||||
terrain_revision: int
|
||||
actor: ActorSpec
|
||||
defaults: dict[str, Any] = field(default_factory=dict)
|
||||
env_cfg: Callable[..., Any] | None = None
|
||||
|
||||
@property
|
||||
def contacts(self) -> bool:
|
||||
return self.terrain_revision >= 4
|
||||
|
||||
@property
|
||||
def courses(self) -> bool:
|
||||
return self.terrain_revision >= 5
|
||||
|
||||
@property
|
||||
def collisions(self) -> bool:
|
||||
return self.terrain_revision >= 7
|
||||
|
||||
def build_env_cfg(self, play: bool, history: int, seed: int):
|
||||
if self.env_cfg is None:
|
||||
raise ValueError(f"Task {self.name} has no environment builder")
|
||||
return self.env_cfg(play=play, history=history, seed=seed)
|
||||
|
||||
|
||||
_REGISTRY: dict[str, TaskSpec] = {}
|
||||
|
||||
|
||||
def register(task: TaskSpec) -> TaskSpec:
|
||||
if task.name in _REGISTRY:
|
||||
raise ValueError(f"Duplicate task name: {task.name}")
|
||||
_REGISTRY[task.name] = task
|
||||
return task
|
||||
|
||||
|
||||
def get_task(name: str) -> TaskSpec:
|
||||
try:
|
||||
return _REGISTRY[name]
|
||||
except KeyError as exc:
|
||||
raise KeyError(f"Unknown task {name!r}; available: {', '.join(task_names())}") from exc
|
||||
|
||||
|
||||
def list_tasks() -> tuple[TaskSpec, ...]:
|
||||
return tuple(_REGISTRY.values())
|
||||
|
||||
|
||||
def task_names() -> tuple[str, ...]:
|
||||
return tuple(_REGISTRY)
|
||||
@@ -0,0 +1,65 @@
|
||||
"""S10 terrain revision 4: legacy 53-value Actor frame with measured wheel speed."""
|
||||
|
||||
from .base import ActorSpec, TaskSpec, register
|
||||
|
||||
ACTOR = ActorSpec(
|
||||
name="s10-v4-53",
|
||||
terms=(3, 3, 12, 12, 4, 16, 3),
|
||||
frame_terms=(
|
||||
{"name": "base_ang_vel", "dim": 3, "scale": 1.0},
|
||||
{"name": "projected_gravity", "dim": 3, "scale": 1.0},
|
||||
{"name": "joint_pos_rel", "dim": 12, "scale": 1.0},
|
||||
{"name": "joint_vel_rel", "dim": 12, "scale": 0.05},
|
||||
{"name": "wheel_vel_rel", "dim": 4, "scale": 0.05},
|
||||
{"name": "last_action", "dim": 16, "scale": 1.0},
|
||||
{"name": "command", "dim": 3, "scale": 1.0},
|
||||
),
|
||||
wheel_vel=True,
|
||||
)
|
||||
|
||||
|
||||
def env_cfg(play: bool, history: int, seed: int):
|
||||
from s10_supermimic.env.terrain import s10_velocity_terrain_env_cfg
|
||||
return s10_velocity_terrain_env_cfg(
|
||||
play=play, history=history, seed=seed, revision=4,
|
||||
actor_wheel_vel=ACTOR.wheel_vel)
|
||||
|
||||
|
||||
TASK = register(TaskSpec(
|
||||
name="s10-v4",
|
||||
description="S10 terrain v4 legacy task: 53-value Actor frame, wheel speed, 20k validated recipe",
|
||||
terrain_revision=4,
|
||||
actor=ACTOR,
|
||||
env_cfg=env_cfg,
|
||||
defaults={
|
||||
"terrain": True,
|
||||
"terrain_revision": 4,
|
||||
"num_envs": 4096,
|
||||
"iterations": 20000,
|
||||
"num_steps_per_env": 1,
|
||||
"buffer_len": 512,
|
||||
"wm_window": 8,
|
||||
"wm_batch": 2048,
|
||||
"wm_updates": 4,
|
||||
"imag_batch": 2048,
|
||||
"horizon": 32,
|
||||
"warmup": 50,
|
||||
"critic_epochs": 5,
|
||||
"critic_minibatches": 4,
|
||||
"save_interval": 500,
|
||||
"log_interval": 10,
|
||||
"command_curriculum": True,
|
||||
"command_curriculum_start": 0,
|
||||
"prioritized_replay": True,
|
||||
"archive_size": 100000,
|
||||
"actor_ramp_iterations": 1000,
|
||||
"actor_hold_iterations": 0,
|
||||
"lr_decay_iterations": 30000,
|
||||
"lr_final_frac": 0.03,
|
||||
"history": 5,
|
||||
"wm_hidden": [512, 512, 512],
|
||||
"actor_hidden": [256, 128, 64],
|
||||
"critic_hidden": [256, 128, 64],
|
||||
"compile_imagination": False,
|
||||
},
|
||||
))
|
||||
@@ -0,0 +1,172 @@
|
||||
"""S10 terrain v4 + upward recovery task.
|
||||
|
||||
This is a separate task, not an inheritance of v4 or v7. It keeps the v4
|
||||
53-value Actor frame, terrain revision, command schedule, terrain proportions,
|
||||
rolling rewards and replay, and changes only:
|
||||
|
||||
* the existing ``upright`` reward into a sign-aware upward reward (no new term);
|
||||
* clamps on the large raw penalty terms;
|
||||
* random roll/pitch/yaw at reset;
|
||||
* a 4-second sustained bad-orientation termination.
|
||||
|
||||
Reference: IsaacLab S10 ``upward_env_cfg.py`` for the upward/recovery state
|
||||
machine only. Its speed ranges, reward set and terrain config are not copied.
|
||||
"""
|
||||
|
||||
import math
|
||||
|
||||
from .base import ActorSpec, TaskSpec, register
|
||||
|
||||
ACTOR = ActorSpec(
|
||||
name="s10-v4-upward-53",
|
||||
terms=(3, 3, 12, 12, 4, 16, 3),
|
||||
frame_terms=(
|
||||
{"name": "base_ang_vel", "dim": 3, "scale": 1.0},
|
||||
{"name": "projected_gravity", "dim": 3, "scale": 1.0},
|
||||
{"name": "joint_pos_rel", "dim": 12, "scale": 1.0},
|
||||
{"name": "joint_vel_rel", "dim": 12, "scale": 0.05},
|
||||
{"name": "wheel_vel_rel", "dim": 4, "scale": 0.05},
|
||||
{"name": "last_action", "dim": 16, "scale": 1.0},
|
||||
{"name": "command", "dim": 3, "scale": 1.0},
|
||||
),
|
||||
wheel_vel=True,
|
||||
)
|
||||
|
||||
|
||||
def env_cfg(play: bool, history: int, seed: int):
|
||||
from mjlab.managers import RewardTermCfg, TerminationTermCfg
|
||||
from mjlab.managers.scene_entity_config import SceneEntityCfg
|
||||
|
||||
from s10_supermimic.core.upward import (
|
||||
BAD_ORIENTATION_WEIGHT, UPWARD_REWARD_CLAMPS, BadOrientationSustained,
|
||||
bad_orientation_penalty, clamped_action_rate, clamped_dof_pos_limits,
|
||||
clamped_joint_torques, clamped_leg_posture, clamped_rolling_cost,
|
||||
clamped_standstill, reset_root_state_upward, upright_upward,
|
||||
upward_curriculum,
|
||||
)
|
||||
from s10_supermimic.env.terrain import s10_velocity_terrain_env_cfg
|
||||
|
||||
cfg = s10_velocity_terrain_env_cfg(
|
||||
play=play, history=history, seed=seed, revision=4, actor_wheel_vel=True)
|
||||
# A 10 s episode lets each env complete 20 episodes per curriculum step
|
||||
# within a 20k-iteration run. The 4 s bad-orientation termination still
|
||||
# leaves 6 s after a successful recovery. Play/evaluation keeps the normal
|
||||
# terrain-course horizon.
|
||||
if not play:
|
||||
cfg.episode_length_s = 10.0
|
||||
|
||||
# Random orientation at reset. Training starts at +/-60 degrees and advances
|
||||
# per environment through upward_curriculum(); play/eval uses the full sphere
|
||||
# so the checkpoint is tested on the final distribution.
|
||||
tilt = math.pi if play else math.radians(60.0)
|
||||
pose = cfg.events["reset_base"].params["pose_range"]
|
||||
pose.update(
|
||||
z=(0.0, 0.2),
|
||||
roll=(-tilt, tilt),
|
||||
pitch=(-tilt, tilt),
|
||||
yaw=(-math.pi, math.pi),
|
||||
)
|
||||
cfg.events["reset_base"].func = reset_root_state_upward
|
||||
|
||||
# Existing upright term becomes the normalized upward reward. Max rate 2.0,
|
||||
# i.e. 0.04 per 20 ms step, the same ratio to tracking (2.0 each) as IsaacLab.
|
||||
upright = cfg.rewards["upright"]
|
||||
upright_asset = upright.params.get(
|
||||
"asset_cfg", SceneEntityCfg("robot", body_names=("base_link",)))
|
||||
upright.func = upright_upward
|
||||
upright.weight = 2.0
|
||||
upright.params = {"asset_cfg": upright_asset}
|
||||
|
||||
# IsaacLab-style recovery shaping: no reward for termination itself, but a
|
||||
# continuous per-step penalty while the orientation is bad. Weight -40 keeps
|
||||
# the same 20x bad/upward ratio as IsaacLab (-80 vs max 4; here -40 vs max 2).
|
||||
cfg.rewards["terminated"].weight = 0.0
|
||||
cfg.rewards["bad_orientation_penalty"] = RewardTermCfg(
|
||||
func=bad_orientation_penalty,
|
||||
weight=BAD_ORIENTATION_WEIGHT,
|
||||
params={"asset_cfg": SceneEntityCfg("robot")},
|
||||
)
|
||||
|
||||
# Clamp the large raw penalties. The weights and the rest of the v4 reward
|
||||
# set are unchanged; the caps only prevent recovery poses from dominating.
|
||||
leg = cfg.rewards["leg_posture"]
|
||||
leg.params = {"asset_cfg": leg.params["asset_cfg"],
|
||||
"max_value": UPWARD_REWARD_CLAMPS["leg_posture"]}
|
||||
leg.func = clamped_leg_posture
|
||||
|
||||
standstill = cfg.rewards["standstill"]
|
||||
standstill.params = {"wheel_cfg": standstill.params["wheel_cfg"],
|
||||
"max_value": UPWARD_REWARD_CLAMPS["standstill"]}
|
||||
standstill.func = clamped_standstill
|
||||
|
||||
cfg.rewards["action_rate_l2"].func = clamped_action_rate
|
||||
cfg.rewards["action_rate_l2"].params = {
|
||||
"max_value": UPWARD_REWARD_CLAMPS["action_rate_l2"]}
|
||||
|
||||
cfg.rewards["joint_torques_l2"].func = clamped_joint_torques
|
||||
cfg.rewards["joint_torques_l2"].params = {
|
||||
"max_value": UPWARD_REWARD_CLAMPS["joint_torques_l2"]}
|
||||
|
||||
limits = cfg.rewards["dof_pos_limits"]
|
||||
limits.params = {"asset_cfg": limits.params["asset_cfg"],
|
||||
"max_value": UPWARD_REWARD_CLAMPS["dof_pos_limits"]}
|
||||
limits.func = clamped_dof_pos_limits
|
||||
|
||||
for name in ("unnecessary_wheel_lift", "rolling_leg_motion"):
|
||||
term = cfg.rewards[name]
|
||||
term.func = clamped_rolling_cost
|
||||
term.params = {"term": name, "max_value": UPWARD_REWARD_CLAMPS[name]}
|
||||
|
||||
# Replace v4's immediate 70-degree fall with the upward recovery state
|
||||
# machine: bad orientation is tolerated for 4 s, then the episode ends.
|
||||
cfg.terminations["fell_over"] = TerminationTermCfg(
|
||||
func=BadOrientationSustained,
|
||||
params={"duration_s": 4.0, "asset_cfg": SceneEntityCfg("robot")},
|
||||
)
|
||||
|
||||
# Each environment advances its own terrain/tilt level from its last 20
|
||||
# recovery outcomes (16/20 promotes); there is no global batch gate.
|
||||
if not play:
|
||||
cfg.curriculum["course"].func = upward_curriculum
|
||||
return cfg
|
||||
|
||||
|
||||
TASK = register(TaskSpec(
|
||||
name="s10-v4-upward",
|
||||
description="S10 v4 recovery task: upward reward, monotonic tilt severity, per-env 20-episode curriculum, horizon 64",
|
||||
terrain_revision=4,
|
||||
actor=ACTOR,
|
||||
env_cfg=env_cfg,
|
||||
defaults={
|
||||
"terrain": True,
|
||||
"terrain_revision": 4,
|
||||
"num_envs": 4096,
|
||||
"iterations": 30000,
|
||||
"num_steps_per_env": 1,
|
||||
"buffer_len": 512,
|
||||
"wm_window": 8,
|
||||
"wm_batch": 2048,
|
||||
"wm_updates": 4,
|
||||
"imag_batch": 2048,
|
||||
"horizon": 64,
|
||||
"warmup": 50,
|
||||
"action_noise": 0.22,
|
||||
"critic_epochs": 5,
|
||||
"critic_minibatches": 4,
|
||||
"save_interval": 500,
|
||||
"log_interval": 10,
|
||||
"command_curriculum": True,
|
||||
"command_curriculum_start": 0,
|
||||
"prioritized_replay": True,
|
||||
"archive_size": 100000,
|
||||
"actor_ramp_iterations": 1000,
|
||||
"actor_hold_iterations": 0,
|
||||
"lr_decay_iterations": 30000,
|
||||
"lr_final_frac": 0.03,
|
||||
"history": 5,
|
||||
"wm_hidden": [512, 512, 512],
|
||||
"actor_hidden": [256, 128, 64],
|
||||
"critic_hidden": [256, 128, 64],
|
||||
"compile_imagination": False,
|
||||
},
|
||||
))
|
||||
@@ -0,0 +1,78 @@
|
||||
"""S10 terrain revision 7: safe low-speed all-terrain locomotion.
|
||||
|
||||
V7 keeps the 53-value v4 Actor frame (measured wheel speed is required for
|
||||
wheel slip/stall feedback), a stair-pass-gated command stage ending at
|
||||
1.5 m/s, 0.8 m/s and 1.2 rad/s with a 0.05 per-axis deadband, a conservative
|
||||
initialization (small x/y/yaw only, no random roll/pitch), a hipx 0.05 /
|
||||
hipy +/-25 deg / knee 1.0 default stance, always-on wheel/body projection
|
||||
penalties (per-wheel outer-1/4 and paired-leg 5/6 with a 0.5 s grace), reduced
|
||||
knee/hipy posture stiffness on stairs, and the v5 stair curriculum plus v7
|
||||
non-wheel collision head (50 N, wheels/knees excluded). Only ``imag_batch`` is
|
||||
reduced to 2048 to fit a 16 GB GPU.
|
||||
"""
|
||||
|
||||
from .base import ActorSpec, TaskSpec, register
|
||||
|
||||
ACTOR = ActorSpec(
|
||||
name="s10-v7-53",
|
||||
terms=(3, 3, 12, 12, 4, 16, 3),
|
||||
frame_terms=(
|
||||
{"name": "base_ang_vel", "dim": 3, "scale": 1.0},
|
||||
{"name": "projected_gravity", "dim": 3, "scale": 1.0},
|
||||
{"name": "joint_pos_rel", "dim": 12, "scale": 1.0},
|
||||
{"name": "joint_vel_rel", "dim": 12, "scale": 0.05},
|
||||
{"name": "wheel_vel_rel", "dim": 4, "scale": 0.05},
|
||||
{"name": "last_action", "dim": 16, "scale": 1.0},
|
||||
{"name": "command", "dim": 3, "scale": 1.0},
|
||||
),
|
||||
wheel_vel=True,
|
||||
)
|
||||
|
||||
|
||||
def env_cfg(play: bool, history: int, seed: int):
|
||||
from s10_supermimic.env.terrain import s10_velocity_terrain_env_cfg
|
||||
return s10_velocity_terrain_env_cfg(
|
||||
play=play, history=history, seed=seed, revision=7,
|
||||
actor_wheel_vel=ACTOR.wheel_vel)
|
||||
|
||||
|
||||
TASK = register(TaskSpec(
|
||||
name="s10-v7",
|
||||
description="S10 terrain v7 safe all-terrain task: 53-value Actor frame, wheel speed, stair-gated low-speed commands, stair curriculum + collision",
|
||||
terrain_revision=7,
|
||||
actor=ACTOR,
|
||||
env_cfg=env_cfg,
|
||||
defaults={
|
||||
"terrain": True,
|
||||
"terrain_revision": 7,
|
||||
"num_envs": 4096,
|
||||
"iterations": 40000,
|
||||
"num_steps_per_env": 1,
|
||||
"buffer_len": 512,
|
||||
"wm_window": 8,
|
||||
"wm_batch": 2048,
|
||||
"wm_updates": 4,
|
||||
# 8192 is the full-size imagination batch; 2048 keeps the 16 GB footprint.
|
||||
"imag_batch": 2048,
|
||||
"horizon": 32,
|
||||
"warmup": 300,
|
||||
"critic_epochs": 5,
|
||||
"critic_minibatches": 4,
|
||||
"save_interval": 500,
|
||||
"log_interval": 50,
|
||||
# V7 uses stair-pass-gated stages, capped at its safe envelope.
|
||||
"command_curriculum": True,
|
||||
"command_curriculum_start": 0,
|
||||
"prioritized_replay": True,
|
||||
"archive_size": 100000,
|
||||
"actor_ramp_iterations": 1000,
|
||||
"actor_hold_iterations": 20000,
|
||||
"lr_decay_iterations": 20000,
|
||||
"lr_final_frac": 0.2,
|
||||
"history": 5,
|
||||
"wm_hidden": [512, 512, 512],
|
||||
"actor_hidden": [256, 128, 64],
|
||||
"critic_hidden": [256, 128, 64],
|
||||
"compile_imagination": True,
|
||||
},
|
||||
))
|
||||
@@ -0,0 +1,162 @@
|
||||
"""Export an S10 SuperMimic actor checkpoint to ONNX + deployment JSON.
|
||||
|
||||
Example:
|
||||
python scripts/export_onnx.py \
|
||||
--checkpoint logs/s10_terrain_v4_20k_full/model_20000.pt \
|
||||
--output logs/s10_terrain_v4_20k_full/model_20000.onnx
|
||||
|
||||
The ONNX file contains only the deterministic actor (RunningNorm + MLP + tanh).
|
||||
Input ``obs`` is [batch, 265] for v4/v4-upward and for the new v7 (wheel speed
|
||||
restored); only legacy v7 checkpoints use [batch, 245]. Output is [batch, 16].
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from mjlab.utils.string import resolve_expr
|
||||
|
||||
from s10_supermimic.core.policy import load_actor
|
||||
from s10_supermimic.robot.config import (
|
||||
LEG_JOINT_NAMES, WHEEL_JOINT_NAMES, get_s10_robot_cfg,
|
||||
)
|
||||
|
||||
|
||||
def task_matches_actor(task, actor_wheel_vel: bool, frame_dim: int) -> bool:
|
||||
return task.actor.wheel_vel == actor_wheel_vel and task.actor.frame_dim == frame_dim
|
||||
|
||||
|
||||
def default_joint_pos(checkpoint, actor_wheel_vel: bool, frame_dim: int) -> list[float]:
|
||||
"""Use the checkpoint task's default pose when its Actor contract matches.
|
||||
|
||||
Legacy 49-value v7 checkpoints predate the ±5° safe stance, so they fall
|
||||
back to the legacy v4 ±20° pose instead of silently exporting a new target.
|
||||
"""
|
||||
robot_cfg = get_s10_robot_cfg()
|
||||
task_name = checkpoint.get("config", {}).get("task")
|
||||
if task_name:
|
||||
from s10_supermimic.tasks import get_task
|
||||
try:
|
||||
task = get_task(task_name)
|
||||
if task_matches_actor(task, actor_wheel_vel, frame_dim):
|
||||
cfg = task.build_env_cfg(
|
||||
play=True, history=int(checkpoint["history"]),
|
||||
seed=int(checkpoint.get("config", {}).get("seed", 0)))
|
||||
robot_cfg = cfg.scene.entities["robot"]
|
||||
except KeyError:
|
||||
pass
|
||||
names = LEG_JOINT_NAMES + WHEEL_JOINT_NAMES
|
||||
return [float(v) for v in resolve_expr(robot_cfg.init_state.joint_pos, names, 0.0)]
|
||||
|
||||
|
||||
def main() -> None:
|
||||
p = argparse.ArgumentParser(description=__doc__)
|
||||
p.add_argument("--checkpoint", type=Path, required=True)
|
||||
p.add_argument("--output", type=Path, default=None,
|
||||
help="ONNX path; defaults to <checkpoint>.onnx")
|
||||
p.add_argument("--opset", type=int, default=17)
|
||||
p.add_argument("--device", default="cpu")
|
||||
args = p.parse_args()
|
||||
|
||||
ckpt = torch.load(args.checkpoint, map_location=args.device, weights_only=True)
|
||||
actor = load_actor(ckpt, args.device).eval()
|
||||
history = int(ckpt["history"])
|
||||
if ckpt.get("format") not in ("s10-flat-policy-v1", "s10-terrain-policy-v1"):
|
||||
raise ValueError(f"Not an S10 policy checkpoint: {ckpt.get('format')}")
|
||||
if history < 1:
|
||||
raise ValueError("checkpoint history must be positive")
|
||||
|
||||
output = args.output or args.checkpoint.with_suffix(".onnx")
|
||||
output.parent.mkdir(parents=True, exist_ok=True)
|
||||
dummy = torch.zeros(1, actor.frame_dim * history, device=args.device)
|
||||
|
||||
with torch.no_grad():
|
||||
reference = actor(dummy).cpu().numpy()
|
||||
|
||||
torch.onnx.export(
|
||||
actor,
|
||||
(dummy,),
|
||||
str(output),
|
||||
input_names=["obs"],
|
||||
output_names=["actions"],
|
||||
opset_version=args.opset,
|
||||
dynamo=False,
|
||||
do_constant_folding=True,
|
||||
dynamic_axes={"obs": {0: "batch"}, "actions": {0: "batch"}},
|
||||
)
|
||||
|
||||
import onnx
|
||||
import onnxruntime as ort
|
||||
onnx.checker.check_model(str(output))
|
||||
session = ort.InferenceSession(str(output), providers=["CPUExecutionProvider"])
|
||||
onnx_out = session.run(["actions"], {"obs": dummy.cpu().numpy()})[0]
|
||||
max_error = float(np.max(np.abs(onnx_out - reference)))
|
||||
if max_error > 1e-5:
|
||||
raise RuntimeError(f"ONNX/PyTorch mismatch: max error {max_error}")
|
||||
|
||||
revision = int(ckpt.get("config", {}).get("terrain_revision", 1))
|
||||
leg_scales = ([0.25, 0.6, 0.8] * 4) if revision >= 2 else ([0.25] * 12)
|
||||
wheel_scale = 35.0 if revision >= 3 else 15.0
|
||||
task_name = ckpt.get("config", {}).get("task")
|
||||
task = None
|
||||
if task_name:
|
||||
from s10_supermimic.tasks import get_task
|
||||
try:
|
||||
task = get_task(task_name)
|
||||
except KeyError:
|
||||
task = None
|
||||
if task is not None and task_matches_actor(task, actor.actor_wheel_vel, actor.frame_dim):
|
||||
actor_schema = task.actor.name
|
||||
else:
|
||||
actor_schema = "s10-v4-53" if actor.actor_wheel_vel else "s10-v7-49"
|
||||
sidecar = {
|
||||
"format": ckpt["format"],
|
||||
"step": ckpt.get("step"),
|
||||
"history": history,
|
||||
"actor_schema": actor_schema,
|
||||
"actor_wheel_vel": actor.actor_wheel_vel,
|
||||
"frame_dim": actor.frame_dim,
|
||||
"observation_dim": actor.frame_dim * history,
|
||||
"action_dim": 16,
|
||||
"dt": ckpt.get("dt", 0.02),
|
||||
"actor_hidden": list(ckpt.get("actor_hidden", ())),
|
||||
"actor_params": sum(x.numel() for x in actor.parameters()),
|
||||
"frame_terms": [
|
||||
{"name": "base_ang_vel", "dim": 3, "scale": 1.0},
|
||||
{"name": "projected_gravity", "dim": 3, "scale": 1.0},
|
||||
{"name": "joint_pos_rel", "dim": 12, "scale": 1.0},
|
||||
{"name": "joint_vel_rel", "dim": 12, "scale": 0.05},
|
||||
*([{"name": "wheel_vel_rel", "dim": 4, "scale": 0.05}]
|
||||
if actor.actor_wheel_vel else []),
|
||||
{"name": "last_action", "dim": 16, "scale": 1.0},
|
||||
{"name": "command", "dim": 3, "scale": 1.0},
|
||||
],
|
||||
"history_layout": "term-wise blocks, oldest to newest (mjlab flattening)",
|
||||
"env_action_order": list(ckpt["env_action_order"]),
|
||||
"leg_scale": leg_scales,
|
||||
"wheel_scale": wheel_scale,
|
||||
"wheel_clip": 35.0,
|
||||
"default_joint_pos": default_joint_pos(
|
||||
ckpt, actor.actor_wheel_vel, actor.frame_dim),
|
||||
"leg_joint_names": list(LEG_JOINT_NAMES),
|
||||
"wheel_joint_names": list(WHEEL_JOINT_NAMES),
|
||||
"onnx": {
|
||||
"input": "obs",
|
||||
"output": "actions",
|
||||
"opset": args.opset,
|
||||
"max_abs_error_vs_torch": max_error,
|
||||
},
|
||||
}
|
||||
sidecar_path = output.with_suffix(".json")
|
||||
sidecar_path.write_text(json.dumps(sidecar, indent=2) + "\n")
|
||||
print(f"Exported ONNX: {output}")
|
||||
print(f"Deployment metadata: {sidecar_path}")
|
||||
print(f"ONNX/PyTorch max error: {max_error:.3e}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,211 @@
|
||||
"""S10 browser sliders with one CPU environment and bounded JPEG storage."""
|
||||
import argparse
|
||||
import gc
|
||||
from io import BytesIO
|
||||
import os
|
||||
from pathlib import Path
|
||||
import threading
|
||||
import time
|
||||
|
||||
from play_keyboard import DEFAULT_CHECKPOINT, keyboard_command_cfg
|
||||
from s10_supermimic.core.gui import ControlState, make_server
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--checkpoint", type=Path, default=DEFAULT_CHECKPOINT)
|
||||
parser.add_argument("--port", type=int, default=8765)
|
||||
parser.add_argument("--device", default="cpu")
|
||||
parser.add_argument("--terrain-kind", choices=("flat", "up", "down", "platform", "slope", "rubble"), default="rubble")
|
||||
parser.add_argument("--terrain-level", type=int, choices=range(6), default=0)
|
||||
parser.add_argument("--fps", type=int, default=50)
|
||||
parser.add_argument("--shadows", action="store_true", help="Enable costly shadows/reflections")
|
||||
parser.add_argument("--ghost-horizon", type=int, default=32,
|
||||
help="World-model future steps drawn as ghost poses")
|
||||
parser.add_argument("--ghost-stride", type=int, default=4,
|
||||
help="Draw every Nth ghost (1 = all, 4 = 8 ghosts)")
|
||||
parser.add_argument("--no-ghost", action="store_true",
|
||||
help="Disable the world-model ghost/error overlay")
|
||||
parser.add_argument("--steps", type=int, default=0, help="Exit after N simulation steps; 0 means until Ctrl+C")
|
||||
args = parser.parse_args()
|
||||
if not args.checkpoint.is_file():
|
||||
parser.error(f"Checkpoint not found: {args.checkpoint}")
|
||||
if not 1 <= args.fps <= 50 or not 1 <= args.port <= 65535 or args.steps < 0:
|
||||
parser.error("Require fps 1..50, port 1..65535 and nonnegative steps")
|
||||
if args.ghost_horizon < 1 or args.ghost_stride < 1:
|
||||
parser.error("ghost-horizon and ghost-stride must be positive")
|
||||
os.environ.setdefault("MUJOCO_GL", "egl")
|
||||
os.environ.setdefault("OMP_NUM_THREADS", "1")
|
||||
import torch
|
||||
from PIL import Image
|
||||
from mjlab.envs import ManagerBasedRlEnv
|
||||
from s10_supermimic.core.bridge import S10Bridge
|
||||
from s10_supermimic.core.policy import load_actor, evaluation_env_cfg
|
||||
from s10_supermimic.core.world_preview import WorldModelPreview
|
||||
|
||||
torch.set_num_threads(1)
|
||||
state = ControlState()
|
||||
server = make_server(state, args.port)
|
||||
env = None
|
||||
worker = None
|
||||
try:
|
||||
checkpoint = torch.load(args.checkpoint, map_location=args.device, weights_only=True)
|
||||
actor = load_actor(checkpoint, args.device)
|
||||
# Keep the world model so the play view can draw the predicted future.
|
||||
checkpoint = {k: checkpoint[k] for k in
|
||||
("format", "history", "config", "env_action_order", "dt", "step",
|
||||
"world_model")}
|
||||
checkpoint["actor_wheel_vel"] = actor.actor_wheel_vel
|
||||
def build_scene(kind, level):
|
||||
cfg = evaluation_env_cfg(checkpoint, level, kind)
|
||||
cfg.scene.num_envs = 1
|
||||
# Playback only needs the actor and the world-model preview; training
|
||||
# rewards/scans are not consumed here.
|
||||
cfg.rewards = {}
|
||||
cfg.metrics = {}
|
||||
cfg.observations.pop("critic", None)
|
||||
cfg.episode_length_s = 1e9
|
||||
cfg.commands["twist"] = keyboard_command_cfg(state)
|
||||
cfg.viewer.width, cfg.viewer.height = 640, 480
|
||||
cfg.viewer.enable_shadows = args.shadows
|
||||
cfg.viewer.enable_reflections = args.shadows
|
||||
cfg.viewer.distance, cfg.viewer.elevation, cfg.viewer.azimuth = 2.5, -20., 140.
|
||||
new_env = ManagerBasedRlEnv(cfg=cfg, device=args.device, render_mode="rgb_array")
|
||||
if checkpoint["format"] == "s10-terrain-policy-v1":
|
||||
from s10_supermimic.core.terrain import TerrainBridge, TerrainField
|
||||
field = TerrainField.from_env(new_env)
|
||||
bridge = TerrainBridge(new_env, field)
|
||||
else:
|
||||
field = None
|
||||
bridge = S10Bridge(new_env)
|
||||
if checkpoint["env_action_order"] != bridge.to_env or abs(checkpoint["dt"] - new_env.step_dt) > 1e-8:
|
||||
new_env.close()
|
||||
raise ValueError("Checkpoint action order/control timestep mismatch")
|
||||
obs, _ = new_env.reset()
|
||||
preview = None
|
||||
if not args.no_ghost:
|
||||
try:
|
||||
preview = WorldModelPreview(
|
||||
checkpoint, new_env, bridge, field, actor,
|
||||
horizon=args.ghost_horizon, device=args.device,
|
||||
env_action_order=checkpoint["env_action_order"])
|
||||
except Exception as exc:
|
||||
print(f"World-model preview disabled: {exc}", flush=True)
|
||||
return new_env, bridge, obs, preview
|
||||
|
||||
kind, level = args.terrain_kind, args.terrain_level
|
||||
env, bridge, obs, preview = build_scene(kind, level)
|
||||
metadata = dict(checkpoint=args.checkpoint.name, iteration=checkpoint["step"],
|
||||
terrain=checkpoint["format"] == "s10-terrain-policy-v1", device=args.device,
|
||||
terrain_revision=checkpoint["config"].get("terrain_revision", 1))
|
||||
metadata.update(scene=f"{kind} · 等级 {level}", terrain_kind=kind, terrain_level=level)
|
||||
state.publish([0., 0., 0.], 0, **metadata)
|
||||
worker = threading.Thread(target=server.serve_forever, daemon=True)
|
||||
worker.start()
|
||||
print(f"S10 GUI ready: http://127.0.0.1:{server.server_port}/", flush=True)
|
||||
print(f"One environment on {args.device}. Ctrl+C exits. Commands expire after 2 seconds without a browser heartbeat.", flush=True)
|
||||
step, next_render = 0, 0.
|
||||
stats_since = time.monotonic()
|
||||
stats_step, frames = 0, 0
|
||||
render_fps = realtime_factor = 0.
|
||||
previous = torch.zeros(1, 16, device=args.device)
|
||||
ghosts = []
|
||||
with torch.inference_mode():
|
||||
while not args.steps or step < args.steps:
|
||||
began = time.monotonic()
|
||||
selection = state.take_scene()
|
||||
if selection is not None:
|
||||
# Release the old simulator before constructing the new one.
|
||||
state.clear()
|
||||
env.close()
|
||||
del bridge, obs
|
||||
d = None
|
||||
env = None
|
||||
gc.collect()
|
||||
if args.device.startswith("cuda"):
|
||||
torch.cuda.empty_cache()
|
||||
kind, level = selection
|
||||
env, bridge, obs, preview = build_scene(kind, level)
|
||||
previous.zero_()
|
||||
ghosts = []
|
||||
metadata.update(scene=f"{kind} · 等级 {level}", terrain_kind=kind, terrain_level=level)
|
||||
state.finish_scene()
|
||||
stats_since = next_render = time.monotonic()
|
||||
stats_step, frames = step, 0
|
||||
if state.take_reset():
|
||||
obs, _ = env.reset()
|
||||
previous.zero_()
|
||||
ghosts = []
|
||||
if preview is not None:
|
||||
preview.clear()
|
||||
with state.lock:
|
||||
paused = state.paused
|
||||
if not paused:
|
||||
state_before, pose_before = bridge.state(), bridge.pose()
|
||||
if preview is not None:
|
||||
preview.update_errors(step, state_before, pose_before)
|
||||
ghosts = preview.predict(
|
||||
obs["actor"], previous, env.command_manager.get_command("twist"),
|
||||
bridge.context(), bridge.robot.data.encoder_bias[:, bridge.legs],
|
||||
state_before, pose_before)
|
||||
preview.push_errors(step, ghosts)
|
||||
action = actor(obs["actor"])
|
||||
if not torch.isfinite(action).all():
|
||||
raise RuntimeError("Non-finite policy action")
|
||||
obs, _, terminated, truncated, _ = env.step(bridge.env_action(action))
|
||||
previous = action
|
||||
if not torch.isfinite(obs["actor"]).all():
|
||||
raise RuntimeError("Non-finite observation")
|
||||
if (terminated | truncated).any():
|
||||
state.clear()
|
||||
previous.zero_()
|
||||
ghosts = []
|
||||
if preview is not None:
|
||||
preview.clear()
|
||||
step += 1
|
||||
d = bridge.robot.data
|
||||
actual = [d.root_link_lin_vel_b[0, 0].item(), d.root_link_lin_vel_b[0, 1].item(),
|
||||
d.root_link_ang_vel_b[0, 2].item()]
|
||||
state.publish(actual, step, render_fps=render_fps, realtime_factor=realtime_factor,
|
||||
wm_error=(preview.summary() if preview is not None else ""),
|
||||
wm_horizon=(args.ghost_horizon if preview is not None else 0),
|
||||
**metadata)
|
||||
if began >= next_render:
|
||||
# Only the simulation thread touches the native renderer camera.
|
||||
with state.lock:
|
||||
azimuth, elevation, distance = state.camera
|
||||
camera = env._offline_renderer._cam
|
||||
camera.azimuth, camera.elevation, camera.distance = azimuth, elevation, distance
|
||||
rgb = (preview.render(ghosts, stride=args.ghost_stride)
|
||||
if preview is not None and ghosts else env.render())
|
||||
buffer = BytesIO()
|
||||
Image.fromarray(rgb).save(buffer, format="JPEG", quality=80)
|
||||
with state.lock:
|
||||
state.jpeg = buffer.getvalue()
|
||||
frames += 1
|
||||
# Schedule from frame start, not completion: rendering time is part
|
||||
# of the frame budget rather than an extra delay on every frame.
|
||||
next_render += 1 / args.fps
|
||||
if next_render < began:
|
||||
next_render = began + 1 / args.fps
|
||||
now = time.monotonic()
|
||||
if now - stats_since >= 1.:
|
||||
render_fps = frames / (now - stats_since)
|
||||
realtime_factor = (step - stats_step) * env.step_dt / (now - stats_since)
|
||||
frames, stats_step, stats_since = 0, step, now
|
||||
time.sleep(max(0., env.step_dt - (time.monotonic() - began)))
|
||||
print(f"Render: {render_fps:.1f} FPS | realtime: {realtime_factor:.2f}x", flush=True)
|
||||
except KeyboardInterrupt:
|
||||
print("Closing S10 GUI.", flush=True)
|
||||
finally:
|
||||
state.clear()
|
||||
if worker is not None:
|
||||
server.shutdown()
|
||||
worker.join(timeout=2.)
|
||||
server.server_close()
|
||||
if env is not None:
|
||||
env.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,13 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
project_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." && pwd)"
|
||||
cd "$project_dir"
|
||||
if [[ -n "${S10_PYTHON:-}" && -x "${S10_PYTHON}" ]]; then
|
||||
gui_python="$S10_PYTHON"
|
||||
elif [[ -x "$project_dir/.venv/bin/python" ]]; then
|
||||
gui_python="$project_dir/.venv/bin/python"
|
||||
else
|
||||
gui_python=python
|
||||
fi
|
||||
export OMP_NUM_THREADS="${OMP_NUM_THREADS:-1}"
|
||||
exec "$gui_python" scripts/play_gui.py "$@"
|
||||
@@ -0,0 +1,260 @@
|
||||
"""Keyboard velocity control of the trained S10 policy (one environment)."""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
from contextlib import ExitStack
|
||||
from dataclasses import dataclass, field
|
||||
import math
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
DEFAULT_CHECKPOINT = ROOT / "logs/s10_terrain_v4_20k_full/model_20000.pt"
|
||||
|
||||
|
||||
@dataclass
|
||||
class KeyboardInput:
|
||||
speed: float = 0.6
|
||||
lateral_speed: float = 0.2
|
||||
turn_speed: float = 0.4
|
||||
keys: set[str] = field(default_factory=set)
|
||||
|
||||
def clear(self):
|
||||
self.keys.clear()
|
||||
|
||||
def command(self):
|
||||
if "space" in self.keys:
|
||||
return (0., 0., 0.)
|
||||
x = float("up" in self.keys) - float("down" in self.keys)
|
||||
y = float("left" in self.keys) - float("right" in self.keys)
|
||||
scale = max(1., math.hypot(x, y))
|
||||
yaw = float("q" in self.keys) - float("e" in self.keys)
|
||||
return (self.speed * x / scale, self.lateral_speed * y / scale, self.turn_speed * yaw)
|
||||
|
||||
|
||||
def configure_rendering(headless, backend="auto"):
|
||||
backend = ("egl" if headless else "glfw") if backend == "auto" else backend
|
||||
os.environ["MUJOCO_GL"] = backend
|
||||
os.environ["SDL_FRAMEBUFFER_ACCELERATION"] = "0"
|
||||
os.environ["SDL_RENDER_DRIVER"] = "software"
|
||||
os.environ.setdefault("PYGAME_HIDE_SUPPORT_PROMPT", "1")
|
||||
return backend
|
||||
|
||||
|
||||
def keyboard_command_cfg(controls):
|
||||
from mjlab.tasks.velocity.mdp.velocity_command import UniformVelocityCommand, UniformVelocityCommandCfg
|
||||
|
||||
class KeyboardVelocity(UniformVelocityCommand):
|
||||
def _resample_command(self, env_ids):
|
||||
self._update_command(env_ids)
|
||||
|
||||
def _update_command(self, env_ids=None):
|
||||
# Commands are in the robot body frame, exactly as during training.
|
||||
for axis, value in enumerate(controls.command()):
|
||||
self.vel_command_b[:, axis] = value
|
||||
|
||||
def reset(self, env_ids=None):
|
||||
controls.clear()
|
||||
return super().reset(env_ids)
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class KeyboardVelocityCfg(UniformVelocityCommandCfg):
|
||||
def build(self, env):
|
||||
return KeyboardVelocity(self, env)
|
||||
|
||||
return KeyboardVelocityCfg(
|
||||
entity_name="robot", resampling_time_range=(1e9, 1e9),
|
||||
heading_command=False, rel_heading_envs=0., rel_standing_envs=0.,
|
||||
rel_forward_envs=0., debug_vis=True,
|
||||
ranges=UniformVelocityCommandCfg.Ranges(
|
||||
lin_vel_x=(-controls.speed, controls.speed),
|
||||
lin_vel_y=(-controls.lateral_speed, controls.lateral_speed),
|
||||
ang_vel_z=(-controls.turn_speed, controls.turn_speed)),
|
||||
)
|
||||
|
||||
|
||||
def main():
|
||||
p = argparse.ArgumentParser(description=__doc__)
|
||||
p.add_argument("--terrain-level", type=int, default=0, choices=range(6))
|
||||
p.add_argument("--terrain-kind", choices=("flat","up","down","platform","slope","rubble"), default="rubble")
|
||||
p.add_argument("--checkpoint", type=Path, default=DEFAULT_CHECKPOINT)
|
||||
p.add_argument("--device", default="cpu", help="One CPU environment by default to limit GPU memory.")
|
||||
p.add_argument("--speed", type=float, default=0.6)
|
||||
p.add_argument("--lateral-speed", type=float, default=0.2)
|
||||
p.add_argument("--turn-speed", type=float, default=0.4)
|
||||
p.add_argument("--width", type=int, default=800)
|
||||
p.add_argument("--height", type=int, default=600)
|
||||
p.add_argument("--headless", action="store_true")
|
||||
p.add_argument("--steps", type=int, default=0, help="0: interactive until Esc; otherwise exit after N steps.")
|
||||
p.add_argument("--demo", action="store_true", help="Automatically exercise forward/stop/left/turn/stop.")
|
||||
p.add_argument("--video", type=Path, help="Stream rendered frames to MP4.")
|
||||
p.add_argument("--ghost-horizon", type=int, default=32,
|
||||
help="World-model future steps drawn as ghost poses")
|
||||
p.add_argument("--ghost-stride", type=int, default=4,
|
||||
help="Draw every Nth ghost (1 = all, 4 = 8 ghosts)")
|
||||
p.add_argument("--no-ghost", action="store_true",
|
||||
help="Disable the world-model ghost/error overlay")
|
||||
p.add_argument("--gl-backend", choices=("auto", "glfw", "egl"), default="auto")
|
||||
args = p.parse_args()
|
||||
if not args.checkpoint.is_file():
|
||||
p.error(f"Checkpoint not found: {args.checkpoint}")
|
||||
if args.steps < 0 or (args.headless and args.steps == 0):
|
||||
p.error("Headless mode requires --steps > 0; steps cannot be negative")
|
||||
if args.ghost_horizon < 1 or args.ghost_stride < 1:
|
||||
p.error("ghost-horizon and ghost-stride must be positive")
|
||||
if min(args.width, args.height) < 320 or args.width % 2 or args.height % 2:
|
||||
p.error("width and height must be even and at least 320")
|
||||
for name, value, maximum in (("speed", args.speed, 1.), ("lateral-speed", args.lateral_speed, 0.2),
|
||||
("turn-speed", args.turn_speed, 0.5)):
|
||||
if not math.isfinite(value) or not 0 < value <= maximum:
|
||||
p.error(f"{name} must be in (0, {maximum}] to stay within the training command range")
|
||||
backend = configure_rendering(args.headless, args.gl_backend)
|
||||
|
||||
import torch
|
||||
from mjlab.envs import ManagerBasedRlEnv
|
||||
from s10_supermimic.core.bridge import S10Bridge
|
||||
from s10_supermimic.core.policy import load_actor
|
||||
from s10_supermimic.core.world_preview import WorldModelPreview, annotate_rgb
|
||||
|
||||
pg = None
|
||||
if not args.headless:
|
||||
try:
|
||||
import pygame as pg
|
||||
except ImportError:
|
||||
p.error("Keyboard window requires pygame: .venv/bin/python -m pip install 'pygame>=2.5'")
|
||||
checkpoint = torch.load(args.checkpoint, map_location=args.device, weights_only=True)
|
||||
actor = load_actor(checkpoint, args.device)
|
||||
controls = KeyboardInput(args.speed, args.lateral_speed, args.turn_speed)
|
||||
from s10_supermimic.core.policy import evaluation_env_cfg
|
||||
cfg = evaluation_env_cfg(checkpoint, args.terrain_level, args.terrain_kind)
|
||||
cfg.scene.num_envs = 1
|
||||
cfg.seed = 0
|
||||
cfg.episode_length_s = 1e9 # No automatic timed teleport during keyboard play.
|
||||
cfg.commands["twist"] = keyboard_command_cfg(controls)
|
||||
cfg.viewer.width, cfg.viewer.height = args.width, args.height
|
||||
cfg.viewer.distance, cfg.viewer.elevation, cfg.viewer.azimuth = 2.5, -20., 140.
|
||||
|
||||
with ExitStack() as stack:
|
||||
env = ManagerBasedRlEnv(cfg=cfg, device=args.device,
|
||||
render_mode="rgb_array" if pg or args.video else None)
|
||||
stack.callback(env.close)
|
||||
obs, _ = env.reset()
|
||||
if checkpoint["format"] == "s10-terrain-policy-v1":
|
||||
from s10_supermimic.core.terrain import TerrainBridge, TerrainField
|
||||
field = TerrainField.from_env(env)
|
||||
bridge = TerrainBridge(env, field)
|
||||
else:
|
||||
field = None
|
||||
bridge = S10Bridge(env)
|
||||
if checkpoint["env_action_order"] != bridge.to_env or abs(checkpoint["dt"] - env.step_dt) > 1e-8:
|
||||
raise ValueError("Checkpoint action order/control timestep differs from this environment")
|
||||
preview = None
|
||||
if not args.no_ghost:
|
||||
try:
|
||||
preview = WorldModelPreview(
|
||||
checkpoint, env, bridge, field, actor,
|
||||
horizon=args.ghost_horizon, device=args.device,
|
||||
env_action_order=checkpoint["env_action_order"])
|
||||
except Exception as exc:
|
||||
print(f"World-model preview disabled: {exc}", flush=True)
|
||||
if pg:
|
||||
stack.callback(pg.quit)
|
||||
# Open SDL only after MuJoCo initialization; SDL displays RGB in software.
|
||||
pg.display.init()
|
||||
pg.font.init()
|
||||
screen = pg.display.set_mode((args.width, args.height))
|
||||
pg.display.set_caption("S10 | Arrows move | Q/E turn | Space stop | R reset | Esc quit")
|
||||
font = pg.font.Font(None, 22)
|
||||
clock = pg.time.Clock()
|
||||
key_names = {pg.K_UP: "up", pg.K_DOWN: "down", pg.K_LEFT: "left", pg.K_RIGHT: "right",
|
||||
pg.K_q: "q", pg.K_e: "e", pg.K_SPACE: "space"}
|
||||
video = None
|
||||
if args.video:
|
||||
import mediapy
|
||||
args.video.parent.mkdir(parents=True, exist_ok=True)
|
||||
video = stack.enter_context(mediapy.VideoWriter(str(args.video), shape=(args.height, args.width),
|
||||
fps=round(1 / env.step_dt)))
|
||||
iteration = checkpoint["step"]
|
||||
del checkpoint # Drop optimizer/world-model tensors; inference uses actor only.
|
||||
print(f"Loaded S10 iteration {iteration}: {args.checkpoint}", flush=True)
|
||||
print(f"One environment | {args.device} | {backend} | {env.step_dt:.3f}s control period", flush=True)
|
||||
print("Arrows: forward/back/strafe | Q/E: turn | release/Space: zero command | R: reset | Esc: quit", flush=True)
|
||||
print("Keyboard changes enter the next observation (one control-step latency).", flush=True)
|
||||
previous = torch.zeros(1, 16, device=args.device)
|
||||
ghosts = []
|
||||
step, running = 0, True
|
||||
with torch.inference_mode():
|
||||
while running and (args.steps == 0 or step < args.steps):
|
||||
if pg:
|
||||
for event in pg.event.get():
|
||||
if event.type == pg.QUIT or (event.type == pg.KEYDOWN and event.key == pg.K_ESCAPE):
|
||||
running = False
|
||||
elif event.type == pg.WINDOWFOCUSLOST:
|
||||
controls.clear()
|
||||
elif event.type == pg.KEYDOWN:
|
||||
if event.key == pg.K_r:
|
||||
controls.clear()
|
||||
obs, _ = env.reset()
|
||||
previous.zero_()
|
||||
ghosts = []
|
||||
if preview is not None:
|
||||
preview.clear()
|
||||
elif event.key in key_names:
|
||||
controls.keys.add(key_names[event.key])
|
||||
elif event.type == pg.KEYUP and event.key in key_names:
|
||||
controls.keys.discard(key_names[event.key])
|
||||
if not running:
|
||||
break
|
||||
if args.demo:
|
||||
phases = ({"up"}, set(), {"left"}, {"q"}, set())
|
||||
controls.keys = set(phases[(step // 100) % len(phases)])
|
||||
state_before, pose_before = bridge.state(), bridge.pose()
|
||||
if preview is not None:
|
||||
preview.update_errors(step, state_before, pose_before)
|
||||
ghosts = preview.predict(
|
||||
obs["actor"], previous, env.command_manager.get_command("twist"),
|
||||
bridge.context(), bridge.robot.data.encoder_bias[:, bridge.legs],
|
||||
state_before, pose_before)
|
||||
preview.push_errors(step, ghosts)
|
||||
action = actor(obs["actor"])
|
||||
if not torch.isfinite(action).all():
|
||||
raise RuntimeError("Non-finite policy action")
|
||||
obs, _, terminated, _, _ = env.step(bridge.env_action(action))
|
||||
previous = action
|
||||
if not torch.isfinite(obs["actor"]).all():
|
||||
raise RuntimeError("Non-finite observation")
|
||||
if terminated.any():
|
||||
controls.clear()
|
||||
previous.zero_()
|
||||
ghosts = []
|
||||
if preview is not None:
|
||||
preview.clear()
|
||||
print("Fall reset; keys cleared.", flush=True)
|
||||
if preview is not None and step % 20 == 0 and step > 0:
|
||||
print("WM " + preview.summary(), flush=True)
|
||||
if pg or video:
|
||||
rgb = (preview.render(ghosts, stride=args.ghost_stride)
|
||||
if preview is not None and ghosts else env.render())
|
||||
if video:
|
||||
lines = ([f"WM {preview.summary()}"] if preview is not None else [])
|
||||
video.add_image(annotate_rgb(rgb, lines))
|
||||
if pg:
|
||||
pg.surfarray.blit_array(screen, rgb.swapaxes(0, 1))
|
||||
vx, vy, yaw = controls.command()
|
||||
actual = bridge.robot.data
|
||||
lines = [f"S10 model {iteration} | Arrows move | Q/E turn | Space stop | R reset | Esc exit",
|
||||
f"Command: vx {vx:+.2f} vy {vy:+.2f} m/s yaw {yaw:+.2f} rad/s",
|
||||
f"Actual: vx {actual.root_link_lin_vel_b[0, 0]:+.2f} vy {actual.root_link_lin_vel_b[0, 1]:+.2f}"
|
||||
f" yaw {actual.root_link_ang_vel_b[0, 2]:+.2f}"]
|
||||
if preview is not None:
|
||||
lines.append("WM " + preview.summary())
|
||||
for i, line in enumerate(lines):
|
||||
screen.blit(font.render(line, True, (240, 240, 240), (20, 24, 30)), (8, 8 + i * 22))
|
||||
pg.display.flip()
|
||||
clock.tick(round(1 / env.step_dt))
|
||||
step += 1
|
||||
print(f"Completed {step} control steps.", flush=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,13 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
project_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." && pwd)"
|
||||
cd "$project_dir"
|
||||
if [[ -n "${S10_PYTHON:-}" && -x "${S10_PYTHON}" ]]; then
|
||||
keyboard_python="$S10_PYTHON"
|
||||
elif [[ -x "$project_dir/.venv/bin/python" ]]; then
|
||||
keyboard_python="$project_dir/.venv/bin/python"
|
||||
else
|
||||
keyboard_python=python
|
||||
fi
|
||||
export OMP_NUM_THREADS="${OMP_NUM_THREADS:-1}"
|
||||
exec "$keyboard_python" scripts/play_keyboard.py "$@"
|
||||
@@ -0,0 +1,29 @@
|
||||
#!/usr/bin/env bash
|
||||
# S10 terrain-v4 50k training from scratch (16 GB recipe).
|
||||
# Override S10_OUTPUT for a fixed output directory.
|
||||
set -euo pipefail
|
||||
|
||||
PROJECT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." && pwd)"
|
||||
cd "$PROJECT_DIR"
|
||||
|
||||
export PYTHONPATH="$PROJECT_DIR"
|
||||
export OMP_NUM_THREADS="${OMP_NUM_THREADS:-1}"
|
||||
export MUJOCO_GL="${MUJOCO_GL:-egl}"
|
||||
|
||||
PY="${S10_PYTHON:-python}"
|
||||
OUT="${S10_OUTPUT:-logs/s10_terrain_v4_50k_$(date +%Y%m%d_%H%M%S)}"
|
||||
mkdir -p "$OUT"
|
||||
|
||||
exec "$PY" scripts/train.py \
|
||||
--task s10-v4 \
|
||||
--terrain --terrain-revision 4 --seed 0 \
|
||||
--num-envs 4096 --num-steps-per-env 1 \
|
||||
--buffer-len 512 --wm-window 8 \
|
||||
--wm-batch 2048 --wm-updates 4 \
|
||||
--imag-batch 2048 --horizon 32 \
|
||||
--warmup 50 --critic-epochs 5 --critic-minibatches 4 \
|
||||
--iterations 50000 --save-interval 500 --log-interval 10 \
|
||||
--command-curriculum --command-curriculum-start 0 \
|
||||
--prioritized-replay \
|
||||
--video --video-length 500 --video-interval 2700 \
|
||||
--output "$OUT"
|
||||
@@ -0,0 +1,25 @@
|
||||
#!/usr/bin/env bash
|
||||
# S10 terrain-v7 safe all-terrain 50k training from the task preset.
|
||||
# Set S10_RESUME to continue from a checkpoint; S10_VIDEO=0 disables video.
|
||||
set -euo pipefail
|
||||
|
||||
PROJECT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." && pwd)"
|
||||
cd "$PROJECT_DIR"
|
||||
|
||||
export PYTHONPATH="$PROJECT_DIR"
|
||||
export OMP_NUM_THREADS="${OMP_NUM_THREADS:-1}"
|
||||
export MUJOCO_GL="${MUJOCO_GL:-egl}"
|
||||
|
||||
PY="${S10_PYTHON:-python}"
|
||||
OUT="${S10_OUTPUT:-logs/s10_terrain_v7_50k_$(date +%Y%m%d_%H%M%S)}"
|
||||
mkdir -p "$OUT"
|
||||
|
||||
args=(--task s10-v7 --iterations 50000 --output "$OUT")
|
||||
if [[ -n "${S10_RESUME:-}" ]]; then
|
||||
args+=(--resume "$S10_RESUME")
|
||||
fi
|
||||
if [[ "${S10_VIDEO:-1}" != "0" ]]; then
|
||||
args+=(--video --video-length 500 --video-interval 2700)
|
||||
fi
|
||||
|
||||
exec "$PY" scripts/train.py "${args[@]}" "$@"
|
||||
@@ -0,0 +1,5 @@
|
||||
"""Independent online S10 model-based training entry point."""
|
||||
from s10_supermimic.core.train import main
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,151 @@
|
||||
"""Standalone S10 MJCF/FK checks that do not require mjlab or torch.
|
||||
|
||||
Run with the system Python:
|
||||
python3 tests/test_s10_model.py
|
||||
|
||||
The test validates:
|
||||
1. The S10 model in ``s10_supermimic/robot/model/s10.xml`` loads.
|
||||
2. Joint names match the training/deployment contract.
|
||||
3. The default pose puts the wheel bottom at ~0.434 m above the ground.
|
||||
4. The analytic wheel-bottom formula in ``core/rolling_rewards.py`` matches
|
||||
MuJoCo forward kinematics for random joint angles and arbitrary base tilt.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from pathlib import Path
|
||||
|
||||
import mujoco
|
||||
import numpy as np
|
||||
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
MODEL = ROOT / "s10_supermimic" / "robot" / "model" / "s10.xml"
|
||||
LEGS = ("fl", "fr", "hl", "hr")
|
||||
LEG_JOINT_NAMES = tuple(
|
||||
f"{leg}_{joint}_joint" for leg in LEGS for joint in ("hipx", "hipy", "knee")
|
||||
)
|
||||
WHEEL_JOINT_NAMES = tuple(f"{leg}_wheel_joint" for leg in LEGS)
|
||||
DEFAULT_Q = {
|
||||
"fl": (0.1, -0.35, 0.65),
|
||||
"fr": (-0.1, -0.35, 0.65),
|
||||
"hl": (0.1, 0.35, -0.65),
|
||||
"hr": (-0.1, 0.35, -0.65),
|
||||
}
|
||||
# Geometry from S10.xml.
|
||||
X0 = {"fl": 0.2277, "fr": 0.2277, "hl": -0.2277, "hr": -0.2277}
|
||||
Y0 = {"fl": 0.0574, "fr": -0.0574, "hl": 0.0574, "hr": -0.0574}
|
||||
SIDE = {"fl": 1.0, "fr": -1.0, "hl": 1.0, "hr": -1.0}
|
||||
HIPY_Y = 0.0779
|
||||
WHEEL_Y = 0.045891
|
||||
LINK = 0.18
|
||||
WHEEL_RADIUS = 0.081
|
||||
WHEEL_HALF_WIDTH = 0.0185
|
||||
|
||||
|
||||
def analytic_wheel_bottom(q, leg, base_z=0.5, up=(0.0, 0.0, 1.0)) -> float:
|
||||
"""NumPy translation of the torch implementation in rolling_rewards.py."""
|
||||
hipx, hipy, knee = (float(v) for v in q)
|
||||
side = SIDE[leg]
|
||||
x = LINK * math.sin(knee)
|
||||
z = -LINK - LINK * math.cos(knee)
|
||||
x3 = math.cos(hipy) * x - math.sin(hipy) * z
|
||||
z3 = math.sin(hipy) * x + math.cos(hipy) * z
|
||||
y3 = side * (HIPY_Y + WHEEL_Y)
|
||||
y4 = math.cos(hipx) * y3 + math.sin(hipx) * z3
|
||||
z4 = -math.sin(hipx) * y3 + math.cos(hipx) * z3
|
||||
point = np.array([X0[leg] + x3, Y0[leg] + y4, z4])
|
||||
up = np.asarray(up, dtype=float)
|
||||
up /= np.linalg.norm(up)
|
||||
axle_z = -up[1] * math.cos(hipx) + up[2] * math.sin(hipx)
|
||||
radius_z = WHEEL_RADIUS * math.sqrt(max(1.0 - axle_z * axle_z, 0.0))
|
||||
radius_z += WHEEL_HALF_WIDTH * abs(axle_z)
|
||||
return base_z + float(point @ up) - radius_z
|
||||
|
||||
|
||||
def wheel_geom_id(model: mujoco.MjModel, leg: str) -> int:
|
||||
body_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, f"{leg}_wheel")
|
||||
for geom_id in range(model.ngeom):
|
||||
if model.geom_bodyid[geom_id] != body_id:
|
||||
continue
|
||||
if model.geom_type[geom_id] == mujoco.mjtGeom.mjGEOM_CYLINDER:
|
||||
return geom_id
|
||||
raise AssertionError(f"No cylinder collision geom for {leg}_wheel")
|
||||
|
||||
|
||||
def measured_wheel_bottom(model, data, geom_id: int) -> float:
|
||||
axis = data.geom_xmat[geom_id].reshape(3, 3)[:, 2]
|
||||
axle_z = abs(float(axis[2]))
|
||||
radius_z = WHEEL_RADIUS * math.sqrt(max(1.0 - axle_z * axle_z, 0.0))
|
||||
radius_z += WHEEL_HALF_WIDTH * axle_z
|
||||
return float(data.geom_xpos[geom_id][2] - radius_z)
|
||||
|
||||
|
||||
def set_pose(model, data, q_map, base_z=0.5, quat=(1.0, 0.0, 0.0, 0.0)):
|
||||
data.qpos[:] = 0.0
|
||||
data.qpos[2] = base_z
|
||||
data.qpos[3:7] = quat
|
||||
for joint_name, value in q_map.items():
|
||||
joint_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, joint_name)
|
||||
data.qpos[model.jnt_qposadr[joint_id]] = value
|
||||
mujoco.mj_forward(model, data)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
assert MODEL.is_file(), MODEL
|
||||
model = mujoco.MjModel.from_xml_path(str(MODEL))
|
||||
data = mujoco.MjData(model)
|
||||
|
||||
joint_names = [model.joint(i).name for i in range(model.njnt)]
|
||||
expected = ("floating_base",) + LEG_JOINT_NAMES + WHEEL_JOINT_NAMES
|
||||
assert set(joint_names) == set(expected), joint_names
|
||||
assert model.nq == 23 and model.nv == 22 and model.njnt == 17
|
||||
assert model.nu == 0, "Robot XML must not contain actuators; mjlab adds them."
|
||||
|
||||
default_q = {
|
||||
f"{leg}_{joint}_joint": value
|
||||
for leg in LEGS
|
||||
for joint, value in zip(("hipx", "hipy", "knee"), DEFAULT_Q[leg])
|
||||
}
|
||||
set_pose(model, data, default_q)
|
||||
for leg in LEGS:
|
||||
geom_id = wheel_geom_id(model, leg)
|
||||
bottom = measured_wheel_bottom(model, data, geom_id)
|
||||
pred = analytic_wheel_bottom(DEFAULT_Q[leg], leg)
|
||||
assert abs(bottom - pred) < 1e-9, (leg, bottom, pred)
|
||||
assert abs((0.5 - bottom) - 0.434145) < 1e-4, (leg, bottom)
|
||||
|
||||
rng = np.random.default_rng(0)
|
||||
max_error = 0.0
|
||||
for _ in range(200):
|
||||
quat = rng.normal(size=4)
|
||||
quat /= np.linalg.norm(quat)
|
||||
q_map = {}
|
||||
q_by_leg = {}
|
||||
for leg in LEGS:
|
||||
q = np.array([
|
||||
rng.uniform(-0.6109, 0.6109),
|
||||
rng.uniform(-2.5307, 2.5307),
|
||||
rng.uniform(-2.7227, 2.7227),
|
||||
])
|
||||
q_by_leg[leg] = q
|
||||
for joint, value in zip(("hipx", "hipy", "knee"), q):
|
||||
q_map[f"{leg}_{joint}_joint"] = value
|
||||
set_pose(model, data, q_map, base_z=0.5, quat=quat)
|
||||
|
||||
rotation = data.xmat[
|
||||
mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "base_link")
|
||||
].reshape(3, 3)
|
||||
up = rotation.T @ np.array([0.0, 0.0, 1.0])
|
||||
for leg in LEGS:
|
||||
geom_id = wheel_geom_id(model, leg)
|
||||
measured = measured_wheel_bottom(model, data, geom_id)
|
||||
pred = analytic_wheel_bottom(q_by_leg[leg], leg, base_z=0.5, up=up)
|
||||
max_error = max(max_error, abs(measured - pred))
|
||||
assert max_error < 1e-9, max_error
|
||||
print(f"S10 model/FK checks passed (max error {max_error:.2e} m)")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,91 @@
|
||||
"""Task registry and dual Actor observation schema checks (no mjlab needed)."""
|
||||
|
||||
import math
|
||||
|
||||
import torch
|
||||
|
||||
from s10_supermimic.core.command_curriculum import stages_for
|
||||
from s10_supermimic.core.policy import Actor, LEGACY_FRAME_DIM, FRAME_DIM, frame, push_history
|
||||
from s10_supermimic.core.wheel_intrusion import pair_body_intrusion, wheel_body_intrusion
|
||||
from s10_supermimic.tasks import get_task, task_names
|
||||
|
||||
|
||||
def main():
|
||||
assert task_names() == ("s10-v4", "s10-v4-upward", "s10-v7"), task_names()
|
||||
v4, v7, upward = get_task("s10-v4"), get_task("s10-v7"), get_task("s10-v4-upward")
|
||||
assert (v4.terrain_revision, v7.terrain_revision) == (4, 7)
|
||||
assert upward.terrain_revision == 4
|
||||
assert upward.actor.frame_dim == LEGACY_FRAME_DIM
|
||||
assert upward.actor.wheel_vel
|
||||
assert upward.courses is False and upward.collisions is False
|
||||
assert upward.defaults["iterations"] == 30000
|
||||
assert upward.defaults["action_noise"] == 0.22
|
||||
assert upward.defaults["horizon"] == 64
|
||||
assert (v4.actor.frame_dim, v7.actor.frame_dim) == (LEGACY_FRAME_DIM, LEGACY_FRAME_DIM)
|
||||
assert (v4.actor.history_dim(5), v7.actor.history_dim(5)) == (265, 265)
|
||||
assert v4.actor.wheel_vel and v7.actor.wheel_vel
|
||||
assert v7.actor.name == "s10-v7-53"
|
||||
assert v4.courses is False and v4.collisions is False
|
||||
assert v7.courses and v7.collisions
|
||||
assert v7.defaults["imag_batch"] == 2048, "v7 keeps the 16 GB batch-size exception"
|
||||
assert v7.defaults["iterations"] == 40000
|
||||
assert v7.defaults["command_curriculum"] is True, "v7 uses stair-gated stages capped at 1.5/0.8/1.2"
|
||||
assert stages_for(7)[-1][1:] == ((-1.5, 1.5), (-.8, .8), (-1.2, 1.2))
|
||||
assert stages_for(4)[-1][1:] == ((-2., 2.), (-1., 1.), (-3., 3.))
|
||||
assert v4.defaults["iterations"] == 20000
|
||||
|
||||
state = torch.zeros(2, 38)
|
||||
action = torch.zeros(2, 16)
|
||||
command = torch.zeros(2, 3)
|
||||
bias = torch.zeros(2, 12)
|
||||
for history, terms, expected in ((5, v4.actor.terms, 265), (5, v7.actor.terms, 265)):
|
||||
actor = Actor(history, actor_wheel_vel=True)
|
||||
latest = frame(state, action, command, bias,
|
||||
wheel_vel=actor.actor_wheel_vel)
|
||||
assert latest.shape == (2, actor.frame_dim)
|
||||
assert push_history(torch.zeros(2, expected), latest, history, terms).shape == (2, expected)
|
||||
legacy = Actor(5, actor_wheel_vel=False)
|
||||
assert legacy.frame_dim == FRAME_DIM, "legacy 49-value checkpoints remain loadable"
|
||||
|
||||
# v7 safe stance: default hipx 0.05 / hipy 25 deg / knee 1.0 has zero
|
||||
# intrusion; a wheel tucked under the body crosses the 1/4 boundary.
|
||||
state = torch.zeros(1, 38)
|
||||
state[0, 0], state[0, 3] = 0.44, -1.0
|
||||
def default_q(hipy, knee=1.0):
|
||||
return torch.tensor([0.05, -hipy, knee, -0.05, -hipy, knee,
|
||||
0.05, hipy, -knee, -0.05, hipy, -knee])
|
||||
assert wheel_body_intrusion(state, default_q(math.radians(25.0))).item() == 0.0
|
||||
# 0.8 rad splay puts the wheel near x=0.134 m: inside the new outer-quarter
|
||||
# boundary (0.15 m), outside the old outer-third boundary (0.10 m).
|
||||
assert wheel_body_intrusion(state, default_q(.8)).item() > 0.0
|
||||
normal = wheel_body_intrusion(state, default_q(.5)).item()
|
||||
tucked = wheel_body_intrusion(state, default_q(1.0)).item()
|
||||
assert normal == 0.0 < tucked, (normal, tucked)
|
||||
|
||||
# Pair 5/6 restriction: zero during the 0.5 s grace, then active only when
|
||||
# both same-side legs are inside. A single leg never triggers it.
|
||||
timers = torch.zeros(1, 4)
|
||||
for _ in range(24):
|
||||
cost, timers = pair_body_intrusion(state, default_q(1.0), timers, .02)
|
||||
assert cost.item() == 0.0
|
||||
cost, timers = pair_body_intrusion(state, default_q(1.0), timers, .02)
|
||||
assert cost.item() > 0.0, cost
|
||||
# Front-left alone stays inside; the pair timer stays unarmed.
|
||||
one_leg = default_q(math.radians(25.0))
|
||||
one_leg[1] = -1.0
|
||||
timers = torch.zeros(1, 4)
|
||||
for _ in range(40):
|
||||
cost, timers = pair_body_intrusion(state, one_leg, timers, .02)
|
||||
assert cost.item() == 0.0
|
||||
assert timers[0, 1].item() == 0.0, timers
|
||||
# Leaving the region resets that wheel's timer.
|
||||
timers = torch.zeros(1, 4)
|
||||
for _ in range(30):
|
||||
_, timers = pair_body_intrusion(state, default_q(1.0), timers, .02)
|
||||
cost, timers = pair_body_intrusion(state, default_q(math.radians(25.0)), timers, .02)
|
||||
assert cost.item() == 0.0 and timers[0, 0].item() == 0.0, (cost, timers)
|
||||
print("Task registry / Actor schema / v7 safety checks passed")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,236 @@
|
||||
"""Pure-torch checks for the upward recovery reward and per-env curriculum."""
|
||||
|
||||
import math
|
||||
|
||||
import torch
|
||||
|
||||
from s10_supermimic.core.imagine import continuation_value
|
||||
from s10_supermimic.core.upward import (
|
||||
BAD_ORIENTATION_START_DEG,
|
||||
HISTORY_LENGTH,
|
||||
PROMOTE_SUCCESSES,
|
||||
TILT_STAGES_DEG,
|
||||
bad_orientation_mask,
|
||||
bad_orientation_value,
|
||||
init_upward_state,
|
||||
reset_root_state_upward,
|
||||
upward_curriculum,
|
||||
upward_curriculum_stats,
|
||||
upward_video_tilt_limit,
|
||||
)
|
||||
from s10_supermimic.core.world_model import GRAVITY
|
||||
|
||||
|
||||
def gravity_at(deg: float, direction: float = 0.0) -> torch.Tensor:
|
||||
"""Projected gravity for a tilt ``deg`` in a fixed horizontal direction."""
|
||||
t = math.radians(deg)
|
||||
return torch.tensor([[
|
||||
math.sin(t) * math.cos(direction),
|
||||
math.sin(t) * math.sin(direction),
|
||||
-math.cos(t),
|
||||
]])
|
||||
|
||||
|
||||
def test_severity_is_monotonic() -> None:
|
||||
values = torch.cat([bad_orientation_value(gravity_at(d)) for d in
|
||||
torch.arange(0.0, 180.5, 0.5).tolist()])
|
||||
assert (values[1:] >= values[:-1] - 1e-6).all(), "severity must be monotonic"
|
||||
assert values.max() <= 1.0 + 1e-5, values.max()
|
||||
assert values.min() >= 0.0, values.min()
|
||||
assert bad_orientation_value(gravity_at(BAD_ORIENTATION_START_DEG)).item() < 1e-6
|
||||
assert bad_orientation_value(gravity_at(180.0)).item() > 0.99
|
||||
v105 = bad_orientation_value(gravity_at(105.0)).item()
|
||||
v110 = bad_orientation_value(gravity_at(110.0)).item()
|
||||
v120 = bad_orientation_value(gravity_at(120.0)).item()
|
||||
assert v105 < v110 < v120, (v105, v110, v120)
|
||||
|
||||
|
||||
def test_severity_gradient_is_finite() -> None:
|
||||
for g in (gravity_at(0.0), gravity_at(90.0), gravity_at(180.0)):
|
||||
x = g.clone().requires_grad_(True)
|
||||
bad_orientation_value(x).sum().backward()
|
||||
assert torch.isfinite(x.grad).all(), x.grad
|
||||
|
||||
|
||||
def test_mask_keeps_bad_set() -> None:
|
||||
assert bad_orientation_mask(gravity_at(60.0)).item()
|
||||
assert not bad_orientation_mask(gravity_at(30.0)).item()
|
||||
assert not bad_orientation_mask(gravity_at(50.0, math.pi / 4)).item()
|
||||
assert bad_orientation_mask(gravity_at(180.0)).item()
|
||||
|
||||
|
||||
def test_upward_continuation_does_not_use_v4_fall_proxy() -> None:
|
||||
state = torch.zeros(3, 38)
|
||||
state[0, GRAVITY] = torch.tensor([0.0, 0.0, -1.0])
|
||||
state[1, GRAVITY] = torch.tensor([0.0, 0.0, 0.0])
|
||||
state[2, GRAVITY] = torch.tensor([0.0, 0.0, 1.0])
|
||||
upward = continuation_value(state, {"upright_mode": "upward"})
|
||||
v4 = continuation_value(state, {})
|
||||
assert torch.allclose(upward, torch.ones(3)), upward
|
||||
assert v4[0] > 0.99 and v4[2] < 1e-4, v4
|
||||
|
||||
|
||||
class _Data:
|
||||
def __init__(self, n: int):
|
||||
self.root_link_pos_w = torch.zeros(n, 3)
|
||||
self.projected_gravity_b = torch.tensor([[0.0, 0.0, -1.0]] * n)
|
||||
|
||||
|
||||
class _Robot:
|
||||
def __init__(self, n: int):
|
||||
self.data = _Data(n)
|
||||
|
||||
|
||||
class _Terrain:
|
||||
def __init__(self, n: int):
|
||||
self.terrain_levels = torch.zeros(n, dtype=torch.long)
|
||||
self.terrain_types = torch.zeros(n, dtype=torch.long)
|
||||
self.terrain_origins = torch.zeros(6, 1, 3)
|
||||
self.env_origins = torch.zeros(n, 3)
|
||||
|
||||
|
||||
class _Scene:
|
||||
def __init__(self, n: int):
|
||||
self.robot = _Robot(n)
|
||||
self.terrain = _Terrain(n)
|
||||
self.env_origins = torch.zeros(n, 3)
|
||||
|
||||
def __getitem__(self, key):
|
||||
return self.robot
|
||||
|
||||
|
||||
class _Termination:
|
||||
def __init__(self, n: int):
|
||||
self.terminated = torch.zeros(n, dtype=torch.bool)
|
||||
|
||||
|
||||
class _Env:
|
||||
def __init__(self, n: int):
|
||||
self.num_envs = n
|
||||
self.device = "cpu"
|
||||
self.scene = _Scene(n)
|
||||
self.episode_length_buf = torch.ones(n, dtype=torch.long)
|
||||
self.termination_manager = _Termination(n)
|
||||
|
||||
|
||||
def _episode(env: _Env, recovered: torch.Tensor) -> None:
|
||||
"""Advance one finished episode for every env and run the curriculum."""
|
||||
env.termination_manager.terminated[:] = ~recovered
|
||||
env.scene.robot.data.projected_gravity_b[:] = torch.where(
|
||||
recovered[:, None],
|
||||
torch.tensor([0.0, 0.0, -1.0]),
|
||||
torch.tensor([0.0, 0.0, 1.0]))
|
||||
upward_curriculum(env, torch.arange(env.num_envs))
|
||||
|
||||
|
||||
def test_per_env_rolling_window_advances_and_demotes() -> None:
|
||||
env = _Env(4)
|
||||
outcomes = torch.zeros(4, dtype=torch.bool)
|
||||
for episode in range(HISTORY_LENGTH):
|
||||
outcomes[0] = True # 20/20 -> advance
|
||||
outcomes[1] = False # 0/20 -> stays at level 0
|
||||
outcomes[2] = episode < 10 # 10/20 -> neutral
|
||||
outcomes[3] = episode < PROMOTE_SUCCESSES # exactly 16/20 -> advance
|
||||
_episode(env, outcomes)
|
||||
stats = upward_curriculum_stats(env)
|
||||
assert env._upward_tilt_levels.tolist() == [1, 0, 0, 1]
|
||||
assert env.scene.terrain.terrain_levels.tolist() == [1, 0, 0, 1]
|
||||
assert stats["tilt_level_q75"] == 1
|
||||
# Advanced envs reset their window and need 20 new episodes; only envs 1
|
||||
# (0/20) and 2 (10/20) are ready after the first block.
|
||||
assert abs(stats["recovery_window_rate"] - 0.25) < 1e-6
|
||||
assert abs(stats["recovery_window_ready"] - 0.5) < 1e-6
|
||||
assert math.isclose(upward_video_tilt_limit(env), math.radians(135.0))
|
||||
|
||||
# The history is per-env: env 1 now gets 20 successes and advances alone;
|
||||
# env 2 stays neutral with 10 successes.
|
||||
for episode in range(HISTORY_LENGTH):
|
||||
outcomes[0] = True
|
||||
outcomes[1] = True
|
||||
outcomes[2] = episode < 10
|
||||
outcomes[3] = True
|
||||
_episode(env, outcomes)
|
||||
assert env._upward_tilt_levels.tolist() == [2, 1, 0, 2]
|
||||
assert env.scene.terrain.terrain_levels.tolist() == [2, 1, 0, 2]
|
||||
|
||||
# A full 20-episode failure window demotes env 0 from level 2 to level 1.
|
||||
for _ in range(HISTORY_LENGTH):
|
||||
outcomes[:] = False
|
||||
outcomes[3] = True # keep env 3 at level 2
|
||||
_episode(env, outcomes)
|
||||
assert env._upward_tilt_levels.tolist() == [1, 0, 0, 2]
|
||||
# Terrain has six levels and is independent of the three tilt levels.
|
||||
assert env.scene.terrain.terrain_levels.tolist() == [1, 0, 0, 3]
|
||||
|
||||
# Before 20 episodes no env is eligible to move.
|
||||
fresh = _Env(4)
|
||||
stats = upward_curriculum_stats(fresh)
|
||||
assert stats["recovery_window_ready"] == 0.0
|
||||
assert stats["tilt_level_q75"] == 0
|
||||
outcomes = torch.ones(4, dtype=torch.bool)
|
||||
for _ in range(HISTORY_LENGTH - 1):
|
||||
_episode(fresh, outcomes)
|
||||
assert fresh._upward_tilt_levels.tolist() == [0, 0, 0, 0]
|
||||
_episode(fresh, outcomes)
|
||||
assert fresh._upward_tilt_levels.tolist() == [1, 1, 1, 1]
|
||||
|
||||
|
||||
class _Asset:
|
||||
def __init__(self, n: int):
|
||||
self.data = type("Data", (), {})()
|
||||
self.data.default_root_state = torch.zeros(n, 13)
|
||||
self.data.default_root_state[:, 3] = 1.0 # identity wxyz quaternion
|
||||
self.pose = None
|
||||
self.velocity = None
|
||||
|
||||
def write_root_link_pose_to_sim(self, pose, env_ids):
|
||||
self.pose = pose.clone()
|
||||
|
||||
def write_root_link_velocity_to_sim(self, velocity, env_ids):
|
||||
self.velocity = velocity.clone()
|
||||
|
||||
|
||||
class _ResetScene:
|
||||
def __init__(self, n: int):
|
||||
self.asset = _Asset(n)
|
||||
self.env_origins = torch.zeros(n, 3)
|
||||
|
||||
def __getitem__(self, key):
|
||||
return self.asset
|
||||
|
||||
|
||||
class _ResetEnv:
|
||||
def __init__(self, n: int):
|
||||
self.num_envs = n
|
||||
self.device = "cpu"
|
||||
self.scene = _ResetScene(n)
|
||||
|
||||
|
||||
def test_reset_supports_per_env_and_play_fallback_limits() -> None:
|
||||
env = _ResetEnv(2)
|
||||
pose_range = {"x": (-0.1, 0.1), "y": (-0.1, 0.1), "z": (0.0, 0.2),
|
||||
"roll": (-math.pi, math.pi), "pitch": (-math.pi, math.pi),
|
||||
"yaw": (-math.pi, math.pi)}
|
||||
ids = torch.arange(2)
|
||||
reset_root_state_upward(env, ids, pose_range)
|
||||
assert env.scene.asset.pose.shape == (2, 7)
|
||||
assert env.scene.asset.velocity.shape == (2, 6)
|
||||
assert torch.isfinite(env.scene.asset.pose).all()
|
||||
env._upward_tilt_limits = torch.tensor(
|
||||
[math.radians(60.0), math.radians(180.0)])
|
||||
reset_root_state_upward(env, ids, pose_range)
|
||||
assert torch.isfinite(env.scene.asset.pose).all()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
test_severity_is_monotonic()
|
||||
test_severity_gradient_is_finite()
|
||||
test_mask_keeps_bad_set()
|
||||
test_upward_continuation_does_not_use_v4_fall_proxy()
|
||||
test_per_env_rolling_window_advances_and_demotes()
|
||||
test_reset_supports_per_env_and_play_fallback_limits()
|
||||
print("Upward reward / continuation / per-env curriculum checks passed")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user