Skip to content
Skild AISkild AISteel

Robot policy

Articulations

Robots pretrained in IsaacLab with RL, run with their policies.

0.00 s
The steel preset (the default), recordedon an NVIDIA RTX PRO 6000 Blackwell.

Run it

Terminal window
cd examples
# the default preset (steel), then the others
uv run shelf play robot_policy
uv run shelf play robot_policy --solver mujoco
# change solver settings live
uv run shelf tune robot_policy

Source: examples/shelf/07_articulations/robot_policy/ · ported from newton/examples/robot/example_robot_policy.py @ skild-ai/newton 7565a26b1

Presets

PresetSolverStatusChecked
steeldefaultsteelworks2026-11-04, newton 7565a26b1, NVIDIA GeForce RTX 3090
mujocomujocoworks2026-10-28, newton 7565a26b1, NVIDIA RTX A6000

Solvers with no row are untested on this example. The full settings are inpresets.toml.

Code

examples/shelf/07_articulations/robot_policy/example.py
# SPDX-FileCopyrightText: Copyright (c) 2025 The Newton Developers
# SPDX-License-Identifier: Apache-2.0
# Modified for Steel Shelf: preset-driven execution and shared runner integration.


###########################################################################
# Example Robot control via keyboard
#
# Shows how to control robots pretrained in IsaacLab with RL. Policies are
# loaded from ONNX files and run with Warp-NN's Warp-backed runtime, so PyTorch
# is not required for policy inference.
#
# Press "p" to reset the robot.
# Press "i", "j", "k", "l", "u", "o" to move the robot.
# Run this example with:
# python -m shelf run robot_policy --robot g1_29dof
# python -m shelf run robot_policy --robot g1_23dof
# python -m shelf run robot_policy --robot go2
# python -m shelf run robot_policy --robot anymal
# python -m shelf run robot_policy --robot anymal --physx
###########################################################################


from dataclasses import dataclass
from typing import Any


import numpy as np
import warp as wp
import yaml
from warp_nn.runtime import OnnxRuntime


import newton
import newton.examples
import newton.utils
from newton import JointTargetMode
from newton.examples.robot.onnx_policy_utils import validate_policy_io_shapes
from common import presets




@dataclass
class RobotConfig:
    """Configuration for a robot including asset paths and policy paths."""


    asset_dir: str
    policy_path: dict[str, str]
    asset_path: str
    yaml_path: str




ROBOT_CONFIGS = {
    "anymal": RobotConfig(
        asset_dir="anybotics_anymal_c",
        policy_path={"mjw": "rl_policies/mjw_anymal.onnx", "physx": "rl_policies/physx_anymal.onnx"},
        asset_path="usd/anymal_c.usda",
        yaml_path="rl_policies/anymal.yaml",
    ),
    "go2": RobotConfig(
        asset_dir="unitree_go2",
        policy_path={"mjw": "rl_policies/mjw_go2.onnx", "physx": "rl_policies/physx_go2.onnx"},
        asset_path="usd/go2.usda",
        yaml_path="rl_policies/go2.yaml",
    ),
    "g1_29dof": RobotConfig(
        asset_dir="unitree_g1",
        policy_path={"mjw": "rl_policies/mjw_g1_29DOF.onnx"},
        asset_path="usd/g1_isaac.usd",
        yaml_path="rl_policies/g1_29dof.yaml",
    ),
    "g1_23dof": RobotConfig(
        asset_dir="unitree_g1",
        policy_path={"mjw": "rl_policies/mjw_g1_23DOF.onnx", "physx": "rl_policies/physx_g1_23DOF.onnx"},
        asset_path="usd/g1_minimal.usd",
        yaml_path="rl_policies/g1_23dof.yaml",
    ),
}




