rendering

WarpCPURenderer

Bases: _BaseWarpRenderer

CPU-side renderer for multi-world MJWarp simulation.

Like flygym.rendering.Renderer, frames are rendered from a private copy of the state of each selected world, with the render geometry recomputed on the CPU.

Source code in src/flygym/warp/rendering.py

class WarpCPURenderer(_BaseWarpRenderer):
    """CPU-side renderer for multi-world MJWarp simulation.

    Like `flygym.rendering.Renderer`, frames are rendered from a private copy of the
    state of each selected world, with the render geometry recomputed on the CPU.
    """

    def _render_setup_impl(self, **kwargs: Any) -> None:
        pass  # use the mj.Renderer and render data inherited from Renderer

    def _render_impl(self, mjw_data: mjw.Data) -> np.ndarray | wp.array:
        rendered_images = np.zeros((*self._buf_dim_per_frame, 3), dtype=np.uint8)
        qpos = mjw_data.qpos.numpy()
        mocap_pos = mjw_data.mocap_pos.numpy()
        mocap_quat = mjw_data.mocap_quat.numpy()

        for wid_among_rendered, world_id in enumerate(self.world_ids):
            self._render_data.qpos[:] = qpos[world_id]
            self._render_data.mocap_pos[:] = mocap_pos[world_id]
            self._render_data.mocap_quat[:] = mocap_quat[world_id]
            update_render_geometry(self.mj_model, self._render_data)

            # Render each enabled camera and store frames
            for cam_name, internal_cam_id in self._cameras_names2id.items():
                cid_among_rendered = self.enabled_cam_names.index(cam_name)

                self.mj_renderer.update_scene(
                    self._render_data, internal_cam_id, self.scene_option
                )
                frame = self.mj_renderer.render()

                if self.buffer_frames:
                    rendered_images[wid_among_rendered, cid_among_rendered] = frame

        return rendered_images

    def _fetch_frames_to_cpu_impl(
        self, world_id_among_rendered: int, cam_id_among_rendered: int
    ) -> list[np.ndarray]:
        frames = []
        for rendered_images in self._frames:
            frame = rendered_images[world_id_among_rendered, cam_id_among_rendered, ...]
            frames.append(frame)
        return frames

WarpGPUBatchRenderer

Bases: _BaseWarpRenderer

GPU-side renderer using MJWarp's GPU batch rendering functionality.

The selected worlds are copied into a separate mjw.Data in which the render geometry is recomputed (see update_render_geometry_gpu). Rendering cost and memory therefore scale with the number of rendered worlds, not simulated ones.

Source code in src/flygym/warp/rendering.py

