Fix 'delete_weld_constraint'.

Author: duburcqa

genesis/engine/solvers/rigid/rigid_solver_decomp.py

@@ -4968,19 +4968,14 @@ def _kernel_add_weld_constraint(
                 self.constraint_solver.ti_n_equalities[i_b] = self.constraint_solver.ti_n_equalities[i_b] + 1
 
     def delete_weld_constraint(self, link1_idx, link2_idx, envs_idx=None, *, unsafe=False):
-        _, link1_idx, _ = self._sanitize_1D_io_variables(
-            None, link1_idx, self.n_links, envs_idx, idx_name="links_idx", skip_allocation=True, unsafe=unsafe
-        )
-        _, link2_idx, envs_idx = self._sanitize_1D_io_variables(
-            None, link2_idx, self.n_links, envs_idx, idx_name="links_idx", skip_allocation=True, unsafe=unsafe
-        )
-        self._kernel_delete_weld_constraint(link1_idx, link2_idx, envs_idx)
+        envs_idx = self._sanitize_envs_idx(envs_idx, unsafe=unsafe)
+        self._kernel_delete_weld_constraint(int(link1_idx), int(link2_idx), envs_idx)
 
     @ti.kernel
     def _kernel_delete_weld_constraint(
         self,
-        link1_idx: ti.types.ndarray(),
-        link2_idx: ti.types.ndarray(),
+        link1_idx: ti.i32,
+        link2_idx: ti.i32,
         envs_idx: ti.types.ndarray(),
     ):
         ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL)
@@ -4989,8 +4984,8 @@ def _kernel_delete_weld_constraint(
             for i_e in range(self.n_equalities, self.constraint_solver.ti_n_equalities[i_b]):
                 if (
                     self.equalities_info[i_e, i_b].eq_type == gs.EQUALITY_TYPE.WELD
-                    and self.equalities_info[i_e, i_b].eq_obj1id == link1_idx[i_b]
-                    and self.equalities_info[i_e, i_b].eq_obj2id == link2_idx[i_b]
+                    and self.equalities_info[i_e, i_b].eq_obj1id == link1_idx
+                    and self.equalities_info[i_e, i_b].eq_obj2id == link2_idx
                 ):
                     if i_e < self.constraint_solver.ti_n_equalities[i_b] - 1:
                         self.equalities_info[i_e, i_b] = self.equalities_info[

tests/test_rigid_physics.py

@@ -408,6 +408,74 @@ def test_equality_weld(gs_sim, mj_sim, gs_solver):
     simulate_and_check_mujoco_consistency(gs_sim, mj_sim, qpos, num_steps=300, tol=tol)
 
 
+def test_dynamic_weld(show_viewer, tol):
+    scene = gs.Scene(
+        show_viewer=show_viewer,
+        show_FPS=False,
+    )
+    plane = scene.add_entity(
+        gs.morphs.Plane(),
+    )
+    cube = scene.add_entity(
+        gs.morphs.Box(
+            size=(0.04, 0.04, 0.04),
+            pos=(0.65, 0.0, 0.02),
+        ),
+        surface=gs.surfaces.Plastic(color=(1, 0, 0)),
+    )
+    robot = scene.add_entity(
+        gs.morphs.MJCF(
+            file="xml/universal_robots_ur5e/ur5e.xml",
+        ),
+    )
+    scene.build(n_envs=4, env_spacing=(3.0, 3.0))
+
+    end_effector = robot.get_link("ee_virtual_link")
+
+    # Compute up and down robot configurations
+    ee_pos_up = np.array((0.65, 0.0, 0.5), dtype=gs.np_float)
+    ee_pos_down = np.array((0.65, 0.0, 0.15), dtype=gs.np_float)
+    qpos_up = robot.inverse_kinematics(
+        link=end_effector,
+        pos=np.tile(ee_pos_up, (4, 1)),
+        quat=np.tile(np.array((0.0, 1.0, 0.0, 0.0), dtype=gs.np_float), (4, 1)),
+    )
+    qpos_down = robot.inverse_kinematics(
+        link=end_effector,
+        pos=np.tile(ee_pos_down, (4, 1)),
+        quat=np.tile(np.array((0.0, 1.0, 0.0, 0.0), dtype=gs.np_float), (4, 1)),
+    )
+
+    # move to pre-grasp pose
+    robot.control_dofs_position(qpos_up)
+    for i in range(120):
+        scene.step()
+
+    # reach
+    robot.control_dofs_position(qpos_down)
+    for i in range(70):
+        scene.step()
+
+    # add weld constraint and move back up
+    scene.sim.rigid_solver.add_weld_constraint(cube.base_link.idx, end_effector.idx, envs_idx=(0, 1, 2))
+    robot.control_dofs_position(qpos_up)
+    for i in range(60):
+        scene.step()
+    cubes_pos, cubes_quat = cube.get_pos(), cube.get_quat()
+    assert_allclose(torch.diff(cubes_quat, dim=0), 0.0, tol=1e-3)
+    assert_allclose(torch.diff(cubes_pos[[0, 1, 2]], dim=0), 0.0, tol=tol)
+    assert_allclose(cubes_pos[-1] - cubes_pos[0], ee_pos_down - ee_pos_up, tol=1e-2)
+
+    # drop
+    scene.sim.rigid_solver.delete_weld_constraint(cube.base_link.idx, end_effector.idx, envs_idx=(0, 1))
+    for i in range(110):
+        scene.step()
+    cubes_pos, cubes_quat = cube.get_pos(), cube.get_quat()
+    assert_allclose(torch.diff(cubes_quat, dim=0), 0.0, tol=1e-3)
+    assert_allclose(torch.diff(cubes_pos[[0, 1, 3]], dim=0), 0.0, tol=1e-2)
+    assert_allclose(cubes_pos[2] - cubes_pos[0], ee_pos_up - ee_pos_down, tol=1e-3)
+
+
 @pytest.mark.required
 @pytest.mark.parametrize("xml_path", ["xml/one_ball_joint.xml"])
 @pytest.mark.parametrize("gs_solver", [gs.constraint_solver.CG, gs.constraint_solver.Newton])
@@ -2300,7 +2368,7 @@ def _check_entity_positions(relative, tol):
     entities[(False, True)].set_dofs_velocity(door_vel, 6)
     entities[(True, False)].set_dofs_velocity(door_vel)
     entities[(True, True)].set_dofs_velocity(door_vel)
-    link_1 = np.array([entities[(True, True)].link_start], dtype=gs.np_int)
+    link_1 = entities[(True, True)].link_start
     for key in ((False, False), (False, True)):
         link_2 = entities[key].link_start
         scene.rigid_solver.add_weld_constraint(link_1, link_2)