[Newton] Graphed Workflows

[AntoineRichard]

Description

Adds graphed workflows as a separate pipeline for now.

Please include a summary of the change and which issue is fixed. Please also include relevant motivation and context..

Type of change

• New feature (non-breaking change which adds functionality)

How to use:

./isaaclab.sh -p scripts/reinforcement_learning/rsl_rl/train.py --task Isaac-Cartpole-Warp-v0 --num_envs 4096 --headless --max_iterations 300 ./isaaclab.sh -p scripts/reinforcement_learning/rsl_rl/train.py --task Isaac-Ant-Warp-v0 --num_envs 4096 --headless --max_iterations 300 ./isaaclab.sh -p scripts/reinforcement_learning/rsl_rl/train.py --task Isaac-Humanoid-Warp-v0 --num_envs 4096 --headless --max_iterations 300

Results

Provides some speed-up on the Isaac front: For Cartpole:

• Torch API:
  • full collection step: Mean: 3398.505688 us, Std: 462.092124 us, N: 4800
  • actuators: Mean: 458.664717 us, Std: 56.303926 us, N: 9600
  • newton step: Mean: 418.564702 us, Std: 27.599639 us, N: 9599
  • post-processing (rl-related calculations): Mean: 1727.790608 us, Std: 427.153509 us, N: 4800
  • Lab overhead: 2560us
• Warp API:
  • full collection step: Mean: 1097.855936 us, Std: 73.003888 us, N: 4800
  • actuators: Mean: 17.609461 us, Std: 4.145735 us, N: 9599
  • newton step: Mean: 370.281310 us, Std: 29.087660 us, N: 9599
  • post-processing (rl-related calculations): Mean: 24.169094 us, Std: 8.389473 us, N: 4800
  • Lab overhead: 357us

For Ant:

• Torch API:
  • full collection step: Mean: 7646.153628 us, Std: 1809.778802 us, N: 32000
  • actuators: 352.259001 us, Mean: 363.953334 us, Std: 29.461383 us, N: 64000
  • newton step: 2011.052002 us, Mean: 2121.656286 us, Std: 321.844311 us, N: 64000
  • post-processing (rl-related calculations): 2433.182002 us, Mean: 2435.709163 us, Std: 1679.963760 us, N: 32000
  • Lab overhead: 3404us
• Warp API:
  • full collection step: Mean: 4603.307170 us, Std: 257.345966 us, N: 32000
  • actuators: Mean: 18.158209 us, Std: 2.621708 us, N: 64000
  • newton step: Mean: 2066.688861 us, Std: 128.275069 us, N: 63999
  • post-processing (rl-related calculations): Mean: 59.394046 us, Std: 5.979035 us, N: 32000
  • Lab overhead: 471us

For Humanoid:

• Torch API:
  • full collection step: Mean: 17507.664525 us, Std: 2403.361261 us, N: 16000
  • actuators: Mean: 404.597538 us, Std: 137.743652 us, N: 31999
  • newton step: Mean: 6812.051563 us, Std: 331.997178 us, N: 31999
  • post-processing (rl-related calculations): Mean: 2780.663383 us, Std: 2256.250511 us, N: 16000
  • Lab overhead: 3883us
• Warp API:
  • full collection step: Mean: 14906.382912 us, Std: 668.738434 us, N: 16000
  • actuators: Mean: 23.225196 us, Std: 3.397187 us, N: 32000
  • newton step: Mean: 7172.681914 us, Std: 355.531817 us, N: 32000
  • post-processing (rl-related calculations): Mean: 80.205396 us, Std: 6.993390 us, N: 16000
  • Lab overhead: 562us

Checklist

• [x] I have read and understood the contribution guidelines
• [x] I have run the pre-commit checks with ./isaaclab.sh --format
• [x] I have made corresponding changes to the documentation
• [x] My changes generate no new warnings
• [ ] I have added tests that prove my fix is effective or that my feature works
• [ ] I have updated the changelog and the corresponding version in the extension's config/extension.toml file
• [x] I have added my name to the CONTRIBUTORS.md or my name already exists there

[ooctipus]

this is amazing

[greptile-apps[bot]]

Greptile Overview

Greptile Summary

This PR introduces a parallel Warp-based graphed workflow pipeline alongside the existing Torch API for Isaac Lab's direct RL environments. The implementation provides significant performance improvements by leveraging Warp's CUDA kernel execution and graph capture capabilities.

Key Changes:

• New DirectRLEnvWarp base class with CUDA graph capture for action and post-processing steps
• Warp-based actuator models (ActuatorBase, IdealPDActuator, DCMotor, ImplicitActuator) with kernel implementations
• Newton-backed Articulation and RigidObject asset classes using Warp arrays
• Example environments (Cartpole, Ant, Humanoid) demonstrating 2-3x speedups
• Sensor implementations for camera and contact sensors compatible with the new workflow