@wp.kernel
def _compute_obs_kernel(
    joint_q: wp.array[float],
    joint_qd: wp.array[float],
    joint_pos_initial: wp.array[float],
    physx_to_mjc_idx: wp.array[int],
    gravity_w: wp.vec3,
    command: wp.vec3,
    prev_act: wp.array2d[float],
    num_dofs: int,
    obs: wp.array2d[float],
):
    q = wp.quat(joint_q[3], joint_q[4], joint_q[5], joint_q[6])


    lin_w = wp.vec3(joint_qd[0], joint_qd[1], joint_qd[2])
    ang_w = wp.vec3(joint_qd[3], joint_qd[4], joint_qd[5])


    vel_b = wp.quat_rotate_inv(q, lin_w)
    avel_b = wp.quat_rotate_inv(q, ang_w)
    grav_b = wp.quat_rotate_inv(q, gravity_w)


    obs[0, 0] = vel_b[0]
    obs[0, 1] = vel_b[1]
    obs[0, 2] = vel_b[2]
    obs[0, 3] = avel_b[0]
    obs[0, 4] = avel_b[1]
    obs[0, 5] = avel_b[2]
    obs[0, 6] = grav_b[0]
    obs[0, 7] = grav_b[1]
    obs[0, 8] = grav_b[2]
    obs[0, 9] = command[0]
    obs[0, 10] = command[1]
    obs[0, 11] = command[2]


    for k in range(num_dofs):
        idx = physx_to_mjc_idx[k]
        obs[0, 12 + k] = joint_q[7 + idx] - joint_pos_initial[idx]
        obs[0, 12 + num_dofs + k] = joint_qd[6 + idx]
        obs[0, 12 + 2 * num_dofs + k] = prev_act[0, k]




@wp.kernel
def _build_joint_target_q_kernel(
    act: wp.array2d[float],
    joint_pos_initial: wp.array[float],
    reorder: wp.array[int],
    action_scale: float,
    num_prefix_zeros: int,
    out: wp.array[float],
):
    i = wp.tid()
    if i < num_prefix_zeros:
        out[i] = 0.0
    else:
        j = i - num_prefix_zeros
        idx = reorder[j]
        out[i] = joint_pos_initial[j] + action_scale * act[0, idx]




def load_policy_and_setup_arrays(example: Any, policy_path: str, num_dofs: int, joint_pos_slice: slice):
    """Load ONNX policy and setup device buffers for the policy step."""
    print("[INFO] Loading policy from:", policy_path)
    example.policy = OnnxRuntime(policy_path, device=example.device)
    example.policy_input_name = example.policy.inputs[0].name
    example.policy_output_name = example.policy.outputs[0].name


    if example.state_0.joint_q is not None:
        example._joint_pos_initial_wp = wp.clone(example.state_0.joint_q[joint_pos_slice])
    else:
        example._joint_pos_initial_wp = wp.zeros(num_dofs, dtype=wp.float32, device=example.device)


    expected = 12 + 3 * num_dofs
    obs_dim = int(example.config["num_observations"]) if "num_observations" in example.config else expected
    if obs_dim != expected:
        raise ValueError(
            f"load_policy_and_setup_arrays: config num_observations={obs_dim} does not match the expected "
            f"layout (12 + 3*num_dofs = {expected})"
        )
    validate_policy_io_shapes(
        policy_path,
        example.policy_input_name,
        example.policy_output_name,
        obs_width=obs_dim,
        action_width=num_dofs,
        context="load_policy_and_setup_arrays",
    )
    example._obs_wp = wp.zeros((1, obs_dim), dtype=wp.float32, device=example.device)
    example._prev_act_wp = wp.zeros((1, num_dofs), dtype=wp.float32, device=example.device)
    example.policy.prepare({example.policy_input_name: example._obs_wp})


    example._physx_to_mjc_wp = wp.array(
        np.asarray(example.physx_to_mjc_indices, dtype=np.int32), dtype=wp.int32, device=example.device
    )
    example._mjc_to_physx_wp = wp.array(
        np.asarray(example.mjc_to_physx_indices, dtype=np.int32), dtype=wp.int32, device=example.device
    )
    example._num_dofs = num_dofs




