grid_rbd (Python bindings)#

The grid-rbd package is the surface most users actually call: it turns a URDF into a cached per-robot .so and hands back a batched handle on the numpy, JAX, or PyTorch backend.

See Python backend interfaces for a source-checked method inventory and CPU-checked input examples for CPU-tested packing and inertial-parameter examples.

from grid_rbd import register_robot

h = register_robot("iiwa", "path/to/iiwa14.urdf")     # numpy handle
tau = h.inverse_dynamics(q, qd, qdd)                  # (batch, nq) float32

hj = register_robot("iiwa", "path/to/iiwa14.urdf", backend="jax")
ht = register_robot("iiwa", "path/to/iiwa14.urdf", backend="torch")

See Python Wrappers (grid-rbd) for the guided tour and bindings/examples/AGENT_INTEGRATION_GUIDE.md for the agent-facing lifecycle notes.

Registration#

grid_rbd.register_robot(name: str, urdf_path: str | None = None, *, urdf_string: str | None = None, floating_base: bool = False, ee_joint_names: list[str] | tuple[str, ...] | None = None, max_batch_size: int = 256, cache_dir: str | Path | None = None, force_rebuild: bool = False, cuda_arch: int | None = None, backend: str = 'numpy', allow_fp64: bool = False, dtype: str = 'float32', runtime_inertia: bool = False, runtime_transform: bool = False, runtime_joint_dynamics: bool = False, use_joint_dynamics: bool = False, enable_tool: bool = False, contact_frames: list[str] | tuple[str, ...] | None = None, output_convention: str = 'pinocchio', algorithm_list: list[str] | tuple[str, ...] | str | None = None, enable_mujoco_kernels: bool = True, _profile_overlay: str | None = 'pybind') → RobotHandle[source]#

Register a robot for fast subsequent calls.

Generates grid.cuh from the URDF, compiles a per-robot .so, and caches it under cache_dir (default ~/.cache/grid-rbd/). Idempotent: if a cache entry matching (urdf, options, grid_rbd version, cuda_arch) already exists, the existing .so is reused — no recompile.

Precision (Phase 8): the numpy backend supports a true fp64 compute tier via dtype="float64" — it builds a SEPARATE .so (-DGRID_WRAPPER_T_DOUBLE + the matching codegen knob that re-derives the shared-mem spill tiers at 2× bytes) and the handle takes/returns float64 numpy arrays computed end-to-end in double precision. fp32 (dtype="float32", default) is unchanged and byte-identical. The fp32 and fp64 .so coexist in the cache (dtype is in the cache key). NOTE: fp64 doubles every arena’s smem footprint, lowering occupancy; some big-robot second-order kernels that already max-spill at fp32 may not fit the device opt-in cap at fp64 — those kernels are left unregistered and raise a clear runtime error when called (no new gating). Wave 2a: dtype="float64" works on ALL backends — the fp64 .so carries fp64 jax and torch surfaces (framework arrays must then be float64). allow_fp64 is the LEGACY fp32-compute upcast convenience (compute in fp32, cast i/o to fp64); prefer dtype="float64" for real double precision. allow_fp64 is ignored when dtype="float64".

Parameter groups — and what each costs. A knob that participates in the .so cache key triggers a full codegen + nvcc rebuild the FIRST time a new value is used (minutes on an arm, tens of minutes on a humanoid); after that the cached .so is reused. The others are free.

Model (re-keys the cache): urdf_path / urdf_string, floating_base, ee_joint_names.
Build contents & size (re-keys the cache): max_batch_size, dtype, algorithm_list, enable_mujoco_kernels (only when False), use_joint_dynamics, enable_tool.
Runtime-mutable tables (re-keys the cache): runtime_inertia, runtime_transform, runtime_joint_dynamics.
Cache & toolchain (no codegen effect): cache_dir, force_rebuild, cuda_arch, backend.
Runtime-only (never rebuilds; not in the cache key): output_convention, allow_fp64, _profile_overlay.

Parameters:

