[BUG FIX] Revert removal of 'zero_velocity' optional argument to 'set_pos', 'set_quat'. (#2047)

* Revert removal of 'zero_velocity' optional argument to 'set_pos', 'set_quat'.
* Rename 'n_equalities_candidate' in 'n_candidate_equalities'.
* Add support of default values in data_oriented. Fix typing.

Author: duburcqa

genesis/engine/entities/rigid_entity/rigid_entity.py

@@ -1998,7 +1998,7 @@ def get_links_invweight(self, links_idx_local=None, envs_idx=None, *, unsafe=Fal
         return self._solver.get_links_invweight(links_idx, envs_idx, unsafe=unsafe)
 
     @gs.assert_built
-    def set_pos(self, pos, envs_idx=None, *, relative=False, unsafe=False):
+    def set_pos(self, pos, envs_idx=None, *, zero_velocity=True, relative=False, unsafe=False):
         """
         Set position of the entity's base link.
 
@@ -2020,13 +2020,14 @@ def set_pos(self, pos, envs_idx=None, *, relative=False, unsafe=False):
             if _pos is not pos:
                 gs.logger.debug(ALLOCATE_TENSOR_WARNING)
             pos = _pos
-        self._solver.set_dofs_velocity(None, self._dofs_idx, envs_idx, skip_forward=True, unsafe=unsafe)
+        if zero_velocity:
+            self._solver.set_dofs_velocity(None, self._dofs_idx, envs_idx, skip_forward=True, unsafe=unsafe)
         self._solver.set_base_links_pos(
             pos.unsqueeze(-2), self._base_links_idx_, envs_idx, relative=relative, unsafe=unsafe
         )
 
     @gs.assert_built
-    def set_quat(self, quat, envs_idx=None, *, relative=False, unsafe=False):
+    def set_quat(self, quat, envs_idx=None, *, zero_velocity=True, relative=False, unsafe=False):
         """
         Set quaternion of the entity's base link.
 
@@ -2048,7 +2049,8 @@ def set_quat(self, quat, envs_idx=None, *, relative=False, unsafe=False):
             if _quat is not quat:
                 gs.logger.debug(ALLOCATE_TENSOR_WARNING)
             quat = _quat
-        self._solver.set_dofs_velocity(None, self._dofs_idx, envs_idx, skip_forward=True, unsafe=unsafe)
+        if zero_velocity:
+            self._solver.set_dofs_velocity(None, self._dofs_idx, envs_idx, skip_forward=True, unsafe=unsafe)
         self._solver.set_base_links_quat(
             quat.unsqueeze(-2), self._base_links_idx_, envs_idx, relative=relative, unsafe=unsafe
         )

genesis/engine/solvers/rigid/constraint_solver_decomp.py

@@ -41,7 +41,7 @@ def __init__(self, rigid_solver: "RigidSolver"):
             4 * rigid_solver.collider._collider_info.max_contact_pairs[None]
             + sum(joint.type in (gs.JOINT_TYPE.REVOLUTE, gs.JOINT_TYPE.PRISMATIC) for joint in self._solver.joints)
             + self._solver.n_dofs
-            + self._solver.n_equalities_candidate * 6
+            + self._solver.n_candidate_equalities_ * 6
         )
         self.len_constraints_ = max(1, self.len_constraints)
 
@@ -351,7 +351,7 @@ def add_weld_constraint(self, link1_idx, link2_idx, envs_idx=None, *, unsafe=Fal
         if overflow:
             gs.logger.warning(
                 "Ignoring dynamically registered weld constraint to avoid exceeding max number of equality constraints"
-                f"({self.rigid_global_info.n_equalities_candidate.to_numpy()}). Please increase the value of "
+                f"({self.rigid_global_info.n_candidate_equalities.to_numpy()}). Please increase the value of "
                 "RigidSolver's option 'max_dynamic_constraints'."
             )
 
@@ -2190,7 +2190,7 @@ def kernel_add_weld_constraint(
     for i_b_ in ti.ndrange(envs_idx.shape[0]):
         i_b = envs_idx[i_b_]
         i_e = constraint_state.ti_n_equalities[i_b]
-        if i_e == rigid_global_info.n_equalities_candidate[None]:
+        if i_e == rigid_global_info.n_candidate_equalities[None]:
             overflow = True
         else:
             shared_pos = links_state.pos[link1_idx, i_b]

genesis/engine/solvers/rigid/rigid_solver_decomp.py

@@ -224,8 +224,7 @@ def build(self):
         self.n_entities_ = max(1, self.n_entities)
         self.n_free_verts_ = max(1, self.n_free_verts)
         self.n_fixed_verts_ = max(1, self.n_fixed_verts)
-
-        self.n_equalities_candidate = max(1, self.n_equalities + self._options.max_dynamic_constraints)
+        self.n_candidate_equalities_ = max(1, self.n_equalities + self._options.max_dynamic_constraints)
 
         # FIXME: AvatarSolver should not inherit from RigidSolver, not to mention that it is completely broken...
         is_rigid_solver = type(self) is RigidSolver
@@ -250,16 +249,7 @@ def build(self):
             self._static_rigid_sim_config = array_class.StructRigidSimStaticConfig(
                 para_level=self.sim._para_level,
                 requires_grad=self.sim.options.requires_grad,
-                use_hibernation=False,
-                batch_links_info=False,
-                batch_dofs_info=False,
-                batch_joints_info=False,
-                enable_mujoco_compatibility=False,
-                enable_multi_contact=True,
                 enable_collision=self._enable_collision,
-                enable_joint_limit=False,
-                box_box_detection=True,
-                sparse_solve=False,
                 integrator=gs.integrator.approximate_implicitfast,
                 solver_type=gs.constraint_solver.CG,
             )

genesis/utils/array_class.py

@@ -1,6 +1,7 @@
 import math
 import dataclasses
 from functools import partial
+from typing_extensions import dataclass_transform  # Made it into standard lib from Python 3.12
 
 import gstaichi as ti
 import numpy as np
@@ -24,35 +25,38 @@ def maybe_shape(shape, is_on):
     return shape if is_on else ()
 
 
+@dataclass_transform(eq_default=True, order_default=True, kw_only_default=False, frozen_default=True)
 class AutoInitMeta(type):
     def __new__(cls, name, bases, namespace):
-        field_names = namespace["__annotations__"].keys()
+        names = tuple(namespace["__annotations__"].keys())
+        defaults = {k: namespace[k] for k in names if k in namespace}
 
         def __init__(self, *args, **kwargs):
-            assigned = {}
+            # Initialize assigned arguments from defaults
+            assigned = defaults.copy()
 
             # Assign positional arguments
-            if len(args) > len(field_names):
-                raise TypeError(f"{name}() takes {len(field_names)} positional arguments but {len(args)} were given")
-            for field, value in zip(field_names, args):
-                assigned[field] = value
+            if len(args) > len(names):
+                raise TypeError(f"{name}() takes {len(names)} positional arguments but {len(args)} were given")
+            for key, value in zip(names, args):
+                assigned[key] = value
 
             # Assign keyword arguments
             for key, value in kwargs.items():
-                if key not in field_names:
+                if key not in names:
                     raise TypeError(f"{name}() got unexpected keyword argument '{key}'")
-                if key in assigned:
+                if key in names[: len(args)]:
                     raise TypeError(f"{name}() got multiple values for argument '{key}'")
                 assigned[key] = value
 
             # Check for missing arguments
-            for field in field_names:
-                if field not in assigned:
-                    raise TypeError(f"{name}() missing required argument: '{field}'")
+            for key in names:
+                if key not in assigned:
+                    raise TypeError(f"{name}() missing required argument: '{key}'")
 
             # Set attributes
-            for field, value in assigned.items():
-                setattr(self, field, value)
+            for key, value in assigned.items():
+                setattr(self, key, value)
 
         namespace["__init__"] = __init__
 
@@ -100,7 +104,7 @@ class StructRigidGlobalInfo(metaclass=BASE_METACLASS):
     noslip_iterations: V_ANNOTATION
     noslip_tolerance: V_ANNOTATION
     n_equalities: V_ANNOTATION
-    n_equalities_candidate: V_ANNOTATION
+    n_candidate_equalities: V_ANNOTATION
     hibernation_thresh_acc: V_ANNOTATION
     hibernation_thresh_vel: V_ANNOTATION
     EPS: V_ANNOTATION
@@ -142,7 +146,7 @@ def get_rigid_global_info(solver):
         noslip_iterations=V_SCALAR_FROM(dtype=gs.ti_int, value=solver._options.noslip_iterations),
         noslip_tolerance=V_SCALAR_FROM(dtype=gs.ti_float, value=solver._options.noslip_tolerance),
         n_equalities=V_SCALAR_FROM(dtype=gs.ti_int, value=solver._n_equalities),
-        n_equalities_candidate=V_SCALAR_FROM(dtype=gs.ti_int, value=solver.n_equalities_candidate),
+        n_candidate_equalities=V_SCALAR_FROM(dtype=gs.ti_int, value=solver.n_candidate_equalities_),
         hibernation_thresh_acc=V_SCALAR_FROM(dtype=gs.ti_float, value=solver._hibernation_thresh_acc),
         hibernation_thresh_vel=V_SCALAR_FROM(dtype=gs.ti_float, value=solver._hibernation_thresh_vel),
         EPS=V_SCALAR_FROM(dtype=gs.ti_float, value=gs.EPS),
@@ -1658,7 +1662,7 @@ class StructEqualitiesInfo(metaclass=BASE_METACLASS):
 
 
 def get_equalities_info(solver):
-    shape = (solver.n_equalities_candidate, solver._B)
+    shape = (solver.n_candidate_equalities_, solver._B)
 
     return StructEqualitiesInfo(
         eq_obj1id=V(dtype=gs.ti_int, shape=shape),
@@ -1725,19 +1729,19 @@ def get_entities_state(solver):
 @ti.data_oriented
 class StructRigidSimStaticConfig(metaclass=AutoInitMeta):
     para_level: int
-    requires_grad: bool
-    use_hibernation: bool
-    batch_links_info: bool
-    batch_dofs_info: bool
-    batch_joints_info: bool
-    enable_mujoco_compatibility: bool
-    enable_multi_contact: bool
     enable_collision: bool
-    enable_joint_limit: bool
-    box_box_detection: bool
-    sparse_solve: bool
-    integrator: int
-    solver_type: int
+    use_hibernation: bool = False
+    batch_links_info: bool = False
+    batch_dofs_info: bool = False
+    batch_joints_info: bool = False
+    enable_mujoco_compatibility: bool = False
+    enable_multi_contact: bool = False
+    enable_joint_limit: bool = False
+    box_box_detection: bool = True
+    sparse_solve: bool = False
+    integrator: int = gs.integrator.approximate_implicitfast
+    solver_type: int = gs.constraint_solver.CG
+    requires_grad: bool = False
 
 
 # =========================================== DataManager ===========================================

tests/test_rigid_physics.py

@@ -2508,8 +2508,8 @@ def test_urdf_capsule(tmp_path, show_viewer, tol):
     for _ in range(40):
         scene.step()
     geom_verts = tensor_to_array(geom.get_verts())
-    assert np.linalg.norm(geom_verts - np.zeros(3), axis=-1, ord=np.inf).min() < 1e-3
-    assert np.linalg.norm(geom_verts - np.array((0.0, 0.0, 0.14)), axis=-1, ord=np.inf).min() < 1e-3
+    assert np.linalg.norm(geom_verts - (0.0, 0.0, 0.0), axis=-1, ord=np.inf).min() < 1e-3
+    assert np.linalg.norm(geom_verts - (0.0, 0.0, 0.14), axis=-1, ord=np.inf).min() < 1e-3
 
 
 @pytest.mark.required