PaperlayerView exact revision on arXiv

Source-linked AI summary

FlashNeRD: Performance-First Contact-Rich Neural Robot Dynamics

Mohammadmehdi Ataei, Mohammad Amin Nabian, Jie Xu, Miles Macklin, Yashraj Narang, Eric Heiden, Octi Zhang, Yunrong Guo, Ruben D'Sa, Oliver Hennigh, Gavriel State, Sanjay Choudhry

arXiv:2610.09130v1cs.RO

TL;DR

NeRD leaves limits in inference speed, contact locations, and interaction with moving objects. FlashNeRD combines streaming inference, variable-contact encoding, and hybrid coupling; across tasks it reports faster dynamics and gains in learning and planning.

  • Problem

    NeRD offers little speedup over its source simulator, accepts contacts only at predefined points, and assumes the surrounding environment is static.

    Source: §I · p. 1

  • Method

    FlashNeRD combines parallel streaming inference, a set-based encoder for detected contacts, and hybrid simulation that couples learned robot motion with analytically simulated objects.

    Source: §I · p. 1, §III · p. 2, E. Two-Way Coupling via Hybrid… · p. 4

  • Results

    FlashNeRD’s model is up to 55× faster than Optimized NeRD with lower rollout error, while results also include climbing all six platforms and cube-reorientation performance within 2% of simulator-trained policies.

    Source: §I · p. 1, A. Model-Only Performance and… · p. 5, §VI · p. 8

  • Takeaways & Limitations

    The results support learned simulation for policy learning and planning, with real-world adaptation offered as a path toward dynamics models that improve with experience.

    Source: §VI · p. 8

    Each model is specific to a robot, actuator model, and timestep; performance on smaller GPUs or onboard hardware remains untested.

    Source: §V · p. 8

Abstract

from arXiv · show

Compared with analytical physics, learned dynamics models promise robot simulation that is faster, inherently differentiable, and easily adaptable to real data. Neural Robot Dynamics (NeRD) pursues this by keeping collision detection analytical and replacing a simulator's numerical dynamics for the robot with a learned model. Three limitations remain. NeRD offers little speedup over the simulator it learned from, accepts contact only at predefined points, and has no two-way coupling with objects it manipulates. FlashNeRD removes all three with a parallel streaming architecture that makes each prediction faster and more accurate, an encoder that accepts contacts wherever they occur, and two-way coupling with objects simulated by analytical solvers. Experiments across five robots show faster and more accurate dynamics, faster policy learning, and faster inference-time planning. Across three robots, FlashNeRD's dynamics model is up to $55\times$ faster than an optimized NeRD and more accurate over long rollouts. This speedup extends to policy learning, where PPO trains an ANYmal locomotion policy in 34 seconds, $3.8\times$ faster than the analytical simulator and $2.7\times$ faster than an optimized NeRD. With DIAL-MPC, a sampling-based MPC method, the robot climbs all six test platforms where fixed-contact NeRD manages one, at approximately half the analytical simulator's planning time. Cube-reorientation policies trained with FlashNeRD complete within 2% of the simulator-trained policy's target count.

I. INTRODUCTION

Robot simulation is a practical bottleneck, while NeRD retains limits in throughput, contact location, and interaction with moving objects. FlashNeRD addresses these limits with streaming inference, variable-contact encoding, and hybrid coupling, and reports gains in model speed, policy learning, climbing, and manipulation.

I. INTRODUCTION

  • Analytical simulation is bottlenecked by iterative contact solves, while NeRD re-encodes full histories, accepts contacts only at predefined probes, and assumes a static environment.Learned robot dynamics can be differentiable, batched, and fine-tuned on real trajectories, but these advantages alone did not resolve NeRD’s three limitations.Source: §I · p. 1
  • Up to 55× model-only throughput versus Optimized NeRD with lower rollout error; FlashNeRD trains locomotion 3.2× faster than the analytical simulator and climbs all six platforms, where NeRD manages one.The method also uses two-way coupling for object manipulation.Source: §I · p. 1

II. RELATED WORK

  • FlashNeRD builds on articulated-body NeRD and pools the analytical detector’s unordered contact set instead of using SDFs, mesh-face graphs, or predicted binary contacts.Efficient transformer inference and solver-coupling methods also inform its design.Source: §II · pp. 1–2

III. METHOD

  • FlashNeRD combines parallel streaming inference, a contact encoder that accepts detected contacts wherever they occur, and hybrid simulation that couples learned robot motion with analytically simulated objects.The hybrid simulation resolves contacts and constraints jointly, allowing forces to flow in both directions.Source: §I · p. 1, §III · p. 2

A. Learning a Robot Transition

FlashNeRD learns a robot’s state transition from recent frames, controls, gravity, and contact inputs. Its parallel branches and direct current-frame readout preserve temporal context while reducing sequential work, and the design improves prediction accuracy and reduces recursive rollout drift.

A. Learning a Robot Transition

  • Each input frame contains robot state, control, and gravity direction; the model predicts a state increment from a window of frames and corresponding contact inputs.Contact can be a fixed-length vector from predefined probes or a variable set of detected contacts; the latter suits changing locations such as edges.Source: A. Learning a Robot Transition · p. 2
  • Training uses randomized physical states and actuator commands independently of any controller, and each model serves RL and planning with the same actuator model and timestep when training covers relevant states and contacts.State, gravity, contacts, and predicted increments are expressed in the robot’s base frame to remove dependence on world position and heading.Source: A. Learning a Robot Transition · p. 2

B. An Architecture for Fast and Accurate Rollouts

  • Attention models history-dependent effects while a parallel feedforward branch transforms current features; one packed projection produces both branches, reducing sequential launches per block.Attention reads the current and preceding frames, and the readout retains current-frame information through a direct affine path alongside temporal features.Source: B. An Architecture for Fast and… · p. 2, Figure 1 · p. 3
  • Controlled experiments show that the parallel design and direct current-frame readout improve prediction accuracy and reduce recursive rollout drift.The readout combines temporal features, the current frame, and a direct affine path to predict the state increment.Source: B. An Architecture for Fast and… · pp. 2–3, Figure 1 · p. 3

C. Streaming Dynamics with Bounded KV Caching

FlashNeRD streams one new frame per step and reuses a bounded cache of past keys and values. Fixed-span circular buffers preserve the trained model’s attention behavior while keeping cache storage independent of rollout length.

C. Streaming Dynamics with Bounded KV Caching

  • A full-window transformer re-embeds all frames at every step, whereas FlashNeRD projects only the newest frame and reuses cached keys and values; an unbounded cache would change attention over untrained pairs and lags.The cache is bounded to preserve predictions of the trained model.Source: C. Streaming Dynamics with Bounded KV… · p. 3
  • Two blocks with span three can cover five raw frames because deeper blocks read hidden states that already summarize earlier frames.Each block caches only its preceding pairs, while its receptive field compounds across blocks.Source: C. Streaming Dynamics with Bounded KV… · p. 3, Figure 2 · p. 3
  • Each block reads cached pairs and the new pair in lag order, overwrites the oldest slot without shifting the buffer, and keeps cache storage independent of rollout length.Training uses the same spans and lag biases, and streaming and full-window evaluation agree to floating-point roundoff.Source: Figure 2 · p. 3, C. Streaming Dynamics with Bounded KV… · p. 3

D. Variable Contact Handling

FlashNeRD replaces fixed contact probes with an encoder for the unordered contacts reported by analytical collision detection. It pools contact features into a fixed-size representation that can be refreshed at every predicted state.

D. Variable Contact Handling

  • Contacts are detected at each predicted state and encoded as an unordered set, allowing contact to occur anywhere rather than only at predefined probes.The detector’s variable number and ordering of contacts are handled without assigning them to fixed robot locations.Source: D. Variable Contact Handling · pp. 3–4
  • Each contact row records the involved bodies and local geometry and motion, including position, normal, lever arm, signed gap, and relative surface velocity.Robot links receive learned embeddings; external bodies receive a static/dynamic flag, and rows are padded to capacity K.Source: D. Variable Contact Handling · pp. 3–4
  • Shared row encoding and soft weighted pooling produce M order-independent summary slots, paired with a normalized contact count to retain information averaging would discard.An empty contact set yields a zero bank and count; pooling costs O(KMD) and uses no contact-to-contact attention.Source: D. Variable Contact Handling · p. 4
  • The pooled contact representation shapes the history embedding and directly informs the readout and increment, while cached hidden states carry earlier contact effects.Detection is rerun at every predicted state, so contacts can appear anywhere without redesigning probes.Source: D. Variable Contact Handling · p. 4

E. Two-Way Coupling via Hybrid Simulation

FlashNeRD couples learned robot dynamics with MJWarp’s analytical object dynamics and joint contact solve. The network predicts pre-contact robot motion, and the resolved robot state feeds the next prediction.

E. Two-Way Coupling via Hybrid Simulation

  • The learned model predicts robot motion before object contact, while MJWarp simulates objects and jointly resolves robot-object contacts and constraints.The resolved robot state, rather than the network’s free prediction, becomes the input to the next step.Source: E. Two-Way Coupling via Hybrid… · p. 4, Figure 3 · p. 4
  • The coupled model learns pre-contact dynamics only, avoiding counting contact effects twice because those effects come from the analytical solve.Training targets are velocity increments before contact, joint constraints, and integration.Source: E. Two-Way Coupling via Hybrid… · p. 4
  • MJWarp solves constraint forces and constrained acceleration using the combined robot-object coordinates, mass matrix, and constraint Jacobian, then integrates velocity and position.Because the learned increment precedes the solve, contact forces accelerate the object and decelerate the robot.Source: E. Two-Way Coupling via Hybrid… · p. 4
  • The interface keeps geometry, inertia, and constraints analytical, advances robot and objects on a common timestep, and assumes a fixed-base robot.A floating base would additionally require predicting base motion.Source: E. Two-Way Coupling via Hybrid… · p. 4

IV. RESULTS

The experiments focus on FlashNeRD’s additions beyond NeRD: the transition model, its use in policy learning and planning, and coupling to an analytically simulated object.

IV. RESULTS

  • NeRD has already shown zero-shot sim-to-real transfer and fine-tuning on real data, so these experiments focus on capabilities FlashNeRD adds.Source: §IV · p. 4
  • The experiments measure the transition model alone before evaluating it in policy learning and planning.Source: §IV · p. 4
  • The final experiment couples FlashNeRD to an analytically simulated object.Source: §IV · p. 4

A. Model-Only Performance and Prediction Accuracy

FlashNeRD increases model-only throughput and improves or preserves recursive-rollout accuracy against Optimized NeRD on CartPole, Franka, and G1. The comparison uses matched data and evaluation conditions, with reported ablations showing comparable CartPole error across several design choices.

A. Model-Only Performance and Prediction Accuracy

  • The comparison matches model parameters, data, splits, controls, ten-frame history, normalization, and evaluation trajectories across CartPole, Franka, and G1.CartPole uses about 9k training trajectories and 2.7M parameters; Franka and G1 use 90k trajectories and 10.9M parameters.Source: A. Model-Only Performance and… · p. 4
  • 25 to 55× faster at batch 512,000, FlashNeRD reduces rollout error by 25.2% to 63.8% on CartPole and 5.4% to 34.2% on G1; Franka ties at H = 1 and improves 9.8% to 20.5% at larger H.At batch 16, it is about 2× faster; Franka’s H = 1 difference is 0.06%.Source: Figure 4 · p. 5, A. Model-Only Performance and… · p. 5
  • Accuracy is measured by prefix-averaged rollout NRMSE, normalized by training-state standard deviations and averaged over channels, steps through H, and trajectories; evaluation starts from groundtruth history with 2,048 trajectories per system at H = 1, 15, 100, and 200.H = 1 measures the first open-loop prediction; larger H averages accumulated error.Source: A. Model-Only Performance and… · p. 5
  • CartPole ablations at H = 500 show comparable NRMSE across tested activations, normalization methods, and precision settings.The tested choices include SwiGLU, GEGLU, ReGLU, GLU, MishGLU; RMSNorm and LayerNorm; and FP32, FP16, and BF16.Source: A. Model-Only Performance and… · p. 5

B. Locomotion Learning with Fixed Contact

With fixed contact, FlashNeRD shortened ANYmal-C PPO training while preserving final policy quality. Its faster dynamics inference also reduced collection latency, leaving other parts of the training loop as the main costs.

B. Locomotion Learning with Fixed Contact

  • All learned models used the same four foot-sphere probes, and PPO compared dynamics backends under shared settings across 16,384 environments and 300 iterations.Timings excluded data collection, dynamics training, and process initialization.Source: B. Locomotion Learning with Fixed… · pp. 5–6
  • 2.6× faster than Optimized NeRD and 3.2× faster than MJWarp, FlashNeRD reached the reward threshold in 34.3 s and completed training in 79.1 s.Fig. 5 plots reward against accumulated training time for the learned models.Source: Figure 5 · p. 6, B. Locomotion Learning with Fixed… · p. 6
  • Collection latency fell from 22.51 to 4.41 ms per step and network latency from 19.63 to 1.32 ms versus Optimized NeRD.PPO updates remained near 0.11 to 0.13 s per iteration, making contact processing and policy inference the largest collection costs.Source: B. Locomotion Learning with Fixed… · p. 6
  • Final MJWarp rewards were 28.0 for FlashNeRD, 27.6 for Optimized NeRD, 27.0 for NeRD, and 27.8 for MJWarp-trained policies.Evaluation used 256 environments, 1,000 steps, mean actions, and the training reward.Source: B. Locomotion Learning with Fixed… · p. 6

C. Variable Contact for Contact-Rich Planning

FlashNeRD’s variable-contact model handled platform-edge contacts that fixed probes missed, improving DIAL-MPC climbing success and planning time. Planning nevertheless remained slower than the 20 ms control interval.

C. Variable Contact for Contact-Rich Planning

  • Platform-edge contacts shift across feet and shins as the robot moves, so sparse fixed probes cannot represent all relevant contacts.DIAL-MPC’s thousands of dynamics queries per decision test both contact fidelity and rollout speed.Source: C. Variable Contact for Contact-Rich… · p. 6
  • Fixed-contact NeRD used 20 probes, while FlashNeRD pooled detected contacts into 16 slots and ran detection at every predicted state.FlashNeRD’s contact capacity was 64 before pooling.Source: C. Variable Contact for Contact-Rich… · p. 6
  • Because both architecture and contact representation changed together, this experiment evaluates the full system rather than isolating contact representation alone.Both models were trained on task-independent data and remained frozen during planning.Source: C. Variable Contact for Contact-Rich… · p. 6
  • All six climbs succeeded with FlashNeRD and MJWarp, versus one with fixed-contact NeRD; median planning time was 1.42× lower than NeRD and 2.05× lower than MJWarp.Table II reports DIAL-MPC success and latency across six terrains; Fig. 7 illustrates a 0.181 m platform climb.Source: C. Variable Contact for Contact-Rich… · pp. 6–7, Table II · p. 7, Figure 7 · p. 7
  • Planning took seconds per decision, above the 20 ms control interval, making the climbing result an offline or slowed-clock result.The illustrated collision detector produced 11 to 40 contact rows per step, while 3 to 16 fixed probes were active.Source: C. Variable Contact for Contact-Rich… · p. 7

D. Coupled Manipulation

FlashNeRD coupled learned hand motion with analytical simulation of a manipulated cube, allowing cube reactions to affect the hand. Two-way coupling brought reorientation performance close to MJWarp-trained PPO, unlike one-way coupling.

D. Coupled Manipulation

  • The learned hand was trained from random state-command samples without the cube, then evaluated on cube reorientation absent from training.The one-way and two-way methods shared the frozen model and differed only in whether cube reaction forces were applied to the hand.Source: D. Coupled Manipulation · p. 7
  • FlashNeRD predicted the hand’s free velocity increment while MJWarp advanced the cube and jointly solved contacts and joint constraints; the resolved hand state fed the next prediction.One-way coupling discarded cube reactions but retained hand joint limits and self-contacts.Source: D. Coupled Manipulation · p. 7
  • In the illustrated evaluation, two-way FlashNeRD and MJWarp policies completed four targets in the same time, while the one-way policy dropped the cube without completing one.The one-way policy dropped the cube at 2.47 s; the other two completed four targets by 2.43 s.Source: Figure 8 · p. 7, D. Coupled Manipulation · p. 7
  • Two-way FlashNeRD averaged 14.26 targets per episode, near MJWarp’s 14.49, while one-way coupling averaged 1.79.The respective 95% bootstrap intervals were [13.96, 14.54], [14.19, 14.78], and [1.65, 1.93].Source: D. Coupled Manipulation · pp. 7–8
  • Discarding cube reactions reduced completed targets eightfold and doubled the drop rate, while reaction-aware policies performed within 2% of MJWarp-trained PPO.The comparison held all else fixed between one-way and two-way coupling.Source: D. Coupled Manipulation · p. 8

V. LIMITATIONS

FlashNeRD combines streaming inference, variable contact, and two-way coupling, with reported performance gains; its model scope and deployment evidence remain bounded.

V. LIMITATIONS

  • Models are specific to a robot, actuator model, and timestep; bounded history and preset contact capacity constrain the modeled dynamics and contacts.Analytical collision detection limits end-to-end speedup; timings use GB200 GPUs, with smaller GPUs and onboard hardware untested.Source: §V · p. 8

VI. CONCLUSION

  • Up to 55× faster than Optimized NeRD with lower rollout error, FlashNeRD also supports climbing where fixed probes fail and cube-reorientation policies within 2% of simulator-trained target counts.These results support learned simulation for policy learning and planning.Source: §VI · p. 8
Loading 2610.09130v1…