name (str) – Human-friendly handle name. Re-registering under the same name with a different URDF overwrites the binding (the old .so lingers in the cache for manual GC).
urdf_path (str | None) – Path to the robot’s URDF file. Mutually exclusive with urdf_string.
urdf_string (str | None, optional) – Inline URDF text (no file on disk). Mutually exclusive with urdf_path. The cache key hashes the URDF bytes, so an inline string and the equivalent file dedupe to the same compiled .so. The string is persisted as entry_dir/robot.urdf for re-runs / debugging.
backend (str, optional) – “numpy” (default) → a numpy RobotHandle; “jax” → a JaxRobotHandle (grid_rbd.jax); “torch” → a TorchRobotHandle (grid_rbd.torch). The jax/torch backends forward to their submodule’s register_robot.
floating_base (bool, optional) – Treat the robot as floating-base. Default False (fixed-base).
ee_joint_names (list[str] | None, optional) – Names of fixed joints to treat as end-effector targets. Default None ⇒ codegen uses all leaf nodes. Currently only the first name is honored (single-target codegen); multi-target support is a v2 concern. Passing a different list changes the cache key, so different target choices land in separate cache entries.
max_batch_size (int, optional) – Compile-time max batch size. Calls with batch <= this run on a single launch; larger batches must be chunked by the caller (a helper will be added in v2).
cache_dir (str | Path | None, optional) – Override the cache root. Default uses platformdirs / $GRID_RBD_CACHE_DIR.
force_rebuild (bool, optional) – Skip the cache hit check and regenerate + recompile.
cuda_arch (int | None, optional) – Compute capability as int (e.g. 120 for sm_120). Default detects via nvidia-smi.
dtype (str, optional) – Compute precision of the .so: "float32" (default) or "float64" (true fp64 tier — see the precision paragraph above). Re-keys the cache; the fp32 and fp64 .so coexist.
allow_fp64 (bool, optional) – LEGACY fp32-compute upcast convenience: accept/return float64 arrays while computing in fp32. Runtime-only (same .so); ignored when dtype="float64". Prefer dtype="float64" for real precision.
enable_tool (bool, optional) – Default False. Build the runtime tool/payload surface (RobotHandle.tool_fext(), attach_tool/detach_tool): a kernel mapping a world-aligned tool-tip wrench at a runtime body to joint-local f_ext rows. Re-keys the cache.
contact_frames (list[str], optional) – Default None. Names of URDF FIXED joints to bake as contact frames (e.g. a quadruped’s foot joints). Builds the multi-contact surface (RobotHandle.contact_fext()): per-frame world-aligned [n_w; f_w] wrenches (moment about the frame origin, LOCAL_WORLD_ALIGNED) -> joint-local f_ext in one kernel, ready to pass as f_ext= to the dynamics ops. Registration order fixes the f_c column order. Re-keys the cache.
runtime_transform (bool, optional) – Like runtime_inertia but for the joint-frame transforms: emits a mutable d_transform_params table ([x,y,z,r,p,y] per joint) + host mutator, and the handle gains set_transform_params / transform_params. Calibration / kinematic-error injection without a recompile. Re-keys the cache.
runtime_joint_dynamics (bool, optional) – Like runtime_inertia but for per-joint damping/friction: emits a mutable [damping(nv) || friction(nv)] table + host mutator and the handle gains set_joint_dynamics_params. Composes with use_joint_dynamics (which bakes the URDF values as the initial table). Re-keys the cache.
_profile_overlay (str | None, optional) – INTERNAL (underscore = not part of the public API; may change without notice). Which launch-config profile block of config/launch_configs/<robot>/<gpu>.json seeds per-algo threads/tier overlays on this handle: "pybind" (default), "torch"/None (what the torch/jax registrars pass — jax’s baked defaults already ARE its tuned config). Runtime-only.
runtime_inertia (bool, optional) – Build the robot with a runtime-mutable inertia table (D.4 / Phase 5, numpy backend only). Default False ⇒ the per-link spatial inertia is baked into the .so (byte-identical to a plain build, same cache key). When True, the codegen emits a d_inertia_params table + on-device 6x6 rebuild + a host mutator, and the handle gains RobotHandle.set_inertia_params() (mutate inertia at runtime, no recompile) and RobotHandle.inertia_params (the baked values to fetch-then-mutate). Re-keys the cache (the runtime-inertia .so coexists with the baked one). With the baked values it reproduces the baked result; mutate to do sysID / domain randomization / payload changes.
use_joint_dynamics (bool, optional) – Model joint-local viscous damping + Coulomb friction in the value paths (inverse_dynamics / forward_dynamics / aba): tau -= damping*qd + friction*sign(qd), per joint, using the damping/friction declared in the URDF. Default False ⇒ the historical no-op build (byte-identical header, same cache key, and consistent with the bare-Pinocchio oracle which ignores damping/friction). When True the bias is emitted ONLY for robots that declare nonzero damping/friction; it re-keys the cache (the damped .so coexists with the baked no-op one). Match against RBDReference(..., use_joint_dynamics=True).
algorithm_list (list[str] | str | None, optional) – Build only a SUBSET of algorithms into the per-robot .so instead of the full default profile. Default None ⇒ the historical full build (every method available; byte-identical header, same cache key, reuses the existing .so). When set (e.g. ["inverse_dynamics", "forward_dynamics"]), only the named algorithms — plus their transitive dependencies, which GRiDCodeGenerator expands automatically (e.g. forward_dynamics_gradient pulls in minv / inverse_dynamics / inverse_dynamics_gradient) — are codegen’d and compiled. This cuts nvcc wall time, peak RAM, and .so size dramatically for big robots with heavy second-order kernels (e.g. fdsva_so on a mid-chain spherical robot is 20+ min / 7 GB). Methods that were NOT built raise a clear runtime error naming the algorithm to add and rebuild — not a segfault. Re-keys the cache (a subset .so coexists with the full build). Recognized names mirror the codegen keys: inverse_dynamics, minv, forward_dynamics, aba, crba, inverse_dynamics_gradient, forward_dynamics_gradient, idsva_so_body_frame, fdsva_so, end_effector_pose``[``_gradient/ _hessian], integrator, integrator_gradient, plus curated profile sets like "dynamics-core". Supported on ALL backends (numpy / jax / torch): the JAX/torch FFI handlers are per-CORE-algo gated, so a subset .so builds only the requested cores on those surfaces too. A method that was NOT built raises the same clean “not built into this robot .so — add to algorithm_list and rebuild” error on jax/torch as it does on numpy.
output_convention (str, optional) – Default IO convention for the returned handle: "pinocchio" (default, GRiD-native) or "mujoco" (mjx parity — wxyz quat, global-linear free-joint velocity). A runtime setting (NOT in the cache key — the .so is identical); it is a byte-identical no-op on a fixed base (mjx≡pinocchio with no free-flyer), so it is accepted on fixed AND floating robots alike for a uniform interface. The value methods (id/fd/aba/crba/minv) and the native-mjx derivative/second-order surfaces (id_gradient / fd_gradient / idsva_so / fdsva_so, floating base) honor it. Equivalent to setting handle.output_convention after registration, or using the per-call thread-safe handle.mujoco view.
enable_mujoco_kernels (bool, optional) – Default True. Set False to build a PIN-ONLY .so: the mjx (MuJoCo-convention) kernel twins are not instantiated and their C-ABI entry points return rc=3 (“not built into this .so”). On a large floating-base non-mimic robot this is the difference between building and running out of memory — the second-order mjx twins are the largest kernels (idsva_so_world_frame was 28x its pin kernel raw; block-parallelizing the epilogue cut it to 2.42x pin on go2-floating), so pin-only builds g1 in ~33 min at ~11 GB peak. Use it if you do not need the MuJoCo output convention (GATO / PDDP second-order DDP, or anything reading Pinocchio-convention derivatives). No-op on fixed-base and mimic robots, which never get mjx twins. On a FLOATING base it is mutually exclusive with output_convention="mujoco" (that needs the twins); on a fixed base output_convention="mujoco" is itself a no-op, so the two compose freely. Codegen-affecting: it participates in the .so cache key only when False, so existing caches stay valid.

Returns:

Ready for forward_dynamics / inverse_dynamics / minv / etc.

Return type:

RobotHandle

One-call convenience: load a URDF and return a ready-to-use handle.

This is the frictionless entry point — no name ceremony, no two-call precompile→get_robot dance. It is a THIN wrapper over register_robot():

Derives a stable, content-addressed name from the URDF bytes (so the caller never types a name, and re-loading the SAME urdf returns the SAME cached robot — no recompile). Override with name= if you want a human-friendly handle.

Calls register_robot (which is itself idempotent on the cache key — a matching .so is reused, no nvcc).

Returns a handle on the requested backend: "numpy" → the pybind RobotHandle, "jax" → a grid_rbd.jax handle, "torch" → a grid_rbd.torch handle.

Parameters:

urdf_path (str | None) – Path to the robot URDF. Mutually exclusive with urdf_string.
backend (str) – "numpy" (default) / "jax" / "torch".
floating_base (bool) – Treat the robot as floating-base (default fixed). Folded into the derived name, so the fixed and floating loads of one URDF get distinct, stable handles (they already compile to distinct cache entries).
urdf_string (str | None) – Inline URDF text instead of a file. Mutually exclusive with urdf_path.
name (str | None) – Override the auto-derived handle name. Default None ⇒ f"{stem}_{floating|fixed}_{sha256(urdf)[:12]}" (collision-resistant: a 48-bit content hash + the base flag).
cache_dir (str | Path | None) – Cache root override (see register_robot()).
**opts – Any other register_robot() keyword (ee_joint_names, max_batch_size, runtime_inertia, runtime_transform, output_convention, algorithm_list, dtype, force_rebuild, cuda_arch, …) — passed straight through. backend-incompatible options raise the same clear error register_robot already gives.

Returns:

Ready for forward_dynamics / inverse_dynamics / minv / etc.

Return type:

RobotHandle | JaxRobotHandle | TorchRobotHandle

The numpy handle#

class grid_rbd.RobotHandle(name: str, so_path: str, meta: dict[str, Any], *, allow_fp64: bool = False)[source]#

Opaque handle to a compiled per-robot GRiD library.

Created by grid_rbd.register_robot(…) and grid_rbd.get_robot(…). Don’t construct directly; the constructor wires up the pybind11 Runner plus the metadata loaded from the cache’s meta.json.

Precision: float32 by default — methods cast inputs to float32 and compute in single precision. A true fp64 tier (Phase 8) is available by registering with dtype="float64": the handle then drives a double- precision .so (_core.RunnerF64), casts inputs to float64, and returns float64 arrays computed end-to-end in double precision (handle.dtype == "float64"). The legacy allow_fp64=True is only an fp32-compute upcast convenience (compute in fp32, cast i/o to fp64, single-precision accuracy caveat); it is off by default and ignored for a true-fp64 handle. The jax/torch handles follow the .so dtype too (dtype="float64" builds carry fp64 jax/torch surfaces; jax additionally needs x64 mode).

Method index (all take/return (B, …) arrays, batch axis first):

dynamics    inverse_dynamics (rnea) · forward_dynamics (fd) · aba ·
            crba · minv · generalized_gravity · nonlinear_effects
gradients   inverse_dynamics_gradient · forward_dynamics_gradient
2nd-order   idsva_so → SecondOrderID · fdsva_so → SecondOrderFD
kinematics  end_effector_pose[_gradient|_hessian] · fk_batched ·
            frame_jacobian[_dot] · com · ccrba · osc_inertia ·
            dccrba · cmm_time_variation