class WarpGPUBatchRenderer(_BaseWarpRenderer):
    """GPU-side renderer using MJWarp's GPU batch rendering functionality.

    The selected worlds are copied into a separate `mjw.Data` in which the render
    geometry is recomputed (see `update_render_geometry_gpu`). Rendering cost and
    memory therefore scale with the number of rendered worlds, not simulated ones.
    """

    def _render_setup_impl(self, **kwargs: Any) -> None:
        if not self._is_scene_option_default(self.scene_option):
            raise RuntimeError(
                "Custom scene options are not supported with WarpGPUBatchRenderer "
                "because it is not implemented in MJWarp batch rendering."
            )

        n_rendered_worlds = len(self.world_ids)
        self._world_ids_gpu = wp.array(self.world_ids, dtype=wp.int32)
        self._rendered_world_ids_gpu = wp.array(
            list(range(n_rendered_worlds)), dtype=wp.int32
        )
        self._render_mjw_data = mjw.put_data(
            self.mj_model, mj.MjData(self.mj_model), nworld=n_rendered_worlds
        )
        self._enabled_cam_ids_gpu = wp.array(
            [self._cameras_names2id[n] for n in self.enabled_cam_names], dtype=wp.int32
        )
        cam_mask = [
            self._cameras_id2name[cid] in self.enabled_cam_names
            for cid in range(self.mj_model.ncam)
        ]

        # Create batch rendering context
        self._rendering_context = mjw.create_render_context(
            mjm=self.mj_model,
            nworld=n_rendered_worlds,
            cam_active=cam_mask,
            cam_res=self.camera_res[::-1],  # MJWarp expects (W, H); we use (H, W)
            **kwargs,
        )

        # Remove normal MjRenderer inherited from CPU Renderer
        self.scene_option = None
        self.mj_renderer = None

    def _render_impl(self, mjw_data: mjw.Data) -> np.ndarray | wp.array:
        render_data = self._render_mjw_data
        for src, dst in (
            (mjw_data.qpos, render_data.qpos),
            (mjw_data.mocap_pos, render_data.mocap_pos),
            (mjw_data.mocap_quat, render_data.mocap_quat),
        ):
            if dst.size > 0:
                wp.launch(
                    wp_gather_worlds_2d,
                    dim=dst.shape,
                    inputs=[src, dst, self._world_ids_gpu],
                )
        update_render_geometry_gpu(self.mjw_model, render_data)
        mjw.refit_bvh(self.mjw_model, render_data, self._rendering_context)
        mjw.render(self.mjw_model, render_data, self._rendering_context)
        rgb_out = wp.zeros(self._buf_dim_per_frame, dtype=wp.vec3f)
        get_rgb_selected_worlds_and_cameras(
            self._rendering_context,
            self._rendered_world_ids_gpu,
            self._enabled_cam_ids_gpu,
            rgb_out,
        )
        return rgb_out

    def _fetch_frames_to_cpu_impl(
        self, world_id_among_rendered: int, cam_id_among_rendered: int
    ) -> list[np.ndarray]:
        frames = []
        for frame_buffer in self._frames:
            frame = frame_buffer[world_id_among_rendered, cam_id_among_rendered, :, :]
            frame = (frame * 255.0).numpy().astype(np.uint8)
            frames.append(frame)
        return frames

    @override
    def close(self):
        return  # nothing to do since we are not using a mj.Renderer context

    @staticmethod
    def _is_scene_option_default(scene_option: mj.MjvOption) -> bool:
        default_option = mj.MjvOption()
        mj.mjv_defaultOption(default_option)
        return scene_option == default_option

modify_world_for_batch_rendering(world)

Modify world MJCF model to make it compatible with MJWarp's GPU batch rendering.

This may reduce texture realism.

Modification happens in place. Returns True if any modifications were made, False otherwise.

Note for developers: Check if anything here can be dropped upon new MJWarp releases.

Source code in src/flygym/warp/rendering.py

def modify_world_for_batch_rendering(world: BaseWorld) -> bool:
    """Modify world MJCF model to make it compatible with MJWarp's GPU batch rendering.

    This may reduce texture realism.

    Modification happens in place. Returns True if any modifications were made, False
    otherwise.

    Note for developers: Check if anything here can be dropped upon new MJWarp releases.
    """
    is_modified = False

    rgb_role = int(mj.mjtTextureRole.mjTEXROLE_RGB)

    # Strip textures from fly body materials
    # (rendering textures on complex meshes causes MJWarp memory corruption)
    for material in world.mjcf_root.materials:
        # Don't touch things that are not part of a Fly
        if material.name.split("/")[0] not in world.fly_lookup:
            continue
        # Make wings half transparent
        if "wing" in material.name:
            material.rgba[3] = 0.5
        # If material has a texture, remove it to reduce memory use
        texture_name = material.textures[rgb_role]
        if texture_name:
            texture_element = world.mjcf_root.texture(texture_name)
            primary_color_rgb = texture_element.rgb1
            material.textures[rgb_role] = ""
            material.rgba[:3] = primary_color_rgb
            is_modified = True

    return is_modified

update_render_geometry_gpu(mjw_model, mjw_data)

GPU counterpart of flygym.rendering.update_render_geometry.

Source code in src/flygym/warp/rendering.py

def update_render_geometry_gpu(mjw_model: mjw.Model, mjw_data: mjw.Data) -> None:
    """GPU counterpart of `flygym.rendering.update_render_geometry`."""
    mjw.kinematics(mjw_model, mjw_data)
    mjw.com_pos(mjw_model, mjw_data)
    mjw.camlight(mjw_model, mjw_data)
    mjw.tendon(mjw_model, mjw_data)