Repository: Genesis-Embodied-AI/Genesis
File: genesis/vis/rasterizer_context.py (Python)

Code before:
                pos = tool_entity.pos[self.sim.cur_substep_local].to_numpy()
                quat = tool_entity.quat[self.sim.cur_substep_local].to_numpy()
                pose = gu.trans_quat_to_T(pos, quat)
                self.set_node_pose(self.static_nodes[tool_entity.uid], pose=pose)

    def set_reflection_mat(self, geom_T):
        height = geom_T[0, 2, 3]
        self.jit.reflection_mat = np.array(
            [
                [1.0, 0.0, 0.0, 0.0],
                [0.0, 1.0, 0.0, 0.0],
                [0.0, 0.0, -1.0, height * 2],
                [0.0, 0.0, 0.0, 1.0],
            ],
            dtype=np.float32,
        )

    def on_rigid(self):
        if self.sim.rigid_solver.is_active():
            # TODO: support dynamic switching in GUI later
            for rigid_entity in self.sim.rigid_solver.entities:
                if rigid_entity.surface.vis_mode == "visual":
                    geoms = rigid_entity.vgeoms
                    geoms_T = self.sim.rigid_solver._vgeoms_render_T
                else:
                    geoms = rigid_entity.geoms
                    geoms_T = self.sim.rigid_solver._geoms_render_T

                for geom in geoms:
                    if "sdf" in rigid_entity.surface.vis_mode:
                        mesh = geom.get_sdf_trimesh()
                    else:
                        mesh = geom.get_trimesh()
                    geom_T = geoms_T[geom.idx][self.rendered_envs_idx]
                    self.rigid_nodes[geom.uid] = self.add_node(
                        pyrender.Mesh.from_trimesh(
                            mesh=mesh,
                            poses=geom_T,
                            smooth=geom.surface.smooth if "collision" not in rigid_entity.surface.vis_mode else False,
                            double_sided=(
                                geom.surface.double_sided if "collision" not in rigid_entity.surface.vis_mode else False
                            ),
                            is_floor=isinstance(rigid_entity._morph, gs.morphs.Plane),
                            env_shared=not self.env_separate_rigid,
                        )
                    )
                    if isinstance(rigid_entity._morph, gs.morphs.Plane):
                        self.set_reflection_mat(geom_T)

    def update_rigid(self, buffer_updates):
        if self.sim.rigid_solver.is_active():

Review comment (performance):
Beware this would trigger a copy since ‘self.rendered_envs_idx’ is not a slice but a list. No sure about the overall performance penalty it induces.

Code after:
            [
                [1.0, 0.0, 0.0, 0.0],
                [0.0, 1.0, 0.0, 0.0],
                [0.0, 0.0, -1.0, height * 2],
                [0.0, 0.0, 0.0, 1.0],
            ],
            dtype=np.float32,
        )

    def on_rigid(self):
        if self.sim.rigid_solver.is_active():
            # TODO: support dynamic switching in GUI later
            for rigid_entity in self.sim.rigid_solver.entities:
                if rigid_entity.surface.vis_mode == "visual":
                    geoms = rigid_entity.vgeoms
                    geoms_T = self.sim.rigid_solver._vgeoms_render_T
                else:
                    geoms = rigid_entity.geoms
                    geoms_T = self.sim.rigid_solver._geoms_render_T

                for geom in geoms:
                    if "sdf" in rigid_entity.surface.vis_mode:
                        mesh = geom.get_sdf_trimesh()
                    else:
                        mesh = geom.get_trimesh()
                    geom_T = geoms_T[geom.idx][self.rendered_envs_idx]
                    self.rigid_nodes[geom.uid] = self.add_node(
                        pyrender.Mesh.from_trimesh(
                            mesh=mesh,
                            poses=geom_T,
                            smooth=geom.surface.smooth if "collision" not in rigid_entity.surface.vis_mode else False,
                            double_sided=(
                                geom.surface.double_sided if "collision" not in rigid_entity.surface.vis_mode else False
                            ),
                            is_floor=isinstance(rigid_entity._morph, gs.morphs.Plane),
                            env_shared=not self.env_separate_rigid,
                        )
                    )
                    if isinstance(rigid_entity._morph, gs.morphs.Plane):
                        self.set_reflection_mat(geom_T)

    def update_rigid(self, buffer_updates):
        if self.sim.rigid_solver.is_active():
            for rigid_entity in self.sim.rigid_solver.entities:
                if rigid_entity.surface.vis_mode == "visual":
                    geoms = rigid_entity.vgeoms
                    geoms_T = self.sim.rigid_solver._vgeoms_render_T
                else:
                    geoms = rigid_entity.geoms
                    geoms_T = self.sim.rigid_solver._geoms_render_T