Integrators and the plant layer#

Overview#

A trajectory optimizer needs more than dynamics: it needs the discrete step x_{k+1} = f(x_k, u_k), the derivatives of that step, and the costs and constraints of the problem, all batched over the knot points. GRiD generates these alongside the dynamics so that a whole shooting or collocation iteration can stay on the GPU. The integrator family lives in the grid namespace; the costs and barriers are emitted as a sibling grid_plant namespace and mirrored on the Python handles. The block is emitted for the default all profile, for any request that carries a plant algorithm (the integrators, fdsva_so, the end-effector pose derivatives, com, ccrba, dccrba), baked cost or contact frames, or a multi-target batch, and whenever emit_plant=True (--plant on the command line) is passed; a header generated for kinematics alone carries no plant block, so a kinematics consumer does not compile the barrier families, the activations and the menu presets it never instantiates. The Python bindings always carry the block.

Signature#

x_next = h.integrator(q, qd, u, dt, integrator_type="rk4")   # (B, NQ + NV)
AB     = h.integrator_gradient(q, qd, u, dt, integrator_type="rk4")   # (B, 2*NV, 3*NV)

x_next = h.plant_step(x, u, dt, integrator_type="euler")     # x = [q; qd], (B, NX)
AB     = h.plant_step_gradient(x, u, dt, integrator_type="euler")    # [A | B]
H      = h.plant_step_hessian(x, u, dt, integrator_type="euler")     # NumPy: (B, 2*NV, 3*NV, 3*NV)

value, grad, Hcost = h.quadratic_state_cost(x, x_des, Q)
value, grad, Hcost = h.quadratic_input_cost(u, u_des, R)
value, grad, Hcost = h.ee_pos_cost(q, p_des, W)       # first EE only
value, grad, Hcost = h.com_cost(q, p_des, W)
value, grad, Hcost = h.momentum_cost(q, qd, h_des, W) # tangent [dq | dv] derivatives
value, grad, Hdiag = h.joint_position_barrier(q, lower, upper, mu)

integrator_type is one of euler, semi_implicit_euler, midpoint, rk4, trapezoidal and constant_acceleration; the default is euler. rk3 and the si_euler alias are removed. dt and gravity=-9.81 are runtime arguments. The gradient is [A | B] = ∂x_{k+1}/∂(x, u) in the tangent space, 2·NV rows by 3·NV columns, and the Hessian is the second-order sensitivity of the same step.

Important

Regenerate previously built GPU artifacts after this contract change. trapezoidal means explicit two-stage Heun, not an implicit solve. The old one-evaluation formula is named constant_acceleration: q_next = integrate(q, dt*v + 0.5*dt**2*a), v_next = v + dt*a.

Midpoint, Heun and RK4 advance the full state. Each stage retracts from the initial configuration using the preceding stage velocity, evaluates forward dynamics at that stage configuration and velocity, and combines both velocities and accelerations with the scheme’s weights. Controls and body-local external forces stay constant throughout the step. Midpoint/Heun have second order and RK4 fourth order on Euclidean configurations. The base-point retraction limits rotational convergence to second order even for RK4; no Munthe–Kaas correction is applied.

Inputs, outputs and scope#

A runnable CPU diagnostic illustrates the full-state position update on constant acceleration.

These examples require a handle built with the relevant algorithms. The integrator calls take q at (B, NQ) and qd, u at (B, NV); the plant_step calls take x of shape (B, NQ+NV) and u of shape (B, NV). Both return the same position-plus-velocity state. NX = NQ + NV; derivative outputs use the 2*NV tangent state, not the ambient quaternion coordinates.

plant_step_hessian is a NumPy/CUDA surface, not a JAX or PyTorch handle method. It supports Euler and semi-implicit Euler on fixed and floating bases; other scheme Hessians are not implemented. Multi-stage integrator gradients are not available for spherical joints. MuJoCo-output integration values and gradients currently support Euler and semi-implicit Euler only. See Python Wrappers (grid-rbd) for backend coverage.