energy      energy · coriolis_matrix ·
            kinetic_energy_regressor · potential_energy_regressor
integration integrator[_gradient]
plant/cost  plant_step[_gradient|_hessian] · quadratic_state_cost ·
            quadratic_input_cost · ee_pos_cost · com_cost ·
            momentum_cost · joint_{position,velocity,torque}_barrier
sysID/param inverse_dynamics_regressor · inertia_params ·
            set_inertia_params (runtime-mutable inertia, no recompile)
runtime-EE  end_effector_pose_runtime[_gradient] (arbitrary target joints)

handle.mujoco is the MuJoCo-convention view of the differentiable methods (mjx free-joint frame, applied per-call; safe alongside pinocchio-convention calls).

Short aliases: rnea → inverse_dynamics(), fd → forward_dynamics() (aba / crba / minv already use their field-standard names).

property name: str#

property num_joints: int#

property num_vel: int#

property num_ees: int#

property nq: int#

7 (pos + quat xyzw) + joints.

Type:: Configuration width
Type:: q is (B, nq). Floating base

property nv: int#

qd/qdd/u INPUTS, dynamics vector OUTPUTS and every matrix/gradient axis are nv wide. Floating base: 6 + joints.

Type:: Tangent/velocity width

property nb: int#

f_ext is (B, 6*nb).

Type:: Body count (incl. the base for floating-base)

property dtype: str#

"float32" (default) or "float64" (Phase 8 true fp64 tier). Inputs are cast to / outputs are returned in this numpy dtype.

Type:: Compute precision of this handle’s .so

property floating_base: bool#

property configuration_layout#: Independent joint blocks mapping public positions to tangent slots.

property output_convention: str#

"pinocchio" (default, GRiD-native — xyzw quat, spatial-local free-joint velocity) or "mujoco" (mjx parity — wxyz quat, global-linear free-joint velocity). Setting "mujoco" makes the value methods take and return MuJoCo-convention q/qd/qdd/M/… for a FLOATING base; it is a byte-identical no-op for a fixed base. See external/RBDReference/equivalents/mujoco_convention.md for the exact transforms. Applies to the VALUE methods (inverse_dynamics, forward_dynamics, aba, crba, minv, …) AND to the gradient / second-order / regressor surfaces (served by the mjx kernel twins when the .so is built with them).

Type:: Output/IO convention

property mujoco: _MujocoView#: MuJoCo-native view (handle.mujoco.inverse_dynamics(qpos, qvel, qacc)): the supported value methods with MuJoCo parameter names and the mjx output convention applied PER CALL — thread-safe and independent of the handle’s output_convention default (it never mutates shared state). Cached.

property runtime_inertia: bool#: True if this robot was registered with runtime_inertia=True (the .so carries a mutable inertia table + set_inertia_params()).

property meta#

sizes, joint names/limits, build facts (algorithm_list / generated_algorithms / dtype / max_batch / cuda_arch / enable_mujoco_kernels), and the launch-config robot key. Mutating the returned dict never affects the handle. Keys added over time; a .so cached by an older grid-rbd simply lacks the newer ones.

Type:: A COPY of the .so’s persisted metadata (meta.json)

property joint_names#: Joint names, index == joint id (the input-vector ordering). None on a .so registered before names were persisted.

capabilities()[source]#

What this .so can do, per python-surface method key: {key: {"built", "mjx", "out_shape", "tier", "threads", "max_threads"}}.

built: the algorithm was in the build’s post-dep-expansion emit set (None on a .so cached before generated_algorithms was persisted — re-register with force_rebuild=True to populate). Best-effort: a robot-CLASS-limited surface (centroidal family on a mimic robot, fk_batched on floating/mimic) can still raise its clean rc=3 error at call time even when requested at build time.
mjx: the MuJoCo-convention twin symbol is REALLY present in the dlopen’d .so (ground truth), for methods that have a twin.
out_shape: trailing per-batch-item output dims (ints).
tier/threads: the tuned launch config baked for this robot/ GPU (None without a config/launch_configs entry).
max_threads: the compiled kernel’s real __launch_bounds__ ceiling via kernel_max_threads() (None where the probe does not apply).
batch_threshold/threads_small: the ARMED E6 batch-switch state (see apply_batch_overlay()) — calls with batch <= batch_threshold launch at threads_small instead of threads. Both None when the switch is not armed for the algo.

property joint_pos_limits#: Per-joint position limits [lower, upper] (index == joint id), from the URDF <limit lower= upper=>. None for a joint with no position limit (fixed / continuous / floating / unspecified); an unbounded side is None. Metadata only — not enforced by any kernel. Returns None on a .so registered before limits were persisted (re-register to populate).

property joint_vel_limits#: Per-joint velocity limits (index == joint id) from the URDF <limit velocity=>; None where unspecified. Metadata only.

property joint_effort_limits#: Per-joint effort (torque) limits (index == joint id) from the URDF <limit effort=>; None where unspecified. Metadata only.

property inertia_params#

The BAKED 10-param-per-body inertia table, shape (num_bodies, 10).

Each row is [m, hx, hy, hz, Ixx, Ixy, Ixz, Iyy, Iyz, Izz] (mass, first moment h = m*c, then the 6 upper-triangle entries of the link’s inertia about its frame origin) in the frozen GRiD/URDF regressor basis — body-indexed, bodies 1..N (the synthetic world-frame link is dropped, mirroring the device table layout). For a FIXED base num_bodies == num_joints; for a FLOATING base the floating trunk is row 0 and num_bodies < num_joints (== num_pos). Fetch this, mutate it, and pass it to set_inertia_params(). Only available on a runtime_inertia build (raises otherwise; the values aren’t persisted for a baked .so).

set_inertia_params(params) → None[source]#

Update the device-resident inertia table at runtime (no recompile).

params is the 10-param-per-body table — either flat (10*num_bodies,) or (num_bodies, 10) — in the same layout / basis as inertia_params (bodies 1..N, each [m, h(3), I_O(6)]). All subsequent algorithm calls (inverse_dynamics, crba, …) reconstruct the per-link spatial inertia from the updated table. The sysID / domain-randomization / payload entry point.

The table is body-indexed by num_bodies (the inertia-body count, == inertia_params rows), NOT num_joints: for a FIXED base they coincide, but for a FLOATING base (or a mimic robot) num_joints == num_pos > num_bodies and the device table is 10*num_bodies long.

Only valid on a robot registered with runtime_inertia=True; raises a clear error otherwise. Passing the baked inertia_params back reproduces the baked result.

attach_tool(joint, *, mass, com=(0.0, 0.0, 0.0), inertia=None, tip_transform=None)[source]#

Weld a rigid tool/payload to the link moved by joint at runtime.

This is the “the robot now has a tool” entry point, and it needs no recompile. It does two things:

Inertia — composes the payload’s spatial inertia (mass [kg], com [m, in the link frame], inertia [3x3 or 6-vector about the payload CoM]) into that link’s spatial inertia and pokes it via set_inertia_params(). The whole dynamics stack (inverse/forward dynamics + gradients) then sees the composite rigid body.
Tip frame — if tip_transform (a 4x4 SE(3) matrix in the joint frame) is given, stores it so subsequent end_effector_pose_runtime() / end_effector_pose_gradient_runtime() calls (with no explicit target) report the SE(3) tool-tip frame. Omit it for a pure payload.

joint is a joint NAME; the tool welds to that joint’s child link, and the tip frame hangs off that joint’s frame — so a tool can be attached ANYWHERE in the chain, not just a leaf. Requires runtime_inertia=True (use register_robot(..., enable_tool=True)). One tool at a time; detach_tool() restores the baked robot.

detach_tool()[source]#: Remove the attached tool: restore the baked inertia for its link and clear the stored tip frame. A no-op if nothing is attached.

