From cbf9711920a1c5a39e54b55b6bd7403ee7b4c724 Mon Sep 17 00:00:00 2001 From: twidmer Date: Fri, 4 Sep 2026 11:08:18 +0200 Subject: [PATCH] Run DR Legs policy with Warp-NN Load the published DR Legs ONNX policy with Warp-NN while retaining the existing Torch observation and control path. Bridge tensors through Warp's zero-copy Torch interoperability. Pin the asset and Warp-NN revisions needed by the converted policy. --- changelog/+kamino-warp-nn-7c31d8a4.changed.md | 1 + .../kamino/examples/rl/example_rl_drlegs.py | 28 +++++----- .../solvers/kamino/examples/rl/onnx_policy.py | 45 ++++++++++++++++ newton/tests/kamino/test_kamino_rl_onnx.py | 52 +++++++++++++++++++ pyproject.toml | 2 +- uv.lock | 7 +-- 6 files changed, 116 insertions(+), 19 deletions(-) create mode 100644 changelog/+kamino-warp-nn-7c31d8a4.changed.md create mode 100644 newton/_src/solvers/kamino/examples/rl/onnx_policy.py create mode 100644 newton/tests/kamino/test_kamino_rl_onnx.py diff --git a/changelog/+kamino-warp-nn-7c31d8a4.changed.md b/changelog/+kamino-warp-nn-7c31d8a4.changed.md new file mode 100644 index 0000000000..27d2b6cada --- /dev/null +++ b/changelog/+kamino-warp-nn-7c31d8a4.changed.md @@ -0,0 +1 @@ +Run the Kamino DR Legs RL example from an ONNX model with Warp-NN instead of loading a PyTorch policy checkpoint. diff --git a/newton/_src/solvers/kamino/examples/rl/example_rl_drlegs.py b/newton/_src/solvers/kamino/examples/rl/example_rl_drlegs.py index 53664c2be8..d0348789da 100644 --- a/newton/_src/solvers/kamino/examples/rl/example_rl_drlegs.py +++ b/newton/_src/solvers/kamino/examples/rl/example_rl_drlegs.py @@ -16,8 +16,8 @@ # + joint_positions (36D) + action_history (24D) # # Usage: -# python example_rl_drlegs.py --policy path/to/model.pt -# python example_rl_drlegs.py --policy path/to/model.pt --mode async +# python example_rl_drlegs.py --policy path/to/model.onnx +# python example_rl_drlegs.py --policy path/to/model.onnx --mode async # python example_rl_drlegs.py --headless --num-steps 200 ########################################################################### @@ -36,10 +36,10 @@ from newton._src.solvers.kamino.examples import run_headless from newton._src.solvers.kamino.examples.rl.joystick import JoystickConfig, JoystickController from newton._src.solvers.kamino.examples.rl.observations import DrlegsBaseObservation +from newton._src.solvers.kamino.examples.rl.onnx_policy import WarpOnnxPolicy from newton._src.solvers.kamino.examples.rl.simulation import RigidBodySim from newton._src.solvers.kamino.examples.rl.simulation_runner import SimulationRunner from newton._src.solvers.kamino.examples.rl.utils import ( - _load_policy_checkpoint, periodic_encoding, quat_inv_mul, quat_rotate_inv, @@ -78,9 +78,11 @@ "control_decimation": 5, "body_pose_offset_z": 0.265, "usd_model": "dr_legs/usd/dr_legs_with_meshes_and_boxes.usda", - "policy_file": "drlegs_walk.pt", + "policy_file": "drlegs_walk.onnx", } +_DRLEGS_ASSET_REF = "a0547548eaa966c2f5478bee496c3cfba1fa98fc" + def _load_drlegs_config(asset_path: Path) -> dict: """Load walk config YAML from assets, falling back to built-in defaults.""" @@ -122,7 +124,7 @@ def __init__( num_worlds = 1 # USD model path - asset_path = newton.utils.download_asset("disneyresearch", ref="261cd1f429619d8ef4f546bd788ab9dea906b5e1") + asset_path = newton.utils.download_asset("disneyresearch", ref=_DRLEGS_ASSET_REF) usd_model_path = str(asset_path / config["usd_model"]) # Create generic articulated body simulator @@ -479,7 +481,7 @@ def render(self): "--sim-dt", type=float, default=None, help="Physics substep duration in seconds (overrides YAML)" ) parser.add_argument( - "--policy", type=str, default=None, help="Path to an rsl_rl checkpoint .pt file (overrides asset default)" + "--policy", type=str, default=None, help="Path to an ONNX policy file (overrides asset default)" ) parser.add_argument( "--mode", @@ -506,11 +508,8 @@ def render(self): msg.info(f"device: {device}") - # Convert warp device to torch device string - torch_device = "cuda" if device.is_cuda else "cpu" - # Load config from YAML (with hardcoded fallback defaults) - asset_path = newton.utils.download_asset("disneyresearch", ref="261cd1f429619d8ef4f546bd788ab9dea906b5e1") + asset_path = newton.utils.download_asset("disneyresearch", ref=_DRLEGS_ASSET_REF) config = _load_drlegs_config(asset_path) # CLI overrides @@ -522,12 +521,15 @@ def render(self): # Load policy: explicit --policy flag > asset default > random actions policy = None if args.policy: - policy = _load_policy_checkpoint(args.policy, device=torch_device) - msg.info(f"Loaded policy from: {args.policy}") + policy_path = Path(args.policy) + if policy_path.suffix.lower() != ".onnx" or not policy_path.is_file(): + raise FileNotFoundError(f"Expected an existing ONNX policy, got '{policy_path}'") + policy = WarpOnnxPolicy(policy_path, device=device, batch_size=1) + msg.info(f"Loaded policy from: {policy_path}") else: default_policy = asset_path / "dr_legs" / "rl_policies" / config["policy_file"] if default_policy.exists(): - policy = _load_policy_checkpoint(str(default_policy), device=torch_device) + policy = WarpOnnxPolicy(default_policy, device=device, batch_size=1) msg.info(f"Loaded default policy from: {default_policy}") else: msg.info(f"No policy at {default_policy} -- using random actions") diff --git a/newton/_src/solvers/kamino/examples/rl/onnx_policy.py b/newton/_src/solvers/kamino/examples/rl/onnx_policy.py new file mode 100644 index 0000000000..e1abf2dd40 --- /dev/null +++ b/newton/_src/solvers/kamino/examples/rl/onnx_policy.py @@ -0,0 +1,45 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 The Newton Developers +# SPDX-License-Identifier: Apache-2.0 + +"""ONNX policy inference using Warp-NN.""" + +from pathlib import Path +from typing import TYPE_CHECKING + +import warp as wp + +if TYPE_CHECKING: + import torch + + +class WarpOnnxPolicy: + """Evaluate a single-input, single-output ONNX policy with Warp-NN.""" + + def __init__(self, path: str | Path, device: wp.DeviceLike, batch_size: int) -> None: + try: + from warp_nn.runtime import OnnxRuntime # noqa: PLC0415 + except ImportError as exc: # pragma: no cover + raise ImportError( + "Kamino ONNX policy inference requires Warp-NN. Install it with `pip install newton[onnx]`." + ) from exc + + self.runtime = OnnxRuntime(str(path), device=device, batch_size=batch_size, input_batch_axes=0) + if len(self.runtime.input_names) != 1 or len(self.runtime.output_names) != 1: + raise ValueError( + f"Policy '{path}' must have exactly one input and one output; got " + f"inputs={self.runtime.input_names}, outputs={self.runtime.output_names}" + ) + self.input_name = self.runtime.input_names[0] + self.output_name = self.runtime.output_names[0] + + def __call__(self, observation: "torch.Tensor") -> "torch.Tensor": + """Evaluate a contiguous float32 Torch observation batch.""" + import torch + + if observation.dtype != torch.float32: + raise TypeError(f"Policy observations must have dtype torch.float32, got {observation.dtype}") + if not observation.is_contiguous(): + raise ValueError("Policy observations must be contiguous for zero-copy Warp inference") + observation_wp = wp.from_torch(observation, dtype=wp.float32) + output_wp = self.runtime({self.input_name: observation_wp})[self.output_name] + return wp.to_torch(output_wp) diff --git a/newton/tests/kamino/test_kamino_rl_onnx.py b/newton/tests/kamino/test_kamino_rl_onnx.py new file mode 100644 index 0000000000..2c186b7437 --- /dev/null +++ b/newton/tests/kamino/test_kamino_rl_onnx.py @@ -0,0 +1,52 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 The Newton Developers +# SPDX-License-Identifier: Apache-2.0 + +import importlib.util +import os +import tempfile +import unittest + +import numpy as np + +_HAS_ONNX = importlib.util.find_spec("onnx") is not None +_HAS_TORCH = importlib.util.find_spec("torch") is not None +_HAS_WARP_NN = importlib.util.find_spec("warp_nn") is not None + +if _HAS_ONNX and _HAS_TORCH and _HAS_WARP_NN: + import onnx + import torch + from onnx import TensorProto, helper, numpy_helper + + from newton._src.solvers.kamino.examples.rl.onnx_policy import WarpOnnxPolicy + + +@unittest.skipUnless(_HAS_ONNX and _HAS_TORCH and _HAS_WARP_NN, "onnx, torch, or warp-nn not installed") +class TestKaminoRlOnnx(unittest.TestCase): + """Test Warp-NN policy inference used by the Kamino RL example.""" + + def test_policy_accepts_torch_tensor(self): + """Evaluate an ONNX policy from a zero-copy Torch input.""" + weights = np.array([[2.0, -1.0], [0.5, 3.0]], dtype=np.float32) + bias = np.array([0.25, -0.5], dtype=np.float32) + graph = helper.make_graph( + [helper.make_node("Gemm", ["observation", "weight", "bias"], ["action"], transB=1)], + "policy", + [helper.make_tensor_value_info("observation", TensorProto.FLOAT, [None, 2])], + [helper.make_tensor_value_info("action", TensorProto.FLOAT, [None, 2])], + [numpy_helper.from_array(weights, "weight"), numpy_helper.from_array(bias, "bias")], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + with tempfile.TemporaryDirectory(dir=os.getcwd()) as tmp_dir: + path = os.path.join(tmp_dir, "policy.onnx") + onnx.save(model, path) + policy = WarpOnnxPolicy(path, device="cpu", batch_size=2) + observation = torch.tensor([[1.0, 2.0], [-1.0, 0.5]], dtype=torch.float32) + actual = policy(observation) + + expected = observation @ torch.from_numpy(weights).T + torch.from_numpy(bias) + torch.testing.assert_close(actual, expected) + + +if __name__ == "__main__": + unittest.main() diff --git a/pyproject.toml b/pyproject.toml index abc5b20eed..a19af0f12f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -42,7 +42,7 @@ sim = [ # Optional ONNX policy inference for neural actuators and RL policy examples. onnx = [ - "warp-nn[onnx]==0.3.1", + "warp-nn[onnx] @ git+https://github.com/NVIDIA/warp-nn.git@d9334ff1b00cb429e8336bb373bdfa1a948f8d49", ] # Asset import and mesh processing dependencies diff --git a/uv.lock b/uv.lock index f6820d4d35..e638ab0d5c 100644 --- a/uv.lock +++ b/uv.lock @@ -3978,7 +3978,7 @@ requires-dist = [ { name = "viser", marker = "extra == 'docs'", specifier = "==1.0.26" }, { name = "viser", marker = "extra == 'notebook'", specifier = "==1.0.26" }, { name = "warp-lang", specifier = ">=1.17.0", index = "https://pypi.nvidia.com/" }, - { name = "warp-nn", extras = ["onnx"], marker = "extra == 'onnx'", specifier = "==0.3.1" }, + { name = "warp-nn", extras = ["onnx"], marker = "extra == 'onnx'", git = "https://github.com/NVIDIA/warp-nn.git?rev=d9334ff1b00cb429e8336bb373bdfa1a948f8d49" }, ] provides-extras = ["sim", "onnx", "importers", "remesh", "examples", "rtx", "torch-cu12", "torch-cu13", "dev", "docs", "notebook"] @@ -7394,13 +7394,10 @@ wheels = [ [[package]] name = "warp-nn" version = "0.3.1" -source = { registry = "https://pypi.org/simple" } +source = { git = "https://github.com/NVIDIA/warp-nn.git?rev=d9334ff1b00cb429e8336bb373bdfa1a948f8d49#d9334ff1b00cb429e8336bb373bdfa1a948f8d49" } dependencies = [ { name = "warp-lang" }, ] -wheels = [ - { url = "https://files.pythonhosted.org/packages/50/36/a692b000d3aa23e49045edab2bad327c21cd6bfb92b485d0c00f83e0b1e3/warp_nn-0.3.1-py3-none-any.whl", hash = "sha256:185fe2a451d70a797cb796472893740c949e2dd38bb93d03c047591602306b76", size = 72099, upload-time = "2026-08-10T07:33:39.461Z" }, -] [package.optional-dependencies] onnx = [