[BUG FIX] Fix camera follow entity.

genesis/utils/geom.py

@@ -1400,20 +1400,21 @@ def rotvec_to_quat(rotvec: np.ndarray, out: np.ndarray | None = None) -> np.ndar
                 and returned, which is slower.
     """
     assert rotvec.ndim >= 1
+    B = rotvec.shape[:-1]
     if out is None:
-        out_ = np.empty((*rotvec.shape[:-1], 4), dtype=rotvec.dtype)
+        out_ = np.empty((*B, 4), dtype=rotvec.dtype)
     else:
-        assert out.shape == (*rotvec.shape[:-1], 4)
+        assert out.shape == (*B, 4)
         out_ = out
 
-    # Compute unit axis and positive angle separately
-    angle = np.sqrt(np.sum(np.square(rotvec), -1))
+    # Split unit axis and positive angle
+    angle = np.sqrt(np.sum(np.square(rotvec.reshape((-1, 3))), -1)).reshape(B)
     # FIXME: Taylor expansion should be used to handle angle ~ 0.0
-    axis = rotvec / np.maximum(angle, gs.EPS)
+    axis = rotvec / np.maximum(angle[..., None], gs.EPS)
 
     # Compute the quaternion representation
     out_[..., 0] = np.cos(0.5 * angle)
-    out_[..., 1:] = np.sin(0.5 * angle) * axis
+    out_[..., 1:] = np.sin(0.5 * angle[..., None]) * axis
 
     return out_
 

genesis/vis/batch_renderer.py

@@ -305,8 +305,8 @@ def build(self):
             lights_intensity_tensor=_make_tensor([light.intensity for light in self._lights]),
         )
 
-    def update_scene(self):
-        self._visualizer._context.update()
+    def update_scene(self, force_render: bool = False):
+        self._visualizer._context.update(force_render)
 
     def render(self, rgb=True, depth=False, segmentation=False, normal=False, antialiasing=False, force_render=False):
         """