property tool#: The currently attached tool dict {joint, row, Xtool} or None.

tool_fext(q, wrench, *, joint=None, offset=None)[source]#

Map a world-aligned tool-tip wrench to a joint-local f_ext array.

wrench is (B, 6) = [n_w; f_w] (moment about the tool tip; world axes) per timestep. Returns (B, 6*num_bodies) joint-local f_ext ([angular;linear] per body) that feeds straight into inverse_dynamics() / forward_dynamics() / aba() as f_ext=... — i.e. the effect of the tool pushing on the world (grinding, pushing, a second gripper finger).

With a tool attached (via attach_tool()) the contact body + tip offset default to that tool; otherwise pass joint (a joint name) and offset (the 3-vector tip point in the joint frame). Needs an enable_tool .so.

property contact_frames#: The registered contact frames [{name, jid, offset}] (registration order == the contact_fext f_c column order) or None.

contact_fext(q, f_c)[source]#

Map per-contact-frame world-aligned wrenches to joint-local f_ext.

f_c is (B, 6*num_contact_frames): per registered frame (in registration order) a 6-vector [n_w; f_w] — WORLD-ALIGNED axes, moment about the contact-frame origin (pinocchio LOCAL_WORLD_ALIGNED; the same wrench convention as tool_fext()). Returns (B, 6*num_bodies) joint-local f_ext ready to pass as f_ext= to inverse_dynamics() / forward_dynamics() / aba() and their gradients — e.g. stance-foot reaction forces on a quadruped/humanoid. Needs a contact_frames=[...] .so.

property runtime_transform: bool#: True if this robot was registered with runtime_transform=True (the .so carries a mutable joint-origin table + set_transform_params()).

property transform_params#

The BAKED 6-param-per-joint origin table, shape (num_joints, 6).

Each row is [x, y, z, roll, pitch, yaw] — the raw URDF <origin> translation + rpy of joint i (joint-indexed, ALL joints 0..NB-1). Fetch this, mutate it, and pass it to set_transform_params(). Only available on a runtime_transform build (raises otherwise).

set_transform_params(params) → None[source]#

Update the device-resident joint-origin table at runtime (no recompile).

params is the 6-param-per-joint table — either flat (6*num_joints,) or (num_joints, 6) — in the same layout / basis as transform_params (joints 0..NB-1, each [x, y, z, roll, pitch, yaw]). All subsequent algorithm calls rebuild each joint’s constant Xfixed from the updated table. The kinematic calibration / domain-randomization entry point for joint frames.

Only valid on a robot registered with runtime_transform=True; raises a clear error otherwise. Passing the baked transform_params back reproduces the baked result.

v1 scope: this mutates the spatial X transforms used by the DYNAMICS (inverse_dynamics, crba, forward_dynamics, Minv, gradients, …). The END-EFFECTOR / homogeneous-transform kinematics (end_effector_pose and its gradient/hessian) still use the BAKED origin and are NOT affected — making them mutable is a planned v2 follow-up.

property runtime_joint_dynamics: bool#: True if registered with runtime_joint_dynamics=True (the mutable damping/friction table backing set_joint_dynamics()).

property joint_damping#: Baked per-v-slot viscous damping (length nv, alpha-folded). Mutate a copy and pass to set_joint_dynamics(). Only on a runtime_joint_dynamics build (raises otherwise).

property joint_friction#: Baked per-v-slot Coulomb friction (length nv, alpha-folded). Only on a runtime_joint_dynamics build (raises otherwise).

set_joint_dynamics(damping=None, friction=None) → None[source]#

Update the device-resident damping/friction table at runtime (no recompile).

Pass either/both as length-nv arrays (v-slot indexed, same basis as joint_damping / joint_friction). An omitted side keeps its current baked value. This is the sysID / domain-randomization entry point for joint dynamics; passing the baked values back reproduces the baked result bit-for-bit. Only on a runtime_joint_dynamics build (raises otherwise).

property max_batch: int#

property max_perf_level_threads: int#

Codegen-time thread-count hint (DOF-aware, warp-rounded).

The default block size for kernel launches. Since v2.0 it is a recommendation, not an enforced floor — callers can override via set_threads_per_block().

property threads_per_block: int#

Active global threads-per-block override.

Returns -1 when no override is set, meaning each algorithm launches at its own autotuned launch_cfg<ALGO>::THREADS baked into grid.cuh (the per-algo default that fixes the FFI thread-default pathology). A value >= 1 is a global override forced via set_threads_per_block() that applies to every algorithm. Do not assume this is a positive block size.

set_threads_per_block(n: int) → None[source]#

Force a single global per-block thread count for all subsequent kernel launches issued through this handle.

The codegen does block-cooperative compute: each block handles one timestep with its threads cooperating via block-stride loops. Batching across timesteps is grid-stride at the block level. Any block size n >= 1 (up to the per-block max, 1024 on current GPUs) is valid; smaller sizes are correct but slower.

By default (no override) each algorithm uses its own autotuned per-algo thread count baked into grid.cuh. Passing n >= 1 overrides that for every algorithm.

kernel_max_threads(algo: str) → int[source]#

Real compiled __launch_bounds__ ceiling of the baked kernel for the short autotune key (id, fd, minv, id_du, ee_pose, idsva_so, …): cudaFuncGetAttributes().maxThreadsPerBlock for grid::<algo>_kernel instantiated at its baked launch_cfg<ALGO>::TIER.

This is the E1 tier-contract read: the FFI autotune uses it to record a self-consistent {tier, threads} instead of guessing a host tier and clamping to it. Returns -1 when the key is unknown/not-built or the .so predates the introspection symbol (the autotune then infers the tier from the swept ceiling — it never crashes on a stale .so).

apply_profile_overlay(profile: str) → int[source]#

Overlay this profile’s per-algo threads (torch_bases / pybind_bases) onto the baked .so at runtime (E6). The torch/pybind surfaces share the baked ffi .so but want different per-algo block sizes; when a profile’s tier matches the baked (ffi) tier the only difference is the block size, settable here with no rebuild.

SKIPS any algo whose profile tier != the baked ffi tier (a tier mismatch needs a profile build, not a runtime overlay) so it never launches a kernel at a tier it wasn’t compiled for. Returns the number of algos overlaid. No-op (returns 0) if the robot has no <profile>_bases.

install_device_pool(alloc_fn) → int[source]#

Framework-allocator integration (2026-09-09): carve GRiD’s entire gridData device arena out of a slab the embedding framework allocates (jax: an XLA-pool jnp buffer; torch: a caching-allocator tensor) instead of raw cudaMalloc — so framework preallocation (XLA’s 75%) and GRiD’s allocations stop fighting (the h1_2 “launch failed” class).

alloc_fn(nbytes) -> (buffer, device_ptr) allocates nbytes on the CUDA device via the framework and returns the keep-alive object plus its raw device address. Must run BEFORE the first kernel call (the arena init is lazy). Honors GRID_WORKSPACE_TIMESTEP_SLOTS; when the framework cannot fit the full-slot slab the slot count is halved down to 1 (mirroring the auto-fit) before giving up. Returns the installed slab size in bytes, or 0 when the pool stays off (the cudaMalloc path with its own VRAM auto-fit remains the fallback). The slab is held by the shared native owner until its last Runner closes. A live arena is never reset to install a pool; late framework views keep using that arena without erasing model updates.

apply_batch_overlay(profile: str = 'ffi') → int[source]#

Arm the E6 batch-switch from this profile’s <profile>_bases_by_n block: for each algo with a small-batch winner recorded there, calls with batch <= the threshold launch at that winner’s thread count instead of the large-batch default. The pick is a stateless per-call threshold compare in the .so (no hysteresis; CUDA-graph capture sees one stable pick since a graph freezes its batch).

Bucket keys are the sweep batch sizes (today just "16"); the threshold is <profile>_by_n_threshold when present, else the geometric midpoint of the smallest bucket and the bake batch (16/256 -> 64). Same safety rules as apply_profile_overlay: enum-index derivation from the descriptor table, algo_count drift refusal, and SKIP on any entry whose tier differs from the baked ffi tier. Returns the number of algos armed; 0 if no by-n block exists.

property num_bodies: int#: Number of bodies/links (incl. the base for floating-base). The external-force array f_ext is shaped (B, 6*num_bodies).

inverse_dynamics(q, qd, qdd=None, *, gravity: float = -9.81, f_ext=None, _convention=None) → ndarray[source]#

Inverse dynamics (RNEA): τ = M(q)·qdd + h(q,qd) − g(q). q is (B, nq), qd and the optional qdd are (B, nv). Returns (B, nv).

With qdd=None (default) this is the bias c = h(q,qd) − g(q) (= RBDReference.inverse_dynamics(q, qd, qdd=0)). Pass a nonzero qdd to get the full RNEA torque including the inertial term M·qdd — the acceleration is now plumbed through (USE_QDD_FLAG=true).

f_ext (optional): per-body external forces, shape (B, 6*num_bodies), body-major, each [angular; linear] in the body’s local frame (subtracted from the per-body force, matching RBDReference.inverse_dynamics(..., f_ext=...)). Default None ⇒ no external force.

With output_convention="mujoco" (floating base) q/qd/qdd are MuJoCo-convention and the returned τ is in the mjx frame.

minv(q, *, _convention=None) → ndarray[source]#

Direct mass-matrix inverse Minv(q). Returns shape (B, NV, NV).

Minv is the tangent-space (pinocchio-convention) inverse mass matrix: NV x NV. For a FIXED base NV == NJ (== num_pos) so the shape is unchanged; for a FLOATING base NV = 6 + n_joints < NJ = 7 + n_joints (the +1 is the quaternion offset in q only). GRiD’s minv kernel writes only the upper triangle (lower zero); we symmetrize on the host before returning so the matrix matches RBDReference.minv(…, output_dense=True). The symmetrization is a single numpy op per call — negligible cost.

With output_convention="mujoco" (floating base) q is MuJoCo-convention and the returned Minv is the mjx-frame inverse mass matrix.

forward_dynamics(q, qd, u, *, gravity: float = -9.81, f_ext=None, _convention=None) → ndarray[source]#

Forward dynamics qdd = M⁻¹·(τ − c). q is (B, nq), qd/u (B, nv). Returns (B, nv).

f_ext (optional): per-body external forces (B, 6*num_bodies), body-major, [angular; linear] local-frame (see inverse_dynamics()). With output_convention="mujoco" (floating base) q/qd/u are MuJoCo-convention and the returned qdd is in the mjx frame.

aba(q, qd, u, *, gravity: float = -9.81, f_ext=None, _convention=None) → ndarray[source]#

Recursive forward dynamics via Articulated Body Algorithm. Returns shape (B, nv). Alternative to forward_dynamics() with the same output but a different implementation.

f_ext (optional): per-body external forces (B, 6*num_bodies), body-major, [angular; linear] local-frame (see inverse_dynamics()). With output_convention="mujoco" (floating base) the IO is mjx-convention.

crba(q, *, gravity: float = -9.81, _convention=None) → ndarray[source]#

Joint-space mass matrix M(q) via Composite Rigid Body Algorithm. Returns shape (B, NV, NV) — the tangent-space (pinocchio-convention) mass matrix. FIXED base: NV == NJ (unchanged); FLOATING base: NV < NJ (the kernel writes NUM_VEL x NUM_VEL). Pass gravity only because the host wrapper takes it; the result doesn’t depend on gravity.

With output_convention="mujoco" (floating base) q is MuJoCo-convention and the returned M is the mjx-frame mass matrix.

end_effector_pose(q, *, _convention=None) → ndarray[source]#

End-effector pose [xyz, rpy] per EE. Returns shape (B, 6*NUM_EES). For multi-EE robots, reshape to (B, NUM_EES, 6) at the caller side.

With output_convention="mujoco" (floating base) q is MuJoCo-convention; the pose itself is frame-INVARIANT, but the native kernel routes q through the mjx quaternion reorder (so the output equals feeding the pin kernel the pin-converted q).

fk_batched(q, *, use_warp: bool = False) → ndarray[source]#

Large-batch forward kinematics, one block (thread variant) or warp (warp variant) per sample.

Input q: (B, NUM_POS) joint positions (batch-major). Output pose7: (B, 7) = [tx, ty, tz, qw, qx, qy, qz] for the leaf EE frame, where the last four are the unit quaternion (w, x, y, z).

use_warp=True runs the warp-cooperative per-sample inner; both variants return identical poses. Fixed AND floating base supported (floating q = [x,y,z, qx,qy,qz,qw, joints], pin convention); mimic joints fold inline. Raises only for spherical-joint or >32-joint robots (those route through end_effector_pose).

end_effector_pose_gradient(q, *, _convention=None) → ndarray[source]#

End-effector pose Jacobian d/dv (TANGENT, pinocchio convention).

Returns shape (B, 6*NUM_EES, NV). Floating-base produces the spatial Jacobian (omega; v) base block, not the older non-standard quaternion-derivative columns. Fixed-base shape unchanged (NV == NJ).

GRiD’s h_end_effector_pose_gradient is stored column-major as (6, NUM_EES*NV) per timestep; we re-orient to (6*NUM_EES, NV) per timestep.

With output_convention="mujoco" (floating base) q is MuJoCo-convention and the returned Jacobian has its base-linear columns reframed into the mjx frame (computed natively in the kernel).

inverse_dynamics_gradient(q, qd, qdd=None, *, gravity: float = -9.81, f_ext=None, out=None, _convention=None) → ndarray[source]#

∂τ/∂(q, qd). Returns shape (B, NV, 2*NV) — concatenated [dc_dq | dc_dqd], tangent-space (pinocchio) convention. Slice with […, :NV] / […, NV:]. FIXED base: NV == NJ (unchanged); FLOATING base: NV < NJ (the kernel writes nv x 2nv).

qdd (optional): joint acceleration. The gradient depends on it (through the M·qdd term); qdd=None (default) ⇒ the bias gradient at qdd=0. Now plumbed through (USE_QDD_FLAG=true) matching inverse_dynamics().

f_ext (optional): per-body external forces (B, 6*num_bodies). f_ext enters RNEA affinely, so for a CONSTANT f_ext the Jacobian ∂c/∂(q,qd) is unchanged; the kwarg is for consistency with inverse_dynamics().

out (optional): a caller-owned (B, 2*NV*NV) array in the compute dtype, C-contiguous and writeable — ideally from pinned_empty() — that the device result is downloaded into directly. The returned (B, NV, 2*NV) array is then a VIEW of out (column-major per item, so not C-contiguous; np.ascontiguousarray it if a consumer needs that). Allocate once, reuse every call: no per-call allocation or host re-layout.

With output_convention="mujoco" (floating base) q/qd/qdd are MuJoCo-convention and the returned gradient is in the mjx frame (the full convention transform — reframe + base-row rotate + ω×v couplings — is baked into the kernel). Requires an explicit qdd (and no f_ext).

forward_dynamics_gradient(q, qd, u, *, gravity: float = -9.81, f_ext=None, out=None, _convention=None) → ndarray[source]#

∂qdd/∂(q, qd). Returns shape (B, NV, 2*NV), tangent-space (pinocchio) convention. FIXED base: NV == NJ (unchanged); FLOATING base: NV < NJ.

f_ext (optional): per-body external forces (B, 6*num_bodies); affine in f_ext so a constant f_ext leaves this Jacobian unchanged.

