From d32d38e7337fcfa69f19941933f43f3906979ee3 Mon Sep 17 00:00:00 2001 From: Alexis Duburcq Date: Sun, 5 Oct 2025 21:16:34 +0200 Subject: [PATCH] Add support of render optional arg 'force_render' for all viewer backends. Fix camera follow entity. --- genesis/utils/geom.py | 13 +++--- genesis/vis/batch_renderer.py | 6 +-- genesis/vis/camera.py | 39 +++++++++++++--- genesis/vis/rasterizer.py | 4 +- genesis/vis/rasterizer_context.py | 6 +-- genesis/vis/raytracer.py | 6 +-- genesis/vis/visualizer.py | 4 +- tests/test_render.py | 75 ++++++++++++++++++++++++++++++- tests/utils.py | 2 +- 9 files changed, 127 insertions(+), 28 deletions(-) diff --git a/genesis/utils/geom.py b/genesis/utils/geom.py index a5ade4978..3c46b2c30 100644 --- a/genesis/utils/geom.py +++ b/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_ diff --git a/genesis/vis/batch_renderer.py b/genesis/vis/batch_renderer.py index 5c25a80c5..1f5e63eed 100644 --- a/genesis/vis/batch_renderer.py +++ b/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) diff --git a/genesis/vis/camera.py b/genesis/vis/camera.py index aeb6d7202..9f1ef5a95 100644 --- a/genesis/vis/camera.py +++ b/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 ) diff --git a/genesis/vis/rasterizer.py b/genesis/vis/rasterizer.py index 6dd832c34..632f3075d 100644 --- a/genesis/vis/rasterizer.py +++ b/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(): diff --git a/genesis/vis/rasterizer_context.py b/genesis/vis/rasterizer_context.py index 98563ecea..c3fa3a856 100644 --- a/genesis/vis/rasterizer_context.py +++ b/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 diff --git a/genesis/vis/raytracer.py b/genesis/vis/raytracer.py index f2e3f60d6..829d901b7 100644 --- a/genesis/vis/raytracer.py +++ b/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(): diff --git a/genesis/vis/visualizer.py b/genesis/vis/visualizer.py index 4f95396f8..1f2c14b48 100644 --- a/genesis/vis/visualizer.py +++ b/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: diff --git a/tests/test_render.py b/tests/test_render.py index b72505e03..f69fff0f3 100644 --- a/tests/test_render.py +++ b/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", diff --git a/tests/utils.py b/tests/utils.py index 3c483e5cc..e8c0c5640 100644 --- a/tests/utils.py +++ b/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")