Skip to content

mjwarp: DR 消费能力(per-world model mutation + payload 消费) #1401

Description

@TATP-233

S5: mjwarp 后端 DR 消费能力(per-world model mutation)

父 roadmap:#1389

目标

补齐 unilab-mjwarp 后端的 DR 消费能力,使其不再是空集(src/unilab/base/backend/mjwarp/backend.py:1055 当前返回空 DomainRandomizationCapabilities()),mjwarp owner 配置可以启用 base_com/push_robot 等 reset/interval 随机化,作为部署策略的最终训练后端。

背景

差距性质已查明:不是框架缺机制,而是 mjwarp 后端没有消费 payload 的写入路径。UniLab 的 DR contract 层(ResetStateTransactiondr/types.py payload 结构、capabilities 过滤、entity bind/write、baseline getters)已 backend-agnostic 且完备;mjlab 已在 MuJoCo Warp 上证明技术路径可行:

  • mjwarp Model 字段 leading dim 声明为 "*",所有 kernel 按 worldid % field.shape[0] 读取——tile 成 (nworld, ...) 即得 per-world 语义,无需改 kernel
  • mjlab sim/randomization.py:20-56 expand_model_fields:warp kernel 把指定字段 tile 到 nworld;之后必须重新 capture CUDA graph(UniLab 后端 mjwarp/backend.py:492-497 注释已预留此约定)。
  • DR 写入 GPU 原位进行,按涉及字段分级调 mjwarp.set_const / set_const_0 重算派生量(body_mass/ipos/inertia/iquat → set_const;dof_armature → set_const_0;geom_friction/kp/kd 无需重算)。UniLab pin 的 mujoco-warp==3.10.0.3 已有这些 API。

范围

  1. 冷路径字段扩展:按声明的 DR 字段集合在构造/materialize 阶段 tile Model 字段到 (nworld, ...),随后重捕获 graph(按需 expansion 控制显存)。
  2. set_state 消费 ResetRandomizationPayload:在现有 reset barrier 内写 per-world 数组的 reset 行 + 分级 set_const* 重算,逐项放开 supported_reset_terms
  3. interval:push/body force 写 data.xfrc_applied(该字段本来就是 per-world);velocity delta 写 _qvel_cache 行后走现有上传+forward 路径。
  4. 恢复 mjwarp owner yaml 中被 null 掉的 base_com/push_robot

验收

  • effect 测试:逐字段验证 mutation 后物理行为变化且 world 间隔离;与 mujoco 后端 DR 的统计一致性对照;microduck mjwarp owner yaml 全量 DR smoke。

规模

数百行(backend 层 + expansion helper)+ effect 测试。框架层与 contract 层零改动。

注意

  • tile 替换数组后必须重捕获 graph;in-place warp.copy 则 graph-safe。
  • 移植 expansion 字段清单时按 mujoco-warp 3.10 的 Model dataclass 核对(mjlab 用 ~=3.11)。

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions