newton.ik · dylanturpin/trajectory-ik

Trajectory IK in Newton

newton.ik.IKSolverTrajectory solves every frame of a joint trajectory jointly — one nonlinear least-squares problem over all frames — instead of frame by frame. Existing per-frame objectives work unchanged; new temporal objectives couple consecutive frames and make the system block-banded, which is what the GPU solve exploits.

python -m newton.examples ik_trajectory · branch dylanturpin/trajectory-ik · numbers on RTX PRO 6000 Blackwell, July 2026

90-frame pick-and-place trajectory re-optimized jointly every viewer frame while the pick waypoint is dragged. Warm-started, CUDA-graph captured.

What was added

IKSolverTrajectory — Levenberg–Marquardt over all frames of one or more trajectories; rows are (trajectory, frame); damping and accept/reject are per-trajectory; the whole step() is CUDA-graph capturable.
Temporal objectives — IKObjectiveSmoothness (velocity / acceleration / jerk in tangent space), IKObjectiveVelocityLimit, IKObjectiveJointReference. They report bounded per-frame-offset Jacobian blocks; no global Jacobian is ever materialized.
Per-frame objectives reused unchanged — position / rotation / joint-limit objectives (and all three Jacobian modes) just get one target per frame. Frames can be pinned exactly via fixed_frames.
Free-joint correctness — floating-base linear rows carry the exact [p]× lever-arm coupling of Newton's world-frame tangent convention (finite-difference-verified far from the origin).

Why it's fast: the problem is block-banded

Per-frame costs touch one frame; smoothness costs touch k+1 consecutive frames. So the Gauss–Newton Hessian is block-banded — tridiagonal for velocity costs, wider for acceleration/jerk — and the Newton step costs O(T), not O(T³).

path 1 · direct

Batched block-tridiagonal Cholesky

Block Thomas: one CUDA block per trajectory, sequential over frames inside a single wp.tile_cholesky kernel — exact solve.
Higher-order stencils reblocked into k·n superblocks, so accel/jerk costs reuse the same tridiagonal kernel.
Sequential-in-T cost hidden by batch parallelism: 64 trajectories cost about the same as one.

path 2 · cg

Block-Jacobi preconditioned CG

Block-banded Hessian in BSR form (warp.sparse, fixed topology, values rewritten in place) + batched warp.optim.linear.cg.
Hand-rolled block-Jacobi preconditioner (per-frame n×n blocks inverted with tile_cholesky_solve), warm-started across LM iterations.
Parallel over frames — wins for long horizons; inexact but agrees with direct to ~1e-6.

Results — trajectory solve

Shared task: Franka FR3, T=120 frames @ 30 fps, 4 trajectories, end-effector path through pick-and-place waypoints, 16 LM iterations, analytic Jacobians. The frame-by-frame baseline is the per-frame IKSolver, warm-started and batched across the 4 trajectories.

Humanoid retargeting (G1, soma-retargeter data)

soma-retargeter today drives newton.ik frame by frame to retarget human motion (it produced the G1 motions in BONES-SEED). Same front end here — SOMA BVH clips scaled to 14 world-frame effector targets — but the whole clip is solved jointly on a floating-base 29-DoF G1 (free joint, so the lever-arm coupling matters).

