diff --git a/genesis/engine/entities/rigid_entity/rigid_entity.py b/genesis/engine/entities/rigid_entity/rigid_entity.py index 24e0dceb30..9573b60939 100644 --- a/genesis/engine/entities/rigid_entity/rigid_entity.py +++ b/genesis/engine/entities/rigid_entity/rigid_entity.py @@ -59,23 +59,25 @@ def wrapper(self, *args, **kwargs): bound = sig.bind(self, *args, **kwargs) bound.apply_defaults() args_dict = dict(tuple(bound.arguments.items())[1:]) - # Key the slot by (method, dofs subset) so same-step calls on distinct subsets (e.g. arm and gripper - # force control) each keep their own entry and gradient path; keyed by method alone, the second call - # would evict the first from the tape. Slices key directly when hashable (Python 3.12 onward) and - # resolve against the entity dof count otherwise. - dofs_idx_local = args_dict.get("dofs_idx_local") - if dofs_idx_local is None: - subset = None - elif isinstance(dofs_idx_local, slice): - if isinstance(dofs_idx_local, Hashable): - subset = dofs_idx_local + # Key the slot by (method, dofs subset, envs subset) so same-step calls on distinct subsets (e.g. arm + # and gripper force control, or per-environment-group commands) each keep their own entry and gradient + # path; a coarser key would let the second call evict the first from the tape. Slices key directly when + # hashable (Python 3.12 onward) and resolve against their dimension size otherwise. + key = [fun.__name__] + for indices, n in ( + (args_dict.get("dofs_idx_local"), self.n_dofs), + (args_dict.get("envs_idx"), self._solver._B), + ): + if indices is None: + subset = None + elif isinstance(indices, slice): + subset = indices if isinstance(indices, Hashable) else tuple(range(*indices.indices(n))) + elif isinstance(indices, torch.Tensor): + subset = tuple(tensor_to_array(indices).reshape(-1).tolist()) else: - subset = tuple(range(*dofs_idx_local.indices(self.n_dofs))) - elif isinstance(dofs_idx_local, torch.Tensor): - subset = tuple(tensor_to_array(dofs_idx_local).reshape(-1).tolist()) - else: - subset = tuple(np.asarray(dofs_idx_local).reshape(-1).tolist()) - self._update_tgt((fun.__name__, subset), args_dict) + subset = tuple(np.asarray(indices).reshape(-1).tolist()) + key.append(subset) + self._update_tgt(tuple(key), args_dict) return fun(self, *args, **kwargs) return wrapper @@ -159,7 +161,15 @@ def __init__( self._load_model() # Initialize target variables and checkpoint - self._tgt_keys = ("set_pos", "set_quat", "set_dofs_velocity", "control_dofs_force") + self._tgt_keys = ( + "set_pos", + "set_quat", + "set_dofs_velocity", + "control_dofs_force", + "control_dofs_velocity", + "control_dofs_position", + "control_dofs_position_velocity", + ) self._tgt = dict() self._tgt_buffer = list() self._ckpt = dict() @@ -1588,17 +1598,11 @@ def process_input(self, in_backward=False): # Do not update [tgt], as input information is finalized at this point self._update_tgt_while_set = False - match key[0]: - case "set_pos": - self.set_pos(**data_kwargs) - case "set_quat": - 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[0]} not in {self._tgt_keys}") + # Every tracked setter replays uniformly from its taped kwargs, so dispatch by name (a rare legitimate + # use of getattr, on our own vetted key set). + if key[0] not in self._tgt_keys: + gs.raise_exception(f"Invalid target key: {key[0]} not in {self._tgt_keys}") + getattr(self, key[0])(**data_kwargs) self._tgt = dict() self._update_tgt_while_set = update_tgt_while_set @@ -1637,6 +1641,15 @@ def process_input_grad(self): self.set_dofs_force_grad, data_kwargs["dofs_idx_local"], data_kwargs["envs_idx"] ) + case "control_dofs_velocity" | "control_dofs_position" | "control_dofs_position_velocity": + # PD control targets are replayed for primal correctness but have no input-gradient path. + for target in (data_kwargs.get("position"), data_kwargs.get("velocity")): + if isinstance(target, torch.Tensor) and target.requires_grad: + gs.raise_exception( + "Gradients with respect to PD control targets are not supported yet. Use " + "'control_dofs_force' for differentiable control inputs." + ) + case _: gs.raise_exception(f"Invalid target key: {key[0]} not in {self._tgt_keys}") @@ -4011,6 +4024,7 @@ def control_dofs_force(self, force, dofs_idx_local=None, envs_idx=None): self._solver.control_dofs_force(force, dofs_idx, envs_idx) @gs.assert_built + @tracked def control_dofs_velocity(self, velocity, dofs_idx_local=None, envs_idx=None): """ Set the PD controller's target velocity for the entity's dofs. This is used for velocity control. @@ -4033,6 +4047,7 @@ def control_dofs_velocity(self, velocity, dofs_idx_local=None, envs_idx=None): self._solver.control_dofs_velocity(velocity, dofs_idx, envs_idx) @gs.assert_built + @tracked def control_dofs_position(self, position, dofs_idx_local=None, envs_idx=None): """ Set the position controller's target position for the entity's dofs. The controller is a proportional term @@ -4056,6 +4071,7 @@ def control_dofs_position(self, position, dofs_idx_local=None, envs_idx=None): self._solver.control_dofs_position(position, dofs_idx, envs_idx) @gs.assert_built + @tracked def control_dofs_position_velocity(self, position, velocity, dofs_idx_local=None, envs_idx=None): """ Set a PD controller's target position and velocity for the entity's dofs. This is used for position control. diff --git a/tests/grad/conftest.py b/tests/grad/conftest.py index 9abedc9374..7e60a914ef 100644 --- a/tests/grad/conftest.py +++ b/tests/grad/conftest.py @@ -3,10 +3,12 @@ import pytest -def _add_hinge_arm(parent, body_name, pos, axis="0 1 0", **joint_kwargs): +def _add_hinge_arm(parent, body_name, pos, axis="0 1 0", joint_pos=None, **joint_kwargs): """Add a 1-DOF hinge arm link (y-axis hinge by default, capsule geom, explicit inertia) and return its body element.""" body = ET.SubElement(parent, "body", name=body_name, pos=pos) + if joint_pos is not None: + joint_kwargs["pos"] = joint_pos ET.SubElement(body, "joint", type="hinge", axis=axis, **joint_kwargs) ET.SubElement(body, "inertial", mass="0.5", pos="0.1 0 0", diaginertia="0.01 0.01 0.01") ET.SubElement(body, "geom", type="capsule", fromto="0 0 0 0.2 0 0", size="0.02", contype="0", conaffinity="0") @@ -170,12 +172,14 @@ def grad_weld_pair(): arm1 = _add_hinge_arm(worldbody, "arm1", "0 0 0", name="j1") _add_hinge_arm(arm1, "arm2", "0.2 0 0", name="j2", axis="1 0 0") equality = ET.SubElement(mjcf, "equality") + # relpose is the pose of body2 in body1's frame at the welded configuration; a wrong sign leaves the weld + # violated at rest and turns the scenario into a violent snap instead of a hold. ET.SubElement( equality, "weld", body1="arm2", body2="arm1", - relpose="0.2 0 0 1 0 0 0", + relpose="-0.2 0 0 1 0 0 0", solimp="0.95 0.99 0.001", solref="0.005 1", ) @@ -185,13 +189,20 @@ def grad_weld_pair(): @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. + # j1 and j2, equality CONNECT between arm3 and arm4, equality WELD between arm6's nested child and arm5. Each + # group acts on a disjoint pair of links so the constraint solver faces a well-posed system within every pair. + # The weld chain nests a skew-axis hinge (see grad_weld_pair) so its rotation rows carry a nonzero velocity bias + # across two separate trees. 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}") + arm = _add_hinge_arm(worldbody, f"arm{i_arm}", f"0 {0.2 * (i_arm - 1):.1f} 0", name=f"j{i_arm}") + if i_arm == 6: + # Skew axis for a nonzero rotation-row velocity bias; the joint position offset keeps the hinge axis + # off the welded point so the weld position rows engage (an axis through it leaves the spin resisted + # only by the weak rotation rows and the chain winds up). + _add_hinge_arm(arm, "arm6b", "0.2 0 0", name="j6b", axis="1 0 0", joint_pos="0.05 0.02 0.06") 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" @@ -199,14 +210,16 @@ def grad_all_eq_fric(): ET.SubElement( equality, "connect", body1="arm3", body2="arm4", anchor="0.2 0 0", solimp="0.95 0.99 0.001", solref="0.005 1" ) + # The softer solref keeps the 6-row weld on this three-dof chain clear of the stabilization stability + # boundary (the shared 0.005 timeconst is clamped to 2 * substep_dt, which is marginal here). ET.SubElement( equality, "weld", - body1="arm5", - body2="arm6", - relpose="0 -0.2 0 1 0 0 0", + body1="arm6b", + body2="arm5", + relpose="-0.2 -0.2 0 1 0 0 0", solimp="0.95 0.99 0.001", - solref="0.005 1", + solref="0.02 1", ) return ET.tostring(mjcf, encoding="unicode") diff --git a/tests/grad/test_rigid_collision.py b/tests/grad/test_rigid_collision.py index 30e8b36d2a..56f455a6d7 100644 --- a/tests/grad/test_rigid_collision.py +++ b/tests/grad/test_rigid_collision.py @@ -92,13 +92,17 @@ def _n_contacts(scene): # only guards the degenerate-FD case. fd_rtol = 1e-6 if precision == "64" else 2e-3 fd_atol = 1e-14 if precision == "64" else 1e-6 - base_force = np.array([0.0, 0.0, -8.0, 0.0, 0.0, 0.0]) + # The small tangential preload keeps the contact manifold away from a gain / loss transition, which the + # tangential finite-difference perturbation would otherwise cross (the in-test contact-count guard trips). + base_force = np.array([0.5, 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): + # Settling under the same base force as the grad window keeps the contact manifold in its loaded + # configuration throughout, so the in-window contact-count guards hold. obj.set_dofs_position(gs.tensor(rest_dofs, dtype=gs.tc_float).sceneless()) for _ in range(n_settle): - obj.control_dofs_force(0.0) + obj.control_dofs_force(base_force) scene.step() scene_ana, obj_ana = _build_contact_scene(shape, grad_capsule, requires_grad=True, show_viewer=show_viewer) @@ -124,18 +128,23 @@ 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)) + obj_fd.control_dofs_force(perturbed[t]) 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()) + # Both the normal (z) and a tangential (x) force component: the tangential adjoint flows through the friction + # rows and their contact-frame chains, which the normal component alone leaves untested. 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=fd_rtol, atol=fd_atol, err_msg=f"contact force.grad mismatch at t={t}") + for i_d in (2, 0): + plus = init_force.copy() + plus[t, i_d] += eps + minus = init_force.copy() + minus[t, i_d] -= eps + fd = (loss_at(plus) - loss_at(minus)) / (2 * eps) + assert_allclose( + ana[t, i_d], fd, rtol=fd_rtol, atol=fd_atol, err_msg=f"contact force.grad mismatch at t={t} dof={i_d}" + ) @pytest.mark.required @@ -341,7 +350,9 @@ def test_constraint_solver_backward_matches_fd(monkeypatch): from genesis.engine.solvers.rigid.constraint.solver import func_solve_body, func_solve_init from genesis.engine.solvers.rigid.rigid_solver import kernel_step_1 - # fp64 is required: the FD perturbation must be small enough for a reliable estimate, which fp32 cannot resolve + # fp64 is required: the parameter-side inputs (aref, efc_D, jac, mass) move the loss by far less than its own + # fp32 resolution at any usable step size, so their finite-difference deltas drown at fp32 + scene = gs.Scene( sim_options=gs.options.SimOptions( requires_grad=True, @@ -413,7 +424,9 @@ def compute_loss(input_mass, input_jac, input_aref, input_efc_D, 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() + # Collapse to a Python scalar right away: the reduction reads a zero-copy view of qacc, and a deferred + # evaluation (MPS is lazy) would otherwise run only after the next call overwrites the buffer in place. + return float(((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) @@ -475,5 +488,5 @@ def compute_loss(input_mass, input_jac, input_aref, input_efc_D, input_force): 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) + error += abs(dL - dL_fd) / 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 index da32ebe9bf..3bd5d6d1f9 100644 --- a/tests/grad/test_rigid_constraints.py +++ b/tests/grad/test_rigid_constraints.py @@ -16,6 +16,7 @@ def test_joint_limit_grad_matches_fd(grad_slider_limit, precision, show_viewer): # Forward: the slider limit must actually bound the cart (it drifts freely when the constraint is off). off = make_diff_scene_pair( grad_slider_limit, + n_envs=2, substeps=4, dt=1.0 / 60.0, gravity=(0.0, 0.0, 0.0), @@ -24,13 +25,14 @@ def test_joint_limit_grad_matches_fd(grad_slider_limit, precision, show_viewer): modes=(False,), ) off.scene_fd.reset() - off.entity_fd.set_dofs_velocity(gs.tensor([100.0], dtype=gs.tc_float)) + off.entity_fd.set_dofs_velocity(100.0) for _ in range(60): off.scene_fd.step() - assert abs(off.scene_fd.rigid_solver.get_state().qpos[0, 0]) > 50.0 + assert (off.scene_fd.rigid_solver.get_state().qpos[:, 0].abs() > 50.0).all() on = make_diff_scene_pair( grad_slider_limit, + n_envs=2, substeps=4, dt=1.0 / 60.0, gravity=(0.0, 0.0, 0.0), @@ -39,47 +41,37 @@ def test_joint_limit_grad_matches_fd(grad_slider_limit, precision, show_viewer): show_viewer=show_viewer, ) on.scene_fd.reset() - on.entity_fd.set_dofs_velocity(gs.tensor([100.0], dtype=gs.tc_float)) + on.entity_fd.set_dofs_velocity(100.0) for _ in range(60): on.scene_fd.step() - assert abs(on.scene_fd.rigid_solver.get_state().qpos[0, 0]) <= 4.5 + assert (on.scene_fd.rigid_solver.get_state().qpos[:, 0].abs() <= 4.5).all() - # 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. + # Backward, with mixed per-environment activity: env 0 drives into the active |x|=4 limit while env 1 stays well + # inside it, so the adjoint solve faces different constraint counts in the same batch. Sanity-check the split. on.scene_fd.reset() - on.entity_fd.set_dofs_velocity(gs.tensor([100.0], dtype=gs.tc_float)) + on.entity_fd.set_dofs_velocity([[100.0], [2.0]]) for _ in range(5): on.scene_fd.step() - assert abs(on.scene_fd.rigid_solver.get_state().qpos[0, 0]) > 3.5 + qpos_end = on.scene_fd.rigid_solver.get_state().qpos + assert abs(qpos_end[0, 0]) > 3.5 + assert abs(qpos_end[1, 0]) < 1.0 assert_grad_matches_fd( on, - [np.array([100.0])], + [np.array([[100.0], [2.0]])], lambda e, x: e.set_dofs_velocity(x), - lambda scene, entity: scene.rigid_solver.get_state().qpos[0, 0] ** 2, + lambda scene, entity: (scene.rigid_solver.get_state().qpos[:, 0] ** 2).sum(), 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, + n_envs=2, substeps=4, dt=1.0 / 60.0, gravity=(0.0, 0.0, 0.0), @@ -90,10 +82,10 @@ def test_joint_limit_grad_matches_fd(grad_slider_limit, precision, show_viewer): 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) + v = gs.tensor([[0.5], [0.3]], 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 = (pair.scene_ana.rigid_solver.get_state().qpos[:, 0] ** 2).sum() loss.backward() grads[key] = tensor_to_array(v.grad) assert_allclose( @@ -131,7 +123,7 @@ def test_per_step_force_into_limit_grad_matches_fd(model_name, request, precisio def setup_fn(scene, entity): if init_pos is not None: - entity.set_dofs_position(gs.tensor(init_pos, dtype=gs.tc_float)) + entity.set_dofs_position(init_pos) def loss_fn(scene, entity): state = scene.rigid_solver.get_state() @@ -140,7 +132,7 @@ def loss_fn(scene, entity): 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.entity_fd.control_dofs_force(force) 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})" @@ -161,6 +153,7 @@ def loss_fn(scene, entity): def test_frictionloss_grad_matches_fd(grad_revolute_frictionloss, precision, show_viewer): pair = make_diff_scene_pair( grad_revolute_frictionloss, + n_envs=2, substeps=4, dt=1.0 / 60.0, gravity=(0.0, 0.0, 0.0), @@ -171,13 +164,13 @@ def test_frictionloss_grad_matches_fd(grad_revolute_frictionloss, precision, sho 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 (qd_to_torch(cs.n_constraints_frictionloss) == 1).all() assert_grad_matches_fd( pair, - [np.array([2.0])], + [np.array([[2.0], [1.2]])], lambda e, x: e.set_dofs_velocity(x), - lambda scene, entity: scene.rigid_solver.get_state().qpos[0, 0] ** 2, + lambda scene, entity: (scene.rigid_solver.get_state().qpos[:, 0] ** 2).sum(), n_steps=10, rtol=1e-10 if precision == "64" else 5e-5, atol=1e-10 if precision == "64" else 5e-5, @@ -198,6 +191,7 @@ def test_frictionloss_grad_matches_fd(grad_revolute_frictionloss, precision, sho def test_equality_grad_matches_fd(model_name, n_rows, request, precision, show_viewer): pair = make_diff_scene_pair( request.getfixturevalue(model_name), + n_envs=2, substeps=4, dt=1.0 / 60.0, gravity=(0.0, 0.0, 0.0), @@ -208,17 +202,19 @@ def test_equality_grad_matches_fd(model_name, n_rows, request, precision, show_v 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 (qd_to_torch(cs.n_constraints_equality) == n_rows).all() # Large initial velocities: the anchor velocity-product bias entering aref is quadratic in velocity, so its # adjoint contribution only clears the fp64 tolerance band when the joints spin fast. + def loss_fn(scene, entity): + qpos = scene.rigid_solver.get_state().qpos + return (qpos[:, 0] ** 2 + 0.7 * qpos[:, 1] ** 2).sum() + assert_grad_matches_fd( pair, - [np.array([4.0, -2.5])], + [np.array([[4.0, -2.5], [-3.0, 2.0]])], 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 - ), + loss_fn, n_steps=10, rtol=1e-10 if precision == "64" else 5e-5, atol=1e-10 if precision == "64" else 5e-5, @@ -245,15 +241,15 @@ def test_all_constraint_groups_grad_matches_fd(grad_all_eq_fric, precision, show 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]) + weights = np.array([1.0, 0.7, 1.3, 0.5, 0.9, 1.1, 0.6]) 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)) + return sum(weights[d] * qpos[d] ** 2 for d in range(7)) assert_grad_matches_fd( pair, - [np.array([0.8, -0.3, 0.5, -0.2, 0.4, -0.6])], + [np.array([0.8, -0.3, 0.5, -0.2, 0.2, -0.3, 0.4])], lambda e, x: e.set_dofs_velocity(x), loss_fn, n_steps=10, diff --git a/tests/grad/test_rigid_dynamics.py b/tests/grad/test_rigid_dynamics.py index 0013cc4297..7ec937ee4a 100644 --- a/tests/grad/test_rigid_dynamics.py +++ b/tests/grad/test_rigid_dynamics.py @@ -1,3 +1,4 @@ +import math import sys import numpy as np @@ -94,10 +95,15 @@ def test_fk_grad_matches_fd(model_name, request, precision, show_viewer): else: step_input = rng.standard_normal((B, n_dofs)) + def apply_vel_per_env(e, x): + # Same-step per-environment commands: each call must keep its own tape slot and gradient path. + e.set_dofs_velocity(x[:1], envs_idx=[0]) + e.set_dofs_velocity(x[1:], envs_idx=[1]) + 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), + "vel": apply_vel_per_env, "force": lambda e, x: e.control_dofs_force(x), }[setter] @@ -195,3 +201,37 @@ def apply_force(entity, force): atol=fp64_tol if precision == "64" else fp32_tol, eps=3e-5 if precision == "64" else 3e-2, ) + + +@pytest.mark.required +@pytest.mark.parametrize("control_mode", ["position", "velocity"]) +def test_per_step_pd_target_grad_matches_fd(control_mode, grad_revolute, precision, show_viewer): + # Per-step-varying PD targets are scenario commands, replayed by the backward unroll; the gradient of a tracked + # initial velocity through the controlled rollout must match finite differences. + pair = make_diff_scene_pair( + grad_revolute, + substeps=4, + show_viewer=show_viewer, + ) + for entity in (pair.entity_ana, pair.entity_fd): + entity.set_dofs_kp(4.0) + entity.set_dofs_kv(0.8) + targets = [0.3 * math.sin(0.7 * t) for t in range(10)] + + def step_fn(entity, i_step): + if control_mode == "position": + entity.control_dofs_position(targets[i_step]) + else: + entity.control_dofs_velocity(targets[i_step]) + + assert_grad_matches_fd( + pair, + [np.array([1.5])], + lambda e, x: e.set_dofs_velocity(x), + lambda scene, entity: scene.rigid_solver.get_state().qpos[0, 0] ** 2, + n_steps=10, + step_fn=step_fn, + rtol=1e-10 if precision == "64" else 5e-5, + atol=1e-10 if precision == "64" else 5e-5, + eps=3e-5 if precision == "64" else 3e-2, + ) diff --git a/tests/grad/test_rigid_optim.py b/tests/grad/test_rigid_optim.py index 787d6be475..8d49d3b6f6 100644 --- a/tests/grad/test_rigid_optim.py +++ b/tests/grad/test_rigid_optim.py @@ -42,13 +42,13 @@ def test_reference_trajectory_recovery_converges(control_target, grad_cartpole, 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)) + robot_ref.set_dofs_velocity(target_ctrl) 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)) + robot_ref.control_dofs_force(target_ctrl[t]) scene_ref.step() ref_state = scene_ref.rigid_solver.get_state() target_qpos = ref_state.qpos.detach().clone() diff --git a/tests/grad/utils.py b/tests/grad/utils.py index f645da0efd..8b8c8502b7 100644 --- a/tests/grad/utils.py +++ b/tests/grad/utils.py @@ -69,16 +69,17 @@ def make_diff_scene_pair( return DiffScenePair(scenes.get(True), entities.get(True), scenes.get(False), entities.get(False)) -def assert_grad_matches_fd(pair, inputs, apply_fn, loss_fn, *, rtol, atol, eps, n_steps=None, setup_fn=None): +def assert_grad_matches_fd( + pair, inputs, apply_fn, loss_fn, *, rtol, atol, eps, n_steps=None, setup_fn=None, step_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.""" + Input i of `inputs` is applied via `apply_fn(entity, x)` before step i and receives its own adjoint. The scene + runs `n_steps` steps (default len(inputs)); extra steps run without re-applying. `setup_fn(scene, entity)` runs + once after reset for untracked initialization; `step_fn(entity, i_step)` runs before every step on both scenes + for undifferentiated per-step scenario commands (e.g. PD targets). The FD reference perturbs each input entry in + turn, so the cost is O(n_steps * total input size). rtol / atol / eps are per-scenario, pinned to the 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 @@ -92,6 +93,8 @@ def assert_grad_matches_fd(pair, inputs, apply_fn, loss_fn, *, rtol, atol, eps, x = gs.tensor(base[i_step], dtype=gs.tc_float, requires_grad=True) x_anas.append(x) apply_fn(pair.entity_ana, x) + if step_fn is not None: + step_fn(pair.entity_ana, i_step) 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" @@ -116,7 +119,9 @@ def assert_grad_matches_fd(pair, inputs, apply_fn, loss_fn, *, rtol, atol, eps, 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)) + apply_fn(pair.entity_fd, inp) + if step_fn is not None: + step_fn(pair.entity_fd, i_step) 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)