@@ -356,7 +356,7 @@ def render(self, rgb=True, depth=False, segmentation=False, normal=False, antial
             return tuple(arr if req else None for req, arr in zip(request, cached))
 
         # Update scene
-        self.update_scene()
+        self.update_scene(force_render)
 
         # Render only what is needed (flags still passed to renderer)
         cameras_pos = torch.stack([camera.get_pos() for camera in self._cameras], dim=1)

genesis/vis/camera.py

@@ -129,6 +129,7 @@ def __init__(
         self._attached_offset_T = None
 
         self._followed_entity = None
+        self._follow_pos_rel = None
         self._follow_fixed_axis = None
         self._follow_smoothing = None
         self._follow_fix_orientation = None
@@ -281,7 +282,23 @@ def follow_entity(self, entity, fixed_axis=(None, None, None), smoothing=None, f
         if self._attached_link is not None:
             gs.raise_exception("Impossible to following an entity with a camera that is already attached.")
 
+        if self._is_built:
+            env_idx = self._env_idx if self._is_batched and self._env_idx is not None else ()
+            pos_rel = self._pos[env_idx] - entity.get_pos(self._env_idx, unsafe=True)
+            if self._env_idx is not None:
+                pos_rel = pos_rel.squeeze(0)
+        else:
+            pos_rel = self._initial_pos - torch.tensor(entity.base_link.pos, dtype=gs.tc_float, device=gs.device)
+            if self._env_idx is None:
+                if self._visualizer._context.rendered_envs_idx is not None:
+                    pos_rel = pos_rel.expand((len(self._visualizer._context.rendered_envs_idx), 1))
+                else:
+                    # Falling back to adding batch dimension to allow broadcasting.
+                    # Note that it is not possible to expand / tile because the batch size is unknown before build.
+                    pos_rel = pos_rel.unsqueeze(0)
+
         self._followed_entity = entity
+        self._follow_pos_rel = pos_rel
         self._follow_fixed_axis = fixed_axis
         self._follow_smoothing = smoothing
         self._follow_fix_orientation = fix_orientation
@@ -293,6 +310,7 @@ def unfollow_entity(self):
         Calling this method has no effect if the camera is not currently following any entity.
         """
         self._followed_entity = None
+        self._follow_pos_rel = None
         self._follow_fixed_axis = None
         self._follow_smoothing = None
         self._follow_fix_orientation = None
@@ -315,16 +333,23 @@ def update_following(self):
             camera_lookat = self._lookat[env_idx].clone()
             camera_pos = self._pos[env_idx].clone()
 
-        # Smooth camera movement with a low-pass filter, in particular Exponential Moving Average (EMA) if requested
+        # Query entity and relative camera positions
         entity_pos = self._followed_entity.get_pos(self._env_idx, unsafe=True)
-        camera_pos -= self._initial_pos
+        follow_pos_rel = self._follow_pos_rel
+        if not self._is_batched:
+            follow_pos_rel = follow_pos_rel.squeeze(0)
+
+        # Smooth camera movement with a low-pass filter, in particular Exponential Moving Average (EMA) if requested
+        camera_pos -= follow_pos_rel
         if self._follow_smoothing is not None:
             camera_pos[:] = self._follow_smoothing * camera_pos + (1.0 - self._follow_smoothing) * entity_pos
             if not self._follow_fix_orientation:
                 camera_lookat[:] = self._follow_smoothing * camera_lookat + (1.0 - self._follow_smoothing) * entity_pos
         else:
             camera_pos[:] = entity_pos
-        camera_pos += self._initial_pos
+            camera_lookat[:] = entity_pos
+
+        camera_pos += follow_pos_rel
 
         # Fix the camera's position along the specified axis if requested
         for i_a, fixed_axis in enumerate(self._follow_fixed_axis):
@@ -428,16 +453,16 @@ def render(
             )
         elif self._raytracer is not None:
             if rgb_:
-                self._raytracer.update_scene()
+                self._raytracer.update_scene(force_render)
                 rgb_arr = self._raytracer.render_camera(self)
 
             if depth or segmentation or normal:
-                self._rasterizer.update_scene()
+                self._rasterizer.update_scene(force_render)
                 _, depth_arr, seg_idxc_arr, normal_arr = self._rasterizer.render_camera(
                     self, False, depth, segmentation, normal=normal
                 )
         else:
-            self._rasterizer.update_scene()
+            self._rasterizer.update_scene(force_render)
             rgb_arr, depth_arr, seg_idxc_arr, normal_arr = self._rasterizer.render_camera(
                 self, rgb_, depth, segmentation, normal=normal
             )
@@ -526,7 +551,7 @@ def render_pointcloud(self, world_frame=True):
             # FIXME: Avoid converting to numpy
             depth_arr = tensor_to_array(depth_arr)
         else:
-            self._rasterizer.update_scene()
+            self._rasterizer.update_scene(force_render=False)
             _, depth_arr, _, _ = self._rasterizer.render_camera(
                 self, rgb=False, depth=True, segmentation=False, normal=False
             )

genesis/vis/rasterizer.py

@@ -142,8 +142,8 @@ def render_camera(self, camera, rgb=True, depth=False, segmentation=False, norma
             normal_arr = retval[int(rgb + depth)]
         return rgb_arr, depth_arr, seg_idxc_arr, normal_arr
 
-    def update_scene(self):
-        self._context.update()
+    def update_scene(self, force_render: bool = False):
+        self._context.update(force_render)
 
     def destroy(self):
         for node in self._camera_nodes.values():

genesis/vis/rasterizer_context.py

@@ -923,9 +923,9 @@ def clear_debug_object(self, obj):
     def clear_debug_objects(self):
         self.clear_external_nodes()
 
-    def update(self):
+    def update(self, force_render: bool = False):
         # Early return if already updated previously
-        if self._t >= self.scene._t:
+        if not force_render and self._t >= self.scene._t:
             return
 
         self._t = self.scene._t
@@ -934,7 +934,7 @@ def update(self):
         self.clear_dynamic_nodes()
 
         # update variables not used in simulation
-        self.visualizer.update_visual_states()
+        self.visualizer.update_visual_states(force_render)
 
         # Reset scene bounds to trigger recomputation. They are involved in shadow map
         self._scene._bounds = None

genesis/vis/raytracer.py

@@ -640,8 +640,8 @@ def update_camera(self, camera):
     def reset(self):
         self._t = -1
 
-    def update_scene(self):
-        if self._t >= self.scene.t:
+    def update_scene(self, force_render: bool = False):
+        if not force_render and self._t >= self.scene.t:
             if self.camera_updated:
                 self._scene.update_scene(time=self._t)
                 self.camera_updated = False
@@ -651,7 +651,7 @@ def update_scene(self):
         self._t = self.scene.t
 
         # update variables not used in simulation
-        self.visualizer.update_visual_states()
+        self.visualizer.update_visual_states(force_render)
 
         # tool entities
         if self.sim.tool_solver.is_active():

genesis/vis/visualizer.py

@@ -200,12 +200,12 @@ def update(self, force=True, auto=None):
             else:
                 gs.raise_exception("Viewer closed.")
 
-    def update_visual_states(self):
+    def update_visual_states(self, force_render: bool = False):
         """
         Update all visualization-only variables here.
         """
         # Early return if already updated previously
-        if self._t >= self.scene._t:
+        if not force_render and self._t >= self.scene._t:
             return
 
         for camera in self._cameras:

tests/test_render.py

@@ -332,7 +332,7 @@ def test_render_api_advanced(tmp_path, n_envs, show_viewer, png_snapshot, render
         )
         cam_1 = scene.add_camera(
             res=CAM_RES,
-            pos=(1.5, -0.5, 1.5),
+            pos=(0.8, -0.5, 0.8),
             lookat=(0.0, 0.0, 0.5),
             fov=45,
             near=0.05,
@@ -589,6 +589,79 @@ def test_segmentation_map(segmentation_level, particle_mode, renderer_type, rend
         assert_array_equal(np.sort(np.unique(seg.flat)), np.arange(0, seg_num))
 
 
+@pytest.mark.parametrize("renderer_type", [RENDERER_TYPE.RASTERIZER])
+@pytest.mark.parametrize("n_envs", [0, 2])
+def test_camera_follow_entity(n_envs, show_viewer):
+    CAM_RES = (100, 100)
+
+    scene = gs.Scene(
+        vis_options=gs.options.VisOptions(
+            rendered_envs_idx=[1] if n_envs else None,
+            segmentation_level="entity",
+        ),
+        show_viewer=False,
+    )
+    for pos in ((1.0, 0.0, 0.0), (-1.0, 0.0, 0.0), (0.0, 1.0, 0.0), (0.0, -1.0, 0.0)):
+        obj = scene.add_entity(
+            gs.morphs.Box(
+                size=(0.1, 0.1, 0.1),
+                pos=pos,
+            ),
+        )
+        cam = scene.add_camera(
+            res=CAM_RES,
+            pos=(0.0, 0.0, 0.0),
+            lookat=(1.0, 0, 0.0),
+            env_idx=1 if n_envs else None,
+            GUI=show_viewer,
+        )
+        cam.follow_entity(obj, smoothing=None)
+
+    scene.build(n_envs=n_envs)
+
+    # First render
+    seg_mask = None
+    for entity_idx, cam in enumerate(scene.visualizer.cameras, 1):
+        _, _, seg, _ = cam.render(rgb=False, segmentation=True)
+        assert (np.unique(seg) == (0, entity_idx)).all()
+        if seg_mask is None:
+            seg_mask = seg != 0
+        else:
+            assert ((seg != 0) == seg_mask).all()
+
+    # Second render - same
+    for i, obj in enumerate(scene.entities):
+        obj.set_pos((10.0, 0.0, i), envs_idx=([1] if n_envs else None))
+    force_render = True
+    for entity_idx, cam in enumerate(scene.visualizer.cameras, 1):
+        _, _, seg, _ = cam.render(rgb=False, segmentation=True, force_render=force_render)
+        assert (np.unique(seg) == (0, entity_idx)).all()
+        assert ((seg != 0) == seg_mask).all()
+        force_render = False
+
+    # Third render - All objects but all different
+    for i, obj in enumerate(scene.entities):
+        obj.set_pos((0.1 * ((i // 2) % 2 - 1), 0.1 * (i % 2), 0.1 * i), envs_idx=([1] if n_envs else None))
+    force_render = True
+    seg_masks = []
+    for cam in scene.visualizer.cameras:
+        _, _, seg, _ = cam.render(rgb=False, segmentation=True, force_render=force_render)
+        assert (np.unique(seg) == np.arange(len(scene.entities) + 1)).all()
+        seg_masks.append(seg != 0)
+        force_render = False
+    assert np.diff(seg_masks, axis=0).any(axis=(1, 2)).all()
+
+    # Track a trajectory over time
+    for i in range(3):
+        pos = 2.0 * (np.random.rand(3) - 0.5)
+        quat = gu.rotvec_to_quat(np.pi * (np.random.rand(3) - 0.5))
+        obj.set_pos(pos + np.array([10.0, 0.0, 0.0]), envs_idx=([1] if n_envs else None))
+        obj.set_quat(quat, envs_idx=([1] if n_envs else None))
+        _, _, seg, _ = cam.render(segmentation=True, force_render=True)
+        assert (np.unique(seg) == (0, entity_idx)).all()
+        assert not seg[tuple([*range(0, res // 3), *range(2 * res // 3, res)] for res in CAM_RES)].any()
+
+
 @pytest.mark.required
 @pytest.mark.parametrize(
     "renderer_type",

tests/utils.py

@@ -32,7 +32,7 @@
 DEFAULT_BRANCH_NAME = "main"
 
 HUGGINGFACE_ASSETS_REVISION = "16e4eae0024312b84518f4b555dd630d6b34095a"
-HUGGINGFACE_SNAPSHOT_REVISION = "a6fd3b99364b927dd5367488e58cd251f254fa94"
+HUGGINGFACE_SNAPSHOT_REVISION = "0db0ca5941d6b64c58d9e9711abe62e3a50738ac"
 
 MESH_EXTENSIONS = (".mtl", *MESH_FORMATS, *GLTF_FORMATS, *USD_FORMATS)
 IMAGE_EXTENSIONS = (".png", ".jpg")