Note

The following are planned after this release: step Hessians for the multi-stage schemes, multi-stage gradients on spherical joints, MuJoCo-convention multi-stage and constant-acceleration integration, and the momentum cost on the largest humanoids (its fused dccrba evaluation runs through the spill tiers). Each of these currently raises an error rather than silently degrading.

Weights are diagonal, supplied as vectors: state x_des and Q have shape (B, NX), input u_des and R have (B, NV), and tracking targets and weights have (B, 3) for position or (B, 6) for momentum. Cost outputs are a tuple of value (B,), gradient and Hessian; the state/position-tracking costs use NX-sized outputs, and input costs use NV. Momentum returns a 2*NV gradient and (2*NV, 2*NV) Gauss–Newton Hessian in tangent [dq | dv] coordinates, including configuration and cross blocks. Position barriers take (B, NQ) bounds and velocity/torque barriers take (B, NV) bounds. Their third output is the Hessian diagonal, not a dense matrix. Finite log-barrier bounds require strictly interior inputs; infinite bounds contribute no term.

Implementation#

The Python references are plant_step, plant_step_gradient, plant_step_hessian and the cost and barrier functions of the _plant.py mixin of RBDReference. The CUDA generators are grid_codegen/algorithms/_integrator.py, _integrator_gradient.py and _plant.py, which emits the grid_plant namespace.

In GRiD#

The integrator composes forward dynamics (the mass-matrix-inverse path) with the chosen scheme inside one kernel, so the four acceleration stages of rk4 do not require four kernel launches. Its gradient uses the analytical forward-dynamics gradient, and the step Hessian composes the second-order FDSVA-SO (Forward Dynamics, Second-Order) with the integration map (including the floating-base retract derivatives). This fused integrator kernel does not make an arbitrary sequence of Python calls a single launch.

On a floating base the integrator applies Pinocchio’s SE(3) update to the base pose. With output_convention="mujoco" the free-joint base position takes MuJoCo’s global additive step instead, the quaternion is reordered, and the returned state is in the MuJoCo frame; this is baked into the kernel, not patched on the host.

The plant layer exposes quadratic state and input costs, an end-effector position cost with a Gauss–Newton Hessian, centre-of- mass and centroidal-momentum costs, and log-barriers on joint positions, velocities and torques with a runtime mu. Quadratic costs have exact ambient-coordinate Hessians; end-effector-position and CoM tracking use Gauss–Newton Hessians. Momentum uses r = A(q)*qd - h_des and the full residual Jacobian J = [(dA/dq)*qd | A] from dccrba. Its gradient is J.T @ (W*r) and Hessian J.T @ diag(W) @ J. This omits residual curvature, not configuration dependence, and is not the exact cost Hessian at a general nonzero residual. MuJoCo momentum derivatives include the configuration dependence of the velocity-frame conversion before forming the gradient and Hessian. On quaternion models, position-tracking derivatives retain tangent blocks embedded in NX-sized outputs; momentum uses the compact 2*NV tangent layout.

Every cost term factors as an activation applied per entry to a residual r with Jacobian J: value sum_i a(r_i), gradient J^T a'(r), Gauss-Newton Hessian J^T diag(a''(r)) J. The quadratic terms above are the quadratic activation a = 1/2 w r^2 over their residuals. The grid_plant namespace carries the activation layer and a menu over it.