def find_physx_mjwarp_mapping(mjwarp_joint_names, physx_joint_names):
    mjc_to_physx = []
    physx_to_mjc = []
    for j in mjwarp_joint_names:
        if j in physx_joint_names:
            mjc_to_physx.append(physx_joint_names.index(j))


    for j in physx_joint_names:
        if j in mjwarp_joint_names:
            physx_to_mjc.append(mjwarp_joint_names.index(j))


    return mjc_to_physx, physx_to_mjc




class Example:
    def __init__(self, viewer, args):
        self.preset = presets.load(__file__)
        if args.robot not in ROBOT_CONFIGS:
            raise ValueError(f"Unknown robot: {args.robot}. Available: {list(ROBOT_CONFIGS.keys())}")
        robot_config = ROBOT_CONFIGS[args.robot]
        print(f"[INFO] Selected robot: {args.robot}")


        asset_directory = str(newton.utils.download_asset(robot_config.asset_dir))
        print(f"[INFO] Asset directory: {asset_directory}")


        yaml_file_path = f"{asset_directory}/{robot_config.yaml_path}"
        with open(yaml_file_path, encoding="utf-8") as f:
            config = yaml.safe_load(f)
        num_dofs = config["num_dofs"]
        print(f"[INFO] Loaded config with {num_dofs} DOFs")


        mjc_to_physx = list(range(num_dofs))
        physx_to_mjc = list(range(num_dofs))


        if args.physx:
            if "physx" not in robot_config.policy_path or "physx_joint_names" not in config:
                physx_robots = [name for name, cfg in ROBOT_CONFIGS.items() if "physx" in cfg.policy_path]
                raise ValueError(
                    f"PhysX policy not available for robot '{args.robot}'. Robots with PhysX support: {physx_robots}"
                )
            policy_path = f"{asset_directory}/{robot_config.policy_path['physx']}"
            mjc_to_physx, physx_to_mjc = find_physx_mjwarp_mapping(
                config["mjw_joint_names"], config["physx_joint_names"]
            )
            if len(mjc_to_physx) != num_dofs or len(physx_to_mjc) != num_dofs:
                missing_mjw = sorted(set(config["physx_joint_names"]) - set(config["mjw_joint_names"]))
                missing_physx = sorted(set(config["mjw_joint_names"]) - set(config["physx_joint_names"]))
                raise ValueError(
                    "PhysX/MJWarp joint mapping is incomplete: "
                    f"expected {num_dofs} DOFs, got {len(mjc_to_physx)} MJWarp-to-PhysX and "
                    f"{len(physx_to_mjc)} PhysX-to-MJWarp entries. "
                    f"Missing from MJWarp: {missing_mjw}; missing from PhysX: {missing_physx}"
                )
        else:
            policy_path = f"{asset_directory}/{robot_config.policy_path['mjw']}"


        fps = self.preset.fps
        self.frame_dt = 1.0e0 / fps
        self.decimation = 4
        self.cycle_time = 1 / fps * self.decimation


        self.sim_time = 0.0
        self.sim_step = 0
        self.sim_substeps = self.preset.substeps
        self.sim_dt = self.frame_dt / self.sim_substeps


        self.viewer = viewer


        self.config = config
        self.robot_config = robot_config


        self.device = wp.get_device()


        builder = newton.ModelBuilder(up_axis=newton.Axis.Z)
        newton.solvers.SolverMuJoCo.register_custom_attributes(builder)
        builder.default_joint_cfg = newton.ModelBuilder.JointDofConfig(
            armature=0.1,
            limit_ke=1.0e2,
            limit_kd=1.0e0,
        )
        builder.default_shape_cfg.ke = 5.0e4
        builder.default_shape_cfg.kd = 5.0e2
        builder.default_shape_cfg.kf = 1.0e3
        builder.default_shape_cfg.mu = 0.75
        builder.rigid_gap = 0.0


        builder.add_usd(
            newton.examples.get_asset(asset_directory + "/" + robot_config.asset_path),
            xform=wp.transform(wp.vec3(0, 0, 0.8)),
            collapse_fixed_joints=False,
            enable_self_collisions=False,
            joint_ordering="dfs",
            hide_collision_shapes=True,
        )
        builder.approximate_meshes("convex_hull")


        builder.add_ground_plane()


        builder.joint_q[:3] = [0.0, 0.0, 0.76]
        builder.joint_q[3:7] = [0.0, 0.0, 0.7071, 0.7071]
        builder.joint_q[7:] = config["mjw_joint_pos"]


        for i in range(len(config["mjw_joint_stiffness"])):
            builder.joint_target_ke[i + 6] = config["mjw_joint_stiffness"][i]
            builder.joint_target_kd[i + 6] = config["mjw_joint_damping"][i]
            builder.joint_armature[i + 6] = config["mjw_joint_armature"][i]
            builder.joint_target_mode[i + 6] = int(JointTargetMode.POSITION)


        self.model = builder.finalize()
        self.model.set_gravity((0.0, 0.0, -9.81))


        self.build_solver()


        self.state_temp = self.model.state()
        self.state_0 = self.model.state()
        self.state_1 = self.model.state()
        self.control = self.model.control()


        self.viewer.set_model(self.model)
        self.viewer.vsync = True


        newton.eval_fk(self.model, self.state_0.joint_q, self.state_0.joint_qd, self.state_0)


        self._initial_joint_q = wp.clone(self.state_0.joint_q)
        self._initial_joint_qd = wp.clone(self.state_0.joint_qd)


        self.physx_to_mjc_indices = np.asarray(physx_to_mjc, dtype=np.int64)
        self.mjc_to_physx_indices = np.asarray(mjc_to_physx, dtype=np.int64)
        self._gravity_w = wp.vec3(0.0, 0.0, -1.0)
        self._command = wp.vec3(0.0, 0.0, 0.0)
        self._reset_key_prev = False


        self.policy = None
        self.policy_input_name = None
        self.policy_output_name = None
        self._joint_pos_initial_wp = None
        self._obs_wp = None
        self._prev_act_wp = None
        self._physx_to_mjc_wp = None
        self._mjc_to_physx_wp = None
        self._num_dofs = None


        load_policy_and_setup_arrays(self, policy_path, config["num_dofs"], slice(7, None))


        self.capture()


    def build_solver(self):
        """Build the solver and everything bound to its instance; the tuner calls this again to swap it."""
        self.solver = self.preset.make_solver(self.model)
        # MuJoCo sizes the contact buffer; other solvers take the collision pipeline's default
        budget = getattr(self.solver, "get_max_contact_count", None)
        self.collision_pipeline = newton.CollisionPipeline(
            self.model, rigid_contact_max=budget() if budget else None, soft_contact_max=0
        )
        self.contacts = self.collision_pipeline.contacts()
        self.solver_observables = self.solver.observables(
            contacts=self.contacts, kinds={self.solver.ObservableKind.CONTACT_F}
        )


    def capture(self):
        self.graph = None
        self.use_graph = False
        if self.device.is_cpu or self.device.is_mempool_enabled:
            print("[INFO] Using graph capture")
            self.use_graph = True
            # Coord layout: 7 slots for the free base (3 position + 4 quaternion).
            self.control.joint_target_q = wp.zeros(self.config["num_dofs"] + 7, dtype=wp.float32, device=self.device)
            with wp.ScopedCapture() as capture:
                self.simulate()
            self.graph = capture.graph


    def simulate(self):
        need_state_copy = self.use_graph and self.sim_substeps % 2 == 1


        for i in range(self.sim_substeps):
            self.state_0.clear_forces()


            self.viewer.apply_forces(self.state_0)
            if self.preset.name == "steel":
                # Steel takes its contacts from Newton's collision pipeline; MuJoCo and Kamino find their own
                self.collision_pipeline.collide(self.state_0, self.contacts)


            self.solver.step(
                self.state_0,
                self.state_1,
                self.control,
                self.contacts,
                self.sim_dt,
                observables=self.solver_observables if i == self.sim_substeps - 1 else None,
            )


            if need_state_copy and i == self.sim_substeps - 1:
                self.state_0.assign(self.state_1)
            else:
                self.state_0, self.state_1 = self.state_1, self.state_0


    def reset(self):
        print("[INFO] Resetting example")
        wp.copy(self.state_0.joint_q, self._initial_joint_q)
        wp.copy(self.state_0.joint_qd, self._initial_joint_qd)
        wp.copy(self.state_1.joint_q, self._initial_joint_q)
        wp.copy(self.state_1.joint_qd, self._initial_joint_qd)
        newton.eval_fk(self.model, self.state_0.joint_q, self.state_0.joint_qd, self.state_0)
        newton.eval_fk(self.model, self.state_1.joint_q, self.state_1.joint_qd, self.state_1)
        if self._prev_act_wp is not None:
            self._prev_act_wp.zero_()


    def step(self):
        fwd = 1.0 if self.viewer.is_key_down("i") else (-1.0 if self.viewer.is_key_down("k") else 0.0)
        lat = 0.5 if self.viewer.is_key_down("j") else (-0.5 if self.viewer.is_key_down("l") else 0.0)
        rot = 1.0 if self.viewer.is_key_down("u") else (-1.0 if self.viewer.is_key_down("o") else 0.0)
        self._command = wp.vec3(float(fwd), float(lat), float(rot))
        reset_down = bool(self.viewer.is_key_down("p"))
        if reset_down and not self._reset_key_prev:
            self.reset()
        self._reset_key_prev = reset_down


        wp.launch(
            _compute_obs_kernel,
            dim=1,
            inputs=[
                self.state_0.joint_q,
                self.state_0.joint_qd,
                self._joint_pos_initial_wp,
                self._physx_to_mjc_wp,
                self._gravity_w,
                self._command,
                self._prev_act_wp,
                self._num_dofs,
                self._obs_wp,
            ],
            device=self.device,
        )
        out = self.policy({self.policy_input_name: self._obs_wp})
        act_wp = out[self.policy_output_name]


        wp.launch(
            _build_joint_target_q_kernel,
            dim=7 + self._num_dofs,
            inputs=[
                act_wp,
                self._joint_pos_initial_wp,
                self._mjc_to_physx_wp,
                float(self.config["action_scale"]),
                7,
                self.control.joint_target_q,
            ],
            device=self.device,
        )


        wp.copy(self._prev_act_wp, act_wp)


        for _ in range(self.decimation):
            if self.graph:
                wp.capture_launch(self.graph)
            else:
                self.simulate()


        self.sim_time += self.frame_dt


    def render(self):
        self.viewer.begin_frame(self.sim_time)
        self.viewer.log_state(self.state_0)
        self.viewer.log_contacts(self.contacts, self.state_0, observables=self.solver_observables)
        self.viewer.end_frame()


    def test_final(self):
        newton.examples.test_body_state(
            self.model,
            self.state_0,
            "all bodies are above the ground",
            lambda q, qd: q[2] > 0.0,
        )


    @staticmethod
    def create_parser():
        parser = newton.examples.create_parser()
        parser.add_argument(
            "--robot",
            type=str,
            default="g1_29dof",
            choices=list(ROBOT_CONFIGS.keys()),
            help="Robot name to load",
        )
        parser.add_argument(
            "--physx",
            action="store_true",
            help="Run physX policy instead of MJWarp.",
        )
        return parser




if __name__ == "__main__":
    parser = Example.create_parser()
    viewer, args = newton.examples.init(parser)
    newton.examples.run(Example(viewer, args), args)

example.py on GitHub

Steel docs and tutorials are licensed under Apache-2.0. Third-party notices. "Skild", "Skild AI" and the Skild logo are trademarks of Skild AI and are not licensed under Apache-2.0.