feat(training-server): release V0.7.3 整合强化学习服务器

.gitignore

@@ -28,3 +28,6 @@ build/
 # Runtime logs
 *.log
 MUJOCO_LOG.TXT
+training_server/rl/logs/
+training_server/rl/wandb/
+training_server/rl/outputs/

CHANGELOG.md

@@ -2,6 +2,16 @@
 
 本项目的重要变更记录在此文件中，版本标签沿用仓库现有的 `V主版本.次版本[.修订版本]` 格式。
 
+## [0.7.3] - 2026-09-01
+
+### 新增
+
+- 将默认 `Unitree-Go2-Flat` 训练所需的 PPO 环境、奖励/观测配置、ONNX 导出逻辑和 Go2 模型资产内置到 `training_server/rl/`。
+
+### 变更
+
+- 训练服务默认使用仓库内置训练器，不再要求单独克隆 `unitree_rl_mjlab`；仍可通过 `--trainer-root` 覆盖。
+
 ## [0.7.2] - 2026-09-01
 
 ### 变更

README.md

@@ -15,7 +15,7 @@
 - V3 地图创作层：认证资产库支持点击添加、拖到画布落位和首个资产自动创建场景，并可通过表单与视口操纵器继续编辑、事务式应用及导出地图 ZIP
 - 浏览器内 Python 控制器（Pyodide）
 - ONNX 强化学习策略推理（ONNX Runtime Web）
-- 可选的本机 mjlab 训练桥接服务
+- 内置 Go2 PPO 任务的本机 mjlab 训练桥接服务
 - 响应式工作区、源码编辑、性能监控和中文诊断
 
 ## 快速开始
@@ -56,7 +56,7 @@ npm run lint:python
 ```text
 .
 ├── web_platform/      React、TypeScript、Three.js 前端与测试夹具
-├── training_server/   可选的本机强化学习训练桥接服务
+├── training_server/   本机训练桥接服务及内置 Go2 PPO 训练代码/资产
 ├── controllers/       可导入平台的 Python 控制器示例
 ├── package.json       根目录统一开发入口
 └── LICENSE            Apache License 2.0

package-lock.json

@@ -1,12 +1,12 @@
 {
   "name": "mujoco-web-platform",
-  "version": "0.7.2",
+  "version": "0.7.3",
   "lockfileVersion": 3,
   "requires": true,
   "packages": {
     "": {
       "name": "mujoco-web-platform",
-      "version": "0.7.2",
+      "version": "0.7.3",
       "license": "Apache-2.0",
       "dependencies": {
         "@monaco-editor/react": "^4.7.0",

package.json

@@ -1,6 +1,6 @@
 {
   "name": "mujoco-web-platform",
-  "version": "0.7.2",
+  "version": "0.7.3",
   "description": "基于 MuJoCo WebAssembly 的本地机器人仿真与控制平台",
   "private": true,
   "type": "module",

pyproject.toml

@@ -1,7 +1,7 @@
 [tool.ruff]
 target-version = "py312"
 line-length = 100
-extend-exclude = [".venv", "build"]
+extend-exclude = [".venv", "build", "training_server/rl"]
 
 [tool.ruff.lint]
 select = ["E", "F", "I", "UP", "B", "SIM"]

training_server/README.md

@@ -1,35 +1,39 @@
 # 本地强化学习训练服务
 
-该服务把 Web 平台发出的受限训练请求转换为本机 `unitree_rl_mjlab/scripts/train.py` 子进程，并提供状态轮询、停止任务和 `policy.onnx` 下载接口。服务只绑定 `127.0.0.1`，不执行前端传入的任意命令。
+该服务把 Web 平台发出的受限训练请求转换为本机训练子进程，并提供状态轮询、停止任务和 `policy.onnx` 下载接口。服务只绑定 `127.0.0.1`，不执行前端传入的任意命令。
+
+仓库已在 [`rl/`](rl/) 内置 `Unitree-Go2-Flat` 所需的 PPO 训练代码、Go2 模型资产和 ONNX 导出逻辑，不再要求另外克隆 `unitree_rl_mjlab`。`mjlab`、PyTorch 等大型运行依赖仍需安装在本机训练环境中。
+
+## 准备训练环境
+
+使用仓库已有的 `.venv` 项目虚拟环境：
+
+```bash
+source .venv/bin/activate
+python -m pip install -r training_server/rl/requirements.txt
+```
 
 ## 启动
 
-必须使用已经安装 `mjlab`、PyTorch 和训练工程依赖的 Python 解释器：
+使用已安装训练依赖的 Python 解释器启动服务：
 
 ```bash
 npm run training-server -- \
-  --trainer-root /path/to/unitree_rl_mjlab \
+  --trainer-python "$PWD/.venv/bin/python"
-  --trainer-python /path/to/training-env/bin/python
-```
-
-如：
-
-```bash
-npm run training-server -- --trainer-root /home/cen/Embodied_Workspace/unitree_rl_mjlab --trainer-python /home/cen/miniconda3/envs/unitree_rl_mjlab/bin/python
 ```
 
 服务启动时会在终端显示一个随机访问令牌。将该令牌填入前端“访问令牌”字段后再连接。令牌只保存在当前浏览器标签页的 `sessionStorage` 中。自动化启动时可固定令牌：
 
 ```bash
 MUJOCO_TRAINING_TOKEN='至少十六个字符的随机令牌' npm run training-server -- \
-  --trainer-root /path/to/unitree_rl_mjlab
+  --trainer-python /path/to/training-env/bin/python
 ```
 
-也可用 `UNITREE_RL_MJLAB_ROOT` 指定工程目录。默认端口是 `8765`。如果前端不是从 `localhost` 或 `127.0.0.1` 提供，可显式添加来源：
+默认训练工程是仓库内的 `training_server/rl`。如需使用包含其他已注册任务的外部训练工程，仍可通过 `--trainer-root /path/to/trainer` 或 `UNITREE_RL_MJLAB_ROOT` 覆盖。默认端口是 `8765`。如果前端不是从 `localhost` 或 `127.0.0.1` 提供，可显式添加来源：
 
 ```bash
 python training_server/server.py \
