Merge branch 'main' into solver-uvs

Author: alelievr

genesis/engine/sensors/base_sensor.py

@@ -48,7 +48,7 @@ class SharedSensorMetadata:
     """
 
     cache_sizes: list[int] = field(default_factory=list)
-    delays_ts: torch.Tensor = make_tensor_field((0, 0), dtype_factory=lambda: gs.tc_int)
+    delays_ts: torch.Tensor = make_tensor_field((0, 0), dtype=gs.tc_int)
 
     def __del__(self):
         try:
@@ -329,7 +329,7 @@ class RigidSensorMetadataMixin:
     """
 
     solver: "RigidSolver | None" = None
-    links_idx: torch.Tensor = make_tensor_field((0,), dtype_factory=lambda: gs.tc_int)
+    links_idx: torch.Tensor = make_tensor_field((0,), dtype=gs.tc_int)
     offsets_pos: torch.Tensor = make_tensor_field((0, 0, 3))
     offsets_quat: torch.Tensor = make_tensor_field((0, 0, 4))
 

genesis/engine/sensors/contact_force.py

@@ -76,7 +76,7 @@ class ContactSensorMetadata(SharedSensorMetadata):
     """
 
     solver: "RigidSolver | None" = None
-    expanded_links_idx: torch.Tensor = make_tensor_field((0,), dtype_factory=lambda: gs.tc_int)
+    expanded_links_idx: torch.Tensor = make_tensor_field((0,), dtype=gs.tc_int)
 
 
 @register_sensor(ContactSensorOptions, ContactSensorMetadata, tuple)

genesis/engine/sensors/imu.py

@@ -70,9 +70,9 @@ class IMUSharedMetadata(RigidSensorMetadataMixin, NoisySensorMetadataMixin, Shar
 
     alignment_rot_matrix: torch.Tensor = make_tensor_field((0, 0, 3, 3))
     magnetic_field_vector: torch.Tensor = make_tensor_field((0, 0, 3))  # added another dimension to match data layout
-    acc_indices: torch.Tensor = make_tensor_field((0, 0), dtype_factory=lambda: gs.tc_int)
-    gyro_indices: torch.Tensor = make_tensor_field((0, 0), dtype_factory=lambda: gs.tc_int)
-    mag_indices: torch.Tensor = make_tensor_field((0, 0), dtype_factory=lambda: gs.tc_int)
+    acc_indices: torch.Tensor = make_tensor_field((0, 0), dtype=gs.tc_int)
+    gyro_indices: torch.Tensor = make_tensor_field((0, 0), dtype=gs.tc_int)
+    mag_indices: torch.Tensor = make_tensor_field((0, 0), dtype=gs.tc_int)
 
 
 class IMUData(NamedTuple):

genesis/engine/sensors/raycaster.py

@@ -293,18 +293,18 @@ class RaycasterSharedMetadata(RigidSensorMetadataMixin, SharedSensorMetadata):
     min_ranges: torch.Tensor = make_tensor_field((0,))
     max_ranges: torch.Tensor = make_tensor_field((0,))
     no_hit_values: torch.Tensor = make_tensor_field((0,))
-    return_world_frame: torch.Tensor = make_tensor_field((0,), dtype_factory=lambda: gs.tc_bool)
+    return_world_frame: torch.Tensor = make_tensor_field((0,), dtype=gs.tc_bool)
 
     patterns: list[RaycastPattern] = field(default_factory=list)
     ray_dirs: torch.Tensor = make_tensor_field((0, 3))
     ray_starts: torch.Tensor = make_tensor_field((0, 3))
     ray_starts_world: torch.Tensor = make_tensor_field((0, 3))
     ray_dirs_world: torch.Tensor = make_tensor_field((0, 3))
 
-    points_to_sensor_idx: torch.Tensor = make_tensor_field((0,), dtype_factory=lambda: gs.tc_int)
-    sensor_cache_offsets: torch.Tensor = make_tensor_field((0,), dtype_factory=lambda: gs.tc_int)
-    sensor_point_offsets: torch.Tensor = make_tensor_field((0,), dtype_factory=lambda: gs.tc_int)
-    sensor_point_counts: torch.Tensor = make_tensor_field((0,), dtype_factory=lambda: gs.tc_int)
+    points_to_sensor_idx: torch.Tensor = make_tensor_field((0,), dtype=gs.tc_int)
+    sensor_cache_offsets: torch.Tensor = make_tensor_field((0,), dtype=gs.tc_int)
+    sensor_point_offsets: torch.Tensor = make_tensor_field((0,), dtype=gs.tc_int)
+    sensor_point_counts: torch.Tensor = make_tensor_field((0,), dtype=gs.tc_int)
 
 
 class RaycasterData(NamedTuple):

genesis/engine/solvers/rigid/collider/gjk.py

@@ -223,10 +223,10 @@ def func_gjk_contact(
             )
 
             if shrink_sphere:
-                # If we shrinked the sphere and capsule to point and line and the distance is larger than the
-                # collision epsilon, it means a shallow penetration. Thus we subtract the radius of the sphere and
-                # the capsule to get the actual distance. If the distance is smaller than the collision epsilon, it
-                # means a deep penetration, which requires the default GJK handling.
+                # If we shrunk the sphere and capsule to point and line and the distance is larger than the collision
+                # epsilon, it means a shallow penetration. Thus we subtract the radius of the sphere and the capsule to
+                # get the actual distance. If the distance is smaller than the collision epsilon, it means a deep
+                # penetration, which requires the default GJK handling.
                 if distance > gjk_info.collision_eps[None]:
                     radius_a, radius_b = 0.0, 0.0
                     if is_sphere_swept_geom_a:

genesis/engine/solvers/rigid/collider/narrowphase.py

@@ -614,7 +614,7 @@ def func_convex_convex_contact(
 
             if multi_contact and is_col_0:
                 # Perturbation axis must not be aligned with the principal axes of inertia the geometry,
-                # otherwise it would be more sensitive to ill-conditionning.
+                # otherwise it would be more sensitive to ill-conditioning.
                 axis = (2 * (i_detection % 2) - 1) * axis_0 + (1 - 2 * ((i_detection // 2) % 2)) * axis_1
                 qrot = gu.ti_rotvec_to_quat(collider_info.mc_perturbation[None] * axis, EPS)
                 func_rotate_frame(i_ga, contact_pos_0, qrot, i_b, geoms_state, geoms_info)
@@ -863,8 +863,8 @@ def func_convex_convex_contact(
                     # First-order correction of the normal direction.
                     # The way the contact normal gets twisted by applying perturbation of geometry poses is
                     # unpredictable as it depends on the final portal discovered by MPR. Alternatively, let compute
-                    # the mininal rotation that makes the corrected twisted normal as closed as possible to the
-                    # original one, up to the scale of the perturbation, then apply first-order Taylor expension of
+                    # the minimal rotation that makes the corrected twisted normal as closed as possible to the
+                    # original one, up to the scale of the perturbation, then apply first-order Taylor expansion of
                     # Rodrigues' rotation formula.
                     twist_rotvec = ti.math.clamp(
                         normal.cross(normal_0),

genesis/utils/misc.py

@@ -343,24 +343,22 @@ def concat_with_tensor(
     return torch.cat([tensor, value], dim=dim)
 
 
-def make_tensor_field(shape: tuple[int, ...] = (), dtype_factory: Callable[[], torch.dtype] | None = None):
+def make_tensor_field(shape: tuple[int, ...] = (), dtype: torch.dtype | None = None):
     """
     Helper method to create a tensor field for dataclasses.
 
     Parameters
     ----------
     shape : tuple
         The shape of the tensor field. It must have zero elements, otherwise it will trigger an exception.
-    dtype_factory : Callable[[], torch.dtype], optional
-        The factory function to create the dtype of the tensor field. Default is gs.tc_float.
-        A factory is used because gs types may not be available at the time of field creation.
+    dtype : torch.dtype, optional
+        Data type of the tensor field. Default is gs.tc_float.
     """
     assert not shape or math.prod(shape) == 0
 
     def _default_factory():
-        nonlocal shape, dtype_factory
-        dtype = dtype_factory() if dtype_factory is not None else gs.tc_float
-        return torch.empty(shape, dtype=dtype, device=gs.device)
+        nonlocal shape, dtype
+        return torch.empty(shape, dtype=dtype or gs.tc_float, device=gs.device)
 
     return field(default_factory=_default_factory)