diff --git a/genesis/engine/entities/rigid_entity/rigid_entity.py b/genesis/engine/entities/rigid_entity/rigid_entity.py index 3a83040c86..c62d9c1079 100644 --- a/genesis/engine/entities/rigid_entity/rigid_entity.py +++ b/genesis/engine/entities/rigid_entity/rigid_entity.py @@ -142,7 +142,7 @@ def __init__( self._load_model() # Initialize target variables and checkpoint - self._tgt_keys = ("pos", "quat", "qpos", "dofs_velocity") + self._tgt_keys = ("pos", "quat", "qpos", "dofs_velocity", "control_dofs_force") self._tgt = dict() self._tgt_buffer = list() self._ckpt = dict() @@ -1536,6 +1536,8 @@ def process_input(self, in_backward=False): self.set_quat(**data_kwargs) case "set_dofs_velocity": self.set_dofs_velocity(**data_kwargs) + case "control_dofs_force": + self.control_dofs_force(**data_kwargs) case _: gs.raise_exception(f"Invalid target key: {key} not in {self._tgt_keys}") @@ -1566,6 +1568,14 @@ def process_input_grad(self): 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: + 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}") @@ -3879,6 +3889,11 @@ def set_dofs_velocity_grad(self, dofs_idx_local, envs_idx, velocity_grad): dofs_idx = self._get_global_idx(dofs_idx_local, self.n_dofs, self._dof_start, unsafe=True) self._solver.set_dofs_velocity_grad(dofs_idx, envs_idx, velocity_grad.data) + @gs.assert_built + def set_dofs_force_grad(self, dofs_idx_local, envs_idx, force_grad): + dofs_idx = self._get_global_idx(dofs_idx_local, self.n_dofs, self._dof_start, unsafe=True) + self._solver.set_dofs_force_grad(dofs_idx, envs_idx, force_grad.data) + # ------------------------------------------------------------------------------------ # ----------------------------- DOF property setters --------------------------------- # ------------------------------------------------------------------------------------ @@ -3912,6 +3927,7 @@ def set_dofs_position(self, position, dofs_idx_local=None, envs_idx=None, *, zer # ------------------------------------------------------------------------------------ @gs.assert_built + @tracked def control_dofs_force(self, force, dofs_idx_local=None, envs_idx=None): """ Control the entity's dofs' motor force. This is used for force/torque control. diff --git a/genesis/engine/scene.py b/genesis/engine/scene.py index 23f50d0cae..0a46c551d3 100644 --- a/genesis/engine/scene.py +++ b/genesis/engine/scene.py @@ -995,13 +995,16 @@ def reset(self, state: SimState | None = None, envs_idx=None): self._reset(state, envs_idx=envs_idx) self._recorder_manager.reset(envs_idx) - def _reset(self, state: SimState | None = None, *, envs_idx=None): + def _reset(self, state: SimState | None = None, *, envs_idx=None, keep_init: bool = False): if self._is_built: if state is None: state = self._init_state else: assert isinstance(state, SimState), "state must be a SimState object" - self._init_state = state + # keep_init=True restores the state while leaving the registered init untouched, so a later bare + # reset() still rewinds to the true initial state. + if not keep_init: + self._init_state = state self._sim.reset(state, envs_idx) else: self._init_state = self._get_state() @@ -1022,6 +1025,40 @@ def _reset(self, state: SimState | None = None, *, envs_idx=None): def _reset_grad(self): self._backward_ready = True + @gs.assert_built + def backward(self, loss: torch.Tensor, *args, **kwargs): + """ + Differentiates `loss` through the recorded rollout and restores the pre-backward physics state. + + Unrolling the gradient tape rewinds the physics state to step 0, so this method snapshots the current state + first, runs the backward pass, and restores the snapshot afterwards. The scene then sits at the same physics + state as before the call, with gradients populated and forward / backward re-armed, ready to continue the + rollout or to be reset. The registered initial state (`reset()` with no argument) is preserved. + + Parameters + ---------- + loss : torch.Tensor + Scalar loss to differentiate. Extra positional and keyword arguments (e.g. `gradient`, `retain_graph`) + are forwarded to `torch.autograd.backward`. + + Returns + ------- + snapshot : SimState + The physics state the scene was restored to. + """ + # Snapshot the current state before the gradient-tape unroll rewinds physics to step 0. + 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) + # 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. + self._reset(snapshot, keep_init=True) + return snapshot + def _get_state(self): return self._sim.get_state() diff --git a/genesis/engine/solvers/kinematic_solver.py b/genesis/engine/solvers/kinematic_solver.py index 9fa422cb1c..c66809d5ef 100644 --- a/genesis/engine/solvers/kinematic_solver.py +++ b/genesis/engine/solvers/kinematic_solver.py @@ -46,6 +46,7 @@ kernel_set_dofs_position, kernel_set_dofs_velocity, kernel_set_dofs_velocity_grad, + kernel_set_dofs_force_grad, kernel_set_dofs_zero_velocity, kernel_set_links_pos, kernel_set_links_quat, @@ -1012,6 +1013,14 @@ def set_dofs_velocity_grad(self, dofs_idx, envs_idx, velocity_grad): velocity_grad_ = velocity_grad_.unsqueeze(0) kernel_set_dofs_velocity_grad(dofs_idx, envs_idx, velocity_grad_, self.dyn_state, self.rigid_config) + def set_dofs_force_grad(self, dofs_idx, envs_idx, force_grad): + force_grad_, dofs_idx, envs_idx = self._sanitize_io_variables( + force_grad, dofs_idx, self.n_dofs, "dofs_idx", envs_idx, skip_allocation=True + ) + if self.n_envs == 0: + force_grad_ = force_grad_.unsqueeze(0) + kernel_set_dofs_force_grad(dofs_idx, envs_idx, force_grad_, self.dyn_state, self.rigid_config) + @mutates(StateChange.GEOMETRY, links=MutatedLinks.ARTICULATED) def set_dofs_position(self, position, dofs_idx=None, envs_idx=None): position, dofs_idx, envs_idx = self._sanitize_io_variables( diff --git a/genesis/engine/solvers/rigid/abd/accessor.py b/genesis/engine/solvers/rigid/abd/accessor.py index 20a65766e0..6b14010e91 100644 --- a/genesis/engine/solvers/rigid/abd/accessor.py +++ b/genesis/engine/solvers/rigid/abd/accessor.py @@ -819,6 +819,20 @@ def kernel_set_dofs_velocity_grad( dyn_state.dofs.vel.grad[dofs_idx[i_d_], envs_idx[i_b_]] = 0.0 +@qd.kernel(fastcache=True) +def kernel_set_dofs_force_grad( + dofs_idx: qd.types.ndarray(), + envs_idx: qd.types.ndarray(), + force_grad: qd.types.ndarray(), + dyn_state: array_class.DynState, + rigid_config: qd.template(), +): + qd.loop_config(serialize=rigid_config.para_level < gs.PARA_LEVEL.ALL) + for i_d_, i_b_ in qd.ndrange(dofs_idx.shape[0], envs_idx.shape[0]): + force_grad[i_b_, i_d_] = dyn_state.dofs.ctrl_force.grad[dofs_idx[i_d_], envs_idx[i_b_]] + dyn_state.dofs.ctrl_force.grad[dofs_idx[i_d_], envs_idx[i_b_]] = 0.0 + + @qd.kernel(fastcache=True) def kernel_set_dofs_zero_velocity( dofs_idx: qd.types.ndarray(), diff --git a/genesis/engine/solvers/rigid/abd/diff.py b/genesis/engine/solvers/rigid/abd/diff.py index 6e096498d2..edf3a98779 100644 --- a/genesis/engine/solvers/rigid/abd/diff.py +++ b/genesis/engine/solvers/rigid/abd/diff.py @@ -19,7 +19,7 @@ import genesis as gs import genesis.utils.geom as gu import genesis.utils.array_class as array_class -from .forward_kinematics import func_update_cartesian_space +from .forward_kinematics import func_update_cartesian_space, func_forward_velocity @qd.func @@ -150,6 +150,7 @@ def kernel_prepare_backward_substep( func_update_cartesian_space( dyn_state, dyn_info, rigid_info, rigid_config, force_update_fixed_geoms=False, is_backward=True ) + func_forward_velocity(dyn_state, dyn_info, rigid_info, rigid_config, is_backward=True) # FIXME: Parameter pruning for ndarray is buggy for now and requires match variable and arg names. # Save results of [update_cartesian_space] to adjoint cache @@ -171,9 +172,9 @@ def kernel_begin_backward_substep( func_copy_next_to_curr_grad(f, dyn_state, rigid_adjoint_cache, rigid_info, rigid_config) if not rigid_config.enable_mujoco_compatibility: - # FIXME: Parameter pruning for ndarray is buggy for now and requires match variable and arg names. - # Save results of [update_cartesian_space] to adjoint cache - func_copy_cartesian_space(dyn_state, dyn_state_adjoint_cache, rigid_config) + # Restore the cartesian space that was overwritten by the post-integrate forward replay in the backward + # substep (see _update_cartesian_grad in rigid_solver.py). + func_copy_cartesian_space(dyn_state_adjoint_cache, dyn_state, rigid_config) return is_grad_valid @@ -264,6 +265,27 @@ def kernel_copy_acc( dyn_state.dofs.acc[i_d, i_b] = rigid_adjoint_cache.dofs_acc[f, i_d, i_b] +@qd.kernel(fastcache=True) +def kernel_copy_next_to_curr_no_check( + dyn_state: array_class.DynState, + 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. + n_qs = rigid_info.qpos.shape[0] + n_dofs = dyn_state.dofs.vel.shape[0] + _B = dyn_state.dofs.vel.shape[1] + + qd.loop_config(serialize=rigid_config.para_level < gs.PARA_LEVEL.ALL) + for i_q, i_b in qd.ndrange(n_qs, _B): + rigid_info.qpos[i_q, i_b] = rigid_info.qpos_next[i_q, i_b] + + qd.loop_config(serialize=rigid_config.para_level < gs.PARA_LEVEL.ALL) + for i_d, i_b in qd.ndrange(n_dofs, _B): + dyn_state.dofs.vel[i_d, i_b] = dyn_state.dofs.vel_next[i_d, i_b] + + @qd.func def func_integrate_dq_entity( i_e, diff --git a/genesis/engine/solvers/rigid/abd/forward_dynamics.py b/genesis/engine/solvers/rigid/abd/forward_dynamics.py index e5f1d01c0c..d187b87348 100644 --- a/genesis/engine/solvers/rigid/abd/forward_dynamics.py +++ b/genesis/engine/solvers/rigid/abd/forward_dynamics.py @@ -27,10 +27,7 @@ def update_qacc_from_qvel_delta( dyn_state: array_class.DynState, rigid_info: array_class.RigidInfo, rigid_config: qd.template(), - is_backward: qd.template(), ): - BW = qd.static(is_backward) - n_dofs = dyn_state.dofs.ctrl_mode.shape[0] _B = dyn_state.dofs.ctrl_mode.shape[1] @@ -52,10 +49,7 @@ def update_qvel( dyn_state: array_class.DynState, rigid_info: array_class.RigidInfo, rigid_config: qd.template(), - is_backward: qd.template(), ): - BW = qd.static(is_backward) - _B = dyn_state.dofs.vel.shape[1] n_dofs = dyn_state.dofs.vel.shape[0] @@ -82,7 +76,7 @@ def kernel_compute_mass_matrix( ): func_compute_mass_matrix(dyn_state, dyn_info, rigid_info, rigid_config, implicit_damping=False, is_backward=False) if decompose: - func_factor_mass(dyn_state, dyn_info, rigid_info, rigid_config, implicit_damping=False, is_backward=False) + func_factor_mass(dyn_state, dyn_info, rigid_info, rigid_config, implicit_damping=False) # @@@@@@@@@ Composer starts here @@ -104,7 +98,7 @@ def func_forward_dynamics( qd.static(rigid_config.integrator == gs.integrator.approximate_implicitfast), is_backward, ) - func_factor_mass(dyn_state, dyn_info, rigid_info, rigid_config, implicit_damping=False, is_backward=is_backward) + func_factor_mass(dyn_state, dyn_info, rigid_info, rigid_config, implicit_damping=False) func_torque_and_passive_force(dyn_state, constraint_state, dyn_info, rigid_info, rigid_config, is_backward) 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) @@ -359,7 +353,6 @@ def func_factor_mass_tiled( scattered into mass_mat_L / mass_mat_D_inv. To avoid a dedicated allocation, that scratch aliases the constraint Hessian buffer nt_H (same shape, and free at mass-factor time since the constraint solve only populates it later in the step); see get_constraint_state. The scratch and mass_mat_L are distinct buffers, so the scatter is race-free. - Backward keeps its own branch in func_factor_mass. """ # Reuse the Hessian's tile width; TileCls is dispatched to match it at the call site, so T and the tile class stay # consistent for either value. In practice this path only runs for mass blocks exceeding shared memory (total @@ -491,333 +484,256 @@ def func_factor_mass( rigid_info: array_class.RigidInfo, rigid_config: qd.template(), implicit_damping: qd.template(), - is_backward: qd.template(), ): - BW = qd.static(is_backward) + n_entities = dyn_info.entities.n_links.shape[0] + _B = dyn_state.dofs.ctrl_mode.shape[1] - if qd.static(not BW): - n_entities = dyn_info.entities.n_links.shape[0] - _B = dyn_state.dofs.ctrl_mode.shape[1] - - if qd.static(rigid_config.enable_register_tiled_mass): - # Register-streaming tiled per-entity factor for the >shared-cap path (same primitive as the constraint - # Hessian). Implies enable_tiled_cholesky_mass_matrix and not mass_matrix_fits_shared; see - # func_factor_mass_tiled. Replaces the cooperative LDL^T in the elif below. - func_factor_mass_tiled( - dyn_state, - dyn_info, - rigid_info, - rigid_config, - implicit_damping, - qd.simt.Tile32x32 if qd.static(rigid_config.cholesky_tile_size == 32) else qd.simt.Tile16x16, - ) - elif qd.static(rigid_config.enable_tiled_cholesky_mass_matrix and not rigid_config.mass_matrix_fits_shared): - # Uncapped cooperative per-entity LDL^T (entity submatrix does not fit shared memory): factors the entity - # mass submatrix in-place in global memory (mass_mat_L) over a block of BLOCK_DIM threads. Each elimination - # step snapshots the pivot row into a small shared vector (O(n_dofs), not O(n_dofs^2)) before updating the - # trailing submatrix, so the parallel per-row updates only READ the pivot row (from shared) -- race-free - # regardless of scheduling. Numerically identical to the scalar branch below; only parallelization differs. - BLOCK_DIM = qd.static(32) - MAX_DOFS_PER_BLOCK = qd.static(rigid_config.tiled_n_dofs_per_block) - - qd.loop_config(name="factor_mass", block_dim=BLOCK_DIM) - for i in range(n_entities * _B * BLOCK_DIM): - tid = i % BLOCK_DIM - i_e = (i // BLOCK_DIM) % n_entities - i_b = i // (BLOCK_DIM * n_entities) - if i_b >= _B: + if qd.static(rigid_config.enable_register_tiled_mass): + # Register-streaming tiled per-entity factor for the >shared-cap path (same primitive as the constraint + # Hessian). Implies enable_tiled_cholesky_mass_matrix and not mass_matrix_fits_shared; see + # func_factor_mass_tiled. Replaces the cooperative LDL^T in the elif below. + func_factor_mass_tiled( + dyn_state, + dyn_info, + rigid_info, + rigid_config, + implicit_damping, + qd.simt.Tile32x32 if qd.static(rigid_config.cholesky_tile_size == 32) else qd.simt.Tile16x16, + ) + elif qd.static(rigid_config.enable_tiled_cholesky_mass_matrix and not rigid_config.mass_matrix_fits_shared): + # Uncapped cooperative per-entity LDL^T (entity submatrix does not fit shared memory): factors the entity + # mass submatrix in-place in global memory (mass_mat_L) over a block of BLOCK_DIM threads. Each elimination + # step snapshots the pivot row into a small shared vector (O(n_dofs), not O(n_dofs^2)) before updating the + # trailing submatrix, so the parallel per-row updates only READ the pivot row (from shared) -- race-free + # regardless of scheduling. Numerically identical to the scalar branch below; only parallelization differs. + BLOCK_DIM = qd.static(32) + MAX_DOFS_PER_BLOCK = qd.static(rigid_config.tiled_n_dofs_per_block) + + qd.loop_config(name="factor_mass", block_dim=BLOCK_DIM) + for i in range(n_entities * _B * BLOCK_DIM): + tid = i % BLOCK_DIM + i_e = (i // BLOCK_DIM) % n_entities + i_b = i // (BLOCK_DIM * n_entities) + if i_b >= _B: + continue + # Skip hibernated entities: their mass matrix is unchanged, so the factor from the last awake step + # stays valid. The slot remaps to an awake entity, so the work scales with the awake entity count. + if qd.static(rigid_config.use_hibernation): + if i_e >= rigid_info.n_awake_entities[i_b]: continue - # Skip hibernated entities: their mass matrix is unchanged, so the factor from the last awake step - # stays valid. The slot remaps to an awake entity, so the work scales with the awake entity count. - if qd.static(rigid_config.use_hibernation): - if i_e >= rigid_info.n_awake_entities[i_b]: - continue - i_e = rigid_info.awake_entities[i_e, i_b] - - if rigid_info.mass_mat_mask[i_e, i_b]: - entity_dof_start = dyn_info.entities.dof_start[i_e] - entity_dof_end = dyn_info.entities.dof_end[i_e] - - pivot_row = qd.simt.block.SharedArray((MAX_DOFS_PER_BLOCK,), gs.qd_float) - - # Factor each mass block rooted in this entity in-place in global memory over its full range, - # block-relative so shared indices stay >= 0; a merged child owns no root and factors nothing. - block_start = entity_dof_start - while block_start < entity_dof_end: - block_end = rigid_info.dofs_mass_block_end[block_start] - if rigid_info.dofs_mass_block_start[block_start] == block_start: - n_block_dofs = block_end - block_start - - # Copy the block's lower triangle into mass_mat_L (+ implicit damping on the diagonal), - # cooperatively. Restricting to the block makes the factorization cost the sum of per-block - # cubes instead of the whole (possibly multi-block) entity cube. - i_d_ = tid - while i_d_ < n_block_dofs: - i_d = block_start + i_d_ - for j_d in range(block_start, i_d + 1): - rigid_info.mass_mat_L[i_d, j_d, i_b] = rigid_info.mass_mat[i_d, j_d, i_b] - if qd.static(implicit_damping): - I_d = [i_d, i_b] if qd.static(rigid_config.batch_dofs_info) else i_d - rigid_info.mass_mat_L[i_d, i_d, i_b] = ( - rigid_info.mass_mat_L[i_d, i_d, i_b] - + dyn_info.dofs.damping[I_d] * rigid_info.substep_dt[None] - ) - if qd.static(rigid_config.integrator == gs.integrator.implicitfast): - if dyn_state.dofs.ctrl_mode[i_d, i_b] <= gs.CTRL_MODE.VELOCITY: - rigid_info.mass_mat_L[i_d, i_d, i_b] = ( - rigid_info.mass_mat_L[i_d, i_d, i_b] - - dyn_info.dofs.act_bias[I_d][2] * rigid_info.substep_dt[None] - ) - i_d_ = i_d_ + BLOCK_DIM - qd.simt.block.sync() + i_e = rigid_info.awake_entities[i_e, i_b] - # In-place LDL^T, eliminating dofs from last to first (matches the scalar branch). - for j in range(n_block_dofs): - i_d = block_end - j - 1 - i_d_local = i_d - block_start - D_inv = 1.0 / rigid_info.mass_mat_L[i_d, i_d, i_b] - if tid == 0: - rigid_info.mass_mat_D_inv[i_d, i_b] = D_inv - - # Phase A: snapshot the (Schur-updated) pivot-row entries below the diagonal into shared. - j_d_ = tid - while j_d_ < i_d_local: - pivot_row[j_d_] = rigid_info.mass_mat_L[i_d, block_start + j_d_, i_b] - j_d_ = j_d_ + BLOCK_DIM - qd.simt.block.sync() - - # Phase B: each lane eliminates one column j_d, updating its own row j_d of the trailing - # submatrix from the read-only snapshot. Distinct rows per lane => no write conflicts, - # and the pivot row is only read (from shared) => no read/write race on row i_d. - j_d_ = tid - while j_d_ < i_d_local: - a = pivot_row[j_d_] * D_inv - j_d = block_start + j_d_ - for k_d_ in range(j_d_ + 1): - rigid_info.mass_mat_L[j_d, block_start + k_d_, i_b] = ( - rigid_info.mass_mat_L[j_d, block_start + k_d_, i_b] - a * pivot_row[k_d_] + if rigid_info.mass_mat_mask[i_e, i_b]: + entity_dof_start = dyn_info.entities.dof_start[i_e] + entity_dof_end = dyn_info.entities.dof_end[i_e] + + pivot_row = qd.simt.block.SharedArray((MAX_DOFS_PER_BLOCK,), gs.qd_float) + + # Factor each mass block rooted in this entity in-place in global memory over its full range, + # block-relative so shared indices stay >= 0; a merged child owns no root and factors nothing. + block_start = entity_dof_start + while block_start < entity_dof_end: + block_end = rigid_info.dofs_mass_block_end[block_start] + if rigid_info.dofs_mass_block_start[block_start] == block_start: + n_block_dofs = block_end - block_start + + # Copy the block's lower triangle into mass_mat_L (+ implicit damping on the diagonal), + # cooperatively. Restricting to the block makes the factorization cost the sum of per-block + # cubes instead of the whole (possibly multi-block) entity cube. + i_d_ = tid + while i_d_ < n_block_dofs: + i_d = block_start + i_d_ + for j_d in range(block_start, i_d + 1): + rigid_info.mass_mat_L[i_d, j_d, i_b] = rigid_info.mass_mat[i_d, j_d, i_b] + if qd.static(implicit_damping): + I_d = [i_d, i_b] if qd.static(rigid_config.batch_dofs_info) else i_d + rigid_info.mass_mat_L[i_d, i_d, i_b] = ( + rigid_info.mass_mat_L[i_d, i_d, i_b] + + dyn_info.dofs.damping[I_d] * rigid_info.substep_dt[None] + ) + if qd.static(rigid_config.integrator == gs.integrator.implicitfast): + if dyn_state.dofs.ctrl_mode[i_d, i_b] <= gs.CTRL_MODE.VELOCITY: + rigid_info.mass_mat_L[i_d, i_d, i_b] = ( + rigid_info.mass_mat_L[i_d, i_d, i_b] + - dyn_info.dofs.act_bias[I_d][2] * rigid_info.substep_dt[None] ) - rigid_info.mass_mat_L[i_d, j_d, i_b] = a - j_d_ = j_d_ + BLOCK_DIM - qd.simt.block.sync() - - # Diagonal coeffs of L are ignored downstream (see scalar branch) but set to 1.0 to match. - if tid == 0: - rigid_info.mass_mat_L[i_d, i_d, i_b] = 1.0 - block_start = block_end - elif qd.static(not rigid_config.enable_tiled_cholesky_mass_matrix or rigid_config.backend == gs.cpu): - qd.loop_config(name="factor_mass", serialize=rigid_config.para_level < gs.PARA_LEVEL.PARTIAL) - for i_slot, i_b in qd.ndrange(n_entities, _B): - # Skip hibernated entities: their mass matrix is unchanged, so the factor from the last awake step - # stays valid. This makes the factorization cost scale with the awake entity count. - i_e = i_slot - if qd.static(rigid_config.use_hibernation): - if i_slot >= rigid_info.n_awake_entities[i_b]: - continue - i_e = rigid_info.awake_entities[i_slot, i_b] - if rigid_info.mass_mat_mask[i_e, i_b]: - # Factor each mass block rooted in this entity, iterated flat over the rooted range with per-DOF - # block bounds (see entities_mass_block_dof_start in array_class.py): elimination never leaves a - # block, so interleaving independent blocks in one descending scan is exact. - blocks_dof_start = rigid_info.entities_mass_block_dof_start[i_e] - blocks_dof_end = rigid_info.entities_mass_block_dof_end[i_e] - for i_d in range(blocks_dof_start, blocks_dof_end): - for j_d in range(rigid_info.dofs_mass_block_start[i_d], i_d + 1): - rigid_info.mass_mat_L[i_d, j_d, i_b] = rigid_info.mass_mat[i_d, j_d, i_b] - - if qd.static(implicit_damping): - I_d = [i_d, i_b] if qd.static(rigid_config.batch_dofs_info) else i_d - rigid_info.mass_mat_L[i_d, i_d, i_b] = ( - rigid_info.mass_mat_L[i_d, i_d, i_b] - + dyn_info.dofs.damping[I_d] * rigid_info.substep_dt[None] - ) - if qd.static(rigid_config.integrator == gs.integrator.implicitfast): - if dyn_state.dofs.ctrl_mode[i_d, i_b] <= gs.CTRL_MODE.VELOCITY: - rigid_info.mass_mat_L[i_d, i_d, i_b] = ( - rigid_info.mass_mat_L[i_d, i_d, i_b] - - dyn_info.dofs.act_bias[I_d][2] * rigid_info.substep_dt[None] - ) - - for i_d_ in range(blocks_dof_end - blocks_dof_start): - i_d = blocks_dof_end - i_d_ - 1 - block_start = rigid_info.dofs_mass_block_start[i_d] - D_inv = 1.0 / rigid_info.mass_mat_L[i_d, i_d, i_b] - rigid_info.mass_mat_D_inv[i_d, i_b] = D_inv - - for j_d_ in range(i_d - block_start): - j_d = i_d - j_d_ - 1 - a = rigid_info.mass_mat_L[i_d, j_d, i_b] * D_inv - for k_d in range(block_start, j_d + 1): - rigid_info.mass_mat_L[j_d, k_d, i_b] -= a * rigid_info.mass_mat_L[i_d, k_d, i_b] - rigid_info.mass_mat_L[i_d, j_d, i_b] = a - - # FIXME: Diagonal coeffs of L are ignored in computations, so no need to update them. - rigid_info.mass_mat_L[i_d, i_d, i_b] = 1.0 - else: - BLOCK_DIM = qd.static(32) - MAX_DOFS_PER_BLOCK = qd.static(rigid_config.tiled_n_dofs_per_block) - WARP_SIZE = qd.static(32) - - qd.loop_config(name="factor_mass", block_dim=BLOCK_DIM) - for i in range(n_entities * _B * BLOCK_DIM): - tid = i % BLOCK_DIM - i_e = (i // BLOCK_DIM) % n_entities - i_b = i // (BLOCK_DIM * n_entities) - if i_b >= _B: - continue - # Skip hibernated entities: their mass matrix is unchanged, so the factor from the last awake step - # stays valid. The slot remaps to an awake entity, so the work scales with the awake entity count. - if qd.static(rigid_config.use_hibernation): - if i_e >= rigid_info.n_awake_entities[i_b]: - continue - i_e = rigid_info.awake_entities[i_e, i_b] - - if rigid_info.mass_mat_mask[i_e, i_b]: - entity_dof_start = dyn_info.entities.dof_start[i_e] - entity_dof_end = dyn_info.entities.dof_end[i_e] - - mass_mat = qd.simt.block.SharedArray((MAX_DOFS_PER_BLOCK, MAX_DOFS_PER_BLOCK + 1), gs.qd_float) - - # Factor each mass block rooted in this entity in shared memory, indexed block-relative so - # shared indices stay >= 0 (a merged child's block starts before its entity); the child owns no - # root and factors nothing, while the parent factors the whole coupled block. - block_start = entity_dof_start - while block_start < entity_dof_end: - block_end = rigid_info.dofs_mass_block_end[block_start] - if rigid_info.dofs_mass_block_start[block_start] == block_start: - n_block_dofs = block_end - block_start - n_lower_tri = n_block_dofs * (n_block_dofs + 1) // 2 - - i_pair = tid - while i_pair < n_lower_tri: - i_d_, j_d_ = linear_to_lower_tri(i_pair) - mass_mat[i_d_, j_d_] = rigid_info.mass_mat[block_start + i_d_, block_start + j_d_, i_b] - i_pair = i_pair + BLOCK_DIM + i_d_ = i_d_ + BLOCK_DIM + qd.simt.block.sync() + + # In-place LDL^T, eliminating dofs from last to first (matches the scalar branch). + for j in range(n_block_dofs): + i_d = block_end - j - 1 + i_d_local = i_d - block_start + D_inv = 1.0 / rigid_info.mass_mat_L[i_d, i_d, i_b] + if tid == 0: + rigid_info.mass_mat_D_inv[i_d, i_b] = D_inv + + # Phase A: snapshot the (Schur-updated) pivot-row entries below the diagonal into shared. + j_d_ = tid + while j_d_ < i_d_local: + pivot_row[j_d_] = rigid_info.mass_mat_L[i_d, block_start + j_d_, i_b] + j_d_ = j_d_ + BLOCK_DIM qd.simt.block.sync() - if qd.static(implicit_damping): - i_d_ = tid - while i_d_ < n_block_dofs: - i_d = block_start + i_d_ - I_d = [i_d, i_b] if qd.static(rigid_config.batch_dofs_info) else i_d - mass_mat[i_d_, i_d_] = ( - mass_mat[i_d_, i_d_] + dyn_info.dofs.damping[I_d] * rigid_info.substep_dt[None] + # Phase B: each lane eliminates one column j_d, updating its own row j_d of the trailing + # submatrix from the read-only snapshot. Distinct rows per lane => no write conflicts, + # and the pivot row is only read (from shared) => no read/write race on row i_d. + j_d_ = tid + while j_d_ < i_d_local: + a = pivot_row[j_d_] * D_inv + j_d = block_start + j_d_ + for k_d_ in range(j_d_ + 1): + rigid_info.mass_mat_L[j_d, block_start + k_d_, i_b] = ( + rigid_info.mass_mat_L[j_d, block_start + k_d_, i_b] - a * pivot_row[k_d_] ) - if qd.static(rigid_config.integrator == gs.integrator.implicitfast): - if dyn_state.dofs.ctrl_mode[i_d, i_b] <= gs.CTRL_MODE.VELOCITY: - mass_mat[i_d_, i_d_] = ( - mass_mat[i_d_, i_d_] - - dyn_info.dofs.act_bias[I_d][2] * rigid_info.substep_dt[None] - ) - i_d_ = i_d_ + BLOCK_DIM - qd.simt.block.sync() - - for j in range(n_block_dofs): - i_d_ = n_block_dofs - j - 1 - i_d = block_end - j - 1 - - D_inv = 1.0 / mass_mat[i_d_, i_d_] - if tid == 0: - rigid_info.mass_mat_D_inv[i_d, i_b] = D_inv - # FIXME: Diagonal coeffs of L are ignored in computations, so no need to update them. - rigid_info.mass_mat_L[i_d, i_d, i_b] = 1.0 - - j_d_ = i_d_ - 1 - tid - while j_d_ >= 0: - a = mass_mat[i_d_, j_d_] * D_inv - for k_d in range(j_d_ + 1): - mass_mat[j_d_, k_d] = mass_mat[j_d_, k_d] - a * mass_mat[i_d_, k_d] - mass_mat[i_d_, j_d_] = a - j_d_ = j_d_ - BLOCK_DIM - if qd.static(rigid_config.backend == gs.cuda): - if i_d_ <= WARP_SIZE: - qd.simt.warp.sync(qd.u32(0xFFFFFFFF)) - else: - qd.simt.block.sync() - else: - qd.simt.block.sync() - - i_pair = tid - n_strict_lower_tri = n_block_dofs * (n_block_dofs - 1) // 2 - while i_pair < n_strict_lower_tri: - i_d_, j_d_ = linear_to_lower_tri(i_pair, strict=True) - rigid_info.mass_mat_L[block_start + i_d_, block_start + j_d_, i_b] = mass_mat[ - i_d_, j_d_ - ] - i_pair = i_pair + BLOCK_DIM + rigid_info.mass_mat_L[i_d, j_d, i_b] = a + j_d_ = j_d_ + BLOCK_DIM qd.simt.block.sync() - block_start = block_end - else: - # Cholesky decomposition that has safe access pattern and robust handling of divide by zero for AD. Even though - # it is logically equivalent to the above block, it shows slightly numerical difference in the result, and thus - # it fails for a unit test ("test_urdf_rope"), while passing all the others. TODO: Investigate if we can fix this - # and only use this block. - # Assume this is the outermost loop - qd.loop_config(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.dofs.ctrl_mode.shape[1]): + # Diagonal coeffs of L are ignored downstream (see scalar branch) but set to 1.0 to match. + if tid == 0: + rigid_info.mass_mat_L[i_d, i_d, i_b] = 1.0 + block_start = block_end + elif qd.static(not rigid_config.enable_tiled_cholesky_mass_matrix or rigid_config.backend == gs.cpu): + qd.loop_config(name="factor_mass", serialize=rigid_config.para_level < gs.PARA_LEVEL.PARTIAL) + for i_slot, i_b in qd.ndrange(n_entities, _B): + # Skip hibernated entities: their mass matrix is unchanged, so the factor from the last awake step + # stays valid. This makes the factorization cost scale with the awake entity count. + i_e = i_slot + if qd.static(rigid_config.use_hibernation): + if i_slot >= rigid_info.n_awake_entities[i_b]: + continue + i_e = rigid_info.awake_entities[i_slot, i_b] if rigid_info.mass_mat_mask[i_e, i_b]: - EPS = rigid_info.EPS[None] - - # Factor each mass block rooted in this entity, iterated flat like the forward scalar arm; the - # per-block index reversal this AD-safe factor uses maps each block onto itself. + # Factor each mass block rooted in this entity, iterated flat over the rooted range with per-DOF + # block bounds (see entities_mass_block_dof_start in array_class.py): elimination never leaves a + # block, so interleaving independent blocks in one descending scan is exact. blocks_dof_start = rigid_info.entities_mass_block_dof_start[i_e] blocks_dof_end = rigid_info.entities_mass_block_dof_end[i_e] for i_d in range(blocks_dof_start, blocks_dof_end): - block_start = rigid_info.dofs_mass_block_start[i_d] - block_end = rigid_info.dofs_mass_block_end[i_d] - i_pr = (block_start + block_end - 1) - i_d - for j_d in range(block_start, i_d + 1): - j_pr = (block_start + block_end - 1) - j_d - rigid_info.mass_mat_L_bw[0, i_pr, j_pr, i_b] = rigid_info.mass_mat[i_d, j_d, i_b] - rigid_info.mass_mat_L_bw[0, j_pr, i_pr, i_b] = rigid_info.mass_mat[i_d, j_d, i_b] + for j_d in range(rigid_info.dofs_mass_block_start[i_d], i_d + 1): + rigid_info.mass_mat_L[i_d, j_d, i_b] = rigid_info.mass_mat[i_d, j_d, i_b] if qd.static(implicit_damping): I_d = [i_d, i_b] if qd.static(rigid_config.batch_dofs_info) else i_d - qd.atomic_add( - rigid_info.mass_mat_L_bw[0, i_pr, i_pr, i_b], - (dyn_info.dofs.damping[I_d] * rigid_info.substep_dt[None]), + rigid_info.mass_mat_L[i_d, i_d, i_b] = ( + rigid_info.mass_mat_L[i_d, i_d, i_b] + + dyn_info.dofs.damping[I_d] * rigid_info.substep_dt[None] ) if qd.static(rigid_config.integrator == gs.integrator.implicitfast): if dyn_state.dofs.ctrl_mode[i_d, i_b] <= gs.CTRL_MODE.VELOCITY: - qd.atomic_add( - rigid_info.mass_mat_L_bw[0, i_pr, i_pr, i_b], - -dyn_info.dofs.act_bias[I_d][2] * rigid_info.substep_dt[None], + rigid_info.mass_mat_L[i_d, i_d, i_b] = ( + rigid_info.mass_mat_L[i_d, i_d, i_b] + - dyn_info.dofs.act_bias[I_d][2] * rigid_info.substep_dt[None] ) - # Cholesky-Banachiewicz algorithm (in the perturbed indices), access pattern is safe for autodiff - # https://en.wikipedia.org/wiki/Cholesky_decomposition - for i_pr in range(blocks_dof_start, blocks_dof_end): - block_start = rigid_info.dofs_mass_block_start[i_pr] - for j_pr in range(block_start, i_pr + 1): - sum = gs.qd_float(0.0) - for k_pr in range(block_start, j_pr): - sum = sum + ( - rigid_info.mass_mat_L_bw[1, i_pr, k_pr, i_b] - * rigid_info.mass_mat_L_bw[1, j_pr, k_pr, i_b] - ) + for i_d_ in range(blocks_dof_end - blocks_dof_start): + i_d = blocks_dof_end - i_d_ - 1 + block_start = rigid_info.dofs_mass_block_start[i_d] + D_inv = 1.0 / rigid_info.mass_mat_L[i_d, i_d, i_b] + rigid_info.mass_mat_D_inv[i_d, i_b] = D_inv + + for j_d_ in range(i_d - block_start): + j_d = i_d - j_d_ - 1 + a = rigid_info.mass_mat_L[i_d, j_d, i_b] * D_inv + for k_d in range(block_start, j_d + 1): + rigid_info.mass_mat_L[j_d, k_d, i_b] -= a * rigid_info.mass_mat_L[i_d, k_d, i_b] + rigid_info.mass_mat_L[i_d, j_d, i_b] = a + + # FIXME: Diagonal coeffs of L are ignored in computations, so no need to update them. + rigid_info.mass_mat_L[i_d, i_d, i_b] = 1.0 + else: + BLOCK_DIM = qd.static(32) + MAX_DOFS_PER_BLOCK = qd.static(rigid_config.tiled_n_dofs_per_block) + WARP_SIZE = qd.static(32) + + qd.loop_config(name="factor_mass", block_dim=BLOCK_DIM) + for i in range(n_entities * _B * BLOCK_DIM): + tid = i % BLOCK_DIM + i_e = (i // BLOCK_DIM) % n_entities + i_b = i // (BLOCK_DIM * n_entities) + if i_b >= _B: + continue + # Skip hibernated entities: their mass matrix is unchanged, so the factor from the last awake step + # stays valid. The slot remaps to an awake entity, so the work scales with the awake entity count. + if qd.static(rigid_config.use_hibernation): + if i_e >= rigid_info.n_awake_entities[i_b]: + continue + i_e = rigid_info.awake_entities[i_e, i_b] - a = rigid_info.mass_mat_L_bw[0, i_pr, j_pr, i_b] - sum - b = qd.math.clamp(rigid_info.mass_mat_L_bw[1, j_pr, j_pr, i_b], EPS, qd.math.inf) - if i_pr == j_pr: - rigid_info.mass_mat_L_bw[1, i_pr, j_pr, i_b] = qd.sqrt(qd.math.clamp(a, EPS, qd.math.inf)) - else: - rigid_info.mass_mat_L_bw[1, i_pr, j_pr, i_b] = a / b + if rigid_info.mass_mat_mask[i_e, i_b]: + entity_dof_start = dyn_info.entities.dof_start[i_e] + entity_dof_end = dyn_info.entities.dof_end[i_e] + + mass_mat = qd.simt.block.SharedArray((MAX_DOFS_PER_BLOCK, MAX_DOFS_PER_BLOCK + 1), gs.qd_float) + + # Factor each mass block rooted in this entity in shared memory, indexed block-relative so + # shared indices stay >= 0 (a merged child's block starts before its entity); the child owns no + # root and factors nothing, while the parent factors the whole coupled block. + block_start = entity_dof_start + while block_start < entity_dof_end: + block_end = rigid_info.dofs_mass_block_end[block_start] + if rigid_info.dofs_mass_block_start[block_start] == block_start: + n_block_dofs = block_end - block_start + n_lower_tri = n_block_dofs * (n_block_dofs + 1) // 2 + + i_pair = tid + while i_pair < n_lower_tri: + i_d_, j_d_ = linear_to_lower_tri(i_pair) + mass_mat[i_d_, j_d_] = rigid_info.mass_mat[block_start + i_d_, block_start + j_d_, i_b] + i_pair = i_pair + BLOCK_DIM + qd.simt.block.sync() - for i_d in range(blocks_dof_start, blocks_dof_end): - block_start = rigid_info.dofs_mass_block_start[i_d] - block_end = rigid_info.dofs_mass_block_end[i_d] - i_pr = (block_start + block_end - 1) - i_d - for j_d in range(block_start, i_d + 1): - j_pr = (block_start + block_end - 1) - j_d + if qd.static(implicit_damping): + i_d_ = tid + while i_d_ < n_block_dofs: + i_d = block_start + i_d_ + I_d = [i_d, i_b] if qd.static(rigid_config.batch_dofs_info) else i_d + mass_mat[i_d_, i_d_] = ( + mass_mat[i_d_, i_d_] + dyn_info.dofs.damping[I_d] * rigid_info.substep_dt[None] + ) + if qd.static(rigid_config.integrator == gs.integrator.implicitfast): + if dyn_state.dofs.ctrl_mode[i_d, i_b] <= gs.CTRL_MODE.VELOCITY: + mass_mat[i_d_, i_d_] = ( + mass_mat[i_d_, i_d_] + - dyn_info.dofs.act_bias[I_d][2] * rigid_info.substep_dt[None] + ) + i_d_ = i_d_ + BLOCK_DIM + qd.simt.block.sync() - a = rigid_info.mass_mat_L_bw[1, i_pr, i_pr, i_b] - rigid_info.mass_mat_L[i_d, j_d, i_b] = rigid_info.mass_mat_L_bw[ - 1, j_pr, i_pr, i_b - ] / qd.math.clamp(a, EPS, qd.math.inf) + for j in range(n_block_dofs): + i_d_ = n_block_dofs - j - 1 + i_d = block_end - j - 1 + + D_inv = 1.0 / mass_mat[i_d_, i_d_] + if tid == 0: + rigid_info.mass_mat_D_inv[i_d, i_b] = D_inv + # FIXME: Diagonal coeffs of L are ignored in computations, so no need to update them. + rigid_info.mass_mat_L[i_d, i_d, i_b] = 1.0 + + j_d_ = i_d_ - 1 - tid + while j_d_ >= 0: + a = mass_mat[i_d_, j_d_] * D_inv + for k_d in range(j_d_ + 1): + mass_mat[j_d_, k_d] = mass_mat[j_d_, k_d] - a * mass_mat[i_d_, k_d] + mass_mat[i_d_, j_d_] = a + j_d_ = j_d_ - BLOCK_DIM + if qd.static(rigid_config.backend == gs.cuda): + if i_d_ <= WARP_SIZE: + qd.simt.warp.sync(qd.u32(0xFFFFFFFF)) + else: + qd.simt.block.sync() + else: + qd.simt.block.sync() - if i_d == j_d: - rigid_info.mass_mat_D_inv[i_d, i_b] = 1.0 / (qd.math.clamp(a**2, EPS, qd.math.inf)) + i_pair = tid + n_strict_lower_tri = n_block_dofs * (n_block_dofs - 1) // 2 + while i_pair < n_strict_lower_tri: + i_d_, j_d_ = linear_to_lower_tri(i_pair, strict=True) + rigid_info.mass_mat_L[block_start + i_d_, block_start + j_d_, i_b] = mass_mat[i_d_, j_d_] + i_pair = i_pair + BLOCK_DIM + qd.simt.block.sync() + block_start = block_end @qd.func @@ -826,14 +742,10 @@ def func_solve_mass_entity( i_b: qd.int32, vec: qd.Tensor, out: qd.Tensor, - out_bw: qd.template(), dyn_info: array_class.DynInfo, rigid_info: array_class.RigidInfo, rigid_config: qd.template(), - is_backward: qd.template(), ): - BW = qd.static(is_backward) - if rigid_info.mass_mat_mask[i_e, i_b]: # Solve M x = y for each mass block rooted in this entity, iterated flat over the rooted range with per-DOF # block bounds (see entities_mass_block_dof_start in array_class.py); the substitutions never cross block @@ -841,44 +753,28 @@ def func_solve_mass_entity( blocks_dof_start = rigid_info.entities_mass_block_dof_start[i_e] blocks_dof_end = rigid_info.entities_mass_block_dof_end[i_e] - # Step 1: Solve w st. L^T @ w = y + # Step 1: Solve w st. L^T @ w = y. Reading out[j_d] (j_d > i_d) from the buffer being written is safe: those + # entries were finalized in earlier (larger i_d) iterations. This func is never auto-reversed; the backward + # pass seeds mass_mat.grad directly via the implicit function theorem (see kernel_manual_compute_qacc_bw in + # manual_bw.py). for i_d_ in range(blocks_dof_end - blocks_dof_start): i_d = blocks_dof_end - i_d_ - 1 block_end = rigid_info.dofs_mass_block_end[i_d] curr_out = vec[i_d, i_b] - if qd.static(BW): - out_bw[0, i_d, i_b] = vec[i_d, i_b] - for j_d in range(i_d + 1, block_end): - # Since we read out[j_d, i_b], and j_d > i_d, which means that out[j_d, i_b] is already - # finalized at this point, we don't need to care about AD mutation rule. - if qd.static(BW): - out_bw[0, i_d, i_b] = ( - out_bw[0, i_d, i_b] - rigid_info.mass_mat_L[j_d, i_d, i_b] * out_bw[0, j_d, i_b] - ) - else: - curr_out = curr_out - rigid_info.mass_mat_L[j_d, i_d, i_b] * out[j_d, i_b] - - if qd.static(not BW): - out[i_d, i_b] = curr_out + curr_out = curr_out - rigid_info.mass_mat_L[j_d, i_d, i_b] * out[j_d, i_b] + out[i_d, i_b] = curr_out # Step 2: z = D^{-1} w for i_d in range(blocks_dof_start, blocks_dof_end): - if qd.static(BW): - out_bw[1, i_d, i_b] = out_bw[0, i_d, i_b] * rigid_info.mass_mat_D_inv[i_d, i_b] - else: - out[i_d, i_b] = out[i_d, i_b] * rigid_info.mass_mat_D_inv[i_d, i_b] + out[i_d, i_b] = out[i_d, i_b] * rigid_info.mass_mat_D_inv[i_d, i_b] # Step 3: Solve x st. L @ x = z for i_d in range(blocks_dof_start, blocks_dof_end): block_start = rigid_info.dofs_mass_block_start[i_d] curr_out = out[i_d, i_b] - if qd.static(BW): - curr_out = out_bw[1, i_d, i_b] - for j_d in range(block_start, i_d): curr_out = curr_out - rigid_info.mass_mat_L[i_d, j_d, i_b] * out[j_d, i_b] - out[i_d, i_b] = curr_out @@ -887,15 +783,10 @@ def func_solve_mass_batch( i_b: qd.int32, vec: qd.Tensor, out: qd.Tensor, - out_bw: qd.template(), dyn_info: array_class.DynInfo, rigid_info: array_class.RigidInfo, rigid_config: qd.template(), - is_backward: qd.template(), ): - BW = qd.static(is_backward) - - # This loop is considered an inner loop qd.loop_config(serialize=qd.static(rigid_config.para_level < gs.PARA_LEVEL.ALL)) for i_0 in ( range(rigid_info.n_awake_entities[i_b]) @@ -903,23 +794,20 @@ def func_solve_mass_batch( else range(dyn_info.entities.n_links.shape[0]) ): i_e = rigid_info.awake_entities[i_0, i_b] if qd.static(rigid_config.use_hibernation) else i_0 - func_solve_mass_entity(i_e, i_b, vec, out, out_bw, dyn_info, rigid_info, rigid_config, is_backward) + func_solve_mass_entity(i_e, i_b, vec, out, dyn_info, rigid_info, rigid_config) @qd.func def func_solve_mass( vec: qd.Tensor, out: qd.Tensor, - out_bw: qd.template(), # None in forward mode, real tensor in backward mode dyn_info: array_class.DynInfo, rigid_info: array_class.RigidInfo, rigid_config: qd.template(), - is_backward: qd.template(), ): - # This loop must be the outermost loop to be differentiable qd.loop_config(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], out.shape[1]): - func_solve_mass_entity(i_e, i_b, vec, out, out_bw, dyn_info, rigid_info, rigid_config, is_backward) + func_solve_mass_entity(i_e, i_b, vec, out, dyn_info, rigid_info, rigid_config) @qd.func @@ -1320,15 +1208,7 @@ def func_compute_qacc( ): BW = qd.static(is_backward) - func_solve_mass( - dyn_state.dofs.force, - dyn_state.dofs.acc_smooth, - dyn_state.dofs.acc_smooth_bw, - dyn_info, - rigid_info, - rigid_config, - 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 qd.loop_config(serialize=qd.static(rigid_config.para_level < gs.PARA_LEVEL.ALL)) @@ -1603,11 +1483,7 @@ def func_integrate( dyn_state.dofs.vel_next[dof_start + 2, i_b], ] ) - # Backward pass requires atomic add - if qd.static(BW): - qd.atomic_add(pos, vel * rigid_info.substep_dt[None]) - else: - pos = pos + vel * rigid_info.substep_dt[None] + pos = pos + vel * rigid_info.substep_dt[None] for j in qd.static(range(3)): rigid_info.qpos_next[q_start + j, i_b] = pos[j] if joint_type == gs.JOINT_TYPE.SPHERICAL or joint_type == gs.JOINT_TYPE.FREE: @@ -1681,6 +1557,10 @@ def kernel_forward_dynamics_without_qacc( rigid_config: qd.template(), is_backward: qd.template(), ): + # Backward-only kernel. func_factor_mass is omitted: its reverse is unneeded since the backward pass seeds + # mass_mat.grad directly via the implicit function theorem (see kernel_manual_compute_qacc_bw in manual_bw.py), + # skipping the LDL^T factor chain. func_compute_mass_matrix is kept so Quadrants autodiff auto-reverses + # mass_mat -> links pos / quat. func_compute_mass_matrix( dyn_state, dyn_info, @@ -1689,7 +1569,6 @@ def kernel_forward_dynamics_without_qacc( qd.static(rigid_config.integrator == gs.integrator.approximate_implicitfast), is_backward, ) - func_factor_mass(dyn_state, dyn_info, rigid_info, rigid_config, implicit_damping=False, is_backward=is_backward) func_torque_and_passive_force(dyn_state, constraint_state, dyn_info, rigid_info, rigid_config, is_backward) 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) @@ -1736,10 +1615,8 @@ def func_implicit_damping( ): rigid_info.mass_mat_mask[i_e, i_b] = True - func_factor_mass(dyn_state, dyn_info, rigid_info, rigid_config, implicit_damping=True, is_backward=is_backward) - func_solve_mass( - dyn_state.dofs.force, dyn_state.dofs.acc, dyn_state.dofs.acc_bw, dyn_info, rigid_info, rigid_config, is_backward - ) + func_factor_mass(dyn_state, dyn_info, rigid_info, rigid_config, implicit_damping=True) + func_solve_mass(dyn_state.dofs.force, dyn_state.dofs.acc, dyn_info, rigid_info, rigid_config) # Disable pre-computed factorization mask right away if qd.static(not rigid_config.enable_mujoco_compatibility or rigid_config.integrator == gs.integrator.Euler): diff --git a/genesis/engine/solvers/rigid/abd/forward_kinematics.py b/genesis/engine/solvers/rigid/abd/forward_kinematics.py index 18aca42e32..fa0af99d85 100644 --- a/genesis/engine/solvers/rigid/abd/forward_kinematics.py +++ b/genesis/engine/solvers/rigid/abd/forward_kinematics.py @@ -459,6 +459,10 @@ def func_forward_kinematics_entity( + dyn_state.joints.xaxis[i_j, i_b] * dyn_state.dofs.pos[dof_start, i_b] ) pos = W(next_I, pos_, dyn_state.links.pos_bw, BW) + # A prismatic joint leaves the link orientation unchanged, but the backward per-joint cache still + # needs the next slot populated: the final R(quat_bw, I_jf, ...) below reads it in backward mode and + # would get uninitialized memory (NaN gradients on qpos) otherwise. + quat = W(next_I, quat, dyn_state.links.quat_bw, BW) # Skip link pose update for fixed root links to let users manually overwrite them I_jf = (i_l, 0 if qd.static(not BW) else n_joints, i_b) @@ -1148,3 +1152,49 @@ def kernel_update_cartesian_space( is_backward: qd.template(), ): func_update_cartesian_space(dyn_state, dyn_info, rigid_info, rigid_config, force_update_fixed_geoms, is_backward) + + +# Standalone forward-replay kernels for the update_cartesian_space sub-stages, used only in the backward pass +# (substep_pre_coupling_grad). The backward unroll replays each sub-stage with is_backward=True (static loops) and +# then reverses it, stage by stage: COM-links and geom-pose updates reverse cleanly through Quadrants autodiff (their +# replay kernel's .grad is called directly), while forward kinematics and forward velocity are reversed manually +# (see kernel_manual_forward_kinematics_bw / kernel_manual_forward_velocity_bw in manual_bw.py) because autodiff +# silently drops their gradient. TODO: once every sub-stage reverses correctly under autodiff, drop these kernels and +# the manual reverses, and differentiate kernel_update_cartesian_space as a whole. +@qd.kernel(fastcache=True) +def kernel_forward_kinematics_replay( + envs_idx: qd.types.ndarray(), + dyn_state: array_class.DynState, + dyn_info: array_class.DynInfo, + rigid_info: array_class.RigidInfo, + rigid_config: qd.template(), + is_backward: qd.template(), +): + for i_b_ in range(envs_idx.shape[0]): + i_b = qd.cast(envs_idx[i_b_], qd.i32) + func_forward_kinematics_batch(i_b, dyn_state, dyn_info, rigid_info, rigid_config, is_backward) + + +@qd.kernel(fastcache=True) +def kernel_update_geoms_replay( + dyn_state: array_class.DynState, + dyn_info: array_class.DynInfo, + rigid_info: array_class.RigidInfo, + rigid_config: qd.template(), + is_backward: qd.template(), +): + func_update_geoms( + dyn_state, dyn_info, rigid_info, rigid_config, force_update_fixed_geoms=False, is_backward=is_backward + ) + + +@qd.kernel(fastcache=True) +def kernel_COM_links_replay( + dyn_state: array_class.DynState, + dyn_info: array_class.DynInfo, + rigid_info: array_class.RigidInfo, + rigid_config: qd.template(), + is_backward: qd.template(), +): + for i_b in range(dyn_state.links.pos.shape[1]): + func_COM_links(i_b, dyn_state, dyn_info, rigid_info, rigid_config, is_backward) diff --git a/genesis/engine/solvers/rigid/abd/manual_bw.py b/genesis/engine/solvers/rigid/abd/manual_bw.py new file mode 100644 index 0000000000..544fc09024 --- /dev/null +++ b/genesis/engine/solvers/rigid/abd/manual_bw.py @@ -0,0 +1,545 @@ +"""Manual reverse-mode kernels for the rigid backward pass. + +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. +""" + +import quadrants as qd + +import genesis as gs +import genesis.utils.array_class as array_class +import genesis.utils.geom as gu + + +@qd.kernel(fastcache=True) +def kernel_manual_forward_kinematics_bw( + dyn_state: array_class.DynState, + 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. + + Iterates each entity's links leaf to root in one launch, so a child's parent pos / quat grad write lands before + the parent's own iteration consumes it, and within each link reverses the full joint chain. A link may carry + several joints (e.g. a planar floating base = slide-x + slide-z + hinge-y on one link); the forward composes them + in sequence and caches the per-joint intermediate pose in dyn_state.links.{pos,quat}_bw[i_l, k]: slot 0 is the + parent pose composed with the link's fixed offset, slot k+1 the pose after joint k, slot n_joints the final link + pose. The reverse seeds the grad on the final pose, reverses joint k for k = n_joints-1 .. 0 (each step consumes + the grad on slot k+1 and emits qpos.grad for that joint plus the grad on slot k), then reverses the base + composition into the parent's pose grad. Each joint also feeds dyn_state.joints.{xanchor,xaxis} downstream + (forward velocity), so their accumulated grads are folded back through slot k as well. + """ + qd.loop_config( + name="manual_fk_only_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]): + n_in_e = dyn_info.entities.n_links[i_e] + 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] + 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 + # backward through the joint chain; after the loop it holds the grad + # on slot 0 (arm base). + g_pos = dyn_state.links.pos.grad[i_l, i_b] + g_quat = dyn_state.links.quat.grad[i_l, i_b] + + for k_rev in range(n_joints): + k = n_joints - 1 - k_rev + i_j = dyn_info.links.joint_start[I_l] + k + 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] + q_start = dyn_info.joints.q_start[I_j] + dof_start = dyn_info.joints.dof_start[I_j] + I_d = [dof_start, i_b] if qd.static(rigid_config.batch_dofs_info) else dof_start + + # Input pose to joint k (slot k), cached by the forward replay. + pos_in = dyn_state.links.pos_bw[i_l, k, i_b] + quat_in = dyn_state.links.quat_bw[i_l, k, i_b] + joint_pos_off = dyn_info.joints.pos[I_j] + xanchor_grad = dyn_state.joints.xanchor.grad[i_j, i_b] + 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]. + 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] + ) + # 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|: + # J^T g = (g - qhat (qhat . g)) / |q|. Copying g straight into the raw entries leaves a spurious + # radial component that disagrees with finite differences even at unit length. + q_raw = qd.Vector( + [ + rigid_info.qpos[q_start + 3, i_b], + rigid_info.qpos[q_start + 4, i_b], + rigid_info.qpos[q_start + 5, i_b], + rigid_info.qpos[q_start + 6, i_b], + ], + dt=gs.qd_float, + ) + q_norm = q_raw.norm() + qhat = q_raw / q_norm + g_quat_raw = (g_quat - qhat * qhat.dot(g_quat)) / 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] + ) + 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) + + elif joint_type == gs.JOINT_TYPE.REVOLUTE: + axis = dyn_info.dofs.motion_ang[I_d] + angle = rigid_info.qpos[q_start, i_b] - rigid_info.qpos0[q_start, i_b] + rotvec = axis * angle + qloc = gu.qd_rotvec_to_quat(rotvec, rigid_info.EPS[None]) + # quat_out = transform_quat_by_quat(qloc, quat_in) = quat_mul(quat_in, qloc) + quat_out = gu.qd_transform_quat_by_quat(qloc, quat_in) + + # pos_out = xanchor - transform(joint_pos_off, quat_out) + # xanchor = transform(joint_pos_off, quat_in) + pos_in + 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) + 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 + + # grad into xanchor = g_pos (from pos_out) + downstream xanchor_grad + g_xanchor = g_pos + xanchor_grad + g_quat_in = ( + g_quat_in_apply + + gu.qd_transform_by_quat_grad_quat(joint_pos_off, quat_in, g_xanchor) + + gu.qd_transform_by_quat_grad_quat(axis, quat_in, xaxis_grad) + ) + g_pos = g_xanchor + g_quat = g_quat_in + + elif joint_type == gs.JOINT_TYPE.PRISMATIC: + axis = dyn_info.dofs.motion_vel[I_d] + displacement = rigid_info.qpos[q_start, i_b] - rigid_info.qpos0[q_start, i_b] + 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 + g_xaxis = qd.Vector( + [ + g_pos[0] * displacement + xaxis_grad[0], + g_pos[1] * displacement + xaxis_grad[1], + g_pos[2] * displacement + xaxis_grad[2], + ], + dt=gs.qd_float, + ) + g_xanchor = g_pos + xanchor_grad + 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) + ) + g_pos = g_xanchor + g_quat = g_quat_in + + elif joint_type == gs.JOINT_TYPE.SPHERICAL: + # qloc = qpos[q_start:q_start+4] (direct); quat_out = quat_mul(quat_in, qloc). + # axis defaults to [0,0,1] (xaxis = transform(axis, quat_in)). + axis = qd.Vector([0.0, 0.0, 1.0], dt=gs.qd_float) + qloc = qd.Vector( + [ + rigid_info.qpos[q_start, i_b], + rigid_info.qpos[q_start + 1, i_b], + rigid_info.qpos[q_start + 2, i_b], + rigid_info.qpos[q_start + 3, i_b], + ], + dt=gs.qd_float, + ) + quat_out = gu.qd_transform_quat_by_quat(qloc, quat_in) + 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) + 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] + g_xanchor = g_pos + xanchor_grad + g_quat_in = ( + g_quat_in_apply + + gu.qd_transform_by_quat_grad_quat(joint_pos_off, quat_in, g_xanchor) + + gu.qd_transform_by_quat_grad_quat(axis, quat_in, xaxis_grad) + ) + g_pos = g_xanchor + g_quat = g_quat_in + + else: # gs.JOINT_TYPE.FIXED - pose passes through unchanged. + pass + + for j in qd.static(range(3)): + dyn_state.joints.xanchor.grad[i_j, i_b][j] = 0.0 + dyn_state.joints.xaxis.grad[i_j, i_b][j] = 0.0 + + # Reverse the arm-base composition (slot 0): + # 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] + 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] + ) + 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] + + 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 + + +@qd.kernel(fastcache=True) +def kernel_manual_forward_velocity_bw( + dyn_state: array_class.DynState, + 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. + + Consumes the grad seeds on cd_vel / cd_ang, their per-joint caches cd_{vel,ang}_bw and cdofd_{ang,vel}, and + accumulates into dyn_state.dofs.{vel,cdof_ang,cdof_vel}.grad plus the parent links' cd_{vel,ang}.grad (the + cross-link chain matching the forward replay's cd_*_bw[i_l, 0] = parent cd_*). + """ + qd.loop_config( + name="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 + 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] + 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 + 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, curr_idx, i_b][k] = ( + dyn_state.links.cd_ang_bw.grad[i_l, curr_idx, i_b][k] + g_cd_ang_next[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.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) +def kernel_manual_compute_qacc_bw( + dyn_state: array_class.DynState, + dyn_info: array_class.DynInfo, + rigid_info: array_class.RigidInfo, + rigid_config: qd.template(), +): + """Manual backward for func_compute_qacc via the implicit function theorem (IFT). + + Forward chain (func_compute_qacc): + acc_smooth = M^{-1} . force (per-block LDL^T solve in func_solve_mass) + acc[i] = acc_smooth[i] (identity copy) + + Reverse chain (manual, by IFT and symmetry of M = L^T D L): + acc_smooth.grad += acc.grad (reverse of the identity copy; acc.grad is then consumed since the forward + copy overwrites acc) + force_contrib = M^{-1} . acc_smooth.grad (M is symmetric, so M^{-T} = M^{-1}) + force.grad += force_contrib + mass_mat[i, i].grad += -force_contrib[i] * acc_smooth[i] + mass_mat[i, j].grad += -(force_contrib[i] * acc_smooth[j] + force_contrib[j] * acc_smooth[i]) (i > j) + mass_mat is stored lower-triangular with the upper half implicit by symmetry, so each off-diagonal parameter + combines the chain terms of both its (i, j) and (j, i) occurrences. The forward factored the dense mass_mat into + mass_mat_L / mass_mat_D_inv already, so only mass_mat.grad is touched here; this kernel is the single place the + backward path populates it, and kernel_forward_dynamics_without_qacc.grad then reverses it into link poses. + + Like func_solve_mass_entity, the triangular solves and the IFT outer product are restricted to the mass blocks + rooted in each entity (see entities_mass_block_dof_start in array_class.py): elimination never crosses a block, + and cross-block mass entries are structural zeros whose grads must stay zero. + """ + qd.loop_config(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.dofs.force.shape[1]): + if rigid_info.mass_mat_mask[i_e, i_b]: + blocks_dof_start = rigid_info.entities_mass_block_dof_start[i_e] + blocks_dof_end = rigid_info.entities_mass_block_dof_end[i_e] + + # Reverse of acc[i] = acc_smooth[i]: drain acc.grad into the acc_smooth.grad seed, stashed in + # acc_smooth_bw[0] as the input of the LDL^T reverse solve. acc.grad is consumed since the forward copy + # overwrites acc. + for i_d in range(blocks_dof_start, blocks_dof_end): + dyn_state.dofs.acc_smooth_bw[0, i_d, i_b] = ( + dyn_state.dofs.acc_smooth.grad[i_d, i_b] + dyn_state.dofs.acc.grad[i_d, i_b] + ) + dyn_state.dofs.acc.grad[i_d, i_b] = 0.0 + dyn_state.dofs.acc_smooth.grad[i_d, i_b] = 0.0 + + # Step 1: solve L^T . u = seed (input from [0], output to [1]) + # u[i] = seed[i] - sum_{j>i} L[j,i] * u[j] + for i_d_ in range(blocks_dof_end - blocks_dof_start): + i_d = blocks_dof_end - i_d_ - 1 + block_end = rigid_info.dofs_mass_block_end[i_d] + curr = dyn_state.dofs.acc_smooth_bw[0, i_d, i_b] + for j_d in range(i_d + 1, block_end): + curr = curr - rigid_info.mass_mat_L[j_d, i_d, i_b] * dyn_state.dofs.acc_smooth_bw[1, j_d, i_b] + dyn_state.dofs.acc_smooth_bw[1, i_d, i_b] = curr + + # Step 2: v = D^{-1} . u (output to [0], overwriting input) + for i_d in range(blocks_dof_start, blocks_dof_end): + dyn_state.dofs.acc_smooth_bw[0, i_d, i_b] = ( + dyn_state.dofs.acc_smooth_bw[1, i_d, i_b] * rigid_info.mass_mat_D_inv[i_d, i_b] + ) + + # Step 3: solve L . delta = v (input from [0], output to [1]) + # delta[i] = v[i] - sum_{j None: self._solver._errno, ) + # Plane-convex contacts come from analytic paths that leave diff_contact_input unfilled; populate it here so + # the differentiable narrow-phase reverse can reconstruct them (see kernel_fill_diff_contact_input_plane). + if self._solver.rigid_config.requires_grad: + narrowphase.kernel_fill_diff_contact_input_plane( + self._solver.dyn_state, self._collider_state, self._solver.dyn_info, self._solver.rigid_config + ) + def get_contacts(self, as_tensor: bool = True, to_torch: bool = True, keep_batch_dim: bool = False): # Early return if already pre-computed contact_data = self._contact_data_cache.setdefault((as_tensor, to_torch), {}) @@ -1109,8 +1116,9 @@ def get_contacts(self, as_tensor: bool = True, to_torch: bool = True, keep_batch def backward(self, dL_dposition, dL_dnormal, dL_dpenetration): func_set_upstream_grad(dL_dposition, dL_dnormal, dL_dpenetration, self._collider_state) + self.backward_narrowphase() - # Compute gradient + def backward_narrowphase(self): func_narrow_phase_diff_convex_vs_convex.grad( self._solver.dyn_state, self._collider_state, diff --git a/genesis/engine/solvers/rigid/collider/diff_gjk.py b/genesis/engine/solvers/rigid/collider/diff_gjk.py index 97227359da..befb93028d 100644 --- a/genesis/engine/solvers/rigid/collider/diff_gjk.py +++ b/genesis/engine/solvers/rigid/collider/diff_gjk.py @@ -829,6 +829,54 @@ def func_differentiable_contact( return contact_pos, contact_normal, penetration, weight +@qd.func +def func_differentiable_plane_contact( + i_ga, + i_gb, + i_b, + i_c, + dyn_state: array_class.DynState, + diff_contact_input: array_class.DiffContactInput, + dyn_info: array_class.DynInfo, +): + """Differentiable plane-convex contact reconstruction. + + Mirrors the analytic plane branch of func_convex_convex_contact: + normal = -normalize(R(quat_plane) @ plane_local_dir) + v_world = R(quat_convex) @ core_local + pos_convex + radius * normal + penetration = normal . (v_world - pos_plane) + contact_pos = v_world - 0.5 * penetration * normal + + [i_ga] is the plane geom and [i_gb] the convex geom. [core_local] (box vertex / sphere center / capsule nearest + endpoint, in the convex geom's local frame) is the pose-independent witness stored by + kernel_fill_diff_contact_input_plane; [radius] and the plane direction come from the geoms info. Gradients flow to + both geom poses through the geoms state pos / quat. For a sphere, [core_local] is the local origin, so the + 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] + + 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) + v_world = core_world + radius * normal + + penetration = normal.dot(v_world - trans_plane) + contact_pos = v_world - 0.5 * penetration * normal + weight = gs.qd_float(1.0) + return contact_pos, normal, penetration, weight + + @qd.func def func_plane_normal(v1, v2, v3): """ diff --git a/genesis/engine/solvers/rigid/collider/narrowphase.py b/genesis/engine/solvers/rigid/collider/narrowphase.py index 240686eb31..017912976c 100644 --- a/genesis/engine/solvers/rigid/collider/narrowphase.py +++ b/genesis/engine/solvers/rigid/collider/narrowphase.py @@ -3015,9 +3015,18 @@ def func_narrow_phase_diff_convex_vs_convex( if is_ref: ref_penetration = -1.0 - contact_pos, contact_normal, penetration, weight = diff_gjk.func_differentiable_contact( - i_ga, i_gb, i_b, i_c, ref_penetration, dyn_state, diff_contact_input, collider_info - ) + contact_pos = gs.qd_vec3(0.0, 0.0, 0.0) + contact_normal = gs.qd_vec3(0.0, 0.0, 0.0) + penetration = gs.qd_float(0.0) + weight = gs.qd_float(0.0) + if dyn_info.geoms.type[i_ga] == gs.GEOM_TYPE.PLANE: + contact_pos, contact_normal, penetration, weight = diff_gjk.func_differentiable_plane_contact( + i_ga, i_gb, i_b, i_c, dyn_state, diff_contact_input, dyn_info + ) + else: + contact_pos, contact_normal, penetration, weight = diff_gjk.func_differentiable_contact( + i_ga, i_gb, i_b, i_c, ref_penetration, dyn_state, diff_contact_input, collider_info + ) collider_state.diff_contact_input.ref_penetration[i_b, i_c] = penetration func_set_contact( @@ -3065,6 +3074,58 @@ def func_narrow_phase_diff_convex_vs_convex( ) +@qd.kernel(fastcache=True) +def kernel_fill_diff_contact_input_plane( + dyn_state: array_class.DynState, + collider_state: array_class.ColliderState, + dyn_info: array_class.DynInfo, + rigid_config: qd.template(), +): + """Populate diff_contact_input for plane-convex contacts. + + The analytic plane paths (func_plane_box_contact, the plane branch of func_convex_convex_contact) leave + diff_contact_input unfilled, so the differentiable narrow-phase reverse would have nothing to reconstruct. Both + paths share the convention contact_pos = v - 0.5 * penetration * normal with + normal = -normalize(R(quat_plane) @ plane_local_dir), so the convex support point is recovered as + v = contact_pos + 0.5 * penetration * normal, and its pose-independent "core" (box vertex / sphere center / + capsule nearest endpoint) as v - radius * normal, stored in the convex geom's local frame. PLANE is the smallest + GEOM_TYPE so it is always geom_a after the canonical type-ordered swap. + """ + _B = collider_state.active_buffer.shape[1] + qd.loop_config(serialize=rigid_config.para_level < gs.PARA_LEVEL.PARTIAL) + for i_c, i_b in qd.ndrange(collider_state.contact_data.pos.shape[0], _B): + if i_c < collider_state.n_contacts[i_b]: + 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] + + penetration = collider_state.contact_data.penetration[i_c, i_b] + contact_pos = collider_state.contact_data.pos[i_c, i_b] + support_pos = contact_pos + 0.5 * penetration * normal + core_world = support_pos - radius * normal + core_local = gu.qd_transform_by_quat(core_world - trans_convex, gu.qd_inv_quat(quat_convex)) + + collider_state.diff_contact_input.geom_a[i_b, i_c] = i_ga + collider_state.diff_contact_input.geom_b[i_b, i_c] = i_gb + collider_state.diff_contact_input.core_local[i_b, i_c] = core_local + collider_state.diff_contact_input.ref_id[i_b, i_c] = i_c + collider_state.diff_contact_input.valid[i_b, i_c] = 1 + + @qd.kernel(fastcache=True) def func_narrow_phase_convex_specializations( geoms_init_AABB: array_class.GeomsInitAABB, diff --git a/genesis/engine/solvers/rigid/constraint/backward.py b/genesis/engine/solvers/rigid/constraint/backward.py index 18e249f6e1..70b3f9bc26 100644 --- a/genesis/engine/solvers/rigid/constraint/backward.py +++ b/genesis/engine/solvers/rigid/constraint/backward.py @@ -2,6 +2,7 @@ import genesis as gs import genesis.utils.array_class as array_class +import genesis.utils.geom as gu @qd.func @@ -52,6 +53,55 @@ def func_matvec_Ap( constraint_state.bw_Ap[i_d, i_b] += constraint_state.jac[i_c, i_d, i_b] * jv +@qd.func +def func_solve_adjoint_u_cg_env( + 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]. + + 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). + """ + n_dofs = constraint_state.bw_u.shape[0] + + # r = g - A*0 = g ; p = r ; u = 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]): + 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]) + + # 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]) + 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] + + @qd.kernel def kernel_solve_adjoint_u( constraint_state: array_class.ConstraintState, @@ -77,70 +127,34 @@ def kernel_solve_adjoint_u( constraint_state.bw_u[i_d, i_b] = 0.0 if qd.static(rigid_config.solver_type == gs.constraint_solver.Newton): - # Since we already have the Cholesky decomposition of A (= L * L^T), we can use it to solve A * u = g. for i_b in range(_B): - # z = L^{-1} g (forward substitution) - # Save solution to bw_r - for i_d in range(n_dofs): - z = constraint_state.dL_dqacc[i_d, i_b] - for j_d in range(i_d): - z -= constraint_state.nt_H[i_b, i_d, j_d] * constraint_state.bw_r[j_d, i_b] - z /= constraint_state.nt_H[i_b, i_d, i_d] - constraint_state.bw_r[i_d, i_b] = z - - # u = L^{-T} z (back substitution) - for i_d_ in range(n_dofs): - i_d = n_dofs - 1 - i_d_ - u = constraint_state.bw_r[i_d, i_b] - for j_d in range(i_d + 1, n_dofs): - u -= constraint_state.nt_H[i_b, j_d, i_d] * constraint_state.bw_u[j_d, i_b] - u /= constraint_state.nt_H[i_b, i_d, i_d] - constraint_state.bw_u[i_d, i_b] = u - else: - # Use CG solver for solving A * u = g. - # 2. Local buffers for solving A * u = g - # Initialize r, p with dL_dqacc - for i_d, i_b in qd.ndrange(n_dofs, _B): - # Residual: g - A * 0 (u = 0) - constraint_state.bw_r[i_d, i_b] = constraint_state.dL_dqacc[i_d, i_b] - # Search direction: p = r - constraint_state.bw_p[i_d, i_b] = constraint_state.bw_r[i_d, i_b] - - # 3. Solve A * u = g, parallelized over batch dimension - for i_b in range(_B): - # Compute Ap for the current search direction - for it in range(rigid_config.iterations): - func_matvec_Ap(i_b, constraint_state, dyn_info, rigid_info, rigid_config) - - # alpha = (r,r)/(p,Hp) - 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]) - - # u += alpha p ; r -= alpha Hp - 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] - - # check tol (optional: per-batch) - # TODO: Might need lower tolerance? - if num < rigid_info.EPS[None]: - break - - # beta = (r_new,r_new)/(r_old,r_old) - num_new = gs.qd_float(0.0) + if constraint_state.n_constraints[i_b] == 0: + # 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) + 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 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]) + z = constraint_state.dL_dqacc[i_d, i_b] + for j_d in range(i_d): + z -= constraint_state.nt_H[i_b, i_d, j_d] * constraint_state.bw_r[j_d, i_b] + z /= constraint_state.nt_H[i_b, i_d, i_d] + constraint_state.bw_r[i_d, i_b] = z - # p = r + beta p - 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] - ) + # u = L^{-T} z (back substitution) + for i_d_ in range(n_dofs): + i_d = n_dofs - 1 - i_d_ + u = constraint_state.bw_r[i_d, i_b] + for j_d in range(i_d + 1, n_dofs): + u -= constraint_state.nt_H[i_b, j_d, i_d] * constraint_state.bw_u[j_d, i_b] + u /= constraint_state.nt_H[i_b, i_d, i_d] + constraint_state.bw_u[i_d, i_b] = 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) @qd.kernel @@ -250,3 +264,1178 @@ def kernel_compute_gradients( val0 = -constraint_state.bw_u[i, i_b] * constraint_state.qacc[j, i_b] val1 = -constraint_state.bw_u[j, i_b] * constraint_state.qacc[i, i_b] constraint_state.dL_dM[i, j, i_b] += (val0 + val1) * 0.5 # symmetrize + + +@qd.kernel(fastcache=True) +def kernel_load_dL_dqacc_from_acc_grad( + dyn_state: array_class.DynState, + constraint_state: array_class.ConstraintState, + rigid_config: qd.template(), +): + """Copy the acc grad into constraint_state.dL_dqacc (the input buffer consumed by kernel_solve_adjoint_u) and + zero the source grad so the downstream implicit-function-theorem path does not re-consume it. + """ + _B = dyn_state.dofs.acc.shape[1] + 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), + ) + 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] + dyn_state.dofs.acc.grad[i_d, i_b] = gs.qd_float(0.0) + + +@qd.kernel(fastcache=True) +def kernel_accumulate_constraint_solver_grads( + dyn_state: array_class.DynState, + constraint_state: array_class.ConstraintState, + rigid_info: array_class.RigidInfo, + rigid_config: qd.template(), +): + """Fold the constraint-solver adjoint outputs into the autodiff grad fields: + dyn_state.dofs.force.grad += constraint_state.dL_dforce + rigid_info.mass_mat.grad += constraint_state.dL_dM + """ + _B = dyn_state.dofs.force.shape[1] + 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), + ) + 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] + for i, j, i_b in qd.ndrange(n_dofs, n_dofs, _B): + rigid_info.mass_mat.grad[i, j, i_b] += constraint_state.dL_dM[i, j, i_b] + + +# --------------------------------------------------------------------------- +# Manual reverses of the inequality constraints (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). +# +# 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 +# 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): +# 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) +# --------------------------------------------------------------------------- +@qd.kernel(fastcache=True) +def kernel_manual_add_joint_limit_constraints_bw( + dyn_state: array_class.DynState, + collider_state: array_class.ColliderState, + constraint_state: array_class.ConstraintState, + dyn_info: array_class.DynInfo, + rigid_info: array_class.RigidInfo, + rigid_config: qd.template(), + enable_collision: qd.template(), +): + """Manual reverse of `add_joint_limit_constraints`. See the section header + above for the shared `n_con` layout and upstream-grad conventions. + + Accumulates into rigid_info.qpos.grad[i_q] and dyn_state.dofs.vel.grad[i_d]. + + Chain rule (per active joint, `pos_delta < 0`): + + Forward: + pos_delta_min = qpos[i_q] - limit_lo + pos_delta_max = limit_hi - qpos[i_q] + pos_delta = min(pos_delta_min, pos_delta_max) + sign = +1 if pos_delta_min < pos_delta_max else -1 + jac_qvel = sign * dofs_vel[i_d] + imp, aref = gu.imp_aref(sol_params, pos_delta, jac_qvel, pos_delta) + diag_raw = invweight * (1 - imp) / imp + diag = max(diag_raw, EPS) + efc_D = 1 / diag + + d(pos_delta) / d(qpos) = sign (chosen branch of `min`) + d(jac_qvel) / d(vel) = sign + + dL/d(imp) = g_aref * d(aref)/d(imp) + g_efc_D * d(efc_D)/d(imp) + g_aref = dL_daref[n_con], g_efc_D = dL_defc_D[n_con] + + dL/d(pos_delta) = g_aref * d(aref)/d(pos_delta)|_direct + + dL/d(imp) * d(imp)/d(imp_x) * d(imp_x)/d(pos_delta) + + dL/d(jac_qvel) = g_aref * d(aref)/d(jac_qvel) = -g_aref * b_coef + + dL/d(qpos) += sign * dL/d(pos_delta) + dL/d(vel) += sign * dL/d(jac_qvel) + """ + EPS = rigid_info.EPS[None] + _B = constraint_state.jac.shape[2] + n_links = dyn_info.links.root_idx.shape[0] + + qd.loop_config( + name="kernel_manual_add_joint_limit_constraints_bw", + serialize=qd.static(rigid_config.para_level < gs.PARA_LEVEL.ALL), + ) + for i_b in range(_B): + # Forward row layout: equality -> frictionloss -> collision -> joint-limit. + # Seed the joint-limit counter past equality + frictionloss (always) + # and collision (when on). + n_con_counter = gs.qd_int( + 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) + + for i_l in range(n_links): + I_l = [i_l, i_b] if qd.static(rigid_config.batch_links_info) else i_l + for i_j in range(dyn_info.links.joint_start[I_l], dyn_info.links.joint_end[I_l]): + I_j = [i_j, i_b] if qd.static(rigid_config.batch_joints_info) else i_j + + if ( + dyn_info.joints.type[I_j] == gs.JOINT_TYPE.REVOLUTE + or dyn_info.joints.type[I_j] == gs.JOINT_TYPE.PRISMATIC + ): + i_q = dyn_info.joints.q_start[I_j] + i_d = dyn_info.joints.dof_start[I_j] + I_d = [i_d, i_b] if qd.static(rigid_config.batch_dofs_info) else i_d + + pos_delta_min = rigid_info.qpos[i_q, i_b] - dyn_info.dofs.limit[I_d][0] + pos_delta_max = dyn_info.dofs.limit[I_d][1] - rigid_info.qpos[i_q, i_b] + pos_delta = qd.min(pos_delta_min, pos_delta_max) + + if pos_delta < 0: + n_con = n_con_counter + n_con_counter = n_con_counter + 1 + + # Replay forward intermediates (cheap, avoids stashing). + sign_pos = (pos_delta_min < pos_delta_max) * 2 - 1 + 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] + 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 + diag = qd.max(diag_raw, EPS) + + # Upstream grads. + g_aref = constraint_state.dL_daref[n_con, i_b] + g_efc_D = constraint_state.dL_defc_D[n_con, i_b] + + # --- Partials of forward outputs w.r.t. intermediates --- + # aref = -b_coef * jac_qvel - k_coef * imp * pos_delta + d_aref_d_imp = -k_coef * pos_delta + 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 + 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 + + # --- Combine --- + dL_d_imp = g_aref * d_aref_d_imp + g_efc_D * d_efc_D_d_imp + dL_d_pos_delta = g_aref * d_aref_d_pos_delta_direct + dL_d_imp * d_imp_d_pos_delta + dL_d_jac_qvel = g_aref * d_aref_d_jac_qvel + + # --- Propagate --- + rigid_info.qpos.grad[i_q, i_b] += sign_f * dL_d_pos_delta + dyn_state.dofs.vel.grad[i_d, i_b] += sign_f * dL_d_jac_qvel + + +@qd.kernel(fastcache=True) +def kernel_manual_add_collision_constraints_bw( + dyn_state: array_class.DynState, + collider_state: array_class.ColliderState, + constraint_state: array_class.ConstraintState, + dyn_info: array_class.DynInfo, + rigid_info: array_class.RigidInfo, + rigid_config: qd.template(), +): + """Manual reverse of `add_collision_constraints`. See the section header + above for the shared `n_con` layout and upstream-grad conventions. + + Produces the gradients w.r.t. the collision constraint's differentiable inputs: + collider_state.contact_data.{pos, normal, penetration}.grad (-> collider.backward) + dyn_state.dofs.{cdof_ang, cdof_vel, vel}.grad + dyn_state.links.root_COM.grad + (cdof / root_COM / vel grads feed the COM / forward-dynamics reverse chain; + contact_data grads feed `collider.backward`.) + + Forward recap (per contact `i_col`, per friction-pyramid row `i` in 0..3): + d1, d2 = qd_orthogonals(normal); d = s_i * (d1 if i<2 else d2), s_i = 2*(i%2)-1 + n = d * friction - normal + jac[n_con, i_d] = sum_chain (sign * vel_motion(i_d)) . n + 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 + """ + EPS = rigid_info.EPS[None] + _B = dyn_state.dofs.ctrl_mode.shape[1] + n_dofs = dyn_state.dofs.ctrl_mode.shape[0] + max_contact_pairs = collider_state.contact_data.link_a.shape[0] + + qd.loop_config( + name="kernel_manual_add_collision_constraints_bw", + serialize=qd.static(rigid_config.para_level < gs.PARA_LEVEL.ALL), + ) + for flat_idx in range(max_contact_pairs * _B): + i_b = flat_idx % _B + i_col_ = flat_idx // _B + if i_col_ < collider_state.n_contacts[i_b]: + # The forward assembles the contact rows in logical (sorted) contact order: row group i_col_ maps to + # physical contact contact_sort_idx[i_col_] (see add_inequality_constraints). + i_col = collider_state.contact_sort_idx[i_col_, i_b] + link_a = collider_state.contact_data.link_a[i_col, i_b] + link_b = collider_state.contact_data.link_b[i_col, i_b] + contact_pos = collider_state.contact_data.pos[i_col, i_b] + normal = collider_state.contact_data.normal[i_col, i_b] + friction = collider_state.contact_data.friction[i_col, i_b] + sol_params = collider_state.contact_data.sol_params[i_col, i_b] + penetration = collider_state.contact_data.penetration[i_col, i_b] + + link_a_maybe_batch = [link_a, i_b] if qd.static(rigid_config.batch_links_info) else link_a + invweight = dyn_info.links.invweight[link_a_maybe_batch][0] + if link_b > -1: + link_b_maybe_batch = [link_b, i_b] if qd.static(rigid_config.batch_links_info) else link_b + invweight = invweight + dyn_info.links.invweight[link_b_maybe_batch][0] + + # --- forward intermediates of qd_orthogonals(normal) --- + # 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 + b_raw = gs.qd_vec3(0.0, 0.0, 0.0) + if 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) + b_raw_norm = b_raw.norm() + b = b_raw / b_raw_norm + 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 + # 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 + 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) + d_efc_D_d_imp = -d_diag_d_imp / (diag * diag) + + # Accumulators for this contact's differentiable inputs. + g_pos = gs.qd_vec3(0.0, 0.0, 0.0) + g_normal = gs.qd_vec3(0.0, 0.0, 0.0) + g_pen = gs.qd_float(0.0) + g_d1 = gs.qd_vec3(0.0, 0.0, 0.0) + g_d2 = gs.qd_vec3(0.0, 0.0, 0.0) + + # Forward row layout: equality -> frictionloss -> collision -> joint-limit. + # Offset past equality + frictionloss. + const_start = ( + constraint_state.n_constraints_equality[i_b] + constraint_state.n_constraints_frictionloss[i_b] + ) + for i in range(4): + 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 + + g_aref = constraint_state.dL_daref[n_con, i_b] + g_efc_D = constraint_state.dL_defc_D[n_con, i_b] + + # aref = -b_coef*jac_qvel + k_coef*imp*penetration (pos arg = -penetration) + d_aref_d_imp = k_coef * penetration + d_aref_d_pen_direct = k_coef * imp + d_aref_d_jac_qvel = -b_coef + + dL_d_imp = g_aref * d_aref_d_imp + g_efc_D * d_efc_D_d_imp + dL_d_pen = g_aref * d_aref_d_pen_direct + dL_d_imp * d_imp_d_imp_x * d_imp_x_d_pen + g_pen += dL_d_pen + dL_d_jac_qvel = g_aref * d_aref_d_jac_qvel + + # 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): + sign = gs.qd_float(-1.0) + link = link_a + if i_ab == 1: + 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_ + + 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) + + # 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 + + link = dyn_info.links.parent_idx[link_mb] + + # n = d*friction - normal + g_normal += -dL_dn + g_d = dL_dn * friction + if i < 2: + g_d1 += s_i * g_d + else: + g_d2 += s_i * g_d + + # Reverse qd_orthogonals: d1 = b x normal, d2 = b, b = normalize(b_raw(normal)). + dL_db = g_d2 + normal.cross(g_d1) + g_normal += g_d1.cross(b) + # 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: + # 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) + g_normal[2] += dL_db_raw[2] * (-n1) + else: + # b_raw = (-n0 n2, -n1 n2, 1 - n2^2) + g_normal[0] += dL_db_raw[0] * (-n2) + g_normal[1] += dL_db_raw[1] * (-n2) + g_normal[2] += dL_db_raw[0] * (-n0) + dL_db_raw[1] * (-n1) + dL_db_raw[2] * (-2.0 * n2) + + for j in qd.static(range(3)): + collider_state.contact_data.pos.grad[i_col, i_b][j] = g_pos[j] + collider_state.contact_data.normal.grad[i_col, i_b][j] = g_normal[j] + collider_state.contact_data.penetration.grad[i_col, i_b] = g_pen + + +@qd.kernel(fastcache=True) +def kernel_manual_add_frictionloss_constraints_bw( + dyn_state: array_class.DynState, + constraint_state: array_class.ConstraintState, + dyn_info: array_class.DynInfo, + rigid_info: array_class.RigidInfo, + rigid_config: qd.template(), +): + """Manual reverse of `add_frictionloss_constraints`. See the section header + above for the shared `n_con` layout and upstream-grad conventions. + + Accumulates into dyn_state.dofs.vel.grad[i_d] only (the conservative, kinematic- + state-only target). Model parameters (frictionloss, sol_params, invweight) + are not differentiated. + + Forward recap (per dof with `frictionloss[I_d] > EPS`, `pos_delta = 0`): + jac[n_con, i_d] = 1.0 + jac_qvel = jac * vel[i_d] = vel[i_d] + imp, aref = imp_aref(sol_params, 0, jac_qvel, 0) + diag = max(invweight * (1 - imp) / imp, EPS); efc_D = 1/diag + + Reverse: pos_delta = 0 kills both the `imp * pos_delta` term in `aref` + and the entire `imp` sensitivity to anything (imp_x = 0 => within_clamp + 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) + """ + EPS = rigid_info.EPS[None] + _B = constraint_state.jac.shape[2] + n_links = dyn_info.links.root_idx.shape[0] + + 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). + serialize=qd.static(rigid_config.para_level < gs.PARA_LEVEL.ALL and rigid_config.backend != gs.metal), + ) + for i_b in range(_B): + # Forward row layout: equality -> frictionloss -> collision -> joint-limit. + # Frictionloss row counter starts past the equality block (which may be + # nonzero when JOINT-type equalities are present). + n_con_counter = gs.qd_int(constraint_state.n_constraints_equality[i_b]) + + for i_l in range(n_links): + I_l = [i_l, i_b] if qd.static(rigid_config.batch_links_info) else i_l + for i_j in range(dyn_info.links.joint_start[I_l], dyn_info.links.joint_end[I_l]): + I_j = [i_j, i_b] if qd.static(rigid_config.batch_joints_info) else i_j + for i_d in range(dyn_info.joints.dof_start[I_j], dyn_info.joints.dof_end[I_j]): + I_d = [i_d, i_b] if qd.static(rigid_config.batch_dofs_info) else i_d + + if dyn_info.dofs.frictionloss[I_d] > EPS: + n_con = n_con_counter + n_con_counter = n_con_counter + 1 + + sol_params = dyn_info.joints.sol_params[I_j] + timeconst = sol_params[0] + dmax = sol_params[3] + b_coef = 2.0 / (dmax * timeconst) + + g_aref = constraint_state.dL_daref[n_con, i_b] + # jac = 1.0 constant => dL/d_vel = dL/d_jac_qvel. + dyn_state.dofs.vel.grad[i_d, i_b] += g_aref * (-b_coef) + + +@qd.kernel(fastcache=True) +def kernel_manual_add_equality_constraints_bw( + dyn_state: array_class.DynState, + constraint_state: array_class.ConstraintState, + dyn_info: array_class.DynInfo, + rigid_info: array_class.RigidInfo, + rigid_config: qd.template(), +): + """Manual reverse of `add_equality_constraints` (JOINT + CONNECT + WELD). + + * JOINT - couples two scalar dofs via a quartic polynomial (1 row). + * CONNECT - 3 rows pinning `global_anchor1` to `global_anchor2` in world. + * WELD - 6 rows: 3 position + 3 orientation, all sharing a single + combined `pos_imp = ||all_error||` (6D). + + Accumulates into kinematic-state grads only (conservative): + 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 + Model parameters (`sol_params`, `eq_data`, `dyn_info.dofs.invweight`, + `dyn_info.links.invweight`) are not differentiated. + + Forward recap (per equality of type JOINT): + diff = qpos[i_qpos2] - qpos0[i_qpos2] + pos_poly = a0 + a1 * diff + a2 * diff^2 + a3 * diff^3 + a4 * diff^4 + pos = qpos[i_qpos1] - qpos0[i_qpos1] - pos_poly + 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] + 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 + """ + EPS = rigid_info.EPS[None] + _B = constraint_state.jac.shape[2] + + qd.loop_config( + name="kernel_manual_add_equality_constraints_bw", + serialize=qd.static(rigid_config.para_level < gs.PARA_LEVEL.ALL), + ) + for i_b in range(_B): + # Equality is the first constraint group; row counter starts at 0. + n_con_counter = gs.qd_int(0) + + for i_e in range(constraint_state.qd_n_equalities[i_b]): + if dyn_info.equalities.eq_type[i_e, i_b] == gs.EQUALITY_TYPE.JOINT: + n_con = n_con_counter + n_con_counter = n_con_counter + 1 + + # ---- Replay forward intermediates ---- + I_joint1 = ( + [dyn_info.equalities.eq_obj1id[i_e, i_b], i_b] + if qd.static(rigid_config.batch_joints_info) + else dyn_info.equalities.eq_obj1id[i_e, i_b] + ) + I_joint2 = ( + [dyn_info.equalities.eq_obj2id[i_e, i_b], i_b] + if qd.static(rigid_config.batch_joints_info) + else dyn_info.equalities.eq_obj2id[i_e, i_b] + ) + i_qpos1 = dyn_info.joints.q_start[I_joint1] + i_qpos2 = dyn_info.joints.q_start[I_joint2] + i_dof1 = dyn_info.joints.dof_start[I_joint1] + i_dof2 = dyn_info.joints.dof_start[I_joint2] + I_dof1 = [i_dof1, i_b] if qd.static(rigid_config.batch_dofs_info) else i_dof1 + I_dof2 = [i_dof2, i_b] if qd.static(rigid_config.batch_dofs_info) else i_dof2 + + pos1 = rigid_info.qpos[i_qpos1, i_b] + pos2 = rigid_info.qpos[i_qpos2, i_b] + ref1 = rigid_info.qpos0[i_qpos1, i_b] + ref2 = rigid_info.qpos0[i_qpos2, i_b] + + a0 = dyn_info.equalities.eq_data[i_e, i_b][0] + a1 = dyn_info.equalities.eq_data[i_e, i_b][1] + a2 = dyn_info.equalities.eq_data[i_e, i_b][2] + a3 = dyn_info.equalities.eq_data[i_e, i_b][3] + a4 = dyn_info.equalities.eq_data[i_e, i_b][4] + + diff = pos2 - ref2 + diff2 = diff * diff + diff3 = diff2 * diff + diff4 = diff3 * diff + pos = pos1 - ref1 - a0 - a1 * diff - a2 * diff2 - a3 * diff3 - a4 * diff4 + deriv = a1 + 2.0 * a2 * diff + 3.0 * a3 * diff2 + 4.0 * a4 * diff3 + d_deriv_d_diff = 2.0 * a2 + 6.0 * a3 * diff + 12.0 * a4 * diff2 + + jac1 = gs.qd_float(1.0) + jac2 = -deriv + vel1 = dyn_state.dofs.vel[i_dof1, i_b] + vel2 = dyn_state.dofs.vel[i_dof2, i_b] + jac_qvel = jac1 * vel1 + jac2 * vel2 + 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] + 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) + + # ---- Upstream grads ---- + g_aref = constraint_state.dL_daref[n_con, i_b] + g_efc_D = constraint_state.dL_defc_D[n_con, i_b] + # jac[dof1] = 1.0 (constant) => no chain through dL_djac[n_con, i_dof1]. + g_jac2 = constraint_state.dL_djac[n_con, i_dof2, i_b] + + # ---- Partials ---- + # 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 + + # diag = max(diag_raw, EPS); efc_D = 1/diag + 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 + + # imp_x = |pos|/width => d_imp_x/d_pos = sign(pos) / width + sign_pos_f = gs.qd_float(1.0) + if pos < 0.0: + sign_pos_f = gs.qd_float(-1.0) + d_imp_d_pos = d_imp_d_imp_x * sign_pos_f / width + + # ---- Combine ---- + 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 + dL_d_pos = g_aref * d_aref_d_pos_direct + dL_d_imp * d_imp_d_pos + # deriv enters via (jac_qvel through jac2 = -deriv) and (jac[dof2] = -deriv). + dL_d_deriv = dL_d_jac_qvel * (-vel2) + g_jac2 * (-1.0) + + # ---- Propagate ---- + dyn_state.dofs.vel.grad[i_dof1, i_b] += dL_d_jac_qvel * jac1 + dyn_state.dofs.vel.grad[i_dof2, i_b] += dL_d_jac_qvel * jac2 + # qpos1 enters only via pos (d_pos/d_pos1 = 1). + rigid_info.qpos.grad[i_qpos1, i_b] += dL_d_pos + # qpos2 enters via pos (d_pos/d_pos2 = -deriv) and deriv (d_deriv/d_diff). + rigid_info.qpos.grad[i_qpos2, i_b] += dL_d_pos * (-deriv) + dL_d_deriv * d_deriv_d_diff + elif dyn_info.equalities.eq_type[i_e, i_b] == gs.EQUALITY_TYPE.CONNECT: + # ---------------------------------------------------------- + # 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] + # 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[n_con, i_d] += jac_i3 + # 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 + # ---------------------------------------------------------- + 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 + + anchor1_local = gs.qd_vec3( + dyn_info.equalities.eq_data[i_e, i_b][0], + dyn_info.equalities.eq_data[i_e, i_b][1], + dyn_info.equalities.eq_data[i_e, i_b][2], + ) + anchor2_local = gs.qd_vec3( + dyn_info.equalities.eq_data[i_e, i_b][3], + dyn_info.equalities.eq_data[i_e, i_b][4], + dyn_info.equalities.eq_data[i_e, i_b][5], + ) + + quat1 = dyn_state.links.quat[link1_idx, i_b] + quat2 = dyn_state.links.quat[link2_idx, i_b] + trans1 = dyn_state.links.pos[link1_idx, i_b] + trans2 = dyn_state.links.pos[link2_idx, i_b] + ga1 = gu.qd_transform_by_trans_quat(pos=anchor1_local, trans=trans1, quat=quat1) + ga2 = gu.qd_transform_by_trans_quat(pos=anchor2_local, trans=trans2, quat=quat2) + pos_diff = ga1 - ga2 + penetration = pos_diff.norm() + + invweight = dyn_info.links.invweight[link1_mb][0] + dyn_info.links.invweight[link2_mb][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] + 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) + + # All 3 rows share imp / penetration / pos_diff. Per-row partials + # only differ in which axis pos_diff[i_3] is used as ref_arg. + 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) + + 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) + + for i_3 in range(3): + n_con = n_con_counter + n_con_counter = n_con_counter + 1 + + g_aref = constraint_state.dL_daref[n_con, i_b] + g_efc_D = constraint_state.dL_defc_D[n_con, i_b] + + 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] + + 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 + + # 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 + if penetration > EPS: + coef_pen = dL_d_imp * d_imp_d_imp_x / (width * penetration) + for j in qd.static(range(3)): + g_pos_diff[j] = g_pos_diff[j] + coef_pen * pos_diff[j] + + # Walk both chains, accumulating cdof / vel / root_COM grads + # and the anchor portion that flows back through t_pos. + g_anchor1_row = gs.qd_vec3(0.0, 0.0, 0.0) + g_anchor2_row = gs.qd_vec3(0.0, 0.0, 0.0) + for i_ab in range(2): + sign = gs.qd_float(1.0) + link = link1_idx + anchor_pos = ga1 + if i_ab == 1: + sign = gs.qd_float(-1.0) + link = link2_idx + 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_ + + 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] + 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) + g_vm[i_3] = g_jac_i3 * sign + + # 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)) + # t_pos = anchor_pos - root_COM[link] + if i_ab == 0: + g_anchor1_row = g_anchor1_row + dt + else: + g_anchor2_row = g_anchor2_row + dt + dyn_state.links.root_COM.grad[link, i_b] += -dt + + link = dyn_info.links.parent_idx[link_mb] + + # 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. + 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 + dyn_state.links.pos.grad[link2_idx, i_b] += g_ga2 + dyn_state.links.quat.grad[link1_idx, i_b] += g_quat1 + dyn_state.links.quat.grad[link2_idx, i_b] += g_quat2 + else: + # ---------------------------------------------------------- + # 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] + # 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 + # 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] + # 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] + # ---------------------------------------------------------- + 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 + + # WELD eq_data layout (per forward comment): + # [0:3] anchor2 (local), [3:6] anchor1 (local), [6:10] relpose, [10] torquescale + anchor1_local = gs.qd_vec3( + dyn_info.equalities.eq_data[i_e, i_b][3], + dyn_info.equalities.eq_data[i_e, i_b][4], + dyn_info.equalities.eq_data[i_e, i_b][5], + ) + anchor2_local = gs.qd_vec3( + dyn_info.equalities.eq_data[i_e, i_b][0], + dyn_info.equalities.eq_data[i_e, i_b][1], + dyn_info.equalities.eq_data[i_e, i_b][2], + ) + relpose = qd.Vector( + [ + dyn_info.equalities.eq_data[i_e, i_b][6], + dyn_info.equalities.eq_data[i_e, i_b][7], + dyn_info.equalities.eq_data[i_e, i_b][8], + dyn_info.equalities.eq_data[i_e, i_b][9], + ], + dt=gs.qd_float, + ) + torquescale = dyn_info.equalities.eq_data[i_e, i_b][10] + + quat_body1 = dyn_state.links.quat[link1_idx, i_b] + quat_body2 = dyn_state.links.quat[link2_idx, i_b] + trans1 = dyn_state.links.pos[link1_idx, i_b] + trans2 = dyn_state.links.pos[link2_idx, i_b] + ga1 = gu.qd_transform_by_trans_quat(pos=anchor1_local, trans=trans1, quat=quat_body1) + ga2 = gu.qd_transform_by_trans_quat(pos=anchor2_local, trans=trans2, quat=quat_body2) + pos_error = ga1 - ga2 + + inv_q2 = gu.qd_inv_quat(quat_body2) + q_var = gu.qd_quat_mul(quat_body1, relpose) + error_quat = gu.qd_quat_mul(inv_q2, q_var) + rot_error = gs.qd_vec3(error_quat[1], error_quat[2], error_quat[3]) * torquescale + + # all_error = (pos_error, rot_error) ; pos_imp = ||all_error|| + pos_imp = qd.sqrt( + pos_error[0] * pos_error[0] + + pos_error[1] * pos_error[1] + + pos_error[2] * pos_error[2] + + rot_error[0] * rot_error[0] + + rot_error[1] * rot_error[1] + + 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] + + 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] + 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 + diag_pos = qd.max(diag_raw_pos, EPS) + d_diag_d_imp_pos = gs.qd_float(0.0) + if diag_raw_pos > EPS: + d_diag_d_imp_pos = -invweight_pos / (imp * imp) + d_efc_D_d_imp_pos = -d_diag_d_imp_pos / (diag_pos * diag_pos) + + diag_raw_rot = invweight_rot * (1.0 - imp) / imp + diag_rot = qd.max(diag_raw_rot, EPS) + d_diag_d_imp_rot = gs.qd_float(0.0) + if diag_raw_rot > EPS: + 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) + g_rot_error = gs.qd_vec3(0.0, 0.0, 0.0) + 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) + + # ---- Position rows (3) -- mirrors CONNECT structure ---- + n_con_orient_base = n_con_counter + 3 # rotation rows start here + for i_3 in range(3): + n_con = n_con_counter + n_con_counter = n_con_counter + 1 + + g_aref = constraint_state.dL_daref[n_con, i_b] + g_efc_D = constraint_state.dL_defc_D[n_con, i_b] + + 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] + + 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 + + # 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) + for i_ab in range(2): + sign = gs.qd_float(1.0) + link = link1_idx + anchor_pos = ga1 + if i_ab == 1: + sign = gs.qd_float(-1.0) + link = link2_idx + 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_ + + 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)) + if i_ab == 0: + g_anchor1_row = g_anchor1_row + dt + else: + g_anchor2_row = g_anchor2_row + dt + dyn_state.links.root_COM.grad[link, i_b] += -dt + + link = dyn_info.links.parent_idx[link_mb] + + # 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 + g_ga2[i_3] = g_ga2[i_3] - g_pos_error_direct + g_ga1 = g_ga1 + g_anchor1_row + g_ga2 = g_ga2 + g_anchor2_row + + # ---- Orientation rows (3) ---- + # Direct contributions: rot_error[i_3] via ref, dL_d_imp via imp. + for i_3 in range(3): + n_con = n_con_counter + n_con_counter = n_con_counter + 1 + g_aref = constraint_state.dL_daref[n_con, i_b] + g_efc_D = constraint_state.dL_defc_D[n_con, i_b] + + 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] + + 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 + + # 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) + g_q = qd.Vector([0.0, 0.0, 0.0, 0.0], dt=gs.qd_float) + for i_ab in range(2): + sign_chain = gs.qd_float(1.0) + link = link1_idx + if i_ab == 1: + sign_chain = gs.qd_float(-1.0) + 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_ + + # 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] + 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 + cdof_ang = dyn_state.dofs.cdof_ang[i_d, i_b] + jac_diff_r_d = sign_chain * cdof_ang + quat2_d = gu.qd_quat_mul_axis(inv_q2, jac_diff_r_d) + + # quat3_d = quat_mul(quat2_d, q_var) + g_quat2_d = gu.qd_quat_mul_grad_lhs(quat2_d, q_var, g_quat3_d) + g_q_contrib = gu.qd_quat_mul_grad_rhs(quat2_d, q_var, g_quat3_d) + g_q = g_q + g_q_contrib + + # quat2_d = quat_mul_axis(inv_q2, jac_diff_r_d) + # = quat_mul(inv_q2, [0, jac_diff_r_d]) + v_padded = qd.Vector( + [0.0, jac_diff_r_d[0], jac_diff_r_d[1], jac_diff_r_d[2]], + dt=gs.qd_float, + ) + g_inv_q2_contrib = gu.qd_quat_mul_grad_lhs(inv_q2, v_padded, g_quat2_d) + g_v_padded = gu.qd_quat_mul_grad_rhs(inv_q2, v_padded, g_quat2_d) + g_inv_q2 = g_inv_q2 + g_inv_q2_contrib + g_jac_diff_r_d = gs.qd_vec3(g_v_padded[1], g_v_padded[2], g_v_padded[3]) + + # 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] + + # Via-penetration contribution (shared across all 6 rows). + if pos_imp > EPS: + coef_pen = dL_d_imp_total * d_imp_d_imp_x / (width * pos_imp) + # pos_error part + for j in qd.static(range(3)): + g_pos_error_via_pen = coef_pen * pos_error[j] + g_ga1[j] = g_ga1[j] + g_pos_error_via_pen + g_ga2[j] = g_ga2[j] - g_pos_error_via_pen + # rot_error part + g_rot_error = g_rot_error + coef_pen * rot_error + + # rot_error = error_quat.xyz * torquescale + # => g_error_quat = (0, g_rot_error * torquescale) + g_error_quat = qd.Vector( + [ + 0.0, + g_rot_error[0] * torquescale, + g_rot_error[1] * torquescale, + g_rot_error[2] * torquescale, + ], + dt=gs.qd_float, + ) + g_inv_q2_eq = gu.qd_quat_mul_grad_lhs(inv_q2, q_var, g_error_quat) + g_q_eq = gu.qd_quat_mul_grad_rhs(inv_q2, q_var, g_error_quat) + g_inv_q2 = g_inv_q2 + g_inv_q2_eq + g_q = g_q + g_q_eq + + # 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( + [g_inv_q2[0], -g_inv_q2[1], -g_inv_q2[2], -g_inv_q2[3]], + dt=gs.qd_float, + ) + + # q_var = quat_mul(quat_body1, relpose); relpose const => drop g_v. + g_quat1_from_q = gu.qd_quat_mul_grad_lhs(quat_body1, relpose, g_q) + + # Anchor chain ga1, ga2 -> dyn_state.links.{pos, quat}. + g_quat1_anchor = gu.qd_transform_by_quat_grad_quat(anchor1_local, quat_body1, g_ga1) + g_quat2_anchor = gu.qd_transform_by_quat_grad_quat(anchor2_local, quat_body2, g_ga2) + + 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 diff --git a/genesis/engine/solvers/rigid/constraint/solver.py b/genesis/engine/solvers/rigid/constraint/solver.py index a2d9529d8a..a7d1ebb64f 100644 --- a/genesis/engine/solvers/rigid/constraint/solver.py +++ b/genesis/engine/solvers/rigid/constraint/solver.py @@ -448,12 +448,13 @@ def delete_weld_constraint(self, link1_idx, link2_idx, envs_idx=None): self._solver.rigid_config, ) - def backward(self, dL_dqacc): + def backward(self): if not self._solver._requires_grad: gs.raise_exception("Please set `requires_grad` to True in SimOptions to enable differentiable mode.") - # Copy upstream gradients - self.constraint_state.dL_dqacc.from_numpy(dL_dqacc) + # 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( @@ -5645,16 +5646,7 @@ def func_update_gradient_batch( ) if qd.static(rigid_config.solver_type == gs.constraint_solver.CG): - func_solve_mass_batch( - i_b, - constraint_state.grad, - constraint_state.Mgrad, - out_bw=None, - dyn_info=dyn_info, - rigid_info=rigid_info, - rigid_config=rigid_config, - is_backward=False, - ) + func_solve_mass_batch(i_b, constraint_state.grad, constraint_state.Mgrad, dyn_info, rigid_info, rigid_config) if qd.static(rigid_config.solver_type == gs.constraint_solver.Newton): if qd.static(rigid_config.enable_per_island_solve): @@ -5725,14 +5717,7 @@ def func_update_gradient_tiled( ) for i_b in range(_B): func_solve_mass_batch( - i_b, - constraint_state.grad, - constraint_state.Mgrad, - out_bw=None, - dyn_info=dyn_info, - rigid_info=rigid_info, - rigid_config=rigid_config, - is_backward=False, + i_b, constraint_state.grad, constraint_state.Mgrad, dyn_info, rigid_info, rigid_config ) if qd.static(rigid_config.solver_type == gs.constraint_solver.Newton): diff --git a/genesis/engine/solvers/rigid/rigid_solver.py b/genesis/engine/solvers/rigid/rigid_solver.py index 8701d58f36..2f9fdb0995 100644 --- a/genesis/engine/solvers/rigid/rigid_solver.py +++ b/genesis/engine/solvers/rigid/rigid_solver.py @@ -32,6 +32,14 @@ from ..kinematic_solver import KinematicSolver, _select_links_offset, _offset_world_shift, _fill_base_link_geom_offsets from .collider import Collider from .constraint import ConstraintSolver +from .constraint.backward import ( + kernel_manual_add_collision_constraints_bw, + kernel_manual_add_frictionloss_constraints_bw, + kernel_manual_add_equality_constraints_bw, + kernel_accumulate_constraint_solver_grads, + kernel_load_dL_dqacc_from_acc_grad, + kernel_manual_add_joint_limit_constraints_bw, +) from .abd.misc import ( func_add_safe_backward, func_apply_coupling_force, @@ -94,7 +102,10 @@ kernel_update_all_verts, kernel_update_geom_aabbs, kernel_update_vgeoms, + kernel_COM_links_replay, kernel_update_cartesian_space, + kernel_forward_kinematics_replay, + kernel_update_geoms_replay, ) from .abd.forward_dynamics import ( func_actuation, @@ -179,6 +190,12 @@ kernel_prepare_backward_substep, kernel_begin_backward_substep, kernel_copy_acc, + kernel_copy_next_to_curr_no_check, +) +from .abd.manual_bw import ( + kernel_manual_compute_qacc_bw, + kernel_manual_forward_kinematics_bw, + kernel_manual_forward_velocity_bw, ) if TYPE_CHECKING: @@ -1231,7 +1248,7 @@ def substep(self, f): ) if isinstance(self.sim.coupler, SAPCoupler): - update_qvel(self.dyn_state, self.rigid_info, self.rigid_config, self._is_backward) + update_qvel(self.dyn_state, self.rigid_info, self.rigid_config) else: self._func_constraint_force() kernel_step_2( @@ -1286,6 +1303,11 @@ 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() @@ -1324,6 +1346,47 @@ def _func_constraint_force(self): self.constraint_solver.add_inequality_constraints() 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) + self.constraint_solver.backward() + kernel_accumulate_constraint_solver_grads( + self.dyn_state, self.constraint_solver.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 + ) + kernel_manual_add_frictionloss_constraints_bw( + self.dyn_state, self.constraint_solver.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) + kernel_manual_add_collision_constraints_bw( + self.dyn_state, + collider_state, + self.constraint_solver.constraint_state, + self.dyn_info, + self.rigid_info, + self.rigid_config, + ) + self.collider.backward_narrowphase() + + if self._options.enable_joint_limit: + kernel_manual_add_joint_limit_constraints_bw( + self.dyn_state, + self.collider._collider_state, + self.constraint_solver.constraint_state, + self.dyn_info, + self.rigid_info, + self.rigid_config, + enable_collision=self._enable_collision, + ) + def _func_forward_dynamics(self): kernel_forward_dynamics( self.dyn_state, self.constraint_solver.constraint_state, self.dyn_info, self.rigid_info, self.rigid_config @@ -1482,6 +1545,37 @@ def reset_grad(self): qd_zero_grad(self.dyn_state_adjoint_cache.geoms) qd_zero_grad(self._rigid_adjoint_cache) + 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. + """ + # Forward replay in dependency order (FK -> COM -> geoms -> velocity). + kernel_forward_kinematics_replay( + envs_idx, self.dyn_state, self.dyn_info, self.rigid_info, self.rigid_config, is_backward=True + ) + kernel_COM_links_replay(self.dyn_state, self.dyn_info, self.rigid_info, self.rigid_config, is_backward=True) + kernel_update_geoms_replay(self.dyn_state, self.dyn_info, self.rigid_info, self.rigid_config, is_backward=True) + kernel_forward_velocity( + 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 + ) + 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 + ) + def substep_pre_coupling_grad(self, f): # Change to backward mode self._is_backward = True @@ -1498,21 +1592,13 @@ def substep_pre_coupling_grad(self, f): self.rigid_config, ) self.substep(f) - # =================== Backward substep ====================== envs_idx = self._scene._sanitize_envs_idx(None) if not self._enable_mujoco_compatibility: - kernel_forward_velocity.grad( - envs_idx, self.dyn_state, self.dyn_info, self.rigid_info, self.rigid_config, is_backward=True - ) - kernel_update_cartesian_space.grad( - self.dyn_state, - self.dyn_info, - self.rigid_info, - self.rigid_config, - force_update_fixed_geoms=False, - is_backward=True, - ) + # The FK backward below builds its Jacobian at the post-integrate qpos / vel, so copy the integrator's + # _next outputs into the current slots first. + kernel_copy_next_to_curr_no_check(self.dyn_state, self.rigid_info, self.rigid_config) + self._update_cartesian_grad(envs_idx) is_grad_valid = kernel_begin_backward_substep( f, @@ -1537,16 +1623,16 @@ def substep_pre_coupling_grad(self, f): errno=self._errno, ) - # We cannot use [kernel_forward_dynamics.grad] because we read [dofs_state.acc] and overwrite it in the kernel, - # which is prohibited (https://docs.taichi-lang.org/docs/differentiable_programming#global-data-access-rules). - # In [kernel_forward_dynamics], we read [acc] in [func_update_acc] and overwrite it in [kernel_compute_qacc]. - # As [kenrel_compute_qacc] is called at the end of [kernel_forward_dynamics], we first backpropagate through - # [kernel_compute_qacc] and then restore the original [acc] from the adjoint cache. This copy operation - # cannot be merged with [kernel_compute_qacc.grad] because .grad function itself is a standalone kernel. - # We could possibly merge this small kernel later if (1) .grad function is regarded as a function instead of a - # kernel, (2) we add another variable to store the new [acc] from [kernel_compute_qacc] and thus can avoid - # the data access violation. However, both of these require major changes. - kernel_compute_qacc.grad(self.dyn_state, self.dyn_info, self.rigid_info, self.rigid_config, is_backward=True) + # Mirror the forward branch in _func_constraint_force: + # (A) _disable_constraint=True: the forward never calls the constraint solver; acc is the smooth-dynamics + # result. Reverse via kernel_manual_compute_qacc_bw (implicit function theorem through M). + # (B) _disable_constraint=False: the forward always calls constraint_solver.resolve. Reverse via + # _constraint_force_grad. + 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) kernel_forward_dynamics_without_qacc.grad( @@ -1560,17 +1646,7 @@ def substep_pre_coupling_grad(self, f): # If it was the very first substep, we need to backpropagate through the initial update of the cartesian space if self._enable_mujoco_compatibility or self._sim.cur_substep_global == 0: - kernel_forward_velocity.grad( - envs_idx, self.dyn_state, self.dyn_info, self.rigid_info, self.rigid_config, is_backward=True - ) - kernel_update_cartesian_space.grad( - self.dyn_state, - self.dyn_info, - self.rigid_info, - self.rigid_config, - force_update_fixed_geoms=False, - is_backward=True, - ) + self._update_cartesian_grad(envs_idx) # Change back to forward mode self._is_backward = False @@ -1582,7 +1658,7 @@ def substep_post_coupling(self, f): return if isinstance(self.sim.coupler, SAPCoupler): - update_qacc_from_qvel_delta(self.dyn_state, self.rigid_info, self.rigid_config, self._is_backward) + update_qacc_from_qvel_delta(self.dyn_state, self.rigid_info, self.rigid_config) kernel_step_2( self.dyn_state, self.collider._collider_state, @@ -1821,9 +1897,11 @@ def save_ckpt(self, ckpt_name): if ckpt_name not in self._ckpt: self._ckpt[ckpt_name] = dict() - self._ckpt[ckpt_name]["qpos"] = qd_to_numpy(self._rigid_adjoint_cache.qpos) - self._ckpt[ckpt_name]["dofs_vel"] = qd_to_numpy(self._rigid_adjoint_cache.dofs_vel) - self._ckpt[ckpt_name]["dofs_acc"] = qd_to_numpy(self._rigid_adjoint_cache.dofs_acc) + # 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) for entity in self._entities: entity.save_ckpt(ckpt_name) diff --git a/genesis/utils/array_class.py b/genesis/utils/array_class.py index 87bcba49da..19a206257e 100644 --- a/genesis/utils/array_class.py +++ b/genesis/utils/array_class.py @@ -87,6 +87,7 @@ class ErrorCode(IntEnum): INVALID_FORCE_NAN = 0b00000000000000000000000000001000 INVALID_ACC_NAN = 0b00000000000000000000000000010000 OVERFLOW_CONTACTS = 0b00000000000000000000000000100000 + MANUAL_BW_UNIMPLEMENTED = 0b00000000000000000000000001000000 # =========================================== RigidInfo =========================================== @@ -109,7 +110,6 @@ class RigidInfo: geoms_init_AABB: qd.Tensor mass_mat: qd.Tensor mass_mat_L: qd.Tensor - mass_mat_L_bw: qd.Tensor mass_mat_D_inv: qd.Tensor mass_mat_tiled_scratch: qd.Tensor mass_mat_mask: qd.Tensor @@ -154,11 +154,6 @@ def get_rigid_info(solver, kinematic_only): f"Mass matrix shape (n_dofs={solver.n_dofs_}, n_dofs={solver.n_dofs_}, n_envs={_B}) is too large." ) requires_grad = solver._requires_grad - mass_mat_shape_bw = maybe_shape((2, *mass_mat_shape), requires_grad) - if math.prod(mass_mat_shape_bw) > np.iinfo(np.int32).max: - gs.raise_exception( - f"Mass matrix buffer shape (2, n_dofs={solver.n_dofs_}, n_dofs={solver.n_dofs_}, n_envs={_B}) is too large." - ) # Batch-first scratch for the register-tiled mass factor (qd.simt tile ops are batch-first, so the factorization # cannot run in place on the batch-last mass_mat_L). Allocated only when that path is enabled, with the constraint @@ -199,7 +194,6 @@ def get_rigid_info(solver, kinematic_only): geoms_init_AABB=V_VEC(3, dtype=gs.qd_float, shape=()), mass_mat=V(dtype=gs.qd_float, shape=()), mass_mat_L=V(dtype=gs.qd_float, shape=()), - mass_mat_L_bw=V(dtype=gs.qd_float, shape=()), mass_mat_D_inv=V(dtype=gs.qd_float, shape=()), mass_mat_tiled_scratch=V(dtype=gs.qd_float, shape=()), mass_mat_mask=V(dtype=gs.qd_bool, shape=()), @@ -240,7 +234,6 @@ def get_rigid_info(solver, kinematic_only): geoms_init_AABB=V_VEC(3, dtype=gs.qd_float, shape=(solver.n_geoms_, 8)), mass_mat=V(dtype=gs.qd_float, shape=mass_mat_shape, layout=mass_mat_layout, needs_grad=requires_grad), mass_mat_L=V(dtype=gs.qd_float, shape=mass_mat_shape, needs_grad=requires_grad), - mass_mat_L_bw=V(dtype=gs.qd_float, shape=mass_mat_shape_bw, needs_grad=requires_grad), mass_mat_D_inv=V(dtype=gs.qd_float, shape=(solver.n_dofs_, _B), needs_grad=requires_grad), mass_mat_tiled_scratch=V(dtype=gs.qd_float, shape=mass_mat_tiled_scratch_shape), mass_mat_mask=V(dtype=gs.qd_bool, shape=(solver.n_entities_, _B)), @@ -749,6 +742,9 @@ class DiffContactInput: # Local positions of the 1 vertex from the two geometries that define the support point for the face above w_local_pos1: qd.Tensor w_local_pos2: qd.Tensor + # Plane-convex contacts only: the convex support "core" (box vertex / sphere center / capsule nearest endpoint) + # in the convex geom's local frame. + core_local: qd.Tensor # Reference id of the contact point, which is needed for the backward pass ref_id: qd.Tensor # Flag whether the contact data can be computed in numerically stable way in both the forward and backward passes @@ -771,6 +767,7 @@ def get_diff_contact_input(_B, max_contacts_per_pair, is_active, requires_grad=F local_pos2_c=V_VEC(3, dtype=gs.qd_float, shape=shape), w_local_pos1=V_VEC(3, dtype=gs.qd_float, shape=shape), w_local_pos2=V_VEC(3, dtype=gs.qd_float, shape=shape), + core_local=V_VEC(3, dtype=gs.qd_float, shape=shape), ref_id=V(dtype=gs.qd_int, shape=shape), valid=V(dtype=gs.qd_int, shape=shape), ref_penetration=V(dtype=gs.qd_float, shape=shape, needs_grad=True), @@ -1678,7 +1675,6 @@ class DofsState: vel_prev: qd.Tensor vel_next: qd.Tensor acc: qd.Tensor - acc_bw: qd.Tensor acc_smooth: qd.Tensor acc_smooth_bw: qd.Tensor qf_smooth: qd.Tensor @@ -1715,7 +1711,6 @@ def get_dofs_state(solver): vel_prev=V(dtype=gs.qd_float, shape=shape, needs_grad=requires_grad), vel_next=V(dtype=gs.qd_float, shape=shape, needs_grad=requires_grad), acc=V(dtype=gs.qd_float, shape=shape, needs_grad=requires_grad), - acc_bw=V(dtype=gs.qd_float, shape=shape_bw, needs_grad=requires_grad), acc_smooth=V(dtype=gs.qd_float, shape=shape, needs_grad=requires_grad), acc_smooth_bw=V(dtype=gs.qd_float, shape=shape_bw, needs_grad=requires_grad), qf_smooth=V(dtype=gs.qd_float, shape=shape, needs_grad=requires_grad), diff --git a/genesis/utils/geom.py b/genesis/utils/geom.py index 5be3b6730c..edb4464910 100644 --- a/genesis/utils/geom.py +++ b/genesis/utils/geom.py @@ -133,6 +133,28 @@ def qd_rotvec_to_quat(rotvec, eps): return quat +@qd.func +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: + 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] + """ + thetasq = rotvec.dot(rotvec) + theta_reg = qd.sqrt(thetasq + eps * eps) + theta_half = 0.5 * theta_reg + sin_half = qd.sin(theta_half) + cos_half = qd.cos(theta_half) + sinc = sin_half / theta_reg + d_sinc_d_theta = (0.5 * cos_half - sinc) / theta_reg + + out_grad_vec = gs.qd_vec3(out_grad[1], out_grad[2], out_grad[3]) + coeff = -0.5 * sin_half / theta_reg * out_grad[0] + d_sinc_d_theta / theta_reg * out_grad_vec.dot(rotvec) + return coeff * rotvec + sinc * out_grad_vec + + @qd.func def qd_quat_to_R(quat, eps): """ @@ -242,6 +264,34 @@ def qd_quat_mul(u, v): return qd.Vector([w, x, y, z], dt=gs.qd_float) +@qd.func +def qd_quat_mul_grad_lhs(u, v, out_grad): + """Adjoint of qd_quat_mul(u, v) with respect to u, given the upstream gradient of its output.""" + return qd.Vector( + [ + out_grad[0] * v[0] + out_grad[1] * v[1] + out_grad[2] * v[2] + out_grad[3] * v[3], + -out_grad[0] * v[1] + out_grad[1] * v[0] - out_grad[2] * v[3] + out_grad[3] * v[2], + -out_grad[0] * v[2] + out_grad[1] * v[3] + out_grad[2] * v[0] - out_grad[3] * v[1], + -out_grad[0] * v[3] - out_grad[1] * v[2] + out_grad[2] * v[1] + out_grad[3] * v[0], + ], + dt=gs.qd_float, + ) + + +@qd.func +def qd_quat_mul_grad_rhs(u, v, out_grad): + """Adjoint of qd_quat_mul(u, v) with respect to v, given the upstream gradient of its output.""" + return qd.Vector( + [ + out_grad[0] * u[0] + out_grad[1] * u[1] + out_grad[2] * u[2] + out_grad[3] * u[3], + -out_grad[0] * u[1] + out_grad[1] * u[0] + out_grad[2] * u[3] - out_grad[3] * u[2], + -out_grad[0] * u[2] - out_grad[1] * u[3] + out_grad[2] * u[0] + out_grad[3] * u[1], + -out_grad[0] * u[3] + out_grad[1] * u[2] - out_grad[2] * u[1] + out_grad[3] * u[0], + ], + dt=gs.qd_float, + ) + + @qd.func def qd_transform_quat_by_quat(v, u): """Transforms quat_v by quat_u. @@ -270,6 +320,53 @@ def qd_transform_by_quat(v, quat): ) / (q_ww + q_xx + q_yy + q_zz) +@qd.func +def qd_transform_by_quat_grad_quat(v, quat, out_grad): + """Adjoint of qd_transform_by_quat(v, quat) with respect to quat, with v held constant. + + The forward is num(q) / |q|^2 where num is the quaternion sandwich q v q*. The full derivative applies the + quotient rule: d(num/D) = dnum/D - num * dD/D^2 with D = |q|^2 and dD/dq = 2 q. Keeping only the numerator term + (dnum) is correct only when q is a fixed unit quaternion; it is wrong whenever q is an optimization variable, + even at unit length (the radial dD term is nonzero), so the denominator term must be included. + """ + q_w, q_x, q_y, q_z = quat + v_x, v_y, v_z = v + + d_out0_d_quat = 2.0 * qd.Vector( + [ + q_w * v_x - q_z * v_y + q_y * v_z, + q_x * v_x + q_y * v_y + q_z * v_z, + -q_y * v_x + q_x * v_y + q_w * v_z, + -q_z * v_x - q_w * v_y + q_x * v_z, + ], + dt=gs.qd_float, + ) + d_out1_d_quat = 2.0 * qd.Vector( + [ + q_z * v_x + q_w * v_y - q_x * v_z, + q_y * v_x - q_x * v_y - q_w * v_z, + q_x * v_x + q_y * v_y + q_z * v_z, + q_w * v_x - q_z * v_y + q_y * v_z, + ], + dt=gs.qd_float, + ) + d_out2_d_quat = 2.0 * qd.Vector( + [ + -q_y * v_x + q_x * v_y + q_w * v_z, + q_z * v_x + q_w * v_y - q_x * v_z, + -q_w * v_x + q_z * v_y - q_y * v_z, + q_x * v_x + q_y * v_y + q_z * v_z, + ], + dt=gs.qd_float, + ) + d_num_d_quat = out_grad[0] * d_out0_d_quat + out_grad[1] * d_out1_d_quat + out_grad[2] * d_out2_d_quat + + # Quotient-rule denominator term: out = num / D with D = |q|^2, so d(out)/dq = dnum/dq / D - out * (2 q) / D. + D = q_w * q_w + q_x * q_x + q_y * q_y + q_z * q_z + out = qd_transform_by_quat(v, quat) + return d_num_d_quat / D - (2.0 / D) * out_grad.dot(out) * quat + + @qd.func def qd_inv_transform_by_quat(v, quat): return qd_transform_by_quat(v, qd_inv_quat(quat)) @@ -383,6 +480,20 @@ def motion_cross_motion(s_ang, s_vel, m_ang, m_vel): return ang, vel +@qd.func +def motion_cross_motion_grad(s_ang, s_vel, m_ang, m_vel, ang_grad, vel_grad): + """Adjoint of motion_cross_motion, returning the additive gradient deltas (s_ang, s_vel, m_ang, m_vel). + + Uses the cross-product adjoint of f = a x b: a.grad += b x f.grad, b.grad += f.grad x a. + """ + return ( + m_ang.cross(ang_grad) + m_vel.cross(vel_grad), + m_ang.cross(vel_grad), + ang_grad.cross(s_ang) + vel_grad.cross(s_vel), + vel_grad.cross(s_ang), + ) + + @qd.func def qd_orthogonals(a): """ diff --git a/genesis/utils/misc.py b/genesis/utils/misc.py index 7a744fb82b..b2db8f2d86 100644 --- a/genesis/utils/misc.py +++ b/genesis/utils/misc.py @@ -702,7 +702,7 @@ def qd_to_torch( except AttributeError: try: tc = value.to_torch(copy=False) - except (ValueError, RuntimeError): + except (ValueError, RuntimeError, TypeError): if copy is False: raise tensor = _maybe_transpose(value.to_torch(), value, transpose) diff --git a/tests/conftest.py b/tests/conftest.py index c55cb07b20..8c62bb32d6 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -100,6 +100,7 @@ def _skip_reason(reason): SKIP_NO_LUISA = _skip_reason("RayTracer is not supported because 'LuisaRenderPy' is not available.") SKIP_NO_VIEWER = _skip_reason("Interactive viewer not supported on this platform.") SKIP_NO_OMNIVERSE_KIT = _skip_reason("omniverse-kit support not available") +SKIP_METAL_GRAD = _skip_reason("Apple Metal GPU computes wrong reverse-mode gradients (Quadrants backward bug).") def is_mem_monitoring_supported(): diff --git a/tests/core/test_grad.py b/tests/core/test_grad.py deleted file mode 100644 index 9c45c73350..0000000000 --- a/tests/core/test_grad.py +++ /dev/null @@ -1,629 +0,0 @@ -import numpy as np -import pytest -import torch - -import genesis as gs -from genesis.utils.geom import R_to_quat -from genesis.utils.misc import qd_to_torch, qd_to_numpy, tensor_to_array -from genesis.utils import set_random_seed - -from ..utils import assert_allclose - - -@pytest.mark.slow # ~350s -@pytest.mark.required -@pytest.mark.parametrize("backend", [gs.cpu, gs.gpu]) -def test_differentiable_push(show_viewer): - HORIZON = 10 - - scene = gs.Scene( - sim_options=gs.options.SimOptions( - dt=2e-3, - substeps=10, - requires_grad=True, - ), - mpm_options=gs.options.MPMOptions( - lower_bound=(0.0, -1.0, 0.0), - upper_bound=(1.0, 1.0, 0.55), - ), - viewer_options=gs.options.ViewerOptions( - camera_pos=(2.5, -0.15, 2.42), - camera_lookat=(0.5, 0.5, 0.1), - ), - show_viewer=show_viewer, - ) - - plane = scene.add_entity( - gs.morphs.URDF( - file="urdf/plane/plane.urdf", - fixed=True, - ) - ) - stick = scene.add_entity( - morph=gs.morphs.Mesh( - file="meshes/stirrer.obj", - scale=0.6, - pos=(0.5, 0.5, 0.05), - euler=(90.0, 0.0, 0.0), - ), - material=gs.materials.Tool( - friction=8.0, - ), - ) - obj = scene.add_entity( - morph=gs.morphs.Box( - lower=(0.2, 0.1, 0.05), - upper=(0.4, 0.3, 0.15), - ), - material=gs.materials.MPM.Elastic( - rho=500, - ), - ) - scene.build(n_envs=2) - - init_pos = gs.tensor([[0.3, 0.1, 0.28], [0.3, 0.1, 0.5]], requires_grad=True) - stick.set_position(init_pos) - pos_obj_init = gs.tensor([0.3, 0.3, 0.1], requires_grad=True) - obj.set_position(pos_obj_init) - v_obj_init = gs.tensor([0.0, -1.0, 0.0], requires_grad=True) - obj.set_velocity(v_obj_init) - goal = gs.tensor([0.5, 0.8, 0.05]) - - loss = 0.0 - v_list = [] - for i in range(HORIZON): - v_i = gs.tensor([[0.0, 1.0, 0.0], [0.0, 1.0, 0.0]], requires_grad=True) - stick.set_velocity(vel=v_i) - v_list.append(v_i) - - scene.step() - - if i == HORIZON // 2: - mpm_particles = scene.get_state().solvers_state[scene.solvers.index(scene.mpm_solver)] - loss += torch.pow(mpm_particles.pos[mpm_particles.active == 1] - goal, 2).sum() - - if i == HORIZON - 2: - state = obj.get_state() - loss += torch.pow(state.pos - goal, 2).sum() - loss.backward() - - # TODO: It would be great to compare the gradient to its analytical or numerical value. - for v_i in v_list[:-1]: - assert (v_i.grad.abs() > gs.EPS).any() - assert (v_list[-1].grad.abs() < gs.EPS).all() - - -@pytest.mark.required -@pytest.mark.precision("64") -@pytest.mark.parametrize("backend", [gs.cpu, gs.gpu]) -def test_diff_contact(): - RTOL = 1e-4 - - scene = gs.Scene( - sim_options=gs.options.SimOptions( - dt=0.01, - # Turn on differentiable mode - requires_grad=True, - ), - show_viewer=False, - ) - - box_size = 0.25 - box_spacing = box_size - 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), - ) - box1 = scene.add_entity( - gs.morphs.Box(size=box_size * vec_one, pos=box_pos_offset + 0.8 * box_spacing * np.array([0, 0, 1])), - ) - scene.build() - solver = scene.sim.rigid_solver - collider = solver.collider - - # Set up initial configuration - x_ang, y_ang, z_ang = 3.0, 3.0, 3.0 - box1.set_quat(R_to_quat(gs.euler_to_R([np.deg2rad(x_ang), np.deg2rad(y_ang), np.deg2rad(z_ang)]))) - - box0_init_pos = box0.get_pos().clone() - box1_init_pos = box1.get_pos().clone() - box0_init_quat = box0.get_quat().clone() - box1_init_quat = box1.get_quat().clone() - - ### Compute the initial loss and compute gradients using differentiable contact detection - # Detect contact - collider.detection() - - # Get contact outputs and their grads - contacts = collider.get_contacts(as_tensor=True, to_torch=True, keep_batch_dim=True) - normal = contacts["normal"].requires_grad_() - position = contacts["position"].requires_grad_() - penetration = contacts["penetration"].requires_grad_() - - loss = ((normal * position).sum(dim=-1) * penetration).sum() - dL_dnormal = torch.autograd.grad(loss, normal, retain_graph=True)[0] - dL_dposition = torch.autograd.grad(loss, position, retain_graph=True)[0] - dL_dpenetration = torch.autograd.grad(loss, penetration)[0] - - # Compute analytical gradients of the geoms position and quaternion - collider.backward(dL_dposition, dL_dnormal, dL_dpenetration) - dL_dpos = qd_to_torch(solver.dyn_state.geoms.pos.grad) - dL_dquat = qd_to_torch(solver.dyn_state.geoms.quat.grad) - - ### Compute directional derivatives along random directions - FD_EPS = 1e-5 - TRIALS = 100 - - def compute_dL_error(dL_dx, x_type): - dL_error_rel = 0.0 - - box0_input_pos = box0_init_pos - box1_input_pos = box1_init_pos - box0_input_quat = box0_init_quat - box1_input_quat = box1_init_quat - - for _ in range(TRIALS): - rand_dx = torch.randn_like(dL_dx) - rand_dx = torch.nn.functional.normalize(rand_dx, dim=-1) - - dL = (rand_dx * dL_dx).sum() - - lossPs = [] - for sign in (1, -1): - # Compute query point - if x_type == "pos": - box0_input_pos = box0_init_pos + sign * rand_dx[0, 0] * FD_EPS - box1_input_pos = box1_init_pos + sign * rand_dx[1, 0] * FD_EPS - else: - # FIXME: The quaternion should be normalized - box0_input_quat = box0_init_quat + sign * rand_dx[0, 0] * FD_EPS - box1_input_quat = box1_init_quat + sign * rand_dx[1, 0] * FD_EPS - - # Update box positions - box0.set_pos(box0_input_pos) - box1.set_pos(box1_input_pos) - box0.set_quat(box0_input_quat) - box1.set_quat(box1_input_quat) - - # Re-detect contact. - # We need to manually reset the contact counter as we are not running the whole sim step. - collider._collider_state.n_contacts.fill(0) - collider.detection() - contacts = collider.get_contacts(as_tensor=True, to_torch=True, keep_batch_dim=True) - normal, position, penetration = contacts["normal"], contacts["position"], contacts["penetration"] - - # Compute loss - loss = ((normal * position).sum(dim=-1) * penetration).sum() - lossPs.append(loss) - - dL_fd = (lossPs[0] - lossPs[1]) / (2 * FD_EPS) - dL_error_rel += (dL - dL_fd).abs() / max(dL.abs(), dL_fd.abs(), gs.EPS) - - dL_error_rel /= TRIALS - return dL_error_rel - - dL_dpos_error_rel = compute_dL_error(dL_dpos, "pos") - assert_allclose(dL_dpos_error_rel, 0.0, atol=RTOL) - dL_dquat_error_rel = compute_dL_error(dL_dquat, "quat") - assert_allclose(dL_dquat_error_rel, 0.0, atol=RTOL) - - -@pytest.mark.required -def test_diff_convex_contact_forward(show_viewer): - # The split narrowphase (GPU-only) used to skip GJK entirely when requires_grad is True, so convex-convex - # contacts were never detected and bodies fell through each other. Differentiable contact detection must - # route through the monolithic diff_gjk path, which produces the same forward contacts as the non-grad path. - scene = gs.Scene( - sim_options=gs.options.SimOptions( - requires_grad=True, - ), - rigid_options=gs.options.RigidOptions( - integrator=gs.integrator.approximate_implicitfast, - # Keep both boxes on the general convex-convex GJK path rather than the specialized box-box detector. - box_box_detection=False, - ), - show_viewer=show_viewer, - ) - - scene.add_entity( - gs.morphs.Plane(), - ) - # Two independent stacks whose x-order (-0.8 < +0.8) is the reverse of their geom-pair detection order (the +0.8 - # stack is added first). On GPU the x-position spatial sort would permute contact_sort_idx into a non-identity - # order; combined with the autodiff backward writing gradients by physical index, that attaches gradients to the - # wrong contacts. The sort must therefore be disabled in autodiff mode, leaving contact_sort_idx the identity. - 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.build() - - for _ in range(20): - scene.step() - - # Each top box rests on its fixed box (top face at z=0.4, half-height 0.2 -> center at 0.6) and never tunnels - # through it. Without contact detection the boxes free-fall to large negative z. - for top, x in zip(tops, (0.8, -0.8)): - assert_allclose(top.get_pos(), (x, 0.0, 0.6), atol=2e-4) - assert_allclose(top.get_dofs_velocity(), 0.0, atol=0.05) - - # In autodiff mode the contact permutation must stay the identity: collider.backward writes upstream gradients - # back by physical contact index, while get_contacts returns them in contact_sort_idx (logical) order. - collider = scene.sim.rigid_solver.collider - assert not collider._collider_static_config.spatial_sort_supported - n_contacts = int(np.atleast_1d(qd_to_numpy(collider._collider_state.n_contacts))[0]) - sort_idx = qd_to_numpy(collider._collider_state.contact_sort_idx)[:n_contacts, 0] - assert_allclose(sort_idx, np.arange(n_contacts), atol=0) - - -@pytest.mark.required -def test_diff_smooth_pair_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), - ), - ) - - # A sphere/ellipsoid pair has an everywhere-curved Minkowski boundary on which diff_gjk's EPA never converges, - # so it would silently tunnel. - with pytest.raises(gs.GenesisException): - scene.build() - - -# We need to use 64-bit precision for this test because we need to use sufficiently small perturbation to get reliable -# gradient estimates through finite difference method. This small perturbation is not supported by 32-bit precision in -# stable way. -@pytest.mark.required -@pytest.mark.precision("64") -@pytest.mark.parametrize("backend", [gs.cpu, gs.gpu]) -def test_diff_solver(monkeypatch): - from genesis.engine.solvers.rigid.constraint.solver import func_solve_init, func_solve_body - from genesis.engine.solvers.rigid.rigid_solver import kernel_step_1 - - RTOL = 1e-4 - - scene = gs.Scene( - sim_options=gs.options.SimOptions( - dt=0.01, - requires_grad=True, - ), - rigid_options=gs.options.RigidOptions( - # We use Newton's method because it converges faster than CG, and therefore gives better gradient estimation - # when using finite difference method - constraint_solver=gs.constraint_solver.Newton, - ), - 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.build() - rigid_solver = scene._sim.rigid_solver - constraint_solver = rigid_solver.constraint_solver - - franka.set_qpos([-1.0124, 1.5559, 1.3662, -1.6878, -1.5799, 1.7757, 1.4602, 0.04, 0.04]) - - # Monkeypatch the constraint resolve function to avoid overwriting the necessary information for computing gradients. - def constraint_solver_resolve(): - func_solve_init( - rigid_solver.dyn_state, - constraint_solver.constraint_state, - rigid_solver.dyn_info, - rigid_solver.rigid_info, - rigid_solver.rigid_config, - is_decomposed=False, - ) - func_solve_body( - rigid_solver.dyn_state, - constraint_solver.constraint_state, - rigid_solver.dyn_info, - rigid_solver.rigid_info, - rigid_solver.rigid_config, - constraint_solver._n_iterations, - ) - - monkeypatch.setattr(constraint_solver, "resolve", constraint_solver_resolve) - - # Step once to compute constraint solver's inputs: [mass], [jac], [aref], [efc_D], [force]. We do not call the - # entire scene.step() because it will overwrite the necessary information that we need to compute the gradients. - kernel_step_1( - rigid_solver.dyn_state, - constraint_solver.constraint_state, - rigid_solver.dyn_info, - rigid_solver.rigid_info, - rigid_solver.rigid_config, - is_forward_pos_updated=True, - is_forward_vel_updated=True, - is_backward=False, - ) - constraint_solver.add_equality_constraints() - rigid_solver.collider.detection() - constraint_solver.add_inequality_constraints() - constraint_solver.resolve() - - # Loss function to compute gradients using finite difference method - def compute_loss(input_mass, input_jac, input_aref, input_efc_D, input_force): - rigid_solver.rigid_info.mass_mat.from_numpy(input_mass) - constraint_solver.constraint_state.jac.from_numpy(input_jac) - constraint_solver.constraint_state.aref.from_numpy(input_aref) - constraint_solver.constraint_state.efc_D.from_numpy(input_efc_D) - rigid_solver.dyn_state.dofs.force.from_numpy(input_force) - - # Recompute acc_smooth from the updated input variables - updated_acc_smooth = np.linalg.solve(input_mass[..., 0], input_force[..., 0]) - rigid_solver.dyn_state.dofs.acc_smooth.from_numpy(updated_acc_smooth[..., None]) - constraint_solver.resolve() - - output_qacc = qd_to_torch(constraint_solver.qacc) - return ((output_qacc - target_qacc) ** 2).mean() - - init_input_mass = qd_to_numpy(rigid_solver.rigid_info.mass_mat, copy=True) - init_input_jac = qd_to_numpy(constraint_solver.constraint_state.jac, copy=True) - init_input_aref = qd_to_numpy(constraint_solver.constraint_state.aref, copy=True) - init_input_efc_D = qd_to_numpy(constraint_solver.constraint_state.efc_D, copy=True) - init_input_force = qd_to_numpy(rigid_solver.dyn_state.dofs.force, copy=True) - - # Initial output of the constraint solver - set_random_seed(0) - init_output_qacc = qd_to_torch(constraint_solver.qacc) - target_qacc = torch.from_numpy(np.random.randn(*init_output_qacc.shape)).to(device=gs.device) - target_qacc = target_qacc * init_output_qacc.abs().mean() - - # Solve the constraint solver and get the output - output_qacc = qd_to_torch(constraint_solver.qacc, copy=True).requires_grad_(True) - - # Compute loss and gradient of the output - loss = ((output_qacc - target_qacc) ** 2).mean() - dL_dqacc = tensor_to_array(torch.autograd.grad(loss, output_qacc)[0]) - - # Compute gradients of the input variables: [mass], [jac], [aref], [efc_D], [force] - constraint_solver.backward(dL_dqacc) - - # Fetch gradients of the input variables - dL_dM = qd_to_numpy(constraint_solver.constraint_state.dL_dM) - dL_djac = qd_to_numpy(constraint_solver.constraint_state.dL_djac) - dL_daref = qd_to_numpy(constraint_solver.constraint_state.dL_daref) - dL_defc_D = qd_to_numpy(constraint_solver.constraint_state.dL_defc_D) - dL_dforce = qd_to_numpy(constraint_solver.constraint_state.dL_dforce) - - ### Compute directional derivatives along random directions - FD_EPS = 1e-3 - TRIALS = 200 - - for dL_dx, x_type in ( - (dL_dforce, "force"), - (dL_daref, "aref"), - (dL_defc_D, "efc_D"), - (dL_djac, "jac"), - (dL_dM, "mass"), - ): - dL_error = 0.0 - for _ in range(TRIALS): - rand_dx = np.random.randn(*dL_dx.shape) - rand_dx = rand_dx / max( - np.linalg.norm(rand_dx, axis=0 if x_type in ("force", "aref", "efc_D") else (0, 1)), gs.EPS - ) - if x_type == "mass": - # Make rand_dx symmetric - rand_dx = (rand_dx + np.moveaxis(rand_dx, 0, 1)) * 0.5 - - dL = (rand_dx * dL_dx).sum() - - input_force = init_input_force - input_aref = init_input_aref - input_efc_D = init_input_efc_D - input_jac = init_input_jac - input_mass = init_input_mass - - # 1 * eps - if x_type == "force": - input_force = init_input_force + rand_dx * FD_EPS - elif x_type == "aref": - input_aref = init_input_aref + rand_dx * FD_EPS - elif x_type == "efc_D": - input_efc_D = init_input_efc_D + rand_dx * FD_EPS - elif x_type == "jac": - input_jac = init_input_jac + rand_dx * FD_EPS - elif x_type == "mass": - input_mass = init_input_mass + rand_dx * FD_EPS - lossP1 = compute_loss(input_mass, input_jac, input_aref, input_efc_D, input_force) - - # -1 * eps - if x_type == "force": - input_force = init_input_force - rand_dx * FD_EPS - elif x_type == "aref": - input_aref = init_input_aref - rand_dx * FD_EPS - elif x_type == "efc_D": - input_efc_D = init_input_efc_D - rand_dx * FD_EPS - elif x_type == "jac": - input_jac = init_input_jac - rand_dx * FD_EPS - elif x_type == "mass": - input_mass = init_input_mass - rand_dx * FD_EPS - - lossP2 = compute_loss(input_mass, input_jac, input_aref, input_efc_D, input_force) - dL_fd = (lossP1 - lossP2) / (2 * FD_EPS) - - dL_error += (dL - dL_fd).abs() / max(abs(dL), abs(dL_fd), gs.EPS) - - dL_error /= TRIALS - assert_allclose(dL_error, 0.0, atol=RTOL) - - -@pytest.mark.slow # ~250s -@pytest.mark.required -@pytest.mark.parametrize("backend", [gs.cpu, gs.gpu]) -def test_differentiable_rigid(show_viewer): - dt = 1e-2 - horizon = 100 - substeps = 1 - 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=dt, - substeps=substeps, - 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), - camera_lookat=(0.5, 0.5, 0.1), - ), - show_viewer=show_viewer, - ) - - box = scene.add_entity( - gs.morphs.Box( - pos=(0, 0, 0), - size=(0.1, 0.1, 0.2), - ), - surface=gs.surfaces.Default( - color=(0.9, 0.0, 0.0, 1.0), - ), - ) - if show_viewer: - target = scene.add_entity( - gs.morphs.Box( - pos=goal_pos, - quat=goal_quat, - size=(0.1, 0.1, 0.2), - ), - surface=gs.surfaces.Default( - color=(0.0, 0.9, 0.0, 0.5), - ), - ) - - scene.build() - - num_iter = 200 - lr = 1e-2 - - init_pos = gs.tensor([0.3, 0.1, 0.28], requires_grad=True) - init_quat = gs.tensor([1.0, 0.0, 0.0, 0.0], requires_grad=True) - optimizer = torch.optim.Adam([init_pos, init_quat], lr=lr) - - scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_iter, eta_min=1e-3) - - for _ in range(num_iter): - scene.reset() - - box.set_pos(init_pos) - box.set_quat(init_quat) - - loss = 0 - for _ in range(horizon): - scene.step() - if show_viewer: - target.set_pos(goal_pos) - target.set_quat(goal_quat) - - box_state = box.get_state() - box_pos = box_state.pos - box_quat = box_state.quat - loss = torch.abs(box_pos - goal_pos).sum() + torch.abs(box_quat - goal_quat).sum() - - optimizer.zero_grad() - loss.backward() # this lets gradient flow all the way back to tensor input - optimizer.step() - scheduler.step() - - with torch.no_grad(): - init_quat.data = init_quat / torch.norm(init_quat, dim=-1, keepdim=True) - - assert_allclose(loss, 0.0, atol=1e-2) - - -@pytest.mark.slow # ~200s -@pytest.mark.required -@pytest.mark.parametrize("backend", [gs.cpu, gs.gpu]) -def test_diff_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, - ), - show_viewer=show_viewer, - ) - robot = scene.add_entity( - gs.morphs.Box( - size=(0.1, 0.1, 0.1), - pos=(0, 0, 0), - ) - ) - scene.build() - - ctrl = gs.tensor(np.random.randn(robot.n_dofs), dtype=gs.tc_float, requires_grad=True) - - grads = [] - for use_sim_state in (False, True): - scene.reset() - - robot.set_dofs_velocity(ctrl) - scene.step() - - if use_sim_state: - solver_state = scene.get_state().solvers_state[scene.solvers.index(scene.rigid_solver)] - chassis_pos = solver_state.links_pos[:, 0].squeeze() - else: - chassis_pos = robot.get_state().pos.squeeze() - - loss = torch.linalg.norm(chassis_pos) - loss.backward() - grad = ctrl.grad.detach().clone() - ctrl.grad.zero_() - - # Basic sanity check - assert (grad[..., :3].abs() > gs.EPS).all() - assert (grad[..., 3:].abs() < gs.EPS).all() - - grads.append(grad) - - assert_allclose(*grads, atol=gs.EPS) diff --git a/tests/grad/__init__.py b/tests/grad/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/grad/conftest.py b/tests/grad/conftest.py new file mode 100644 index 0000000000..bdae65f549 --- /dev/null +++ b/tests/grad/conftest.py @@ -0,0 +1,255 @@ +import xml.etree.ElementTree as ET + +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.""" + 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, "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 + + +@pytest.fixture(scope="session") +def grad_free(): + mjcf = ET.Element("mujoco", model="free") + worldbody = ET.SubElement(mjcf, "worldbody") + body = ET.SubElement(worldbody, "body", name="chassis", pos="0 0 0") + ET.SubElement(body, "freejoint") + ET.SubElement(body, "inertial", mass="1.0", pos="0 0 0", diaginertia="0.1 0.1 0.1") + ET.SubElement(body, "geom", type="box", size="0.1 0.1 0.1", contype="0", conaffinity="0") + return ET.tostring(mjcf, encoding="unicode") + + +@pytest.fixture(scope="session") +def grad_revolute(): + mjcf = ET.Element("mujoco", model="revolute") + worldbody = ET.SubElement(mjcf, "worldbody") + _add_hinge_arm(worldbody, "arm", "0 0 0") + return ET.tostring(mjcf, encoding="unicode") + + +@pytest.fixture(scope="session") +def grad_revolute_frictionloss(): + mjcf = ET.Element("mujoco", model="revolute_frictionloss") + worldbody = ET.SubElement(mjcf, "worldbody") + _add_hinge_arm(worldbody, "arm", "0 0 0", frictionloss="0.5") + return ET.tostring(mjcf, encoding="unicode") + + +@pytest.fixture(scope="session") +def grad_prismatic(): + mjcf = ET.Element("mujoco", model="prismatic") + worldbody = ET.SubElement(mjcf, "worldbody") + body = ET.SubElement(worldbody, "body", name="slider", pos="0 0 0") + ET.SubElement(body, "joint", type="slide", axis="1 0 0") + ET.SubElement(body, "inertial", mass="0.5", pos="0 0 0", diaginertia="0.01 0.01 0.01") + ET.SubElement(body, "geom", type="box", size="0.05 0.05 0.05", contype="0", conaffinity="0") + return ET.tostring(mjcf, encoding="unicode") + + +@pytest.fixture(scope="session") +def grad_spherical(): + mjcf = ET.Element("mujoco", model="spherical") + worldbody = ET.SubElement(mjcf, "worldbody") + body = ET.SubElement(worldbody, "body", name="ball", pos="0 0 0") + ET.SubElement(body, "joint", type="ball") + 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 ET.tostring(mjcf, encoding="unicode") + + +@pytest.fixture(scope="session") +def grad_capsule(): + mjcf = ET.Element("mujoco", model="capsule") + ET.SubElement(mjcf, "compiler", angle="degree") + worldbody = ET.SubElement(mjcf, "worldbody") + body = ET.SubElement(worldbody, "body", name="capsule", pos="0 0 0") + ET.SubElement(body, "geom", type="capsule", size="0.1 0.2") + ET.SubElement(body, "joint", name="capsule_joint", type="free") + return ET.tostring(mjcf, encoding="unicode") + + +@pytest.fixture(scope="session") +def grad_free_with_revolute(): + mjcf = ET.Element("mujoco", model="free_with_child") + worldbody = ET.SubElement(mjcf, "worldbody") + chassis = ET.SubElement(worldbody, "body", name="chassis", pos="0 0 0") + ET.SubElement(chassis, "freejoint") + ET.SubElement(chassis, "inertial", mass="1.0", pos="0 0 0", diaginertia="0.1 0.1 0.1") + ET.SubElement(chassis, "geom", type="box", size="0.1 0.1 0.1", contype="0", conaffinity="0") + arm = ET.SubElement(chassis, "body", name="arm", pos="0.2 0 0") + ET.SubElement(arm, "joint", type="hinge", axis="0 1 0") + ET.SubElement(arm, "inertial", mass="0.5", pos="0.1 0 0", diaginertia="0.01 0.01 0.01") + ET.SubElement(arm, "geom", type="capsule", fromto="0 0 0 0.2 0 0", size="0.02", contype="0", conaffinity="0") + return ET.tostring(mjcf, encoding="unicode") + + +@pytest.fixture(scope="session") +def grad_revolute_chain3(): + 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") + 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 + return ET.tostring(mjcf, encoding="unicode") + + +@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") + ET.SubElement(body, "inertial", pos="0 0 0", mass="1.0", diaginertia="1.0 1.0 1.0") + ET.SubElement(body, "geom", type="box", size="0.25 0.25 0.1", contype="0", conaffinity="0") + return ET.tostring(mjcf, encoding="unicode") + + +@pytest.fixture(scope="session") +def grad_hinge_pair_joint_eq_linear(): + mjcf = ET.Element("mujoco", model="hinge_pair_joint_eq_linear") + worldbody = ET.SubElement(mjcf, "worldbody") + 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, "joint", joint1="j1", joint2="j2", polycoef="0 1 0 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_hinge_pair_joint_eq_quadratic(): + mjcf = ET.Element("mujoco", model="hinge_pair_joint_eq_quadratic") + worldbody = ET.SubElement(mjcf, "worldbody") + 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, "joint", joint1="j1", joint2="j2", polycoef="0 1 0.5 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_connect_loop(): + 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") + 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" + ) + return ET.tostring(mjcf, encoding="unicode") + + +@pytest.fixture(scope="session") +def grad_weld_pair(): + 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") + equality = ET.SubElement(mjcf, "equality") + ET.SubElement( + equality, + "weld", + body1="arm1", + body2="arm2", + relpose="0 -0.3 0 1 0 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_all_eq_fric(): + # Integration scene exercising every differentiated constraint group: frictionloss on j1, equality JOINT between + # j1 and j2, equality CONNECT between arm3 and arm4, equality WELD between arm5 and arm6. Each group acts on a + # disjoint pair of links so the constraint solver faces a well-posed system within every pair. + mjcf = ET.Element("mujoco", model="all_eq_fric") + worldbody = ET.SubElement(mjcf, "worldbody") + _add_hinge_arm(worldbody, "arm1", "0 0 0", name="j1", frictionloss="0.5") + for i_arm in range(2, 7): + _add_hinge_arm(worldbody, f"arm{i_arm}", f"0 {0.2 * (i_arm - 1):.1f} 0", name=f"j{i_arm}") + equality = ET.SubElement(mjcf, "equality") + ET.SubElement( + equality, "joint", joint1="j1", joint2="j2", polycoef="0 1 0 0 0", solimp="0.95 0.99 0.001", solref="0.005 1" + ) + ET.SubElement( + equality, "connect", body1="arm3", body2="arm4", anchor="0.2 0 0", solimp="0.95 0.99 0.001", solref="0.005 1" + ) + ET.SubElement( + equality, + "weld", + body1="arm5", + body2="arm6", + relpose="0 -0.2 0 1 0 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_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") + ET.SubElement(cart, "inertial", pos="0 0 0", mass="1.0", diaginertia="1.0 1.0 1.0") + ET.SubElement(cart, "geom", type="box", size="0.25 0.25 0.1", contype="0", conaffinity="0", rgba="0 0 0.8 1") + pole = ET.SubElement(cart, "body", name="pole", pos="0 0 0") + ET.SubElement(pole, "joint", name="hinge", type="hinge", axis="0 1 0", damping="0.0") + ET.SubElement(pole, "inertial", pos="0 0 0.5", mass="10.0", diaginertia="1.0 1.0 1.0") + ET.SubElement( + pole, "geom", type="box", pos="0 0 0.5", size="0.025 0.025 0.5", contype="0", conaffinity="0", rgba="1 1 1 1" + ) + return ET.tostring(mjcf, encoding="unicode") + + +@pytest.fixture(scope="session") +def grad_hopper(): + mjcf = ET.Element("mujoco", model="hopper") + ET.SubElement(mjcf, "compiler", angle="radian") + default = ET.SubElement(mjcf, "default") + ET.SubElement(default, "joint", limited="true", armature="1", damping="1") + ET.SubElement(default, "geom", condim="3", friction="0.9 0.005 0.0001") + worldbody = ET.SubElement(mjcf, "worldbody") + torso = ET.SubElement(worldbody, "body", name="torso", pos="0 0 1.25") + root_kwargs = {"pos": "0 0 0", "limited": "false", "armature": "0", "damping": "0"} + ET.SubElement(torso, "joint", name="rootx", axis="1 0 0", type="slide", **root_kwargs) + ET.SubElement(torso, "joint", name="rootz", axis="0 0 1", type="slide", **root_kwargs) + ET.SubElement(torso, "joint", name="rooty", axis="0 1 0", type="hinge", **root_kwargs) + ET.SubElement(torso, "geom", name="torso_geom", type="capsule", size="0.05 0.2") + thigh = ET.SubElement(torso, "body", name="thigh", pos="0 0 -0.2") + ET.SubElement(thigh, "joint", name="thigh_joint", pos="0 0 0", axis="0 -1 0", type="hinge", range="-2.61799 0") + ET.SubElement(thigh, "geom", name="thigh_geom", type="capsule", size="0.05 0.225", pos="0 0 -0.225") + leg = ET.SubElement(thigh, "body", name="leg", pos="0 0 -0.7") + ET.SubElement(leg, "joint", name="leg_joint", pos="0 0 0.25", axis="0 -1 0", type="hinge", range="-2.61799 0") + ET.SubElement(leg, "geom", name="leg_geom", type="capsule", size="0.04 0.25") + foot = ET.SubElement(leg, "body", name="foot", pos="0 0 -0.25") + ET.SubElement( + foot, "joint", name="foot_joint", pos="0 0 0", axis="0 -1 0", type="hinge", range="-0.785398 0.785398" + ) + ET.SubElement( + foot, + "geom", + name="foot_geom", + type="capsule", + size="0.06 0.195", + pos="0.06 0 0", + quat="0.707107 0 -0.707107 0", + friction="2 0.005 0.0001", + ) + return ET.tostring(mjcf, encoding="unicode") diff --git a/tests/grad/test_grad_tape.py b/tests/grad/test_grad_tape.py new file mode 100644 index 0000000000..2b06fac6c0 --- /dev/null +++ b/tests/grad/test_grad_tape.py @@ -0,0 +1,123 @@ +# 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 ..utils import assert_allclose +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): + 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) + return pair.scene_ana, pair.entity_ana + + def run_segment(scene, entity, velocity): + entity.set_dofs_velocity(velocity) + for _ in range(horizon): + scene.step() + 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) + + # Scene A: one scene, snapshot + reset between two horizons. + scene_a, robot_a = build(show_viewer=show_viewer) + v1 = np.random.default_rng(seed=101).standard_normal((B, robot_a.n_dofs)) + v2 = np.random.default_rng(seed=202).standard_normal((B, robot_a.n_dofs)) + scene_a.reset() + v1_a = gs.tensor(v1, dtype=gs.tc_float, requires_grad=True) + loss_h1_a = run_segment(scene_a, robot_a, v1_a) + qpos_mid_a = read_qpos(scene_a) + 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 + grad1_a = tensor_to_array(v1_a.grad).copy() + + v2_a = gs.tensor(v2, dtype=gs.tc_float, requires_grad=True) + loss_h2_a = run_segment(scene_a, robot_a, v2_a) + qpos_end_a = read_qpos(scene_a) + scene_a.backward(loss_h2_a) + grad2_a = tensor_to_array(v2_a.grad).copy() + + # Scene B: horizon 1 only; its backward returns the terminal snapshot Scene C resumes from. + scene_b, robot_b = build() + scene_b.reset() + v1_b = gs.tensor(v1, dtype=gs.tc_float, requires_grad=True) + loss_h1_b = run_segment(scene_b, robot_b, v1_b) + qpos_mid_b = read_qpos(scene_b) + 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_allclose(grad1_a, grad1_b, **tol) + + # Scene C: fresh scene resumed from B's mid-trajectory snapshot. + scene_c, robot_c = build() + scene_c.reset(snapshot_b) + v2_c = gs.tensor(v2, dtype=gs.tc_float, requires_grad=True) + loss_h2_c = run_segment(scene_c, robot_c, v2_c) + qpos_end_c = read_qpos(scene_c) + 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_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): + 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, + ), + show_viewer=show_viewer, + ) + robot = scene.add_entity( + gs.morphs.Box( + size=(0.1, 0.1, 0.1), + pos=(0, 0, 0), + ) + ) + scene.build() + + ctrl = gs.tensor(np.random.randn(robot.n_dofs), dtype=gs.tc_float, requires_grad=True) + + grads = [] + for is_sim_state_source in (False, True): + scene.reset() + robot.set_dofs_velocity(ctrl) + scene.step() + if is_sim_state_source: + chassis_pos = scene.rigid_solver.get_state().links_pos[:, 0].squeeze() + else: + chassis_pos = robot.get_state().pos.squeeze() + loss = torch.linalg.norm(chassis_pos) + loss.backward() + grads.append(ctrl.grad.detach().clone()) + ctrl.grad.zero_() + assert (grads[-1][..., :3].abs() > gs.EPS).all() + assert_allclose(grads[-1][..., 3:], 0.0, atol=gs.EPS) + + assert_allclose(*grads, atol=gs.EPS) diff --git a/tests/grad/test_hybrid_push.py b/tests/grad/test_hybrid_push.py new file mode 100644 index 0000000000..c64097445b --- /dev/null +++ b/tests/grad/test_hybrid_push.py @@ -0,0 +1,81 @@ +# 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 + +import genesis as gs + + +@pytest.mark.slow # ~350s +@pytest.mark.required +@pytest.mark.debug(False) +def test_hybrid_mpm_tool_push_grad(show_viewer): + HORIZON = 10 + + scene = gs.Scene( + sim_options=gs.options.SimOptions( + dt=2e-3, + substeps=10, + requires_grad=True, + ), + mpm_options=gs.options.MPMOptions( + lower_bound=(0.0, -1.0, 0.0), + upper_bound=(1.0, 1.0, 0.55), + ), + viewer_options=gs.options.ViewerOptions( + camera_pos=(2.5, -0.15, 2.42), + camera_lookat=(0.5, 0.5, 0.1), + ), + show_viewer=show_viewer, + ) + scene.add_entity( + gs.morphs.URDF( + file="urdf/plane/plane.urdf", + fixed=True, + ) + ) + stick = scene.add_entity( + morph=gs.morphs.Mesh( + file="meshes/stirrer.obj", + scale=0.6, + pos=(0.5, 0.5, 0.05), + euler=(90.0, 0.0, 0.0), + ), + material=gs.materials.Tool( + friction=8.0, + ), + ) + obj = scene.add_entity( + morph=gs.morphs.Box( + lower=(0.2, 0.1, 0.05), + upper=(0.4, 0.3, 0.15), + ), + material=gs.materials.MPM.Elastic( + rho=500, + ), + ) + scene.build(n_envs=2) + + stick.set_position(gs.tensor([[0.3, 0.1, 0.28], [0.3, 0.1, 0.5]], requires_grad=True)) + obj.set_position(gs.tensor([0.3, 0.3, 0.1], requires_grad=True)) + obj.set_velocity(gs.tensor([0.0, -1.0, 0.0], requires_grad=True)) + goal = gs.tensor([0.5, 0.8, 0.05]) + + loss = 0.0 + v_list = [] + for i in range(HORIZON): + v_i = gs.tensor([[0.0, 1.0, 0.0], [0.0, 1.0, 0.0]], requires_grad=True) + stick.set_velocity(vel=v_i) + v_list.append(v_i) + scene.step() + if i == HORIZON // 2: + mpm_particles = scene.get_state().solvers_state[scene.solvers.index(scene.mpm_solver)] + loss += torch.pow(mpm_particles.pos[mpm_particles.active == 1] - goal, 2).sum() + if i == HORIZON - 2: + loss += torch.pow(obj.get_state().pos - goal, 2).sum() + loss.backward() + + # Every step but the last must move the object (non-zero velocity gradient); the last step cannot affect the loss. + for v_i in v_list[:-1]: + assert (v_i.grad.abs() > gs.EPS).any() + assert (v_list[-1].grad.abs() < gs.EPS).all() diff --git a/tests/grad/test_rigid_collision.py b/tests/grad/test_rigid_collision.py new file mode 100644 index 0000000000..50dc41d440 --- /dev/null +++ b/tests/grad/test_rigid_collision.py @@ -0,0 +1,373 @@ +# 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 + +import genesis as gs +from genesis.utils import set_random_seed +from genesis.utils.geom import R_to_quat +from genesis.utils.misc import qd_to_numpy, qd_to_torch, tensor_to_array + +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): + # 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] + rest_dofs = [0.0, 0.0, rest_z, 0.0, 0.0, 0.0] + 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 + 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) + scene.step() + + scene_ana, obj_ana = _build_contact_scene(shape, grad_capsule, requires_grad=True, show_viewer=show_viewer) + scene_ana.reset() + settle(scene_ana, obj_ana) + nc = _n_contacts(scene_ana) + assert nc > 0, f"setup error: {shape} not in contact after settle (n_contacts={nc})" + forces = [gs.tensor(init_force[t], dtype=gs.tc_float, requires_grad=True) for t in range(n_steps)] + for t in range(n_steps): + obj_ana.control_dofs_force(forces[t]) + scene_ana.step() + assert _n_contacts(scene_ana) == nc, "contact set changed during grad window - FD invalid" + loss = (scene_ana.rigid_solver.get_state().qpos[0, :3] ** 2).sum() + scene_ana.backward(loss) + ana = np.stack([tensor_to_array(f.grad) for f in forces]) + + scene_fd, obj_fd = _build_contact_scene(shape, grad_capsule, requires_grad=True) + + def loss_at(perturbed): + scene_fd.reset() + settle(scene_fd, obj_fd) + for t in range(n_steps): + obj_fd.control_dofs_force(gs.tensor(perturbed[t], dtype=gs.tc_float)) + scene_fd.step() + assert _n_contacts(scene_fd) == nc, "contact set changed under FD perturbation" + return float((scene_fd.rigid_solver.get_state().qpos[0, :3] ** 2).sum().detach()) + + for t in range(n_steps): + plus = init_force.copy() + plus[t, 2] += eps + 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}") + + +@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. + 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, + ), + 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.build() + + for _ in range(20): + scene.step() + + for top, x in zip(tops, (0.8, -0.8)): + assert_allclose(top.get_pos(), (x, 0.0, 0.6), atol=2e-4) + assert_allclose(top.get_dofs_velocity(), 0.0, atol=0.05) + + +@pytest.mark.required +def test_rigid_diff_contact_pair_unsupported_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))) + + # 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() + + +@pytest.mark.required +@pytest.mark.precision("64") +@pytest.mark.debug(False) +def test_rigid_contact_detection_jacobian_matches_fd(): + scene = gs.Scene( + sim_options=gs.options.SimOptions( + dt=0.01, + requires_grad=True, + ), + show_viewer=False, + ) + 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)) + 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])) + ) + scene.build() + collider = scene.sim.rigid_solver.collider + + box1.set_quat(R_to_quat(gs.euler_to_R([np.deg2rad(3.0), np.deg2rad(3.0), np.deg2rad(3.0)]))) + box0_init_pos = box0.get_pos().clone() + box1_init_pos = box1.get_pos().clone() + box0_init_quat = box0.get_quat().clone() + box1_init_quat = box1.get_quat().clone() + + collider.detection() + contacts = collider.get_contacts(as_tensor=True, to_torch=True, keep_batch_dim=True) + normal = contacts["normal"].requires_grad_() + position = contacts["position"].requires_grad_() + penetration = contacts["penetration"].requires_grad_() + loss = ((normal * position).sum(dim=-1) * penetration).sum() + dL_dnormal = torch.autograd.grad(loss, normal, retain_graph=True)[0] + dL_dposition = torch.autograd.grad(loss, position, retain_graph=True)[0] + dL_dpenetration = torch.autograd.grad(loss, penetration)[0] + + collider.backward(dL_dposition, dL_dnormal, dL_dpenetration) + dL_dpos = qd_to_torch(scene.sim.rigid_solver.dyn_state.geoms.pos.grad) + dL_dquat = qd_to_torch(scene.sim.rigid_solver.dyn_state.geoms.quat.grad) + + fd_eps = 1e-5 + trials = 100 + + def directional_error(dL_dx, x_type): + error_rel = 0.0 + for _ in range(trials): + rand_dx = torch.nn.functional.normalize(torch.randn_like(dL_dx), dim=-1) + dL = (rand_dx * dL_dx).sum() + losses = [] + for sign in (1, -1): + if x_type == "pos": + box0.set_pos(box0_init_pos + sign * rand_dx[0, 0] * fd_eps) + box1.set_pos(box1_init_pos + sign * rand_dx[1, 0] * fd_eps) + box0.set_quat(box0_init_quat) + box1.set_quat(box1_init_quat) + else: + box0.set_pos(box0_init_pos) + box1.set_pos(box1_init_pos) + box0.set_quat(box0_init_quat + sign * rand_dx[0, 0] * fd_eps) + box1.set_quat(box1_init_quat + sign * rand_dx[1, 0] * fd_eps) + collider._collider_state.n_contacts.fill(0) + collider.detection() + c = collider.get_contacts(as_tensor=True, to_torch=True, keep_batch_dim=True) + losses.append(((c["normal"] * c["position"]).sum(dim=-1) * c["penetration"]).sum()) + dL_fd = (losses[0] - losses[1]) / (2 * fd_eps) + error_rel += (dL - dL_fd).abs() / max(dL.abs(), dL_fd.abs(), gs.EPS) + return error_rel / trials + + assert_allclose(directional_error(dL_dpos, "pos"), 0.0, atol=1e-4) + assert_allclose(directional_error(dL_dquat, "quat"), 0.0, atol=1e-4) + + +@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 + from genesis.engine.solvers.rigid.rigid_solver import kernel_step_1 + + scene = gs.Scene( + sim_options=gs.options.SimOptions( + dt=0.01, + requires_grad=True, + ), + rigid_options=gs.options.RigidOptions( + constraint_solver=gs.constraint_solver.Newton, + ), + 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.build() + rigid_solver = scene._sim.rigid_solver + constraint_solver = rigid_solver.constraint_solver + + franka.set_qpos([-1.0124, 1.5559, 1.3662, -1.6878, -1.5799, 1.7757, 1.4602, 0.04, 0.04]) + + def constraint_solver_resolve(): + func_solve_init( + rigid_solver.dyn_state, + constraint_solver.constraint_state, + rigid_solver.dyn_info, + rigid_solver.rigid_info, + rigid_solver.rigid_config, + is_decomposed=False, + ) + func_solve_body( + rigid_solver.dyn_state, + constraint_solver.constraint_state, + rigid_solver.dyn_info, + rigid_solver.rigid_info, + rigid_solver.rigid_config, + constraint_solver._n_iterations, + ) + + monkeypatch.setattr(constraint_solver, "resolve", constraint_solver_resolve) + + kernel_step_1( + rigid_solver.dyn_state, + constraint_solver.constraint_state, + rigid_solver.dyn_info, + rigid_solver.rigid_info, + rigid_solver.rigid_config, + is_forward_pos_updated=True, + is_forward_vel_updated=True, + is_backward=False, + ) + constraint_solver.add_equality_constraints() + rigid_solver.collider.detection() + constraint_solver.add_inequality_constraints() + constraint_solver.resolve() + + def compute_loss(input_mass, input_jac, input_aref, input_efc_D, input_force): + rigid_solver.rigid_info.mass_mat.from_numpy(input_mass) + constraint_solver.constraint_state.jac.from_numpy(input_jac) + constraint_solver.constraint_state.aref.from_numpy(input_aref) + constraint_solver.constraint_state.efc_D.from_numpy(input_efc_D) + rigid_solver.dyn_state.dofs.force.from_numpy(input_force) + updated_acc_smooth = np.linalg.solve(input_mass[..., 0], input_force[..., 0]) + rigid_solver.dyn_state.dofs.acc_smooth.from_numpy(updated_acc_smooth[..., None]) + constraint_solver.resolve() + return ((qd_to_torch(constraint_solver.qacc) - target_qacc) ** 2).mean() + + init_input_mass = qd_to_numpy(rigid_solver.rigid_info.mass_mat, copy=True) + init_input_jac = qd_to_numpy(constraint_solver.constraint_state.jac, copy=True) + init_input_aref = qd_to_numpy(constraint_solver.constraint_state.aref, copy=True) + init_input_efc_D = qd_to_numpy(constraint_solver.constraint_state.efc_D, copy=True) + init_input_force = qd_to_numpy(rigid_solver.dyn_state.dofs.force, copy=True) + + set_random_seed(0) + init_output_qacc = qd_to_torch(constraint_solver.qacc) + target_qacc = torch.from_numpy(np.random.randn(*init_output_qacc.shape)).to(device=gs.device) + target_qacc = target_qacc * init_output_qacc.abs().mean() + + output_qacc = qd_to_torch(constraint_solver.qacc, copy=True).requires_grad_(True) + loss = ((output_qacc - target_qacc) ** 2).mean() + dL_dqacc = tensor_to_array(torch.autograd.grad(loss, output_qacc)[0]) + constraint_solver.constraint_state.dL_dqacc.from_numpy(dL_dqacc) + constraint_solver.backward() + + dL_dM = qd_to_numpy(constraint_solver.constraint_state.dL_dM) + dL_djac = qd_to_numpy(constraint_solver.constraint_state.dL_djac) + dL_daref = qd_to_numpy(constraint_solver.constraint_state.dL_daref) + dL_defc_D = qd_to_numpy(constraint_solver.constraint_state.dL_defc_D) + dL_dforce = qd_to_numpy(constraint_solver.constraint_state.dL_dforce) + + fd_eps = 1e-3 + trials = 200 + for dL_dx, x_type in ( + (dL_dforce, "force"), + (dL_daref, "aref"), + (dL_defc_D, "efc_D"), + (dL_djac, "jac"), + (dL_dM, "mass"), + ): + error = 0.0 + for _ in range(trials): + rand_dx = np.random.randn(*dL_dx.shape) + rand_dx = rand_dx / max( + np.linalg.norm(rand_dx, axis=0 if x_type in ("force", "aref", "efc_D") else (0, 1)), gs.EPS + ) + if x_type == "mass": + rand_dx = (rand_dx + np.moveaxis(rand_dx, 0, 1)) * 0.5 + dL = (rand_dx * dL_dx).sum() + + inputs = dict( + input_mass=init_input_mass, + input_jac=init_input_jac, + input_aref=init_input_aref, + input_efc_D=init_input_efc_D, + input_force=init_input_force, + ) + key = { + "force": "input_force", + "aref": "input_aref", + "efc_D": "input_efc_D", + "jac": "input_jac", + "mass": "input_mass", + }[x_type] + init_x = inputs[key] + loss_p = compute_loss(**{**inputs, key: init_x + rand_dx * fd_eps}) + loss_m = compute_loss(**{**inputs, key: init_x - rand_dx * fd_eps}) + dL_fd = (loss_p - loss_m) / (2 * fd_eps) + error += (dL - dL_fd).abs() / max(abs(dL), abs(dL_fd), gs.EPS) + assert_allclose(error / trials, 0.0, atol=1e-4) diff --git a/tests/grad/test_rigid_constraints.py b/tests/grad/test_rigid_constraints.py new file mode 100644 index 0000000000..aa9f1cf638 --- /dev/null +++ b/tests/grad/test_rigid_constraints.py @@ -0,0 +1,265 @@ +# 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 +import pytest + +import genesis as gs +from genesis.utils.misc import qd_to_torch, tensor_to_array + +from ..utils import assert_allclose +from .utils import assert_grad_matches_fd, make_diff_scene_pair + + +@pytest.mark.required +@pytest.mark.debug(False) +def test_rigid_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, + substeps=4, + dt=1.0 / 60.0, + gravity=(0.0, 0.0, 0.0), + enable_joint_limit=False, + disable_constraint=True, + ) + off.scene_fd.reset() + off.entity_fd.set_dofs_velocity(gs.tensor([100.0], dtype=gs.tc_float)) + for _ in range(60): + off.scene_fd.step() + assert abs(off.scene_fd.rigid_solver.get_state().qpos[0, 0]) > 50.0 + + on = make_diff_scene_pair( + grad_slider_limit, + substeps=4, + dt=1.0 / 60.0, + gravity=(0.0, 0.0, 0.0), + enable_joint_limit=True, + disable_constraint=False, + show_viewer=show_viewer, + ) + on.scene_fd.reset() + on.entity_fd.set_dofs_velocity(gs.tensor([100.0], dtype=gs.tc_float)) + for _ in range(60): + on.scene_fd.step() + assert abs(on.scene_fd.rigid_solver.get_state().qpos[0, 0]) <= 4.5 + + # Backward: a rollout that drives the cart into the active |x|=4 limit, so the gradient flows through the + # constraint correction. Sanity-check that the cart actually reaches the band first. + on.scene_fd.reset() + on.entity_fd.set_dofs_velocity(gs.tensor([100.0], dtype=gs.tc_float)) + for _ in range(5): + on.scene_fd.step() + assert abs(on.scene_fd.rigid_solver.get_state().qpos[0, 0]) > 3.5 + + assert_grad_matches_fd( + on, + [np.array([100.0])], + lambda e, x: e.set_dofs_velocity(x), + lambda scene, entity: scene.rigid_solver.get_state().qpos[0, 0] ** 2, + n_steps=5, + rtol=1e-10 if precision == "64" else 5e-4, + atol=1e-10 if precision == "64" else 5e-4, + eps=3e-4 if precision == "64" else 3e-2, + ) + + # Inside-limit single step: the cart stays well inside the range so the limit is present but inactive; the + # constraint-inclusive forward+backward chain must still satisfy central FD (a smoother path than the crossing). + assert_grad_matches_fd( + on, + [np.array([2.0])], + lambda e, x: e.set_dofs_velocity(x), + lambda scene, entity: scene.rigid_solver.get_state().qpos[0, 0] ** 2, + n_steps=1, + rtol=1e-10 if precision == "64" else 5e-4, + atol=1e-10 if precision == "64" else 5e-4, + eps=3e-4 if precision == "64" else 3e-2, + ) + + # Inactive-path parity: with the limit enabled but never hit, the adjoint must equal the no-limit baseline - the + # inactive constraint branch must inject no spurious gradient. + off_solver = make_diff_scene_pair( + grad_slider_limit, + substeps=4, + dt=1.0 / 60.0, + gravity=(0.0, 0.0, 0.0), + enable_joint_limit=False, + disable_constraint=False, + ) + grads = {} + for pair, key in ((off_solver, "off"), (on, "on")): + pair.scene_ana.reset() + v = gs.tensor([0.5], dtype=gs.tc_float, requires_grad=True) + pair.entity_ana.set_dofs_velocity(v) + pair.scene_ana.step() + loss = pair.scene_ana.rigid_solver.get_state().qpos[0, 0] ** 2 + loss.backward() + grads[key] = tensor_to_array(v.grad) + assert_allclose( + grads["on"], + grads["off"], + rtol=1e-6 if precision == "64" else 1e-4, + atol=1e-9 if precision == "64" else 1e-6, + ) + + +@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): + # 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). + 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), + }[model_name] + + pair = make_diff_scene_pair( + request.getfixturevalue(model_name), + substeps=4, + dt=1.0 / 60.0, + gravity=gravity, + enable_joint_limit=True, + disable_constraint=False, + show_viewer=show_viewer, + ) + forces = [np.array(per_step_force) for _ in range(n_steps)] + + def setup_fn(scene, entity): + if init_pos is not None: + entity.set_dofs_position(gs.tensor(init_pos, dtype=gs.tc_float)) + + def loss_fn(scene, entity): + state = scene.rigid_solver.get_state() + return (state.links_pos.reshape(-1) ** 2).sum() if is_links_loss else state.qpos[0, 0] ** 2 + + pair.scene_fd.reset() + setup_fn(pair.scene_fd, pair.entity_fd) + for force in forces: + pair.entity_fd.control_dofs_force(gs.tensor(force, dtype=gs.tc_float)) + pair.scene_fd.step() + reached = abs(pair.scene_fd.rigid_solver.get_state().qpos[0, sanity_dof]) + assert reached > sanity_thresh, f"setup error: {model_name} did not reach its limit band (q={reached})" + + assert_grad_matches_fd( + pair, + forces, + 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, + eps=3e-2, + ) + + +@pytest.mark.required +def test_rigid_frictionloss_grad_matches_fd(grad_revolute_frictionloss, precision, show_viewer): + pair = make_diff_scene_pair( + grad_revolute_frictionloss, + substeps=4, + dt=1.0 / 60.0, + gravity=(0.0, 0.0, 0.0), + enable_joint_limit=False, + disable_constraint=False, + show_viewer=show_viewer, + ) + pair.scene_fd.reset() + pair.scene_fd.step() + cs = pair.scene_fd.rigid_solver.constraint_solver.constraint_state + assert qd_to_torch(cs.n_constraints_frictionloss)[0] == 1 + + assert_grad_matches_fd( + pair, + [np.array([2.0])], + lambda e, x: e.set_dofs_velocity(x), + lambda scene, entity: scene.rigid_solver.get_state().qpos[0, 0] ** 2, + n_steps=10, + rtol=1e-10 if precision == "64" else 5e-5, + atol=1e-10 if precision == "64" else 5e-5, + eps=3e-6 if precision == "64" else 1e-2, + ) + + +@pytest.mark.required +@pytest.mark.parametrize( + "model_name, n_rows", + [ + ("grad_hinge_pair_joint_eq_linear", 1), + ("grad_hinge_pair_joint_eq_quadratic", 1), + ("grad_connect_loop", 3), + ("grad_weld_pair", 6), + ], +) +def test_rigid_equality_grad_matches_fd(model_name, n_rows, request, precision, show_viewer): + pair = make_diff_scene_pair( + request.getfixturevalue(model_name), + substeps=4, + dt=1.0 / 60.0, + gravity=(0.0, 0.0, 0.0), + enable_joint_limit=False, + disable_constraint=False, + show_viewer=show_viewer, + ) + pair.scene_fd.reset() + pair.scene_fd.step() + cs = pair.scene_fd.rigid_solver.constraint_solver.constraint_state + assert qd_to_torch(cs.n_constraints_equality)[0] == n_rows + + assert_grad_matches_fd( + pair, + [np.array([0.8, -0.3])], + 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, + 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): + # 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( + grad_all_eq_fric, + substeps=4, + dt=1.0 / 60.0, + gravity=(0.0, 0.0, 0.0), + enable_joint_limit=False, + disable_constraint=False, + show_viewer=show_viewer, + ) + pair.scene_fd.reset() + pair.scene_fd.step() + cs = pair.scene_fd.rigid_solver.constraint_solver.constraint_state + assert qd_to_torch(cs.n_constraints_equality)[0] == 10 + assert qd_to_torch(cs.n_constraints_frictionloss)[0] == 1 + + weights = np.array([1.0, 0.7, 1.3, 0.5, 0.9, 1.1]) + + def loss_fn(scene, entity): + qpos = scene.rigid_solver.get_state().qpos[0] + return sum(weights[d] * qpos[d] ** 2 for d in range(6)) + + assert_grad_matches_fd( + pair, + [np.array([0.8, -0.3, 0.5, -0.2, 0.4, -0.6])], + 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, + 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 new file mode 100644 index 0000000000..30e671c839 --- /dev/null +++ b/tests/grad/test_rigid_dynamics.py @@ -0,0 +1,168 @@ +# 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 +import pytest + +import genesis as gs + +from ..conftest import SKIP_METAL_GRAD +from .utils import assert_grad_matches_fd, make_diff_scene_pair + + +@pytest.mark.required +@pytest.mark.parametrize( + "backend", + [ + gs.cpu, + # FIXME: Quadrants' released native-Metal reverse-mode autodiff collapses per-env adjoints (fixed upstream, + # see Quadrants issue #805). Re-enable once the fix ships in a Quadrants release. + pytest.param(gs.gpu, marks=pytest.mark.skipif(sys.platform == "darwin", reason=SKIP_METAL_GRAD)), + ], +) +@pytest.mark.parametrize( + "model_name", + [ + "grad_free", + "grad_revolute", + "grad_prismatic", + "grad_spherical", + "grad_free_with_revolute", + "grad_revolute_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) + 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") + # (setter, output, input_seed) per joint, plus the position and quaternion target seeds. Single-step + # force->position is only checked on cartpole, where its sensitivity clears the finite-difference floor. + checks_by_joint = { + "grad_free": ((("pos", "pos", 10), ("quat", "quat", 11), ("vel", "pos", 12), ("vel", "quat", 13)), 1, 2), + "grad_revolute": ((("vel", "pos", 30), ("vel", "quat", 31), ("force", "quat", 32)), 21, 22), + "grad_prismatic": ((("vel", "pos", 50),), 41, 0), + "grad_spherical": ((("vel", "pos", 70), ("vel", "quat", 71), ("force", "quat", 72)), 61, 62), + "grad_free_with_revolute": ( + (("pos", "pos", 70), ("quat", "quat", 71), ("vel", "pos", 72), ("vel", "quat", 73)), + 61, + 62, + ), + "grad_revolute_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), + }[model_name] + + pos_shape = (B, 3) if is_single_link else (B, n_links, 3) + quat_shape = (B, 4) if is_single_link else (B, n_links, 4) + tgt_pos = gs.tensor(np.random.RandomState(pos_seed).standard_normal(pos_shape), dtype=gs.tc_float).reshape(-1) + tgt_quat = gs.tensor(np.random.RandomState(quat_seed).standard_normal(quat_shape), dtype=gs.tc_float).reshape(-1) + for setter, output, input_seed in checks: + rng = np.random.default_rng(input_seed) + if setter == "pos": + step_input = rng.standard_normal((B, 3)) + elif setter == "quat": + step_input = np.broadcast_to(np.array([1.0, 0.0, 0.0, 0.0]), (B, 4)).copy() + step_input = step_input + 0.05 * rng.standard_normal((B, 4)) + step_input = step_input / np.linalg.norm(step_input, axis=-1, keepdims=True) + else: + step_input = rng.standard_normal((B, n_dofs)) + + apply_fn = { + "pos": lambda e, x: e.set_pos(x), + "quat": lambda e, x: e.set_quat(x), + "vel": lambda e, x: e.set_dofs_velocity(x), + "force": lambda e, x: e.control_dofs_force(x), + }[setter] + + target = tgt_pos if output == "pos" else tgt_quat + + def loss_fn(scene, entity, tgt=target, out=output, sl=is_single_link): + if sl: + pose = entity.get_state().pos if out == "pos" else entity.get_state().quat + else: + state = scene.rigid_solver.get_state() + pose = state.links_pos if out == "pos" else state.links_quat + return ((pose.reshape(-1) - tgt) ** 2).sum() + + assert_grad_matches_fd( + pair, + [step_input], + apply_fn, + loss_fn, + rtol=1e-9 if precision == "64" else fp32_tol, + atol=1e-9 if precision == "64" else fp32_tol, + eps=3e-5 if precision == "64" else fp32_eps, + ) + + +@pytest.mark.required +@pytest.mark.parametrize( + "model_name", + [ + "grad_free", + "grad_revolute", + "grad_prismatic", + "grad_free_with_revolute", + "grad_revolute_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): + # 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, 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_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_cartpole": ("links", (2, 3), 167, 2e-5), + "grad_hopper": ("links", (5, 3), 168, 2e-4), + }[model_name] + pair = make_diff_scene_pair(request.getfixturevalue(model_name), n_envs=0, substeps=4, show_viewer=show_viewer) + 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 + return ((pose.reshape(-1) - target) ** 2).sum() + + # fp32 needs a large step to clear the state-noise floor; fp64 needs a small step to bound truncation error. + assert_grad_matches_fd( + pair, + inputs, + lambda e, x: e.control_dofs_force(x), + loss_fn, + rtol=5e-9 if precision == "64" else fp32_tol, + atol=5e-9 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 new file mode 100644 index 0000000000..9ffd7f5f7b --- /dev/null +++ b/tests/grad/test_rigid_optim.py @@ -0,0 +1,156 @@ +# 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 +import pytest +import torch + +import genesis as gs +from genesis.utils.misc import tensor_to_array + +from ..conftest import SKIP_METAL_GRAD +from ..utils import assert_allclose +from .utils import make_diff_scene_pair + + +@pytest.mark.required +@pytest.mark.parametrize( + "backend", + [ + gs.cpu, + pytest.param(gs.gpu, marks=pytest.mark.skipif(sys.platform == "darwin", reason=SKIP_METAL_GRAD)), + ], +) +@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): + # 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 + 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) + 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) + + scene_ref.reset() + if control_target == "init_vel": + target_ctrl = rng.normal(size=(B, N_DOFS)) * 0.5 + robot_ref.set_dofs_velocity(gs.tensor(target_ctrl, dtype=gs.tc_float)) + for _ in range(N_STEPS): + scene_ref.step() + else: + target_ctrl = rng.normal(size=(N_STEPS, B, N_DOFS)) * 0.2 + for t in range(N_STEPS): + robot_ref.control_dofs_force(gs.tensor(target_ctrl[t], dtype=gs.tc_float)) + scene_ref.step() + ref_state = scene_ref.rigid_solver.get_state() + target_qpos = ref_state.qpos.detach().clone() + target_vel = ref_state.dofs_vel.detach().clone() + + if control_target == "init_vel": + init_vel = gs.tensor(target_ctrl + rng.normal(size=(B, N_DOFS)) * 0.3, dtype=gs.tc_float, requires_grad=True) + params = [init_vel] + else: + forces = [ + gs.tensor(target_ctrl[t] + rng.normal(size=(B, N_DOFS)) * 0.1, dtype=gs.tc_float, requires_grad=True) + for t in range(N_STEPS) + ] + params = forces + optimizer = torch.optim.Adam(params, lr=LR) + + loss_history = [] + for _ in range(N_ITER): + optimizer.zero_grad(set_to_none=False) + scene_opt.reset() + if control_target == "init_vel": + robot_opt.set_dofs_velocity(init_vel) + for _ in range(N_STEPS): + scene_opt.step() + else: + for t in range(N_STEPS): + robot_opt.control_dofs_force(forces[t]) + scene_opt.step() + state = scene_opt.rigid_solver.get_state() + diff_pos = (state.qpos - target_qpos).reshape(B, -1) + diff_vel = (state.dofs_vel - target_vel).reshape(B, -1) + loss_per_env = (diff_pos**2).sum(dim=-1) + (diff_vel**2).sum(dim=-1) + loss_history.append(tensor_to_array(loss_per_env).copy()) + loss_per_env.sum().backward() + optimizer.step() + + history = np.asarray(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})") + assert_allclose(final, 0.0, atol=ABS_THRESHOLD, err_msg="final loss above absolute threshold") + + +@pytest.mark.slow +@pytest.mark.required +@pytest.mark.debug(False) +def test_rigid_optim_reach_goal_pose(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), + camera_lookat=(0.5, 0.5, 0.1), + ), + show_viewer=show_viewer, + ) + box = scene.add_entity( + gs.morphs.Box( + pos=(0, 0, 0), + size=(0.1, 0.1, 0.2), + ), + surface=gs.surfaces.Default( + color=(0.9, 0.0, 0.0, 1.0), + ), + ) + scene.build() + + init_pos = gs.tensor([0.3, 0.1, 0.28], requires_grad=True) + init_quat = gs.tensor([1.0, 0.0, 0.0, 0.0], requires_grad=True) + optimizer = torch.optim.Adam([init_pos, init_quat], lr=1e-2) + scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=200, eta_min=1e-3) + + for _ in range(200): + scene.reset() + box.set_pos(init_pos) + box.set_quat(init_quat) + for _ in range(100): + scene.step() + box_state = box.get_state() + loss = torch.abs(box_state.pos - goal_pos).sum() + torch.abs(box_state.quat - goal_quat).sum() + optimizer.zero_grad() + loss.backward() + optimizer.step() + scheduler.step() + with torch.no_grad(): + init_quat.data = init_quat / torch.norm(init_quat, dim=-1, keepdim=True) + + assert_allclose(loss, 0.0, atol=1e-2) diff --git a/tests/grad/utils.py b/tests/grad/utils.py new file mode 100644 index 0000000000..9af357ab3d --- /dev/null +++ b/tests/grad/utils.py @@ -0,0 +1,130 @@ +from typing import NamedTuple + +import numpy as np +import torch + +import genesis as gs +from genesis.engine.entities import RigidEntity +from genesis.utils.misc import tensor_to_array + +from ..utils import assert_allclose + + +class DiffScenePair(NamedTuple): + scene_ana: "gs.Scene" + entity_ana: RigidEntity + scene_fd: "gs.Scene" + entity_fd: RigidEntity + + +def make_diff_scene_pair( + mjcf, + *, + n_envs=0, + 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, + 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 + + scenes = [] + entities = [] + for requires_grad in (True, False): + scene = gs.Scene( + sim_options=gs.options.SimOptions( + dt=dt, + substeps=substeps, + gravity=gravity, + requires_grad=requires_grad, + ), + rigid_options=gs.options.RigidOptions(**rigid_kwargs), + 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)) + scene.build(n_envs=n_envs) + scenes.append(scene) + entities.append(entity) + return DiffScenePair(scenes[0], entities[0], scenes[1], entities[1]) + + +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. + base = [np.array(inp, dtype=np.float64) for inp in inputs] + total_steps = len(base) if n_steps is None else n_steps + + # Analytical pass on the diff-mode scene: apply each tracked input before its step, then backprop the loss. + pair.scene_ana.reset() + if setup_fn is not None: + setup_fn(pair.scene_ana, pair.entity_ana) + x_anas = [] + for i_step in range(total_steps): + if i_step < len(base): + x = gs.tensor(base[i_step], dtype=gs.tc_float, requires_grad=True) + x_anas.append(x) + apply_fn(pair.entity_ana, x) + pair.scene_ana.step() + loss = loss_fn(pair.scene_ana, pair.entity_ana) + assert loss.requires_grad, "loss does not require grad - output is not grad-aware" + loss.backward() + ana_grads = [] + for i_step, x in enumerate(x_anas): + assert x.grad is not None, f"input {i_step}: x.grad is None after backward" + ana_grads.append(tensor_to_array(x.grad)) + + # Finite-difference reference on the production scene: perturb each entry of each input by +/- eps, re-run the + # full trajectory for both signs, and central-difference the loss. + for i_input in range(len(base)): + fd_grad = np.zeros_like(base[i_input]) + for i_entry in range(base[i_input].size): + perturbed = [] + for sign in (+1, -1): + pair.scene_fd.reset() + if setup_fn is not None: + setup_fn(pair.scene_fd, pair.entity_fd) + for i_step in range(total_steps): + if i_step < len(base): + inp = base[i_step].copy() + if i_step == i_input: + inp.reshape(-1)[i_entry] += sign * eps + apply_fn(pair.entity_fd, gs.tensor(inp, dtype=gs.tc_float)) + pair.scene_fd.step() + 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), + rtol=rtol, + atol=atol, + err_msg=f"input {i_input}: FD vs analytical mismatch", + )