out (optional): caller-owned (B, 2*NV*NV) buffer, as in inverse_dynamics_gradient() — the result is a view of it.

With output_convention="mujoco" (floating base) q/qd/u are MuJoCo-convention and the returned gradient is in the mjx frame (the full convention transform — reframe + base-row rotate + ω×v couplings — is baked into the kernel). Requires no f_ext.

end_effector_pose_hessian(q, *, _convention=None) → ndarray[source]#

End-effector pose Hessian ∂²(pose)/∂v² (tangent-space, pinocchio convention). Returns shape (B, 6*NUM_EES, NV, NV). For fixed-base NV == NJ; for floating-base the (NV, NV) block indexes spatial twist components.

With output_convention="mujoco" (floating base) q is MuJoCo-convention and the returned Hessian is the symmetric coordinate Hessian along the mjx retract (double column-reframe + symmetrized base-rotation frame term, baked in-kernel).

pinned_empty(shape, dtype=None) → ndarray[source]#: A page-locked (cudaMallocHost) numpy array of shape in this robot’s compute dtype, to be reused as out= by the gradient methods, idsva_so() / fdsva_so(). Allocate ONCE, reuse every call: the device->host copy then runs at the PCIe rate and no host-side copy follows (g1 idsva_so at batch 1024, 702 MB: ~40 ms vs 120 ms through a fresh pageable array, measured 2026-10-01). The array owns its memory and keeps the robot library loaded until it is collected. Page-locked memory is a limited resource: do not allocate per call.

is_pinned(arr) → bool[source]#: True when arr’s buffer is page-locked (as returned by pinned_empty()).

idsva_so(q, qd, qdd=None, *, gravity: float = -9.81, out=None, _convention=None) → SecondOrderID[source]#

Second-order inverse dynamics. Returns a SecondOrderID NamedTuple (d2tau_dq, d2tau_dqd, d2tau_cross, dM_dq), each tensor shape (B, NV, NV, NV). (NamedTuple is a plain tuple — positional unpacking and indexing still work.)

Uses the codegen-time dispatcher: body-frame for fixed-base, world-frame for floating-base.

out (optional): a caller-owned (B, 4*NV**3) array in the compute dtype, C-contiguous and writeable — ideally from pinned_empty() — that receives the result directly; the returned tensors are views of it. With allow_fp64 the returned tensors are float64 upcast copies, but out is still filled.

With output_convention="mujoco" (floating base) q/qd/qdd are MuJoCo-convention and all four 2nd-order tensors are returned in the mjx frame (the explicit-analytic SO transform + dM_dq closed form, baked in-kernel).

inverse_dynamics_regressor(q, qd, qdd=None, *, gravity: float = -9.81, _convention=None) → ndarray[source]#

Inverse-dynamics inertial-parameter regressor Y with tau = Y . pi (pi = the stacked 10-param spatial inertia of each link). Returns (B, NV, 10*NUM_BODIES). Mirrors RBDReference.inverse_dynamics_regressor(q, qd, qdd).

With output_convention="mujoco" (floating base) q/qd/qdd are MuJoCo-convention and the base-linear ROWS (0:3) of Y are rotated by R in-kernel (the rows transform like a generalized force, Y_mjx = G Y_pin).

fdsva_so(q, qd, u, *, gravity: float = -9.81, out=None, _convention=None) → SecondOrderFD[source]#

Second-order forward dynamics. Returns a SecondOrderFD NamedTuple of 4 tensors each shape (B, NV, NV, NV) (a plain tuple, so positional unpacking / indexing still work). out: see idsva_so() (a (B, 4*NV**3) caller-owned buffer, e.g. from pinned_empty()).

With output_convention="mujoco" (floating base) q/qd/u are MuJoCo-convention and all four 2nd-order tensors are returned in the mjx frame (explicit-analytic SO transform, contravector output-map, baked in-kernel).

integrator(q, qd, u, dt, *, integrator_type: str = 'euler', gravity: float = -9.81, _convention=None)[source]#

One integration step x_{k+1} = integrator(x_k, u, dt).

Returns shape (B, NUM_POS + NUM_VEL) — concatenated [q_new, v_new]. dt is the runtime timestep; gravity is the signed gravitational acceleration (default -9.81). integrator_type is one of euler / semi_implicit_euler / constant_acceleration / trapezoidal (Heun) / midpoint / rk4.

With output_convention="mujoco" (floating base) the free-joint base position takes a GLOBAL additive step (the MuJoCo retract) rather than pinocchio’s SE(3) update; q/qd are MuJoCo-convention and the returned q_new is in the mjx frame (quaternion wxyz). Baked into the kernel.

integrator_gradient(q, qd, u, dt, *, integrator_type: str = 'euler', gravity: float = -9.81, _convention=None)[source]#

Gradient of the integrator step. Returns shape (B, 2*NV, 3*NV) — column blocks [d/dq | d/dqd | d/du] in tangent space.

dt is the runtime timestep; gravity is the signed gravitational acceleration (default -9.81).

With output_convention="mujoco" (floating base, EULER/SI-EULER) q/qd/u are MuJoCo-convention and the returned state-transition Jacobian is in the mjx tangent (global-add retract rows + G velocity reframe + input couplings, in-kernel).

quadratic_state_cost(x, x_des, Q, *, _convention=None)[source]#

1/2 * sum_i Q_i (x_i - x_des_i)^2 over the full state x = [q; qd].

x / x_des / Q are (B, NUM_POS + NUM_VEL). Returns:: value (B,), grad (B, NX), hess = diag(Q) (B, NX, NX).

With output_convention="mujoco" (floating base) x is MuJoCo-convention: the kernel input-converts the velocity base-linear block (global->local) before differencing against the mjx-frame x_des/Q, so the VALUE is convention-DEPENDENT; the qd-block grad/hess are reframed in-kernel.

quadratic_input_cost(u, u_des, R)[source]#

1/2 * sum_i R_i (u_i - u_des_i)^2 over the input u (size NUM_VEL).

u / u_des / R are (B, NUM_VEL). Returns:: value (B,), grad (B, NV), hess = diag(R) (B, NV, NV).

ee_pos_cost(q, p_des, W, *, _convention=None)[source]#

End-effector position cost over the 3 position axes (EE 0).

q is (B, NUM_POS); p_des / W are (B, 3). Returns:: value (B,), grad_x (B, NX) = [J_p^T (W·r); 0], GN hess_x (B, NX, NX) with the top-left NV×NV q-block = J_p^T diag(W) J_p.

The hessian is returned in the kernel’s column-major layout; since the GN hessian J_p^T W J_p is symmetric the row/col-major distinction is immaterial.

With output_convention="mujoco" (floating base) q is MuJoCo-convention; the EE position (and thus the value) is invariant, and the q-block grad/hess are reframed to the mjx tangent (covector / congruence) in-kernel.

joint_position_barrier(var, lower, upper, mu)[source]#

Log-barrier b = -mu·(log(x-lo)+log(hi-x)) over NUM_POS positions.

var / lower / upper are (B, NUM_POS); an ±inf bound contributes zero. Returns (value (B,), grad (B, NUM_POS), hess_diag (B, NUM_POS)).

joint_velocity_barrier(var, lower, upper, mu)[source]#: Log-barrier over NUM_VEL velocities. See joint_position_barrier.

joint_torque_barrier(var, lower, upper, mu)[source]#: Log-barrier over NUM_VEL torques. See joint_position_barrier.

plant_step(x, u, dt, *, integrator_type: str = 'euler', gravity: float = -9.81, _convention=None)[source]#

x_{k+1} = integrator(x_k, u_k, dt). Thin wrapper over grid::integrator.

x is (B, NUM_POS + NUM_VEL); u is (B, NUM_VEL). Returns (B, NX). integrator_type is one of euler / semi_implicit_euler / constant_acceleration / trapezoidal (Heun) / midpoint / rk4 (same codes as integrator()).

With output_convention="mujoco" (floating base, EULER/SI-EULER) x/u are MuJoCo-convention and the returned next state uses the mjx global-add retract (base position) + reordered quaternion.

plant_step_gradient(x, u, dt, *, integrator_type: str = 'euler', gravity: float = -9.81, _convention=None)[source]#

[A | B] = d x_{k+1}/d(x,u) = the integrator-gradient s_dAB surface.

x is (B, NUM_POS + NUM_VEL); u is (B, NUM_VEL). Returns (B, 2*NV, 3*NV) with column blocks [d/dq | d/dqd | d/du] in tangent space. Pass-through to grid::integrator_gradient (the value is byte-identical to integrator_gradient()). Matches RBDReference.plant_step_gradient (= integrator_gradient). integrator_type is one of euler / semi_implicit_euler / constant_acceleration / trapezoidal / midpoint / rk4.

com_cost(q, p_des, W, *, _convention=None)[source]#

Center-of-mass tracking cost over the 3 CoM axes.

q is (B, NUM_POS); p_des / W are (B, 3). Returns:: value (B,), grad_x (B, NX) = [J_com^T (W·r); 0], GN hess_x (B, NX, NX) with the top-left NV×NV q-block = J_com^T diag(W) J_com.

Matches RBDReference.com_cost(q, p_des, W).

With output_convention="mujoco" (floating base) the CoM position (and value) is invariant; the q-block grad/hess are reframed to the mjx tangent in-kernel.

momentum_cost(q, qd, h_des, W, *, _convention=None)[source]#

Centroidal-momentum tracking cost over the 6 momentum components.

q is (B, NUM_POS); qd is (B, NUM_VEL); h_des / W are (B, 6). Returns:: value (B,), grad (B, 2*NV) = Jᵀ (W·r), GN hess (B, 2*NV, 2*NV) = Jᵀ diag(W) J in tangent [dq | dv] order with J = [(∂A/∂q)·qd | A] — configuration and cross blocks included (the kernel evaluates dccrba once). An exact cost Hessian is not implied.

Matches RBDReference.momentum_cost(q, qd, h_des, W).

With output_convention="mujoco" (floating base) the centroidal momentum h (and value) is invariant; the derivatives are pulled back through the full input-state Jacobian in-kernel.

com(q, *, _convention=None)[source]#

Center-of-mass world position p_com (3,) and CoM Jacobian J_com.

Returns (p_com, J_com) where p_com is (B, 3) and J_com is (B, 3, NV) = d(p_com)/dv. Matches RBDReference.com(q) (= p_com) and RBDReference.jacobian_com(q) (= J_com).

With output_convention="mujoco" (floating base) q is MuJoCo-convention; p_com is invariant and the J_com columns are reframed (computed natively in the kernel).

ccrba(q, qd, *, _convention=None)[source]#

Centroidal momentum matrix A (6 x NV) and momentum h = A·qd (6,).

Returns (A, h) where A is (B, 6, NV) and h is (B, 6), in the Pinocchio convention ([linear; angular] at the CoM, world aligned). Matches RBDReference.ccrba(q, qd).

With output_convention="mujoco" (floating base) q/qd are MuJoCo-convention; h is invariant and the A columns are reframed (computed natively in the kernel).

energy(q, qd, *, gravity: float = -9.81, _convention=None)[source]#

Kinetic / potential / mechanical energy. Returns (B, 3) = [KE, PE, KE+PE]. Matches RBDReference.kinetic_energy / potential_energy / mechanical_energy (PE uses gravity).

With output_convention="mujoco" (floating base) q/qd are MuJoCo-convention; the energies are frame-INVARIANT, so the native kernel only converts the inputs and the output equals the pin result.

generalized_gravity(q, *, gravity: float = -9.81, _convention=None)[source]#

Generalized gravity torque g(q) = RNEA(q, 0, 0). Returns (B, NV). Matches RBDReference.generalized_gravity(q, GRAVITY=gravity).

With output_convention="mujoco" (floating base) q is MuJoCo-convention and the returned g is in the mjx frame (the base rows are rotated in-kernel). Output shape is invariant (B, NV).

nonlinear_effects(q, qd, *, gravity: float = -9.81, _convention=None)[source]#

Nonlinear (bias) effects c(q,qd) = RNEA(q, qd, 0) = C(q,qd)·qd + g(q). Returns (B, NV). Matches RBDReference.nonlinear_effects(q, qd, GRAVITY=gravity).

With output_convention="mujoco" (floating base) q/qd are MuJoCo-convention and the returned bias matches MuJoCo’s qfrc_bias (the floating-root accel-couple −ω×v is injected and the base rows rotated in-kernel). Output shape is invariant (B, NV).

coriolis_matrix(q, qd, *, gravity: float = -9.81, _convention=None)[source]#

Coriolis matrix C(q,qd). Returns (B, NV, NV) row-major, with C·qd + g(q) = nonlinear_effects(q, qd). Matches RBDReference.coriolis_matrix(q, qd) (gravity is unused by C; the kwarg mirrors the host wrapper signature).

With output_convention="mujoco" (floating base) q/qd are MuJoCo-convention and the returned C is the mjx-frame Coriolis matrix (congruence G C Gᵀ, computed natively in the kernel).

kinetic_energy_regressor(q, qd, *, gravity: float = -9.81, _convention=None)[source]#

Kinetic-energy regressor y_KE, length 10*num_bodies, with KE = y_KE · π (π = stacked per-link inertial parameters, body-major, 10 params/body). Returns (B, 10*num_bodies). Matches RBDReference.kinetic_energy_regressor(q, qd) (gravity unused).

With output_convention="mujoco" (floating base) q/qd are MuJoCo-convention; the regressor is frame-INVARIANT, so the native kernel only converts the inputs and the output equals the pin result.

potential_energy_regressor(q, *, gravity: float = -9.81, _convention=None)[source]#

Potential-energy regressor y_PE, length 10*num_bodies, with PE = y_PE · π. Returns (B, 10*num_bodies). Matches RBDReference.potential_energy_regressor(q, GRAVITY=gravity) (PE uses gravity; only the mass + first-moment columns are nonzero).

With output_convention="mujoco" (floating base) q is MuJoCo-convention; the regressor is frame-INVARIANT, so the native kernel only converts the input and the output equals the pin result.

dccrba(q, *, _convention=None)[source]#

dCCRBA tensor ∂A/∂q, shape (B, 6, NV, NV) indexed [:, :, k, i] = ∂A[:, k]/∂q_i (Pinocchio centroidal convention, [linear; angular] at the CoM, world-aligned). Matches RBDReference.dccrba(q) (which returns (6, NV, NV) per sample).

Runs on all robots — including mimic and big floating-base (the centroidal sweep pool spills to global memory at the spilled tiers). A clear RuntimeError is raised only in the rare case where even the most-spilled tier’s centroidal pool exceeds this GPU’s shared-memory cap.

With output_convention="mujoco" (floating base) q is MuJoCo-convention and the returned tensor is the dA/dq of the mjx CMM (double G^{-1} reframe of the qd-column and q-tangent indices + base-rotation frame term, baked in-kernel). Same (B, 6, NV, NV) layout.

cmm_time_variation(q, qd, *, _convention=None)[source]#

Centroidal-momentum-matrix time variation Ȧ = dA(q(t))/dt, shape (B, 6, NV) (Pinocchio convention, [linear; angular] at the CoM, world-aligned) = Σ_i (∂A/∂q_i)·qd_i. Matches RBDReference.cmm_time_variation(q, qd).

Runs on all robots (mimic + big floating-base via centroidal-pool spill); raises a clear RuntimeError only on the rare oversized-pool case (see dccrba()).

