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

Code before:
        self.link_frame_shown = False
        self.camera_frustum_shown = False

        self.world_frame_mesh = mu.create_frame(
            origin_radius=0.012,
            axis_radius=0.005,
            axis_length=self.world_frame_size,
            head_radius=0.01,
            head_length=0.03,
        )

        self.link_frame_mesh = trimesh.creation.axis(origin_size=0.03, axis_radius=0.025, axis_length=1.0)
        self.link_frame_mesh.visual.face_colors[:, :3] = (0.7 * self.link_frame_mesh.visual.face_colors[:, :3]).astype(
            int
        )
        self.link_frame_mesh.vertices *= self.link_frame_size

    def build(self, scene):
        self.scene = scene
        self.sim = scene.sim
        self.visualizer = scene.visualizer
        self.visualizer.update_visual_states()

        if self.rendered_envs_idx is None:
            self.rendered_envs_idx = list(range(self.sim._B))
        self.rendered_envs_idx_slice = slice(self.rendered_envs_idx[0], self.rendered_envs_idx[-1] + 1)

        # pyrender scene
        self._scene = pyrender.Scene(
            ambient_light=self.ambient_light,
            bg_color=self.background_color,
            n_envs=len(self.rendered_envs_idx),
        )

        if gs.platform != "Windows":
            self.jit = JITRenderer(self._scene, [], [])
        else:
            from genesis.ext.pyrender.non_jit_renderer import SimpleNonJITRenderer

            self.jit = SimpleNonJITRenderer(self._scene, [], [])

        # nodes
        self.world_frame_node = None
        self.link_frame_nodes = dict()
        self.frustum_nodes = dict()  # nodes camera frustums
        self.rigid_nodes = dict()
        self.static_nodes = dict()  # used across all frames
        self.dynamic_nodes = list()  # nodes that live within single frame
        self.external_nodes = dict()  # nodes added by external user

        self.on_lights()

Review comment (suggestion):
I don't think the assumption of consecutive `rendered_envs_idx` is good. What we should do instead is to measure the overhead introduced by list indexing and verify it won't add more overhead compared to the current version. 

Code after:
        self.world_frame_mesh = mu.create_frame(
            origin_radius=0.012,
            axis_radius=0.005,
            axis_length=self.world_frame_size,
            head_radius=0.01,
            head_length=0.03,
        )

        self.link_frame_mesh = trimesh.creation.axis(origin_size=0.03, axis_radius=0.025, axis_length=1.0)
        self.link_frame_mesh.visual.face_colors[:, :3] = (0.7 * self.link_frame_mesh.visual.face_colors[:, :3]).astype(
            int
        )
        self.link_frame_mesh.vertices *= self.link_frame_size

    def build(self, scene):
        self.scene = scene
        self.sim = scene.sim
        self.visualizer = scene.visualizer
        self.visualizer.update_visual_states()

        if self.rendered_envs_idx is None:
            self.rendered_envs_idx = list(range(self.sim._B))

        # pyrender scene
        self._scene = pyrender.Scene(
            ambient_light=self.ambient_light,
            bg_color=self.background_color,
            n_envs=len(self.rendered_envs_idx),
        )

        self.jit = JITRenderer(self._scene, [], [])

        # nodes
        self.world_frame_node = None
        self.link_frame_nodes = dict()
        self.frustum_nodes = dict()  # nodes camera frustums
        self.rigid_nodes = dict()
        self.static_nodes = dict()  # used across all frames
        self.dynamic_nodes = list()  # nodes that live within single frame
        self.external_nodes = dict()  # nodes added by external user

        self.on_lights()

        if self.show_world_frame:
            self.on_world_frame()
        if self.show_link_frame:
            self.on_link_frame()
        if self.show_cameras:
            self.on_camera_frustum()

        self.on_tool()