-  --trainer-root /path/to/unitree_rl_mjlab \
+  --trainer-python /path/to/training-env/bin/python \
   --allow-origin http://192.168.1.10:5173
 ```
 
@@ -43,7 +47,7 @@ python training_server/server.py \
 - `DELETE /api/training/jobs/{id}`：停止训练；
 - `GET /api/training/jobs/{id}/artifacts/policy.onnx`：下载本次生成的策略。
 
-任务保存在服务内存中，服务重启后历史任务状态会丢失，但 mjlab 日志和 checkpoint 仍保留在训练工程中。
+任务保存在服务内存中，服务重启后历史任务状态会丢失；训练日志、checkpoint 和 ONNX 产物保留在 `training_server/rl/logs/rsl_rl/`，使用外部训练工程时则保留在对应工程中。
 
 ## 测试
 

training_server/rl/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright 2024 The Isaac Lab Project Developers
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

training_server/rl/README.md

@@ -0,0 +1,28 @@
+# 内置 Go2 强化学习训练器
+
+这里包含 `training_server` 默认任务 `Unitree-Go2-Flat` 所需的训练代码与资产：
+
+- RSL-RL/PPO 训练入口；
+- Go2 速度跟踪环境、观测、奖励、终止条件和课程配置；
+- Go2 MJCF 与网格资产；
+- checkpoint 保存及 `policy.onnx` 导出 runner。
+
+这是 [unitreerobotics/unitree_rl_mjlab](https://github.com/unitreerobotics/unitree_rl_mjlab) 的最小内置子集，来源提交为 `1425b15f73bd4095f0df53709d7c389c3eb9e790`，按上游 Apache License 2.0 许可证分发，见 [`LICENSE`](LICENSE)。为避免复制与默认训练任务无关的机器人、真机部署和运动模仿代码，只保留了 Go2 速度任务依赖。
+
+## 环境
+
+从仓库根目录使用已有的 `.venv` 项目虚拟环境：
+
+```bash
+source .venv/bin/activate
+python -m pip install -r training_server/rl/requirements.txt
+```
+
+训练服务会直接从本目录启动 `scripts/train.py`，不需要安装本目录自身。训练产物写入本目录的 `logs/rsl_rl/`（已被仓库忽略）。
+
+可以脱离服务验证任务是否成功注册：
+
+```bash
+cd training_server/rl
+python scripts/train.py Unitree-Go2-Flat --help
+```

training_server/rl/requirements.txt

@@ -0,0 +1,9 @@
+# 与内置 unitree_rl_mjlab 子集验证一致的训练栈。
+# mjlab 1.2.0 会使用 MuJoCo 3.5 的 mjENBL_MULTICCD API，不能放任
+# pip 将 mujoco/warp-lang 升级到较新的不兼容版本。
+mjlab==1.2.0
+mujoco==3.5.0
+mujoco-warp==3.5.0
+warp-lang==1.15.0
+# RSL-RL 5.0.1 仍向 wandb.Settings 传递 start_method，0.29 已删除该字段。
+wandb==0.28.2

training_server/rl/scripts/train.py

@@ -0,0 +1,235 @@
+"""Script to train RL agent with RSL-RL."""
+
+import logging
+import os
+import sys
+from dataclasses import asdict, dataclass, field
+from datetime import datetime
+from pathlib import Path
+from typing import Literal, cast
+
+# 训练器作为仓库内置子集直接从 scripts/ 启动，不要求额外执行 pip install -e。
+TRAINER_ROOT = Path(__file__).resolve().parents[1]
+if str(TRAINER_ROOT) not in sys.path:
+  sys.path.insert(0, str(TRAINER_ROOT))
+
+import tyro
+import warp as wp
+
+# mjlab 1.2.0 的 GPU CUDA-graph 检查仍访问旧公开路径 wp.context；
+# Warp 1.15 已把实现移到 warp._src.context，但保留了相同 runtime 契约。
+if not hasattr(wp, "context"):
+  from warp._src import context as warp_context
+
+  wp.context = warp_context  # type: ignore[attr-defined]
+
+from mjlab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg
+from mjlab.rl import MjlabOnPolicyRunner, RslRlBaseRunnerCfg, RslRlVecEnvWrapper
+from mjlab.tasks.registry import list_tasks, load_env_cfg, load_rl_cfg, load_runner_cls
+from mjlab.tasks.tracking.mdp import MotionCommandCfg
+from mjlab.utils.gpu import select_gpus
+from mjlab.utils.os import dump_yaml, get_checkpoint_path
+from mjlab.utils.torch import configure_torch_backends
+from mjlab.utils.wrappers import VideoRecorder
+
+
+@dataclass(frozen=True)
+class TrainConfig:
+  env: ManagerBasedRlEnvCfg
+  agent: RslRlBaseRunnerCfg
+  motion_file: str | None = None
+  video: bool = False
+  video_length: int = 200
+  video_interval: int = 2000
+  enable_nan_guard: bool = False
+  torchrunx_log_dir: str | None = None
+  gpu_ids: list[int] | Literal["all"] | None = field(default_factory=lambda: [0])
+
+  @staticmethod
+  def from_task(task_id: str) -> "TrainConfig":
+    env_cfg = load_env_cfg(task_id)
+    agent_cfg = load_rl_cfg(task_id)
+    return TrainConfig(env=env_cfg, agent=agent_cfg)
+
+
+def run_train(task_id: str, cfg: TrainConfig, log_dir: Path) -> None:
+  cuda_visible = os.environ.get("CUDA_VISIBLE_DEVICES", "")
+  if cuda_visible == "":
+    device = "cpu"
+    seed = cfg.agent.seed
+    rank = 0
+  else:
+    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    rank = int(os.environ.get("RANK", "0"))
+    # Set EGL device to match the CUDA device.
+    os.environ["MUJOCO_EGL_DEVICE_ID"] = str(local_rank)
+    device = f"cuda:{local_rank}"
+    # Set seed to have diversity in different processes.
+    seed = cfg.agent.seed + local_rank
+
+  configure_torch_backends()
+
+  cfg.agent.seed = seed
+  cfg.env.seed = seed
+
+  print(f"[INFO] Training with: device={device}, seed={seed}, rank={rank}")
+
+  # Check if this is a tracking task by checking for motion command.
+  is_tracking_task = "motion" in cfg.env.commands and isinstance(
+    cfg.env.commands["motion"], MotionCommandCfg
+  )
+
+  if is_tracking_task:
+    if not cfg.motion_file:
+      raise ValueError("For tracking tasks, --motion-file must be set ...")
+    motion_path = Path(cfg.motion_file).expanduser().resolve()
+    if not motion_path.exists():
+      raise FileNotFoundError(f"Motion file not found: {motion_path}")
+    motion_cmd = cfg.env.commands["motion"]
+    assert isinstance(motion_cmd, MotionCommandCfg)
+    motion_cmd.motion_file = str(motion_path)
+    print(f"[INFO] Using motion file: {motion_cmd.motion_file}")
+
+    # Check if motion_file is already set (e.g., via CLI --env.commands.motion.motion-file).
+    if motion_cmd.motion_file and Path(motion_cmd.motion_file).exists():
+      print(f"[INFO] Using local motion file: {motion_cmd.motion_file}")
+
+  # Enable NaN guard if requested.
+  if cfg.enable_nan_guard:
+    cfg.env.sim.nan_guard.enabled = True
+    print(f"[INFO] NaN guard enabled, output dir: {cfg.env.sim.nan_guard.output_dir}")
+
+  if rank == 0:
+    print(f"[INFO] Logging experiment in directory: {log_dir}")
+
+  env = ManagerBasedRlEnv(
+    cfg=cfg.env, device=device, render_mode="rgb_array" if cfg.video else None
+  )
+
+  log_root_path = log_dir.parent  # Go up from specific run dir to experiment dir.
+
+  resume_path: Path | None = None
+  if cfg.agent.resume:
+      # Load checkpoint from local filesystem.
+      resume_path = get_checkpoint_path(
+        log_root_path, cfg.agent.load_run, cfg.agent.load_checkpoint
+      )
+
+  # Only record videos on rank 0 to avoid multiple workers writing to the same files.
+  if cfg.video and rank == 0:
+    env = VideoRecorder(
+      env,
+      video_folder=Path(log_dir) / "videos" / "train",
+      step_trigger=lambda step: step % cfg.video_interval == 0,
+      video_length=cfg.video_length,
+      disable_logger=True,
+    )
+    print("[INFO] Recording videos during training.")
+
+  env = RslRlVecEnvWrapper(env, clip_actions=cfg.agent.clip_actions)
+
+  agent_cfg = asdict(cfg.agent)
+  env_cfg = asdict(cfg.env)
+
+  runner_cls = load_runner_cls(task_id)
+  if runner_cls is None:
+    runner_cls = MjlabOnPolicyRunner
+
+  runner_kwargs = {}
+  runner = runner_cls(env, agent_cfg, str(log_dir), device, **runner_kwargs)
+
+  runner.add_git_repo_to_log(__file__)
+  if resume_path is not None:
+    print(f"[INFO]: Loading model checkpoint from: {resume_path}")
+    runner.load(str(resume_path))
+
+  # Only write config files from rank 0 to avoid race conditions.
+  if rank == 0:
+    dump_yaml(log_dir / "params" / "env.yaml", env_cfg)
+    dump_yaml(log_dir / "params" / "agent.yaml", agent_cfg)
+
+  runner.learn(
+    num_learning_iterations=cfg.agent.max_iterations, init_at_random_ep_len=True
+  )
+
+  env.close()
+
+
+def launch_training(task_id: str, args: TrainConfig | None = None):
+  args = args or TrainConfig.from_task(task_id)
+
+  # Create log directory once before launching workers.
+  log_root_path = Path("logs") / "rsl_rl" / args.agent.experiment_name
+  log_root_path.resolve()
+  log_dir_name = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+  if args.agent.run_name:
+    log_dir_name += f"_{args.agent.run_name}"
+  log_dir = log_root_path / log_dir_name
+
+  # Select GPUs based on CUDA_VISIBLE_DEVICES and user specification.
+  selected_gpus, num_gpus = select_gpus(args.gpu_ids)
+
+  # Set environment variables for all modes.
+  if selected_gpus is None:
+    os.environ["CUDA_VISIBLE_DEVICES"] = ""
+  else:
+    os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(map(str, selected_gpus))
+  os.environ["MUJOCO_GL"] = "egl"
+
+  if num_gpus <= 1:
+    # CPU or single GPU: run directly without torchrunx.
+    run_train(task_id, args, log_dir)
+  else:
+    # Multi-GPU: use torchrunx.
+    import torchrunx
+
+    # torchrunx redirects stdout to logging.
+    logging.basicConfig(level=logging.INFO)
+
+    # Configure torchrunx logging directory.
+    # Priority: 1) existing env var, 2) user flag, 3) default to {log_dir}/torchrunx.
+    if "TORCHRUNX_LOG_DIR" not in os.environ:
+      if args.torchrunx_log_dir is not None:
+        # User specified a value via flag (could be "" to disable).
+        os.environ["TORCHRUNX_LOG_DIR"] = args.torchrunx_log_dir
+      else:
+        # Default: put logs in training directory.
+        os.environ["TORCHRUNX_LOG_DIR"] = str(log_dir / "torchrunx")
+
+    print(f"[INFO] Launching training with {num_gpus} GPUs", flush=True)
+    torchrunx.Launcher(
+      hostnames=["localhost"],
+      workers_per_host=num_gpus,
+      backend=None,  # Let rsl_rl handle process group initialization.
+      copy_env_vars=torchrunx.DEFAULT_ENV_VARS_FOR_COPY + ("MUJOCO*",),
+    ).run(run_train, task_id, args, log_dir)
+
+
+def main():
+  # Parse first argument to choose the task.
+  # Import tasks to populate the registry.
+  import mjlab.tasks  # noqa: F401
+  import src.tasks
+
+  all_tasks = list_tasks()
+  chosen_task, remaining_args = tyro.cli(
+    tyro.extras.literal_type_from_choices(all_tasks),
+    add_help=False,
+    return_unknown_args=True,
+    config=mjlab.TYRO_FLAGS,
+  )
+
+  args = tyro.cli(
+    TrainConfig,
+    args=remaining_args,
+    default=TrainConfig.from_task(chosen_task),
+    prog=sys.argv[0] + f" {chosen_task}",
+    config=mjlab.TYRO_FLAGS,
+  )
+  del remaining_args
+
+  launch_training(task_id=chosen_task, args=args)
+
+
+if __name__ == "__main__":
+  main()

training_server/rl/src/__init__.py

@@ -0,0 +1,4 @@
+from pathlib import Path
+
+
+SRC_PATH: Path = Path(__file__).parent

training_server/rl/src/assets/robots/__init__.py

@@ -0,0 +1,3 @@
+"""仓库内置训练器支持的机器人资产。"""
+
+from .unitree_go2.go2_constants import get_go2_robot_cfg as get_go2_robot_cfg

training_server/rl/src/assets/robots/unitree_go2/__init__.py

@@ -0,0 +1 @@
+"""Unitree Go2 quadruped."""

training_server/rl/src/assets/robots/unitree_go2/go2_constants.py

@@ -0,0 +1,148 @@
+"""Unitree Go2 constants."""
+
+from pathlib import Path
+
+import mujoco
+
+from src import SRC_PATH
+from mjlab.actuator import BuiltinPositionActuatorCfg
+from mjlab.entity import EntityArticulationInfoCfg, EntityCfg
+from mjlab.utils.actuator import ElectricActuator, reflected_inertia
+from mjlab.utils.os import update_assets
+from mjlab.utils.spec_config import CollisionCfg
+
+##
+# MJCF and assets.
+##
+
+GO2_XML: Path = (
+  SRC_PATH / "assets" / "robots" / "unitree_go2" / "xmls" / "go2.xml"
+)
+assert GO2_XML.exists()
+
+
+def get_assets(meshdir: str) -> dict[str, bytes]:
+  assets: dict[str, bytes] = {}
+  update_assets(assets, GO2_XML.parent / "assets", meshdir)
+  return assets
+
+
+def get_spec() -> mujoco.MjSpec:
+  spec = mujoco.MjSpec.from_file(str(GO2_XML))
+  spec.assets = get_assets(spec.meshdir)
+  return spec
+
+
+##
+# Actuator config.
+##
+
+GO2_ACTUATOR_HIP = BuiltinPositionActuatorCfg(
+  target_names_expr=(
+    ".*hip_.*",
+  ),
+  stiffness=20.0,
+  damping=1.0,
+  effort_limit=23.5,
+  armature=0.01,
+)
+GO2_ACTUATOR_THIGH = BuiltinPositionActuatorCfg(
+  target_names_expr=(
+    ".*thigh_.*",
+  ),
+  stiffness=20.0,
+  damping=1.0,
+  effort_limit=23.5,
+  armature=0.01,
+)
+GO2_ACTUATOR_CALF = BuiltinPositionActuatorCfg(
+  target_names_expr=(
+    ".*calf_.*",
+  ),
+  stiffness=40.0,
+  damping=2.0,
+  effort_limit=45,
+  armature=0.02,
+)
+
+##
+# Keyframes.
+##
+
+
+INIT_STATE = EntityCfg.InitialStateCfg(
+  pos=(0.0, 0.0, 0.32),
+  joint_pos={
+    ".*thigh_joint": 0.9,
+    ".*calf_joint": -1.8,
+    ".*R_hip_joint": 0.1,
+    ".*L_hip_joint": -0.1,
+  },
+  joint_vel={".*": 0.0},
+)
+
+##
+# Collision config.
+##
+
+_foot_regex = "^[FR][LR]_foot_collision$"
+
+# This disables all collisions except the feet.
+# Furthermore, feet self collisions are disabled.
+FEET_ONLY_COLLISION = CollisionCfg(
+  geom_names_expr=(_foot_regex,),
+  contype=0,
+  conaffinity=1,
+  condim=3,
+  priority=1,
+  friction=(0.6,),
+  solimp=(0.9, 0.95, 0.023),
+)
+
+# This enables all collisions, excluding self collisions.
+# Foot collisions are given custom condim, friction and solimp.
+FULL_COLLISION = CollisionCfg(
+  geom_names_expr=(".*_collision",),
+  condim={_foot_regex: 3, ".*_collision": 1},
+  priority={_foot_regex: 1},
+  friction={_foot_regex: (0.6,)},
+  solimp={_foot_regex: (0.9, 0.95, 0.023)},
+  contype=1,
+  conaffinity=0,
+)
+
+##
+# Final config.
+##
+
+GO2_ARTICULATION = EntityArticulationInfoCfg(
+  actuators=(
+    GO2_ACTUATOR_HIP,
+    GO2_ACTUATOR_THIGH,
+    GO2_ACTUATOR_CALF,
+  ),
+  soft_joint_pos_limit_factor=0.9,
+)
+
+
+def get_go2_robot_cfg() -> EntityCfg:
+  """Get a fresh Go2 robot configuration instance.
+
+  Returns a new EntityCfg instance each time to avoid mutation issues when
+  the config is shared across multiple places.
+  """
+  return EntityCfg(
+    init_state=INIT_STATE,
+    collisions=(FULL_COLLISION,),
+    spec_fn=get_spec,
+    articulation=GO2_ARTICULATION,
+  )
+
+if __name__ == "__main__":
+  import mujoco.viewer as viewer
+
+  from mjlab.entity.entity import Entity
+
+  robot = Entity(get_go2_robot_cfg())
+
+  viewer.launch(robot.spec.compile())

training_server/rl/src/assets/robots/unitree_go2/xmls/go2.xml

@@ -0,0 +1,161 @@
+<mujoco model="go2">
+  <compiler angle="radian" meshdir="assets" autolimits="true" />
+
+  <default>
+    <default class="go2">
+      <default class="visual">
+        <geom type="mesh" contype="0" conaffinity="0" density="0" group="2"/>
+      </default>
+      <default class="collision">
+        <geom priority="1" condim="6" group="3"/>
+      </default>
+      <site rgba="1 0 0 1" group="5"/>
+    </default>
+  </default>
+
+  <asset>
+    <material name="metal" rgba=".9 .95 .95 1" />
+    <material name="black" rgba="0 0 0 1" />
+    <material name="white" rgba="1 1 1 1" />
+    <material name="gray" rgba="0.671705 0.692426 0.774270 1" />
+
+    <mesh file="base_0.obj" />
+    <mesh file="base_1.obj" />
+    <mesh file="base_2.obj" />
+    <mesh file="base_3.obj" />
+    <mesh file="base_4.obj" />
+    <mesh file="hip_0.obj" />
+    <mesh file="hip_1.obj" />
+    <mesh file="thigh_0.obj" />
+    <mesh file="thigh_1.obj" />
+    <mesh file="thigh_mirror_0.obj" />
+    <mesh file="thigh_mirror_1.obj" />
+    <mesh file="calf_0.obj" />
+    <mesh file="calf_1.obj" />
+    <mesh file="calf_mirror_0.obj" />
+    <mesh file="calf_mirror_1.obj" />
+    <mesh file="foot.obj" />
+  </asset>
+
+  <worldbody>
+    <body name="base_link" pos="0 0 0.445" childclass="go2">
+      <inertial pos="0.021112 0 -0.005366" quat="-0.000543471 0.713435 -0.00173769 0.700719" mass="6.921" diaginertia="0.107027 0.0980771 0.0244531" />
+      <freejoint name="floating_base_joint"/>
+      <geom mesh="base_0" material="black" class="visual" />
+      <geom mesh="base_1" material="black" class="visual" />
+      <geom mesh="base_2" material="black" class="visual" />
+      <geom mesh="base_3" material="white" class="visual" />
+      <geom mesh="base_4" material="gray" class="visual" />
+      <geom name="base1_collision" size="0.1881 0.04675 0.057" type="box" class="collision" />
+      <geom name="base2_collision" size="0.05 0.045" pos="0.285 0 0.01" type="cylinder" class="collision" />
+      <geom name="base3_collision" size="0.047" pos="0.293 0 -0.06" type="sphere" class="collision" />
+      <site name="imu" pos="-0.02557 0 0.04232" group="5"/>
+      <body name="FL_hip" pos="0.1934 0.0465 0">
+        <inertial pos="-0.0054 0.00194 -0.000105" quat="0.497014 0.499245 0.505462 0.498237" mass="0.678" diaginertia="0.00088403 0.000596003 0.000479967" />
+        <joint name="FL_hip_joint" axis="1 0 0" range="-1.0472 1.0472"/>
+        <geom mesh="hip_0" material="metal" class="visual" />
+        <geom mesh="hip_1" material="gray" class="visual" />
+        <geom name="FL_hip_collision" size="0.046 0.02" pos="0 0.08 0" quat="1 1 0 0" type="cylinder" class="collision" />
+        <body name="FL_thigh" pos="0 0.0955 0">
+          <inertial pos="-0.00374 -0.0223 -0.0327" quat="0.829533 0.0847635 -0.0200632 0.551623" mass="1.152" diaginertia="0.00594973 0.00584149 0.000878787" />
+          <joint name="FL_thigh_joint" axis="0 1 0" range="-1.5708 3.4907"/>
+          <geom mesh="thigh_0" material="metal" class="visual" />
+          <geom mesh="thigh_1" material="gray" class="visual" />
+          <geom name="FL_thigh_collision" size="0.1065 0.01225 0.017" pos="0 0 -0.1065" quat="0.707107 0 0.707107 0" type="box" class="collision" />
+          <body name="FL_calf" pos="0 0 -0.213">
+            <inertial pos="0.00629595 -0.000622121 -0.141417" quat="0.710672 0.00154099 -0.00450087 0.703508" mass="0.241352" diaginertia="0.0014901 0.00146356 5.31397e-05" />
+            <joint name="FL_calf_joint" axis="0 1 0" range="-2.7227 -0.83776"/>
+            <geom mesh="calf_0" material="gray" class="visual" />
+            <geom mesh="calf_1" material="black" class="visual" />
+            <geom name="FL_calf1_collision" size="0.012 0.06" pos="0.008 0 -0.06" quat="0.994493 0 -0.104807 0" type="cylinder" class="collision" />
+            <geom name="FL_calf2_collision" size="0.011 0.0325" pos="0.02 0 -0.148" quat="0.999688 0 0.0249974 0" type="cylinder" class="collision" />
+            <geom pos="0 0 -0.213" mesh="foot" class="visual" material="black" />
+            <geom name="FL_foot_collision" pos="0 0 -0.213" type="sphere" size="0.022" class="collision"/>
+            <site name="FL" pos="0 0 -0.213" type="sphere" size="0.022"/>
+          </body>
+        </body>
+      </body>
+      <body name="FR_hip" pos="0.1934 -0.0465 0">
+        <inertial pos="-0.0054 -0.00194 -0.000105" quat="0.498237 0.505462 0.499245 0.497014" mass="0.678" diaginertia="0.00088403 0.000596003 0.000479967" />
+        <joint name="FR_hip_joint" axis="1 0 0" range="-1.0472 1.0472"/>
+        <geom mesh="hip_0" material="metal" class="visual" quat="4.63268e-05 1 0 0" />
+        <geom mesh="hip_1" material="gray" class="visual" quat="4.63268e-05 1 0 0" />
+        <geom name="FR_hip_collision" size="0.046 0.02" pos="0 -0.08 0" quat="0.707107 0.707107 0 0" type="cylinder" class="collision" />
+        <body name="FR_thigh" pos="0 -0.0955 0">
+          <inertial pos="-0.00374 0.0223 -0.0327" quat="0.551623 -0.0200632 0.0847635 0.829533" mass="1.152" diaginertia="0.00594973 0.00584149 0.000878787" />
+          <joint name="FR_thigh_joint" axis="0 1 0" range="-1.5708 3.4907"/>
+          <geom mesh="thigh_mirror_0" material="metal" class="visual" />
+          <geom mesh="thigh_mirror_1" material="gray" class="visual" />
+          <geom name="FR_thigh_collision" size="0.1065 0.01225 0.017" pos="0 0 -0.1065" quat="0.707107 0 0.707107 0" type="box" class="collision" />
+          <body name="FR_calf" pos="0 0 -0.213">
+            <inertial pos="0.00629595 0.000622121 -0.141417" quat="0.703508 -0.00450087 0.00154099 0.710672" mass="0.241352" diaginertia="0.0014901 0.00146356 5.31397e-05" />
+            <joint name="FR_calf_joint" axis="0 1 0" range="-2.7227 -0.83776"/>
+            <geom mesh="calf_mirror_0" material="gray" class="visual" />
+            <geom mesh="calf_mirror_1" material="black" class="visual" />
+            <geom name="FR_calf1_collision" size="0.013 0.06" pos="0.01 0 -0.06" quat="0.995004 0 -0.0998334 0" type="cylinder" class="collision" />
+            <geom name="FR_calf2_collision" size="0.011 0.0325" pos="0.02 0 -0.148" quat="0.999688 0 0.0249974 0" type="cylinder" class="collision" />
+            <geom pos="0 0 -0.213" mesh="foot" class="visual" material="black" />
+            <geom name="FR_foot_collision" pos="0 0 -0.213" type="sphere" size="0.022" class="collision"/>
+            <site name="FR" pos="0 0 -0.213" type="sphere" size="0.022"/>
+          </body>
+        </body>
+      </body>
+      <body name="RL_hip" pos="-0.1934 0.0465 0">
+        <inertial pos="0.0054 0.00194 -0.000105" quat="0.505462 0.498237 0.497014 0.499245" mass="0.678" diaginertia="0.00088403 0.000596003 0.000479967" />
+        <joint name="RL_hip_joint" axis="1 0 0" range="-1.0472 1.0472"/>
+        <geom mesh="hip_0" material="metal" class="visual" quat="4.63268e-05 0 1 0" />
+        <geom mesh="hip_1" material="gray" class="visual" quat="4.63268e-05 0 1 0" />
+        <geom name="RL_hip_collision" size="0.046 0.02" pos="0 0.08 0" quat="0.707107 0.707107 0 0" type="cylinder" class="collision" />
+        <body name="RL_thigh" pos="0 0.0955 0">
+          <inertial pos="-0.00374 -0.0223 -0.0327" quat="0.829533 0.0847635 -0.0200632 0.551623" mass="1.152" diaginertia="0.00594973 0.00584149 0.000878787" />
+          <joint name="RL_thigh_joint" axis="0 1 0" range="-0.5236 4.5379"/>
+          <geom mesh="thigh_0" material="metal" class="visual" />
+          <geom mesh="thigh_1" material="gray" class="visual" />
+          <geom name="RL_thigh_collision" size="0.1065 0.01225 0.017" pos="0 0 -0.1065" quat="0.707107 0 0.707107 0" type="box" class="collision" />
+          <body name="RL_calf" pos="0 0 -0.213">
+            <inertial pos="0.00629595 -0.000622121 -0.141417" quat="0.710672 0.00154099 -0.00450087 0.703508" mass="0.241352" diaginertia="0.0014901 0.00146356 5.31397e-05" />
+            <joint name="RL_calf_joint" axis="0 1 0" range="-2.7227 -0.83776"/>
+            <geom mesh="calf_0" material="gray" class="visual" />
+            <geom mesh="calf_1" material="black" class="visual" />
+            <geom name="RL_calf1_collision" size="0.013 0.06" pos="0.01 0 -0.06" quat="0.995004 0 -0.0998334 0" type="cylinder" class="collision" />
+            <geom name="RL_calf2_collision" size="0.011 0.0325" pos="0.02 0 -0.148" quat="0.999688 0 0.0249974 0" type="cylinder" class="collision" />
+            <geom pos="0 0 -0.213" mesh="foot" class="visual" material="black" />
+            <geom name="RL_foot_collision" pos="0 0 -0.213" type="sphere" size="0.022" class="collision"/>
+            <site name="RL" pos="0 0 -0.213" type="sphere" size="0.022"/>
+          </body>
+        </body>
+      </body>
+      <body name="RR_hip" pos="-0.1934 -0.0465 0">
+        <inertial pos="0.0054 -0.00194 -0.000105" quat="0.499245 0.497014 0.498237 0.505462" mass="0.678" diaginertia="0.00088403 0.000596003 0.000479967" />
+        <joint name="RR_hip_joint" axis="1 0 0" range="-1.0472 1.0472"/>
+        <geom mesh="hip_0" material="metal" class="visual" quat="2.14617e-09 4.63268e-05 4.63268e-05 -1" />
+        <geom mesh="hip_1" material="gray" class="visual" quat="2.14617e-09 4.63268e-05 4.63268e-05 -1" />
+        <geom name="RR_hip_collision" size="0.046 0.02" pos="0 -0.08 0" quat="0.707107 0.707107 0 0" type="cylinder" class="collision" />
+        <body name="RR_thigh" pos="0 -0.0955 0">
+          <inertial pos="-0.00374 0.0223 -0.0327" quat="0.551623 -0.0200632 0.0847635 0.829533" mass="1.152" diaginertia="0.00594973 0.00584149 0.000878787" />
+          <joint name="RR_thigh_joint" axis="0 1 0" range="-0.5236 4.5379"/>
+          <geom mesh="thigh_mirror_0" material="metal" class="visual" />
+          <geom mesh="thigh_mirror_1" material="gray" class="visual" />
+          <geom name="RR_thigh_collision" size="0.1065 0.01225 0.017" pos="0 0 -0.1065" quat="0.707107 0 0.707107 0" type="box" class="collision" />
+          <body name="RR_calf" pos="0 0 -0.213">
+            <inertial pos="0.00629595 0.000622121 -0.141417" quat="0.703508 -0.00450087 0.00154099 0.710672" mass="0.241352" diaginertia="0.0014901 0.00146356 5.31397e-05" />
+            <joint name="RR_calf_joint" axis="0 1 0" range="-2.7227 -0.83776"/>
+            <geom mesh="calf_mirror_0" material="gray" class="visual" />
+            <geom mesh="calf_mirror_1" material="black" class="visual" />
+            <geom name="RR_calf1_collision" size="0.013 0.06" pos="0.01 0 -0.06" quat="0.995004 0 -0.0998334 0" type="cylinder" class="collision" />
+            <geom name="RR_calf2_collision" size="0.011 0.0325" pos="0.02 0 -0.148" quat="0.999688 0 0.0249974 0" type="cylinder" class="collision" />
+            <geom pos="0 0 -0.213" mesh="foot" class="visual" material="black" />
+            <geom name="RR_foot_collision" pos="0 0 -0.213" type="sphere" size="0.022" class="collision"/>
+            <site name="RR" pos="0 0 -0.213" type="sphere" size="0.022"/>
+          </body>
+        </body>
+      </body>
+    </body>
+  </worldbody>
+  <sensor>
+    <gyro name="imu_ang_vel" site="imu"/>
+    <velocimeter name="imu_lin_vel" site="imu"/>
+    <accelerometer name="imu_lin_acc" site="imu"/>
+    <subtreeangmom name="root_angmom" body="base_link"/>
+  </sensor>
+</mujoco>

training_server/rl/src/assets/robots/unitree_go2/xmls/scene_go2.xml

@@ -0,0 +1,306 @@
+<mujoco model="scene_go2">
+  <compiler angle="radian" meshdir="assets" autolimits="true" />
+
+  <option cone="elliptic" impratio="100" />
+
+  <default>
+    <default class="go2">
+      <geom friction="0.4" margin="0.001" condim="1"/>
+      <joint axis="0 1 0" damping="0.1" armature="0.01" frictionloss="0.2"/>
+      <motor ctrlrange="-23.7 23.7"/>
+      <default class="abduction">
+        <joint axis="1 0 0" range="-1.0472 1.0472"/>
+      </default>
+      <default class="hip">
+        <default class="front_hip">
+          <joint range="-1.5708 3.4907"/>
+        </default>
+        <default class="back_hip">
+          <joint range="-0.5236 4.5379"/>
+        </default>
+      </default>
+      <default class="knee">
+        <joint range="-2.7227 -0.83776"/>
+        <motor ctrlrange="-45.43 45.43"/>
+      </default>
+      <default class="visual">
+        <geom type="mesh" contype="0" conaffinity="0" group="2"/>
+      </default>
+      <default class="collision">
+        <geom group="3"/>
+        <default class="foot">
+          <geom size="0.022" pos="-0.002 0 -0.213" priority="1" condim="6"
+            friction="0.4 0.02 0.01"/>
+        </default>
+      </default>
+    </default>
+  </default>
+
+  <asset>
+    <material name="metal" rgba=".9 .95 .95 1" />
+    <material name="black" rgba="0 0 0 1" />
+    <material name="white" rgba="1 1 1 1" />
+    <material name="gray" rgba="0.671705 0.692426 0.774270 1" />
+
+    <mesh file="base_0.obj" />
+    <mesh file="base_1.obj" />
+    <mesh file="base_2.obj" />
+    <mesh file="base_3.obj" />
+    <mesh file="base_4.obj" />
+    <mesh file="hip_0.obj" />
+    <mesh file="hip_1.obj" />
+    <mesh file="thigh_0.obj" />
+    <mesh file="thigh_1.obj" />
+    <mesh file="thigh_mirror_0.obj" />
+    <mesh file="thigh_mirror_1.obj" />
+    <mesh file="calf_0.obj" />
+    <mesh file="calf_1.obj" />
+    <mesh file="calf_mirror_0.obj" />
+    <mesh file="calf_mirror_1.obj" />
+    <mesh file="foot.obj" />
+  </asset>
+
+  <worldbody>
+    <body name="base_link" pos="0 0 0.445" childclass="go2">
+      <inertial pos="0.021112 0 -0.005366" quat="-0.000543471 0.713435 -0.00173769 0.700719"
+        mass="6.921"
+        diaginertia="0.107027 0.0980771 0.0244531" />
+      <freejoint />
+      <geom mesh="base_0" material="black" class="visual" />
+      <geom mesh="base_1" material="black" class="visual" />
+      <geom mesh="base_2" material="black" class="visual" />
+      <geom mesh="base_3" material="white" class="visual" />
+      <geom mesh="base_4" material="gray" class="visual" />
+      <geom size="0.1881 0.04675 0.057" type="box" class="collision" />
+      <geom size="0.05 0.045" pos="0.285 0 0.01" type="cylinder" class="collision" />
+      <geom size="0.047" pos="0.293 0 -0.06" class="collision" />
+      <site name="imu" pos="-0.02557 0 0.04232" />
+      <body name="FL_hip" pos="0.1934 0.0465 0">
+        <inertial pos="-0.0054 0.00194 -0.000105" quat="0.497014 0.499245 0.505462 0.498237"
+          mass="0.678"
+          diaginertia="0.00088403 0.000596003 0.000479967" />
+        <joint name="FL_hip_joint" class="abduction" />
+        <geom mesh="hip_0" material="metal" class="visual" />
+        <geom mesh="hip_1" material="gray" class="visual" />
+        <geom size="0.046 0.02" pos="0 0.08 0" quat="1 1 0 0" type="cylinder" class="collision" />
+        <body name="FL_thigh" pos="0 0.0955 0">
+          <inertial pos="-0.00374 -0.0223 -0.0327" quat="0.829533 0.0847635 -0.0200632 0.551623"
+            mass="1.152"
+            diaginertia="0.00594973 0.00584149 0.000878787" />
+          <joint name="FL_thigh_joint" class="front_hip" />
+          <geom mesh="thigh_0" material="metal" class="visual" />
+          <geom mesh="thigh_1" material="gray" class="visual" />
+          <geom size="0.1065 0.01225 0.017" pos="0 0 -0.1065" quat="0.707107 0 0.707107 0"
+            type="box" class="collision" />
+          <body name="FL_calf" pos="0 0 -0.213">
+            <inertial pos="0.00629595 -0.000622121 -0.141417"
+              quat="0.710672 0.00154099 -0.00450087 0.703508"
+              mass="0.241352" diaginertia="0.0014901 0.00146356 5.31397e-05" />
+            <joint name="FL_calf_joint" class="knee" />
+            <geom mesh="calf_0" material="gray" class="visual" />
+            <geom mesh="calf_1" material="black" class="visual" />
+            <geom size="0.012 0.06" pos="0.008 0 -0.06" quat="0.994493 0 -0.104807 0"
+              type="cylinder" class="collision" />
+            <geom size="0.011 0.0325" pos="0.02 0 -0.148" quat="0.999688 0 0.0249974 0"
+              type="cylinder" class="collision" />
+            <geom pos="0 0 -0.213" mesh="foot" class="visual" material="black" />
+            <geom name="FL" class="foot" />
+            <body name="FL_foot" pos="0 0 -0.213" />
+          </body>
+        </body>
+      </body>
+      <body name="FR_hip" pos="0.1934 -0.0465 0">
+        <inertial pos="-0.0054 -0.00194 -0.000105" quat="0.498237 0.505462 0.499245 0.497014"
+          mass="0.678"
+          diaginertia="0.00088403 0.000596003 0.000479967" />
+        <joint name="FR_hip_joint" class="abduction" />
+        <geom mesh="hip_0" material="metal" class="visual" quat="4.63268e-05 1 0 0" />
+        <geom mesh="hip_1" material="gray" class="visual" quat="4.63268e-05 1 0 0" />
+        <geom size="0.046 0.02" pos="0 -0.08 0" quat="0.707107 0.707107 0 0" type="cylinder"
+          class="collision" />
+        <body name="FR_thigh" pos="0 -0.0955 0">
+          <inertial pos="-0.00374 0.0223 -0.0327" quat="0.551623 -0.0200632 0.0847635 0.829533"
+            mass="1.152"
+            diaginertia="0.00594973 0.00584149 0.000878787" />
+          <joint name="FR_thigh_joint" class="front_hip" />
+          <geom mesh="thigh_mirror_0" material="metal" class="visual" />
+          <geom mesh="thigh_mirror_1" material="gray" class="visual" />
+          <geom size="0.1065 0.01225 0.017" pos="0 0 -0.1065" quat="0.707107 0 0.707107 0"
+            type="box" class="collision" />
+          <body name="FR_calf" pos="0 0 -0.213">
+            <inertial pos="0.00629595 0.000622121 -0.141417"
+              quat="0.703508 -0.00450087 0.00154099 0.710672"
+              mass="0.241352" diaginertia="0.0014901 0.00146356 5.31397e-05" />
+            <joint name="FR_calf_joint" class="knee" />
+            <geom mesh="calf_mirror_0" material="gray" class="visual" />
+            <geom mesh="calf_mirror_1" material="black" class="visual" />
+            <geom size="0.013 0.06" pos="0.01 0 -0.06" quat="0.995004 0 -0.0998334 0"
+              type="cylinder" class="collision" />
+            <geom size="0.011 0.0325" pos="0.02 0 -0.148" quat="0.999688 0 0.0249974 0"
+              type="cylinder" class="collision" />
+            <geom pos="0 0 -0.213" mesh="foot" class="visual" material="black" />
+            <geom name="FR" class="foot" />
+            <body name="FR_foot" pos="0 0 -0.213" />
+          </body>
+        </body>
+      </body>
+      <body name="RL_hip" pos="-0.1934 0.0465 0">
+        <inertial pos="0.0054 0.00194 -0.000105" quat="0.505462 0.498237 0.497014 0.499245"
+          mass="0.678"
+          diaginertia="0.00088403 0.000596003 0.000479967" />
+        <joint name="RL_hip_joint" class="abduction" />
+        <geom mesh="hip_0" material="metal" class="visual" quat="4.63268e-05 0 1 0" />
+        <geom mesh="hip_1" material="gray" class="visual" quat="4.63268e-05 0 1 0" />
+        <geom size="0.046 0.02" pos="0 0.08 0" quat="0.707107 0.707107 0 0" type="cylinder"
+          class="collision" />
+        <body name="RL_thigh" pos="0 0.0955 0">
+          <inertial pos="-0.00374 -0.0223 -0.0327" quat="0.829533 0.0847635 -0.0200632 0.551623"
+            mass="1.152"
+            diaginertia="0.00594973 0.00584149 0.000878787" />
+          <joint name="RL_thigh_joint" class="back_hip" />
+          <geom mesh="thigh_0" material="metal" class="visual" />
+          <geom mesh="thigh_1" material="gray" class="visual" />
+          <geom size="0.1065 0.01225 0.017" pos="0 0 -0.1065" quat="0.707107 0 0.707107 0"
+            type="box" class="collision" />
+          <body name="RL_calf" pos="0 0 -0.213">
+            <inertial pos="0.00629595 -0.000622121 -0.141417"
+              quat="0.710672 0.00154099 -0.00450087 0.703508"
+              mass="0.241352" diaginertia="0.0014901 0.00146356 5.31397e-05" />
+            <joint name="RL_calf_joint" class="knee" />
+            <geom mesh="calf_0" material="gray" class="visual" />
+            <geom mesh="calf_1" material="black" class="visual" />
+            <geom size="0.013 0.06" pos="0.01 0 -0.06" quat="0.995004 0 -0.0998334 0"
+              type="cylinder" class="collision" />
+            <geom size="0.011 0.0325" pos="0.02 0 -0.148" quat="0.999688 0 0.0249974 0"
+              type="cylinder" class="collision" />
+            <geom pos="0 0 -0.213" mesh="foot" class="visual" material="black" />
+            <geom name="RL" class="foot" />
+            <body name="RL_foot" pos="0 0 -0.213" />
+          </body>
+        </body>
+      </body>
+      <body name="RR_hip" pos="-0.1934 -0.0465 0">
+        <inertial pos="0.0054 -0.00194 -0.000105" quat="0.499245 0.497014 0.498237 0.505462"
+          mass="0.678"
+          diaginertia="0.00088403 0.000596003 0.000479967" />
+        <joint name="RR_hip_joint" class="abduction" />
+        <geom mesh="hip_0" material="metal" class="visual"
+          quat="2.14617e-09 4.63268e-05 4.63268e-05 -1" />
+        <geom mesh="hip_1" material="gray" class="visual"
+          quat="2.14617e-09 4.63268e-05 4.63268e-05 -1" />
+        <geom size="0.046 0.02" pos="0 -0.08 0" quat="0.707107 0.707107 0 0" type="cylinder"
+          class="collision" />
+        <body name="RR_thigh" pos="0 -0.0955 0">
+          <inertial pos="-0.00374 0.0223 -0.0327" quat="0.551623 -0.0200632 0.0847635 0.829533"
+            mass="1.152"
+            diaginertia="0.00594973 0.00584149 0.000878787" />
+          <joint name="RR_thigh_joint" class="back_hip" />
+          <geom mesh="thigh_mirror_0" material="metal" class="visual" />
+          <geom mesh="thigh_mirror_1" material="gray" class="visual" />
+          <geom size="0.1065 0.01225 0.017" pos="0 0 -0.1065" quat="0.707107 0 0.707107 0"
+            type="box" class="collision" />
+          <body name="RR_calf" pos="0 0 -0.213">
+            <inertial pos="0.00629595 0.000622121 -0.141417"
+              quat="0.703508 -0.00450087 0.00154099 0.710672"
+              mass="0.241352" diaginertia="0.0014901 0.00146356 5.31397e-05" />
+            <joint name="RR_calf_joint" class="knee" />
+            <geom mesh="calf_mirror_0" material="gray" class="visual" />
+            <geom mesh="calf_mirror_1" material="black" class="visual" />
+            <geom size="0.013 0.06" pos="0.01 0 -0.06" quat="0.995004 0 -0.0998334 0"
+              type="cylinder" class="collision" />
+            <geom size="0.011 0.0325" pos="0.02 0 -0.148" quat="0.999688 0 0.0249974 0"
+              type="cylinder" class="collision" />
+            <geom pos="0 0 -0.213" mesh="foot" class="visual" material="black" />
+            <geom name="RR" class="foot" />
+            <body name="RR_foot" pos="0 0 -0.213" />
+          </body>
+        </body>
+      </body>
+    </body>
+  </worldbody>
+
+  <actuator>
+    <motor class="abduction" name="FR_hip" joint="FR_hip_joint" />
+    <motor class="hip" name="FR_thigh" joint="FR_thigh_joint" />
+    <motor class="knee" name="FR_calf" joint="FR_calf_joint" />
+    <motor class="abduction" name="FL_hip" joint="FL_hip_joint" />
+    <motor class="hip" name="FL_thigh" joint="FL_thigh_joint" />
+    <motor class="knee" name="FL_calf" joint="FL_calf_joint" />
+    <motor class="abduction" name="RR_hip" joint="RR_hip_joint" />
+    <motor class="hip" name="RR_thigh" joint="RR_thigh_joint" />
+    <motor class="knee" name="RR_calf" joint="RR_calf_joint" />
+    <motor class="abduction" name="RL_hip" joint="RL_hip_joint" />
+    <motor class="hip" name="RL_thigh" joint="RL_thigh_joint" />
+    <motor class="knee" name="RL_calf" joint="RL_calf_joint" />
+  </actuator>
+
+  <sensor>
+    <jointpos name="FR_hip_pos" joint="FR_hip_joint" />
+    <jointpos name="FR_thigh_pos" joint="FR_thigh_joint" />
+    <jointpos name="FR_calf_pos" joint="FR_calf_joint" />
+    <jointpos name="FL_hip_pos" joint="FL_hip_joint" />
+    <jointpos name="FL_thigh_pos" joint="FL_thigh_joint" />
+    <jointpos name="FL_calf_pos" joint="FL_calf_joint" />
+    <jointpos name="RR_hip_pos" joint="RR_hip_joint" />
+    <jointpos name="RR_thigh_pos" joint="RR_thigh_joint" />
+    <jointpos name="RR_calf_pos" joint="RR_calf_joint" />
+    <jointpos name="RL_hip_pos" joint="RL_hip_joint" />
+    <jointpos name="RL_thigh_pos" joint="RL_thigh_joint" />
+    <jointpos name="RL_calf_pos" joint="RL_calf_joint" />
+
+    <jointvel name="FR_hip_vel" joint="FR_hip_joint" />
+    <jointvel name="FR_thigh_vel" joint="FR_thigh_joint" />
+    <jointvel name="FR_calf_vel" joint="FR_calf_joint" />
+    <jointvel name="FL_hip_vel" joint="FL_hip_joint" />
+    <jointvel name="FL_thigh_vel" joint="FL_thigh_joint" />
+    <jointvel name="FL_calf_vel" joint="FL_calf_joint" />
+    <jointvel name="RR_hip_vel" joint="RR_hip_joint" />
+    <jointvel name="RR_thigh_vel" joint="RR_thigh_joint" />
+    <jointvel name="RR_calf_vel" joint="RR_calf_joint" />
+    <jointvel name="RL_hip_vel" joint="RL_hip_joint" />
+    <jointvel name="RL_thigh_vel" joint="RL_thigh_joint" />
+    <jointvel name="RL_calf_vel" joint="RL_calf_joint" />
+
+    <jointactuatorfrc name="FR_hip_torque" joint="FR_hip_joint" noise="0.01" />
+    <jointactuatorfrc name="FR_thigh_torque" joint="FR_thigh_joint" noise="0.01" />
+    <jointactuatorfrc name="FR_calf_torque" joint="FR_calf_joint" noise="0.01" />
+    <jointactuatorfrc name="FL_hip_torque" joint="FL_hip_joint" noise="0.01" />
+    <jointactuatorfrc name="FL_thigh_torque" joint="FL_thigh_joint" noise="0.01" />
+    <jointactuatorfrc name="FL_calf_torque" joint="FL_calf_joint" noise="0.01" />
+    <jointactuatorfrc name="RR_hip_torque" joint="RR_hip_joint" noise="0.01" />
+    <jointactuatorfrc name="RR_thigh_torque" joint="RR_thigh_joint" noise="0.01" />
+    <jointactuatorfrc name="RR_calf_torque" joint="RR_calf_joint" noise="0.01" />
+    <jointactuatorfrc name="RL_hip_torque" joint="RL_hip_joint" noise="0.01" />
+    <jointactuatorfrc name="RL_thigh_torque" joint="RL_thigh_joint" noise="0.01" />
+    <jointactuatorfrc name="RL_calf_torque" joint="RL_calf_joint" noise="0.01" />
+
+    <framequat name="imu_quat" objtype="site" objname="imu" />
+    <gyro name="imu_gyro" site="imu" />
+    <accelerometer name="imu_acc" site="imu" />
+
+    <framepos name="frame_pos" objtype="site" objname="imu" />
+    <framelinvel name="frame_vel" objtype="site" objname="imu" />
+  </sensor>
+
+  <keyframe>
+    <key name="home" qpos="0 0 0.27 1 0 0 0 0 0.9 -1.8 0 0.9 -1.8 0 0.9 -1.8 0 0.9 -1.8"
+      ctrl="0 0.9 -1.8 0 0.9 -1.8 0 0.9 -1.8 0 0.9 -1.8" />
+  </keyframe>
+
+  <!-- setup scene -->
+  <statistic center="1.0 1.0 1.0" extent="0.8"/>
+  <visual>
+    <headlight diffuse="0.6 0.6 0.6" ambient="0.1 0.1 0.1" specular="0.9 0.9 0.9"/>
+    <rgba haze="0.15 0.25 0.35 1"/>
+    <global azimuth="-140" elevation="-20"/>
+  </visual>
+  <asset>
+    <texture type="skybox" builtin="flat" rgb1="0 0 0" rgb2="0 0 0" width="512" height="3072"/>
+    <texture type="2d" name="groundplane" builtin="checker" mark="edge" rgb1="0.2 0.3 0.4" rgb2="0.1 0.2 0.3" markrgb="0.8 0.8 0.8" width="300" height="300"/>
+    <material name="groundplane" texture="groundplane" texuniform="true" texrepeat="5 5" reflectance="0.2"/>
+  </asset>
+  <worldbody>
+    <light pos="1 0 3.5" dir="0 0 -1" directional="true"/>
+    <geom name="floor" size="0 0 0.05" type="plane" material="groundplane"/>
+  </worldbody>
+</mujoco>

training_server/rl/src/tasks/__init__.py

@@ -0,0 +1,5 @@
+from mjlab.utils.lab_api.tasks.importer import import_packages
+
+_BLACKLIST_PKGS = ["utils", ".mdp"]
+
+import_packages(__name__, _BLACKLIST_PKGS)

training_server/rl/src/tasks/velocity/__init__.py

@@ -0,0 +1 @@
+"""Velocity tracking environments for legged robots."""

training_server/rl/src/tasks/velocity/config/go2/__init__.py

@@ -0,0 +1,24 @@
+from mjlab.tasks.registry import register_mjlab_task
+from src.tasks.velocity.rl import VelocityOnPolicyRunner
+
+from .env_cfgs import (
+  unitree_go2_flat_env_cfg,
+  unitree_go2_rough_env_cfg,
+)
+from .rl_cfg import unitree_go2_ppo_runner_cfg
+
+register_mjlab_task(
+  task_id="Unitree-Go2-Rough",
+  env_cfg=unitree_go2_rough_env_cfg(),
+  play_env_cfg=unitree_go2_rough_env_cfg(play=True),
+  rl_cfg=unitree_go2_ppo_runner_cfg(),
+  runner_cls=VelocityOnPolicyRunner,
+)
+
+register_mjlab_task(
+  task_id="Unitree-Go2-Flat",
+  env_cfg=unitree_go2_flat_env_cfg(),
+  play_env_cfg=unitree_go2_flat_env_cfg(play=True),
+  rl_cfg=unitree_go2_ppo_runner_cfg(),
+  runner_cls=VelocityOnPolicyRunner,
+)

training_server/rl/src/tasks/velocity/config/go2/env_cfgs.py

@@ -0,0 +1,170 @@
+"""Unitree Go2 velocity environment configurations."""
+
+from typing import Literal
+
+from src.assets.robots import (
+  get_go2_robot_cfg,
+)
+from mjlab.envs import ManagerBasedRlEnvCfg
+from mjlab.envs import mdp as envs_mdp
+from mjlab.envs.mdp.actions import JointPositionActionCfg
+from mjlab.managers import TerminationTermCfg
+from mjlab.managers.event_manager import EventTermCfg
+from mjlab.sensor import ContactMatch, ContactSensorCfg, RayCastSensorCfg
+from mjlab.tasks.velocity import mdp
+from mjlab.tasks.velocity.mdp import UniformVelocityCommandCfg
+
+from src.tasks.velocity.velocity_env_cfg import make_velocity_env_cfg
+
+TerrainType = Literal["rough", "obstacles"]
+
+
+def unitree_go2_rough_env_cfg(
+  play: bool = False,
+) -> ManagerBasedRlEnvCfg:
+  """Create Unitree Go2 rough terrain velocity configuration."""
+  cfg = make_velocity_env_cfg()
+
+  cfg.sim.mujoco.ccd_iterations = 500
+  cfg.sim.contact_sensor_maxmatch = 500
+
+  cfg.scene.entities = {"robot": get_go2_robot_cfg()}
+
+  # Set raycast sensor frame to Go2 base_link.
+  for sensor in cfg.scene.sensors or ():
+    if sensor.name == "terrain_scan":
+      assert isinstance(sensor, RayCastSensorCfg)
+      sensor.frame.name = "base_link"
+
+  foot_names = ("FR", "FL", "RR", "RL")
+  site_names = ("FR", "FL", "RR", "RL")
+  geom_names = tuple(f"{name}_foot_collision" for name in foot_names)
+
+  feet_ground_cfg = ContactSensorCfg(
+    name="feet_ground_contact",
+    primary=ContactMatch(mode="geom", pattern=geom_names, entity="robot"),
+    secondary=ContactMatch(mode="body", pattern="terrain"),
+    fields=("found", "force"),
+    reduce="netforce",
+    num_slots=1,
+    track_air_time=True,
+  )
+  nonfoot_ground_cfg = ContactSensorCfg(
+    name="nonfoot_ground_touch",
+    primary=ContactMatch(
+      mode="geom",
+      entity="robot",
+      # Grab all collision geoms...
+      pattern=r".*_collision\d*$",
+      # Except for the foot geoms.
+      exclude=tuple(geom_names),
+    ),
+    secondary=ContactMatch(mode="body", pattern="terrain"),
+    fields=("found", "force"),
+    reduce="none",
+    num_slots=1,
+    history_length=4,
+  )
+  cfg.scene.sensors = (cfg.scene.sensors or ()) + (
+    feet_ground_cfg,
+    nonfoot_ground_cfg,
+  )
+
+  if cfg.scene.terrain is not None and cfg.scene.terrain.terrain_generator is not None:
+    cfg.scene.terrain.terrain_generator.curriculum = True
+
+  joint_pos_action = cfg.actions["joint_pos"]
+  assert isinstance(joint_pos_action, JointPositionActionCfg)
+
+  cfg.viewer.body_name = "base_link"
+  cfg.viewer.distance = 1.5
+  cfg.viewer.elevation = -10.0
+
+  cfg.observations["critic"].terms["foot_height"].params["asset_cfg"].site_names = site_names
+
+  cfg.events["foot_friction"].params["asset_cfg"].geom_names = geom_names
+  cfg.events["base_com"].params["asset_cfg"].body_names = ("base_link",)
+
+  cfg.rewards["pose"].params["std_standing"] = {
+    r".*(FR|FL|RR|RL)_hip_joint.*": 0.05,
+    r".*(FR|FL|RR|RL)_thigh_joint.*": 0.1,
+    r".*(FR|FL|RR|RL)_calf_joint.*": 0.15,
+  }
+  cfg.rewards["pose"].params["std_walking"] = {
+    r".*(FR|FL|RR|RL)_hip_joint.*": 0.15,
+    r".*(FR|FL|RR|RL)_thigh_joint.*": 0.35,
+    r".*(FR|FL|RR|RL)_calf_joint.*": 0.5,
+  }
+  cfg.rewards["pose"].params["std_running"] = {
+    r".*(FR|FL|RR|RL)_hip_joint.*": 0.15,
+    r".*(FR|FL|RR|RL)_thigh_joint.*": 0.35,
+    r".*(FR|FL|RR|RL)_calf_joint.*": 0.5,
+  }
+
+  cfg.rewards["foot_gait"].params["offset"] = [0.0, 0.5, 0.5, 0.0]
+  cfg.rewards["body_orientation_l2"].params["asset_cfg"].body_names = ("base_link",)
+  cfg.rewards["body_ang_vel"].params["asset_cfg"].body_names = ("base_link",)
+  cfg.rewards["foot_clearance"].params["asset_cfg"].site_names = site_names
+  cfg.rewards["foot_slip"].params["asset_cfg"].site_names = site_names
+
+  cfg.terminations["illegal_contact"] = TerminationTermCfg(
+    func=mdp.illegal_contact,
+    params={"sensor_name": nonfoot_ground_cfg.name, "force_threshold": 10.0},
+  )
+
+  # Apply play mode overrides.
+  if play:
+    # Effectively infinite episode length.
+    cfg.episode_length_s = int(1e9)
+
+    cfg.observations["actor"].enable_corruption = False
+    cfg.events.pop("push_robot", None)
+    cfg.curriculum = {}
+    cfg.events["randomize_terrain"] = EventTermCfg(
+      func=envs_mdp.randomize_terrain,
+      mode="reset",
+      params={},
+    )
+
+    if cfg.scene.terrain is not None:
+      if cfg.scene.terrain.terrain_generator is not None:
+        cfg.scene.terrain.terrain_generator.curriculum = False
+        cfg.scene.terrain.terrain_generator.num_cols = 5
+        cfg.scene.terrain.terrain_generator.num_rows = 5
+        cfg.scene.terrain.terrain_generator.border_width = 10.0
+
+  return cfg
+
+
+def unitree_go2_flat_env_cfg(play: bool = False) -> ManagerBasedRlEnvCfg:
+  """Create Unitree Go2 flat terrain velocity configuration."""
+  cfg = unitree_go2_rough_env_cfg(play=play)
+
+  cfg.sim.njmax = 300
+  cfg.sim.mujoco.ccd_iterations = 50
+  cfg.sim.contact_sensor_maxmatch = 64
+  cfg.sim.nconmax = None
+
+  # Switch to flat terrain.
+  assert cfg.scene.terrain is not None
+  cfg.scene.terrain.terrain_type = "plane"
+  cfg.scene.terrain.terrain_generator = None
+
+  # Remove raycast sensor and height scan (no terrain to scan).
+  cfg.scene.sensors = tuple(
+    s for s in (cfg.scene.sensors or ()) if s.name != "terrain_scan"
+  )
+  del cfg.observations["actor"].terms["height_scan"]
+  del cfg.observations["critic"].terms["height_scan"]
+
+  # Disable terrain curriculum (not present in play mode since rough clears all).
+  cfg.curriculum.pop("terrain_levels", None)
+
+  if play:
+    twist_cmd = cfg.commands["twist"]
+    assert isinstance(twist_cmd, UniformVelocityCommandCfg)
+    twist_cmd.ranges.lin_vel_x = (-0.5, 1.0)
+    twist_cmd.ranges.lin_vel_y = (-0.5, 0.5)
+    twist_cmd.ranges.ang_vel_z = (-0.5, 0.5)
+
+  return cfg

training_server/rl/src/tasks/velocity/config/go2/rl_cfg.py

@@ -0,0 +1,46 @@
+"""RL configuration for Unitree Go2 velocity task."""
+
+from mjlab.rl import (
+  RslRlModelCfg,
+  RslRlOnPolicyRunnerCfg,
+  RslRlPpoAlgorithmCfg,
+)
+
+
+def unitree_go2_ppo_runner_cfg() -> RslRlOnPolicyRunnerCfg:
+  """Create RL runner configuration for Unitree Go2 velocity task."""
+  return RslRlOnPolicyRunnerCfg(
+    actor=RslRlModelCfg(
+      hidden_dims=(512, 256, 128),
+      activation="elu",
+      obs_normalization=True,
+      distribution_cfg={
+        "class_name": "GaussianDistribution",
+        "init_std": 1.0,
+        "std_type": "scalar",
+      },
+    ),
+    critic=RslRlModelCfg(
+      hidden_dims=(512, 256, 128),
+      activation="elu",
+      obs_normalization=True,
+    ),
+    algorithm=RslRlPpoAlgorithmCfg(
+      value_loss_coef=1.0,
+      use_clipped_value_loss=True,
+      clip_param=0.2,
+      entropy_coef=0.01,
+      num_learning_epochs=5,
+      num_mini_batches=4,
+      learning_rate=1.0e-3,
+      schedule="adaptive",
+      gamma=0.99,
+      lam=0.95,
+      desired_kl=0.01,
+      max_grad_norm=1.0,
+    ),
+    experiment_name="go2_velocity",
+    save_interval=100,
+    num_steps_per_env=24,
+    max_iterations=10001,
+  )

training_server/rl/src/tasks/velocity/mdp/__init__.py

@@ -0,0 +1,7 @@
+from mjlab.envs.mdp import *  # noqa: F401, F403
+
+from .curriculums import *  # noqa: F403
+from .observations import *  # noqa: F403
+from .rewards import *  # noqa: F403
+from .terminations import *  # noqa: F403
+from .velocity_command import *  # noqa: F403

training_server/rl/src/tasks/velocity/mdp/curriculums.py

@@ -0,0 +1,107 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, TypedDict, cast
+
+import torch
+
+from mjlab.entity import Entity
+from mjlab.managers.scene_entity_config import SceneEntityCfg
+
+from .velocity_command import UniformVelocityCommandCfg
+
+if TYPE_CHECKING:
+  from mjlab.envs import ManagerBasedRlEnv
+
+_DEFAULT_SCENE_CFG = SceneEntityCfg("robot")
+
+
+class VelocityStage(TypedDict):
+  step: int
+  lin_vel_x: tuple[float, float] | None
+  lin_vel_y: tuple[float, float] | None
+  ang_vel_z: tuple[float, float] | None
+
+
+class RewardWeightStage(TypedDict):
+  step: int
+  weight: float
+
+
+def terrain_levels_vel(
+  env: ManagerBasedRlEnv,
+  env_ids: torch.Tensor,
+  command_name: str,
+  asset_cfg: SceneEntityCfg = _DEFAULT_SCENE_CFG,
+) -> torch.Tensor:
+  asset: Entity = env.scene[asset_cfg.name]
+
+  terrain = env.scene.terrain
+  assert terrain is not None
+  terrain_generator = terrain.cfg.terrain_generator
+  assert terrain_generator is not None
+
+  command = env.command_manager.get_command(command_name)
+  assert command is not None
+
+  # Compute the distance the robot walked.
+  distance = torch.norm(
+    asset.data.root_link_pos_w[env_ids, :2] - env.scene.env_origins[env_ids, :2], dim=1
+  )
+
+  # Robots that walked far enough progress to harder terrains.
+  move_up = distance > terrain_generator.size[0] / 2
+
+  # Robots that walked less than half of their required distance go to simpler
+  # terrains.
+  move_down = (
+    distance < torch.norm(command[env_ids, :2], dim=1) * env.max_episode_length_s * 0.5
+  )
+  move_down *= ~move_up
+
+  # Update terrain levels.
+  terrain.update_env_origins(env_ids, move_up, move_down)
+
+  return torch.mean(terrain.terrain_levels.float())
+
+
+def commands_vel(
+  env: ManagerBasedRlEnv,
+  env_ids: torch.Tensor,
+  command_name: str,
+  velocity_stages: list[VelocityStage],
+) -> dict[str, torch.Tensor]:
+  del env_ids  # Unused.
+  command_term = env.command_manager.get_term(command_name)
+  assert command_term is not None
+  cfg = cast(UniformVelocityCommandCfg, command_term.cfg)
+  for stage in velocity_stages:
+    if env.common_step_counter > stage["step"]:
+      if "lin_vel_x" in stage and stage["lin_vel_x"] is not None:
+        cfg.ranges.lin_vel_x = stage["lin_vel_x"]
+      if "lin_vel_y" in stage and stage["lin_vel_y"] is not None:
+        cfg.ranges.lin_vel_y = stage["lin_vel_y"]
+      if "ang_vel_z" in stage and stage["ang_vel_z"] is not None:
+        cfg.ranges.ang_vel_z = stage["ang_vel_z"]
+  return {
+    # "lin_vel_x_min": torch.tensor(cfg.ranges.lin_vel_x[0]),
+    # "lin_vel_x_max": torch.tensor(cfg.ranges.lin_vel_x[1]),
+    # "lin_vel_y_min": torch.tensor(cfg.ranges.lin_vel_y[0]),
+    # "lin_vel_y_max": torch.tensor(cfg.ranges.lin_vel_y[1]),
+    # "ang_vel_z_min": torch.tensor(cfg.ranges.ang_vel_z[0]),
+    # "ang_vel_z_max": torch.tensor(cfg.ranges.ang_vel_z[1]),
+  }
+
+
+def reward_weight(
+  env: ManagerBasedRlEnv,
+  env_ids: torch.Tensor,
+  reward_name: str,
+  weight_stages: list[RewardWeightStage],
+) -> torch.Tensor:
+  """Update a reward term's weight based on training step stages."""
+  del env_ids  # Unused.
+  reward_term_cfg = env.reward_manager.get_term_cfg(reward_name)
+  for stage in weight_stages:
+    if env.common_step_counter > stage["step"]:
+      reward_term_cfg.weight = stage["weight"]
+  return torch.tensor([reward_term_cfg.weight])

training_server/rl/src/tasks/velocity/mdp/observations.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from mjlab.entity import Entity
+from mjlab.managers.scene_entity_config import SceneEntityCfg
+from mjlab.sensor import ContactSensor
+
+if TYPE_CHECKING:
+  from mjlab.envs import ManagerBasedRlEnv
+
+_DEFAULT_ASSET_CFG = SceneEntityCfg("robot")
+
+
+def foot_height(
+  env: ManagerBasedRlEnv, asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG
+) -> torch.Tensor:
+  asset: Entity = env.scene[asset_cfg.name]
+  return asset.data.site_pos_w[:, asset_cfg.site_ids, 2]  # (num_envs, num_sites)
+
+
+def foot_air_time(env: ManagerBasedRlEnv, sensor_name: str) -> torch.Tensor:
+  sensor: ContactSensor = env.scene[sensor_name]
+  sensor_data = sensor.data
+  current_air_time = sensor_data.current_air_time
+  assert current_air_time is not None
+  return current_air_time
+
+
+def foot_contact(env: ManagerBasedRlEnv, sensor_name: str) -> torch.Tensor:
+  sensor: ContactSensor = env.scene[sensor_name]
+  sensor_data = sensor.data
+  assert sensor_data.found is not None
+  return (sensor_data.found > 0).float()
+
+
+def foot_contact_forces(env: ManagerBasedRlEnv, sensor_name: str) -> torch.Tensor:
+  sensor: ContactSensor = env.scene[sensor_name]
+  sensor_data = sensor.data
+  assert sensor_data.force is not None
+  forces_flat = sensor_data.force.flatten(start_dim=1)  # [B, N*3]
+  return torch.sign(forces_flat) * torch.log1p(torch.abs(forces_flat))
+
+
+def phase(env: ManagerBasedRlEnv, period: float, command_name: str) -> torch.Tensor:
+    global_phase = (env.episode_length_buf * env.step_dt) % period / period
+    phase = torch.zeros(env.num_envs, 2, device=env.device)
+    phase[:, 0] = torch.sin(global_phase * torch.pi * 2.0)
+    phase[:, 1] = torch.cos(global_phase * torch.pi * 2.0)
+    stand_mask = torch.linalg.norm(env.command_manager.get_command(command_name), dim=1) < 0.1
+    phase = torch.where(stand_mask.unsqueeze(1), torch.zeros_like(phase), phase)
+    return phase

training_server/rl/src/tasks/velocity/mdp/rewards.py

@@ -0,0 +1,427 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from mjlab.entity import Entity
+from mjlab.managers.reward_manager import RewardTermCfg
+from mjlab.managers.scene_entity_config import SceneEntityCfg
+from mjlab.sensor import BuiltinSensor, ContactSensor
+from mjlab.utils.lab_api.math import quat_apply_inverse
+from mjlab.utils.lab_api.string import (
+  resolve_matching_names_values,
+)
+
+if TYPE_CHECKING:
+  from mjlab.envs import ManagerBasedRlEnv
+
+
+_DEFAULT_ASSET_CFG = SceneEntityCfg("robot")
+
+
+def track_linear_velocity(
+  env: ManagerBasedRlEnv,
+  std: float,
+  command_name: str,
+  asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG,
+) -> torch.Tensor:
+  """Reward for tracking the commanded base linear velocity.
+
+  The commanded z velocity is assumed to be zero.
+  """
+  asset: Entity = env.scene[asset_cfg.name]
+  command = env.command_manager.get_command(command_name)
+  assert command is not None, f"Command '{command_name}' not found."
+  actual = asset.data.root_link_lin_vel_b
+  xy_error = torch.sum(torch.square(command[:, :2] - actual[:, :2]), dim=1)
+  z_error = torch.square(actual[:, 2])
+  lin_vel_error = xy_error + (2 * z_error)
+  return torch.exp(-lin_vel_error / std**2)
+
+
+def track_angular_velocity(
+  env: ManagerBasedRlEnv,
+  std: float,
+  command_name: str,
+  asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG,
+) -> torch.Tensor:
+  """Reward heading error for heading-controlled envs, angular velocity for others.
+
+  The commanded xy angular velocities are assumed to be zero.
+  """
+  asset: Entity = env.scene[asset_cfg.name]
+  command = env.command_manager.get_command(command_name)
+  assert command is not None, f"Command '{command_name}' not found."
+  actual = asset.data.root_link_ang_vel_b
+  z_error = torch.square(command[:, 2] - actual[:, 2])
+  xy_error = torch.sum(torch.square(actual[:, :2]), dim=1)
+  ang_vel_error = z_error + (0.05 * xy_error)
+  return torch.exp(-ang_vel_error / std**2)
+
+
+def body_orientation_l2(
+  env: ManagerBasedRlEnv,
+  asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG,
+) -> torch.Tensor:
+  """Reward flat base orientation (robot being upright).
+
+  If asset_cfg has body_ids specified, computes the projected gravity
+  for that specific body. Otherwise, uses the root link projected gravity.
+  """
+  asset: Entity = env.scene[asset_cfg.name]
+
+  # If body_ids are specified, compute projected gravity for that body.
+  if asset_cfg.body_ids:
+    body_quat_w = asset.data.body_link_quat_w[:, asset_cfg.body_ids, :]  # [B, N, 4]
+    body_quat_w = body_quat_w.squeeze(1)  # [B, 4]
+    gravity_w = asset.data.gravity_vec_w  # [3]
+    projected_gravity_b = quat_apply_inverse(body_quat_w, gravity_w)  # [B, 3]
+    xy_squared = torch.sum(torch.square(projected_gravity_b[:, :2]), dim=1)
+  else:
+    # Use root link projected gravity.
+    xy_squared = torch.sum(torch.square(asset.data.projected_gravity_b[:, :2]), dim=1)
+  return xy_squared
+
+
+def self_collision_cost(
+  env: ManagerBasedRlEnv,
+  sensor_name: str,
+  force_threshold: float = 10.0,
+) -> torch.Tensor:
+  """Penalize self-collisions.
+
+  When the sensor provides force history (from ``history_length > 0``),
+  counts substeps where any contact force exceeds *force_threshold*.
+  Falls back to the instantaneous ``found`` count otherwise.
+  """
+  sensor: ContactSensor = env.scene[sensor_name]
+  data = sensor.data
+  if data.force_history is not None:
+    # force_history: [B, N, H, 3]
+    force_mag = torch.norm(data.force_history, dim=-1)  # [B, N, H]
+    hit = (force_mag > force_threshold).any(dim=1)  # [B, H]
+    return hit.sum(dim=-1).float()  # [B]
+  assert data.found is not None
+  return data.found.squeeze(-1)
+
+
+def body_angular_velocity_penalty(
+  env: ManagerBasedRlEnv,
+  asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG,
+) -> torch.Tensor:
+  """Penalize excessive body angular velocities."""
+  asset: Entity = env.scene[asset_cfg.name]
+  ang_vel = asset.data.body_link_ang_vel_w[:, asset_cfg.body_ids, :]
+  ang_vel = ang_vel.squeeze(1)
+  ang_vel_xy = ang_vel[:, :2]  # Don't penalize z-angular velocity.
+  return torch.sum(torch.square(ang_vel_xy), dim=1)
+
+
+def angular_momentum_penalty(
+  env: ManagerBasedRlEnv,
+  sensor_name: str,
+) -> torch.Tensor:
+  """Penalize whole-body angular momentum to encourage natural arm swing."""
+  angmom_sensor: BuiltinSensor = env.scene[sensor_name]
+  angmom = angmom_sensor.data
+  angmom_magnitude_sq = torch.sum(torch.square(angmom), dim=-1)
+  angmom_magnitude = torch.sqrt(angmom_magnitude_sq)
+  env.extras["log"]["Metrics/angular_momentum_mean"] = torch.mean(angmom_magnitude)
+  return angmom_magnitude_sq
+
+
+def feet_air_time(
+  env: ManagerBasedRlEnv,
+  sensor_name: str,
+  threshold: float = 0.4,
+  command_name: str | None = None,
+  command_threshold: float = 0.1,
+) -> torch.Tensor:
+  """Reward feet air time."""
+  sensor: ContactSensor = env.scene[sensor_name]
+  sensor_data = sensor.data
+  air_time = sensor_data.current_air_time
+  contact_time = sensor_data.current_contact_time
+  in_contact = contact_time > 0.0
+  in_mode_time = torch.where(in_contact, contact_time, air_time)
+  single_stance = torch.mean(in_contact.float(), dim=1) == 0.5
+  mode_time = torch.min(torch.where(single_stance.unsqueeze(-1), in_mode_time, 0.0), dim=1)[0]
+  error = torch.abs(mode_time - threshold)
+  reward = torch.clamp(threshold - error, min=0.0)
+  if command_name is not None:
+    command = env.command_manager.get_command(command_name)
+    if command is not None:
+      linear_norm = torch.norm(command[:, :2], dim=1)
+      angular_norm = torch.abs(command[:, 2])
+      total_command = linear_norm + angular_norm
+      scale = (total_command > command_threshold).float()
+      reward *= scale
+  return reward
+
+
+def feet_clearance(
+  env: ManagerBasedRlEnv,
+  target_height: float,
+  command_name: str | None = None,
+  command_threshold: float = 0.1,
+  asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG,
+) -> torch.Tensor:
+  """Penalize deviation from target clearance height, weighted by foot velocity."""
+  asset: Entity = env.scene[asset_cfg.name]
+  foot_z = asset.data.site_pos_w[:, asset_cfg.site_ids, 2]  # [B, N]
+  foot_vel_xy = asset.data.site_lin_vel_w[:, asset_cfg.site_ids, :2]  # [B, N, 2]
+  vel_norm = torch.norm(foot_vel_xy, dim=-1)  # [B, N]
+  delta = torch.abs(foot_z - target_height)  # [B, N]
+  cost = torch.sum(delta * vel_norm, dim=1)  # [B]
+  if command_name is not None:
+    command = env.command_manager.get_command(command_name)
+    if command is not None:
+      linear_norm = torch.norm(command[:, :2], dim=1)
+      angular_norm = torch.abs(command[:, 2])
+      total_command = linear_norm + angular_norm
+      active = (total_command > command_threshold).float()
+      cost = cost * active
+  return cost
+
+
+def feet_gait(
+        env: ManagerBasedRlEnv,
+        period: float,
+        offset: list[float],
+        threshold: float,
+        command_threshold: float,
+        command_name: str,
+        sensor_name: str,
+) -> torch.Tensor:
+    sensor: ContactSensor = env.scene[sensor_name]
+    is_contact = sensor.data.current_contact_time > 0
+    global_phase = ((env.episode_length_buf * env.step_dt) / period).unsqueeze(1)
+    offsets = torch.as_tensor(offset, device=env.device, dtype=global_phase.dtype).view(1, -1)
+    leg_phase = (global_phase + offsets) % 1.0
+    is_stance = (leg_phase < threshold)
+    reward = (is_stance == is_contact).float().mean(dim=1)
+    if command_name is not None:
+        command = env.command_manager.get_command(command_name)
+        if command is not None:
+            linear_norm = torch.norm(command[:, :2], dim=1)
+            angular_norm = torch.abs(command[:, 2])
+            total_command = linear_norm + angular_norm
+            scale = (total_command > command_threshold).float()
+            reward *= scale
+    return reward
+
+
+class feet_swing_height:
+  """Penalize deviation from target swing height, evaluated at landing."""
+
+  def __init__(self, cfg: RewardTermCfg, env: ManagerBasedRlEnv):
+    self.sensor_name = cfg.params["sensor_name"]
+    self.site_names = cfg.params["asset_cfg"].site_names
+    self.peak_heights = torch.zeros(
+      (env.num_envs, len(self.site_names)), device=env.device, dtype=torch.float32
+    )
+    self.step_dt = env.step_dt
+
+  def __call__(
+    self,
+    env: ManagerBasedRlEnv,
+    sensor_name: str,
+    target_height: float,
+    command_name: str,
+    command_threshold: float,
+    asset_cfg: SceneEntityCfg,
+  ) -> torch.Tensor:
+    asset: Entity = env.scene[asset_cfg.name]
+    contact_sensor: ContactSensor = env.scene[sensor_name]
+    command = env.command_manager.get_command(command_name)
+    assert command is not None
+    foot_heights = asset.data.site_pos_w[:, asset_cfg.site_ids, 2]
+    in_air = contact_sensor.data.found == 0
+    self.peak_heights = torch.where(
+      in_air,
+      torch.maximum(self.peak_heights, foot_heights),
+      self.peak_heights,
+    )
+    first_contact = contact_sensor.compute_first_contact(dt=self.step_dt)
+    linear_norm = torch.norm(command[:, :2], dim=1)
+    angular_norm = torch.abs(command[:, 2])
+    total_command = linear_norm + angular_norm
+    active = (total_command > command_threshold).float()
+    error = self.peak_heights / target_height - 1.0
+    cost = torch.sum(torch.square(error) * first_contact.float(), dim=1) * active
+    num_landings = torch.sum(first_contact.float())
+    peak_heights_at_landing = self.peak_heights * first_contact.float()
+    mean_peak_height = torch.sum(peak_heights_at_landing) / torch.clamp(
+      num_landings, min=1
+    )
+    env.extras["log"]["Metrics/peak_height_mean"] = mean_peak_height
+    self.peak_heights = torch.where(
+      first_contact,
+      torch.zeros_like(self.peak_heights),
+      self.peak_heights,
+    )
+    return cost
+
+
+def feet_slip(
+  env: ManagerBasedRlEnv,
+  sensor_name: str,
+  command_name: str,
+  command_threshold: float = 0.01,
+  asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG,
+) -> torch.Tensor:
+  """Penalize foot sliding (xy velocity while in contact)."""
+  asset: Entity = env.scene[asset_cfg.name]
+  contact_sensor: ContactSensor = env.scene[sensor_name]
+  command = env.command_manager.get_command(command_name)
+  assert command is not None
+  linear_norm = torch.norm(command[:, :2], dim=1)
+  angular_norm = torch.abs(command[:, 2])
+  total_command = linear_norm + angular_norm
+  active = (total_command > command_threshold).float()
+  assert contact_sensor.data.found is not None
+  in_contact = (contact_sensor.data.found > 0).float()  # [B, N]
+  foot_vel_xy = asset.data.site_lin_vel_w[:, asset_cfg.site_ids, :2]  # [B, N, 2]
+  vel_xy_norm = torch.norm(foot_vel_xy, dim=-1)  # [B, N]
+  vel_xy_norm_sq = torch.square(vel_xy_norm)  # [B, N]
+  cost = torch.sum(vel_xy_norm_sq * in_contact, dim=1) * active
+  num_in_contact = torch.sum(in_contact)
+  mean_slip_vel = torch.sum(vel_xy_norm * in_contact) / torch.clamp(
+    num_in_contact, min=1
+  )
+  env.extras["log"]["Metrics/slip_velocity_mean"] = mean_slip_vel
+  return cost
+
+
+def soft_landing(
+  env: ManagerBasedRlEnv,
+  sensor_name: str,
+  command_name: str | None = None,
+  command_threshold: float = 0.05,
+) -> torch.Tensor:
+  """Penalize high impact forces at landing to encourage soft footfalls."""
+  contact_sensor: ContactSensor = env.scene[sensor_name]
+  sensor_data = contact_sensor.data
+  assert sensor_data.force is not None
+  forces = sensor_data.force  # [B, N, 3]
+  force_magnitude = torch.norm(forces, dim=-1)  # [B, N]
+  first_contact = contact_sensor.compute_first_contact(dt=env.step_dt)  # [B, N]
+  landing_impact = force_magnitude * first_contact.float()  # [B, N]
+  cost = torch.sum(landing_impact, dim=1)  # [B]
+  num_landings = torch.sum(first_contact.float())
+  mean_landing_force = torch.sum(landing_impact) / torch.clamp(num_landings, min=1)
+  env.extras["log"]["Metrics/landing_force_mean"] = mean_landing_force
+  if command_name is not None:
+    command = env.command_manager.get_command(command_name)
+    if command is not None:
+      linear_norm = torch.norm(command[:, :2], dim=1)
+      angular_norm = torch.abs(command[:, 2])
+      total_command = linear_norm + angular_norm
+      active = (total_command > command_threshold).float()
+      cost = cost * active
+  return cost
+
+
+class variable_posture:
+  """Penalize deviation from default pose with speed-dependent tolerance.
+
+  Uses per-joint standard deviations to control how much each joint can deviate
+  from default pose. Smaller std = stricter (less deviation allowed), larger
+  std = more forgiving. The reward is: exp(-mean(error² / std²))
+
+  Three speed regimes (based on linear + angular command velocity):
+    - std_standing (speed < walking_threshold): Tight tolerance for holding pose.
+    - std_walking (walking_threshold <= speed < running_threshold): Moderate.
+    - std_running (speed >= running_threshold): Loose tolerance for large motion.
+
+  Tune std values per joint based on how much motion that joint needs at each
+  speed. Map joint name patterns to std values, e.g. {".*knee.*": 0.35}.
+  """
+
+  def __init__(self, cfg: RewardTermCfg, env: ManagerBasedRlEnv):
+    asset: Entity = env.scene[cfg.params["asset_cfg"].name]
+    default_joint_pos = asset.data.default_joint_pos
+    assert default_joint_pos is not None
+    self.default_joint_pos = default_joint_pos
+
+    _, joint_names = asset.find_joints(cfg.params["asset_cfg"].joint_names)
+
+    _, _, std_standing = resolve_matching_names_values(
+      data=cfg.params["std_standing"],
+      list_of_strings=joint_names,
+    )
+    self.std_standing = torch.tensor(
+      std_standing, device=env.device, dtype=torch.float32
+    )
+
+    _, _, std_walking = resolve_matching_names_values(
+      data=cfg.params["std_walking"],
+      list_of_strings=joint_names,
+    )
+    self.std_walking = torch.tensor(std_walking, device=env.device, dtype=torch.float32)
+
+    _, _, std_running = resolve_matching_names_values(
+      data=cfg.params["std_running"],
+      list_of_strings=joint_names,
+    )
+    self.std_running = torch.tensor(std_running, device=env.device, dtype=torch.float32)
+
+  def __call__(
+    self,
+    env: ManagerBasedRlEnv,
+    std_standing,
+    std_walking,
+    std_running,
+    asset_cfg: SceneEntityCfg,
+    command_name: str,
+    walking_threshold: float = 0.5,
+    running_threshold: float = 1.5,
+  ) -> torch.Tensor:
+    del std_standing, std_walking, std_running  # Unused.
+
+    asset: Entity = env.scene[asset_cfg.name]
+    command = env.command_manager.get_command(command_name)
+    assert command is not None
+
+    linear_speed = torch.norm(command[:, :2], dim=1)
+    angular_speed = torch.abs(command[:, 2])
+    total_speed = linear_speed + angular_speed
+
+    standing_mask = (total_speed < walking_threshold).float()
+    walking_mask = (
+      (total_speed >= walking_threshold) & (total_speed < running_threshold)
+    ).float()
+    running_mask = (total_speed >= running_threshold).float()
+
+    std = (
+      self.std_standing * standing_mask.unsqueeze(1)
+      + self.std_walking * walking_mask.unsqueeze(1)
+      + self.std_running * running_mask.unsqueeze(1)
+    )
+
+    current_joint_pos = asset.data.joint_pos[:, asset_cfg.joint_ids]
+    desired_joint_pos = self.default_joint_pos[:, asset_cfg.joint_ids]
+    error_squared = torch.square(current_joint_pos - desired_joint_pos)
+
+    return torch.exp(-torch.mean(error_squared / (std**2), dim=1))
+
+
+def stand_still(
+        env: ManagerBasedRlEnv,
+        command_name: str,
+        command_threshold: float = 0.1,
+        asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG
+) -> torch.Tensor:
+    asset: Entity = env.scene[asset_cfg.name]
+    diff_angle = asset.data.joint_pos[:, asset_cfg.joint_ids] - asset.data.default_joint_pos[:, asset_cfg.joint_ids]
+    reward = torch.sum(torch.square(diff_angle), dim=1)
+    if command_name is not None:
+        command = env.command_manager.get_command(command_name)
+        if command is not None:
+            linear_norm = torch.norm(command[:, :2], dim=1)
+            angular_norm = torch.abs(command[:, 2])
+            total_command = linear_norm + angular_norm
+            scale = (total_command <= command_threshold).float()
+            reward *= scale
+    return reward

training_server/rl/src/tasks/velocity/mdp/terminations.py

@@ -0,0 +1,25 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from mjlab.sensor import ContactSensor
+
+if TYPE_CHECKING:
+  from mjlab.envs import ManagerBasedRlEnv
+
+
+def illegal_contact(
+  env: ManagerBasedRlEnv,
+  sensor_name: str,
+  force_threshold: float = 10.0,
+) -> torch.Tensor:
+  sensor: ContactSensor = env.scene[sensor_name]
+  data = sensor.data
+  if data.force_history is not None:
+    # force_history: [B, N, H, 3]
+    force_mag = torch.norm(data.force_history, dim=-1)  # [B, N, H]
+    return (force_mag > force_threshold).any(dim=-1).any(dim=-1)  # [B]
+  assert data.found is not None
+  return torch.any(data.found, dim=-1)

training_server/rl/src/tasks/velocity/mdp/velocity_command.py

@@ -0,0 +1,282 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+
+import numpy as np
+import torch
+
+from mjlab.entity import Entity
+from mjlab.managers.command_manager import CommandTerm, CommandTermCfg
+from mjlab.utils.lab_api.math import (
+  matrix_from_quat,
+  quat_apply,
+  wrap_to_pi,
+)
+
+if TYPE_CHECKING:
+  import viser
+
+  from mjlab.envs.manager_based_rl_env import ManagerBasedRlEnv
+  from mjlab.viewer.debug_visualizer import DebugVisualizer
+
+
+class UniformVelocityCommand(CommandTerm):
+  cfg: UniformVelocityCommandCfg
+
+  def __init__(self, cfg: UniformVelocityCommandCfg, env: ManagerBasedRlEnv):
+    super().__init__(cfg, env)
+
+    if self.cfg.heading_command and self.cfg.ranges.heading is None:
+      raise ValueError("heading_command=True but ranges.heading is set to None.")
+    if self.cfg.ranges.heading and not self.cfg.heading_command:
+      raise ValueError("ranges.heading is set but heading_command=False.")
+
+    self.robot: Entity = env.scene[cfg.entity_name]
+
+    self.vel_command_b = torch.zeros(self.num_envs, 3, device=self.device)
+    self.heading_target = torch.zeros(self.num_envs, device=self.device)
+    self.heading_error = torch.zeros(self.num_envs, device=self.device)
+    self.is_heading_env = torch.zeros(
+      self.num_envs, dtype=torch.bool, device=self.device
+    )
+    self.is_standing_env = torch.zeros_like(self.is_heading_env)
+
+    self.metrics["error_vel_xy"] = torch.zeros(self.num_envs, device=self.device)
+    self.metrics["error_vel_yaw"] = torch.zeros(self.num_envs, device=self.device)
+
+    # Set by create_gui() when the viewer is active.
+    self._joystick_enabled: viser.GuiCheckboxHandle | None = None
+    self._joystick_sliders: list[viser.GuiSliderHandle] = []
+    self._joystick_get_env_idx: Callable[[], int] | None = None
+
+  @property
+  def command(self) -> torch.Tensor:
+    return self.vel_command_b
+
+  def _update_metrics(self) -> None:
+    max_command_time = self.cfg.resampling_time_range[1]
+    max_command_step = max_command_time / self._env.step_dt
+    self.metrics["error_vel_xy"] += (
+      torch.norm(
+        self.vel_command_b[:, :2] - self.robot.data.root_link_lin_vel_b[:, :2], dim=-1
+      )
+      / max_command_step
+    )
+    self.metrics["error_vel_yaw"] += (
+      torch.abs(self.vel_command_b[:, 2] - self.robot.data.root_link_ang_vel_b[:, 2])
+      / max_command_step
+    )
+
+  def _resample_command(self, env_ids: torch.Tensor) -> None:
+    r = torch.empty(len(env_ids), device=self.device)
+    self.vel_command_b[env_ids, 0] = r.uniform_(*self.cfg.ranges.lin_vel_x)
+    self.vel_command_b[env_ids, 1] = r.uniform_(*self.cfg.ranges.lin_vel_y)
+    self.vel_command_b[env_ids, 2] = r.uniform_(*self.cfg.ranges.ang_vel_z)
+    self.vel_command_b[env_ids, :] *= (torch.norm(self.vel_command_b[env_ids, :], dim=1) > 0.1).unsqueeze(1)
+    if self.cfg.heading_command:
+      assert self.cfg.ranges.heading is not None
+      self.heading_target[env_ids] = r.uniform_(*self.cfg.ranges.heading)
+      self.is_heading_env[env_ids] = r.uniform_(0.0, 1.0) <= self.cfg.rel_heading_envs
+    self.is_standing_env[env_ids] = r.uniform_(0.0, 1.0) <= self.cfg.rel_standing_envs
+
+    init_vel_mask = r.uniform_(0.0, 1.0) < self.cfg.init_velocity_prob
+    init_vel_env_ids = env_ids[init_vel_mask]
+    if len(init_vel_env_ids) > 0:
+      root_pos = self.robot.data.root_link_pos_w[init_vel_env_ids]
+      root_quat = self.robot.data.root_link_quat_w[init_vel_env_ids]
+      lin_vel_b = self.robot.data.root_link_lin_vel_b[init_vel_env_ids]
+      lin_vel_b[:, :2] = self.vel_command_b[init_vel_env_ids, :2]
+      root_lin_vel_w = quat_apply(root_quat, lin_vel_b)
+      root_ang_vel_b = self.robot.data.root_link_ang_vel_b[init_vel_env_ids]
+      root_ang_vel_b[:, 2] = self.vel_command_b[init_vel_env_ids, 2]
+      root_state = torch.cat(
+        [root_pos, root_quat, root_lin_vel_w, root_ang_vel_b], dim=-1
+      )
+      self.robot.write_root_state_to_sim(root_state, init_vel_env_ids)
+
+  def _update_command(self) -> None:
+    if self.cfg.heading_command:
+      self.heading_error = wrap_to_pi(self.heading_target - self.robot.data.heading_w)
+      env_ids = self.is_heading_env.nonzero(as_tuple=False).flatten()
+      self.vel_command_b[env_ids, 2] = torch.clip(
+        self.cfg.heading_control_stiffness * self.heading_error[env_ids],
+        min=self.cfg.ranges.ang_vel_z[0],
+        max=self.cfg.ranges.ang_vel_z[1],
+      )
+    standing_env_ids = self.is_standing_env.nonzero(as_tuple=False).flatten()
+    self.vel_command_b[standing_env_ids, :] = 0.0
+
+  # GUI.
+
+  def create_gui(
+    self,
+    name: str,
+    server: "viser.ViserServer",
+    get_env_idx: Callable[[], int],
+  ) -> None:
+    """Create velocity joystick sliders in the Viser viewer."""
+    from viser import Icon
+
+    ranges = self.cfg.ranges
+
+    axes = [
+      ("lin_vel_x", ranges.lin_vel_x[1]),
+      ("lin_vel_y", ranges.lin_vel_y[1]),
+      ("ang_vel_z", ranges.ang_vel_z[1]),
+    ]
+    sliders: list = []
+
+    with server.gui.add_folder(name.capitalize()):
+      enabled = server.gui.add_checkbox("Enable", initial_value=False)
+
+      for label, max_val in axes:
+        max_input = server.gui.add_slider(
+          f"Max {label}",
+          initial_value=max_val,
+          step=0.1,
+          min=0.1,
+          max=10.0,
+        )
+        slider = server.gui.add_slider(
+          label,
+          min=-max_val,
+          max=max_val,
+          step=0.05,
+          initial_value=0.0,
+        )
+
+        @max_input.on_update
+        def _(_ev, _s=slider, _m=max_input) -> None:
+          _s.min = -_m.value
+          _s.max = _m.value
+
+        sliders.append(slider)
+
+      zero_btn = server.gui.add_button("Zero", icon=Icon.SQUARE_X)
+
+      @zero_btn.on_click
+      def _(_) -> None:
+        for s in sliders:
+          s.value = 0.0
+
+    # Store GUI state for compute() override.
+    self._joystick_enabled = enabled
+    self._joystick_sliders = sliders
+    self._joystick_get_env_idx = get_env_idx
+
+  def compute(self, dt: float) -> None:
+    super().compute(dt)
+    if self._joystick_enabled is not None and self._joystick_enabled.value:
+      assert self._joystick_get_env_idx is not None
+      idx = self._joystick_get_env_idx()
+      for i, s in enumerate(self._joystick_sliders):
+        self.vel_command_b[idx, i] = s.value
+
+  # Visualization.
+
+  def _debug_vis_impl(self, visualizer: "DebugVisualizer") -> None:
+    """Draw velocity command and actual velocity arrows."""
+    env_indices = visualizer.get_env_indices(self.num_envs)
+    if not env_indices:
+      return
+
+    cmds = self.command.cpu().numpy()
+    base_pos_ws = self.robot.data.root_link_pos_w.cpu().numpy()
+    base_quat_w = self.robot.data.root_link_quat_w
+    base_mat_ws = matrix_from_quat(base_quat_w).cpu().numpy()
+    lin_vel_bs = self.robot.data.root_link_lin_vel_b.cpu().numpy()
+    ang_vel_bs = self.robot.data.root_link_ang_vel_b.cpu().numpy()
+
+    scale = self.cfg.viz.scale
+    z_offset = self.cfg.viz.z_offset
+
+    for batch in env_indices:
+      base_pos_w = base_pos_ws[batch]
+      base_mat_w = base_mat_ws[batch]
+      cmd = cmds[batch]
+      lin_vel_b = lin_vel_bs[batch]
+      ang_vel_b = ang_vel_bs[batch]
+
+      # Skip if robot appears uninitialized (at origin).
+      if np.linalg.norm(base_pos_w) < 1e-6:
+        continue
+
+      # Helper to transform local to world coordinates.
+      def local_to_world(
+        vec: np.ndarray, pos: np.ndarray = base_pos_w, mat: np.ndarray = base_mat_w
+      ) -> np.ndarray:
+        return pos + mat @ vec
+
+      # Command linear velocity arrow (blue).
+      cmd_lin_from = local_to_world(np.array([0, 0, z_offset]) * scale)
+      cmd_lin_to = local_to_world(
+        (np.array([0, 0, z_offset]) + np.array([cmd[0], cmd[1], 0])) * scale
+      )
+      visualizer.add_arrow(
+        cmd_lin_from, cmd_lin_to, color=(0.2, 0.2, 0.6, 0.6), width=0.015
+      )
+
+      # Command angular velocity arrow (green).
+      cmd_ang_from = cmd_lin_from
+      cmd_ang_to = local_to_world(
+        (np.array([0, 0, z_offset]) + np.array([0, 0, cmd[2]])) * scale
+      )
+      visualizer.add_arrow(
+        cmd_ang_from, cmd_ang_to, color=(0.2, 0.6, 0.2, 0.6), width=0.015
+      )
+
+      # Actual linear velocity arrow (cyan).
+      act_lin_from = local_to_world(np.array([0, 0, z_offset]) * scale)
+      act_lin_to = local_to_world(
+        (np.array([0, 0, z_offset]) + np.array([lin_vel_b[0], lin_vel_b[1], 0])) * scale
+      )
+      visualizer.add_arrow(
+        act_lin_from, act_lin_to, color=(0.0, 0.6, 1.0, 0.7), width=0.015
+      )
+
+      # Actual angular velocity arrow (light green).
+      act_ang_from = act_lin_from
+      act_ang_to = local_to_world(
+        (np.array([0, 0, z_offset]) + np.array([0, 0, ang_vel_b[2]])) * scale
+      )
+      visualizer.add_arrow(
+        act_ang_from, act_ang_to, color=(0.0, 1.0, 0.4, 0.7), width=0.015
+      )
+
+
+@dataclass(kw_only=True)
+class UniformVelocityCommandCfg(CommandTermCfg):
+  entity_name: str
+  heading_command: bool = False
+  heading_control_stiffness: float = 1.0
+  rel_standing_envs: float = 0.0
+  rel_heading_envs: float = 1.0
+  init_velocity_prob: float = 0.0
+
+  @dataclass
+  class Ranges:
+    lin_vel_x: tuple[float, float]
+    lin_vel_y: tuple[float, float]
+    ang_vel_z: tuple[float, float]
+    heading: tuple[float, float] | None = None
+
+  ranges: Ranges
+
+  @dataclass
+  class VizCfg:
+    z_offset: float = 0.2
+    scale: float = 0.5
+
+  viz: VizCfg = field(default_factory=VizCfg)
+
+  def build(self, env: ManagerBasedRlEnv) -> UniformVelocityCommand:
+    return UniformVelocityCommand(self, env)
+
+  def __post_init__(self):
+    if self.heading_command and self.ranges.heading is None:
+      raise ValueError(
+        "The velocity command has heading commands active (heading_command=True) but "
+        "the `ranges.heading` parameter is set to None."
+      )

training_server/rl/src/tasks/velocity/rl/__init__.py

@@ -0,0 +1,3 @@
+from .runner import (
+  VelocityOnPolicyRunner as VelocityOnPolicyRunner,
+)

training_server/rl/src/tasks/velocity/rl/runner.py

@@ -0,0 +1,28 @@
+import os
+
+import wandb
+
+from mjlab.rl import RslRlVecEnvWrapper
+from mjlab.rl.exporter_utils import (
+  attach_metadata_to_onnx,
+  get_base_metadata,
+)
+from mjlab.rl.runner import MjlabOnPolicyRunner
+
+
+class VelocityOnPolicyRunner(MjlabOnPolicyRunner):
+  env: RslRlVecEnvWrapper
+
+  def save(self, path: str, infos=None):
+    super().save(path, infos)
+    policy_path = path.split("model")[0]
+    filename = "policy.onnx"
+    self.export_policy_to_onnx(policy_path, filename)
+    run_name: str = (
+      wandb.run.name if self.logger.logger_type == "wandb" and wandb.run else "local"
+    )  # type: ignore[assignment]
+    onnx_path = os.path.join(policy_path, filename)
+    metadata = get_base_metadata(self.env.unwrapped, run_name)
+    attach_metadata_to_onnx(onnx_path, metadata)
+    if self.logger.logger_type in ["wandb"]:
+      wandb.save(policy_path + filename, base_path=os.path.dirname(policy_path))

training_server/rl/src/tasks/velocity/velocity_env_cfg.py

@@ -0,0 +1,431 @@
+"""Velocity task configuration.
+
+This module provides a factory function to create a base velocity task config.
+Robot-specific configurations call the factory and customize as needed.
+"""
+
+import math
+from dataclasses import replace
+
+from mjlab.envs import ManagerBasedRlEnvCfg
+from mjlab.envs import mdp as envs_mdp
+from mjlab.envs.mdp import dr
+from mjlab.envs.mdp.actions import JointPositionActionCfg
+from mjlab.managers.action_manager import ActionTermCfg
+from mjlab.managers.command_manager import CommandTermCfg
+from mjlab.managers.curriculum_manager import CurriculumTermCfg
+from mjlab.managers.event_manager import EventTermCfg
+from mjlab.managers.metrics_manager import MetricsTermCfg
+from mjlab.managers.observation_manager import ObservationGroupCfg, ObservationTermCfg
+from mjlab.managers.reward_manager import RewardTermCfg
+from mjlab.managers.scene_entity_config import SceneEntityCfg
+from mjlab.managers.termination_manager import TerminationTermCfg
+from mjlab.scene import SceneCfg
+from mjlab.sensor import GridPatternCfg, ObjRef, RayCastSensorCfg
+from mjlab.sim import MujocoCfg, SimulationCfg
+from mjlab.tasks.velocity import mdp
+from mjlab.tasks.velocity.mdp import UniformVelocityCommandCfg
+from mjlab.terrains import TerrainEntityCfg
+from mjlab.terrains.config import ROUGH_TERRAINS_CFG
+from mjlab.utils.noise import UniformNoiseCfg as Unoise
+from mjlab.viewer import ViewerConfig
+
+import src.tasks.velocity.mdp as mdp
+
+
+def make_velocity_env_cfg() -> ManagerBasedRlEnvCfg:
+  """Create base velocity tracking task configuration."""
+
+  ##
+  # Sensors
+  ##
+
+  terrain_scan = RayCastSensorCfg(
+    name="terrain_scan",
+    frame=ObjRef(type="body", name="", entity="robot"),  # Set per-robot.
+    ray_alignment="yaw",
+    pattern=GridPatternCfg(size=(1.6, 1.0), resolution=0.1),
+    max_distance=5.0,
+    exclude_parent_body=True,
+    debug_vis=True,
+    viz=RayCastSensorCfg.VizCfg(show_normals=True),
+  )
+
+  ##
+  # Observations
+  ##
+
+  actor_terms = {
+    "base_ang_vel": ObservationTermCfg(
+      func=mdp.builtin_sensor,
+      params={"sensor_name": "robot/imu_ang_vel"},
+      noise=Unoise(n_min=-0.2, n_max=0.2),
+    ),
+    "projected_gravity": ObservationTermCfg(
+      func=mdp.projected_gravity,
+      noise=Unoise(n_min=-0.05, n_max=0.05),
+    ),
+    "command": ObservationTermCfg(
+      func=mdp.generated_commands,
+      params={"command_name": "twist"},
+    ),
+    "phase": ObservationTermCfg(
+      func=mdp.phase,
+      params={"period": 0.6, "command_name": "twist"},
+    ),
+    "joint_pos": ObservationTermCfg(
+      func=mdp.joint_pos_rel,
+      noise=Unoise(n_min=-0.01, n_max=0.01),
+    ),
+    "joint_vel": ObservationTermCfg(
+      func=mdp.joint_vel_rel,
+      noise=Unoise(n_min=-1.5, n_max=1.5),
+    ),
+    "actions": ObservationTermCfg(func=mdp.last_action),
+    "height_scan": ObservationTermCfg(
+      func=envs_mdp.height_scan,
+      params={"sensor_name": "terrain_scan"},
+      noise=Unoise(n_min=-0.1, n_max=0.1),
+      scale=1 / terrain_scan.max_distance,
+    ),
+  }
+
+  critic_terms = {
+    **actor_terms,
+    "base_lin_vel": ObservationTermCfg(
+      func=mdp.builtin_sensor,
+      params={"sensor_name": "robot/imu_lin_vel"},
+      noise=Unoise(n_min=-0.5, n_max=0.5),
+    ),
+    "height_scan": ObservationTermCfg(
+      func=envs_mdp.height_scan,
+      params={"sensor_name": "terrain_scan"},
+      scale=1 / terrain_scan.max_distance,
+    ),
+    "foot_height": ObservationTermCfg(
+      func=mdp.foot_height,
+      params={"asset_cfg": SceneEntityCfg("robot", site_names=())},  # Set per-robot.
+    ),
+    "foot_air_time": ObservationTermCfg(
+      func=mdp.foot_air_time,
+      params={"sensor_name": "feet_ground_contact"},
+    ),
+    "foot_contact": ObservationTermCfg(
+      func=mdp.foot_contact,
+      params={"sensor_name": "feet_ground_contact"},
+    ),
+    "foot_contact_forces": ObservationTermCfg(
+      func=mdp.foot_contact_forces,
+      params={"sensor_name": "feet_ground_contact"},
+    ),
+  }
+
+  observations = {
+    "actor": ObservationGroupCfg(
+      terms=actor_terms,
+      concatenate_terms=True,
+      enable_corruption=True,
+      history_length=1,
+    ),
+    "critic": ObservationGroupCfg(
+      terms=critic_terms,
+      concatenate_terms=True,
+      enable_corruption=False,
+      history_length=1,
+    ),
+  }
+
+  ##
+  # Metrics
+  ##
+
+  metrics = {
+    "mean_action_acc": MetricsTermCfg(
+      func=mdp.mean_action_acc,
+    ),
+  }
+
+  ##
+  # Actions
+  ##
+
+  actions: dict[str, ActionTermCfg] = {
+    "joint_pos": JointPositionActionCfg(
+      entity_name="robot",
+      actuator_names=(".*",),
+      scale=0.25,  # Override per-robot.
+      use_default_offset=True,
+    )
+  }
+
+  ##
+  # Commands
+  ##
+
+  commands: dict[str, CommandTermCfg] = {
+    "twist": UniformVelocityCommandCfg(
+      entity_name="robot",
+      resampling_time_range=(3.0, 8.0),
+      rel_standing_envs=0.05,
+      heading_command=True,
+      heading_control_stiffness=0.5,
+      debug_vis=True,
+      ranges=UniformVelocityCommandCfg.Ranges(
+        lin_vel_x=(-1.0, 2.0),
+        lin_vel_y=(-1.0, 1.0),
+        ang_vel_z=(-1.0, 1.0),
+        heading=(-math.pi, math.pi),
+      ),
+    )
+  }
+
+  ##
+  # Events
+  ##
+
+  events = {
+    "reset_base": EventTermCfg(
+      func=mdp.reset_root_state_uniform,
+      mode="reset",
+      params={
+        "pose_range": {
+          "x": (-0.5, 0.5),
+          "y": (-0.5, 0.5),
+          "z": (0.0, 0.0),
+          "yaw": (-3.14, 3.14),
+        },
+        "velocity_range": {},
+      },
+    ),
+    "reset_robot_joints": EventTermCfg(
+      func=mdp.reset_joints_by_offset,
+      mode="reset",
+      params={
+        "position_range": (-0.0, 0.0),
+        "velocity_range": (-0.0, 0.0),
+        "asset_cfg": SceneEntityCfg("robot", joint_names=(".*",)),
+      },
+    ),
+    "push_robot": EventTermCfg(
+      func=mdp.push_by_setting_velocity,
+      mode="interval",
+      interval_range_s=(5.0, 6.0),
+      params={
+        "velocity_range": {
+          "x": (-0.5, 0.5),
+          "y": (-0.5, 0.5),
+          "z": (-0.4, 0.4),
+          "roll": (-0.52, 0.52),
+          "pitch": (-0.52, 0.52),
+          "yaw": (-0.78, 0.78),
+        },
+      },
+    ),
+    "foot_friction": EventTermCfg(
+      mode="startup",
+      func=dr.geom_friction,
+      params={
+        "asset_cfg": SceneEntityCfg("robot", geom_names=()),  # Set per-robot.
+        "operation": "abs",
+        "ranges": (0.3, 1.6),
+        "shared_random": True,  # All foot geoms share the same friction.
+      },
+    ),
+    "encoder_bias": EventTermCfg(
+      mode="startup",
+      func=dr.encoder_bias,
+      params={
+        "asset_cfg": SceneEntityCfg("robot"),
+        "bias_range": (-0.015, 0.015),
+      },
+    ),
+    "base_com": EventTermCfg(
+      mode="startup",
+      func=dr.body_com_offset,
+      params={
+        "asset_cfg": SceneEntityCfg("robot", body_names=()),  # Set per-robot.
+        "operation": "add",
+        "ranges": {
+          0: (-0.05, 0.05),
+          1: (-0.05, 0.05),
+          2: (-0.05, 0.05),
+        },
+      },
+    ),
+  }
+
+  ##
+  # Rewards
+  ##
+
+  rewards = {
+    "track_linear_velocity": RewardTermCfg(
+      func=mdp.track_linear_velocity,
+      weight=1.0,
+      params={"command_name": "twist", "std": math.sqrt(0.25)},
+    ),
+    "track_angular_velocity": RewardTermCfg(
+      func=mdp.track_angular_velocity,
+      weight=1.0,
+      params={"command_name": "twist", "std": math.sqrt(0.5)},
+    ),
+    "body_orientation_l2": RewardTermCfg(
+      func=mdp.body_orientation_l2,
+      weight=-1.0,
+      params={"asset_cfg": SceneEntityCfg("robot", body_names=())},  # Set per-robot.
+    ),
+    "pose": RewardTermCfg(
+      func=mdp.variable_posture,
+      weight=1.0,
+      params={
+        "asset_cfg": SceneEntityCfg("robot", joint_names=".*"),
+        "command_name": "twist",
+        "std_standing": {},  # Set per-robot.
+        "std_walking": {},  # Set per-robot.
+        "std_running": {},  # Set per-robot.
+        "walking_threshold": 0.1,
+        "running_threshold": 1.5,
+      },
+    ),
+    "body_ang_vel": RewardTermCfg(
+      func=mdp.body_angular_velocity_penalty,
+      weight=-0.05,  # Override per-robot
+      params={"asset_cfg": SceneEntityCfg("robot", body_names=())},  # Set per-robot.
+    ),
+    "angular_momentum": RewardTermCfg(
+      func=mdp.angular_momentum_penalty,
+      weight=-0.025,  # Override per-robot
+      params={"sensor_name": "robot/root_angmom"},
+    ),
+    "is_terminated": RewardTermCfg(func=mdp.is_terminated, weight=-200.0),
+    "joint_acc_l2": RewardTermCfg(func=mdp.joint_acc_l2, weight=-2.5e-7),
+    "joint_pos_limits": RewardTermCfg(func=mdp.joint_pos_limits, weight=-10.0),
+    "action_rate_l2": RewardTermCfg(func=mdp.action_rate_l2, weight=-0.05),
+    "foot_gait": RewardTermCfg(
+      func=mdp.feet_gait,
+      weight=0.5,
+      params={
+        "period": 0.6,
+        "offset": [0.0, 0.5],
+        "threshold": 0.56,
+        "command_threshold": 0.1,
+        "command_name": "twist",
+        "sensor_name": "feet_ground_contact",
+      }
+    ),
+    "foot_clearance": RewardTermCfg(
+      func=mdp.feet_clearance,
+      weight=-1.0,
+      params={
+        "target_height": 0.10,
+        "command_name": "twist",
+        "command_threshold": 0.1,
+        "asset_cfg": SceneEntityCfg("robot", site_names=()),  # Set per-robot.
+      },
+    ),
+    "foot_slip": RewardTermCfg(
+      func=mdp.feet_slip,
+      weight=-0.25,
+      params={
+        "sensor_name": "feet_ground_contact",
+        "command_name": "twist",
+        "command_threshold": 0.1,
+        "asset_cfg": SceneEntityCfg("robot", site_names=()),  # Set per-robot.
+      },
+    ),
+    "soft_landing": RewardTermCfg(
+      func=mdp.soft_landing,
+      weight=-1e-3,
+      params={
+        "sensor_name": "feet_ground_contact",
+        "command_name": "twist",
+        "command_threshold": 0.1,
+      },
+    ),
+    "stand_still": RewardTermCfg(
+      func=mdp.stand_still,
+      weight=-1.0,
+      params={
+        "command_name": "twist",
+        "command_threshold": 0.1,
+        "asset_cfg": SceneEntityCfg("robot", joint_names=".*"),
+      },
+    ),
+  }
+
+  ##
+  # Terminations
+  ##
+
+  terminations = {
+    "time_out": TerminationTermCfg(func=mdp.time_out, time_out=True),
+    "fell_over": TerminationTermCfg(
+      func=mdp.bad_orientation,
+      params={"limit_angle": math.radians(70.0)},
+    ),
+  }
+
+  ##
+  # Curriculum
+  ##
+
+  curriculum = {
+    "terrain_levels": CurriculumTermCfg(
+      func=mdp.terrain_levels_vel,
+      params={"command_name": "twist"},
+    ),
+    "command_vel": CurriculumTermCfg(
+      func=mdp.commands_vel,
+      params={
+        "command_name": "twist",
+        "velocity_stages": [
+          {"step": 0, "lin_vel_x": (-0.5, 1.0), "lin_vel_y": (-0.5, 0.5), "ang_vel_z": (-1.0, 1.0)},
+          {"step": 5000 * 24, "lin_vel_x": (-1.0, 2.0), "lin_vel_y": (-1.0, 1.0)},
+        ],
+      },
+    ),
+  }
+
+  ##
+  # Assemble and return
+  ##
+
+  return ManagerBasedRlEnvCfg(
+    scene=SceneCfg(
+      terrain=TerrainEntityCfg(
+        terrain_type="generator",
+        terrain_generator=replace(ROUGH_TERRAINS_CFG),
+        max_init_terrain_level=5,
+      ),
+      sensors=(terrain_scan,),
+      num_envs=1,
+      extent=2.0,
+    ),
+    observations=observations,
+    actions=actions,
+    commands=commands,
+    events=events,
+    rewards=rewards,
+    terminations=terminations,
+    curriculum=curriculum,
+    metrics=metrics,
+    viewer=ViewerConfig(
+      origin_type=ViewerConfig.OriginType.ASSET_BODY,
+      entity_name="robot",
+      body_name="",  # Set per-robot.
+      distance=3.0,
+      elevation=-5.0,
+      azimuth=90.0,
+    ),
+    sim=SimulationCfg(
+      nconmax=35,
+      njmax=1500,
+      mujoco=MujocoCfg(
+        timestep=0.005,
+        iterations=10,
+        ls_iterations=20,
+      ),
+    ),
+    decimation=4,
+    episode_length_s=20.0,
+  )

training_server/server.py

@@ -25,7 +25,7 @@ from pathlib import Path
 from typing import Any
 from urllib.parse import unquote, urlsplit
 
-VERSION = "0.2.0"
+VERSION = "0.3.0"
 # 浏览器当前 ONNX 运行时只实现 Go2 的 47→12 部署契约；其他任务须由服务启动参数显式放行。
 DEFAULT_TASKS = ("Unitree-Go2-Flat",)
 ACTIVE_STATES = {"queued", "running"}
@@ -130,7 +130,7 @@ class TrainingManager:
         if self.check_environment and self._environment_error is False:
             probe = (
                 "import importlib.util,sys; "
-                "missing=[m for m in ('mjlab','torch','tyro') "
+                "missing=[m for m in ('mjlab','torch','tyro','wandb') "
                 "if importlib.util.find_spec(m) is None]; "
                 "print(','.join(missing)); sys.exit(bool(missing))"
             )
@@ -414,7 +414,7 @@ class TrainingRequestHandler(BaseHTTPRequestHandler):
     manager: TrainingManager
     allowed_origins: tuple[str, ...] = ()
     access_token = ""
-    server_version = "MuJoCoLocalTraining/0.2"
+    server_version = "MuJoCoLocalTraining/0.3"
 
     def log_message(self, format: str, *args: Any) -> None:
         sys.stderr.write(f"[{self.log_date_time_string()}] {format % args}\n")
@@ -553,8 +553,7 @@ def default_trainer_root() -> Path:
     configured = os.environ.get("UNITREE_RL_MJLAB_ROOT")
     if configured:
         return Path(configured)
-    repository = Path(__file__).resolve().parents[2]
+    return Path(__file__).resolve().parent / "rl"
-    return repository.parent.parent / "unitree_rl_mjlab"
 
 
 def parse_args() -> argparse.Namespace:
@@ -570,7 +569,7 @@ def parse_args() -> argparse.Namespace:
         "--trainer-root",
         type=Path,
         default=default_trainer_root(),
-        help="unitree_rl_mjlab 工程目录",
+        help="训练工程目录；默认使用仓库内置的 Go2 训练器",
     )
     parser.add_argument(
         "--trainer-python", default=sys.executable, help="已安装 mjlab/torch 的 Python 解释器"

training_server/tests/test_server.py

@@ -14,6 +14,7 @@ from server import (  # noqa: E402
     TrainingJob,
     TrainingManager,
     TrainingRequestHandler,
+    default_trainer_root,
     termination_signal_handler,
 )
 
@@ -57,6 +58,22 @@ out.write_bytes(b'onnx')
         value.update(patch)
         return value
 
+    def test_default_trainer_is_bundled_with_go2_assets(self):
+        trainer_root = default_trainer_root()
+        self.assertEqual(trainer_root, Path(__file__).resolve().parents[1] / "rl")
+        self.assertTrue((trainer_root / "scripts" / "train.py").is_file())
+        self.assertTrue(
+            (
+                trainer_root
+                / "src"
+                / "assets"
+                / "robots"
+                / "unitree_go2"
+                / "xmls"
+                / "go2.xml"
+            ).is_file()
+        )
+
     def test_validates_allowlist_and_limits(self):
         with self.assertRaises(ApiError):
             self.manager.parse_config(self.payload(taskId="shell injection"))

web_platform/README.md

@@ -112,17 +112,16 @@ Python 控制器是可信的单文件脚本，必须同步定义 `step(ctx, stat
 
 ## 本地强化学习训练
 
-训练仍由本机 Python/mjlab 进程执行，但可以从右侧“控制 → 本地强化学习训练”直接发起和管理。先使用安装了 mjlab、PyTorch 及训练依赖的 Python 启动本地桥接服务：
+训练仍由本机 Python/mjlab 进程执行，但可以从右侧“控制 → 本地强化学习训练”直接发起和管理。仓库已经内置默认 Go2 任务的训练代码与资产；先按 [`training_server/README.md`](../training_server/README.md) 安装训练依赖，再使用对应 Python 启动本地桥接服务：
 
 ```bash
 npm run training-server -- \
-  --trainer-root /path/to/unitree_rl_mjlab \
   --trainer-python /path/to/training-env/bin/python
 ```
 
 服务启动时会在终端输出一个随机访问令牌；在界面中填写该令牌后连接。令牌仅保存在当前标签页的 `sessionStorage`。界面默认连接 `http://127.0.0.1:8765`，可选择服务端允许的任务、并行环境数、训练迭代、随机种子、CPU/GPU、GPU 编号和实验记录方式。W&B 默认为本地离线模式，无需登录或 API Key；也可完全禁用，只有明确选择在线模式时才会联网登录。训练期间页面轮询迭代进度与最近日志，可以停止任务；训练成功后点击“导入策略”，生成的 `policy.onnx` 会进入现有 ONNX 加载流程。
 
-桥接服务只监听本机回环地址，并检查 Host、Origin 和 Bearer Token；仅接受允许列表中的任务和经过范围校验的参数，不执行前端提供的 Shell 命令；一次只运行一个训练进程。当前任务使用 `unitree_rl_mjlab` 自带的机器人资产与环境配置，**不会自动把浏览器中临时编辑的 MJCF/URDF 作为训练环境**。自定义浏览器模型训练需要先在 mjlab 中注册对应 task。服务配置、接口和安全边界见 [`../training_server/README.md`](../training_server/README.md)。
+桥接服务只监听本机回环地址，并检查 Host、Origin 和 Bearer Token；仅接受允许列表中的任务和经过范围校验的参数，不执行前端提供的 Shell 命令；一次只运行一个训练进程。默认任务使用仓库内置的 Go2 机器人资产与环境配置，**不会自动把浏览器中临时编辑的 MJCF/URDF 作为训练环境**。自定义浏览器模型训练需要在兼容的外部训练工程中注册 task，并通过服务的 `--trainer-root` 指定该工程。服务配置、接口和安全边界见 [`../training_server/README.md`](../training_server/README.md)。
 
 ## ONNX 强化学习策略