grid_rbd (Python bindings)#

The grid-rbd package is the surface most users 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', dtype: str = 'float32', runtime_inertia: bool = False, runtime_transform: bool = False, runtime_joint_dynamics: bool = False, use_joint_dynamics: bool = False, friction_model: str = 'sign', friction_width=None, enable_tool: bool = False, contact_frames: list[str] | tuple[str, ...] | None = None, cost_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. The validated path is fp32: the .so computes in float on the GPU, wrapper inputs of any float dtype are cast to float32 and outputs are float32. dtype="float64" selects the fp64 tier, an alpha feature: it builds a separate .so (-DGRID_WRAPPER_T_DOUBLE plus the codegen knob that re-derives the shared-memory spill tiers for 8-byte values) whose handle takes and returns float64 arrays computed end to end in double on the numpy, jax and torch backends (framework arrays must then be float64). The fp64 tier is not validated by the signed receipt (its tests run in a separate lane, make test-fp64) and is not guaranteed to build or fit on every robot: double arenas halve the headroom under the per-block shared-memory cap, so big-robot kernels that already max-spill at fp32 may be left unregistered and raise a clear runtime error when called. Selecting it emits one UserWarning per process. Full fp64 support is on the roadmap. The fp32 and fp64 .so coexist in the cache (dtype is in the cache key).

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, friction_model and friction_width (only when non-default), 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, _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, the validated path) or "float64" (the alpha fp64 tier; see the precision paragraph above). Re-keys the cache; the fp32 and fp64 .so coexist.
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 grid_assets/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).
friction_model (str, optional) – The velocity shape of the Coulomb term under use_joint_dynamics. "sign" (default) applies friction*sign(qd), whose derivatives are zero almost everywhere. "tanh" applies the smooth friction*tanh(qd/w): the inverse-dynamics gradient gains (friction/w)*sech^2(qd/w) on the velocity diagonal, the second-order kernels gain -(2*friction/w^2)*tanh(qd/w)*sech^2(qd/w) on the qd-qd triple diagonal, and the forward-dynamics derivatives follow by composition. Re-keys the cache only when "tanh". Match against RBDReference(..., use_joint_dynamics=True, friction_model="tanh", friction_width=w).
cost_frames (list[str], optional) – Named joints or fixed joints baked as cost frames for the plant menu’s frame position, orientation and twist terms (grid_plant::frame_eval and tracking_cost_menu); the end-effector target is always frame 0. Pulls the frame Jacobian family into the build and re-keys the cache.
friction_width (float | sequence of float | None, optional) – The tanh width w in rad/s (m/s for prismatic joints): one scalar for every velocity slot or one value per velocity slot (length nv). None (default) is 0.03. Baked as a literal; on a runtime_joint_dynamics build it is the third row of the mutable table and RobotHandle.set_joint_dynamics() can change it. Requires friction_model="tanh".
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

grid_rbd.get_robot(name: str, cache_dir: str | Path | None = None, *, backend: str = 'numpy', output_convention: str = 'pinocchio', urdf_path: str | None = None, _profile_overlay: str | None = 'pybind', **register_opts: Any) → RobotHandle[source]#

Look up a previously-registered robot by name.

Raises RobotNotRegisteredError if name isn’t in the manifest — unless urdf_path is given: then a miss compiles and registers that URDF under name (just-in-time; **register_opts such as floating_base=True are forwarded to register_robot()) and a hit ignores it. A name alone is never enough to build from: the URDF bytes are the identity. backend (‘numpy’ / ‘jax’ / ‘torch’) mirrors register_robot(): it returns the matching backend handle for the same cached .so (the cache is shared across backends). Default ‘numpy’ keeps the historical return type. output_convention (‘pinocchio’ or ‘mujoco’) is a runtime IO setting mirroring register_robot(); it can also be set later via handle.output_convention.

Cache tools#

grid_rbd.precompile(name: str, urdf_path: str | None = None, *, urdf_string: str | None = None, tiers: Iterable[dict[str, Any]] | None = None, floating_base: bool = False, ee_joint_names: list[str] | tuple[str, ...] | None = None, max_batch_size: int = 256, backends: Iterable[str] = ('numpy',), cache_dir: str | Path | None = None, cuda_arch: int | None = None) → list[dict[str, Any]][source]#

Ahead-of-time: build + populate the persistent cache for a robot so every later register_robot / get_robot / jax.jit is an instant cache hit.

Each entry in tiers is a dict of codegen-affecting overrides applied on top of the defaults (floating_base / ee_joint_names / max_batch_size) — e.g. tiers=[{}, {"floating_base": True}] prebuilds both the fixed- and floating-base .so. tiers=None builds the single default tier. Each requested backend (“numpy”/”jax”/”torch”) warms that surface’s artifacts on the same cached .so.

This is a thin, idempotent driver over register_robot(): a tier already in the cache is a no-op (no nvcc); a missing tier compiles once and populates the cache. Returns the manifest entry for each (tier, backend) built, in order. Build offline once, ship/keep the cache dir, and every later run starts in well under a second.

grid_rbd.build_plan(name: str, urdf_path: str | None = None, **opts: Any) → dict[str, Any][source]#

Resolve a registration WITHOUT building it: the effective codegen options, the build identity (toolchain / GLASS content / env knobs), the stage-1 input key, and whether the cache already holds a sound artifact for it (cached / content_key / stale_reasons).

No generation, no nvcc, no dlopen, no CUDA context: it needs only the URDF and a cuda_arch= (or nvidia-smi on PATH to detect one). Takes the same keyword arguments as warm_robot() / register_robot() (build-affecting ones only). Use it to explain a rebuild, to pre-check a deployment’s cache, or to see what a build box needs before it builds.

grid_rbd.gc_cache(cache_dir: str | Path | None = None, *, execute: bool = False, objects_days: float | None = None) → dict[str, Any][source]#: Scan (and with execute=True delete) the cache garbage described in the module docstring. Returns a report dict; sizes in bytes.

Consumer artifacts#

A consumer’s own CUDA sources compiled against the header GRiD generates for a robot, under the store’s rules, in a namespace of the cache directory that belongs to the consumer (<cache_dir>/consumers/<consumer>/). GRiD never binds the consumer’s symbols; the caller loads the artifact itself.

grid_rbd.build_artifact(consumer: str, name: str, urdf_path: str | Path, *, sources: Sequence[str | Path], headers: Sequence[str | Path] = (), include_dirs: Sequence[str | Path] = (), defines: Sequence[str] | dict[str, Any] | None = None, extra_flags: Sequence[str] = (), link_flags: Sequence[str] = (), artifact_name: str = 'artifact.so', cache_dir: str | Path | None = None, cuda_arch: int | None = None, force_rebuild: bool = False, plan_only: bool = False, floating_base: bool = False, ee_joint_names: Sequence[str] | None = None, dtype: str = 'float32', algorithm_list: Sequence[str] | str | None = None, enable_mujoco_kernels: bool = True, cost_frames: Sequence[str] | None = None, contact_frames: Sequence[str] | None = None, use_joint_dynamics: bool = False, friction_model: str = 'sign', friction_width=None, runtime_inertia: bool = False, runtime_transform: bool = False, runtime_joint_dynamics: bool = False, emit_plant: bool | None = None, collision: bool = False, collision_res: float | Sequence[float] = 0.05, collision_native: bool = False, spherized_urdf: str | Path | None = None, launch_config_profile: str | None = None) → ConsumerArtifact | dict[str, Any][source]#

Build (or reuse) the consumer’s artifact for one robot and return where it is.

consumer names the namespace in the cache directory (one per consuming project); name the artifact within it (one per robot and configuration). sources are the consumer’s translation units, each compiled with -c against the generated grid.cuh (plus the static collision header and the GLASS include directories) and linked with –shared into artifact_name; headers are the consumer’s own header files, folded into the identity by content (an include directory is folded by path only, so list the headers whose edits must re-key the build); defines are NAME or NAME=VALUE; extra_flags go to every unit compile, link_flags to the link.

The codegen keywords (floating_base, ee_joint_names, dtype, algorithm_list, enable_mujoco_kernels, cost_frames, contact_frames, the joint-dynamics and runtime knobs) mean what they mean on register_robot; emit_plant is the generator’s plant-block switch (None: the generator’s rule), and collision, collision_res, collision_native and spherized_urdf select the generated grid_collision namespace (a sphere model is validated against the URDF). launch_config_profile (default: the bindings’ “ffi” profile) names the baked launch-configuration profile.

plan_only=True reports what a build would use and whether it is already there, without generating or compiling. force_rebuild=True recompiles every unit.

grid_rbd.load_artifact(consumer: str, name: str, cache_dir: str | Path | None = None, cuda_arch: int | None = None) → ConsumerArtifact[source]#: The prepared artifact name of consumer, without a toolchain: the manifest names it, the entry must be whole and recorded for this device (or one of a fat build’s devices) and key schema. With no detectable GPU and no cuda_arch, the recorded device is trusted.

grid_rbd.list_artifacts(consumer: str, cache_dir: str | Path | None = None) → list[dict[str, Any]][source]#: Every artifact the consumer’s manifest names, with its metadata.

class grid_rbd.ConsumerArtifact(consumer: 'str', name: 'str', path: 'Path', header: 'Path', store_dir: 'Path', content_key: 'str', input_key: 'str | None', cuda_arch: 'int', meta: 'dict[str, Any]', compiled: 'bool')[source]#

consumer: str#

name: str#

path: Path#

header: Path#

store_dir: Path#

content_key: str#

input_key: str | None#

cuda_arch: int#

meta: dict[str, Any]#

compiled: bool#

__init__(consumer: str, name: str, path: Path, header: Path, store_dir: Path, content_key: str, input_key: str | None, cuda_arch: int, meta: dict[str, Any], compiled: bool) → None#

exception grid_rbd.ConsumerArtifactError[source]#: Base of the artifact errors.

exception grid_rbd.ArtifactNotFoundError[source]#: No artifact of that name in the consumer’s manifest.

exception grid_rbd.ArtifactIncompatibleError[source]#: The recorded artifact cannot be loaded here (another device or key schema).

The numpy handle#

class grid_rbd.RobotHandle(name: str, so_path: str, meta: dict[str, Any])[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 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" (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 grid_assets/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).

property friction_model: str#

"sign" (the default, derivatives zero almost everywhere) or "tanh" (the smooth friction*tanh(qd/w) with its derivatives in the gradient and second-order kernels). "sign" on a build without joint dynamics.

Type:: The Coulomb friction shape of a use_joint_dynamics build

property joint_friction_width#: Baked per-v-slot tanh friction width w (length nv, rad/s). Only on a runtime_joint_dynamics build with friction_model="tanh" and a robot that declares friction (raises otherwise).

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

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

Pass any of the three as length-nv arrays (v-slot indexed, same basis as joint_damping / joint_friction / joint_friction_width). An omitted side keeps its current baked value. friction_width is accepted only on a friction_model="tanh" build, where the table carries a third row. 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 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

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.