[FEATURE] Add full support of tensor broadcasting in getters. (#2051)

* Remove tensor allocation warnings.
* Unify tensor sanitization.
* Remove 'unsafe' support now that zero-copy is inherently safe.
* Raise exception if invalid markers at passed in unit test expr.
* Add fully support of (enhanced) tensor broadcasting in setters (w/ and w/o zero-copy).
* Rename 'verts_idx' in 'verts_local_idx' for clarity.
* More robust batched camera and sensor debug draw.

Author: duburcqa

examples/coupling/fem_cube_linked_with_arm.py

@@ -115,10 +115,7 @@ def main():
 
         print("cube init pos", cube.init_positions)
         pin_idx = [1, 5]
-        cube.set_vertex_constraints(
-            verts_idx=pin_idx,
-            link=end_joint.link,
-        )
+        cube.set_vertex_constraints(verts_idx_local=pin_idx, link=end_joint.link)
         print("Cube initial positions:", cube.init_positions[pin_idx])
         scene.draw_debug_spheres(poss=cube.init_positions[pin_idx], radius=0.02, color=(1.0, 0.0, 1.0, 0.8))
 

examples/fem_hard_and_soft_constraint.py

@@ -106,30 +106,19 @@ def get_next_circle_position():
     try:
         target_positions = blob.init_positions[pinned_idx]
         scene.draw_debug_spheres(poss=target_positions, radius=0.02, color=(1, 0, 1, 0.8))
-        blob.set_vertex_constraints(
-            verts_idx=pinned_idx,
-            target_poss=target_positions,
-            is_soft_constraint=True,
-            stiffness=1e4,
-        )
+        blob.set_vertex_constraints(pinned_idx, target_positions, is_soft_constraint=True, stiffness=1e4)
 
         target_positions = get_next_circle_position()
         debug_circle = scene.draw_debug_spheres(poss=target_positions, radius=0.02, color=(0, 1, 0, 0.8))
-        cube.set_vertex_constraints(
-            verts_idx=pinned_idx,
-            target_poss=target_positions,
-        )
+        cube.set_vertex_constraints(pinned_idx, target_positions)
 
         for step in tqdm(range(total_steps), total=total_steps):
             if debug_circle is not None:
                 scene.clear_debug_object(debug_circle)
 
             new_pos = get_next_circle_position()
             debug_circle = scene.draw_debug_spheres(poss=new_pos, radius=0.02, color=(0, 1, 0, 0.8))
-            cube.update_constraint_targets(
-                verts_idx=pinned_idx,
-                target_poss=new_pos,
-            )
+            cube.update_constraint_targets(pinned_idx, new_pos)
 
             scene.step()
 

examples/sap_coupling/fem_fixed_constraint.py

@@ -63,7 +63,7 @@ def main():
         target_poss = cube.init_positions[verts_idx] + torch.tensor(
             (0.15 * (math.cos(0.04 * i) - 1.0), 0.15 * math.sin(0.04 * i), 0.0)
         )
-        cube.set_vertex_constraints(verts_idx=verts_idx, target_poss=target_poss)
+        cube.set_vertex_constraints(verts_idx, target_poss)
         scene.step(update_visualizer=False)
         if args.vis:
             scene.visualizer.context.draw_debug_sphere(

genesis/__init__.py

@@ -33,7 +33,6 @@
 from .logging import Logger
 from .version import __version__
 from .utils import redirect_libc_stderr, set_random_seed, get_platform, get_device
-from .utils.misc import ALLOCATE_TENSOR_WARNING
 
 
 # Global state
@@ -109,9 +108,6 @@ def init(
     logger.info(f"~<│{wave}>~ ~~~~<Genesis>~~~~ ~<{wave}│>~")
     logger.info(f"~<╰{'─'*(bar_width)}╯>~")
 
-    # FIXME: Disable this warning for now, because it is not useful without printing the entire traceback
-    logger.addFilter(lambda record: record.msg != ALLOCATE_TENSOR_WARNING)
-
     # Get concrete device and backend
     global device
     device, device_name, total_mem, backend = get_device(backend)

genesis/engine/entities/drone_entity.py

@@ -5,7 +5,7 @@
 import gstaichi as ti
 
 import genesis as gs
-import genesis.utils.misc as mu
+from genesis.utils.misc import get_assets_dir, broadcast_tensor
 
 from .rigid_entity import RigidEntity
 
@@ -16,14 +16,14 @@ def _load_scene(self, morph, surface):
         super()._load_scene(morph, surface)
 
         # additional drone specific attributes
-        properties = ET.parse(os.path.join(mu.get_assets_dir(), morph.file)).getroot()[0].attrib
+        properties = ET.parse(os.path.join(get_assets_dir(), morph.file)).getroot()[0].attrib
         self._KF = float(properties["kf"])
         self._KM = float(properties["km"])
 
         self._n_propellers = len(morph.propellers_link_name)
 
         propellers_link = gs.List([self.get_link(name) for name in morph.propellers_link_name])
-        self._propellers_link_idxs = torch.tensor(
+        self._propellers_link_idx = torch.tensor(
             [link.idx for link in propellers_link], dtype=gs.tc_int, device=gs.device
         )
         try:
@@ -64,18 +64,22 @@ def set_propellels_rpm(self, propellels_rpm):
             gs.raise_exception("`set_propellels_rpm` can only be called once per step.")
         self._prev_prop_t = self.sim.cur_step_global
 
-        propellels_rpm = self.solver._process_dim(
-            torch.as_tensor(propellels_rpm, dtype=gs.tc_float, device=gs.device)
-        ).T.contiguous()
-        if len(propellels_rpm) != len(self._propellers_link_idxs):
-            gs.raise_exception("Last dimension of `propellels_rpm` does not match `entity.n_propellers`.")
-        if torch.any(propellels_rpm < 0):
-            gs.raise_exception("`propellels_rpm` cannot be negative.")
+        assert propellels_rpm is not None
+        propellels_rpm, *_ = self._solver._sanitize_io_variables(
+            propellels_rpm, self._propellers_link_idx, self._n_propellers, "propellers_link_idx"
+        )
+        if self._scene.n_envs == 0:
+            propellels_rpm = propellels_rpm[None]
+
+        # FIXME: This check is too expensive
+        # if (propellels_rpm < 0.0).any():
+        #     gs.raise_exception("`propellels_rpm` cannot be negative.")
+
         self._propellers_revs = (self._propellers_revs + propellels_rpm) % (60 / self.solver.dt)
 
         self.solver.set_drone_rpm(
             self._n_propellers,
-            self._propellers_link_idxs,
+            self._propellers_link_idx,
             propellels_rpm,
             self._propellers_spin,
             self.KF,
@@ -122,7 +126,7 @@ def COM_link_idx(self):
     @property
     def propellers_idx(self):
         """The indices of the drone's propeller links."""
-        return self._propellers_link_idxs
+        return self._propellers_link_idx
 
     @property
     def propellers_spin(self):

genesis/engine/entities/fem_entity.py

@@ -13,7 +13,7 @@
 from genesis.engine.couplers import SAPCoupler
 from genesis.engine.states.cache import QueriedStates
 from genesis.engine.states.entities import FEMEntityState
-from genesis.utils.misc import ALLOCATE_TENSOR_WARNING, to_gs_tensor, tensor_to_array
+from genesis.utils.misc import to_gs_tensor, tensor_to_array, broadcast_tensor
 
 from .base_entity import Entity
 
@@ -126,6 +126,34 @@ def __init__(self, scene, solver, material, morph, surface, idx, v_start=0, el_s
     # ----------------------------------- basic entity ops -------------------------------
     # ------------------------------------------------------------------------------------
 
+    def _sanitize_verts_idx_local(self, verts_idx_local=None, envs_idx=None):
+        if verts_idx_local is None:
+            verts_idx_local = range(self.n_vertices)
+
+        if envs_idx is None:
+            verts_idx_local_ = broadcast_tensor(verts_idx_local, gs.tc_int, (-1,), ("envs_idx",))
+        else:
+            verts_idx_local_ = broadcast_tensor(verts_idx_local, gs.tc_int, (len(envs_idx), -1), ("envs_idx", ""))
+
+        # FIXME: This check is too expensive
+        # if not (0 <= verts_idx_local_ & verts_idx_local_ < self.n_vertices).all():
+        #     gs.raise_exception("Elements of `verts_idx_local' are out-of-range.")
+
+        return verts_idx_local_.contiguous()
+
+    def _sanitize_verts_tensor(self, tensor, dtype, verts_idx=None, envs_idx=None, element_shape=(), *, batched=True):
+        n_vertices = verts_idx.shape[-1] if verts_idx is not None else self.n_vertices
+        if batched:
+            assert envs_idx is not None
+            batch_shape = (len(envs_idx), n_vertices)
+            dim_names = ("envs_idx", "verts_idx", *("" for _ in element_shape))
+        else:
+            batch_shape = (n_vertices,)
+            dim_names = ("verts_idx", *("" for _ in element_shape))
+        tensor_shape = (*batch_shape, *element_shape)
+
+        return broadcast_tensor(tensor, dtype, tensor_shape, dim_names).contiguous()
+
     def set_position(self, pos):
         """
         Set the target position(s) for the FEM entity.
@@ -153,14 +181,14 @@ def set_position(self, pos):
         if pos.ndim == 1:
             if pos.shape == (3,):
                 pos = self.init_positions_COM_offset + pos
-                self._tgt["pos"] = pos.unsqueeze(0).tile((self._sim._B, 1, 1))
+                self._tgt["pos"] = pos[None].tile((self._sim._B, 1, 1))
                 is_valid = True
         elif pos.ndim == 2:
             if pos.shape == (self.n_vertices, 3):
-                self._tgt["pos"] = pos.unsqueeze(0).tile((self._sim._B, 1, 1))
+                self._tgt["pos"] = pos[None].tile((self._sim._B, 1, 1))
                 is_valid = True
             elif pos.shape == (self._sim._B, 3):
-                pos = self.init_positions_COM_offset.unsqueeze(0) + pos.unsqueeze(1)
+                pos = self.init_positions_COM_offset[None] + pos[:, None]
                 self._tgt["pos"] = pos
                 is_valid = True
         elif pos.ndim == 3:
@@ -200,10 +228,10 @@ def set_velocity(self, vel):
                 is_valid = True
         elif vel.ndim == 2:
             if vel.shape == (self.n_vertices, 3):
-                self._tgt["vel"] = vel.unsqueeze(0).tile((self._sim._B, 1, 1))
+                self._tgt["vel"] = vel[None].tile((self._sim._B, 1, 1))
                 is_valid = True
             elif vel.shape == (self._sim._B, 3):
-                self._tgt["vel"] = vel.unsqueeze(1).tile((1, self.n_vertices, 1))
+                self._tgt["vel"] = vel[:, None].tile((1, self.n_vertices, 1))
                 is_valid = True
         elif vel.ndim == 3:
             if vel.shape == (self._sim._B, self.n_vertices, 3):
@@ -241,7 +269,7 @@ def set_actuation(self, actu):
             is_valid = True
         elif actu.ndim == 1:
             if actu.shape == (n_groups,):
-                self._tgt["actu"] = actu.unsqueeze(0).tile((self._sim._B, 1))
+                self._tgt["actu"] = actu[None].tile((self._sim._B, 1))
                 is_valid = True
             elif actu.shape == (self.n_elements,):
                 gs.raise_exception("Cannot set per-element actuation.")
@@ -814,46 +842,16 @@ def set_muscle_direction(self, muscle_direction):
             muscle_direction=muscle_direction,
         )
 
-    def _sanitize_input_tensor(self, tensor, dtype, unbatched_ndim=1):
-        _tensor = torch.as_tensor(tensor, dtype=dtype, device=gs.device)
-
-        if _tensor.ndim < unbatched_ndim + 1:
-            _tensor = _tensor.repeat((self._sim._B, *((1,) * max(1, _tensor.ndim))))
-            if self._sim._B > 1:
-                gs.logger.debug(ALLOCATE_TENSOR_WARNING)
-        else:
-            _tensor = _tensor.contiguous()
-            if _tensor is not tensor:
-                gs.logger.debug(ALLOCATE_TENSOR_WARNING)
-
-            if len(_tensor) != self._sim._B:
-                gs.raise_exception("Input tensor batch size must match the number of environments.")
-
-        if _tensor.ndim != unbatched_ndim + 1:
-            gs.raise_exception(f"Input tensor ndim is {_tensor.ndim}, should be {unbatched_ndim + 1}.")
-
-        return _tensor
-
-    def _sanitize_input_verts_idx(self, verts_idx_local):
-        verts_idx = self._sanitize_input_tensor(verts_idx_local, dtype=gs.tc_int, unbatched_ndim=1) + self._v_start
-        assert ((verts_idx >= 0) & (verts_idx < self._solver.n_vertices)).all(), "Vertex indices out of bounds."
-        return verts_idx
-
-    def _sanitize_input_poss(self, poss):
-        poss = self._sanitize_input_tensor(poss, dtype=gs.tc_float, unbatched_ndim=2)
-        assert poss.ndim == 3 and poss.shape[2] == 3, "Position tensor must have shape (B, num_verts, 3)."
-        return poss
-
     def set_vertex_constraints(
-        self, verts_idx, target_poss=None, link=None, is_soft_constraint=False, stiffness=0.0, envs_idx=None
+        self, verts_idx_local, target_poss=None, link=None, is_soft_constraint=False, stiffness=0.0, envs_idx=None
     ):
         """
         Set vertex constraints for specified vertices.
 
         Parameters
         ----------
-            verts_idx : array_like
-                List of vertex indices to constrain.
+            verts_idx_local : array_like
+                List of local vertex indices to constrain.
             target_poss : array_like, shape (len(verts_idx), 3), optional
                 List of target positions [x, y, z] for each vertex. If not provided, the initial positions are used.
             link : RigidLink
@@ -874,30 +872,22 @@ def set_vertex_constraints(
         if not self._solver._constraints_initialized:
             self._solver.init_constraints()
 
+        use_current_poss = target_poss is None
         envs_idx = self._scene._sanitize_envs_idx(envs_idx)
-        verts_idx = self._sanitize_input_verts_idx(verts_idx)
+        verts_idx_local = self._sanitize_verts_idx_local(verts_idx_local, envs_idx)
+        verts_idx = verts_idx_local + self._v_start
+        target_poss = self._sanitize_verts_tensor(target_poss, gs.tc_float, verts_idx, envs_idx, (3,))
 
-        if target_poss is None:
-            target_poss = torch.zeros(
-                (verts_idx.shape[0], verts_idx.shape[1], 3), dtype=gs.tc_float, device=gs.device, requires_grad=False
-            )
+        if use_current_poss:
             self._kernel_get_verts_pos(self._sim.cur_substep_local, target_poss, verts_idx)
-        target_poss = self._sanitize_input_poss(target_poss)
-
-        assert (
-            len(envs_idx) == len(target_poss) == len(verts_idx)
-        ), "First dimension should match number of environments."
-        assert target_poss.shape[1] == verts_idx.shape[1], "Target position should be provided for each vertex."
 
         if link is None:
             link_init_pos = torch.zeros((self._sim._B, 3), dtype=gs.tc_float, device=gs.device)
             link_init_quat = torch.zeros((self._sim._B, 4), dtype=gs.tc_float, device=gs.device)
             link_idx = -1
         else:
             assert isinstance(link, RigidLink), "Only RigidLink is supported for vertex constraints."
-            link_init_pos = self._sanitize_input_tensor(link.get_pos(), dtype=gs.tc_float)
-            link_init_quat = self._sanitize_input_tensor(link.get_quat(), dtype=gs.tc_float)
-            link_idx = link.idx
+            link_init_pos, link_init_quat, link_idx = link.get_pos(), link.get_quat(), link.idx
 
         self._solver._kernel_set_vertex_constraints(
             self._sim.cur_substep_local,
@@ -911,31 +901,36 @@ def set_vertex_constraints(
             envs_idx,
         )
 
-    def update_constraint_targets(self, verts_idx, target_poss, envs_idx=None):
+    def update_constraint_targets(self, verts_idx_local, target_poss, envs_idx=None):
         """Update target positions for existing constraints."""
         if not self._solver._constraints_initialized:
             gs.logger.warning("Ignoring update_constraint_targets; constraints have not been initialized.")
             return
 
+        assert target_poss is not None
         envs_idx = self._scene._sanitize_envs_idx(envs_idx)
-        verts_idx = self._sanitize_input_verts_idx(verts_idx)
-        target_poss = self._sanitize_input_poss(target_poss)
-        assert target_poss.shape[1] == verts_idx.shape[1], "Target position should be provided for each vertex."
+        verts_idx_local = self._sanitize_verts_idx_local(verts_idx_local, envs_idx)
+        verts_idx = verts_idx_local + self._v_start
+        target_poss = self._sanitize_verts_tensor(target_poss, gs.tc_float, verts_idx, envs_idx, (3,))
 
         self._solver._kernel_update_constraint_targets(verts_idx, target_poss, envs_idx)
 
-    def remove_vertex_constraints(self, verts_idx=None, envs_idx=None):
+    def remove_vertex_constraints(self, verts_idx_local=None, envs_idx=None):
         """Remove constraints from specified vertices, or all if None."""
         if not self._solver._constraints_initialized:
             gs.logger.warning("Ignoring remove_vertex_constraints; constraints have not been initialized.")
             return
 
-        if verts_idx is None:
+        # FIXME: GsTaichi 'fill' method is very inefficient. Try using zero-copy if possible.
+        if verts_idx_local is None:
             self._solver.vertex_constraints.is_constrained.fill(0)
-        else:
-            verts_idx = self._sanitize_input_verts_idx(verts_idx)
-            envs_idx = self._scene._sanitize_envs_idx(envs_idx)
-            self._solver._kernel_remove_specific_constraints(verts_idx, envs_idx)
+            return
+
+        envs_idx = self._scene._sanitize_envs_idx(envs_idx)
+        verts_idx_local = self._sanitize_verts_idx_local(verts_idx_local, envs_idx)
+        verts_idx = verts_idx_local + self._v_start
+
+        self._solver._kernel_remove_specific_constraints(verts_idx, envs_idx)
 
     @ti.kernel
     def _kernel_get_verts_pos(self, f: ti.i32, pos: ti.types.ndarray(), verts_idx: ti.types.ndarray()):
@@ -954,14 +949,8 @@ def get_el2v(self):
         el2v : gs.Tensor
             Tensor of shape (n_elements, 4) mapping each element to its vertex indices.
         """
-
         el2v = gs.zeros((self.n_elements, 4), dtype=int, requires_grad=False, scene=self.scene)
-        self._solver._kernel_get_el2v(
-            element_el_start=self._el_start,
-            n_elements=self.n_elements,
-            el2v=el2v,
-        )
-
+        self._solver._kernel_get_el2v(element_el_start=self._el_start, n_elements=self.n_elements, el2v=el2v)
         return el2v
 
     @ti.kernel