Activations (GRID_PLANT_ACTIVATION_*, with activation_value / activation_slope / activation_curvature): quadratic; pseudo-Huber a = w (sqrt(r^2 + eps^2) - eps) with eps > 0 (the weighted absolute value as eps goes to zero, quadratic for |r| small against eps); the relaxed log barrier of Grandia et al. on the feature value v in [lo, hi] (w is mu, eps the relaxation delta: the log barrier above delta and its quadratic extension below, so an infeasible iterate keeps a finite value and slope); the one-sided quadratic barrier 1/2 w (max(0, lo - v)^2 + max(0, v - hi)^2); and the smooth hinge w (sqrt(d^2 + eps^2) - eps) on the bound violation d = max(0, lo - v) + max(0, v - hi) with eps > 0, which is zero inside the bounds and pseudo-Huber in the violation (a smooth form of the bound-violation norm of STORM; its curvature is zero inside and w eps^2 / s^3 outside). An infinite bound contributes nothing; every activated term takes per-entry bound buffers that may be null when the activation is not a barrier. GRID_PLANT_ACTIVATION_NONE skips a term in the preset.

Residuals: state_cost and input_cost on x - x_des and u - u_des; linear_feature_cost<T, M, N> on A v + b - target over the state (N = NX) or the input (N = NU) with a caller matrix; and, on a baked cost frame, frame_position_cost on p_f(q) - p_des, frame_orientation_cost on the log-map residual log(R_des^T R_f(q)) with Jacobian Jr(e)^-1 times the frame’s body-axes angular Jacobian (no Euler chart, regular at every orientation), and frame_velocity_cost on the twist J(q) qd - v_des with the residual Jacobian [d(J qd)/dq | J] (fixed base). Cost frames are named joints or fixed joints baked with gen_all_code(cost_frames=[...]) or register_robot(cost_frames=[...]) (the fixed joint’s offset applied; the end-effector target is always frame 0); frame_eval is the caller-scratch evaluator behind the three terms, built on grid::frame_jacobian_inner. gravity_torque is the caller-scratch RNEA(q, 0, 0) a consumer uses to set a knot’s torque target to the gravity-holding torque in place. tracking_cost_menu<T, MX, MU, CF, CF2> composes the frame position, state, input, the three joint barriers, the linear features, the frame orientation and twist and a second frame position behind cost_menu_params<T>, every target, weight, bound and activation parameter a runtime buffer the caller writes per knot; with quadratic activations on the first three terms and the rest off it is tracking_cost. The reference is RBDReference.state_cost, input_cost, linear_feature_cost, frame_position_cost, frame_orientation_cost, frame_velocity_cost and gravity_torque.

On a floating base (a free-flyer root with single-axis joints; a spherical joint is refused by name in the header) the preset is in the tangent layout: the state term is state_cost_tangent on the residual [difference(q_des, q); qd - qd_des] under any activation, with tangent-indexed bounds whose six base-pose entries carry no bound; the input terms, the torque barrier and the linear u feature act on the NUM_VEL - 6 actuated inputs; the position barrier bounds the joint positions and the velocity barrier the tangent velocities; the linear x feature reads [q[7:]; qd]; the gradient is 2 NUM_VEL wide and the Hessian 2 NUM_VEL square, column-major. The frame terms use the same tangent layout on both bases, and frame_velocity_cost takes its configuration derivative along the tangent, the root entering as six columns under the same rigid-motion rules as a joint. GRID_PLANT_MENU_TANGENT, GRID_PLANT_MENU_NU and GRID_PLANT_MENU_NLX state the layout, and one GRID_PLANT_HAS_* constant per emitted piece (ACTIVATIONS, STATE_COST, STATE_COST_TANGENT, INPUT_COST, LINEAR_FEATURE_COST, FRAME_EVAL, FRAME_POSITION_COST, FRAME_ORIENTATION_COST, FRAME_VELOCITY_COST, COST_MENU, GRAVITY_TORQUE, COST_FRAMES) lets a consumer assert on what the header emitted. With zero base weights under the quadratic activation the tangent state term is quadratic_state_cost_tangent.

Select integrator / integrator_gradient for step values / gradients and fdsva_so for the step Hessian dependency; the plant generator only emits operations whose dependencies exist. CoM costs require com and ccrba; full-state momentum costs require dccrba. CUDA momentum callers supply a tier-sized dccrba arena and a workspace slot for spills. Use the wrapper guide’s available-operation checks after registration rather than assuming every handle includes the full plant layer.