diff --git a/CLAUDE.md b/CLAUDE.md index dff1e2138..9e7e8b074 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -46,6 +46,9 @@ - **Always pass `transpose=True`** to `qd_to_numpy` / `qd_to_torch` to move the batch dimension to front (`[B, n, dim]`), aligning with public getter conventions. - **Hoist conversions before loops.** Call `qd_to_numpy` / `qd_to_torch` once, then index the local array. - **`qd_to_numpy` with `copy=False` fails on CUDA.** GPU data always requires a copy to numpy. +- **Store `qd_to_torch(..., copy=False)` views in named locals**, one local per view; never chain the in-place op on the call expression. Drop `transpose=True` for element-wise in-place ops between two field views. +- **Metal stream sync is the caller's job**: after a batch of torch zero-copy writes, call `torch.mps.synchronize()` once before the next quadrants kernel reads the buffers. Helpers performing such writes (e.g. `qd_zero_grad`) never sync internally. +- **Branch zero-copy paths on `gs.use_zerocopy` alone; no try/except probing.** - **Do not pass `copy=` at all** to `qd_to_torch` / `qd_to_numpy` when the value returned to the caller is fresh arithmetic (nothing internal aliases out) — let it zero-copy or copy only as needed. Reserve `copy=True` for when a raw field view would otherwise escape and could be mutated. - **Store Python references at build time** (entity, link objects in dicts) for runtime lookup, rather than re-deriving from solver fields every step. - **Never pass numpy scalar types** (e.g., `numpy.float64`) to kernel arguments. Always cast to native Python types (`float()`, `int()`). Numpy scalars break quadrants fastcache (weak reference). @@ -58,6 +61,7 @@ - **Rigid kernel and func calls are positional for self-named arguments**, relying on the canonical parameter order below; anonymous constants (bare literals such as trailing static flags) are passed by keyword. A func call that passes a struct member alongside its parent struct stays keyword-only, because quadrants' positional func-argument expansion duplicates the member. - **Canonical parameter order, for every kernel and func:** indices of any kind first (loop indices, index tensors like `envs_idx`, index-valued scalars like `joint_idx` or range starts/ends), then dynamic native Python scalars and fixed-size quadrants vectors (per-call values: positions, quaternions, penetrations), then the kernel/func-specific tensors, then all state structs, then all info structs, then the static configs, then the kernel/func-specific constants-in-practice (shape integers, eps, tolerances - info by nature even when not declared as such), and the kernel/func-specific static compilation flags (`qd.template()` booleans) at the very end - except `errno`, which systematically takes the last position. Within the state/info/config groups the component order is fixed - dyn, rigid, collider (mpr, gjk, support_field, sdf), constraint - with leaf structs at their aggregate's slot in its declared field order. Call sites pass arguments in signature order. - **Rigid solver** is the reference implementation for kernel and code quality standards. +- **No dtype in quadrants fixed-size vector annotations**: `qd.types.vector(3)`, never `qd.types.vector(3, dtype=gs.qd_float)`. - **Pure read-write data accessors on the hot path go through zero-copy views, keeping the kernel only as fallback.** A tiny kernel that merely shuffles or rewrites a field (payload remap, flag sweep) pays full kernel dispatch for microseconds of work: when `gs.use_zerocopy` holds, do it in place through a `qd_to_torch(..., copy=False)` / `qd_to_numpy(..., copy=False)` view, and fall back to the kernel otherwise. Reserve dedicated kernels for genuine computation. ## API Design @@ -114,6 +118,7 @@ - **Step budget (hard).** Use the strict minimum steps needed, measured not guessed; if the checked quantity is FP-sensitive, pad by whichever is larger: +10% or round up to 50. Per-test horizon tiers: `<100` fine, `100-300` grey (needs justification), `>300` near-prohibited (only 4-5 such tests in the whole suite). Also bound total env-steps (`steps * n_envs`) - a 500-step x 16-env settle (8000) is unacceptable; delete it or fit the budget. CI cost multiplies every step across the matrix. - **Validate the math, not scale.** One constraint proves the theory (works for 1 => works for 1k); no large/high-DOF scene is needed for correctness. A real multi-body scene earns its steps only when it validates something a single constraint cannot (end-to-end stability, e.g. the bowl tower). - **Prefer FP-robust checks.** A single constraint-solve comparison (one step from a fixed state) is robust; a multi-step trajectory comparison across numerically-distinct codepaths (different solver arms/factor paths) is not - fp32 compounds and correct paths diverge. Compare a single solve, or an analytical physical property. Cross-engine (MuJoCo) consistency = the real-world check; analytical closed-form = the math check. Regression tests are reactive (add when a bug surfaces), not speculative. +- **Gradient-vs-FD tolerances are pinned to measured floors**: floor T = max|ana - fd| / (1 + |fd|) at the configured eps, worst across cpu and BOTH GPU archs; tolerance = 1.5x-5x the floor, values only {1, 2, 5}e-X; do-not-chase floors 1e-10 (fp64) / 5e-5 (fp32); eps per precision (large for fp32, small for fp64). A floor of exactly 0 = vacuous check - fix the loss, not the tolerance. An eps-independent residual is an analytical bug, never an FD artifact. - **Tighten toward physically-exact conditions; slack hides bugs.** A padded safety factor or generous tolerance that makes a scenario pass trivially also masks real behavioral differences (across backends/arms/factor paths). Prefer the tightest thresholds the physics justifies. When tightening a test makes it start failing, investigate the root, don't back off. And verify a fix against the EXACT assertion path (full horizon, all phases), never a cheaper proxy - the part a proxy drops is where behavior diverges. ### Scene construction in tests and examples diff --git a/CODING_GUIDELINES.md b/CODING_GUIDELINES.md index 6b65e6304..15b17b45a 100644 --- a/CODING_GUIDELINES.md +++ b/CODING_GUIDELINES.md @@ -5,7 +5,7 @@ This document describes how we write, test, and review code in Genesis. It is ai Two things to keep in mind throughout: - **The rigid solver is the reference implementation.** When in doubt about naming, kernel structure, or code organization, open `genesis/engine/solvers/rigid/` and imitate what you find there. The historical rationale for the naming scheme is documented in [PR #1053](https://github.com/Genesis-Embodied-AI/Genesis/pull/1053). -- **Match the established sibling pattern; do not invent structure.** Before adding any helper, fixture, option, cast, comment, or file layout, look at how a sibling in the same directory already does it, and copy that. Genuine novelty is rarely the problem in reviews; deviation from an existing convention almost always is. +- **Match the established sibling pattern; do not invent structure.** Before adding any helper, fixture, option, cast, comment, or file layout, look at how a sibling in the same directory already does it, and copy that. Genuine novelty is rarely the problem in reviews; deviation from an existing convention almost always is. When a sibling contradicts an explicit rule in this document, the rule wins: legacy remnants are not a license. ## 1. Philosophy and priorities @@ -136,6 +136,9 @@ Genesis strongly resists the accumulation of small wrappers. Helpers accrete, ge - **Always pass `transpose=True`** so the batch dimension comes first (`[B, n, dim]`), matching the public getter convention. - **Hoist conversions out of loops.** Convert once, then index the local array. If only a slice is needed, pass the slice as an argument to the conversion rather than slicing the full result afterward. - Do not pass `copy=` at all when the returned value is fresh arithmetic; let the function zero-copy or copy as needed. Reserve `copy=True` for the case where a raw field view would otherwise escape and could be mutated later. Note that `copy=False` fails on CUDA: GPU data always requires a copy to reach numpy. +- **Store `qd_to_torch(..., copy=False)` views in named locals** before operating on them, one local per view; never chain the in-place op on the call expression. The named local documents what the view aliases and keeps each zero-copy acquisition visible on its own line. For element-wise in-place ops between two field views, drop `transpose=True`: both sides share the raw field layout, so transposing both is a no-op that only makes the views non-contiguous (batch-first is for values returned to consumers). +- **Metal stream sync is the caller's job.** torch (MPS) and quadrants share no compute stream on Metal: after a batch of torch zero-copy writes, call `torch.mps.synchronize()` once before the next quadrants kernel reads the buffers. Helpers performing such writes (e.g. `qd_zero_grad`) never sync internally, so a caller can batch many of them under a single flush. +- **Branch zero-copy paths on `gs.use_zerocopy` alone; no try/except probing.** Every platform gate (backend, torch version, device consistency) is resolved at `gs.init`, array_class buffers are standalone dense allocations with DLPack-supported dtypes, and allocations are capped at int32 element count at build time - so for these buffers the flag decides availability exactly, and exception-driven fallbacks are speculative defensive code. - **Zero-copy aliasing hazard:** on the CPU backend, `qd_to_numpy` returns a live view into the field buffer. If you store the result across simulation steps, `.copy()` it first (or reduce it to Python scalars immediately) - otherwise a later "delta" between stored and current values is identically zero, because both names alias the same memory. - Prefer torch-based zero-copy to feed state into other solvers; numpy zero-copy only exists on the CPU backend. - Handle all environments at once. Calling getters/setters per environment index in a loop is inefficient; batch the call, falling back to a loop only where an accessor fundamentally requires it. Avoid branching code on batched-versus-unbatched; handle it at the argument level (`envs_idx=env_idx if batched else None`). @@ -148,6 +151,7 @@ Genesis strongly resists the accumulation of small wrappers. Helpers accrete, ge - **New kernels are free functions decorated with `@qd.kernel`** - no `@qd.data_oriented` classes. Use `V_ANNOTATION` from `genesis.utils.array_class` for type-polymorphic parameters. The FEM solver is the historical exception: it follows the old `@qd.data_oriented` method pattern, and kernels added there must stay consistent with it. Never put `@qd.data_oriented` on non-solver classes (materials, couplers, entities). - **Call rigid kernels and funcs positionally for self-named arguments.** The canonical parameter order below makes those call sites unambiguous without keyword names. Anonymous constants (bare literals, typically the trailing static flags) are passed by keyword so their meaning shows at the call site, and every argument after the first keyword one is keyword too. Exception: numeric constants in module-local qd.func calls stay positional - forcing the keyword there cascades onto every following argument and wrecks the compactness of internal call chains for no readability gain. One exception: a func call that passes a struct member alongside its parent struct (e.g. `constraint_state.qacc` next to `constraint_state`) stays keyword-only, because quadrants' positional func-argument expansion duplicates the member in the flattened call. Kernel calls from Python are unaffected. - **Parameters follow the canonical component order.** Indices of any kind come first - loop indices, index tensors (`envs_idx`, `dofs_idx`), index-valued scalars (`joint_idx`, geom ids, range starts/ends) - then the dynamic native Python scalars and fixed-size quadrants vectors (per-call values: positions, quaternions, penetrations, step sizes), then the tensors specific to the kernel/func (raw inputs/outputs, errno), then every state struct, then every info struct, then the static configs, then the kernel/func-specific constants-in-practice (shape integers, eps, tolerances - info by nature even when not declared as such), and finally the kernel/func-specific static compilation flags (`qd.template()` booleans such as `is_backward`). One exception: `errno` systematically takes the last position. The order keeps what matters most in front and sorts the rest by constant-ness (how likely each argument is to change between calls); since the constants-in-practice cluster at the tail, an anonymous constant can take its keyword without forcing keyword form onto the arguments that matter. Within the state/info/config groups the component order is fixed - dyn, rigid, collider (mpr, gjk, support_field, sdf), constraint - with leaf structs at their aggregate's slot in its declared field order, and adjoint-cache instances right after their primary. Within a group the guideline mandates nothing further (beyond the fixed component order of the struct groups): keep the existing relative order, and never reshuffle a signature more than the guideline requires. One predictable order makes every signature readable at a glance, keeps positional call sites unambiguous, and turns argument mismatches into immediate compile errors instead of silent misbindings. +- **No dtype in quadrants fixed-size vector annotations**: `qd.types.vector(3)`, never `qd.types.vector(3, dtype=gs.qd_float)`. Annotations are evaluated once at import and never re-evaluated after destroy / re-init, so a captured dtype goes stale when precision changes between inits. - **The outermost `for` loop of a kernel is auto-parallelized.** Outer-scope locals mutated inside that loop are a single shared memory slot that all threads race on; the value after the loop is undefined and run-to-run unstable on GPU. The CPU backend often looks correct only because its default executor is single-threaded - do not let that fool you. To fix: serialize the loop (wrap the body in `for _ in range(1):`, or use `qd.loop_config(serialize=True)`), or use `qd.atomic_min`/`qd.atomic_max`/`qd.atomic_add` for pure reductions. When something "works on CPU but fails on Metal", check for this pattern before suspecting a compiler bug. - **Loop-variable scoping:** a variable assigned in the outer scope of a kernel (for example a lane-strided `i_c_ = tid; while i_c_ < n: ...`) cannot be reused as a `for` loop variable in the same kernel; the compiler rejects it. The trailing-underscore convention for relative indices conveniently keeps these names distinct. - **Do not pass the same buffer under two different argument names of a single kernel.** Intra-kernel argument aliasing miscompiles. Use dedicated buffers. @@ -227,6 +231,7 @@ Genesis strongly resists the accumulation of small wrappers. Helpers accrete, ge - If your change makes an existing test fail and you believe the test's expectation is what should change, raise it with the maintainers with your analysis; do not silently rewrite assertions. - **Never remove or weaken an existing assertion or measurement to silence a failure**, including a CI-only failure. A threshold that only holds on one machine is a calibration problem to raise with data, not a license to delete the check. - **Tighten toward physically-exact conditions; slack hides bugs.** A padded safety factor or a generous tolerance that makes a scenario pass trivially also masks real behavioral differences - between solver backends, arms, or factor paths. Prefer the tightest thresholds the physics justifies: an over-pushed or loosely-checked scenario passes on every backend while a genuine discrepancy sits underneath it. When tightening a test toward the exact expected behavior makes it start failing, that failure is the test working as intended - a newly-observable issue to investigate at its root, not a calibration to back off. +- **Gradient-vs-FD tolerances are pinned to measured floors.** Floor T = max|ana - fd| / (1 + |fd|) at the configured eps, worst across platforms - cpu and BOTH GPU archs, whose fp32 floors can differ by 4x. Tolerance = 1.5x-5x the floor, values only {1, 2, 5}e-X, with do-not-chase floors 1e-10 (fp64) / 5e-5 (fp32) and per-precision eps (fp32 needs a large step to clear its noise floor, fp64 a small one to bound truncation). A floor of exactly 0 means the check is vacuous (constant observable) - fix the loss, not the tolerance. An eps-independent residual is an analytical bug (missing adjoint term or wrong backward primal), never an FD artifact. - **Verify a fix against the exact assertion path, not a cheaper proxy.** A reduced reproduction - fewer steps, only the first phase, a simpler scene - can pass while the full test still fails, because the part it drops is exactly where the behavior diverges (a longer settle, a second load phase). Confirm on the real test at the real horizon before declaring a fix validated. ### 12.6 Step budget and minimality @@ -243,6 +248,6 @@ Genesis strongly resists the accumulation of small wrappers. Helpers accrete, ge - **Commit messages are a single title line**: a plain sentence ending with a period, without the PR bracket tag. No body, no implementation details, no co-author trailers. - **No AI attribution anywhere.** No "Generated with ..." footers, no AI co-author trailers, in commits, PR bodies, issues, or comments. - **PR titles carry exactly one bracket tag from a closed set and end with a period.** The tags are `[BUG FIX]`, `[FEATURE]`, `[MISC]` (everything else, including performance and refactor work: `[MISC] Speed up ...`), `[CHANGING]` (a default-behavior change), and `[BREAKING]` (an API break). Inventing a tag (`[PERF]`) is a review-blocking error. -- **PR titles state the benefit for end users, not the implementation.** PR titles are used to generate the changelog automatically, so each title must be clear and enjoyable to read for end users and let them understand what improvement to expect: `[MISC] Speed up forward kinematics on GPU.`, not `[MISC] Remove unnecessary atomic add from root COM accumulation.`. The mechanism belongs in the PR description. +- **PR titles state the benefit for end users, not the implementation.** PR titles are used to generate the changelog automatically, so each title must be clear and enjoyable to read for end users and let them understand what improvement to expect: `[MISC] Speed up forward kinematics on GPU.`, not `[MISC] Remove unnecessary atomic add from root COM accumulation.`. The mechanism belongs in the PR description, and so does every testing / validation detail (tolerances, platforms, suites) - never the title. - **PR descriptions are terse and go straight to the point.** State the change and the one fact that justifies it; cut every sentence a reviewer can infer. Long descriptions do not get read. Follow the PR template (`.github/pull_request_template.md`), keeping each section tight and removing sections that are empty. - **Search for existing related issues and track them in the PR description**, even when you are not the author of those issues. Link each one in the `Related Issue` section with the accurate keyword: `Resolves` only when the PR fully closes it, otherwise `Addresses` with a word on which part is covered. This keeps the issue tracker in sync with the work actually landing. diff --git a/genesis/engine/entities/rigid_entity/rigid_entity.py b/genesis/engine/entities/rigid_entity/rigid_entity.py index 343340e19..24e0dceb3 100644 --- a/genesis/engine/entities/rigid_entity/rigid_entity.py +++ b/genesis/engine/entities/rigid_entity/rigid_entity.py @@ -2,7 +2,7 @@ import math import os from itertools import chain -from typing import TYPE_CHECKING, Literal, Any +from typing import TYPE_CHECKING, Literal, Any, Hashable from functools import wraps import quadrants as qd @@ -21,7 +21,7 @@ from genesis.utils import mjcf as mju from genesis.utils import terrain as tu from genesis.utils import urdf as uu -from genesis.utils.misc import DeprecationError, broadcast_tensor, qd_to_numpy, qd_to_torch +from genesis.utils.misc import DeprecationError, broadcast_tensor, qd_to_numpy, qd_to_torch, tensor_to_array from genesis.typing import UnitVec4FType, Vec3FType from genesis.engine.states.entities import RigidEntityState @@ -59,7 +59,23 @@ def wrapper(self, *args, **kwargs): bound = sig.bind(self, *args, **kwargs) bound.apply_defaults() args_dict = dict(tuple(bound.arguments.items())[1:]) - self._update_tgt(fun.__name__, args_dict) + # Key the slot by (method, dofs subset) so same-step calls on distinct subsets (e.g. arm and gripper + # force control) each keep their own entry and gradient path; keyed by method alone, the second call + # would evict the first from the tape. Slices key directly when hashable (Python 3.12 onward) and + # resolve against the entity dof count otherwise. + dofs_idx_local = args_dict.get("dofs_idx_local") + if dofs_idx_local is None: + subset = None + elif isinstance(dofs_idx_local, slice): + if isinstance(dofs_idx_local, Hashable): + subset = dofs_idx_local + else: + subset = tuple(range(*dofs_idx_local.indices(self.n_dofs))) + elif isinstance(dofs_idx_local, torch.Tensor): + subset = tuple(tensor_to_array(dofs_idx_local).reshape(-1).tolist()) + else: + subset = tuple(np.asarray(dofs_idx_local).reshape(-1).tolist()) + self._update_tgt((fun.__name__, subset), args_dict) return fun(self, *args, **kwargs) return wrapper @@ -143,7 +159,7 @@ def __init__( self._load_model() # Initialize target variables and checkpoint - self._tgt_keys = ("pos", "quat", "qpos", "dofs_velocity", "control_dofs_force") + self._tgt_keys = ("set_pos", "set_quat", "set_dofs_velocity", "control_dofs_force") self._tgt = dict() self._tgt_buffer = list() self._ckpt = dict() @@ -1426,6 +1442,8 @@ def attach( """ if self._is_attached: gs.raise_exception("Entity already attached.") + if self._solver._requires_grad: + gs.raise_exception("Attach is not supported yet when requires_grad is True.") is_mounting = pos is not None or quat is not None if is_mounting: @@ -1570,7 +1588,7 @@ def process_input(self, in_backward=False): # Do not update [tgt], as input information is finalized at this point self._update_tgt_while_set = False - match key: + match key[0]: case "set_pos": self.set_pos(**data_kwargs) case "set_quat": @@ -1580,7 +1598,7 @@ def process_input(self, in_backward=False): case "control_dofs_force": self.control_dofs_force(**data_kwargs) case _: - gs.raise_exception(f"Invalid target key: {key} not in {self._tgt_keys}") + gs.raise_exception(f"Invalid target key: {key[0]} not in {self._tgt_keys}") self._tgt = dict() self._update_tgt_while_set = update_tgt_while_set @@ -1590,35 +1608,37 @@ def process_input_grad(self): for key in reversed(self._tgt_buffer[index].keys()): data_kwargs = self._tgt_buffer[index][key] - match key: - # We need to unpack the data_kwargs because [_backward_from_qd] only supports positional arguments + match key[0]: + # We need to unpack the data_kwargs because [_backward_from_qd] only supports positional arguments. + # Inputs are stored on the tape as passed by the user, so scalars and array-likes (valid setter + # inputs that cannot carry a gradient) are filtered out with the tensor check. case "set_pos": pos = data_kwargs.pop("pos") - if pos.requires_grad: + if isinstance(pos, torch.Tensor) and pos.requires_grad: pos._backward_from_qd(self.set_pos_grad, data_kwargs["envs_idx"], data_kwargs["relative"]) case "set_quat": quat = data_kwargs.pop("quat") - if quat.requires_grad: + if isinstance(quat, torch.Tensor) and quat.requires_grad: quat._backward_from_qd(self.set_quat_grad, data_kwargs["envs_idx"], data_kwargs["relative"]) case "set_dofs_velocity": velocity = data_kwargs.pop("velocity") # [velocity] could be None when we want to zero the velocity (see set_dofs_velocity of RigidSolver) - if velocity is not None and velocity.requires_grad: + if isinstance(velocity, torch.Tensor) and velocity.requires_grad: velocity._backward_from_qd( self.set_dofs_velocity_grad, data_kwargs["dofs_idx_local"], data_kwargs["envs_idx"] ) case "control_dofs_force": force = data_kwargs.pop("force") - if force.requires_grad: + if isinstance(force, torch.Tensor) and force.requires_grad: force._backward_from_qd( self.set_dofs_force_grad, data_kwargs["dofs_idx_local"], data_kwargs["envs_idx"] ) case _: - gs.raise_exception(f"Invalid target key: {key} not in {self._tgt_keys}") + gs.raise_exception(f"Invalid target key: {key[0]} not in {self._tgt_keys}") def save_ckpt(self, ckpt_name): if ckpt_name not in self._ckpt: diff --git a/genesis/engine/scene.py b/genesis/engine/scene.py index 0a46c551d..30815b3a9 100644 --- a/genesis/engine/scene.py +++ b/genesis/engine/scene.py @@ -1050,12 +1050,13 @@ def backward(self, loss: torch.Tensor, *args, **kwargs): snapshot = self.get_state() # The sim unroll (self._backward) re-enters the torch graph from each step's queried states, so the graph # must survive the initial autograd pass. - kwargs.setdefault("retain_graph", True) + if kwargs.setdefault("retain_graph", True) is not True: + gs.raise_exception("'retain_graph' must be left unset: scene.backward requires the graph to survive.") # The functional torch.autograd.backward fills torch and queried-state grads while leaving the sim unroll to # the explicit self._backward call below, keeping gs.Tensor.backward's automatic scene._backward out of it. torch.autograd.backward(loss, *args, **kwargs) self._backward() - # keep_init=True preserves the registered initial state so a later bare reset() still rewinds to it. + # keep_init semantics: see _reset self._reset(snapshot, keep_init=True) return snapshot diff --git a/genesis/engine/solvers/kinematic_solver.py b/genesis/engine/solvers/kinematic_solver.py index c66809d5e..72addf95b 100644 --- a/genesis/engine/solvers/kinematic_solver.py +++ b/genesis/engine/solvers/kinematic_solver.py @@ -644,6 +644,9 @@ def reset_grad(self): qd_zero_grad(self.dyn_state.dofs) qd_zero_grad(self.dyn_state.joints) qd_zero_grad(self.rigid_info) + # One flush for the zeroing batch; see qd_zero_grad in misc.py. + if gs.use_zerocopy and gs.backend == gs.metal: + torch.mps.synchronize() for entity in self._entities: entity.reset_grad() self._queried_states.clear() diff --git a/genesis/engine/solvers/rigid/abd/diff.py b/genesis/engine/solvers/rigid/abd/diff.py index edf3a9877..bf374cdc4 100644 --- a/genesis/engine/solvers/rigid/abd/diff.py +++ b/genesis/engine/solvers/rigid/abd/diff.py @@ -271,8 +271,8 @@ def kernel_copy_next_to_curr_no_check( rigid_info: array_class.RigidInfo, rigid_config: qd.template(), ): - # Unguarded copy of the _next slots to current, used in the backward substep right before the forward replay so - # the backward kernels see the post-integrate qpos / vel. + """Unguarded copy of the _next slots to current, used in the backward substep right before the forward replay so + the backward kernels see the post-integrate qpos / vel.""" n_qs = rigid_info.qpos.shape[0] n_dofs = dyn_state.dofs.vel.shape[0] _B = dyn_state.dofs.vel.shape[1] diff --git a/genesis/engine/solvers/rigid/abd/forward_dynamics.py b/genesis/engine/solvers/rigid/abd/forward_dynamics.py index d187b8734..fc43d7fd4 100644 --- a/genesis/engine/solvers/rigid/abd/forward_dynamics.py +++ b/genesis/engine/solvers/rigid/abd/forward_dynamics.py @@ -103,7 +103,7 @@ def func_forward_dynamics( func_update_acc(dyn_state, dyn_info, rigid_info, rigid_config, update_cacc=False, is_backward=is_backward) func_update_force(dyn_state, dyn_info, rigid_info, rigid_config, is_backward) func_bias_force(dyn_state, dyn_info, rigid_info, rigid_config, is_backward) - func_compute_qacc(dyn_state, dyn_info, rigid_info, rigid_config, is_backward) + func_compute_qacc(dyn_state, dyn_info, rigid_info, rigid_config) @qd.kernel(fastcache=True) @@ -1187,27 +1187,13 @@ def func_bias_force( dyn_state.dofs.qf_smooth[i_d, i_b] = dyn_state.dofs.force[i_d, i_b] -@qd.kernel -def kernel_compute_qacc( - dyn_state: array_class.DynState, - dyn_info: array_class.DynInfo, - rigid_info: array_class.RigidInfo, - rigid_config: qd.template(), - is_backward: qd.template(), -): - func_compute_qacc(dyn_state, dyn_info, rigid_info, rigid_config, is_backward) - - @qd.func def func_compute_qacc( dyn_state: array_class.DynState, dyn_info: array_class.DynInfo, rigid_info: array_class.RigidInfo, rigid_config: qd.template(), - is_backward: qd.template(), ): - BW = qd.static(is_backward) - func_solve_mass(dyn_state.dofs.force, dyn_state.dofs.acc_smooth, dyn_info, rigid_info, rigid_config) # Assume this is the outermost loop @@ -1581,10 +1567,7 @@ def func_implicit_damping( dyn_info: array_class.DynInfo, rigid_info: array_class.RigidInfo, rigid_config: qd.template(), - is_backward: qd.template(), ): - BW = qd.static(is_backward) - EPS = rigid_info.EPS[None] n_entities = dyn_info.entities.dof_start.shape[0] diff --git a/genesis/engine/solvers/rigid/abd/manual_bw.py b/genesis/engine/solvers/rigid/abd/manual_bw.py index 544fc0902..1eb397370 100644 --- a/genesis/engine/solvers/rigid/abd/manual_bw.py +++ b/genesis/engine/solvers/rigid/abd/manual_bw.py @@ -2,10 +2,8 @@ They hand-compute the Jacobian-transpose products of forward stages whose reverse Quadrants automatic differentiation silently drops (forward kinematics, forward velocity) or cannot express (the mass solve, reversed via the implicit -function theorem). Hibernation support is pending and flags errno (ErrorCode.MANUAL_BW_UNIMPLEMENTED) so the host -halts instead of silently corrupting gradients. The chain-rule building blocks (quaternion product, quaternion -transform, rotation-vector conversion, motion cross product) live next to their forwards in genesis/utils/geom.py as -the *_grad_* adjoint funcs. +function theorem). The chain-rule building blocks (quaternion product, quaternion transform, rotation-vector +conversion, motion cross product) live next to their forwards in genesis/utils/geom.py as the *_grad_* adjoint funcs. """ import quadrants as qd @@ -21,7 +19,6 @@ def kernel_manual_forward_kinematics_bw( dyn_info: array_class.DynInfo, rigid_info: array_class.RigidInfo, rigid_config: qd.template(), - errno: qd.Tensor, ): """Manual reverse of kernel_forward_kinematics_replay: link pos / quat grads back to qpos and parent-pose grads. @@ -44,7 +41,7 @@ def kernel_manual_forward_kinematics_bw( for i_l_rev in range(n_in_e): i_l = dyn_info.entities.link_end[i_e] - 1 - i_l_rev I_l = [i_l, i_b] if qd.static(rigid_config.batch_links_info) else i_l - parent_idx = dyn_info.links.parent_idx[I_l] + i_parent = dyn_info.links.parent_idx[I_l] n_joints = dyn_info.links.joint_end[I_l] - dyn_info.links.joint_start[I_l] # Grad seeded on the final link pose (= slot n_joints). Carried @@ -70,11 +67,15 @@ def kernel_manual_forward_kinematics_bw( xaxis_grad = dyn_state.joints.xaxis.grad[i_j, i_b] if joint_type == gs.JOINT_TYPE.FREE: - # Final pose is set absolutely from qpos (slot in unused); - # xanchor = qpos[0:3]. + # Final pose is set absolutely from qpos (slot in unused); xanchor = qpos[0:3]. The forward also + # publishes the pose as dofs.pos = (pos, quat_to_xyz(quat)), whose grad (accumulated by the + # force-stage reverses through position control / joint stiffness) folds back here. for j in qd.static(range(3)): rigid_info.qpos.grad[q_start + j, i_b] = ( - rigid_info.qpos.grad[q_start + j, i_b] + g_pos[j] + xanchor_grad[j] + rigid_info.qpos.grad[q_start + j, i_b] + + g_pos[j] + + xanchor_grad[j] + + dyn_state.dofs.pos.grad[dof_start + j, i_b] ) # The forward normalizes the raw qpos quaternion (quat_ = q / |q|) before writing the link # quaternion, so the gradient must pass through the Jacobian of q / |q|: @@ -91,11 +92,22 @@ def kernel_manual_forward_kinematics_bw( ) q_norm = q_raw.norm() qhat = q_raw / q_norm - g_quat_raw = (g_quat - qhat * qhat.dot(g_quat)) / q_norm + g_xyz = qd.Vector( + [ + dyn_state.dofs.pos.grad[dof_start + 3, i_b], + dyn_state.dofs.pos.grad[dof_start + 4, i_b], + dyn_state.dofs.pos.grad[dof_start + 5, i_b], + ], + dt=gs.qd_float, + ) + g_quat_eff = g_quat + gu.qd_quat_to_xyz_grad_quat(qhat, rigid_info.EPS[None], g_xyz) + g_quat_raw = (g_quat_eff - qhat * qhat.dot(g_quat_eff)) / q_norm for j in qd.static(range(4)): rigid_info.qpos.grad[q_start + 3 + j, i_b] = ( rigid_info.qpos.grad[q_start + 3 + j, i_b] + g_quat_raw[j] ) + for j in qd.static(range(6)): + dyn_state.dofs.pos.grad[dof_start + j, i_b] = 0.0 g_pos = qd.Vector([0.0, 0.0, 0.0], dt=gs.qd_float) g_quat = qd.Vector([0.0, 0.0, 0.0, 0.0], dt=gs.qd_float) @@ -114,7 +126,12 @@ def kernel_manual_forward_kinematics_bw( g_quat_in_apply = gu.qd_quat_mul_grad_lhs(quat_in, qloc, gq_out) rotvec_grad = gu.qd_rotvec_to_quat_grad_rotvec(rotvec, rigid_info.EPS[None], g_qloc) angle_grad = axis[0] * rotvec_grad[0] + axis[1] * rotvec_grad[1] + axis[2] * rotvec_grad[2] - rigid_info.qpos.grad[q_start, i_b] = rigid_info.qpos.grad[q_start, i_b] + angle_grad + # The forward also publishes the angle as dofs.pos[dof_start] = qpos - qpos0, whose grad + # (accumulated by the force-stage reverses) folds back linearly. + rigid_info.qpos.grad[q_start, i_b] = ( + rigid_info.qpos.grad[q_start, i_b] + angle_grad + dyn_state.dofs.pos.grad[dof_start, i_b] + ) + dyn_state.dofs.pos.grad[dof_start, i_b] = 0.0 # grad into xanchor = g_pos (from pos_out) + downstream xanchor_grad g_xanchor = g_pos + xanchor_grad @@ -132,7 +149,12 @@ def kernel_manual_forward_kinematics_bw( xaxis = gu.qd_transform_by_quat(axis, quat_in) # pos_out = pos_in + xaxis * displacement ; quat_out = quat_in displacement_grad = xaxis[0] * g_pos[0] + xaxis[1] * g_pos[1] + xaxis[2] * g_pos[2] - rigid_info.qpos.grad[q_start, i_b] = rigid_info.qpos.grad[q_start, i_b] + displacement_grad + # The forward also publishes the displacement as dofs.pos[dof_start] = qpos - qpos0, whose grad + # (accumulated by the force-stage reverses) folds back linearly. + rigid_info.qpos.grad[q_start, i_b] = ( + rigid_info.qpos.grad[q_start, i_b] + displacement_grad + dyn_state.dofs.pos.grad[dof_start, i_b] + ) + dyn_state.dofs.pos.grad[dof_start, i_b] = 0.0 g_xaxis = qd.Vector( [ g_pos[0] * displacement + xaxis_grad[0], @@ -141,13 +163,15 @@ def kernel_manual_forward_kinematics_bw( ], dt=gs.qd_float, ) - g_xanchor = g_pos + xanchor_grad + # pos_out never reads xanchor here (pos_out = pos_in + xaxis * displacement), so only the + # downstream xanchor consumer's grad flows through xanchor = T(joint_pos_off, quat_in) + pos_in + # into quat_in; g_pos itself reaches pos_in directly. g_quat_in = ( g_quat + gu.qd_transform_by_quat_grad_quat(axis, quat_in, g_xaxis) - + gu.qd_transform_by_quat_grad_quat(joint_pos_off, quat_in, g_xanchor) + + gu.qd_transform_by_quat_grad_quat(joint_pos_off, quat_in, xanchor_grad) ) - g_pos = g_xanchor + g_pos = g_pos + xanchor_grad g_quat = g_quat_in elif joint_type == gs.JOINT_TYPE.SPHERICAL: @@ -167,8 +191,21 @@ def kernel_manual_forward_kinematics_bw( gq_out = g_quat - gu.qd_transform_by_quat_grad_quat(joint_pos_off, quat_out, g_pos) g_qloc = gu.qd_quat_mul_grad_rhs(quat_in, qloc, gq_out) g_quat_in_apply = gu.qd_quat_mul_grad_lhs(quat_in, qloc, gq_out) + # The forward also publishes the joint coordinate as dofs.pos = quat_to_xyz(qloc), whose grad + # (accumulated by the force-stage reverses) folds back through the Euler extraction. + g_xyz = qd.Vector( + [ + dyn_state.dofs.pos.grad[dof_start, i_b], + dyn_state.dofs.pos.grad[dof_start + 1, i_b], + dyn_state.dofs.pos.grad[dof_start + 2, i_b], + ], + dt=gs.qd_float, + ) + g_qloc = g_qloc + gu.qd_quat_to_xyz_grad_quat(qloc, rigid_info.EPS[None], g_xyz) for j in qd.static(range(4)): rigid_info.qpos.grad[q_start + j, i_b] = rigid_info.qpos.grad[q_start + j, i_b] + g_qloc[j] + for j in qd.static(range(3)): + dyn_state.dofs.pos.grad[dof_start + j, i_b] = 0.0 g_xanchor = g_pos + xanchor_grad g_quat_in = ( g_quat_in_apply @@ -189,27 +226,58 @@ def kernel_manual_forward_kinematics_bw( # arm_base_pos = parent_pos + transform(link_offset_pos, parent_quat) # arm_base_quat = quat_mul(parent_quat, link_offset_quat) # propagating slot-0 grad (g_pos, g_quat) into the parent's pose grad. - if parent_idx != -1: - parent_quat = dyn_state.links.quat[parent_idx, i_b] + if i_parent != -1: + parent_quat = dyn_state.links.quat[i_parent, i_b] link_off_pos = dyn_info.links.pos[I_l] link_off_quat = dyn_info.links.quat[I_l] parent_quat_grad_from_pos = gu.qd_transform_by_quat_grad_quat(link_off_pos, parent_quat, g_pos) parent_quat_grad_from_quat = gu.qd_quat_mul_grad_lhs(parent_quat, link_off_quat, g_quat) for j in qd.static(range(3)): - dyn_state.links.pos.grad[parent_idx, i_b][j] = ( - dyn_state.links.pos.grad[parent_idx, i_b][j] + g_pos[j] - ) + dyn_state.links.pos.grad[i_parent, i_b][j] = dyn_state.links.pos.grad[i_parent, i_b][j] + g_pos[j] for j in qd.static(range(4)): - dyn_state.links.quat.grad[parent_idx, i_b][j] = ( - dyn_state.links.quat.grad[parent_idx, i_b][j] + dyn_state.links.quat.grad[i_parent, i_b][j] = ( + dyn_state.links.quat.grad[i_parent, i_b][j] + parent_quat_grad_from_pos[j] + parent_quat_grad_from_quat[j] ) - for j in qd.static(range(3)): - dyn_state.links.pos.grad[i_l, i_b][j] = 0.0 - for j in qd.static(range(4)): - dyn_state.links.quat.grad[i_l, i_b][j] = 0.0 + # The forward skips writing fixed root links' pose so users can overwrite it (see + # func_forward_kinematics_entity): their pose is persistent input state whose grad must keep + # accumulating across substeps instead of being consumed here. + if not (i_parent == -1 and dyn_info.links.is_fixed[I_l]): + for j in qd.static(range(3)): + dyn_state.links.pos.grad[i_l, i_b][j] = 0.0 + for j in qd.static(range(4)): + dyn_state.links.quat.grad[i_l, i_b][j] = 0.0 + + +@qd.func +def func_cd_contraction_bw( + i_b, + i_l, + i_slot, + dof_lo, + dof_hi, + dyn_state: array_class.DynState, +): + """Reverse of the forward-velocity contraction atomic_add(cd_*_bw[i_l, i_slot], cdof_*[i_d] * vel[i_d]) over + dofs [dof_lo, dof_hi): accumulates cd_*_bw[i_slot].grad into cdof_{ang,vel}.grad and vel.grad.""" + g_cd_vel = dyn_state.links.cd_vel_bw.grad[i_l, i_slot, i_b] + g_cd_ang = dyn_state.links.cd_ang_bw.grad[i_l, i_slot, i_b] + for i_d in range(dof_lo, dof_hi): + dof_vel = dyn_state.dofs.vel[i_d, i_b] + cdof_vel = dyn_state.dofs.cdof_vel[i_d, i_b] + cdof_ang = dyn_state.dofs.cdof_ang[i_d, i_b] + for k in qd.static(range(3)): + dyn_state.dofs.cdof_vel.grad[i_d, i_b][k] = ( + dyn_state.dofs.cdof_vel.grad[i_d, i_b][k] + g_cd_vel[k] * dof_vel + ) + dyn_state.dofs.cdof_ang.grad[i_d, i_b][k] = ( + dyn_state.dofs.cdof_ang.grad[i_d, i_b][k] + g_cd_ang[k] * dof_vel + ) + dot_vel = cdof_vel[0] * g_cd_vel[0] + cdof_vel[1] * g_cd_vel[1] + cdof_vel[2] * g_cd_vel[2] + dot_ang = cdof_ang[0] * g_cd_ang[0] + cdof_ang[1] * g_cd_ang[1] + cdof_ang[2] * g_cd_ang[2] + dyn_state.dofs.vel.grad[i_d, i_b] = dyn_state.dofs.vel.grad[i_d, i_b] + dot_vel + dot_ang @qd.kernel(fastcache=True) @@ -218,7 +286,6 @@ def kernel_manual_forward_velocity_bw( dyn_info: array_class.DynInfo, rigid_info: array_class.RigidInfo, rigid_config: qd.template(), - errno: qd.Tensor, ): """Manual reverse of kernel_forward_velocity: link velocity grads back to dof velocity and cdof grads. @@ -231,226 +298,126 @@ def kernel_manual_forward_velocity_bw( serialize=qd.static(rigid_config.para_level < gs.PARA_LEVEL.PARTIAL), ) for i_e, i_b in qd.ndrange(dyn_info.entities.n_links.shape[0], dyn_state.links.pos.shape[1]): - if qd.static(rigid_config.use_hibernation): - errno[i_b] = errno[i_b] | array_class.ErrorCode.MANUAL_BW_UNIMPLEMENTED - else: - n_in_e = dyn_info.entities.n_links[i_e] - # Leaf -> root iteration so each link's cd_*_bw[0].grad (which - # accumulates into parent.cd_*.grad) is propagated *before* the - # parent's own iteration uses it. - for i_l_rev in range(n_in_e): - i_l = dyn_info.entities.link_end[i_e] - 1 - i_l_rev - I_l = [i_l, i_b] if qd.static(rigid_config.batch_links_info) else i_l - n_joints = dyn_info.links.joint_end[I_l] - dyn_info.links.joint_start[I_l] - i_parent = dyn_info.links.parent_idx[I_l] - - # --- Step 1 reverse: cd_*[i_l].grad -> cd_*_bw[i_l, n_joints].grad + n_in_e = dyn_info.entities.n_links[i_e] + # Leaf -> root iteration so each link's cd_*_bw[0].grad (which + # accumulates into parent.cd_*.grad) is propagated *before* the + # parent's own iteration uses it. + for i_l_rev in range(n_in_e): + i_l = dyn_info.entities.link_end[i_e] - 1 - i_l_rev + I_l = [i_l, i_b] if qd.static(rigid_config.batch_links_info) else i_l + n_joints = dyn_info.links.joint_end[I_l] - dyn_info.links.joint_start[I_l] + i_parent = dyn_info.links.parent_idx[I_l] + + # --- Step 1 reverse: cd_*[i_l].grad -> cd_*_bw[i_l, n_joints].grad + for k in qd.static(range(3)): + dyn_state.links.cd_vel_bw.grad[i_l, n_joints, i_b][k] = ( + dyn_state.links.cd_vel_bw.grad[i_l, n_joints, i_b][k] + dyn_state.links.cd_vel.grad[i_l, i_b][k] + ) + dyn_state.links.cd_ang_bw.grad[i_l, n_joints, i_b][k] = ( + dyn_state.links.cd_ang_bw.grad[i_l, n_joints, i_b][k] + dyn_state.links.cd_ang.grad[i_l, i_b][k] + ) + # consume cd_vel/cd_ang.grad[i_l] + for k in qd.static(range(3)): + dyn_state.links.cd_vel.grad[i_l, i_b][k] = 0.0 + dyn_state.links.cd_ang.grad[i_l, i_b][k] = 0.0 + + # --- Step 2: iterate joints in reverse + for i_j_rev in range(n_joints): + i_j_ = n_joints - 1 - i_j_rev + i_j = i_j_ + dyn_info.links.joint_start[I_l] + I_j = [i_j, i_b] if qd.static(rigid_config.batch_joints_info) else i_j + joint_type = dyn_info.joints.type[I_j] + dof_start = dyn_info.joints.dof_start[I_j] + dof_end = dyn_info.joints.dof_end[I_j] + i_slot = i_j_ + i_slot_next = i_j_ + 1 + + # --- Reverse the joint-space contraction into cd_*_bw[next]. + # Forward (FREE contracts only its angular dofs here; the linear ones go through slot curr): + # _vel = cdof_vel[d] * vel[d]; atomic_add(cd_vel_bw[next], _vel) + # _ang = cdof_ang[d] * vel[d]; atomic_add(cd_ang_bw[next], _ang) + dof_lo = dof_start + 3 if joint_type == gs.JOINT_TYPE.FREE else dof_start + func_cd_contraction_bw(i_b, i_l, i_slot_next, dof_lo, dof_end, dyn_state) + g_cd_vel_next = dyn_state.links.cd_vel_bw.grad[i_l, i_slot_next, i_b] + g_cd_ang_next = dyn_state.links.cd_ang_bw.grad[i_l, i_slot_next, i_b] + + # --- Reverse the slot copy cd_*_bw[next] = cd_*_bw[curr] for k in qd.static(range(3)): - dyn_state.links.cd_vel_bw.grad[i_l, n_joints, i_b][k] = ( - dyn_state.links.cd_vel_bw.grad[i_l, n_joints, i_b][k] + dyn_state.links.cd_vel.grad[i_l, i_b][k] + dyn_state.links.cd_vel_bw.grad[i_l, i_slot, i_b][k] = ( + dyn_state.links.cd_vel_bw.grad[i_l, i_slot, i_b][k] + g_cd_vel_next[k] ) - dyn_state.links.cd_ang_bw.grad[i_l, n_joints, i_b][k] = ( - dyn_state.links.cd_ang_bw.grad[i_l, n_joints, i_b][k] + dyn_state.links.cd_ang.grad[i_l, i_b][k] + dyn_state.links.cd_ang_bw.grad[i_l, i_slot, i_b][k] = ( + dyn_state.links.cd_ang_bw.grad[i_l, i_slot, i_b][k] + g_cd_ang_next[k] ) - # consume cd_vel/cd_ang.grad[i_l] + # consume next for k in qd.static(range(3)): - dyn_state.links.cd_vel.grad[i_l, i_b][k] = 0.0 - dyn_state.links.cd_ang.grad[i_l, i_b][k] = 0.0 - - # --- Step 2: iterate joints in reverse - for i_j_rev in range(n_joints): - i_j_ = n_joints - 1 - i_j_rev - i_j = i_j_ + dyn_info.links.joint_start[I_l] - I_j = [i_j, i_b] if qd.static(rigid_config.batch_joints_info) else i_j - joint_type = dyn_info.joints.type[I_j] - dof_start = dyn_info.joints.dof_start[I_j] - dof_end = dyn_info.joints.dof_end[I_j] - curr_idx = i_j_ - next_idx = i_j_ + 1 - - # --- [d-rev] cd_*_bw[next].grad -> cdof_*.grad / vel.grad - # Forward (FREE angular: i_3=0..2 at d=dof_start+3+i_3; else: d in ds..de): - # _vel = cdof_vel[d] * vel[d]; atomic_add(cd_vel_bw[next], _vel) - # _ang = cdof_ang[d] * vel[d]; atomic_add(cd_ang_bw[next], _ang) - g_cd_vel_next = dyn_state.links.cd_vel_bw.grad[i_l, next_idx, i_b] - g_cd_ang_next = dyn_state.links.cd_ang_bw.grad[i_l, next_idx, i_b] - if joint_type == gs.JOINT_TYPE.FREE: - for i_3 in qd.static(range(3)): - i_d = dof_start + 3 + i_3 - dof_vel = dyn_state.dofs.vel[i_d, i_b] - cdof_vel = dyn_state.dofs.cdof_vel[i_d, i_b] - cdof_ang = dyn_state.dofs.cdof_ang[i_d, i_b] - for k in qd.static(range(3)): - dyn_state.dofs.cdof_vel.grad[i_d, i_b][k] = ( - dyn_state.dofs.cdof_vel.grad[i_d, i_b][k] + g_cd_vel_next[k] * dof_vel - ) - dyn_state.dofs.cdof_ang.grad[i_d, i_b][k] = ( - dyn_state.dofs.cdof_ang.grad[i_d, i_b][k] + g_cd_ang_next[k] * dof_vel - ) - dot_vel = ( - cdof_vel[0] * g_cd_vel_next[0] - + cdof_vel[1] * g_cd_vel_next[1] - + cdof_vel[2] * g_cd_vel_next[2] - ) - dot_ang = ( - cdof_ang[0] * g_cd_ang_next[0] - + cdof_ang[1] * g_cd_ang_next[1] - + cdof_ang[2] * g_cd_ang_next[2] - ) - dyn_state.dofs.vel.grad[i_d, i_b] = dyn_state.dofs.vel.grad[i_d, i_b] + dot_vel + dot_ang - else: - for i_d in range(dof_start, dof_end): - dof_vel = dyn_state.dofs.vel[i_d, i_b] - cdof_vel = dyn_state.dofs.cdof_vel[i_d, i_b] - cdof_ang = dyn_state.dofs.cdof_ang[i_d, i_b] - for k in qd.static(range(3)): - dyn_state.dofs.cdof_vel.grad[i_d, i_b][k] = ( - dyn_state.dofs.cdof_vel.grad[i_d, i_b][k] + g_cd_vel_next[k] * dof_vel - ) - dyn_state.dofs.cdof_ang.grad[i_d, i_b][k] = ( - dyn_state.dofs.cdof_ang.grad[i_d, i_b][k] + g_cd_ang_next[k] * dof_vel - ) - dot_vel = ( - cdof_vel[0] * g_cd_vel_next[0] - + cdof_vel[1] * g_cd_vel_next[1] - + cdof_vel[2] * g_cd_vel_next[2] - ) - dot_ang = ( - cdof_ang[0] * g_cd_ang_next[0] - + cdof_ang[1] * g_cd_ang_next[1] - + cdof_ang[2] * g_cd_ang_next[2] - ) - dyn_state.dofs.vel.grad[i_d, i_b] = dyn_state.dofs.vel.grad[i_d, i_b] + dot_vel + dot_ang - - # --- [c-rev] cd_*_bw[next] = cd_*_bw[curr] -> curr.grad += next.grad + dyn_state.links.cd_vel_bw.grad[i_l, i_slot_next, i_b][k] = 0.0 + dyn_state.links.cd_ang_bw.grad[i_l, i_slot_next, i_b][k] = 0.0 + + # --- Reverse motion_cross_motion: + # Forward: (cdofd_ang[i_d], cdofd_vel[i_d]) = + # motion_cross_motion(cd_ang_bw[curr], cd_vel_bw[curr], cdof_ang[i_d], cdof_vel[i_d]) + # over the same dof range as the contraction above (FREE applies it to its angular dofs only). + s_ang_primal = dyn_state.links.cd_ang_bw[i_l, i_slot, i_b] + s_vel_primal = dyn_state.links.cd_vel_bw[i_l, i_slot, i_b] + for i_d in range(dof_lo, dof_end): + g_cdofd_ang = dyn_state.dofs.cdofd_ang.grad[i_d, i_b] + g_cdofd_vel = dyn_state.dofs.cdofd_vel.grad[i_d, i_b] + cdof_ang = dyn_state.dofs.cdof_ang[i_d, i_b] + cdof_vel = dyn_state.dofs.cdof_vel[i_d, i_b] + g_cd_ang, g_cd_vel, g_cdof_ang, g_cdof_vel = gu.motion_cross_motion_grad( + s_ang_primal, s_vel_primal, cdof_ang, cdof_vel, g_cdofd_ang, g_cdofd_vel + ) for k in qd.static(range(3)): - dyn_state.links.cd_vel_bw.grad[i_l, curr_idx, i_b][k] = ( - dyn_state.links.cd_vel_bw.grad[i_l, curr_idx, i_b][k] + g_cd_vel_next[k] + dyn_state.links.cd_ang_bw.grad[i_l, i_slot, i_b][k] = ( + dyn_state.links.cd_ang_bw.grad[i_l, i_slot, i_b][k] + g_cd_ang[k] ) - dyn_state.links.cd_ang_bw.grad[i_l, curr_idx, i_b][k] = ( - dyn_state.links.cd_ang_bw.grad[i_l, curr_idx, i_b][k] + g_cd_ang_next[k] + dyn_state.links.cd_vel_bw.grad[i_l, i_slot, i_b][k] = ( + dyn_state.links.cd_vel_bw.grad[i_l, i_slot, i_b][k] + g_cd_vel[k] ) - # consume next - for k in qd.static(range(3)): - dyn_state.links.cd_vel_bw.grad[i_l, next_idx, i_b][k] = 0.0 - dyn_state.links.cd_ang_bw.grad[i_l, next_idx, i_b][k] = 0.0 - - # --- [b-rev] motion_cross_motion reverse: - # Forward: (cdofd_ang[i_d], cdofd_vel[i_d]) = - # motion_cross_motion(cd_ang_bw[curr], cd_vel_bw[curr], cdof_ang[i_d], cdof_vel[i_d]) - # Reverse via gu.motion_cross_motion_grad(s_ang, s_vel, m_ang, m_vel, g_cdofd_ang, g_cdofd_vel) - s_ang_primal = dyn_state.links.cd_ang_bw[i_l, curr_idx, i_b] - s_vel_primal = dyn_state.links.cd_vel_bw[i_l, curr_idx, i_b] - if joint_type == gs.JOINT_TYPE.FREE: - # Angular dofs i_3=0..2 at i_d = dof_start + 3 + i_3 (linear cdofd_* are explicit 0) - for i_3 in qd.static(range(3)): - i_d = dof_start + 3 + i_3 - g_cdofd_ang = dyn_state.dofs.cdofd_ang.grad[i_d, i_b] - g_cdofd_vel = dyn_state.dofs.cdofd_vel.grad[i_d, i_b] - cdof_ang = dyn_state.dofs.cdof_ang[i_d, i_b] - cdof_vel = dyn_state.dofs.cdof_vel[i_d, i_b] - g_cd_ang, g_cd_vel, g_cdof_ang, g_cdof_vel = gu.motion_cross_motion_grad( - s_ang_primal, s_vel_primal, cdof_ang, cdof_vel, g_cdofd_ang, g_cdofd_vel - ) - for k in qd.static(range(3)): - dyn_state.links.cd_ang_bw.grad[i_l, curr_idx, i_b][k] = ( - dyn_state.links.cd_ang_bw.grad[i_l, curr_idx, i_b][k] + g_cd_ang[k] - ) - dyn_state.links.cd_vel_bw.grad[i_l, curr_idx, i_b][k] = ( - dyn_state.links.cd_vel_bw.grad[i_l, curr_idx, i_b][k] + g_cd_vel[k] - ) - dyn_state.dofs.cdof_ang.grad[i_d, i_b][k] = ( - dyn_state.dofs.cdof_ang.grad[i_d, i_b][k] + g_cdof_ang[k] - ) - dyn_state.dofs.cdof_vel.grad[i_d, i_b][k] = ( - dyn_state.dofs.cdof_vel.grad[i_d, i_b][k] + g_cdof_vel[k] - ) - # consume cdofd_*.grad[i_d] - for k in qd.static(range(3)): - dyn_state.dofs.cdofd_ang.grad[i_d, i_b][k] = 0.0 - dyn_state.dofs.cdofd_vel.grad[i_d, i_b][k] = 0.0 - # Linear dofs (i_3=0..2 at i_d = dof_start + i_3): cdofd_* set to 0 - # (constant), reverse is no-op; just consume to mirror P8. - for i_3 in qd.static(range(3)): - i_d = dof_start + i_3 - for k in qd.static(range(3)): - dyn_state.dofs.cdofd_ang.grad[i_d, i_b][k] = 0.0 - dyn_state.dofs.cdofd_vel.grad[i_d, i_b][k] = 0.0 - else: - for i_d in range(dof_start, dof_end): - g_cdofd_ang = dyn_state.dofs.cdofd_ang.grad[i_d, i_b] - g_cdofd_vel = dyn_state.dofs.cdofd_vel.grad[i_d, i_b] - cdof_ang = dyn_state.dofs.cdof_ang[i_d, i_b] - cdof_vel = dyn_state.dofs.cdof_vel[i_d, i_b] - g_cd_ang, g_cd_vel, g_cdof_ang, g_cdof_vel = gu.motion_cross_motion_grad( - s_ang_primal, s_vel_primal, cdof_ang, cdof_vel, g_cdofd_ang, g_cdofd_vel - ) - for k in qd.static(range(3)): - dyn_state.links.cd_ang_bw.grad[i_l, curr_idx, i_b][k] = ( - dyn_state.links.cd_ang_bw.grad[i_l, curr_idx, i_b][k] + g_cd_ang[k] - ) - dyn_state.links.cd_vel_bw.grad[i_l, curr_idx, i_b][k] = ( - dyn_state.links.cd_vel_bw.grad[i_l, curr_idx, i_b][k] + g_cd_vel[k] - ) - dyn_state.dofs.cdof_ang.grad[i_d, i_b][k] = ( - dyn_state.dofs.cdof_ang.grad[i_d, i_b][k] + g_cdof_ang[k] - ) - dyn_state.dofs.cdof_vel.grad[i_d, i_b][k] = ( - dyn_state.dofs.cdof_vel.grad[i_d, i_b][k] + g_cdof_vel[k] - ) - for k in qd.static(range(3)): - dyn_state.dofs.cdofd_ang.grad[i_d, i_b][k] = 0.0 - dyn_state.dofs.cdofd_vel.grad[i_d, i_b][k] = 0.0 - - # --- [a-rev] (FREE only) cd_*_bw[curr].grad -> linear cdof_*.grad / vel.grad - # Forward (FREE linear pre-motion_cross_motion): for i_3=0..2 at i_d = dof_start + i_3, - # _vel = cdof_vel[i_d] * vel[i_d]; atomic_add(cd_vel_bw[curr], _vel) - # _ang = cdof_ang[i_d] * vel[i_d]; atomic_add(cd_ang_bw[curr], _ang) - # (cdof_vel[linear] = e_i_3 constant; cdof_ang[linear] = 0 constant) - if joint_type == gs.JOINT_TYPE.FREE: - g_cd_vel_curr = dyn_state.links.cd_vel_bw.grad[i_l, curr_idx, i_b] - g_cd_ang_curr = dyn_state.links.cd_ang_bw.grad[i_l, curr_idx, i_b] - for i_3 in qd.static(range(3)): - i_d = dof_start + i_3 - dof_vel = dyn_state.dofs.vel[i_d, i_b] - cdof_vel = dyn_state.dofs.cdof_vel[i_d, i_b] - cdof_ang = dyn_state.dofs.cdof_ang[i_d, i_b] - for k in qd.static(range(3)): - dyn_state.dofs.cdof_vel.grad[i_d, i_b][k] = ( - dyn_state.dofs.cdof_vel.grad[i_d, i_b][k] + g_cd_vel_curr[k] * dof_vel - ) - dyn_state.dofs.cdof_ang.grad[i_d, i_b][k] = ( - dyn_state.dofs.cdof_ang.grad[i_d, i_b][k] + g_cd_ang_curr[k] * dof_vel - ) - dot_vel = ( - cdof_vel[0] * g_cd_vel_curr[0] - + cdof_vel[1] * g_cd_vel_curr[1] - + cdof_vel[2] * g_cd_vel_curr[2] - ) - dot_ang = ( - cdof_ang[0] * g_cd_ang_curr[0] - + cdof_ang[1] * g_cd_ang_curr[1] - + cdof_ang[2] * g_cd_ang_curr[2] - ) - dyn_state.dofs.vel.grad[i_d, i_b] = dyn_state.dofs.vel.grad[i_d, i_b] + dot_vel + dot_ang - - # --- Step 1 (initial cvel setup) reverse: - # Forward: cd_*_bw[i_l, 0, i_b] = parent.cd_*[i_parent, i_b] (if i_parent != -1) else 0 - # Reverse: parent.cd_*.grad[i_parent] += cd_*_bw[i_l, 0].grad; consume slot 0 - g_cd_vel_slot0 = dyn_state.links.cd_vel_bw.grad[i_l, 0, i_b] - g_cd_ang_slot0 = dyn_state.links.cd_ang_bw.grad[i_l, 0, i_b] - if i_parent != -1: - for k in qd.static(range(3)): - dyn_state.links.cd_vel.grad[i_parent, i_b][k] = ( - dyn_state.links.cd_vel.grad[i_parent, i_b][k] + g_cd_vel_slot0[k] + dyn_state.dofs.cdof_ang.grad[i_d, i_b][k] = ( + dyn_state.dofs.cdof_ang.grad[i_d, i_b][k] + g_cdof_ang[k] ) - dyn_state.links.cd_ang.grad[i_parent, i_b][k] = ( - dyn_state.links.cd_ang.grad[i_parent, i_b][k] + g_cd_ang_slot0[k] + dyn_state.dofs.cdof_vel.grad[i_d, i_b][k] = ( + dyn_state.dofs.cdof_vel.grad[i_d, i_b][k] + g_cdof_vel[k] ) - # consume slot 0 + # consume cdofd_*.grad[i_d] + for k in qd.static(range(3)): + dyn_state.dofs.cdofd_ang.grad[i_d, i_b][k] = 0.0 + dyn_state.dofs.cdofd_vel.grad[i_d, i_b][k] = 0.0 + if joint_type == gs.JOINT_TYPE.FREE: + # Linear dofs: the forward writes constant-zero cdofd, so the reverse is a no-op; consume the + # grads so they do not leak into the next substep. + for i_3 in qd.static(range(3)): + i_d = dof_start + i_3 + for k in qd.static(range(3)): + dyn_state.dofs.cdofd_ang.grad[i_d, i_b][k] = 0.0 + dyn_state.dofs.cdofd_vel.grad[i_d, i_b][k] = 0.0 + + # --- Reverse the FREE linear contraction into cd_*_bw[curr]. + # Forward (FREE linear pre-motion_cross_motion): for i_3=0..2 at i_d = dof_start + i_3, + # _vel = cdof_vel[i_d] * vel[i_d]; atomic_add(cd_vel_bw[curr], _vel) + # _ang = cdof_ang[i_d] * vel[i_d]; atomic_add(cd_ang_bw[curr], _ang) + if joint_type == gs.JOINT_TYPE.FREE: + func_cd_contraction_bw(i_b, i_l, i_slot, dof_start, dof_start + 3, dyn_state) + + # --- Step 1 (initial cvel setup) reverse: + # Forward: cd_*_bw[i_l, 0, i_b] = parent.cd_*[i_parent, i_b] (if i_parent != -1) else 0 + # Reverse: parent.cd_*.grad[i_parent] += cd_*_bw[i_l, 0].grad; consume slot 0 + g_cd_vel_slot0 = dyn_state.links.cd_vel_bw.grad[i_l, 0, i_b] + g_cd_ang_slot0 = dyn_state.links.cd_ang_bw.grad[i_l, 0, i_b] + if i_parent != -1: for k in qd.static(range(3)): - dyn_state.links.cd_vel_bw.grad[i_l, 0, i_b][k] = 0.0 - dyn_state.links.cd_ang_bw.grad[i_l, 0, i_b][k] = 0.0 + dyn_state.links.cd_vel.grad[i_parent, i_b][k] = ( + dyn_state.links.cd_vel.grad[i_parent, i_b][k] + g_cd_vel_slot0[k] + ) + dyn_state.links.cd_ang.grad[i_parent, i_b][k] = ( + dyn_state.links.cd_ang.grad[i_parent, i_b][k] + g_cd_ang_slot0[k] + ) + # consume slot 0 + for k in qd.static(range(3)): + dyn_state.links.cd_vel_bw.grad[i_l, 0, i_b][k] = 0.0 + dyn_state.links.cd_ang_bw.grad[i_l, 0, i_b][k] = 0.0 @qd.kernel(fastcache=True) diff --git a/genesis/engine/solvers/rigid/collider/contact.py b/genesis/engine/solvers/rigid/collider/contact.py index 58a1ad303..ed02b735b 100644 --- a/genesis/engine/solvers/rigid/collider/contact.py +++ b/genesis/engine/solvers/rigid/collider/contact.py @@ -542,7 +542,7 @@ def func_contact_orthogonals( @qd.func def func_rotate_frame( pos: qd.types.vector(3), quat: qd.types.vector(4), contact_pos: qd.types.vector(3), qrot: qd.types.vector(4) -) -> tuple[qd.types.vector(3, dtype=gs.qd_float), qd.types.vector(4, dtype=gs.qd_float)]: +) -> tuple[qd.types.vector(3), qd.types.vector(4)]: """ Instead of modifying geoms_state in place, this function takes thread-local pos/quat and returns the updated values. diff --git a/genesis/engine/solvers/rigid/collider/diff_gjk.py b/genesis/engine/solvers/rigid/collider/diff_gjk.py index befb93028..88c0bccb1 100644 --- a/genesis/engine/solvers/rigid/collider/diff_gjk.py +++ b/genesis/engine/solvers/rigid/collider/diff_gjk.py @@ -829,6 +829,30 @@ def func_differentiable_contact( return contact_pos, contact_normal, penetration, weight +@qd.func +def func_plane_contact_frame( + i_b, + i_ga, + i_gb, + dyn_state: array_class.DynState, + dyn_info: array_class.DynInfo, +): + """World contact normal and convex-side radius of a plane [i_ga] vs convex [i_gb] pair: normal = + -normalize(R(quat_plane) @ plane_local_dir), radius = data[0] for SPHERE / CAPSULE and 0 otherwise. + + Both the forward witness capture and the differentiable reconstruction share this frame (and the convention + contact_pos = v - 0.5 * penetration * normal built on it).""" + plane_dir = gs.qd_vec3(dyn_info.geoms.data[i_ga][0], dyn_info.geoms.data[i_ga][1], dyn_info.geoms.data[i_ga][2]) + plane_dir = gu.qd_transform_by_quat(plane_dir, dyn_state.geoms.quat[i_ga, i_b]) + normal = -plane_dir.normalized() + + radius = gs.qd_float(0.0) + geom_type = dyn_info.geoms.type[i_gb] + if geom_type == gs.GEOM_TYPE.SPHERE or geom_type == gs.GEOM_TYPE.CAPSULE: + radius = dyn_info.geoms.data[i_gb][0] + return normal, radius + + @qd.func def func_differentiable_plane_contact( i_ga, @@ -854,18 +878,10 @@ def func_differentiable_plane_contact( orientation gradient is zero, matching the rotation-invariant forward contact. """ trans_plane = dyn_state.geoms.pos[i_ga, i_b] - quat_plane = dyn_state.geoms.quat[i_ga, i_b] trans_convex = dyn_state.geoms.pos[i_gb, i_b] quat_convex = dyn_state.geoms.quat[i_gb, i_b] - plane_dir = gs.qd_vec3(dyn_info.geoms.data[i_ga][0], dyn_info.geoms.data[i_ga][1], dyn_info.geoms.data[i_ga][2]) - plane_dir = gu.qd_transform_by_quat(plane_dir, quat_plane) - normal = -plane_dir.normalized() - - radius = gs.qd_float(0.0) - geom_type = dyn_info.geoms.type[i_gb] - if geom_type == gs.GEOM_TYPE.SPHERE or geom_type == gs.GEOM_TYPE.CAPSULE: - radius = dyn_info.geoms.data[i_gb][0] + normal, radius = func_plane_contact_frame(i_b, i_ga, i_gb, dyn_state, dyn_info) core_local = diff_contact_input.core_local[i_b, i_c] core_world = gu.qd_transform_by_trans_quat(core_local, trans_convex, quat_convex) diff --git a/genesis/engine/solvers/rigid/collider/narrowphase.py b/genesis/engine/solvers/rigid/collider/narrowphase.py index 017912976..eb5e83c29 100644 --- a/genesis/engine/solvers/rigid/collider/narrowphase.py +++ b/genesis/engine/solvers/rigid/collider/narrowphase.py @@ -3098,20 +3098,10 @@ def kernel_fill_diff_contact_input_plane( i_ga = collider_state.contact_data.geom_a[i_c, i_b] i_gb = collider_state.contact_data.geom_b[i_c, i_b] if dyn_info.geoms.type[i_ga] == gs.GEOM_TYPE.PLANE: - quat_plane = dyn_state.geoms.quat[i_ga, i_b] trans_convex = dyn_state.geoms.pos[i_gb, i_b] quat_convex = dyn_state.geoms.quat[i_gb, i_b] - plane_dir = gs.qd_vec3( - dyn_info.geoms.data[i_ga][0], dyn_info.geoms.data[i_ga][1], dyn_info.geoms.data[i_ga][2] - ) - plane_dir = gu.qd_transform_by_quat(plane_dir, quat_plane) - normal = -plane_dir.normalized() - - radius = gs.qd_float(0.0) - geom_type = dyn_info.geoms.type[i_gb] - if geom_type == gs.GEOM_TYPE.SPHERE or geom_type == gs.GEOM_TYPE.CAPSULE: - radius = dyn_info.geoms.data[i_gb][0] + normal, radius = diff_gjk.func_plane_contact_frame(i_b, i_ga, i_gb, dyn_state, dyn_info) penetration = collider_state.contact_data.penetration[i_c, i_b] contact_pos = collider_state.contact_data.pos[i_c, i_b] diff --git a/genesis/engine/solvers/rigid/constraint/backward.py b/genesis/engine/solvers/rigid/constraint/backward.py index 70b3f9bc2..d92ab2f74 100644 --- a/genesis/engine/solvers/rigid/constraint/backward.py +++ b/genesis/engine/solvers/rigid/constraint/backward.py @@ -4,6 +4,8 @@ import genesis.utils.array_class as array_class import genesis.utils.geom as gu +from . import solver + @qd.func def func_matvec_Ap( @@ -54,14 +56,14 @@ def func_matvec_Ap( @qd.func -def func_solve_adjoint_u_cg_env( +def func_solve_adjoint_u_cg_batch( i_b, constraint_state: array_class.ConstraintState, dyn_info: array_class.DynInfo, rigid_info: array_class.RigidInfo, rigid_config: qd.template(), ): - """CG solve of A u = g for a single environment [i_b]. + """Conjugate-gradient (CG) solve of A u = g for a single environment [i_b]. A = M + J^T diag(D) J is applied implicitly by func_matvec_Ap, which reads rigid_info.mass_mat directly and loops only over the active constraints, so this also solves the unconstrained case A = M (empty J term). @@ -69,35 +71,43 @@ def func_solve_adjoint_u_cg_env( n_dofs = constraint_state.bw_u.shape[0] # r = g - A*0 = g ; p = r ; u = 0 + num = gs.qd_float(0.0) for i_d in range(n_dofs): constraint_state.bw_u[i_d, i_b] = 0.0 constraint_state.bw_r[i_d, i_b] = constraint_state.dL_dqacc[i_d, i_b] constraint_state.bw_p[i_d, i_b] = constraint_state.bw_r[i_d, i_b] - - for it in range(rigid_info.iterations[None]): + num += constraint_state.bw_r[i_d, i_b] * constraint_state.bw_r[i_d, i_b] + + # The stopping target is relative to the seed |g|^2: an absolute threshold either exits at a huge relative + # residual when g is small (each backward substep shrinks the upstream gradient by roughly the loss scale) or, + # past convergence, lets the clamped alpha / beta denominators inject garbage steps that corrupt u. The + # denominator break exits once p collapses to the round-off floor, where p^T A p underflows for a positive + # semi-definite (PSD) A; alpha and beta then never need clamping. + num_target = num * rigid_info.EPS[None] * rigid_info.EPS[None] + for _ in range(rigid_info.iterations[None]): + if num <= num_target: + break func_matvec_Ap(i_b, constraint_state, dyn_info, rigid_info, rigid_config) # alpha = (r,r)/(p,Ap) - num = gs.qd_float(0.0) den = gs.qd_float(0.0) for i_d in range(n_dofs): - num += constraint_state.bw_r[i_d, i_b] * constraint_state.bw_r[i_d, i_b] den += constraint_state.bw_p[i_d, i_b] * constraint_state.bw_Ap[i_d, i_b] - alpha = num / qd.max(den, rigid_info.EPS[None]) + if den <= 0.0: + break + alpha = num / den # u += alpha p ; r -= alpha Ap for i_d in range(n_dofs): constraint_state.bw_u[i_d, i_b] += alpha * constraint_state.bw_p[i_d, i_b] constraint_state.bw_r[i_d, i_b] -= alpha * constraint_state.bw_Ap[i_d, i_b] - if num < rigid_info.EPS[None]: - break - # beta = (r_new,r_new)/(r_old,r_old) ; p = r + beta p num_new = gs.qd_float(0.0) for i_d in range(n_dofs): num_new += constraint_state.bw_r[i_d, i_b] * constraint_state.bw_r[i_d, i_b] - beta = num_new / qd.max(num, rigid_info.EPS[None]) + beta = num_new / num + num = num_new for i_d in range(n_dofs): constraint_state.bw_p[i_d, i_b] = constraint_state.bw_r[i_d, i_b] + beta * constraint_state.bw_p[i_d, i_b] @@ -132,7 +142,7 @@ def kernel_solve_adjoint_u( # No active constraint: A = M. The forward's constrained-Hessian Cholesky nt_H is unreliable for # these envs (the GPU tiled factorization skips them), so solve M u = g via CG, which reads mass_mat # directly and never touches nt_H. - func_solve_adjoint_u_cg_env(i_b, constraint_state, dyn_info, rigid_info, rigid_config) + func_solve_adjoint_u_cg_batch(i_b, constraint_state, dyn_info, rigid_info, rigid_config) else: # Reuse the forward's Cholesky decomposition A = L * L^T to solve A u = g. # z = L^{-1} g (forward substitution); saved to bw_r @@ -154,7 +164,7 @@ def kernel_solve_adjoint_u( else: # CG solver for A * u = g (parallelized over the batch dimension). for i_b in range(_B): - func_solve_adjoint_u_cg_env(i_b, constraint_state, dyn_info, rigid_info, rigid_config) + func_solve_adjoint_u_cg_batch(i_b, constraint_state, dyn_info, rigid_info, rigid_config) @qd.kernel @@ -279,7 +289,7 @@ def kernel_load_dL_dqacc_from_acc_grad( n_dofs = dyn_state.dofs.acc.shape[0] qd.loop_config( name="kernel_load_dL_dqacc_from_acc_grad", - serialize=qd.static(rigid_config.para_level < gs.PARA_LEVEL.ALL), + serialize=qd.static(rigid_config.para_level < gs.PARA_LEVEL.PARTIAL), ) for i_d, i_b in qd.ndrange(n_dofs, _B): constraint_state.dL_dqacc[i_d, i_b] = dyn_state.dofs.acc.grad[i_d, i_b] @@ -301,7 +311,7 @@ def kernel_accumulate_constraint_solver_grads( n_dofs = dyn_state.dofs.force.shape[0] qd.loop_config( name="kernel_accumulate_constraint_solver_grads", - serialize=qd.static(rigid_config.para_level < gs.PARA_LEVEL.ALL), + serialize=qd.static(rigid_config.para_level < gs.PARA_LEVEL.PARTIAL), ) for i_d, i_b in qd.ndrange(n_dofs, _B): dyn_state.dofs.force.grad[i_d, i_b] += constraint_state.dL_dforce[i_d, i_b] @@ -310,33 +320,29 @@ def kernel_accumulate_constraint_solver_grads( # --------------------------------------------------------------------------- -# Manual reverses of the inequality constraints (frictionloss, collision, -# joint-limit). Shared conventions for the kernels below. +# Manual reverses of the constraint-row assembly (equality, frictionloss, collision, joint-limit). Shared +# conventions for the kernels below. # -# Why manual (not autograd): the constraint rows are built inside the forward -# solver with a data-dependent count and ordering -- n_con is assigned by -# atomic_add as active constraints are discovered -- which autograd cannot -# differentiate cleanly (the row index is not a static, taped quantity). +# These reverses are manual because the constraint rows are built inside the forward solver with a data-dependent +# count and ordering - n_con is assigned by atomic_add as active constraints are discovered - and the row index is +# a runtime quantity outside what autograd can tape. # -# Upstream grads: kernel_compute_gradients populates, per constraint row -# n_con, constraint_state.dL_daref[n_con] (dL/d aref), dL_defc_D[n_con] -# (dL/d efc_D), and dL_djac[n_con, i_d] (dL/d jac). The collision reverse uses -# dL_djac; the frictionloss and joint-limit reverses ignore it (their jac -# entries are constants -- frictionloss is 1.0, joint-limit is piecewise +-1 -- -# so the sub-gradient w.r.t. jac is 0). Each kernel consumes these and +# Upstream grads: kernel_compute_gradients populates, per constraint row n_con, constraint_state.dL_daref[n_con] +# (dL/d aref), dL_defc_D[n_con] (dL/d efc_D), and dL_djac[n_con, i_d] (dL/d jac). The collision and equality +# reverses use dL_djac; the frictionloss and joint-limit reverses ignore it (their jac entries are the constants +# 1.0 and piecewise +-1 respectively, so the sub-gradient w.r.t. jac is 0). Each kernel consumes these and # accumulates into its own differentiable inputs. # # n_con row layout: the forward adds constraints in the order equality -> frictionloss -> collision -> joint-limit -# (see add_equality_constraints / add_inequality_constraints in solver.py). Equality sub-types CONNECT and WELD are -# rejected host-side (not yet differentiated); JOINT is differentiated by -# kernel_manual_add_equality_constraints_bw. The manual reverses re-walk the same forward loops deterministically to -# recover their own n_con (no atomic_add, no n_constraints reset): +# (see add_equality_constraints / add_inequality_constraints in solver.py). All three equality sub-types (JOINT, +# CONNECT, WELD) are differentiated by kernel_manual_add_equality_constraints_bw. The manual reverses re-walk the +# same forward loops deterministically to recover their own n_con: # n_eq = constraint_state.n_constraints_equality[i_b] # n_fric = constraint_state.n_constraints_frictionloss[i_b] -# equality (JOINT) : seed counter at 0 -# frictionloss : seed counter at n_eq -# collision : n_con = n_eq + n_fric + i_col_ * 4 + i, with i_col_ the logical (sorted) contact index -# joint-limit : seed counter at n_eq + n_fric (+ 4 * n_contacts if collision on) +# equality : seed counter at 0 +# frictionloss : seed counter at n_eq +# collision : n_con = n_eq + n_fric + i_col_ * 4 + i, with i_col_ the logical (sorted) contact index +# joint-limit : seed counter at n_eq + n_fric (+ 4 * n_contacts if collision on) # --------------------------------------------------------------------------- @qd.kernel(fastcache=True) def kernel_manual_add_joint_limit_constraints_bw( @@ -396,7 +402,9 @@ def kernel_manual_add_joint_limit_constraints_bw( constraint_state.n_constraints_equality[i_b] + constraint_state.n_constraints_frictionloss[i_b] ) if qd.static(enable_collision): - n_con_counter = n_con_counter + gs.qd_int(collider_state.n_contacts[i_b] * 4) + n_con_counter = n_con_counter + gs.qd_int( + collider_state.n_contacts[i_b] * qd.static(rigid_config.rows_per_contact) + ) for i_l in range(n_links): I_l = [i_l, i_b] if qd.static(rigid_config.batch_links_info) else i_l @@ -424,26 +432,8 @@ def kernel_manual_add_joint_limit_constraints_bw( sign_f = gs.qd_float(sign_pos) sol_params = dyn_info.joints.sol_params[I_j] - timeconst = sol_params[0] - dampratio = sol_params[1] - dmin = sol_params[2] - dmax = sol_params[3] + imp, b_coef, k_coef, d_imp_d_imp_x = gu.imp_aref_grad(sol_params, pos_delta) width = sol_params[4] - mid = sol_params[5] - power = sol_params[6] - - imp_x = qd.abs(pos_delta) / width - imp_a_coef = 1.0 / mid ** (power - 1.0) - imp_b_coef = 1.0 / (1.0 - mid) ** (power - 1.0) - imp_a = imp_a_coef * imp_x**power - imp_b = 1.0 - imp_b_coef * (1.0 - imp_x) ** power - imp_y = imp_a if imp_x < mid else imp_b - imp_raw = dmin + imp_y * (dmax - dmin) - imp_clamped = qd.math.clamp(imp_raw, dmin, dmax) - imp = dmax if imp_x > 1.0 else imp_clamped - - b_coef = 2.0 / (dmax * timeconst) - k_coef = 1.0 / (dmax * dmax * timeconst * timeconst * dampratio * dampratio) invweight = dyn_info.dofs.invweight[I_d] diag_raw = invweight * (1.0 - imp) / imp @@ -459,28 +449,15 @@ def kernel_manual_add_joint_limit_constraints_bw( d_aref_d_jac_qvel = -b_coef d_aref_d_pos_delta_direct = -k_coef * imp - # diag_raw = invweight*(1-imp)/imp => d(diag_raw)/d(imp) = -invweight/imp^2 - # diag = max(diag_raw, EPS); efc_D = 1/diag - # d(efc_D)/d(imp) = -1/diag^2 * d(diag)/d(imp), 0 if clamped to EPS + # diag_raw = invweight * (1-imp)/imp gives d(diag_raw)/d(imp) = -invweight/imp^2; + # diag = max(diag_raw, EPS) and efc_D = 1/diag, with zero derivative when clamped at EPS. d_diag_d_imp = gs.qd_float(0.0) if diag_raw > EPS: d_diag_d_imp = -invweight / (imp * imp) d_efc_D_d_imp = -d_diag_d_imp / (diag * diag) - # d(imp)/d(imp_x): active only inside the smooth clamp band. - within_clamp = (imp_raw > dmin) and (imp_raw < dmax) and (imp_x <= 1.0) - d_imp_y_d_imp_x = gs.qd_float(0.0) - if imp_x < mid: - d_imp_y_d_imp_x = power * imp_a_coef * imp_x ** (power - 1.0) - else: - d_imp_y_d_imp_x = power * imp_b_coef * (1.0 - imp_x) ** (power - 1.0) - d_imp_d_imp_x = gs.qd_float(0.0) - if within_clamp: - d_imp_d_imp_x = (dmax - dmin) * d_imp_y_d_imp_x - - # d(imp_x)/d(pos_delta) = sign(pos_delta)/width; pos_delta < 0 => -1/width - d_imp_x_d_pos_delta = -1.0 / width - d_imp_d_pos_delta = d_imp_d_imp_x * d_imp_x_d_pos_delta + # d(imp_x)/d(pos_delta) = sign(pos_delta)/width, and pos_delta < 0 here + d_imp_d_pos_delta = d_imp_d_imp_x * (-1.0 / width) # --- Combine --- dL_d_imp = g_aref * d_aref_d_imp + g_efc_D * d_efc_D_d_imp @@ -518,7 +495,7 @@ def kernel_manual_add_collision_constraints_bw( vel_motion = cdof_vel - t_pos x cdof_ang, t_pos = contact_pos - root_COM[link] jac_qvel = sum_chain jac[n_con, i_d] * dofs_vel[i_d] imp, aref = imp_aref(sol_params, -penetration, jac_qvel, -penetration) - diag = (invweight + friction^2 invweight) * 2 friction^2 (1-imp)/imp ; efc_D = 1/diag + diag = (invweight + mu2 invweight) * 2 mu2 (1-imp)/imp, mu2 = friction^2/impratio ; efc_D = 1/diag """ EPS = rigid_info.EPS[None] _B = dyn_state.dofs.ctrl_mode.shape[1] @@ -527,7 +504,9 @@ def kernel_manual_add_collision_constraints_bw( qd.loop_config( name="kernel_manual_add_collision_constraints_bw", - serialize=qd.static(rigid_config.para_level < gs.PARA_LEVEL.ALL), + # Per-contact reverses are independent (grad writes accumulate atomically); same gate as the forward + # per-contact assembly in solver.py. + serialize=qd.static(rigid_config.para_level < gs.PARA_LEVEL.PARTIAL), ) for flat_idx in range(max_contact_pairs * _B): i_b = flat_idx % _B @@ -554,9 +533,9 @@ def kernel_manual_add_collision_constraints_bw( # b_raw branches on |normal[1]| < 0.5; b = normalize(b_raw) # d1 = b x normal, d2 = b n0, n1, n2 = normal[0], normal[1], normal[2] - branch_a = qd.abs(n1) < 0.5 + is_branch_a = qd.abs(n1) < 0.5 b_raw = gs.qd_vec3(0.0, 0.0, 0.0) - if branch_a: + if is_branch_a: b_raw = gs.qd_vec3(-n0 * n1, 1.0 - n1 * n1, -n2 * n1) else: b_raw = gs.qd_vec3(-n0 * n2, -n1 * n2, 1.0 - n2 * n2) @@ -565,48 +544,19 @@ def kernel_manual_add_collision_constraints_bw( d1 = b.cross(normal) d2 = b - sol_timeconst = sol_params[0] - sol_dampratio = sol_params[1] - sol_dmin = sol_params[2] - sol_dmax = sol_params[3] - sol_width = sol_params[4] - sol_mid = sol_params[5] - sol_power = sol_params[6] - neg_pen = -penetration - imp_x = qd.abs(neg_pen) / sol_width + imp, b_coef, k_coef, d_imp_d_imp_x = gu.imp_aref_grad(sol_params, neg_pen) # d(imp_x)/d(penetration) = -sign(neg_pen)/width sign_neg = gs.qd_float(1.0) if neg_pen >= 0 else gs.qd_float(-1.0) - d_imp_x_d_pen = -sign_neg / sol_width - - imp_a_coef = 1.0 / sol_mid ** (sol_power - 1.0) - imp_b_coef = 1.0 / (1.0 - sol_mid) ** (sol_power - 1.0) - imp_a = imp_a_coef * imp_x**sol_power - imp_b = 1.0 - imp_b_coef * (1.0 - imp_x) ** sol_power - imp_y = imp_a if imp_x < sol_mid else imp_b - imp_raw = sol_dmin + imp_y * (sol_dmax - sol_dmin) - imp_clamped = qd.math.clamp(imp_raw, sol_dmin, sol_dmax) - imp = sol_dmax if imp_x > 1.0 else imp_clamped - - b_coef = 2.0 / (sol_dmax * sol_timeconst) - # k_coef matches gu.imp_aref's k = 1/(dmax^2 timeconst^2 dampratio^2) - k_coef = 1.0 / (sol_dmax * sol_dmax * sol_timeconst * sol_timeconst * sol_dampratio * sol_dampratio) - - # diag = C0 * (1-imp)/imp, C0 = 2 friction^2 invweight (1 + friction^2) - C0 = (invweight + friction * friction * invweight) * 2.0 * friction * friction + d_imp_x_d_pen = -sign_neg / sol_params[4] + + # diag = C0 * (1-imp)/imp with the impratio-regularized cone coefficient of the forward: + # friction_sq_reg = friction^2 / impratio (see add_collision_constraints in solver.py). + friction_sq_reg = friction * friction / rigid_info.impratio[None] + C0 = (invweight + friction_sq_reg * invweight) * 2.0 * friction_sq_reg diag_raw = C0 * (1.0 - imp) / imp diag = qd.max(diag_raw, EPS) - within_clamp = (imp_raw > sol_dmin) and (imp_raw < sol_dmax) and (imp_x <= 1.0) - d_imp_y_d_imp_x = gs.qd_float(0.0) - if imp_x < sol_mid: - d_imp_y_d_imp_x = sol_power * imp_a_coef * imp_x ** (sol_power - 1.0) - else: - d_imp_y_d_imp_x = sol_power * imp_b_coef * (1.0 - imp_x) ** (sol_power - 1.0) - d_imp_d_imp_x = gs.qd_float(0.0) - if within_clamp: - d_imp_d_imp_x = (sol_dmax - sol_dmin) * d_imp_y_d_imp_x - d_diag_d_imp = gs.qd_float(0.0) if diag_raw > EPS: d_diag_d_imp = -C0 / (imp * imp) @@ -624,11 +574,11 @@ def kernel_manual_add_collision_constraints_bw( const_start = ( constraint_state.n_constraints_equality[i_b] + constraint_state.n_constraints_frictionloss[i_b] ) - for i in range(4): + for i in range(qd.static(rigid_config.rows_per_contact)): s_i = gs.qd_float(2 * (i % 2) - 1) d = s_i * d1 if i < 2 else s_i * d2 n = d * friction - normal - n_con = const_start + i_col_ * 4 + i + n_con = const_start + i_col_ * qd.static(rigid_config.rows_per_contact) + i g_aref = constraint_state.dL_daref[n_con, i_b] g_efc_D = constraint_state.dL_defc_D[n_con, i_b] @@ -643,6 +593,13 @@ def kernel_manual_add_collision_constraints_bw( g_pen += dL_d_pen dL_d_jac_qvel = g_aref * d_aref_d_jac_qvel + # d(jac_qvel)/d(vel[i_d]) = jac[n_con, i_d]: accumulate once per relevant dof recorded by the + # forward row assembly - the chain walk below visits shared ancestor dofs of same-root pairs + # twice (see _append_relevant_dof in solver.py). + for k in range(constraint_state.jac_n_dofs[n_con, i_b]): + i_d = constraint_state.jac_dofs_idx[n_con, k, i_b] + dyn_state.dofs.vel.grad[i_d, i_b] += dL_d_jac_qvel * constraint_state.jac[n_con, i_d, i_b] + # Reverse jac[n_con, i_d] over the kinematic chain. dL_dn = gs.qd_vec3(0.0, 0.0, 0.0) for i_ab in range(2): @@ -652,33 +609,33 @@ def kernel_manual_add_collision_constraints_bw( sign = gs.qd_float(1.0) link = link_b while link > -1: - link_mb = [link, i_b] if qd.static(rigid_config.batch_links_info) else link - for i_d_ in range(dyn_info.links.n_dofs[link_mb]): - i_d = dyn_info.links.dof_end[link_mb] - 1 - i_d_ + link_maybe_batch = [link, i_b] if qd.static(rigid_config.batch_links_info) else link + for i_d_ in range(dyn_info.links.n_dofs[link_maybe_batch]): + i_d = dyn_info.links.dof_end[link_maybe_batch] - 1 - i_d_ cdof_ang = dyn_state.dofs.cdof_ang[i_d, i_b] cdof_vel = dyn_state.dofs.cdof_vel[i_d, i_b] t_pos = contact_pos - dyn_state.links.root_COM[link, i_b] vel_motion = cdof_vel - t_pos.cross(cdof_ang) - jac_stored = constraint_state.jac[n_con, i_d, i_b] g_jac = ( constraint_state.dL_djac[n_con, i_d, i_b] + dL_d_jac_qvel * dyn_state.dofs.vel[i_d, i_b] ) - dyn_state.dofs.vel.grad[i_d, i_b] += dL_d_jac_qvel * jac_stored # jac_contrib = (sign * vel_motion) . n dL_dn += g_jac * sign * vel_motion - g_vm = g_jac * sign * n # dL/d(vel_motion) + # dL/d(vel_motion) + g_vm = g_jac * sign * n # vel_motion = cdof_vel - t_pos x cdof_ang dyn_state.dofs.cdof_vel.grad[i_d, i_b] += g_vm dyn_state.dofs.cdof_ang.grad[i_d, i_b] += t_pos.cross(g_vm) - dt = -(cdof_ang.cross(g_vm)) # dL/d(t_pos) - g_pos += dt - dyn_state.links.root_COM.grad[link, i_b] += -dt + # dL/d(t_pos) + g_t_pos = -cdof_ang.cross(g_vm) + g_pos += g_t_pos + dyn_state.links.root_COM.grad[link, i_b] += -g_t_pos - link = dyn_info.links.parent_idx[link_mb] + link = dyn_info.links.parent_idx[link_maybe_batch] # n = d*friction - normal g_normal += -dL_dn @@ -694,7 +651,7 @@ def kernel_manual_add_collision_constraints_bw( # b = b_raw / |b_raw| dL_db_raw = (dL_db - dL_db.dot(b) * b) / b_raw_norm # b_raw(normal) branch Jacobian - if branch_a: + if is_branch_a: # b_raw = (-n0 n1, 1 - n1^2, -n2 n1) g_normal[0] += dL_db_raw[0] * (-n1) g_normal[1] += dL_db_raw[0] * (-n0) + dL_db_raw[1] * (-2.0 * n1) + dL_db_raw[2] * (-n2) @@ -737,7 +694,7 @@ def kernel_manual_add_frictionloss_constraints_bw( is False => d_imp / d_anything = 0). What survives is the direct `aref = -b_coef * jac_qvel` term, so - dL/d_vel[i_d] += dL_daref[n_con] * (-b_coef) + dL/d_vel[i_d] += dL_daref[n_con] * (-b_coef) """ EPS = rigid_info.EPS[None] _B = constraint_state.jac.shape[2] @@ -745,8 +702,7 @@ def kernel_manual_add_frictionloss_constraints_bw( qd.loop_config( name="kernel_manual_add_frictionloss_constraints_bw", - # Mirror the forward's serialize condition (frictionloss forward has a - # Metal-specific quirk; keep parity to make the loop walk identical). + # Same serialize condition as the forward; see add_frictionloss_constraints in solver.py for the Metal gate serialize=qd.static(rigid_config.para_level < gs.PARA_LEVEL.ALL and rigid_config.backend != gs.metal), ) for i_b in range(_B): @@ -776,6 +732,80 @@ def kernel_manual_add_frictionloss_constraints_bw( dyn_state.dofs.vel.grad[i_d, i_b] += g_aref * (-b_coef) +@qd.func +def func_cddb_ang_bw( + i_b, + link, + g_cddb_ang, + dyn_state: array_class.DynState, + dyn_info: array_class.DynInfo, + rigid_config: qd.template(), +): + """Reverse of the chain contraction cddb_ang = sum_d cdofd_ang[d] * vel[d] (see func_equality_jdotv in + solver.py), accumulating into cdofd_ang.grad and vel.grad over the link's ancestor dofs.""" + i_l = link + while i_l > -1: + I_l = [i_l, i_b] if qd.static(rigid_config.batch_links_info) else i_l + for i_d in range(dyn_info.links.dof_start[I_l], dyn_info.links.dof_end[I_l]): + dyn_state.dofs.cdofd_ang.grad[i_d, i_b] += g_cddb_ang * dyn_state.dofs.vel[i_d, i_b] + dyn_state.dofs.vel.grad[i_d, i_b] += g_cddb_ang.dot(dyn_state.dofs.cdofd_ang[i_d, i_b]) + i_l = dyn_info.links.parent_idx[I_l] + + +@qd.func +def func_equality_jdotv_bw( + i_b, + link, + anchor_pos, + g_jdotv, + dyn_state: array_class.DynState, + dyn_info: array_class.DynInfo, + rigid_config: qd.template(), +): + """Reverse of func_equality_jdotv (see solver.py) for one chain, given the upstream gradient g_jdotv of its + linear Jdot @ qvel. + + Accumulates into the chain dofs' cdofd_ang / cdofd_vel / vel grads and the link's cd_ang / cd_vel / root_COM + grads, and returns the gradient w.r.t. the world anchor position (the caller owns the anchor -> link pos / quat + chain). Returns zero for the world (link == -1).""" + g_anchor = gs.qd_vec3(0.0, 0.0, 0.0) + if link > -1: + # Replay the chain contraction; only cddb_ang is consumed by the adjoint below + _jdotv, cddb_ang = solver.func_equality_jdotv(i_b, link, anchor_pos, dyn_state, dyn_info, rigid_config) + offset = anchor_pos - dyn_state.links.root_COM[link, i_b] + pvel = dyn_state.links.cd_vel[link, i_b] + dyn_state.links.cd_ang[link, i_b].cross(offset) + + # jdotv = cddb_vel + cddb_ang x offset + cd_ang x pvel, using g_u = w x g and g_w = g x u for c = u x w + g_cddb_vel = g_jdotv + g_cddb_ang = offset.cross(g_jdotv) + g_offset = g_jdotv.cross(cddb_ang) + g_cd_ang = pvel.cross(g_jdotv) + g_pvel = g_jdotv.cross(dyn_state.links.cd_ang[link, i_b]) + + # pvel = cd_vel + cd_ang x offset + g_cd_vel = g_pvel + g_cd_ang = g_cd_ang + offset.cross(g_pvel) + g_offset = g_offset + g_pvel.cross(dyn_state.links.cd_ang[link, i_b]) + + # offset = anchor_pos - root_COM[link] + g_anchor = g_offset + dyn_state.links.root_COM.grad[link, i_b] += -g_offset + dyn_state.links.cd_ang.grad[link, i_b] += g_cd_ang + dyn_state.links.cd_vel.grad[link, i_b] += g_cd_vel + + # cddb_{ang,vel} = sum_d cdofd_{ang,vel}[d] * vel[d] over the ancestor chain + i_l = link + while i_l > -1: + I_l = [i_l, i_b] if qd.static(rigid_config.batch_links_info) else i_l + for i_d in range(dyn_info.links.dof_start[I_l], dyn_info.links.dof_end[I_l]): + dyn_state.dofs.cdofd_ang.grad[i_d, i_b] += g_cddb_ang * dyn_state.dofs.vel[i_d, i_b] + dyn_state.dofs.cdofd_vel.grad[i_d, i_b] += g_cddb_vel * dyn_state.dofs.vel[i_d, i_b] + dyn_state.dofs.vel.grad[i_d, i_b] += g_cddb_ang.dot(dyn_state.dofs.cdofd_ang[i_d, i_b]) + dyn_state.dofs.vel.grad[i_d, i_b] += g_cddb_vel.dot(dyn_state.dofs.cdofd_vel[i_d, i_b]) + i_l = dyn_info.links.parent_idx[I_l] + return g_anchor + + @qd.kernel(fastcache=True) def kernel_manual_add_equality_constraints_bw( dyn_state: array_class.DynState, @@ -795,9 +825,9 @@ def kernel_manual_add_equality_constraints_bw( JOINT: rigid_info.qpos.grad[i_qpos1], qpos.grad[i_qpos2] dyn_state.dofs.vel.grad[i_dof1], vel.grad[i_dof2] - CONNECT: - dyn_state.links.{pos, quat, root_COM}.grad[link1 / link2] - dyn_state.dofs.{cdof_ang, cdof_vel, vel}.grad over each link chain + CONNECT / WELD: + dyn_state.links.{pos, quat, root_COM, cd_ang, cd_vel}.grad[link1 / link2] + dyn_state.dofs.{cdof_ang, cdof_vel, cdofd_ang, cdofd_vel, vel}.grad over each link chain Model parameters (`sol_params`, `eq_data`, `dyn_info.dofs.invweight`, `dyn_info.links.invweight`) are not differentiated. @@ -808,10 +838,10 @@ def kernel_manual_add_equality_constraints_bw( deriv = d(pos_poly)/d(diff) = a1 + 2 * a2 * diff + 3 * a3 * diff^2 + 4 * a4 * diff^3 jac[n_con, i_dof1] = 1.0 jac[n_con, i_dof2] = -deriv - jac_qvel = vel[i_dof1] - deriv * vel[i_dof2] + jac_qvel = vel[i_dof1] - deriv * vel[i_dof2] imp, aref = imp_aref(sol_params, -|pos|, jac_qvel, pos) - aref = -b * jac_qvel - k * imp * pos - diag = max(invweight * (1 - imp) / imp, EPS); efc_D = 1/diag + aref = -b * jac_qvel - k * imp * pos + diag = max(invweight * (1 - imp) / imp, EPS); efc_D = 1/diag """ EPS = rigid_info.EPS[None] _B = constraint_state.jac.shape[2] @@ -874,27 +904,8 @@ def kernel_manual_add_equality_constraints_bw( invweight = dyn_info.dofs.invweight[I_dof1] + dyn_info.dofs.invweight[I_dof2] sol_params = dyn_info.equalities.sol_params[i_e, i_b] - timeconst = sol_params[0] - dampratio = sol_params[1] - dmin = sol_params[2] - dmax = sol_params[3] + imp, b_coef, k_coef, d_imp_d_imp_x = gu.imp_aref_grad(sol_params, pos) width = sol_params[4] - mid = sol_params[5] - power = sol_params[6] - - # imp_x = |pos_delta_arg| / width = |-|pos|| / width = |pos|/width - imp_x = qd.abs(pos) / width - imp_a_coef = 1.0 / mid ** (power - 1.0) - imp_b_coef = 1.0 / (1.0 - mid) ** (power - 1.0) - imp_a = imp_a_coef * imp_x**power - imp_b = 1.0 - imp_b_coef * (1.0 - imp_x) ** power - imp_y = imp_a if imp_x < mid else imp_b - imp_raw = dmin + imp_y * (dmax - dmin) - imp_clamped = qd.math.clamp(imp_raw, dmin, dmax) - imp = dmax if imp_x > 1.0 else imp_clamped - - b_coef = 2.0 / (dmax * timeconst) - k_coef = 1.0 / (dmax * dmax * timeconst * timeconst * dampratio * dampratio) diag_raw = invweight * (1.0 - imp) / imp diag = qd.max(diag_raw, EPS) @@ -906,7 +917,7 @@ def kernel_manual_add_equality_constraints_bw( g_jac2 = constraint_state.dL_djac[n_con, i_dof2, i_b] # ---- Partials ---- - # aref = -b * jac_qvel - k * imp * pos + # aref = -b * jac_qvel - k * imp * pos d_aref_d_jac_qvel = -b_coef d_aref_d_pos_direct = -k_coef * imp d_aref_d_imp = -k_coef * pos @@ -917,17 +928,6 @@ def kernel_manual_add_equality_constraints_bw( d_diag_d_imp = -invweight / (imp * imp) d_efc_D_d_imp = -d_diag_d_imp / (diag * diag) - # d(imp)/d(imp_x) active only inside the smooth clamp band. - within_clamp = (imp_raw > dmin) and (imp_raw < dmax) and (imp_x <= 1.0) - d_imp_y_d_imp_x = gs.qd_float(0.0) - if imp_x < mid: - d_imp_y_d_imp_x = power * imp_a_coef * imp_x ** (power - 1.0) - else: - d_imp_y_d_imp_x = power * imp_b_coef * (1.0 - imp_x) ** (power - 1.0) - d_imp_d_imp_x = gs.qd_float(0.0) - if within_clamp: - d_imp_d_imp_x = (dmax - dmin) * d_imp_y_d_imp_x - # imp_x = |pos|/width => d_imp_x/d_pos = sign(pos) / width sign_pos_f = gs.qd_float(1.0) if pos < 0.0: @@ -953,24 +953,25 @@ def kernel_manual_add_equality_constraints_bw( # CONNECT: 3 rows pin global_anchor1 == global_anchor2. # # Forward recap (per row i_3 in {0,1,2}): - # ga1 = trans(dyn_state.links.pos[link1], dyn_state.links.quat[link1]) * eq_data[0:3] - # ga2 = trans(dyn_state.links.pos[link2], dyn_state.links.quat[link2]) * eq_data[3:6] + # ga1 = trans(dyn_state.links.pos[link1], dyn_state.links.quat[link1]) * eq_data[0:3] + # ga2 = trans(dyn_state.links.pos[link2], dyn_state.links.quat[link2]) * eq_data[3:6] # For each link in (link1, link2) chain, for each dof on that link: # t_pos = ga_link - root_COM[link] # vel_motion = cdof_vel - t_pos x cdof_ang - # jac_i3 = sign * vel_motion[i_3] (sign = +1 for link1, -1 for link2) + # jac_i3 = sign * vel_motion[i_3] (sign = +1 for link1, -1 for link2) # jac[n_con, i_d] += jac_i3 - # jac_qvel += jac_i3 * vel[i_d] + # jac_qvel += jac_i3 * vel[i_d] # pos_diff = ga1 - ga2 # penetration = ||pos_diff|| # imp, aref = imp_aref(sol_params, -penetration, jac_qvel, pos_diff[i_3]) - # aref = -b * jac_qvel - k * imp * pos_diff[i_3] - # diag = max(invweight * (1 - imp) / imp, EPS); efc_D = 1/diag + # aref = -b * jac_qvel - k * imp * pos_diff[i_3] + # stored aref[n_con] = aref - jdotv[i_3], jdotv = jdotv1 - jdotv2 (func_equality_jdotv) + # diag = max(invweight * (1 - imp) / imp, EPS); efc_D = 1/diag # ---------------------------------------------------------- link1_idx = dyn_info.equalities.eq_obj1id[i_e, i_b] link2_idx = dyn_info.equalities.eq_obj2id[i_e, i_b] - link1_mb = [link1_idx, i_b] if qd.static(rigid_config.batch_links_info) else link1_idx - link2_mb = [link2_idx, i_b] if qd.static(rigid_config.batch_links_info) else link2_idx + link1_maybe_batch = [link1_idx, i_b] if qd.static(rigid_config.batch_links_info) else link1_idx + link2_maybe_batch = [link2_idx, i_b] if qd.static(rigid_config.batch_links_info) else link2_idx anchor1_local = gs.qd_vec3( dyn_info.equalities.eq_data[i_e, i_b][0], @@ -992,30 +993,13 @@ def kernel_manual_add_equality_constraints_bw( pos_diff = ga1 - ga2 penetration = pos_diff.norm() - invweight = dyn_info.links.invweight[link1_mb][0] + dyn_info.links.invweight[link2_mb][0] + invweight = ( + dyn_info.links.invweight[link1_maybe_batch][0] + dyn_info.links.invweight[link2_maybe_batch][0] + ) sol_params = dyn_info.equalities.sol_params[i_e, i_b] - timeconst = sol_params[0] - dampratio = sol_params[1] - dmin = sol_params[2] - dmax = sol_params[3] + imp, b_coef, k_coef, d_imp_d_imp_x = gu.imp_aref_grad(sol_params, -penetration) width = sol_params[4] - mid = sol_params[5] - power = sol_params[6] - - # imp_x = |-penetration| / width = penetration / width - imp_x = penetration / width - imp_a_coef = 1.0 / mid ** (power - 1.0) - imp_b_coef = 1.0 / (1.0 - mid) ** (power - 1.0) - imp_a = imp_a_coef * imp_x**power - imp_b = 1.0 - imp_b_coef * (1.0 - imp_x) ** power - imp_y = imp_a if imp_x < mid else imp_b - imp_raw = dmin + imp_y * (dmax - dmin) - imp_clamped = qd.math.clamp(imp_raw, dmin, dmax) - imp = dmax if imp_x > 1.0 else imp_clamped - - b_coef = 2.0 / (dmax * timeconst) - k_coef = 1.0 / (dmax * dmax * timeconst * timeconst * dampratio * dampratio) diag_raw = invweight * (1.0 - imp) / imp diag = qd.max(diag_raw, EPS) @@ -1027,20 +1011,13 @@ def kernel_manual_add_equality_constraints_bw( d_diag_d_imp = -invweight / (imp * imp) d_efc_D_d_imp = -d_diag_d_imp / (diag * diag) - within_clamp = (imp_raw > dmin) and (imp_raw < dmax) and (imp_x <= 1.0) - d_imp_y_d_imp_x = gs.qd_float(0.0) - if imp_x < mid: - d_imp_y_d_imp_x = power * imp_a_coef * imp_x ** (power - 1.0) - else: - d_imp_y_d_imp_x = power * imp_b_coef * (1.0 - imp_x) ** (power - 1.0) - d_imp_d_imp_x = gs.qd_float(0.0) - if within_clamp: - d_imp_d_imp_x = (dmax - dmin) * d_imp_y_d_imp_x - # Accumulate dL/d_ga over the 3 rows so we propagate to # dyn_state.links.{pos,quat} only once per anchor. g_ga1 = gs.qd_vec3(0.0, 0.0, 0.0) g_ga2 = gs.qd_vec3(0.0, 0.0, 0.0) + # The stored rows are aref - jdotv[i_3] (see func_equality_connect in solver.py); collect the bias + # gradient per row and reverse it through both chains after the row loop. + g_jdotv = gs.qd_vec3(0.0, 0.0, 0.0) for i_3 in range(3): n_con = n_con_counter @@ -1052,10 +1029,17 @@ def kernel_manual_add_equality_constraints_bw( d_aref_d_jac_qvel = -b_coef d_aref_d_pos_diff_i3_direct = -k_coef * imp d_aref_d_imp = -k_coef * pos_diff[i_3] + g_jdotv[i_3] = -g_aref dL_d_imp = g_aref * d_aref_d_imp + g_efc_D * d_efc_D_d_imp dL_d_jac_qvel = g_aref * d_aref_d_jac_qvel + # d(jac_qvel)/d(vel[i_d]) over the forward-recorded dof support (see the collision reverse + # above for why the chain walk cannot accumulate this). + for k in range(constraint_state.jac_n_dofs[n_con, i_b]): + i_d = constraint_state.jac_dofs_idx[n_con, k, i_b] + dyn_state.dofs.vel.grad[i_d, i_b] += dL_d_jac_qvel * constraint_state.jac[n_con, i_d, i_b] + # dL/d_pos_diff: (a) direct axis-i_3 term, (b) via penetration / imp. g_pos_diff = gs.qd_vec3(0.0, 0.0, 0.0) g_pos_diff[i_3] = g_aref * d_aref_d_pos_diff_i3_direct @@ -1078,22 +1062,20 @@ def kernel_manual_add_equality_constraints_bw( anchor_pos = ga2 while link > -1: - link_mb = [link, i_b] if qd.static(rigid_config.batch_links_info) else link - for i_d_ in range(dyn_info.links.n_dofs[link_mb]): - i_d = dyn_info.links.dof_end[link_mb] - 1 - i_d_ + link_maybe_batch = [link, i_b] if qd.static(rigid_config.batch_links_info) else link + for i_d_ in range(dyn_info.links.n_dofs[link_maybe_batch]): + i_d = dyn_info.links.dof_end[link_maybe_batch] - 1 - i_d_ cdof_ang = dyn_state.dofs.cdof_ang[i_d, i_b] cdof_vel = dyn_state.dofs.cdof_vel[i_d, i_b] t_pos = anchor_pos - dyn_state.links.root_COM[link, i_b] - # jac_i3 = sign * vel_motion[i_3] - # upstream: dL_djac[n_con, i_d] + dL_d_jac_qvel * vel[i_d] - jac_stored = constraint_state.jac[n_con, i_d, i_b] + # jac_i3 = sign * vel_motion[i_3] + # upstream: dL_djac[n_con, i_d] + dL_d_jac_qvel * vel[i_d] g_jac_i3 = ( constraint_state.dL_djac[n_con, i_d, i_b] + dL_d_jac_qvel * dyn_state.dofs.vel[i_d, i_b] ) - dyn_state.dofs.vel.grad[i_d, i_b] += dL_d_jac_qvel * jac_stored # dL/d_vel_motion (only component i_3) g_vm = gs.qd_vec3(0.0, 0.0, 0.0) @@ -1102,22 +1084,26 @@ def kernel_manual_add_equality_constraints_bw( # vel_motion = cdof_vel - t_pos x cdof_ang dyn_state.dofs.cdof_vel.grad[i_d, i_b] += g_vm dyn_state.dofs.cdof_ang.grad[i_d, i_b] += t_pos.cross(g_vm) - dt = -(cdof_ang.cross(g_vm)) + g_t_pos = -cdof_ang.cross(g_vm) # t_pos = anchor_pos - root_COM[link] if i_ab == 0: - g_anchor1_row = g_anchor1_row + dt + g_anchor1_row = g_anchor1_row + g_t_pos else: - g_anchor2_row = g_anchor2_row + dt - dyn_state.links.root_COM.grad[link, i_b] += -dt + g_anchor2_row = g_anchor2_row + g_t_pos + dyn_state.links.root_COM.grad[link, i_b] += -g_t_pos - link = dyn_info.links.parent_idx[link_mb] + link = dyn_info.links.parent_idx[link_maybe_batch] # pos_diff = ga1 - ga2 => g_ga1 += g_pos_diff, g_ga2 += -g_pos_diff. g_ga1 = g_ga1 + g_pos_diff + g_anchor1_row g_ga2 = g_ga2 - g_pos_diff + g_anchor2_row - # Propagate accumulated g_aref grads back to dyn_state.links.{pos, quat}. - # g_aref = trans + R(quat) * anchor_local; anchor_local is model param. + # Reverse the velocity-product bias jdotv = jdotv1 - jdotv2 through both chains + g_ga1 = g_ga1 + func_equality_jdotv_bw(i_b, link1_idx, ga1, g_jdotv, dyn_state, dyn_info, rigid_config) + g_ga2 = g_ga2 + func_equality_jdotv_bw(i_b, link2_idx, ga2, -g_jdotv, dyn_state, dyn_info, rigid_config) + + # Propagate accumulated g_ga grads back to dyn_state.links.{pos, quat}. + # ga = trans + R(quat) * anchor_local; anchor_local is model param. g_quat1 = gu.qd_transform_by_quat_grad_quat(anchor1_local, quat1, g_ga1) g_quat2 = gu.qd_transform_by_quat_grad_quat(anchor2_local, quat2, g_ga2) dyn_state.links.pos.grad[link1_idx, i_b] += g_ga1 @@ -1126,31 +1112,31 @@ def kernel_manual_add_equality_constraints_bw( dyn_state.links.quat.grad[link2_idx, i_b] += g_quat2 else: # ---------------------------------------------------------- - # WELD: 6 rows -- 3 position + 3 orientation, all sharing + # WELD: 6 rows - 3 position + 3 orientation, all sharing # a single combined pos_imp = ||all_error|| (6D). # # Forward recap: - # ga1 = trans(dyn_state.links.pos[link1], dyn_state.links.quat[link1]) * eq_data[3:6] - # ga2 = trans(dyn_state.links.pos[link2], dyn_state.links.quat[link2]) * eq_data[0:3] + # ga1 = trans(dyn_state.links.pos[link1], dyn_state.links.quat[link1]) * eq_data[3:6] + # ga2 = trans(dyn_state.links.pos[link2], dyn_state.links.quat[link2]) * eq_data[0:3] # pos_error = ga1 - ga2 # inv_q2 = inv_quat(dyn_state.links.quat[link2]) # q = quat_mul(dyn_state.links.quat[link1], relpose) # error_quat = quat_mul(inv_q2, q) - # rot_error = error_quat.xyz * torquescale + # rot_error = error_quat.xyz * torquescale # all_error = (pos_error, rot_error); pos_imp = ||all_error|| # # Position rows (i_3 = 0..2): same chain structure as CONNECT # with invweight[0]; ref_arg = pos_error[i_3]. - # Orientation rows (i_3 = 0..2): jac_phase1 = sign * cdof_ang[i_d] + # Orientation rows (i_3 = 0..2): jac_phase1 = sign * cdof_ang[i_d] # then quat post-process # quat2_d = qd_quat_mul_axis(inv_q2, jac_phase1[i_d]) # quat3_d = qd_quat_mul(quat2_d, q) - # jac[n_con+3+i_3, i_d] = 0.5 * torquescale * quat3_d[i_3+1] + # jac[n_con+3+i_3, i_d] = 0.5 * torquescale * quat3_d[i_3+1] # ---------------------------------------------------------- link1_idx = dyn_info.equalities.eq_obj1id[i_e, i_b] link2_idx = dyn_info.equalities.eq_obj2id[i_e, i_b] - link1_mb = [link1_idx, i_b] if qd.static(rigid_config.batch_links_info) else link1_idx - link2_mb = [link2_idx, i_b] if qd.static(rigid_config.batch_links_info) else link2_idx + link1_maybe_batch = [link1_idx, i_b] if qd.static(rigid_config.batch_links_info) else link1_idx + link2_maybe_batch = [link2_idx, i_b] if qd.static(rigid_config.batch_links_info) else link2_idx # WELD eq_data layout (per forward comment): # [0:3] anchor2 (local), [3:6] anchor1 (local), [6:10] relpose, [10] torquescale @@ -1198,31 +1184,16 @@ def kernel_manual_add_equality_constraints_bw( + rot_error[2] * rot_error[2] ) - invweight_pos = dyn_info.links.invweight[link1_mb][0] + dyn_info.links.invweight[link2_mb][0] - invweight_rot = dyn_info.links.invweight[link1_mb][1] + dyn_info.links.invweight[link2_mb][1] + invweight_pos = ( + dyn_info.links.invweight[link1_maybe_batch][0] + dyn_info.links.invweight[link2_maybe_batch][0] + ) + invweight_rot = ( + dyn_info.links.invweight[link1_maybe_batch][1] + dyn_info.links.invweight[link2_maybe_batch][1] + ) sol_params = dyn_info.equalities.sol_params[i_e, i_b] - timeconst = sol_params[0] - dampratio = sol_params[1] - dmin = sol_params[2] - dmax = sol_params[3] + imp, b_coef, k_coef, d_imp_d_imp_x = gu.imp_aref_grad(sol_params, -pos_imp) width = sol_params[4] - mid = sol_params[5] - power = sol_params[6] - - # imp_x = |-pos_imp| / width = pos_imp/width (all rows share) - imp_x = pos_imp / width - imp_a_coef = 1.0 / mid ** (power - 1.0) - imp_b_coef = 1.0 / (1.0 - mid) ** (power - 1.0) - imp_a = imp_a_coef * imp_x**power - imp_b = 1.0 - imp_b_coef * (1.0 - imp_x) ** power - imp_y = imp_a if imp_x < mid else imp_b - imp_raw = dmin + imp_y * (dmax - dmin) - imp_clamped = qd.math.clamp(imp_raw, dmin, dmax) - imp = dmax if imp_x > 1.0 else imp_clamped - - b_coef = 2.0 / (dmax * timeconst) - k_coef = 1.0 / (dmax * dmax * timeconst * timeconst * dampratio * dampratio) # Per-group diag/efc_D depend on invweight; same imp. diag_raw_pos = invweight_pos * (1.0 - imp) / imp @@ -1239,16 +1210,6 @@ def kernel_manual_add_equality_constraints_bw( d_diag_d_imp_rot = -invweight_rot / (imp * imp) d_efc_D_d_imp_rot = -d_diag_d_imp_rot / (diag_rot * diag_rot) - within_clamp = (imp_raw > dmin) and (imp_raw < dmax) and (imp_x <= 1.0) - d_imp_y_d_imp_x = gs.qd_float(0.0) - if imp_x < mid: - d_imp_y_d_imp_x = power * imp_a_coef * imp_x ** (power - 1.0) - else: - d_imp_y_d_imp_x = power * imp_b_coef * (1.0 - imp_x) ** (power - 1.0) - d_imp_d_imp_x = gs.qd_float(0.0) - if within_clamp: - d_imp_d_imp_x = (dmax - dmin) * d_imp_y_d_imp_x - # Accumulators across all 6 rows. g_ga1 = gs.qd_vec3(0.0, 0.0, 0.0) g_ga2 = gs.qd_vec3(0.0, 0.0, 0.0) @@ -1256,6 +1217,11 @@ def kernel_manual_add_equality_constraints_bw( dL_d_imp_total = gs.qd_float(0.0) # Per-row jac_qvel grad (for the orientation chain walk below). dL_d_jac_qvel_orient = gs.qd_vec3(0.0, 0.0, 0.0) + # The stored position rows are aref - jdotv[i_3] and the rotation rows carry the rotational bias + # 0.5 * (t1 + t2 + t3)[i_3 + 1] * torquescale (see func_equality_weld in solver.py); collect both + # bias gradients per row and reverse them after the row loops. + g_jdotv = gs.qd_vec3(0.0, 0.0, 0.0) + g_t_quat = qd.Vector([0.0, 0.0, 0.0, 0.0], dt=gs.qd_float) # ---- Position rows (3) -- mirrors CONNECT structure ---- n_con_orient_base = n_con_counter + 3 # rotation rows start here @@ -1269,12 +1235,19 @@ def kernel_manual_add_equality_constraints_bw( d_aref_d_jac_qvel = -b_coef d_aref_d_ref_direct = -k_coef * imp d_aref_d_imp = -k_coef * pos_error[i_3] + g_jdotv[i_3] = -g_aref dL_d_imp_total = dL_d_imp_total + g_aref * d_aref_d_imp + g_efc_D * d_efc_D_d_imp_pos dL_d_jac_qvel = g_aref * d_aref_d_jac_qvel # Direct ref-axis contribution (pos_error[i_3]): g_pos_error_direct = g_aref * d_aref_d_ref_direct + # d(jac_qvel)/d(vel[i_d]) over the forward-recorded dof support (see the collision reverse + # above for why the chain walk cannot accumulate this). + for k in range(constraint_state.jac_n_dofs[n_con, i_b]): + i_d = constraint_state.jac_dofs_idx[n_con, k, i_b] + dyn_state.dofs.vel.grad[i_d, i_b] += dL_d_jac_qvel * constraint_state.jac[n_con, i_d, i_b] + # Chain walk (same shape as CONNECT pos chain): g_anchor1_row = gs.qd_vec3(0.0, 0.0, 0.0) g_anchor2_row = gs.qd_vec3(0.0, 0.0, 0.0) @@ -1288,34 +1261,32 @@ def kernel_manual_add_equality_constraints_bw( anchor_pos = ga2 while link > -1: - link_mb = [link, i_b] if qd.static(rigid_config.batch_links_info) else link - for i_d_ in range(dyn_info.links.n_dofs[link_mb]): - i_d = dyn_info.links.dof_end[link_mb] - 1 - i_d_ + link_maybe_batch = [link, i_b] if qd.static(rigid_config.batch_links_info) else link + for i_d_ in range(dyn_info.links.n_dofs[link_maybe_batch]): + i_d = dyn_info.links.dof_end[link_maybe_batch] - 1 - i_d_ cdof_ang = dyn_state.dofs.cdof_ang[i_d, i_b] cdof_vel = dyn_state.dofs.cdof_vel[i_d, i_b] t_pos = anchor_pos - dyn_state.links.root_COM[link, i_b] - jac_stored = constraint_state.jac[n_con, i_d, i_b] g_jac_i3 = ( constraint_state.dL_djac[n_con, i_d, i_b] + dL_d_jac_qvel * dyn_state.dofs.vel[i_d, i_b] ) - dyn_state.dofs.vel.grad[i_d, i_b] += dL_d_jac_qvel * jac_stored g_vm = gs.qd_vec3(0.0, 0.0, 0.0) g_vm[i_3] = g_jac_i3 * sign dyn_state.dofs.cdof_vel.grad[i_d, i_b] += g_vm dyn_state.dofs.cdof_ang.grad[i_d, i_b] += t_pos.cross(g_vm) - dt = -(cdof_ang.cross(g_vm)) + g_t_pos = -cdof_ang.cross(g_vm) if i_ab == 0: - g_anchor1_row = g_anchor1_row + dt + g_anchor1_row = g_anchor1_row + g_t_pos else: - g_anchor2_row = g_anchor2_row + dt - dyn_state.links.root_COM.grad[link, i_b] += -dt + g_anchor2_row = g_anchor2_row + g_t_pos + dyn_state.links.root_COM.grad[link, i_b] += -g_t_pos - link = dyn_info.links.parent_idx[link_mb] + link = dyn_info.links.parent_idx[link_maybe_batch] # pos_error = ga1 - ga2 => direct g splits to g_ga1, g_ga2 oppositely. g_ga1[i_3] = g_ga1[i_3] + g_pos_error_direct @@ -1323,6 +1294,10 @@ def kernel_manual_add_equality_constraints_bw( g_ga1 = g_ga1 + g_anchor1_row g_ga2 = g_ga2 + g_anchor2_row + # Reverse the velocity-product bias jdotv = jdotv1 - jdotv2 through both chains + g_ga1 = g_ga1 + func_equality_jdotv_bw(i_b, link1_idx, ga1, g_jdotv, dyn_state, dyn_info, rigid_config) + g_ga2 = g_ga2 + func_equality_jdotv_bw(i_b, link2_idx, ga2, -g_jdotv, dyn_state, dyn_info, rigid_config) + # ---- Orientation rows (3) ---- # Direct contributions: rot_error[i_3] via ref, dL_d_imp via imp. for i_3 in range(3): @@ -1334,11 +1309,22 @@ def kernel_manual_add_equality_constraints_bw( d_aref_d_jac_qvel = -b_coef d_aref_d_ref_direct = -k_coef * imp d_aref_d_imp = -k_coef * rot_error[i_3] + # Rotational bias: stored aref = aref - 0.5 * (t1 + t2 + t3)[i_3 + 1] * torquescale, the same + # upstream gradient reaching each of t1 / t2 / t3. + g_t_quat[i_3 + 1] = -0.5 * torquescale * g_aref dL_d_imp_total = dL_d_imp_total + g_aref * d_aref_d_imp + g_efc_D * d_efc_D_d_imp_rot dL_d_jac_qvel_orient[i_3] = g_aref * d_aref_d_jac_qvel g_rot_error[i_3] = g_rot_error[i_3] + g_aref * d_aref_d_ref_direct + # d(jac_qvel)/d(vel[i_d]) over the forward-recorded dof support (see the collision reverse + # above for why the chain walk cannot accumulate this). + for k in range(constraint_state.jac_n_dofs[n_con, i_b]): + i_d = constraint_state.jac_dofs_idx[n_con, k, i_b] + dyn_state.dofs.vel.grad[i_d, i_b] += ( + dL_d_jac_qvel_orient[i_3] * constraint_state.jac[n_con, i_d, i_b] + ) + # Orientation chain walk: per i_d on chain, build g_quat3_d from # the 3 orient rows, then back-prop through quat_mul/quat_mul_axis. g_inv_q2 = qd.Vector([0.0, 0.0, 0.0, 0.0], dt=gs.qd_float) @@ -1351,21 +1337,19 @@ def kernel_manual_add_equality_constraints_bw( link = link2_idx while link > -1: - link_mb = [link, i_b] if qd.static(rigid_config.batch_links_info) else link - for i_d_ in range(dyn_info.links.n_dofs[link_mb]): - i_d = dyn_info.links.dof_end[link_mb] - 1 - i_d_ + link_maybe_batch = [link, i_b] if qd.static(rigid_config.batch_links_info) else link + for i_d_ in range(dyn_info.links.n_dofs[link_maybe_batch]): + i_d = dyn_info.links.dof_end[link_maybe_batch] - 1 - i_d_ # Build g_quat3_d (only xyz components feed jac). g_quat3_d = qd.Vector([0.0, 0.0, 0.0, 0.0], dt=gs.qd_float) for i_3 in qd.static(range(3)): row = n_con_orient_base + i_3 - jac_stored = constraint_state.jac[row, i_d, i_b] gjac = ( constraint_state.dL_djac[row, i_d, i_b] + dL_d_jac_qvel_orient[i_3] * dyn_state.dofs.vel[i_d, i_b] ) - dyn_state.dofs.vel.grad[i_d, i_b] += dL_d_jac_qvel_orient[i_3] * jac_stored - # jac[row, i_d] = 0.5 * torquescale * quat3_d[i_3+1] + # jac[row, i_d] = 0.5 * torquescale * quat3_d[i_3+1] g_quat3_d[i_3 + 1] = g_quat3_d[i_3 + 1] + gjac * 0.5 * torquescale # Replay quat3_d, quat2_d, jac_diff_r_d @@ -1392,7 +1376,7 @@ def kernel_manual_add_equality_constraints_bw( # jac_diff_r_d = sign_chain * cdof_ang[i_d] dyn_state.dofs.cdof_ang.grad[i_d, i_b] += sign_chain * g_jac_diff_r_d - link = dyn_info.links.parent_idx[link_mb] + link = dyn_info.links.parent_idx[link_maybe_batch] # Via-penetration contribution (shared across all 6 rows). if pos_imp > EPS: @@ -1421,6 +1405,78 @@ def kernel_manual_add_equality_constraints_bw( g_inv_q2 = g_inv_q2 + g_inv_q2_eq g_q = g_q + g_q_eq + # ---- Rotational bias t1 + t2 + t3 (see func_equality_weld in solver.py) ---- + # Replay the bias intermediates from the same state the forward read. + omega1 = dyn_state.links.cd_ang[link1_idx, i_b] + omega2 = dyn_state.links.cd_ang[link2_idx, i_b] + domega = omega1 - omega2 + p_omega1 = qd.Vector([0.0, omega1[0], omega1[1], omega1[2]], dt=gs.qd_float) + p_omega2 = qd.Vector([0.0, omega2[0], omega2[1], omega2[2]], dt=gs.qd_float) + p_domega = qd.Vector([0.0, domega[0], domega[1], domega[2]], dt=gs.qd_float) + qdot_body1 = 0.5 * gu.qd_quat_mul(p_omega1, quat_body1) + qdot0r = gu.qd_quat_mul(qdot_body1, relpose) + qdot_body2 = 0.5 * gu.qd_quat_mul(p_omega2, quat_body2) + inv_qdot2 = gu.qd_inv_quat(qdot_body2) + m1 = gu.qd_quat_mul_axis(inv_qdot2, domega) + m3 = gu.qd_quat_mul_axis(inv_q2, domega) + + g_omega1 = gs.qd_vec3(0.0, 0.0, 0.0) + g_omega2 = gs.qd_vec3(0.0, 0.0, 0.0) + g_domega = gs.qd_vec3(0.0, 0.0, 0.0) + + # t2 = quat_mul(quat_mul_axis(inv_q2, djrdv), q_var), with quat_mul_axis(u, a) = quat_mul(u, [0, a]) + # and djrdv = cddb1_ang - cddb2_ang the difference of the chains' angular Jdot @ qvel. + _jdotv1, cddb1_ang = solver.func_equality_jdotv(i_b, link1_idx, ga1, dyn_state, dyn_info, rigid_config) + _jdotv2, cddb2_ang = solver.func_equality_jdotv(i_b, link2_idx, ga2, dyn_state, dyn_info, rigid_config) + djrdv = cddb1_ang - cddb2_ang + p_djrdv = qd.Vector([0.0, djrdv[0], djrdv[1], djrdv[2]], dt=gs.qd_float) + m2 = gu.qd_quat_mul_axis(inv_q2, djrdv) + g_m2 = gu.qd_quat_mul_grad_lhs(m2, q_var, g_t_quat) + g_q = g_q + gu.qd_quat_mul_grad_rhs(m2, q_var, g_t_quat) + g_inv_q2 = g_inv_q2 + gu.qd_quat_mul_grad_lhs(inv_q2, p_djrdv, g_m2) + g_p_djrdv = gu.qd_quat_mul_grad_rhs(inv_q2, p_djrdv, g_m2) + g_djrdv = gs.qd_vec3(g_p_djrdv[1], g_p_djrdv[2], g_p_djrdv[3]) + func_cddb_ang_bw(i_b, link1_idx, g_djrdv, dyn_state, dyn_info, rigid_config) + func_cddb_ang_bw(i_b, link2_idx, -g_djrdv, dyn_state, dyn_info, rigid_config) + + # t3 = quat_mul(quat_mul_axis(inv_q2, domega), qdot0r) + g_m3 = gu.qd_quat_mul_grad_lhs(m3, qdot0r, g_t_quat) + g_qdot0r = gu.qd_quat_mul_grad_rhs(m3, qdot0r, g_t_quat) + g_inv_q2 = g_inv_q2 + gu.qd_quat_mul_grad_lhs(inv_q2, p_domega, g_m3) + g_p_domega3 = gu.qd_quat_mul_grad_rhs(inv_q2, p_domega, g_m3) + for j in qd.static(range(3)): + g_domega[j] = g_domega[j] + g_p_domega3[j + 1] + # qdot0r = quat_mul(qdot_body1, relpose); relpose is a model param + g_qdot_body1 = gu.qd_quat_mul_grad_lhs(qdot_body1, relpose, g_qdot0r) + # qdot_body1 = 0.5 * quat_mul([0, omega1], quat_body1) + g_p_omega1 = gu.qd_quat_mul_grad_lhs(p_omega1, quat_body1, 0.5 * g_qdot_body1) + g_quat1_bias = gu.qd_quat_mul_grad_rhs(p_omega1, quat_body1, 0.5 * g_qdot_body1) + for j in qd.static(range(3)): + g_omega1[j] = g_omega1[j] + g_p_omega1[j + 1] + + # t1 = quat_mul(quat_mul_axis(inv_quat(qdot_body2), domega), q_var) + g_m1 = gu.qd_quat_mul_grad_lhs(m1, q_var, g_t_quat) + g_q = g_q + gu.qd_quat_mul_grad_rhs(m1, q_var, g_t_quat) + g_inv_qdot2 = gu.qd_quat_mul_grad_lhs(inv_qdot2, p_domega, g_m1) + g_p_domega1 = gu.qd_quat_mul_grad_rhs(inv_qdot2, p_domega, g_m1) + for j in qd.static(range(3)): + g_domega[j] = g_domega[j] + g_p_domega1[j + 1] + # inv_quat flips the xyz signs of the incoming gradient + g_qdot_body2 = qd.Vector( + [g_inv_qdot2[0], -g_inv_qdot2[1], -g_inv_qdot2[2], -g_inv_qdot2[3]], dt=gs.qd_float + ) + # qdot_body2 = 0.5 * quat_mul([0, omega2], quat_body2) + g_p_omega2 = gu.qd_quat_mul_grad_lhs(p_omega2, quat_body2, 0.5 * g_qdot_body2) + g_quat2_bias = gu.qd_quat_mul_grad_rhs(p_omega2, quat_body2, 0.5 * g_qdot_body2) + for j in qd.static(range(3)): + g_omega2[j] = g_omega2[j] + g_p_omega2[j + 1] + + # domega = omega1 - omega2; omega_c = links.cd_ang[link_c] + g_omega1 = g_omega1 + g_domega + g_omega2 = g_omega2 - g_domega + dyn_state.links.cd_ang.grad[link1_idx, i_b] += g_omega1 + dyn_state.links.cd_ang.grad[link2_idx, i_b] += g_omega2 + # inv_q2 = inv_quat(quat_body2): (w, x, y, z) -> (w, -x, -y, -z) # => g_quat_body2 (from inv_q2 chain) = (g_inv_q2[0], -g_inv_q2[1], -g_inv_q2[2], -g_inv_q2[3]) g_quat2_from_inv = qd.Vector( @@ -1437,5 +1493,5 @@ def kernel_manual_add_equality_constraints_bw( dyn_state.links.pos.grad[link1_idx, i_b] += g_ga1 dyn_state.links.pos.grad[link2_idx, i_b] += g_ga2 - dyn_state.links.quat.grad[link1_idx, i_b] += g_quat1_anchor + g_quat1_from_q - dyn_state.links.quat.grad[link2_idx, i_b] += g_quat2_anchor + g_quat2_from_inv + dyn_state.links.quat.grad[link1_idx, i_b] += g_quat1_anchor + g_quat1_from_q + g_quat1_bias + dyn_state.links.quat.grad[link2_idx, i_b] += g_quat2_anchor + g_quat2_from_inv + g_quat2_bias diff --git a/genesis/engine/solvers/rigid/constraint/solver.py b/genesis/engine/solvers/rigid/constraint/solver.py index a7d1ebb64..5ebfd000f 100644 --- a/genesis/engine/solvers/rigid/constraint/solver.py +++ b/genesis/engine/solvers/rigid/constraint/solver.py @@ -449,13 +449,13 @@ def delete_weld_constraint(self, link1_idx, link2_idx, envs_idx=None): ) def backward(self): + """Adjoint solve of the constraint force computation. + + The caller must pre-populate the upstream gradient constraint_state.dL_dqacc (see + RigidSolver._constraint_force_grad).""" if not self._solver._requires_grad: gs.raise_exception("Please set `requires_grad` to True in SimOptions to enable differentiable mode.") - # Upstream gradient dL_dqacc is expected to be pre-populated in - # constraint_state.dL_dqacc by the caller (see - # kernel_load_dL_dqacc_from_acc_grad). - # 1. We first need to find a solution to A^T * u = g system. backward_constraint_solver.kernel_solve_adjoint_u( self.constraint_state, self._solver.dyn_info, self._solver.rigid_info, self._solver.rigid_config diff --git a/genesis/engine/solvers/rigid/rigid_solver.py b/genesis/engine/solvers/rigid/rigid_solver.py index 2f9fdb099..b0ea2464e 100644 --- a/genesis/engine/solvers/rigid/rigid_solver.py +++ b/genesis/engine/solvers/rigid/rigid_solver.py @@ -126,7 +126,6 @@ kernel_compute_mass_matrix, kernel_forward_dynamics, kernel_update_acc, - kernel_compute_qacc, kernel_forward_dynamics_without_qacc, update_qacc_from_qvel_delta, update_qvel, @@ -736,8 +735,10 @@ def _build_static_config(self): f"{type(self.sim.coupler).__name__} is not supported yet when requires_grad is True." ) - if getattr(self._options, "noslip_iterations", 0) > 0: + if self._options.noslip_iterations > 0: gs.raise_exception("Noslip is not supported yet when requires_grad is True.") + if self._options.enable_torsional_friction or self._options.enable_rolling_friction: + gs.raise_exception("Torsional and rolling friction are not supported yet when requires_grad is True.") def _create_data_manager(self): # We initialize data even if the solver is not active because the coupler needs arguments like @@ -1303,11 +1304,6 @@ def check_errno(self): gs.raise_exception("Invalid accelerations causing 'nan'. Please decrease Rigid simulation timestep.") if errno & array_class.ErrorCode.OVERFLOW_HIBERNATION_ISLANDS: gs.raise_exception("Contact island buffer overflow. Please increase RigidOptions 'max_collision_pairs'.") - if errno & array_class.ErrorCode.MANUAL_BW_UNIMPLEMENTED: - gs.raise_exception( - "Differentiable mode encountered a configuration (e.g. hibernation) that the manual backward " - "kernels do not support yet. Please disable it in this scene." - ) def _kernel_detect_collision(self): self.collider.clear() @@ -1347,40 +1343,68 @@ def _func_constraint_force(self): self.constraint_solver.resolve() def _constraint_force_grad(self): - # Backward pass for the constraint solver. - kernel_load_dL_dqacc_from_acc_grad(self.dyn_state, self.constraint_solver.constraint_state, self.rigid_config) + """Backward pass for the constraint solver: seed dL_dqacc from acc.grad, run the adjoint solve, fold its + outputs back into the autodiff grad fields, then reverse the constraint-row assembly.""" + constraint_state = self.constraint_solver.constraint_state + # Pure grad shuffles: in-place through zero-copy views when supported, kernel dispatch otherwise (see + # "Pure read-write data accessors on the hot path" in CLAUDE.md). gs.use_zerocopy encodes every zero-copy + # availability condition for these buffers (standalone dense allocations with DLPack-supported dtypes; the + # platform gates live in gs.init). torch (MPS) and quadrants do not share a compute stream on Metal, so the + # writes are flushed before the next kernels (see set_base_links_quat). + if gs.use_zerocopy: + acc_grad = qd_to_torch(self.dyn_state.dofs.acc.grad, copy=False) + dL_dqacc = qd_to_torch(constraint_state.dL_dqacc, copy=False) + dL_dqacc.copy_(acc_grad) + acc_grad.zero_() + if gs.backend == gs.metal: + torch.mps.synchronize() + else: + kernel_load_dL_dqacc_from_acc_grad(self.dyn_state, constraint_state, self.rigid_config) self.constraint_solver.backward() - kernel_accumulate_constraint_solver_grads( - self.dyn_state, self.constraint_solver.constraint_state, self.rigid_info, self.rigid_config - ) + if gs.use_zerocopy: + force_grad = qd_to_torch(self.dyn_state.dofs.force.grad, copy=False) + dL_dforce = qd_to_torch(constraint_state.dL_dforce, copy=False) + force_grad.add_(dL_dforce) + mass_mat_grad = qd_to_torch(self.rigid_info.mass_mat.grad, copy=False) + dL_dM = qd_to_torch(constraint_state.dL_dM, copy=False) + mass_mat_grad.add_(dL_dM) + if gs.backend == gs.metal: + torch.mps.synchronize() + else: + kernel_accumulate_constraint_solver_grads( + self.dyn_state, constraint_state, self.rigid_info, self.rigid_config + ) kernel_manual_add_equality_constraints_bw( - self.dyn_state, self.constraint_solver.constraint_state, self.dyn_info, self.rigid_info, self.rigid_config + self.dyn_state, constraint_state, self.dyn_info, self.rigid_info, self.rigid_config ) kernel_manual_add_frictionloss_constraints_bw( - self.dyn_state, self.constraint_solver.constraint_state, self.dyn_info, self.rigid_info, self.rigid_config + self.dyn_state, constraint_state, self.dyn_info, self.rigid_info, self.rigid_config ) if self._enable_collision: collider_state = self.collider._collider_state - collider_state.contact_data.pos.grad.fill(0.0) - collider_state.contact_data.normal.grad.fill(0.0) - collider_state.contact_data.penetration.grad.fill(0.0) + qd_zero_grad(collider_state.contact_data.pos) + qd_zero_grad(collider_state.contact_data.normal) + qd_zero_grad(collider_state.contact_data.penetration) + # One flush for the zeroing batch; see qd_zero_grad in misc.py. + if gs.use_zerocopy and gs.backend == gs.metal: + torch.mps.synchronize() kernel_manual_add_collision_constraints_bw( self.dyn_state, collider_state, - self.constraint_solver.constraint_state, + constraint_state, self.dyn_info, self.rigid_info, self.rigid_config, ) self.collider.backward_narrowphase() - if self._options.enable_joint_limit: + if self._enable_joint_limit: kernel_manual_add_joint_limit_constraints_bw( self.dyn_state, self.collider._collider_state, - self.constraint_solver.constraint_state, + constraint_state, self.dyn_info, self.rigid_info, self.rigid_config, @@ -1544,13 +1568,17 @@ def reset_grad(self): qd_zero_grad(self.dyn_state_adjoint_cache.joints) qd_zero_grad(self.dyn_state_adjoint_cache.geoms) qd_zero_grad(self._rigid_adjoint_cache) + # One flush for the zeroing batch; see qd_zero_grad in misc.py. + if gs.use_zerocopy and gs.backend == gs.metal: + torch.mps.synchronize() def _update_cartesian_grad(self, envs_idx): """Forward-replay the post-integrate cartesian-space update (FK -> COM -> geom poses -> velocity) under - is_backward=True, then reverse it stage by stage: velocity and forward kinematics are reversed manually - (kernel_manual_*_bw in manual_bw.py), while COM and the link->geom transform are reversed by Quadrants - autodiff (.grad). Shared by the post-integrate reverse and the first-substep initial-state reverse in - substep_pre_coupling_grad. + is_backward=True, then reverse it stage by stage. + + Velocity and forward kinematics are reversed manually (kernel_manual_*_bw in manual_bw.py), while COM and + the link->geom transform are reversed by Quadrants autodiff (.grad). Shared by the post-integrate reverse + and the first-substep initial-state reverse in substep_pre_coupling_grad. """ # Forward replay in dependency order (FK -> COM -> geoms -> velocity). kernel_forward_kinematics_replay( @@ -1562,19 +1590,16 @@ def _update_cartesian_grad(self, envs_idx): envs_idx, self.dyn_state, self.dyn_info, self.rigid_info, self.rigid_config, is_backward=True ) - # Reverse in opposite order (velocity -> COM -> geoms -> FK). - kernel_manual_forward_velocity_bw( - self.dyn_state, self.dyn_info, self.rigid_info, self.rigid_config, self._errno - ) + # Reverse the stages: velocity first, forward kinematics last. COM and geoms both consume only FK + # outputs, so their mutual order is free. + kernel_manual_forward_velocity_bw(self.dyn_state, self.dyn_info, self.rigid_info, self.rigid_config) kernel_COM_links_replay.grad( self.dyn_state, self.dyn_info, self.rigid_info, self.rigid_config, is_backward=True ) kernel_update_geoms_replay.grad( self.dyn_state, self.dyn_info, self.rigid_info, self.rigid_config, is_backward=True ) - kernel_manual_forward_kinematics_bw( - self.dyn_state, self.dyn_info, self.rigid_info, self.rigid_config, self._errno - ) + kernel_manual_forward_kinematics_bw(self.dyn_state, self.dyn_info, self.rigid_info, self.rigid_config) def substep_pre_coupling_grad(self, f): # Change to backward mode @@ -1631,7 +1656,6 @@ def substep_pre_coupling_grad(self, f): if not self._disable_constraint: self._constraint_force_grad() else: - # Manual backward for func_compute_qacc via the implicit function theorem. kernel_manual_compute_qacc_bw(self.dyn_state, self.dyn_info, self.rigid_info, self.rigid_config) kernel_copy_acc(f, self.dyn_state, self._rigid_adjoint_cache, self.rigid_config) @@ -1897,8 +1921,8 @@ def save_ckpt(self, ckpt_name): if ckpt_name not in self._ckpt: self._ckpt[ckpt_name] = dict() - # copy=True required: with the zerocopy backend qd_to_numpy returns a - # view, so later substeps would overwrite this ckpt's buffer in place. + # copy=True required: with the zerocopy backend qd_to_numpy returns a view, so later substeps would + # overwrite this ckpt's buffer in place. self._ckpt[ckpt_name]["qpos"] = qd_to_numpy(self._rigid_adjoint_cache.qpos, copy=True) self._ckpt[ckpt_name]["dofs_vel"] = qd_to_numpy(self._rigid_adjoint_cache.dofs_vel, copy=True) self._ckpt[ckpt_name]["dofs_acc"] = qd_to_numpy(self._rigid_adjoint_cache.dofs_acc, copy=True) @@ -1913,6 +1937,10 @@ def load_ckpt(self, ckpt_name): self.dyn_state.dofs.acc.from_numpy(self._ckpt[ckpt_name]["dofs_acc"][0]) if not self._enable_mujoco_compatibility: + # Mirror the post-integrate refresh of kernel_step_2: the replayed substeps skip their own cartesian / + # velocity updates (is_forward_pos_updated / is_forward_vel_updated), so both must be recomputed here + # from the restored qpos / vel. A stale link velocity would corrupt the velocity-product terms of every + # backward primal in the window. kernel_update_cartesian_space( self.dyn_state, self.dyn_info, @@ -1921,6 +1949,14 @@ def load_ckpt(self, ckpt_name): force_update_fixed_geoms=False, is_backward=False, ) + kernel_forward_velocity( + self._scene._sanitize_envs_idx(None), + self.dyn_state, + self.dyn_info, + self.rigid_info, + self.rigid_config, + is_backward=False, + ) for entity in self._entities: entity.load_ckpt(ckpt_name) @@ -3200,7 +3236,7 @@ def kernel_step_2( func_update_acc(dyn_state, dyn_info, rigid_info, rigid_config, update_cacc=True, is_backward=is_backward) if qd.static(rigid_config.integrator != gs.integrator.approximate_implicitfast): - func_implicit_damping(dyn_state, dyn_info, rigid_info, rigid_config, is_backward) + func_implicit_damping(dyn_state, dyn_info, rigid_info, rigid_config) func_integrate(dyn_state, dyn_info, rigid_info, rigid_config, is_backward) diff --git a/genesis/utils/array_class.py b/genesis/utils/array_class.py index 19a206257..e623bd825 100644 --- a/genesis/utils/array_class.py +++ b/genesis/utils/array_class.py @@ -87,7 +87,6 @@ class ErrorCode(IntEnum): INVALID_FORCE_NAN = 0b00000000000000000000000000001000 INVALID_ACC_NAN = 0b00000000000000000000000000010000 OVERFLOW_CONTACTS = 0b00000000000000000000000000100000 - MANUAL_BW_UNIMPLEMENTED = 0b00000000000000000000000001000000 # =========================================== RigidInfo =========================================== diff --git a/genesis/utils/geom.py b/genesis/utils/geom.py index edb446491..030ef130d 100644 --- a/genesis/utils/geom.py +++ b/genesis/utils/geom.py @@ -137,8 +137,11 @@ def qd_rotvec_to_quat(rotvec, eps): def qd_rotvec_to_quat_grad_rotvec(rotvec, eps, out_grad): """Adjoint of qd_rotvec_to_quat(rotvec, eps) with respect to rotvec, given the upstream quaternion gradient. - With theta_reg = sqrt(|rotvec|^2 + eps^2), c = cos(theta_reg / 2) and sinc = sin(theta_reg / 2) / theta_reg, the - forward is quat = (c, sinc * rotvec). Chain rule through theta_reg gives, per component i: + Differentiates the eps-regularized surrogate quat = (c, sinc * rotvec) with theta_reg = sqrt(|rotvec|^2 + eps^2), + c = cos(theta_reg / 2) and sinc = sin(theta_reg / 2) / theta_reg. The surrogate deviates from the exact forward + (which branches at |rotvec| = eps and applies a first-order renormalization) by O(eps^2) in the main branch and + smooths over the constant-identity branch, trading exactness at the branch point for a gradient that stays + finite and continuous through rotvec = 0. Chain rule through theta_reg gives, per component i: d(quat[0])/d(rotvec[i]) = -0.5 * sin(theta_reg / 2) * rotvec[i] / theta_reg d(quat[1+j])/d(rotvec[i]) = delta(i, j) * sinc + rotvec[j] * (0.5 * c - sinc) / theta_reg^2 * rotvec[i] """ @@ -215,6 +218,93 @@ def qd_quat_to_xyz(quat, eps): return qd.Vector([roll, pitch, yaw], dt=gs.qd_float) +@qd.func +def qd_quat_to_xyz_grad_quat(quat, eps, out_grad): + """Adjoint of qd_quat_to_xyz(quat, eps) with respect to quat, given the upstream Euler-angle gradient. + + Mirrors the forward branch structure: each angle is an atan2 of quadratic forms u = s * P(quat) with + s = 2 / |quat|^2, so du/dq = s * (dP/dq - 2 P q / |quat|^2) and datan2(y, x) = (x dy - y dx) / (x^2 + y^2). + The gradient is zero in the degenerate |quat|^2 <= eps branch (constant forward), and the roll / yaw terms + switch to the gimbal-lock forms when cosp <= eps, matching the forward exactly. + """ + g_quat = qd.Vector.zero(gs.qd_float, 4) + quat_norm_sqr = quat.norm_sqr() + if quat_norm_sqr > eps: + s = 2.0 / quat_norm_sqr + q_w, q_x, q_y, q_z = quat + + # u = s * P with dP/dq per quadratic form; du/dq folds the ds/dq = -2 s q / |quat|^2 term + p_siny = q_w * q_z - q_x * q_y + p_cosy = q_y * q_y + q_z * q_z + p_pitch = q_x * q_z + q_w * q_y + p_roll_num = q_w * q_x - q_y * q_z + p_roll_den = q_x * q_x + q_y * q_y + p_gimbal_num = q_w * q_z + q_x * q_y + p_gimbal_den = q_x * q_x + q_z * q_z + u_siny = s * p_siny + u_cosy = 1.0 - s * p_cosy + u_pitch = s * p_pitch + u_roll_num = s * p_roll_num + u_roll_den = 1.0 - s * p_roll_den + cosp = qd.sqrt(u_cosy * u_cosy + u_siny * u_siny) + + g_roll = out_grad[0] + g_pitch = out_grad[1] + g_yaw = out_grad[2] + + # pitch = atan2(u_pitch, cosp) with cosp = |(u_siny, u_cosy)| + denom_pitch = u_pitch * u_pitch + cosp * cosp + g_u_pitch = g_pitch * cosp / denom_pitch + g_cosp = -g_pitch * u_pitch / denom_pitch + g_u_siny = gs.qd_float(0.0) + g_u_cosy = gs.qd_float(0.0) + if cosp > eps: + g_u_siny = g_cosp * u_siny / cosp + g_u_cosy = g_cosp * u_cosy / cosp + + g_u_roll_num = gs.qd_float(0.0) + g_u_roll_den = gs.qd_float(0.0) + g_u_gimbal_num = gs.qd_float(0.0) + g_u_gimbal_den = gs.qd_float(0.0) + if cosp > eps: + # roll = atan2(u_roll_num, u_roll_den); yaw = atan2(u_siny, u_cosy) + denom_roll = u_roll_num * u_roll_num + u_roll_den * u_roll_den + g_u_roll_num = g_roll * u_roll_den / denom_roll + g_u_roll_den = -g_roll * u_roll_num / denom_roll + denom_yaw = u_siny * u_siny + u_cosy * u_cosy + g_u_siny = g_u_siny + g_yaw * u_cosy / denom_yaw + g_u_cosy = g_u_cosy - g_yaw * u_siny / denom_yaw + else: + # Gimbal lock: roll = 0 and yaw = atan2(s * (wz + xy), 1 - s * (xx + zz)) + u_gimbal_num = s * p_gimbal_num + u_gimbal_den = 1.0 - s * p_gimbal_den + denom_gimbal = u_gimbal_num * u_gimbal_num + u_gimbal_den * u_gimbal_den + g_u_gimbal_num = g_yaw * u_gimbal_den / denom_gimbal + g_u_gimbal_den = -g_yaw * u_gimbal_num / denom_gimbal + + # Fold each u = s * P (or 1 - s * P) back to quat: du/dq_k = +-s * (dP/dq_k - 2 P q_k / |quat|^2) + two_over_n = 2.0 / quat_norm_sqr + d_p_siny = qd.Vector([q_z, -q_y, -q_x, q_w], dt=gs.qd_float) + d_p_cosy = qd.Vector([0.0, 0.0, 2.0 * q_y, 2.0 * q_z], dt=gs.qd_float) + d_p_pitch = qd.Vector([q_y, q_z, q_w, q_x], dt=gs.qd_float) + d_p_roll_num = qd.Vector([q_x, q_w, -q_z, -q_y], dt=gs.qd_float) + d_p_roll_den = qd.Vector([0.0, 2.0 * q_x, 2.0 * q_y, 0.0], dt=gs.qd_float) + d_p_gimbal_num = qd.Vector([q_z, q_y, q_x, q_w], dt=gs.qd_float) + d_p_gimbal_den = qd.Vector([0.0, 2.0 * q_x, 0.0, 2.0 * q_z], dt=gs.qd_float) + for j in qd.static(range(4)): + q_j = quat[j] + g_quat[j] = ( + g_u_siny * s * (d_p_siny[j] - p_siny * two_over_n * q_j) + - g_u_cosy * s * (d_p_cosy[j] - p_cosy * two_over_n * q_j) + + g_u_pitch * s * (d_p_pitch[j] - p_pitch * two_over_n * q_j) + + g_u_roll_num * s * (d_p_roll_num[j] - p_roll_num * two_over_n * q_j) + - g_u_roll_den * s * (d_p_roll_den[j] - p_roll_den * two_over_n * q_j) + + g_u_gimbal_num * s * (d_p_gimbal_num[j] - p_gimbal_num * two_over_n * q_j) + - g_u_gimbal_den * s * (d_p_gimbal_den[j] - p_gimbal_den * two_over_n * q_j) + ) + return g_quat + + @qd.func def qd_quat_to_rotvec(quat, eps): q_w, q_x, q_y, q_z = quat @@ -535,6 +625,39 @@ def imp_aref(params, neg_penetration, vel, pos): return imp, aref +@qd.func +def imp_aref_grad(params, neg_penetration): + """Replay of imp_aref's impedance together with its derivative w.r.t. the impedance driver. + + Returns (imp, b, k, d_imp_d_x) with x = |neg_penetration| / width: the shared factors every manual constraint + reverse needs to chain gradients through imp_aref, whose outputs are aref = -b * vel - k * imp * pos and the + impedance imp feeding diag / efc_D. d_imp_d_x is zero outside the smooth band (imp_raw clamped at dmin / dmax or + x > 1), matching the piecewise-flat regions of the forward. + """ + timeconst, dampratio, dmin, dmax, width, mid, power = params + + imp_x = qd.abs(neg_penetration) / width + imp_a_coef = 1.0 / mid ** (power - 1.0) + imp_b_coef = 1.0 / (1.0 - mid) ** (power - 1.0) + imp_a = imp_a_coef * imp_x**power + imp_b = 1.0 - imp_b_coef * (1.0 - imp_x) ** power + imp_y = imp_a if imp_x < mid else imp_b + imp_raw = dmin + imp_y * (dmax - dmin) + imp = qd.math.clamp(imp_raw, dmin, dmax) + imp = dmax if imp_x > 1.0 else imp + + b = 2.0 / (dmax * timeconst) + k = 1.0 / (dmax * dmax * timeconst * timeconst * dampratio * dampratio) + + d_imp_y_d_x = power * imp_a_coef * imp_x ** (power - 1.0) + if imp_x >= mid: + d_imp_y_d_x = power * imp_b_coef * (1.0 - imp_x) ** (power - 1.0) + d_imp_d_x = gs.qd_float(0.0) + if imp_raw > dmin and imp_raw < dmax and imp_x <= 1.0: + d_imp_d_x = (dmax - dmin) * d_imp_y_d_x + return imp, b, k, d_imp_d_x + + # ------------------------------------------------------------------------------------ # -------------------------------- torch and numpy ----------------------------------- # ------------------------------------------------------------------------------------ diff --git a/genesis/utils/misc.py b/genesis/utils/misc.py index b2db8f2d8..751bfeca9 100644 --- a/genesis/utils/misc.py +++ b/genesis/utils/misc.py @@ -793,8 +793,10 @@ def qd_zero_grad(value) -> None: Reverse-mode accumulation in Genesis writes through `qd.atomic_add`, so adjoint buffers must start at zero between consecutive `loss.backward()` calls. Solvers call this from `reset_grad` to clear all owned adjoint storage without - enumerating fields by name. Zeroing goes through `qd_to_torch(grad, copy=False).zero_()`, a contiguous in-place - memset on the underlying device memory - no Quadrants kernel launch. + enumerating fields by name. Zeroing goes through an in-place `zero_()` on the zero-copy torch view of each grad + buffer, a contiguous memset on the underlying device memory. The writes are left unsynchronized so a caller can + batch many calls under a single flush: on Metal, call `torch.mps.synchronize()` after the batch and before the + next quadrants kernel reads the buffers (see set_base_links_quat). """ if value is None: return @@ -804,10 +806,12 @@ def qd_zero_grad(value) -> None: grad = value.grad if gs.use_zerocopy: try: - qd_to_torch(grad, copy=False).zero_() + grad_view = qd_to_torch(grad, copy=False) + grad_view.zero_() except ValueError: - # No zero-copy view for this buffer (e.g. a field past 2**31 bytes in its SNode tree); fill it in - # place through quadrants instead. + # No zero-copy view for this buffer (e.g. an interleaved AOS struct member, or a field whose + # in-tree byte offset the installed torch cannot carry through DLPack); fill it in place through + # quadrants instead. grad.fill(0.0) else: grad.fill(0.0) diff --git a/tests/grad/conftest.py b/tests/grad/conftest.py index bdae65f54..9abedc937 100644 --- a/tests/grad/conftest.py +++ b/tests/grad/conftest.py @@ -3,10 +3,11 @@ import pytest -def _add_hinge_arm(parent, body_name, pos, **joint_kwargs): - """Add a 1-DOF hinge arm link (y-axis hinge, capsule geom, explicit inertia) and return its body element.""" +def _add_hinge_arm(parent, body_name, pos, axis="0 1 0", **joint_kwargs): + """Add a 1-DOF hinge arm link (y-axis hinge by default, capsule geom, explicit inertia) and return its body + element.""" body = ET.SubElement(parent, "body", name=body_name, pos=pos) - ET.SubElement(body, "joint", type="hinge", axis="0 1 0", **joint_kwargs) + ET.SubElement(body, "joint", type="hinge", axis=axis, **joint_kwargs) ET.SubElement(body, "inertial", mass="0.5", pos="0.1 0 0", diaginertia="0.01 0.01 0.01") ET.SubElement(body, "geom", type="capsule", fromto="0 0 0 0.2 0 0", size="0.02", contype="0", conaffinity="0") return body @@ -27,7 +28,7 @@ def grad_free(): def grad_revolute(): mjcf = ET.Element("mujoco", model="revolute") worldbody = ET.SubElement(mjcf, "worldbody") - _add_hinge_arm(worldbody, "arm", "0 0 0") + _add_hinge_arm(worldbody, "arm", "0 0 0", stiffness="2.0") return ET.tostring(mjcf, encoding="unicode") @@ -88,13 +89,18 @@ def grad_free_with_revolute(): @pytest.fixture(scope="session") -def grad_revolute_chain3(): +def grad_chain3(): + # Hinge -> offset slide -> hinge: the middle slide joint carries a position offset so its anchor depends on the + # moving parent orientation, and the chain mixes joint types within one entity. mjcf = ET.Element("mujoco", model="chain3") worldbody = ET.SubElement(mjcf, "worldbody") parent = worldbody for name in ("l1", "l2", "l3"): body = ET.SubElement(parent, "body", name=name, pos="0 0 0" if name == "l1" else "0.2 0 0") - ET.SubElement(body, "joint", type="hinge", axis="0 1 0") + if name == "l2": + ET.SubElement(body, "joint", type="slide", axis="1 0 0", pos="0.05 0.02 0.03") + else: + ET.SubElement(body, "joint", type="hinge", axis="0 1 0") ET.SubElement(body, "inertial", mass="0.3", pos="0.1 0 0", diaginertia="0.005 0.005 0.005") ET.SubElement(body, "geom", type="capsule", fromto="0 0 0 0.2 0 0", size="0.02", contype="0", conaffinity="0") parent = body @@ -104,7 +110,6 @@ def grad_revolute_chain3(): @pytest.fixture(scope="session") def grad_slider_limit(): mjcf = ET.Element("mujoco", model="slider_limit") - ET.SubElement(mjcf, "option", gravity="0 0 0") worldbody = ET.SubElement(mjcf, "worldbody") body = ET.SubElement(worldbody, "body", name="cart", pos="0 0 0") ET.SubElement(body, "joint", name="slider", type="slide", axis="1 0 0", range="-4 4", damping="0.0") @@ -141,30 +146,36 @@ def grad_hinge_pair_joint_eq_quadratic(): @pytest.fixture(scope="session") def grad_connect_loop(): + # arm2 hangs off arm1 and the connect closes the loop within one kinematic tree: the constraint rows then share + # arm1's dof across both chains (dedup path) and both anchors move with the chain (velocity-product bias). mjcf = ET.Element("mujoco", model="connect_loop") worldbody = ET.SubElement(mjcf, "worldbody") - _add_hinge_arm(worldbody, "arm1", "0 0 0", name="j1") - _add_hinge_arm(worldbody, "arm2", "0 0.3 0", name="j2") + arm1 = _add_hinge_arm(worldbody, "arm1", "0 0 0", name="j1") + _add_hinge_arm(arm1, "arm2", "0.2 0 0", name="j2") equality = ET.SubElement(mjcf, "equality") ET.SubElement( - equality, "connect", body1="arm1", body2="arm2", anchor="0.2 0 0", solimp="0.95 0.99 0.001", solref="0.005 1" + equality, "connect", body1="arm2", body2="arm1", anchor="0.2 0 0", solimp="0.95 0.99 0.001", solref="0.005 1" ) return ET.tostring(mjcf, encoding="unicode") @pytest.fixture(scope="session") def grad_weld_pair(): + # arm2 hangs off arm1 about a skew axis and the weld ties it back to arm1: the constraint rows share arm1's dof + # across both chains (dedup path), and the nested hinge's angular velocity-product bias (parent angular velocity + # cross child axis) is nonzero. A weld between parallel-axis or single-hinge chains would leave the rotation + # rows' velocity bias identically zero and its adjoint untested. mjcf = ET.Element("mujoco", model="weld_pair") worldbody = ET.SubElement(mjcf, "worldbody") - _add_hinge_arm(worldbody, "arm1", "0 0 0", name="j1") - _add_hinge_arm(worldbody, "arm2", "0 0.3 0", name="j2") + arm1 = _add_hinge_arm(worldbody, "arm1", "0 0 0", name="j1") + _add_hinge_arm(arm1, "arm2", "0.2 0 0", name="j2", axis="1 0 0") equality = ET.SubElement(mjcf, "equality") ET.SubElement( equality, "weld", - body1="arm1", - body2="arm2", - relpose="0 -0.3 0 1 0 0 0", + body1="arm2", + body2="arm1", + relpose="0.2 0 0 1 0 0 0", solimp="0.95 0.99 0.001", solref="0.005 1", ) @@ -203,7 +214,6 @@ def grad_all_eq_fric(): @pytest.fixture(scope="session") def grad_cartpole(): mjcf = ET.Element("mujoco", model="cartpole") - ET.SubElement(mjcf, "option", gravity="0 0 -9.81") worldbody = ET.SubElement(mjcf, "worldbody") cart = ET.SubElement(worldbody, "body", name="cart", pos="0 0 0") ET.SubElement(cart, "joint", name="slider", type="slide", axis="1 0 0", range="-4 4", damping="0.0") diff --git a/tests/grad/test_grad_tape.py b/tests/grad/test_grad_tape.py index 2b06fac6c..9f58fcfa6 100644 --- a/tests/grad/test_grad_tape.py +++ b/tests/grad/test_grad_tape.py @@ -1,26 +1,24 @@ -# Autodiff-tape behavior: horizon truncation (snapshot + reset between two backward windows must reproduce the -# gradients and states of independent fresh scenes), and gradient-source parity (entity state vs rigid-solver state). import numpy as np import pytest import torch import genesis as gs -from genesis.utils.misc import qd_to_numpy, tensor_to_array +from genesis.utils.misc import tensor_to_array -from ..utils import assert_allclose +from ..utils import assert_allclose, assert_equal from .utils import make_diff_scene_pair @pytest.mark.required @pytest.mark.parametrize("model_name", ["grad_free", "grad_revolute", "grad_free_with_revolute"]) -def test_grad_horizon_truncation_matches_independent_scenes(model_name, request, show_viewer): +def test_horizon_truncation_matches_independent_scenes(model_name, request, show_viewer): mjcf = request.getfixturevalue(model_name) tol = dict(atol=1e-5, rtol=1e-4) horizon = 5 B = 2 def build(show_viewer=False): - pair = make_diff_scene_pair(mjcf, n_envs=2, substeps=4, gravity=(0.0, 0.0, 0.0), show_viewer=show_viewer) + pair = make_diff_scene_pair(mjcf, n_envs=B, substeps=4, gravity=(0.0, 0.0, 0.0), show_viewer=show_viewer) return pair.scene_ana, pair.entity_ana def run_segment(scene, entity, velocity): @@ -30,7 +28,7 @@ def run_segment(scene, entity, velocity): return (scene.rigid_solver.get_state().qpos ** 2).sum() def read_qpos(scene): - return qd_to_numpy(scene.rigid_solver.rigid_info.qpos, copy=True) + return tensor_to_array(scene.rigid_solver.get_state().qpos) # Scene A: one scene, snapshot + reset between two horizons. scene_a, robot_a = build(show_viewer=show_viewer) @@ -43,7 +41,7 @@ def read_qpos(scene): scene_a.backward(loss_h1_a) # backward consumes the input buffer, so the step / substep counters (which index it) reset to 0 while the # restored physics state carries over; horizon 2 records a fresh tape from 0. - assert scene_a._t == 0 and scene_a._sim._cur_substep_global == 0 + assert scene_a.t == 0 and scene_a._sim._cur_substep_global == 0 grad1_a = tensor_to_array(v1_a.grad).copy() v2_a = gs.tensor(v2, dtype=gs.tc_float, requires_grad=True) @@ -61,8 +59,8 @@ def read_qpos(scene): snapshot_b = scene_b.backward(loss_h1_b) grad1_b = tensor_to_array(v1_b.grad).copy() - assert_allclose(qpos_mid_a, qpos_mid_b, atol=0, rtol=0) - assert_allclose(float(tensor_to_array(loss_h1_a)), float(tensor_to_array(loss_h1_b)), atol=0, rtol=0) + assert_equal(qpos_mid_a, qpos_mid_b) + assert_equal(loss_h1_a, loss_h1_b) assert_allclose(grad1_a, grad1_b, **tol) # Scene C: fresh scene resumed from B's mid-trajectory snapshot. @@ -74,31 +72,32 @@ def read_qpos(scene): scene_c.backward(loss_h2_c) grad2_c = tensor_to_array(v2_c.grad).copy() - assert_allclose(qpos_end_a, qpos_end_c, atol=0, rtol=0) - assert_allclose(float(tensor_to_array(loss_h2_a)), float(tensor_to_array(loss_h2_c)), atol=0, rtol=0) + assert_equal(qpos_end_a, qpos_end_c) + assert_equal(loss_h2_a, loss_h2_c) assert_allclose(grad2_a, grad2_c, **tol) @pytest.mark.slow @pytest.mark.required -@pytest.mark.parametrize("backend", [gs.cpu, gs.gpu]) -def test_rigid_sim_vs_solver_state_grad_parity(show_viewer): +def test_sim_vs_solver_state_grad_parity(show_viewer): scene = gs.Scene( sim_options=gs.options.SimOptions( - dt=0.01, gravity=(0.0, 0.0, 0.0), requires_grad=True, ), rigid_options=gs.options.RigidOptions( enable_collision=False, ), + viewer_options=gs.options.ViewerOptions( + camera_pos=(1.2, -1.2, 0.8), + camera_lookat=(0.0, 0.0, 0.2), + ), show_viewer=show_viewer, ) robot = scene.add_entity( gs.morphs.Box( size=(0.1, 0.1, 0.1), - pos=(0, 0, 0), - ) + ), ) scene.build() diff --git a/tests/grad/test_hybrid_push.py b/tests/grad/test_hybrid_push.py index c64097445..673b3985c 100644 --- a/tests/grad/test_hybrid_push.py +++ b/tests/grad/test_hybrid_push.py @@ -1,5 +1,3 @@ -# Reverse-mode gradient through a hybrid rigid-tool / MPM-object push: the per-step stick velocities that move the -# deformable object toward a goal must carry non-zero gradients, except the final step which cannot affect the loss. import pytest import torch @@ -9,7 +7,7 @@ @pytest.mark.slow # ~350s @pytest.mark.required @pytest.mark.debug(False) -def test_hybrid_mpm_tool_push_grad(show_viewer): +def test_mpm_tool_push_grad(show_viewer): HORIZON = 10 scene = gs.Scene( diff --git a/tests/grad/test_rigid_collision.py b/tests/grad/test_rigid_collision.py index 50dc41d44..30e8b36d2 100644 --- a/tests/grad/test_rigid_collision.py +++ b/tests/grad/test_rigid_collision.py @@ -1,6 +1,3 @@ -# Differentiable-contact gradient checks: per-step contact-force adjoints (box-box and plane-convex, verified -# against contact-preserving finite differences), the forward convex-contact detection path, the unsupported -# smooth-pair guard, and the low-level contact-detection and constraint-solver backward passes. import numpy as np import pytest import torch @@ -13,54 +10,76 @@ from ..utils import assert_allclose -def _build_contact_scene(shape, mjcf_capsule, *, requires_grad, show_viewer=False): - scene = gs.Scene( - sim_options=gs.options.SimOptions( - dt=0.01, - substeps=2, - gravity=(0.0, 0.0, -9.81), - requires_grad=requires_grad, - ), - rigid_options=gs.options.RigidOptions( - enable_collision=True, - enable_self_collision=False, - enable_joint_limit=False, - disable_constraint=False, - use_hibernation=False, - use_contact_island=False, - box_box_detection=False, - ), - viewer_options=gs.options.ViewerOptions( - camera_pos=(1.2, -1.2, 0.8), - camera_lookat=(0.0, 0.0, 0.2), - ), - show_viewer=show_viewer, - ) - if shape == "ground_box": - scene.add_entity(gs.morphs.Box(size=(2.0, 2.0, 0.2), pos=(0.0, 0.0, 0.1), fixed=True)) - obj = scene.add_entity(gs.morphs.Box(size=(0.4, 0.4, 0.4), pos=(0.0, 0.0, 0.4))) - else: - scene.add_entity(gs.morphs.Plane()) - if shape == "box": - obj = scene.add_entity(gs.morphs.Box(size=(0.4, 0.4, 0.4), pos=(0.0, 0.0, 0.3))) - elif shape == "sphere": - obj = scene.add_entity(gs.morphs.Sphere(radius=0.2, pos=(0.0, 0.0, 0.3))) - elif shape == "capsule": - obj = scene.add_entity(gs.morphs.MJCF(file=mjcf_capsule, align=False)) - else: - raise ValueError(shape) - scene.build(n_envs=0) - return scene, obj - - -def _n_contacts(scene): - return qd_to_numpy(scene.rigid_solver.collider._collider_state.n_contacts)[0] - - @pytest.mark.required @pytest.mark.parametrize("backend", [gs.cpu, gs.gpu]) @pytest.mark.parametrize("shape", ["ground_box", "box", "sphere", "capsule"]) -def test_rigid_contact_per_step_force_grad_matches_fd(shape, grad_capsule, precision, show_viewer): +def test_contact_per_step_force_grad_matches_fd(shape, grad_capsule, precision, show_viewer): + def _build_contact_scene(shape, mjcf_capsule, *, requires_grad, show_viewer=False): + scene = gs.Scene( + sim_options=gs.options.SimOptions( + substeps=2, + requires_grad=requires_grad, + ), + rigid_options=gs.options.RigidOptions( + # A non-neutral impratio makes the regularized cone coefficient observable to finite differences + impratio=2.0, + ), + viewer_options=gs.options.ViewerOptions( + camera_pos=(1.2, -1.2, 0.8), + camera_lookat=(0.0, 0.0, 0.2), + ), + show_viewer=show_viewer, + ) + if shape == "ground_box": + scene.add_entity( + gs.morphs.Box( + size=(2.0, 2.0, 0.2), + pos=(0.0, 0.0, 0.1), + fixed=True, + ), + vis_mode="collision", + ) + obj = scene.add_entity( + gs.morphs.Box( + size=(0.4, 0.4, 0.4), + pos=(0.0, 0.0, 0.4), + ), + vis_mode="collision", + ) + else: + scene.add_entity(gs.morphs.Plane()) + if shape == "box": + obj = scene.add_entity( + gs.morphs.Box( + size=(0.4, 0.4, 0.4), + pos=(0.0, 0.0, 0.3), + ), + vis_mode="collision", + ) + elif shape == "sphere": + obj = scene.add_entity( + gs.morphs.Sphere( + radius=0.2, + pos=(0.0, 0.0, 0.3), + ), + vis_mode="collision", + ) + elif shape == "capsule": + obj = scene.add_entity( + gs.morphs.MJCF( + file=mjcf_capsule, + align=False, + ), + vis_mode="collision", + ) + else: + raise ValueError(shape) + scene.build(n_envs=0) + return scene, obj + + def _n_contacts(scene): + return qd_to_numpy(scene.rigid_solver.collider._collider_state.n_contacts)[0] + # Rest z puts the body's lowest point on its support: box / sphere half extent 0.2, upright capsule # radius 0.1 + half length 0.2 = 0.3, box-on-ground centered at 0.40. rest_z = {"ground_box": 0.40, "box": 0.20, "sphere": 0.20, "capsule": 0.30}[shape] @@ -68,16 +87,18 @@ def test_rigid_contact_per_step_force_grad_matches_fd(shape, grad_capsule, preci n_settle = 12 n_steps = 2 eps = 1e-2 - # Contact force gradients are tiny (stiff contact barely moves); fp32 tolerates the coarser finite-difference floor. - fd_atol = 1e-10 if precision == "64" else 5e-7 + # Contact force gradients are tiny (stiff contact barely moves): at fp32 the FD reference mostly rounds to + # zero, so the absolute tolerance carries the check; at fp64 the relative tolerance does and the absolute one + # only guards the degenerate-FD case. + fd_rtol = 1e-6 if precision == "64" else 2e-3 + fd_atol = 1e-14 if precision == "64" else 1e-6 base_force = np.array([0.0, 0.0, -8.0, 0.0, 0.0, 0.0]) init_force = np.broadcast_to(base_force, (n_steps, 6)).copy() def settle(scene, obj): obj.set_dofs_position(gs.tensor(rest_dofs, dtype=gs.tc_float).sceneless()) - zero = gs.tensor([0.0] * 6, dtype=gs.tc_float) for _ in range(n_settle): - obj.control_dofs_force(zero) + obj.control_dofs_force(0.0) scene.step() scene_ana, obj_ana = _build_contact_scene(shape, grad_capsule, requires_grad=True, show_viewer=show_viewer) @@ -94,6 +115,9 @@ def settle(scene, obj): scene_ana.backward(loss) ana = np.stack([tensor_to_array(f.grad) for f in forces]) + # The FD reference must also run in diff mode: differentiable scenes route contacts through diff_gjk, whose + # contact set and forces differ slightly from the production narrowphase, so a production-mode reference would + # measure that forward gap instead of the gradient. scene_fd, obj_fd = _build_contact_scene(shape, grad_capsule, requires_grad=True) def loss_at(perturbed): @@ -111,29 +135,44 @@ def loss_at(perturbed): minus = init_force.copy() minus[t, 2] -= eps fd_z = (loss_at(plus) - loss_at(minus)) / (2 * eps) - assert_allclose(ana[t, 2], fd_z, rtol=2e-3, atol=fd_atol, err_msg=f"contact force.grad mismatch at t={t}") + assert_allclose(ana[t, 2], fd_z, rtol=fd_rtol, atol=fd_atol, err_msg=f"contact force.grad mismatch at t={t}") @pytest.mark.required -def test_rigid_contact_no_tunneling_forward(show_viewer): - # Differentiable contact detection must route convex-convex pairs through the monolithic diff_gjk path; the split - # narrowphase used to skip GJK under requires_grad, so stacked boxes fell through each other. Observable: each top - # box stays on its support and comes to rest instead of tunneling. +def test_contact_no_tunneling_forward(show_viewer): + # Differentiable contact detection must route convex-convex pairs through the monolithic diff_gjk path so + # stacked boxes keep colliding under requires_grad. Observable: each top box stays on its support and comes to + # rest instead of tunneling. scene = gs.Scene( sim_options=gs.options.SimOptions( requires_grad=True, ), - rigid_options=gs.options.RigidOptions( - integrator=gs.integrator.approximate_implicitfast, - box_box_detection=False, + viewer_options=gs.options.ViewerOptions( + camera_pos=(2.0, -2.0, 1.2), + camera_lookat=(0.0, 0.0, 0.4), ), show_viewer=show_viewer, ) scene.add_entity(gs.morphs.Plane()) tops = [] for x in (0.8, -0.8): - scene.add_entity(gs.morphs.Box(size=(0.6, 0.6, 0.4), pos=(x, 0.0, 0.2), fixed=True)) - tops.append(scene.add_entity(gs.morphs.Box(size=(0.4, 0.4, 0.4), pos=(x, 0.0, 0.6)))) + scene.add_entity( + gs.morphs.Box( + size=(0.6, 0.6, 0.4), + pos=(x, 0.0, 0.2), + fixed=True, + ), + vis_mode="collision", + ) + tops.append( + scene.add_entity( + gs.morphs.Box( + size=(0.4, 0.4, 0.4), + pos=(x, 0.0, 0.6), + ), + vis_mode="collision", + ) + ) scene.build() for _ in range(20): @@ -145,29 +184,81 @@ def test_rigid_contact_no_tunneling_forward(show_viewer): @pytest.mark.required -def test_rigid_diff_contact_pair_unsupported_raises(): +def test_diff_unsupported_configuration_raises(): scene = gs.Scene( sim_options=gs.options.SimOptions( requires_grad=True, ), show_viewer=False, ) - scene.add_entity(gs.morphs.Sphere(radius=0.2, pos=(0.0, 0.0, 0.2), fixed=True)) - scene.add_entity(gs.morphs.Sphere(radius=0.2, pos=(0.0, 0.0, 0.5))) + scene.add_entity( + gs.morphs.Sphere( + radius=0.2, + pos=(0.0, 0.0, 0.2), + fixed=True, + ), + ) + scene.add_entity( + gs.morphs.Sphere( + radius=0.2, + pos=(0.0, 0.0, 0.5), + ), + ) # A sphere/sphere pair has an everywhere-curved Minkowski boundary on which diff_gjk's EPA never converges, so # it would silently tunnel; the build must reject it instead. with pytest.raises(gs.GenesisException): scene.build() + # Torsional / rolling friction rows have no manual reverse; the build must reject them + scene = gs.Scene( + sim_options=gs.options.SimOptions( + requires_grad=True, + ), + rigid_options=gs.options.RigidOptions( + enable_torsional_friction=True, + ), + show_viewer=False, + ) + scene.add_entity( + gs.morphs.Box( + size=(0.1, 0.1, 0.1), + pos=(0.0, 0.0, 0.5), + ), + ) + with pytest.raises(gs.GenesisException): + scene.build() + + # Attach merges kinematic trees across entities, which the per-entity backward kernels cannot reverse + scene = gs.Scene( + sim_options=gs.options.SimOptions( + requires_grad=True, + ), + show_viewer=False, + ) + base = scene.add_entity( + gs.morphs.Box( + size=(0.1, 0.1, 0.1), + pos=(0.0, 0.0, 0.5), + fixed=True, + ), + ) + child = scene.add_entity( + gs.morphs.Box( + size=(0.1, 0.1, 0.1), + pos=(0.0, 0.0, 0.8), + ), + ) + with pytest.raises(gs.GenesisException): + child.attach(base) + @pytest.mark.required @pytest.mark.precision("64") @pytest.mark.debug(False) -def test_rigid_contact_detection_jacobian_matches_fd(): +def test_contact_detection_jacobian_matches_fd(): scene = gs.Scene( sim_options=gs.options.SimOptions( - dt=0.01, requires_grad=True, ), show_viewer=False, @@ -175,9 +266,17 @@ def test_rigid_contact_detection_jacobian_matches_fd(): box_size = 0.25 vec_one = np.array([1.0, 1.0, 1.0]) box_pos_offset = (0.0, 0.0, 0.0) + 0.5 * box_size * vec_one - box0 = scene.add_entity(gs.morphs.Box(size=box_size * vec_one, pos=box_pos_offset)) + box0 = scene.add_entity( + gs.morphs.Box( + size=box_size * vec_one, + pos=box_pos_offset, + ), + ) box1 = scene.add_entity( - gs.morphs.Box(size=box_size * vec_one, pos=box_pos_offset + 0.8 * box_size * np.array([0, 0, 1])) + gs.morphs.Box( + size=box_size * vec_one, + pos=box_pos_offset + 0.8 * box_size * np.array([0, 0, 1]), + ), ) scene.build() collider = scene.sim.rigid_solver.collider @@ -237,16 +336,14 @@ def directional_error(dL_dx, x_type): @pytest.mark.required @pytest.mark.precision("64") @pytest.mark.debug(False) -def test_rigid_constraint_solver_backward_matches_fd(monkeypatch): - # fp64 is required: the FD perturbation must be small enough for a reliable estimate, which fp32 cannot resolve. - # These internal solver symbols are imported locally to keep a mismatch with the installed engine from breaking - # collection of the whole module. - from genesis.engine.solvers.rigid.constraint.solver import func_solve_init, func_solve_body +def test_constraint_solver_backward_matches_fd(monkeypatch): + # Engine modules resolve gs dtypes at import time, which requires gs.init: import after initialization + from genesis.engine.solvers.rigid.constraint.solver import func_solve_body, func_solve_init from genesis.engine.solvers.rigid.rigid_solver import kernel_step_1 + # fp64 is required: the FD perturbation must be small enough for a reliable estimate, which fp32 cannot resolve scene = gs.Scene( sim_options=gs.options.SimOptions( - dt=0.01, requires_grad=True, ), rigid_options=gs.options.RigidOptions( @@ -254,9 +351,18 @@ def test_rigid_constraint_solver_backward_matches_fd(monkeypatch): ), show_viewer=False, ) - scene.add_entity(gs.morphs.Plane(pos=(0, 0, 0))) - scene.add_entity(gs.morphs.Box(size=(1, 1, 1), pos=(10, 10, 0.49))) - franka = scene.add_entity(gs.morphs.MJCF(file="xml/franka_emika_panda/panda.xml")) + scene.add_entity(gs.morphs.Plane()) + scene.add_entity( + gs.morphs.Box( + size=(1, 1, 1), + pos=(10, 10, 0.49), + ), + ) + franka = scene.add_entity( + gs.morphs.MJCF( + file="xml/franka_emika_panda/panda.xml", + ), + ) scene.build() rigid_solver = scene._sim.rigid_solver constraint_solver = rigid_solver.constraint_solver diff --git a/tests/grad/test_rigid_constraints.py b/tests/grad/test_rigid_constraints.py index aa9f1cf63..da32ebe9b 100644 --- a/tests/grad/test_rigid_constraints.py +++ b/tests/grad/test_rigid_constraints.py @@ -1,5 +1,3 @@ -# FD-vs-analytical reverse-mode gradient checks through the rigid constraint solver: joint limits, dof frictionloss, -# and the three equality types (joint / connect / weld), plus an all-groups integration scene. import math import numpy as np @@ -14,7 +12,7 @@ @pytest.mark.required @pytest.mark.debug(False) -def test_rigid_joint_limit_grad_matches_fd(grad_slider_limit, precision, show_viewer): +def test_joint_limit_grad_matches_fd(grad_slider_limit, precision, show_viewer): # Forward: the slider limit must actually bound the cart (it drifts freely when the constraint is off). off = make_diff_scene_pair( grad_slider_limit, @@ -23,6 +21,7 @@ def test_rigid_joint_limit_grad_matches_fd(grad_slider_limit, precision, show_vi gravity=(0.0, 0.0, 0.0), enable_joint_limit=False, disable_constraint=True, + modes=(False,), ) off.scene_fd.reset() off.entity_fd.set_dofs_velocity(gs.tensor([100.0], dtype=gs.tc_float)) @@ -86,6 +85,7 @@ def test_rigid_joint_limit_grad_matches_fd(grad_slider_limit, precision, show_vi gravity=(0.0, 0.0, 0.0), enable_joint_limit=False, disable_constraint=False, + modes=(True,), ) grads = {} for pair, key in ((off_solver, "off"), (on, "on")): @@ -107,18 +107,15 @@ def test_rigid_joint_limit_grad_matches_fd(grad_slider_limit, precision, show_vi @pytest.mark.required @pytest.mark.parametrize("model_name", ["grad_slider_limit", "grad_cartpole", "grad_hopper"]) @pytest.mark.debug(False) -def test_rigid_per_step_force_grad_matches_fd(model_name, request, precision, show_viewer): +def test_per_step_force_into_limit_grad_matches_fd(model_name, request, precision, show_viewer): # Per-step control-force adjoint driving a joint into its limit, across three topologies. A constant force over # the horizon pushes the tracked dof into the active band; the setup-sanity assert guards against a vacuous run. - hopper_force = np.zeros(6) - hopper_force[5] = 200.0 # (gravity, n_steps, per-step force, loss reads links_pos, sanity dof, sanity threshold, initial dof pose, - # fp32 tolerance). The fp32 floor depends on the topology: the slider limit is the noisiest, the hopper the - # cleanest; fp64 clears 5e-5 for all three (the cartpole limit kink sets that floor). + # fp32 tolerance). gravity, n_steps, per_step_force, is_links_loss, sanity_dof, sanity_thresh, init_pos, fp32_tol = { - "grad_slider_limit": ((0.0, 0.0, 0.0), 10, [500.0], False, 0, 3.5, None, 1e-3), - "grad_cartpole": ((0.0, 0.0, -9.81), 15, [2000.0, 0.0], False, 0, 3.5, [0.0, -math.pi], 5e-4), - "grad_hopper": ((0.0, 0.0, 0.0), 10, hopper_force, True, 5, 0.7, None, 2e-4), + "grad_slider_limit": ((0.0, 0.0, 0.0), 10, [500.0], False, 0, 3.5, None, 1e-4), + "grad_cartpole": ((0.0, 0.0, -9.81), 15, [2000.0, 0.0], False, 0, 3.5, [0.0, -math.pi], 2e-4), + "grad_hopper": ((0.0, 0.0, 0.0), 10, [0.0, 0.0, 0.0, 0.0, 0.0, 200.0], True, 5, 0.7, None, 5e-5), }[model_name] pair = make_diff_scene_pair( @@ -154,14 +151,14 @@ def loss_fn(scene, entity): lambda e, x: e.control_dofs_force(x), loss_fn, setup_fn=setup_fn, - rtol=5e-5 if precision == "64" else fp32_tol, - atol=5e-5 if precision == "64" else fp32_tol, + rtol=1e-10 if precision == "64" else fp32_tol, + atol=1e-10 if precision == "64" else fp32_tol, eps=3e-2, ) @pytest.mark.required -def test_rigid_frictionloss_grad_matches_fd(grad_revolute_frictionloss, precision, show_viewer): +def test_frictionloss_grad_matches_fd(grad_revolute_frictionloss, precision, show_viewer): pair = make_diff_scene_pair( grad_revolute_frictionloss, substeps=4, @@ -198,7 +195,7 @@ def test_rigid_frictionloss_grad_matches_fd(grad_revolute_frictionloss, precisio ("grad_weld_pair", 6), ], ) -def test_rigid_equality_grad_matches_fd(model_name, n_rows, request, precision, show_viewer): +def test_equality_grad_matches_fd(model_name, n_rows, request, precision, show_viewer): pair = make_diff_scene_pair( request.getfixturevalue(model_name), substeps=4, @@ -213,23 +210,24 @@ def test_rigid_equality_grad_matches_fd(model_name, n_rows, request, precision, cs = pair.scene_fd.rigid_solver.constraint_solver.constraint_state assert qd_to_torch(cs.n_constraints_equality)[0] == n_rows + # Large initial velocities: the anchor velocity-product bias entering aref is quadratic in velocity, so its + # adjoint contribution only clears the fp64 tolerance band when the joints spin fast. assert_grad_matches_fd( pair, - [np.array([0.8, -0.3])], + [np.array([4.0, -2.5])], lambda e, x: e.set_dofs_velocity(x), lambda scene, entity: ( scene.rigid_solver.get_state().qpos[0, 0] ** 2 + 0.7 * scene.rigid_solver.get_state().qpos[0, 1] ** 2 ), n_steps=10, - rtol=2e-9 if precision == "64" else 5e-5, - atol=2e-9 if precision == "64" else 5e-5, + rtol=1e-10 if precision == "64" else 5e-5, + atol=1e-10 if precision == "64" else 5e-5, eps=1e-3 if precision == "64" else 3e-2, ) @pytest.mark.required -@pytest.mark.parametrize("backend", [gs.cpu, gs.gpu]) -def test_rigid_all_constraints_grad_matches_fd(grad_all_eq_fric, precision, show_viewer): +def test_all_constraint_groups_grad_matches_fd(grad_all_eq_fric, precision, show_viewer): # Integration scene: frictionloss + equality joint + connect + weld on disjoint link pairs. Guards row-offset # bookkeeping across every differentiated constraint group at once; per-group formulas are pinned elsewhere. pair = make_diff_scene_pair( @@ -259,7 +257,7 @@ def loss_fn(scene, entity): lambda e, x: e.set_dofs_velocity(x), loss_fn, n_steps=10, - rtol=5e-10 if precision == "64" else 5e-5, - atol=5e-10 if precision == "64" else 5e-5, + rtol=1e-10 if precision == "64" else 5e-5, + atol=1e-10 if precision == "64" else 5e-5, eps=3e-4 if precision == "64" else 3e-3, ) diff --git a/tests/grad/test_rigid_dynamics.py b/tests/grad/test_rigid_dynamics.py index 30e671c83..0013cc429 100644 --- a/tests/grad/test_rigid_dynamics.py +++ b/tests/grad/test_rigid_dynamics.py @@ -1,5 +1,3 @@ -# Finite-difference vs analytical reverse-mode gradient checks for rigid forward kinematics (constraints off), one -# packed test per joint topology exercising every tracked setter, and a multi-step control-force adjoint check. import sys import numpy as np @@ -29,16 +27,24 @@ "grad_prismatic", "grad_spherical", "grad_free_with_revolute", - "grad_revolute_chain3", + "grad_chain3", "grad_cartpole", "grad_hopper", ], ) -def test_rigid_fk_grad_matches_fd(model_name, request, precision, show_viewer): - pair = make_diff_scene_pair(request.getfixturevalue(model_name), n_envs=2, substeps=4, show_viewer=show_viewer) +def test_fk_grad_matches_fd(model_name, request, precision, show_viewer): + is_tall = model_name in ("grad_cartpole", "grad_hopper") + B = 2 + pair = make_diff_scene_pair( + request.getfixturevalue(model_name), + n_envs=B, + substeps=4, + show_viewer=show_viewer, + camera_pos=(2.5, -2.5, 1.8) if is_tall else (1.2, -1.2, 0.8), + camera_lookat=(0.0, 0.0, 0.9) if is_tall else (0.0, 0.0, 0.2), + ) n_dofs = pair.entity_ana.n_dofs n_links = pair.entity_ana.n_links - B = 2 # Single-link joints read the entity pose; multi-link topologies read the rigid-solver per-link pose. is_single_link = model_name in ("grad_free", "grad_revolute", "grad_prismatic", "grad_spherical") @@ -54,24 +60,23 @@ def test_rigid_fk_grad_matches_fd(model_name, request, precision, show_viewer): 61, 62, ), - "grad_revolute_chain3": ((("vel", "pos", 90), ("vel", "quat", 91)), 81, 82), + "grad_chain3": ((("vel", "pos", 90), ("vel", "quat", 91)), 81, 82), "grad_cartpole": ((("vel", "pos", 190), ("vel", "quat", 191), ("force", "pos", 192)), 181, 182), "grad_hopper": ((("vel", "pos", 210), ("vel", "quat", 211)), 201, 202), } checks, pos_seed, quat_seed = checks_by_joint[model_name] - # Per-topology fp32 finite-difference floor (tolerance, step): quaternion and multi-link chain topologies - # (free, free_with_revolute, hopper) are noisier at fp32 and need a smaller step, while the single-DOF cases - # are far cleaner and pin down to ~1e-5. fp64 clears 1e-9 for every topology, so it stays a single band. - fp32_tol, fp32_eps = { - "grad_free": (2e-4, 1e-2), - "grad_revolute": (5e-5, 3e-2), - "grad_prismatic": (5e-6, 3e-2), - "grad_spherical": (1e-4, 3e-2), - "grad_free_with_revolute": (5e-4, 1e-2), - "grad_revolute_chain3": (2e-4, 3e-2), - "grad_cartpole": (2e-4, 3e-2), - "grad_hopper": (5e-4, 3e-2), + # Per-topology finite-difference floors (fp64 tolerance, fp32 tolerance, fp32 step): free-joint topologies + # are the noisiest at both precisions and need a smaller fp32 step. + fp64_tol, fp32_tol, fp32_eps = { + "grad_free": (1e-9, 2e-4, 1e-2), + "grad_revolute": (1e-10, 5e-5, 3e-2), + "grad_prismatic": (1e-10, 5e-6, 3e-2), + "grad_spherical": (1e-10, 1e-4, 3e-2), + "grad_free_with_revolute": (1e-9, 2e-4, 1e-2), + "grad_chain3": (2e-10, 1e-4, 3e-2), + "grad_cartpole": (1e-10, 5e-5, 3e-2), + "grad_hopper": (5e-10, 2e-4, 3e-2), }[model_name] pos_shape = (B, 3) if is_single_link else (B, n_links, 3) @@ -111,8 +116,8 @@ def loss_fn(scene, entity, tgt=target, out=output, sl=is_single_link): [step_input], apply_fn, loss_fn, - rtol=1e-9 if precision == "64" else fp32_tol, - atol=1e-9 if precision == "64" else fp32_tol, + rtol=fp64_tol if precision == "64" else fp32_tol, + atol=fp64_tol if precision == "64" else fp32_tol, eps=3e-5 if precision == "64" else fp32_eps, ) @@ -125,44 +130,68 @@ def loss_fn(scene, entity, tgt=target, out=output, sl=is_single_link): "grad_revolute", "grad_prismatic", "grad_free_with_revolute", - "grad_revolute_chain3", + "grad_chain3", "grad_spherical", "grad_cartpole", "grad_hopper", ], ) @pytest.mark.debug(False) -def test_rigid_fk_multistep_force_grad_matches_fd(model_name, request, precision, show_viewer): +def test_fk_multistep_force_grad_matches_fd(model_name, request, precision, show_viewer): # Ten distinct per-step control forces, each of which must receive an independent adjoint across the unroll. - # (output kind: entity state vs rigid-solver links, per-link output shape, target seed, fp32 tolerance). The - # per-topology fp32 floor spans 2e-6 (prismatic) to 2e-4 (hopper), tracking how far the ten-step unroll - # amplifies fp32 noise; fp64 clears 5e-9 for all, set by the chaotic hopper/chain3 chains. + # (output kind: entity state pos / quat or rigid-solver links, per-link output shape, target seed, fp32 + # tolerance). Anchored single-joint topologies (revolute, spherical) read the quaternion: their base link + # position is pinned at the joint anchor, so a position loss would be constant and the check vacuous. output, output_shape, seed, fp32_tol = { - "grad_free": ("state", (3,), 161, 2e-5), - "grad_revolute": ("state", (3,), 162, 1e-5), - "grad_prismatic": ("state", (3,), 163, 2e-6), + "grad_free": ("state_pos", (3,), 161, 5e-6), + "grad_revolute": ("state_quat", (4,), 162, 5e-6), + "grad_prismatic": ("state_pos", (3,), 163, 2e-6), "grad_free_with_revolute": ("links", (2, 3), 164, 5e-5), - "grad_revolute_chain3": ("links", (3, 3), 165, 1e-4), - "grad_spherical": ("state", (3,), 166, 1e-5), + "grad_chain3": ("links", (3, 3), 165, 2e-5), + "grad_spherical": ("state_quat", (4,), 166, 1e-4), "grad_cartpole": ("links", (2, 3), 167, 2e-5), - "grad_hopper": ("links", (5, 3), 168, 2e-4), + "grad_hopper": ("links", (5, 3), 168, 1e-4), }[model_name] - pair = make_diff_scene_pair(request.getfixturevalue(model_name), n_envs=0, substeps=4, show_viewer=show_viewer) + is_tall = model_name in ("grad_cartpole", "grad_hopper") + pair = make_diff_scene_pair( + request.getfixturevalue(model_name), + n_envs=0, + substeps=4, + show_viewer=show_viewer, + camera_pos=(2.5, -2.5, 1.8) if is_tall else (1.2, -1.2, 0.8), + camera_lookat=(0.0, 0.0, 0.9) if is_tall else (0.0, 0.0, 0.2), + ) n_dofs = pair.entity_ana.n_dofs target = gs.tensor(np.random.RandomState(seed).standard_normal((1, *output_shape)), dtype=gs.tc_float).reshape(-1) inputs = [np.random.default_rng(seed * 100 + t).standard_normal((n_dofs,)) for t in range(10)] def loss_fn(scene, entity): - pose = entity.get_state().pos if output == "state" else scene.rigid_solver.get_state().links_pos + if output == "state_pos": + pose = entity.get_state().pos + elif output == "state_quat": + pose = entity.get_state().quat + else: + pose = scene.rigid_solver.get_state().links_pos return ((pose.reshape(-1) - target) ** 2).sum() + def apply_force(entity, force): + # Split the same-step control across two dof subsets (the standard arm + gripper pattern): each call must + # keep its own tape slot and gradient path. The second subset is passed as a slice, a valid index form the + # tape key must accept on every backend. + if n_dofs == 1: + entity.control_dofs_force(force) + else: + entity.control_dofs_force(force[..., :1], dofs_idx_local=[0]) + entity.control_dofs_force(force[..., 1:], dofs_idx_local=slice(1, n_dofs)) + # fp32 needs a large step to clear the state-noise floor; fp64 needs a small step to bound truncation error. + fp64_tol = 5e-10 if model_name == "grad_hopper" else 1e-10 assert_grad_matches_fd( pair, inputs, - lambda e, x: e.control_dofs_force(x), + apply_force, loss_fn, - rtol=5e-9 if precision == "64" else fp32_tol, - atol=5e-9 if precision == "64" else fp32_tol, + rtol=fp64_tol if precision == "64" else fp32_tol, + atol=fp64_tol if precision == "64" else fp32_tol, eps=3e-5 if precision == "64" else 3e-2, ) diff --git a/tests/grad/test_rigid_optim.py b/tests/grad/test_rigid_optim.py index 9ffd7f5f7..787d6be47 100644 --- a/tests/grad/test_rigid_optim.py +++ b/tests/grad/test_rigid_optim.py @@ -1,5 +1,3 @@ -# End-to-end reverse-mode optimization on the rigid solver: Adam recovers a cartpole reference trajectory (through -# either the initial velocity or a per-step force sequence) and drives a free box to a goal pose. import sys import numpy as np @@ -24,18 +22,19 @@ ) @pytest.mark.parametrize("control_target", ["init_vel", "control_force"]) @pytest.mark.debug(False) -def test_rigid_optim_cartpole_converges(control_target, grad_cartpole, show_viewer): +def test_reference_trajectory_recovery_converges(control_target, grad_cartpole, show_viewer): # Reproduce a reference cartpole trajectory by optimizing either the initial velocity or the per-step control # forces; every env must drive its per-env loss below both a relative-reduction and an absolute threshold. The # reference is exactly reproducible, so a correct gradient lets Adam crush the loss to its optimizer plateau # (fp32 and fp64 reach the same floor). init_vel (4 params) converges far below control_force (128 per-step # forces), hence the per-target thresholds - each pinned just above the measured converged loss. - N_STEPS, N_ITER, LR, N_DOFS, B = 32, 150, 1e-2, 2, 2 + N_STEPS, N_ITER, LR, B = 32, 150, 1e-2, 2 REL_REDUCTION, ABS_THRESHOLD = { "init_vel": (2e-6, 2e-7), "control_force": (1e-3, 5e-7), }[control_target] - pair = make_diff_scene_pair(grad_cartpole, n_envs=2, show_viewer=show_viewer) + pair = make_diff_scene_pair(grad_cartpole, n_envs=B, show_viewer=show_viewer) + N_DOFS = pair.entity_ana.n_dofs scene_ref, robot_ref = pair.scene_fd, pair.entity_fd scene_opt, robot_opt = pair.scene_ana, pair.entity_ana rng = np.random.default_rng(seed=11 if control_target == "init_vel" else 23) @@ -86,7 +85,7 @@ def test_rigid_optim_cartpole_converges(control_target, grad_cartpole, show_view loss_per_env.sum().backward() optimizer.step() - history = np.asarray(loss_history) + history = np.array(loss_history) initial, final = history[0], history[-1] rel_ratios = final / initial assert_allclose(rel_ratios, 0.0, atol=REL_REDUCTION, err_msg=f"loss reduction insufficient (initial={initial})") @@ -96,25 +95,19 @@ def test_rigid_optim_cartpole_converges(control_target, grad_cartpole, show_view @pytest.mark.slow @pytest.mark.required @pytest.mark.debug(False) -def test_rigid_optim_reach_goal_pose(show_viewer): +def test_goal_pose_optimization_converges(show_viewer): goal_pos = gs.tensor([0.7, 1.0, 0.05]) goal_quat = gs.tensor([0.3, 0.2, 0.1, 0.9]) goal_quat = goal_quat / torch.norm(goal_quat, dim=-1, keepdim=True) scene = gs.Scene( sim_options=gs.options.SimOptions( - dt=1e-2, - substeps=1, requires_grad=True, gravity=(0, 0, -1), ), rigid_options=gs.options.RigidOptions( enable_collision=False, - enable_self_collision=False, - enable_joint_limit=False, disable_constraint=True, - use_contact_island=False, - use_hibernation=False, ), viewer_options=gs.options.ViewerOptions( camera_pos=(2.5, -0.15, 2.42), diff --git a/tests/grad/utils.py b/tests/grad/utils.py index 9af357ab3..f645da0ef 100644 --- a/tests/grad/utils.py +++ b/tests/grad/utils.py @@ -24,32 +24,22 @@ def make_diff_scene_pair( substeps=1, dt=0.01, gravity=(0.0, 0.0, -9.81), - enable_collision=False, enable_joint_limit=False, disable_constraint=True, - box_box_detection=None, + modes=(True, False), show_viewer=False, camera_pos=(1.2, -1.2, 0.8), camera_lookat=(0.0, 0.0, 0.2), ): - # Builds a diff-mode scene (scene_ana, the only one backward() runs on) and a production-mode reference - # (scene_fd, the one finite differences perturb) from the same MJCF with identical config. The two kernels - # produce bit-identical forward states, so FD on scene_fd is a valid reference for scene_ana's analytical - # gradient. - rigid_kwargs = dict( - enable_collision=enable_collision, - enable_self_collision=False, - enable_joint_limit=enable_joint_limit, - disable_constraint=disable_constraint, - use_hibernation=False, - use_contact_island=False, - ) - if box_box_detection is not None: - rigid_kwargs["box_box_detection"] = box_box_detection + """Build a diff-mode scene (scene_ana, the only one backward() runs on) and a production-mode reference + (scene_fd, the one finite differences perturb) from the same MJCF with identical config. - scenes = [] - entities = [] - for requires_grad in (True, False): + The two kernels produce bit-identical forward states, so FD on scene_fd is a valid reference for scene_ana's + analytical gradient. `modes` selects which of the pair to build (True = diff-mode scene_ana, False = production + scene_fd); tests consuming only one side skip the other's kernel compilation.""" + scenes = {} + entities = {} + for requires_grad in modes: scene = gs.Scene( sim_options=gs.options.SimOptions( dt=dt, @@ -57,29 +47,38 @@ def make_diff_scene_pair( gravity=gravity, requires_grad=requires_grad, ), - rigid_options=gs.options.RigidOptions(**rigid_kwargs), + rigid_options=gs.options.RigidOptions( + enable_collision=False, + enable_joint_limit=enable_joint_limit, + disable_constraint=disable_constraint, + ), viewer_options=gs.options.ViewerOptions( camera_pos=camera_pos, camera_lookat=camera_lookat, ), show_viewer=show_viewer and requires_grad, ) - entity = scene.add_entity(gs.morphs.MJCF(file=mjcf)) + entity = scene.add_entity( + gs.morphs.MJCF( + file=mjcf, + ), + ) scene.build(n_envs=n_envs) - scenes.append(scene) - entities.append(entity) - return DiffScenePair(scenes[0], entities[0], scenes[1], entities[1]) + scenes[requires_grad] = scene + entities[requires_grad] = entity + return DiffScenePair(scenes.get(True), entities.get(True), scenes.get(False), entities.get(False)) def assert_grad_matches_fd(pair, inputs, apply_fn, loss_fn, *, rtol, atol, eps, n_steps=None, setup_fn=None): - # Central finite-difference check of a tracked setter's reverse-mode gradient. `inputs` holds one array per - # applied input; input i is applied via `apply_fn(entity, x)` before step i and must receive an independent - # adjoint from the backward unroll. The scene runs `n_steps` steps (default len(inputs)); when it exceeds the - # number of inputs the remaining steps run without re-applying (a single input driving an N-step rollout). - # `setup_fn(scene, entity)` runs once after reset for untracked initialization (e.g. an initial pose). The FD - # reference perturbs each entry of each input in turn and re-runs the full trajectory on the production-mode - # scene, so the cost is O(n_steps * total input size). rtol / atol / eps are required and per-scenario: each - # test pins them to its own measured finite-difference floor. + """Central finite-difference check of a tracked setter's reverse-mode gradient. + + `inputs` holds one array per applied input; input i is applied via `apply_fn(entity, x)` before step i and must + receive an independent adjoint from the backward unroll. The scene runs `n_steps` steps (default len(inputs)); + when it exceeds the number of inputs the remaining steps run without re-applying (a single input driving an + N-step rollout). `setup_fn(scene, entity)` runs once after reset for untracked initialization (e.g. an initial + pose). The FD reference perturbs each entry of each input in turn and re-runs the full trajectory on the + production-mode scene, so the cost is O(n_steps * total input size). rtol / atol / eps are required and + per-scenario: each test pins them to its own measured finite-difference floor.""" base = [np.array(inp, dtype=np.float64) for inp in inputs] total_steps = len(base) if n_steps is None else n_steps @@ -122,8 +121,8 @@ def assert_grad_matches_fd(pair, inputs, apply_fn, loss_fn, *, rtol, atol, eps, perturbed.append(float(loss_fn(pair.scene_fd, pair.entity_fd))) fd_grad.reshape(-1)[i_entry] = (perturbed[0] - perturbed[1]) / (2.0 * eps) assert_allclose( - torch.from_numpy(ana_grads[i_input]), - torch.from_numpy(fd_grad), + ana_grads[i_input], + fd_grad, rtol=rtol, atol=atol, err_msg=f"input {i_input}: FD vs analytical mismatch",