With output_convention="mujoco" (floating base) q/qd are MuJoCo-convention and the returned Ȧ has its columns reframed into the mjx frame (computed natively in the kernel).

frame_jacobian(q, *, target_jid=None, reference_frame=None, _convention=None)[source]#

Geometric Jacobian (6 x NV, [linear; angular]) of a frame. Returns (B, 6, NV). Matches RBDReference.frame_jacobian(q, frame_name, reference_frame).

target_jid selects the frame’s joint id (default: the leaf end-effector joint baked at codegen time). reference_frame is 'LOCAL' (0), 'WORLD' (1), or 'LOCAL_WORLD_ALIGNED' (2, the default), or the equivalent int. Both are now RUNTIME parameters of the GPU surface.

With output_convention="mujoco" (floating base) q is MuJoCo-convention and the returned Jacobian is column-reframed J G⁻¹ (mjx frame).

frame_jacobian_dot(q, qd, *, target_jid=None, reference_frame=None, _convention=None)[source]#

Time derivative Jdot of frame_jacobian() along v = qd (6 x NV, [linear; angular]). Returns (B, 6, NV). Matches RBDReference.frame_jacobian_dot(q, qd, frame_name, reference_frame).

target_jid / reference_frame are RUNTIME parameters (default: leaf-EE joint / LOCAL_WORLD_ALIGNED); see frame_jacobian().

With output_convention="mujoco" (floating base) q/qd are MuJoCo-convention and Jdot is column-reframed (mjx frame).

osc_inertia(q, *, _convention=None)[source]#

Operational-space (task) inertia Lambda = (J·M⁻¹·Jᵀ)⁻¹ (6 x 6) for the leaf-EE frame (LWA). Returns (B, 6, 6). Matches RBDReference.osc_inertia(q).

Lambda is frame-INVARIANT, but with output_convention="mujoco" the MuJoCo q (wxyz quaternion) must be reordered before the kinematics build — the mjx kernel does that, so mjx q routes through it.

end_effector_pose_runtime(q, ee_joint_names=None, ee_offsets=None, *, _convention=None)[source]#

Runtime-target end-effector pose [xyz; rpy] at an offset point.

Mirrors RBDReference.end_effector_pose(q, ee_joint_names, ee_offsets): ee_joint_names (None => all leaf joints, or a str / list of joint names) selects the EE frames, ee_offsets (None => frame origin, or one [x,y,z] / [x,y,z,1] per EE) shifts the measurement point. The single-target GPU kernel is looped over the resolved jid list and the results stacked.

Returns (B, NUM_EE, 6) where each row is [x, y, z, roll, pitch, yaw].

With output_convention="mujoco" (floating base) q is MuJoCo-convention (quat reordered in-kernel); the pose VALUE is frame-invariant.

end_effector_pose_gradient_runtime(q, ee_joint_names=None, ee_offsets=None, *, _convention=None)[source]#

Runtime-target end-effector pose gradient d[xyz; rpy]/dv (6 x NV) at an offset point. Same ee_joint_names / ee_offsets semantics as end_effector_pose_runtime(); mirrors RBDReference.end_effector_pose_gradient(q, ee_joint_names, ee_offsets).

Returns (B, NUM_EE, 6, NV).

With output_convention="mujoco" (floating base) q is MuJoCo-convention; the base-linear columns are reframed by R^T in-kernel (column-reframe class).

rnea(q, qd, qdd=None, *, gravity: float = -9.81, f_ext=None, _convention=None) → ndarray#

Inverse dynamics (RNEA): τ = M(q)·qdd + h(q,qd) − g(q). q is (B, nq), qd and the optional qdd are (B, nv). Returns (B, nv).

With qdd=None (default) this is the bias c = h(q,qd) − g(q) (= RBDReference.inverse_dynamics(q, qd, qdd=0)). Pass a nonzero qdd to get the full RNEA torque including the inertial term M·qdd — the acceleration is now plumbed through (USE_QDD_FLAG=true).

f_ext (optional): per-body external forces, shape (B, 6*num_bodies), body-major, each [angular; linear] in the body’s local frame (subtracted from the per-body force, matching RBDReference.inverse_dynamics(..., f_ext=...)). Default None ⇒ no external force.

With output_convention="mujoco" (floating base) q/qd/qdd are MuJoCo-convention and the returned τ is in the mjx frame.

fd(q, qd, u, *, gravity: float = -9.81, f_ext=None, _convention=None) → ndarray#

Forward dynamics qdd = M⁻¹·(τ − c). q is (B, nq), qd/u (B, nv). Returns (B, nv).

f_ext (optional): per-body external forces (B, 6*num_bodies), body-major, [angular; linear] local-frame (see inverse_dynamics()). With output_convention="mujoco" (floating base) q/qd/u are MuJoCo-convention and the returned qdd is in the mjx frame.

close() → None[source]#: Release the underlying .so handle (and, for a handle made by context(), close its runtime context first: no new admissions, drain, free). After close(), method calls will fail. Idempotent.

property ctx_id: int#: The runtime-context id this handle dispatches to. 0 = the artifact’s DEFAULT context (created lazily on the first call, shared by every handle over this .so that did not ask for its own); a handle made by context() carries its own salted id.

context(*, workspace_slots: int = 0) → RobotHandle[source]#: A NEW runtime context on the same compiled artifact: its own arena (allocator pool, streams, robot tables, plant staging, launch overrides), independent of the default context and of every other context. Returns a handle bound to it; closing that handle closes the context. workspace_slots caps the per-block workspace slot count (0 = the auto-fit). Contexts are the unit of isolation for concurrent pipelines on the single GPU; they are NOT a multi-GPU mechanism.

property model_version: int#

starts at 1 and increments on every runtime-parameter mutation — set_inertia_params(), set_transform_params(), set_joint_dynamics(), hence attach_tool() / detach_tool(). Launch overrides (set_threads_per_block() & co.) do not bump it. A mutation is ordered after every admitted call and before every later one (exclusive admission). The torch / JAX autograd forwards stamp the version on device at execution time and their backwards refuse to run against a different version (“model mutated between forward and backward”) — recompute the forward after mutating.

Type:: This handle’s context MODEL VERSION (W04-B B2)

property device_profile: dict#: The device-profile record captured when this handle’s context was created: device and artifact compute capability, total/free device memory at init, the arena bytes and workspace slot count that were fitted, max_batch, the opt-in shared-memory cap, and whether the arena was carved from a caller-owned slab. Creates the default context if this handle uses it and it does not exist yet.

Step Hessian#

grid_rbd.RobotHandle.plant_step_hessian(x, u, dt, *, integrator_type='euler', gravity=-9.81)#: Return the explicit step Hessian with shape (B, 2*NV, 3*NV, 3*NV) for x of shape (B, NQ+NV) and u of shape (B, NV). Input derivative axes are [dq; dqd; du] and output rows are position tangent followed by velocity. Euler and semi-implicit Euler are supported on fixed and floating bases; multi-stage RK Hessians are unsupported. This method is available on the NumPy handle, not the JAX/PyTorch handles. The generated artifact must include its fdsva_so dependency. See Integrators and the plant layer.

JAX and torch handles#

JaxRobotHandle (grid_rbd.jax) and TorchRobotHandle (grid_rbd.torch) expose framework-native value, explicit derivative, integrator and selected plant methods. They are not exact mirrors: fk_batched and plant_step_hessian are NumPy-only; capture is PyTorch-only. Supported differentiable operations use custom_vjp / autograd.Function; explicit Hessian outputs do not imply support for nested higher-order autodiff. Their modules import jax / torch at load time, so they are documented in the guided tour (Python Wrappers (grid-rbd)) rather than autodoc’d here; the docstrings in bindings/grid_rbd/jax/__init__.py and bindings/grid_rbd/torch/__init__.py are the authoritative per-method reference.