diff --git a/docs/sphinx/source/en/2-user_guide/3-backends/7-newton.md b/docs/sphinx/source/en/2-user_guide/3-backends/7-newton.md index 1de946d91..eb7042e83 100644 --- a/docs/sphinx/source/en/2-user_guide/3-backends/7-newton.md +++ b/docs/sphinx/source/en/2-user_guide/3-backends/7-newton.md @@ -103,6 +103,10 @@ as `env.*` fields in the owner YAML: raises an explicit capacity error when a bound is too small; it never silently truncates constraints. - `newton_capacity_check_steps`: how often capacity is checked (default 1). +- `newton_use_cuda_graph`: CUDA graph replay (default `true`). A graph-capable + UniSim runtime captures both Newton state parities after cold capacity + calibration; ineligible CUDA drivers or capture failures warn and fall back + to eager execution. ## Playback and Rendering diff --git a/docs/sphinx/source/zh_CN/2-user_guide/3-backends/7-newton.md b/docs/sphinx/source/zh_CN/2-user_guide/3-backends/7-newton.md index 6c57cbe58..b6c4fe9d7 100644 --- a/docs/sphinx/source/zh_CN/2-user_guide/3-backends/7-newton.md +++ b/docs/sphinx/source/zh_CN/2-user_guide/3-backends/7-newton.md @@ -96,6 +96,9 @@ Newton/MuJoCo-Warp 3.11 要求显式的设备与存储容量,对应 owner YAML 320 / 512)。适配器在冷路径标定 solver 计数,容量不足时抛出明确的 capacity 错误,绝不静默截断约束。 - `newton_capacity_check_steps`:容量检查步数(默认 1)。 +- `newton_use_cuda_graph`:CUDA graph replay(默认 `true`)。支持 graph 的 + UniSim runtime 会在冷路径容量校准后捕获两个 Newton state 奇偶;CUDA + 驱动不满足条件或捕获失败时发出警告并回退 eager 执行。 ## Playback 与渲染 diff --git a/pyproject.rocm.toml b/pyproject.rocm.toml index b987516f9..34b5e1fda 100644 --- a/pyproject.rocm.toml +++ b/pyproject.rocm.toml @@ -25,8 +25,9 @@ dependencies = [ "numpy", # Physics implementations are provided by the independently released # unisim-core package. Version 1.5.0 provides the M2 physical entity, - # immutable variant, layout and selected-reset contracts. - "unisim-core>=1.5.0", + # immutable variant, layout and selected-reset contracts; 1.6.0 adds + # Newton CUDA graph execution. + "unisim-core>=1.6.0", # RL algorithms and async runtimes live in the independently released # uni-rl package (distribution name ``unilab-rl``); see pyproject.toml. "unilab-rl==1.2.1", diff --git a/pyproject.toml b/pyproject.toml index 69d2f55bc..ceb3ac065 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,9 +41,9 @@ dependencies = [ # Physics implementations are provided by the independently released # unisim-core package. Version 1.5.0 provides the M2 physical entity, # immutable variant, layout and selected-reset contracts; 1.5.1 adds - # get_motion_body_ids on the newton/genesis adapters and fixes the - # subprocess motion-id offset (unilabsim/unisim#137). - "unisim-core>=1.5.1", + # get_motion_body_ids on the newton/genesis adapters; 1.6.0 adds Newton + # CUDA graph execution. + "unisim-core>=1.6.0", # RL algorithms and async runtimes (PPO/APPO/SAC/TD3 runners, # collectors, IPC, logging) live in the independently released uni-rl # package (distribution name ``unilab-rl``), consumed via the injected @@ -171,7 +171,7 @@ viser = ["viser>=1.0.26", "trimesh>=3.21.7"] # required-environments; elsewhere the extra is empty and the CLI reports a # targeted runtime diagnostic. superdex = [ - "unisim-core[superdex]>=1.5.1 ; python_version >= '3.12' and sys_platform == 'linux' and platform_machine == 'x86_64'", + "unisim-core[superdex]>=1.6.0 ; python_version >= '3.12' and sys_platform == 'linux' and platform_machine == 'x86_64'", ] [dependency-groups] diff --git a/src/unilab/base/backend_factory.py b/src/unilab/base/backend_factory.py index b4718e606..56e808b51 100644 --- a/src/unilab/base/backend_factory.py +++ b/src/unilab/base/backend_factory.py @@ -57,6 +57,7 @@ def env_backend_kwargs(cfg: "EnvCfg") -> dict[str, Any]: # its own device. if cfg.genesis_device_id is not None: result["genesis_device_id"] = cfg.genesis_device_id + result["newton_use_cuda_graph"] = cfg.newton_use_cuda_graph return result diff --git a/src/unilab/base/base.py b/src/unilab/base/base.py index 225743fe7..a6408f33c 100644 --- a/src/unilab/base/base.py +++ b/src/unilab/base/base.py @@ -72,6 +72,7 @@ class EnvCfg: newton_nconmax: Optional[int] = None newton_njmax: Optional[int] = None newton_capacity_check_steps: int = 1 + newton_use_cuda_graph: bool = True # ``isaacgym`` runs physics in a Python 3.8 worker subprocess (Preview 4 is # EOL and incompatible with the main environment). ``None`` keeps the # backend defaults (device 0, generous handshake/step timeout). @@ -180,6 +181,10 @@ def validate(self): "newton_capacity_check_steps must be a positive integer, " f"got {self.newton_capacity_check_steps!r}" ) + if not isinstance(self.newton_use_cuda_graph, bool): + raise ValueError( + f"newton_use_cuda_graph must be bool, got {self.newton_use_cuda_graph!r}" + ) if self.isaacgym_device_id is not None and ( isinstance(self.isaacgym_device_id, bool) or not isinstance(self.isaacgym_device_id, int) diff --git a/src/unilab/conf/flashsac/task/g1_motion_tracking/newton.yaml b/src/unilab/conf/flashsac/task/g1_motion_tracking/newton.yaml index 6812345a4..4d8ece50c 100644 --- a/src/unilab/conf/flashsac/task/g1_motion_tracking/newton.yaml +++ b/src/unilab/conf/flashsac/task/g1_motion_tracking/newton.yaml @@ -2,8 +2,8 @@ # FlashSAC Newton owner for G1 WBT. Newton's runtime is optional and shares # the MuJoCo 3.11 / MuJoCo-Warp 3.11 / Warp 1.16 line with the mujoco/mjwarp # extras; process device binding supplies the rank-local CUDA device when -# ``newton_device`` is left unset. Requires unisim-core>=1.5.1 for the motion -# body id capability (unilabsim/unisim#137). Native ViewerGL rendering +# ``newton_device`` is left unset. Requires unisim-core>=1.6.0 for the motion +# body id capability and CUDA graph execution. Native ViewerGL rendering # (offline record + interactive) is provided by the backend; playback inherits # auto mode. defaults: @@ -24,3 +24,4 @@ env: newton_nconmax: 320 newton_njmax: 512 newton_capacity_check_steps: 1 + newton_use_cuda_graph: true diff --git a/src/unilab/conf/ppo/task/g1_walk_flat/newton.yaml b/src/unilab/conf/ppo/task/g1_walk_flat/newton.yaml index f8650a882..61034dbc0 100644 --- a/src/unilab/conf/ppo/task/g1_walk_flat/newton.yaml +++ b/src/unilab/conf/ppo/task/g1_walk_flat/newton.yaml @@ -36,6 +36,7 @@ env: newton_nconmax: 320 newton_njmax: 512 newton_capacity_check_steps: 1 + newton_use_cuda_graph: true events: # Newton currently rejects runtime PD gain randomization; keep the owner # fail-closed until that adapter capability is added. @@ -43,4 +44,3 @@ env: play_profile: enabled: false - diff --git a/src/unilab/conf/sac/task/g1_walk_flat/newton.yaml b/src/unilab/conf/sac/task/g1_walk_flat/newton.yaml index 2532efa12..a27df36fa 100644 --- a/src/unilab/conf/sac/task/g1_walk_flat/newton.yaml +++ b/src/unilab/conf/sac/task/g1_walk_flat/newton.yaml @@ -38,6 +38,7 @@ env: newton_nconmax: 320 newton_njmax: 512 newton_capacity_check_steps: 1 + newton_use_cuda_graph: true render_spacing: 2.0 events: # Newton currently rejects runtime PD gain randomization; keep the owner diff --git a/tests/base/backend/test_newton_backend_options.py b/tests/base/backend/test_newton_backend_options.py new file mode 100644 index 000000000..7342f3e88 --- /dev/null +++ b/tests/base/backend/test_newton_backend_options.py @@ -0,0 +1,27 @@ +import pytest + +from unilab.base.backend_factory import env_backend_kwargs +from unilab.base.base import EnvCfg + + +def test_newton_cuda_graph_option_defaults_to_enabled() -> None: + cfg = EnvCfg() + cfg.validate() + + assert cfg.newton_use_cuda_graph is True + assert env_backend_kwargs(cfg)["newton_use_cuda_graph"] is True + + +def test_newton_cuda_graph_option_routes_to_backend_factory() -> None: + cfg = EnvCfg(newton_use_cuda_graph=False) + cfg.validate() + + assert env_backend_kwargs(cfg)["newton_use_cuda_graph"] is False + + +@pytest.mark.parametrize("value", [None, 1, "true"]) +def test_newton_cuda_graph_option_requires_bool(value: object) -> None: + cfg = EnvCfg(newton_use_cuda_graph=value) + + with pytest.raises(ValueError, match="newton_use_cuda_graph must be bool"): + cfg.validate() diff --git a/tests/base/test_isaacgym_fixed_variants.py b/tests/base/test_isaacgym_fixed_variants.py index 6f78d907c..61cd703f7 100644 --- a/tests/base/test_isaacgym_fixed_variants.py +++ b/tests/base/test_isaacgym_fixed_variants.py @@ -158,7 +158,10 @@ def test_public_layout_drift_fails_closed( ).replace( '', '' - '', + '' + '' + '' + "", ), encoding="utf-8", ) diff --git a/tests/envs/locomotion/g1/test_g1_owner_contract.py b/tests/envs/locomotion/g1/test_g1_owner_contract.py index 2f99a643a..47a7bfd2d 100644 --- a/tests/envs/locomotion/g1/test_g1_owner_contract.py +++ b/tests/envs/locomotion/g1/test_g1_owner_contract.py @@ -479,6 +479,7 @@ def test_g1_owner_materializes_complete_plain_manager_cfg( assert env_cfg.newton_nconmax == 320 assert env_cfg.newton_njmax == 512 assert env_cfg.newton_capacity_check_steps == 1 + assert env_cfg.newton_use_cuda_graph is True # Native ViewerGL rendering (interactive viewer + offscreen record) # is supported; playback stays on the base config's auto mode. assert hydra_cfg.training.play_render_mode == "auto" diff --git a/tests/scripts/test_train_scripts.py b/tests/scripts/test_train_scripts.py index e28e579df..41515c667 100644 --- a/tests/scripts/test_train_scripts.py +++ b/tests/scripts/test_train_scripts.py @@ -2224,6 +2224,8 @@ def test_offpolicy_flashsac_g1_motion_tracking_task_composes(backend: str) -> No assert cfg.training.sim_backend == backend assert cfg.algo.num_envs == 2048 assert cfg.algo.max_iterations == 25000 + if backend == "newton": + assert cfg.env.newton_use_cuda_graph is True def test_offpolicy_rejects_algo_argument_mismatch(): diff --git a/uv.lock b/uv.lock index de163ea49..fd6826944 100644 --- a/uv.lock +++ b/uv.lock @@ -5240,8 +5240,8 @@ requires-dist = [ { name = "trimesh", marker = "extra == 'viser'", specifier = ">=3.21.7" }, { name = "typing-extensions" }, { name = "unilab-rl", specifier = "==1.2.1" }, - { name = "unisim-core", specifier = ">=1.5.1" }, - { name = "unisim-core", extras = ["superdex"], marker = "python_full_version >= '3.12' and platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'superdex'", specifier = ">=1.5.1" }, + { name = "unisim-core", specifier = ">=1.6.0" }, + { name = "unisim-core", extras = ["superdex"], marker = "python_full_version >= '3.12' and platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'superdex'", specifier = ">=1.6.0" }, { name = "viser", marker = "extra == 'viser'", specifier = ">=1.0.26" }, { name = "wandb" }, { name = "warp-lang", marker = "extra == 'mjwarp'", specifier = "==1.16.0" }, @@ -5284,13 +5284,13 @@ wheels = [ [[package]] name = "unisim-core" -version = "1.5.1" +version = "1.6.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/ce/a5/aaa2505237af593ee495963a537893e71a7452fb376e6881d8fb5e4d7695/unisim_core-1.5.1.tar.gz", hash = "sha256:49c5b178ba2e513a0eeebd0c42048a56d72fa1d563f2c66457268130a26ba5af", size = 354707, upload-time = "2026-09-17T04:56:46.441Z" } +sdist = { url = "https://files.pythonhosted.org/packages/df/d2/6802ddbfea15f27e165f8565faae39eeb6ab31a484384b64a3fc459f269a/unisim_core-1.6.0.tar.gz", hash = "sha256:72dbc11170e9516909a9ec81c353aad17222b52a0aa24f0fc63f5290aee3cf51", size = 363157, upload-time = "2026-09-17T10:25:30.036Z" } [package.optional-dependencies] superdex = [ diff --git a/uv.rocm.lock b/uv.rocm.lock index 5c9753d59..8d0c6f706 100644 --- a/uv.rocm.lock +++ b/uv.rocm.lock @@ -3790,7 +3790,7 @@ requires-dist = [ { name = "triton-rocm", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'", specifier = "==3.6.0", index = "https://download.pytorch.org/whl/rocm7.2" }, { name = "typing-extensions" }, { name = "unilab-rl", specifier = "==1.2.1" }, - { name = "unisim-core", specifier = ">=1.5.0" }, + { name = "unisim-core", specifier = ">=1.6.0" }, { name = "viser", marker = "extra == 'viser'", specifier = ">=1.0.26" }, { name = "wandb" }, ] @@ -3830,13 +3830,13 @@ wheels = [ [[package]] name = "unisim-core" -version = "1.5.0" +version = "1.6.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/1c/d8/e142031c7aeb29653ba04b882f5f8df57a7fc5603a7bae682b08121e005f/unisim_core-1.5.0.tar.gz", hash = "sha256:4fd0f276af7f1905812f9189e11422889993423e00a331abdecaa00a8cdd118c", size = 354365, upload-time = "2026-09-17T04:39:54.318Z" } +sdist = { url = "https://files.pythonhosted.org/packages/df/d2/6802ddbfea15f27e165f8565faae39eeb6ab31a484384b64a3fc459f269a/unisim_core-1.6.0.tar.gz", hash = "sha256:72dbc11170e9516909a9ec81c353aad17222b52a0aa24f0fc63f5290aee3cf51", size = 363157, upload-time = "2026-09-17T10:25:30.036Z" } [[package]] name = "urllib3"