Dance clip retargeted to G1 from the same 14 effector targets. The robot with the orange target markers is the joint trajectory solve; the other replays the raw frame-by-frame loop (soma's production approach, without its post-processing) — watch its arms and knees twitch.

Cross-library check: single-frame IK still holds up

Re-run of the earlier single-frame comparison (Newton IK paper protocol: random reachable targets from FK of uniform-random configurations; success = position < 5 mm and orientation < 0.05 rad). Original numbers were measured on an RTX 4090; these are on this box's RTX PRO 6000, so compare ratios, not absolutes.

Appendix

Design space: who solves the banded system how

Everyone in trajectory optimization sits at one of four points on how they solve the block-banded Gauss–Newton system (bandwidth 2(k+1)n−1 for k-order costs; the banded Newton step is O(T·k²n³) — linear in T, same structure as Riccati/DDP and factor-graph message passing):

Approach	Who	Linear algebra	Parallel over T?	Batch?
Direct banded / Riccati	KOMO, GPMP2, Crocoddyl/Aligator, OCS2	banded/block-tridiag Cholesky, exact	no (sequential sweep)	rare (CPU)
Iterative CG on normal eqs	pyroki/jaxls, MPCGPU	matrix-free or BSR SpMV	yes	yes
Quasi-Newton, no linear solve	cuRobo (L-BFGS on B-spline knots)	two-loop recursion only	yes	excellent
Sampling	STOMP, MPPI, MJPC	none	rollouts parallel	excellent

Newton's prototype ships the first two as selectable backends (direct is new territory for GPU-batch; jaxls has no banded direct solver — its only sparse direct option is CHOLMOD via a CPU host-callback). cuRobo never forms a Jacobian: temporal structure lives in its B-spline knots (16 knots × 7 DoF = 112 variables, small enough that sparsity stops mattering), with linear-ish convergence and ~100 fixed iterations in exchange. Parallel-in-time direct methods (associative-scan LQR, cyclic reduction) exist but only pay off for single very long trajectories; batch parallelism is the cheaper GPU win.

Warp assessment: what worked, gaps, proposed extensions

Worked out of the box (pinned 1.15.0.dev)

tile_cholesky / tile_lower_solve / tile_matmul compose into a batched banded factorization inside one kernel with a runtime-length sequential loop — including CUDA graph capture (~0.9 ms for 256 T=240 n=8 tridiagonal solves).
BSR construction + in-place value rewrites (bsr_from_triplets, bsr_block_index), batched CG (batch_offsets), capturable check_every=0 loop. No warp fork needed.

Gaps & bugs hit

preconditioner(A, "diag") is point-Jacobi even for block BSR — block-Jacobi was hand-rolled (~60 lines).
Bug (pinned version; fixed on warp main): iterative solvers' internal reductions don't pass device= — CPU solves return NaNs when the default device is CUDA. Worked around with wp.ScopedDevice.
No direct sparse factorization of any kind in warp — the tile API is the only route (fine for banded, rules out general sparsity).
Tile shapes are compile-time: one specialization per (DoFs, residuals, bandwidth); first-use JIT is 30–60 s per combination.
Shared memory bounds the superblock (k·n ≲ 64–96 fp32) — jerk-order smoothing on high-DoF humanoids needs the non-reblocked banded variant or the CG path.

Small warp extensions worth proposing

preconditioner(A, "block_diag") for BSR.
block_tridiag_cholesky(D, L, b, x) on (B,T,n,n) arrays — the canonical trajopt/Kalman/spline kernel this branch effectively contains.
bsr_diag_add(A, λ) for LM damping; deterministic transpose SpMV; per-batch early exit in batched iterative solvers.

Benchmark protocol & caveats

Hardware: all new numbers from this page: RTX PRO 6000 Blackwell (sm_120), driver 595.84. The reference single-frame numbers (shown hollow) are from the Newton IK paper on an RTX 4090 — absolute times are not comparable across the two GPUs; the cross-solver ratios are the signal.
Single-frame: Newton = 64 Roberts seeds, 16 LM iterations, analytic Jacobians, CUDA graph (asv protocol). cuRobo = v0.7.8 source build, its own published benchmark config (franka.yml, 16 seeds, CUDA graphs, collision checking off). PyRoKi = current pyroki/jaxls, same URDFs as Newton (FR3; plain 19-DoF H1 with 4 EE targets: both elbows + both ankles). Median of 10 runs after warmup; JIT/compile excluded.
Robot models differ slightly across libraries (cuRobo uses its packaged Panda; Newton/PyRoKi use FR3) — inherited from each library's standard config, same as the original comparison.
Trajectory task is identical for Newton and PyRoKi (same URDF, targets, weights, T=120@30fps, frame 0 pinned; Newton 16 LM iterations, jaxls run to its own convergence). cuRobo's trajectory optimizer solves start→goal motion generation on B-spline knots — not full-path tracking — so it is reported as context, not a head-to-head.
Trajectory tracking error: the joint solve trades tracking against smoothness by design (weights: pos 1.0, rot 0.5, limits 10, vel 0.01, accel 5e-4, rest-pose 1e-3). The frame-by-frame baseline tracks near-exactly but with ~3× the peak acceleration and jerk. Longer horizons traverse the same path more slowly, so both error and smoothness improve with T.
G1 retargeting: soma-retargeter front end (SOMA BVH → scaled 14-effector targets, its production ik_map weights), 30 fps downsample, floating base. Frame-by-frame baseline = 24 warm-started iterations per frame (soma production setting); trajectory = 16 LM iterations, velocity+acceleration smoothness replacing the per-frame smoothing filter, without soma's foot-contact post-processing on either side.

Pointers

Branch: dylanturpin/newton-collab @ dylanturpin/trajectory-ik
Demo: python -m newton.examples ik_trajectory
Design notes: TRAJECTORY_IK_NOTES.md (design space, warp assessment, next steps)
Solver: newton/_src/sim/ik/ik_trajectory_solver.py · objectives: ik_trajectory_objectives.py · tests: test_ik_trajectory.py
Benchmarks & this page: traj_ik_report/ on the same branch (git-excluded)