Performance Results: The benchmarks show substantial overhead reduction:

• Cartpole: Lab overhead reduced from 2560µs to 357µs (7.2x improvement)
• Ant: Lab overhead reduced from 3404µs to 471µs (7.2x improvement)
• Humanoid: Lab overhead reduced from 3883µs to 562µs (6.9x improvement)

Critical Issues Found:

1. DC motor clipping formula is mathematically incorrect (lines 72-73 in kernels.py), missing velocity normalization by vel_limit
2. Gravity normalization divides by sum of components instead of vector magnitude in both articulation and rigid object data classes
3. Action conversion in step() creates a new tensor each iteration (wp.from_torch(action) at line 375), preventing full CUDA graph capture of the training loop

Incomplete Features:

• Event manager and observation/action noise models are commented out throughout
• Multiple TODO comments indicating work in progress
• Missing test coverage (checklist item unchecked)

Confidence Score: 2/5

• This PR has critical mathematical errors in the DC motor model and prevents full CUDA graph capture, limiting the performance benefits
• Score of 2 reflects multiple critical issues: (1) DC motor torque-speed curve formula is mathematically incorrect and will produce wrong motor behaviors, (2) gravity normalization bug affects physics accuracy for non-axis-aligned gravity, (3) action tensor allocation prevents achieving the main goal of full CUDA graph capture. The commented-out event manager and noise model code indicates the feature is incomplete. While the overall architecture is sound and performance improvements are demonstrated, these logical errors need resolution before merge
• source/isaaclab/isaaclab/newton/actuators/kernels.py (DC motor formula), source/isaaclab/isaaclab/envs/direct_rl_env_warp.py (action conversion blocking graph capture), and both data classes with gravity normalization bugs

Important Files Changed

File Analysis

Filename	Score	Overview
source/isaaclab/isaaclab/newton/actuators/kernels.py	2/5	DC motor clipping formula incorrect: doesn't normalize velocity by vel_limit, causing wrong torque limits
source/isaaclab/isaaclab/envs/direct_rl_env_warp.py	3/5	Action conversion creates new tensor each step (line 375), preventing CUDA graph capture; event manager and noise models commented out
source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/cartpole_warp_env.py	5/5	Clean implementation using Warp kernels for observations, rewards, and resets
source/isaaclab/isaaclab/newton/assets/articulation/articulation_data.py	2/5	Gravity normalization uses sum of components instead of vector magnitude (line 91), causing incorrect results for non-axis-aligned gravity
source/isaaclab/isaaclab/newton/assets/rigid_object/rigid_object_data.py	2/5	Same gravity normalization bug as articulation_data.py (line 84), dividing by sum instead of magnitude

Sequence Diagram

sequenceDiagram
    participant RL as RL Training Loop
    participant Env as DirectRLEnvWarp
    participant Scene as InteractiveScene
    participant Asset as Articulation/RigidObject
    participant Actuator as ActuatorBase (Warp)
    participant Newton as Newton Solver
    participant Kernels as Warp Kernels

    Note over RL,Kernels: Environment Step (with CUDA Graph Capture)
    
    RL->>Env: step(action: torch.Tensor)
    Env->>Env: wp.from_torch(action)
    Note over Env: ⚠️ Creates new tensor each step<br/>Prevents full graph capture
    
    loop Decimation Steps
        Env->>Kernels: _pre_physics_step (Warp kernel)
        Kernels->>Kernels: update_actions kernel
        
        Note over Env,Newton: Graph Captured: step_warp_action()
        Env->>Asset: _apply_action()
        Asset->>Actuator: compute() - PD/DC Motor kernels
        Actuator->>Kernels: compute_pd_actuator kernel
        Actuator->>Kernels: clip_efforts_dc_motor kernel
        Note over Actuator,Kernels: ⚠️ DC motor clipping formula<br/>missing velocity normalization
        
        Env->>Scene: write_data_to_sim()
        Scene->>Newton: Apply joint efforts/external forces
        Newton->>Newton: step() - MuJoCo solver
        Scene->>Asset: update(dt)
    end
    
    Note over Env,Kernels: Graph Captured: step_warp_end()
    Env->>Kernels: add_to_env (episode length)
    Env->>Kernels: _get_dones() - Warp kernel
    Env->>Kernels: _get_rewards() - Warp kernel
    Env->>Kernels: _reset_idx() - Reset mask
    Env->>Kernels: _get_observations() - Warp kernel
    
    Env->>Env: wp.to_torch() bindings
    Env-->>RL: obs, rewards, dones (torch.Tensor)