Run DR Legs policy with Warp-NN #4158
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1 @@ | ||
| Run the Kamino DR Legs RL example from an ONNX model with Warp-NN instead of loading a PyTorch policy checkpoint. | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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: | ||
|
Contributor
There was a problem hiding this comment. 🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift 🔎 Supported by static analysis🏁 Script executed: #!/bin/bash
set -eu
printf '%s\n' '--- target file ---'
cat -n newton/_src/solvers/kamino/examples/rl/onnx_policy.py
printf '%s\n' '--- direct references ---'
rg -n -C 4 'WarpOnnxPolicy|onnx_policy|OnnxRuntime|94|12' newton/_src/solvers/kamino
printf '%s\n' '--- applicable guideline files ---'
fd -HI 'REVIEW_GUIDELINES\.rst|CODING_GUIDELINES\.rst' .Repository: newton-physics/newton Length of output: 50377 🤖 get_repo_knowledge executed:
Length of output: 33578 🏁 Script executed: #!/bin/bash
set -eu
printf '%s\n' '--- DR Legs policy call and action assignment ---'
sed -n '180,215p' newton/_src/solvers/kamino/examples/rl/example_rl_drlegs.py
sed -n '425,455p' newton/_src/solvers/kamino/examples/rl/example_rl_drlegs.py
printf '%s\n' '--- ONNX/Warp-NN references and dependency declarations ---'
rg -n -C 3 'warp[-_]nn|Warp-NN|onnxruntime|onnx' pyproject.toml uv.lock setup.cfg requirements* newton/_src/solvers/kamino 2>/dev/null || true
printf '%s\n' '--- repository guideline files and relevant rules ---'
fd -HI 'REVIEW_GUIDELINES\.rst|CODING_GUIDELINES\.rst' . --exec sh -c 'echo "--- $1"; rg -n -C 2 "model|shape|boundary|validation|Warp-NN|optional" "$1"' sh {}Repository: newton-physics/newton Length of output: 28915 🌐 Web query:
💡 Result: Warp-NN does not use the standard Citations:
🌐 Web query:
💡 Result: In ONNX Runtime, the InferenceSession class does not have an attribute named input_names. To retrieve the names of the input nodes for a model, you should access the get_inputs method of the InferenceSession object [1]. The get_inputs method returns a list of NodeArg objects, each of which contains the name of an input. You can extract these names using a list comprehension: input_names = [input.name for input in session.get_inputs] This approach is the standard way to inspect model inputs in the ONNX Runtime Python API [1]. Citations: 🏁 Script executed: #!/bin/bash
set -eu
printf '%s\n' '--- exact action dimensions and policy construction context ---'
sed -n '1,25p' newton/_src/solvers/kamino/examples/rl/example_rl_drlegs.py
sed -n '335,445p' newton/_src/solvers/kamino/examples/rl/example_rl_drlegs.py
printf '%s\n' '--- pinned Warp-NN tree entries ---'
python3 - <<'PY'
import json, urllib.request
sha = "d9334ff1b00cb429e8336bb373bdfa1a948f8d49"
url = f"https://api.github.com/repos/NVIDIA/warp-nn/git/trees/{sha}?recursive=1"
with urllib.request.urlopen(url) as r:
data = json.load(r)
for item in data.get("tree", []):
p = item.get("path", "")
if "runtime" in p.lower() or "onnx" in p.lower():
print(p)
PYRepository: newton-physics/newton Length of output: 8683 🏁 Script executed: #!/bin/bash
set -eu
sha=d9334ff1b00cb429e8336bb373bdfa1a948f8d49
curl -ksSf "https://api.github.com/repos/NVIDIA/warp-nn/git/trees/${sha}?recursive=1" |
jq -r '.tree[] | select(.path | test("runtime|onnx"; "i")) | .path'Repository: newton-physics/newton Length of output: 284 🏁 Script executed: #!/bin/bash
set -eu
base=https://raw.githubusercontent.com/NVIDIA/warp-nn/d9334ff1b00cb429e8336bb373bdfa1a948f8d49
printf '%s\n' '--- pinned OnnxRuntime implementation ---'
curl -ksSf "$base/warp_nn/runtime/onnx_runtime.py" |
nl -ba |
rg -n -C 5 'class OnnxRuntime|input_names|output_names|shape|batch|def __call__|def forward'
printf '%s\n' '--- pinned runtime tests ---'
curl -ksSf "$base/tests/runtime/test_onnx_runtime.py" |
nl -ba |
rg -n -C 4 'input_names|output_names|shape|batch_size|OnnxRuntime'Repository: newton-physics/newton Length of output: 337 🏁 Script executed: #!/bin/bash
set -eu
sha=d9334ff1b00cb429e8336bb373bdfa1a948f8d49
for path in warp_nn/runtime/onnx_runtime.py tests/runtime/test_onnx_runtime.py; do
printf '%s\n' "--- $path ---"
curl -ksSf "https://api.github.com/repos/NVIDIA/warp-nn/contents/${path}?ref=${sha}" |
jq -r '.content' | tr -d '\n' | base64 -d |
grep -n -E -C 5 'class OnnxRuntime|input_names|output_names|shape|batch|def __call__|def forward'
doneRepository: newton-physics/newton Length of output: 49502 🏁 Script executed: #!/bin/bash
set -eu
rg -n -C 3 'num_worlds|batch_size=1|ArgumentParser|add_argument' \
newton/_src/solvers/kamino/examples/rl/example_rl_drlegs.pyRepository: newton-physics/newton Length of output: 5263 Validate the ONNX output shape at policy load.
🤖 Prompt for AI AgentsSource: Path instructions |
||
| 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) | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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) | ||
|
Comment on lines
+27
to
+48
Contributor
There was a problem hiding this comment. 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win 🔎 Supported by static analysis🤖 get_repo_knowledge executed:
Length of output: 44096 🏁 Script executed: #!/bin/bash
set -eu
printf '%s\n' '--- changed test file ---'
cat -n newton/tests/kamino/test_kamino_rl_onnx.py | sed -n '1,180p'
printf '%s\n' '--- policy definitions and related tests ---'
rg -n --glob '*.py' 'class WarpOnnxPolicy|WarpOnnxPolicy\(|non-contiguous|contiguous|float32|validation|shape' newtonRepository: newton-physics/newton Length of output: 50378 🏁 Script executed: #!/bin/bash
set -eu
printf '%s\n' '--- test file ---'
cat -n newton/tests/kamino/test_kamino_rl_onnx.py | sed -n '1,180p'
printf '%s\n' '--- WarpOnnxPolicy bindings ---'
rg -n --glob '*.py' 'class WarpOnnxPolicy|WarpOnnxPolicy' newton/_src newton/tests | head -80Repository: newton-physics/newton Length of output: 3445 🏁 Script executed: #!/bin/bash
set -eu
cat -n newton/_src/solvers/kamino/examples/rl/onnx_policy.py | sed -n '1,240p'
printf '%s\n' '--- nearby ONNX test files ---'
git ls-files '*onnx*' '*kamino*' | grep -E '(^|/)test_.*onnx|onnx.*test|kamino.*test' | head -80Repository: newton-physics/newton Length of output: 7702 Add tests for the Cover non- 🤖 Prompt for AI AgentsSource: Path instructions |
||
|
|
||
|
|
||
| if __name__ == "__main__": | ||
| unittest.main() | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -42,7 +42,7 @@ sim = [ | |
|
|
||
| # Optional ONNX policy inference for neural actuators and RL policy examples. | ||
| onnx = [ | ||
| "warp-nn[onnx] @ git+https://github.com/NVIDIA/warp-nn.git@d9334ff1b00cb429e8336bb373bdfa1a948f8d49", | ||
|
Contributor
There was a problem hiding this comment. 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win 🔎 Supported by static analysis🌐 Web query:
💡 Result: Yes, PyPI rejects uploads of packages that contain direct URL dependency references (often referred to as PEP 440 or PEP 508 direct references) in their metadata [1][2][3]. When you attempt to upload a distribution (such as a wheel or sdist) to PyPI that includes a direct URL dependency in its Requires-Dist metadata field, the upload will fail with an error similar to "400 Client Error: Bad Request - Packages with direct (URL) references in Requires-Dist are not allowed" [2][4][3]. This restriction is enforced by PyPI's backend software (Warehouse) [4][3]. While various Python Enhancement Proposals (PEPs) regarding metadata have historically permitted index servers to decide whether to allow or prohibit these references [5][6], PyPI explicitly prohibits them to maintain security, reliability, and independence from external, potentially unreliable sources [7][8][3]. Direct references are intended as a tool for software integrators to use locally or in private environments, rather than for packages intended for broad distribution via public indices [7][1][8]. Citations:
🤖 get_repo_knowledge executed:
Length of output: 43387 🏁 Script executed: #!/bin/bash
set -eu
printf '%s\n' '--- pyproject.toml (dependency and metadata sections) ---'
sed -n '1,110p' pyproject.toml
printf '%s\n' '--- repository references to warp-nn and publish configuration ---'
rg -n -i --glob '!*.lock' 'warp-nn|pypi|twine|build|publish|license' .github pyproject.toml setup.cfg setup.py tox.ini 2>/dev/null || true
printf '%s\n' '--- tracked packaging files ---'
git ls-files | rg '(^|/)(pyproject\.toml|setup\.cfg|setup\.py|.*lock|.*license.*|.*publish.*|.*release.*)$' || trueRepository: newton-physics/newton Length of output: 11503 🌐 Web query:
💡 Result: Warp-NN is an open-source library for building and training neural networks for Physical AI workflows using NVIDIA Warp [1][2]. The project is licensed under the Apache License 2.0 [3][4]. Regarding your specific queries: Packaging and Configuration The project uses a pyproject.toml file for configuration and packaging [4]. It is published on PyPI as warp-nn [3]. ONNX Support Warp-NN includes a minimal ONNX inference runtime designed to run on Warp kernels, which avoids a dependency on the full ONNX Runtime [4][5]. It offers an optional "onnx" extra that can be installed for this purpose [4]: pip install warp-nn[onnx] The dependency included in this optional extra is onnx>=1.16.0 [4]. Regarding the specific identifier "d9334ff1b00cb429e8336bb373bdfa1a948f8d49" provided in your query, there is no evidence that this string is a documented license, version, or specific artifact identifier associated with the Warp-NN codebase. It may be a truncated or unrelated hash; the project's commit hashes and versions are tracked via standard Git and PyPI mechanisms [3][6][7]. Top results: [1] https://github.com/NVIDIA/warp-nn [3] https://pypi.org/project/warp-nn/ [4] https://github.com/NVIDIA/warp-nn/blob/develop/pyproject.toml [5] NVIDIA/warp-nn#3 Citations:
🏁 Script executed: #!/bin/bash
set -eu
sha='d9334ff1b00cb429e8336bb373bdfa1a948f8d49'
printf '%s\n' '--- release workflow ---'
sed -n '1,75p' .github/workflows/release.yml
printf '%s\n' '--- locked warp-nn entry ---'
rg -n -A18 -B4 'name = "warp-nn"|warp-nn' uv.lock
printf '%s\n' '--- pinned commit identity ---'
curl -fsSL "https://api.github.com/repos/NVIDIA/warp-nn/commits/$sha" |
jq '{sha: .sha, message: .commit.message, date: .commit.author.date}'
printf '%s\n' '--- pinned commit packaging metadata ---'
curl -fsSL "https://raw.githubusercontent.com/NVIDIA/warp-nn/$sha/pyproject.toml" |
sed -n '1,180p'
printf '%s\n' '--- pinned commit license files ---'
for f in LICENSE LICENSE.md COPYING; do
if curl -fsSL "https://raw.githubusercontent.com/NVIDIA/warp-nn/$sha/$f" >/tmp/warp-nn-license; then
printf '%s\n' "--- $f ---"
sed -n '1,12p' /tmp/warp-nn-license
fi
doneRepository: newton-physics/newton Length of output: 19035 🌐 Web query:
💡 Result: In the warp-nn library, the ONNX support is provided through an optional dependency defined by the extra named "onnx" [1][2]. As of the latest release, version 0.3.1 (released August 10, 2026), this extra requires the onnx package to be at version 1.16.0 or higher [1]. Users can install this support via pip using the command: pip install warp-nn[onnx]. Note that while warp-nn is part of the NVIDIA ecosystem, it is maintained separately from the core warp-lang library [3]. Citations:
Replace the direct 🤖 Prompt for AI AgentsSource: Path instructions |
||
| ] | ||
|
|
||
| # Asset import and mesh processing dependencies | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Add migration guidance for PyTorch policy users.
This changed fragment states that PyTorch checkpoint loading is replaced, but it does not tell users how to provide the ONNX policy or complete the required Warp-NN setup. Add concise migration guidance.
As per path instructions: “Changed, deprecated, and removed entries include migration guidance.”
🤖 Prompt for AI Agents
Source: Path instructions