Skip to main content

mujoco_rs/
renderer.rs

1//! Module related to implementation of the [`MjRenderer`].
2use crate::wrappers::mj_visualization::MjvScene;
3use crate::wrappers::mj_rendering::MjrContext;
4
5#[cfg(target_os = "linux")]
6use crate::renderer::egl::GlStateEgl;
7
8use crate::vis_common::{sync_geoms, flip_image_vertically, write_png};
9use crate::builder_setters;
10use crate::prelude::*;
11
12use bitflags::bitflags;
13
14use std::io::{self, BufWriter, Write};
15use std::fmt::Display;
16use std::error::Error;
17use std::num::NonZero;
18use std::ops::Deref;
19use std::path::Path;
20use std::fs::File;
21
22/// Exported `png` crate for convenience.
23pub use png;
24
25/// Scale factor for converting normalized [0..1] depth to u16.
26const DEPTH_U16_SCALE: f32 = u16::MAX as f32;
27
28#[cfg(feature = "renderer-winit-fallback")]
29mod universal;
30
31#[cfg(feature = "renderer-winit-fallback")]
32use universal::GlStateWinit;
33
34#[cfg(target_os = "linux")]
35mod egl;
36
37
38const EXTRA_INTERNAL_VISUAL_GEOMS: u32 = 100;
39
40/// Compute the depth near/far clip planes in metres from the model.
41fn model_near_far(model: &MjModel) -> (f32, f32) {
42    let map = &model.vis().map;
43    let extent = model.stat().extent as f32;
44    (map.znear * extent, map.zfar * extent)
45}
46
47
48/// GlState enum wrapper. By default, headless implementation will be used
49/// when supported. Only on failure will an invisible winit window be used.
50#[derive(Debug)]
51#[allow(clippy::large_enum_variant)]
52pub(crate) enum GlState {
53    #[cfg(feature = "renderer-winit-fallback")] Winit(GlStateWinit),
54    #[cfg(target_os = "linux")] Egl(egl::GlStateEgl),
55}
56
57impl GlState {
58    /// Creates a new [`GlState`], which by default tries to use
59    /// an offscreen implementation. As a fallback, winit will be used.
60    pub(crate) fn new(width: NonZero<u32>, height: NonZero<u32>) -> Result<Self, RendererError> {
61        #[cfg(target_os = "linux")]
62        #[allow(unused_variables)]
63        let egl_err = match GlStateEgl::new(width, height) {
64            Ok(egl_state) => return Ok(Self::Egl(egl_state)),
65            Err(e) => e,
66        };
67
68        #[cfg(feature = "renderer-winit-fallback")]
69        match GlStateWinit::new(width, height) {
70            Ok(winit_state) => return Ok(Self::Winit(winit_state)),
71            #[cfg(not(target_os = "linux"))]
72            Err(e) => {
73                return Err(e);
74            },
75
76            #[cfg(target_os = "linux")]
77            _ => {}
78        }
79
80        #[cfg(target_os = "linux")]
81        Err(RendererError::GlutinError(egl_err))
82    }
83
84    /// Makes the internal OpenGL context current for the calling thread.
85    pub(crate) fn make_current(&self) -> glutin::error::Result<()> {
86        match self {
87            #[cfg(target_os = "linux")]
88            Self::Egl(egl_state) => egl_state.make_current(),
89            #[cfg(feature = "renderer-winit-fallback")]
90            Self::Winit(winit_state) => winit_state.make_current()
91        }
92    }
93}
94
95
96/// A builder for [`MjRenderer`].
97#[derive(Debug)]
98pub struct MjRendererBuilder {
99    width: u32,
100    height: u32,
101    num_visual_internal_geom: u32,
102    num_visual_user_geom: u32,
103    rgb: bool,
104    depth: bool,
105    png_compression: png::Compression,
106    font_scale: MjtFontScale,
107    camera: MjvCamera,
108    opts: MjvOption,
109}
110
111impl MjRendererBuilder {
112    /// Create a builder with default configuration.
113    /// Defaults are:
114    /// - `width` and `height`: use offwidth and offheight of MuJoCo's visual/global settings from the model,
115    /// - `num_visual_internal_geom`: 100,
116    /// - `num_visual_user_geom`: 0,
117    /// - `rgb`: true,
118    /// - `depth`: false,
119    /// - `png_compression`: [`png::Compression::NoCompression`] (fastest, largest files),
120    /// - `font_scale`: [`MjtFontScale::mjFONTSCALE_100`],
121    /// - `camera`: [`MjvCamera::default()`] (MuJoCo's `mjv_defaultCamera`),
122    /// - `opts`: [`MjvOption::default()`] (MuJoCo's `mjv_defaultOption`).
123    pub fn new() -> Self {
124        Self {
125            width: 0, height: 0,
126            num_visual_internal_geom: EXTRA_INTERNAL_VISUAL_GEOMS, num_visual_user_geom: 0,
127            rgb: true, depth: false, png_compression: png::Compression::NoCompression,
128            font_scale: MjtFontScale::mjFONTSCALE_100,
129            camera: MjvCamera::default(), opts: MjvOption::default(),
130        }
131    }
132
133    builder_setters! {
134        width: u32; "
135image width.
136
137<div class=\"warning\">
138
139The width must be less or equal to the offscreen buffer width,
140which can be configured at the top of the model's XML like so:
141
142```xml
143<visual>
144    <global offwidth=\"1920\" .../>
145</visual>
146```
147
148</div>";
149
150        height: u32; "\
151image height.
152
153<div class=\"warning\">
154
155The height must be less or equal to the offscreen buffer height,
156which can be configured at the top of the model's XML like so:
157
158```xml
159<visual>
160    <global offheight=\"1080\" .../>
161</visual>
162```
163
164</div>";
165
166        num_visual_internal_geom: u32; "\
167            maximum number of additional visual-only internal geoms to allocate for.
168            Note that the total number of geoms in the internal scene will be
169            `model.ngeom` + `num_visual_internal_geom` + `num_visual_user_geom`.";
170
171        num_visual_user_geom: u32;      "maximum number of additional visual-only user geoms (drawn by the user).";
172        rgb: bool;                      "RGB rendering enabled (true) or disabled (false).";
173        depth: bool;                    "depth rendering enabled (true) or disabled (false).";
174        png_compression: png::Compression; "PNG compression level used by [`MjRenderer::save_rgb`] and [`MjRenderer::save_depth`].";
175        font_scale: MjtFontScale;       "font scale of drawn text (with [MjrContext]).";
176        camera: MjvCamera;              "camera used for drawing.";
177        opts: MjvOption;                "visualization options.";
178    }
179
180    /// Builds a [`MjRenderer`].
181    /// # Returns
182    /// On success, returns [`Ok`] variant containing the [`MjRenderer`].
183    /// # Errors
184    /// - [`RendererError::ZeroDimension`] if the width or height is zero.
185    /// - [`RendererError::GlutinError`] if OpenGL initialization fails.
186    /// - [`RendererError::EventLoopError`] if the event loop fails to initialize
187    ///   (feature `renderer-winit-fallback`).
188    /// - [`RendererError::GlInitFailed`] if the fallback window initialization fails
189    ///   (feature `renderer-winit-fallback`).
190    pub fn build<M: Deref<Target = MjModel>>(self, model: M) -> Result<MjRenderer, RendererError> {
191        // Assume model's maximum should be used
192        let mut height = self.height;
193        let mut width = self.width;
194        if width == 0 && height == 0 {
195            let global = &model.vis().global;
196            height = global.offheight as u32;
197            width = global.offwidth as u32;
198        }
199
200        let gl_state = GlState::new(
201            NonZero::new(width).ok_or(RendererError::ZeroDimension)?,
202            NonZero::new(height).ok_or(RendererError::ZeroDimension)?,
203        )?;
204
205        // Initialize the rendering context to render to the offscreen buffer.
206        // SAFETY: gl_state was just created above, establishing a current GL context.
207        let mut context = unsafe { MjrContext::new(&model) };
208        context.offscreen();
209        context.change_font(self.font_scale);
210
211        let extra_geom = self.num_visual_internal_geom as usize + self.num_visual_user_geom as usize;
212        let (near, far) = model_near_far(&model);
213
214        // The 3D scene for visualization
215        let scene = MjvScene::new(
216            &*model,
217            model.ffi().ngeom as usize + extra_geom
218        );
219
220        let user_scene = MjvScene::new(
221            &*model,
222            self.num_visual_user_geom as usize
223        );
224
225        // Construct the renderer and create allocated buffers.
226        let renderer = MjRenderer {
227            scene, user_scene, context, camera: self.camera, option: self.opts,
228            flags: RendererFlags::empty(), rgb: None, depth: None,
229            width: width as usize, height: height as usize, gl_state,
230            png_compression: self.png_compression, font_scale: self.font_scale,
231            near, far, extra_geom,
232        }   // These require special care
233            .with_rgb_rendering(self.rgb)
234            .with_depth_rendering(self.depth);
235
236        Ok(renderer)
237    }
238}
239
240
241/// Delegates to [`MjRendererBuilder::new`].
242impl Default for MjRendererBuilder {
243    fn default() -> Self {
244        Self::new()
245    }
246}
247
248/// A renderer for rendering 3D scenes.
249/// By default, RGB rendering is enabled and depth rendering is disabled.
250#[derive(Debug)]
251pub struct MjRenderer {
252    scene: MjvScene,
253    user_scene: MjvScene,
254    context: MjrContext,
255
256    /* OpenGL */
257    gl_state: GlState,
258
259    /* Configuration */
260    camera: MjvCamera,
261    option: MjvOption,
262    flags: RendererFlags,
263    png_compression: png::Compression,
264    font_scale: MjtFontScale,
265
266    /* Cached from the current model */
267    // Depth near/far clip planes in metres, derived from model.vis().map and model.stat().
268    near: f32,
269    far: f32,
270    // Scene capacity headroom beyond ngeom: preserved across model switches.
271    extra_geom: usize,
272
273    /* Storage */
274    // Use Box to allow less space to be used
275    // when rgb or depth rendering is disabled
276    rgb: Option<Box<[u8]>>,
277    depth: Option<Box<[f32]>>,
278
279    width: usize,
280    height: usize,
281}
282
283impl MjRenderer {
284    /// Construct a new renderer.
285    /// The `max_user_geom` parameter
286    /// defines how much space will be allocated for additional, user-defined visual-only geoms.
287    /// It can thus be set to 0 if no additional geoms will be drawn by the user.
288    /// # Scene allocation
289    /// The renderer uses two scenes:
290    /// - the internal scene: used by the renderer to draw the model's state.
291    /// - the user scene: used by the user to add additional geoms to the internal scene
292    ///
293    /// The **internal scene** allocates the amount of space needed to fit every pre-existing
294    /// model geom + user visual-only geoms + additional visual-only geoms that aren't from the user (e.g., tendons).
295    /// By default, the renderer reserves 100 extra geom slots for drawing the additional visual-only geoms.
296    /// If that is not enough or it is too much, you can construct [`MjRenderer`] via its builder
297    /// ([`MjRenderer::builder`]), which allows more configuration.
298    ///
299    /// <div class="warning">
300    ///
301    /// Parameters `width` and `height` must be less or equal to the offscreen buffer size,
302    /// which can be configured at the top of the model's XML like so:
303    ///
304    /// ```xml
305    /// <visual>
306    ///    <global offwidth="1920" offheight="1080"/>
307    /// </visual>
308    /// ```
309    ///
310    /// </div>
311    /// # Returns
312    /// On success, returns [`Ok`] variant containing the [`MjRenderer`].
313    /// # Errors
314    /// - [`RendererError::ZeroDimension`] if the width or height is zero.
315    /// - [`RendererError::GlutinError`] if OpenGL initialization fails.
316    /// - [`RendererError::EventLoopError`] if the event loop fails to initialize
317    ///   (feature `renderer-winit-fallback`).
318    /// - [`RendererError::GlInitFailed`] if the fallback window initialization fails
319    ///   (feature `renderer-winit-fallback`).
320    pub fn new<M: Deref<Target = MjModel>>(model: M, width: usize, height: usize, max_user_geom: usize) -> Result<Self, RendererError> {
321        MjRendererBuilder::new()
322            .width(width as u32).height(height as u32).num_visual_user_geom(max_user_geom as u32)
323            .build(model)
324    }
325
326    /// Create a [`MjRendererBuilder`] to configure [`MjRenderer`].
327    pub fn builder() -> MjRendererBuilder {
328        MjRendererBuilder::new()
329    }
330
331    /// Return an immutable reference to the internal scene.
332    pub fn scene(&self) -> &MjvScene {
333        &self.scene
334    }
335
336    /// Return an immutable reference to a user scene for drawing custom visual-only geoms.
337    pub fn user_scene(&self) -> &MjvScene {
338        &self.user_scene
339    }
340
341    /// Return a mutable reference to a user scene for drawing custom visual-only geoms.
342    pub fn user_scene_mut(&mut self) -> &mut MjvScene {
343        &mut self.user_scene
344    }
345
346    /// Return an immutable reference to visualization options.
347    pub fn opts(&self) -> &MjvOption {
348        &self.option
349    }
350
351    /// Return a mutable reference to visualization options.
352    pub fn opts_mut(&mut self) -> &mut MjvOption {
353        &mut self.option
354    }
355
356    /// Return an immutable reference to the camera.
357    pub fn camera(&self) -> &MjvCamera {
358        &self.camera
359    }
360
361    /// Return a mutable reference to the camera.
362    pub fn camera_mut(&mut self) -> &mut MjvCamera {
363        &mut self.camera
364    }
365
366    /// Check if RGB rendering is enabled.
367    pub fn rgb_enabled(&self) -> bool {
368        self.flags.contains(RendererFlags::RENDER_RGB)
369    }
370
371    /// Check if depth rendering is enabled.
372    pub fn depth_enabled(&self) -> bool {
373        self.flags.contains(RendererFlags::RENDER_DEPTH)
374    }
375
376    /// Sets the font scale.
377    ///
378    /// # Errors
379    /// Returns [`RendererError::GlutinError`] if the OpenGL context cannot be made current.
380    pub fn set_font_scale(&mut self, font_scale: MjtFontScale) -> Result<(), RendererError> {
381        self.gl_state.make_current().map_err(RendererError::GlutinError)?;
382        self.font_scale = font_scale;
383        self.context.change_font(font_scale);
384        Ok(())
385    }
386
387    /// Update the visualization options.
388    pub fn set_opts(&mut self, options: MjvOption) {
389        self.option = options;
390    }
391
392    fn update_x_from(
393        &self,
394        model: &MjModel,
395        id: usize,
396        upload: fn(&MjrContext, &MjModel, usize) -> Result<(), crate::error::MjrContextError>,
397    ) -> Result<(), RendererError>
398    {
399        self.prepare_upload(model)?;
400        upload(&self.context, model, id)?;
401        Ok(())
402    }
403
404    fn prepare_upload(&self, model: &MjModel) -> Result<(), RendererError> {
405        if model.signature() != self.scene.signature() {
406            return Err(RendererError::SignatureMismatch);
407        }
408        self.gl_state.make_current().map_err(RendererError::GlutinError)
409    }
410
411    fn update_all_from_impl(
412        &self,
413        model: &MjModel,
414        n: usize,
415        upload: fn(&MjrContext, &MjModel, usize) -> Result<(), crate::error::MjrContextError>,
416    ) -> Result<(), RendererError>
417    {
418        self.prepare_upload(model)?;
419        for id in 0..n {
420            upload(&self.context, model, id)?;
421        }
422        Ok(())
423    }
424
425    /// Re-uploads the texture with `texture_id` from `model` to the GPU immediately.
426    ///
427    /// # Errors
428    /// - [`RendererError::SignatureMismatch`] if `model`'s signature does not match the renderer's scene.
429    /// - [`RendererError::ContextError`] if `texture_id >= model.ntex()`.
430    /// - [`RendererError::GlutinError`] if the OpenGL context cannot be made current.
431    pub fn update_texture_from(&self, model: &MjModel, texture_id: usize) -> Result<(), RendererError> {
432        self.update_x_from(model, texture_id, MjrContext::upload_texture)
433    }
434
435    /// Re-uploads all textures from `model` to the GPU immediately.
436    ///
437    /// # Errors
438    /// - [`RendererError::SignatureMismatch`] if `model`'s signature does not match the renderer's scene.
439    /// - [`RendererError::GlutinError`] if the OpenGL context cannot be made current.
440    pub fn update_textures_from(&self, model: &MjModel) -> Result<(), RendererError> {
441        self.update_all_from_impl(model, model.ntex() as usize, MjrContext::upload_texture)
442    }
443
444    /// Re-uploads the mesh with `mesh_id` from `model` to the GPU immediately.
445    ///
446    /// All data arrays read by `mjr_uploadMesh` are copied: vertex positions
447    /// (`mesh_vert`), per-vertex normals (`mesh_normal`), UV texture coordinates
448    /// (`mesh_texcoord`), face--vertex indices (`mesh_face`), face--normal indices
449    /// (`mesh_facenormal`), face--texcoord indices (`mesh_facetexcoord`), and convex
450    /// hull graph data (`mesh_graph`). Layout fields (address and count arrays) are
451    /// not copied because they are fixed by the model signature.
452    ///
453    /// # Errors
454    /// - [`RendererError::SignatureMismatch`] if `model`'s signature does not match the renderer's scene.
455    /// - [`RendererError::ContextError`] if `mesh_id >= model.nmesh()`.
456    /// - [`RendererError::GlutinError`] if the OpenGL context cannot be made current.
457    pub fn update_mesh_from(&self, model: &MjModel, mesh_id: usize) -> Result<(), RendererError> {
458        self.update_x_from(model, mesh_id, MjrContext::upload_mesh)
459    }
460
461    /// Re-uploads all meshes from `model` to the GPU immediately.
462    ///
463    /// All data arrays read by `mjr_uploadMesh` are bulk-uploaded: vertex positions
464    /// (`mesh_vert`), per-vertex normals (`mesh_normal`), UV texture coordinates
465    /// (`mesh_texcoord`), face--vertex indices (`mesh_face`), face--normal indices
466    /// (`mesh_facenormal`), face--texcoord indices (`mesh_facetexcoord`), and convex
467    /// hull graph data (`mesh_graph`). Layout fields (address and count arrays) are
468    /// not copied because they are fixed by the model signature.
469    ///
470    /// # Errors
471    /// - [`RendererError::SignatureMismatch`] if `model`'s signature does not match the renderer's scene.
472    /// - [`RendererError::GlutinError`] if the OpenGL context cannot be made current.
473    pub fn update_meshes_from(&self, model: &MjModel) -> Result<(), RendererError> {
474        self.update_all_from_impl(model, model.nmesh() as usize, MjrContext::upload_mesh)
475    }
476
477    /// Re-uploads the heightfield with `hfield_id` from `model` to the GPU immediately.
478    ///
479    /// # Errors
480    /// - [`RendererError::SignatureMismatch`] if `model`'s signature does not match the renderer's scene.
481    /// - [`RendererError::ContextError`] if `hfield_id >= model.nhfield()`.
482    /// - [`RendererError::GlutinError`] if the OpenGL context cannot be made current.
483    pub fn update_hfield_from(&self, model: &MjModel, hfield_id: usize) -> Result<(), RendererError> {
484        self.update_x_from(model, hfield_id, MjrContext::upload_hfield)
485    }
486
487    /// Re-uploads all heightfields from `model` to the GPU immediately.
488    ///
489    /// # Errors
490    /// - [`RendererError::SignatureMismatch`] if `model`'s signature does not match the renderer's scene.
491    /// - [`RendererError::GlutinError`] if the OpenGL context cannot be made current.
492    pub fn update_hfields_from(&self, model: &MjModel) -> Result<(), RendererError> {
493        self.update_all_from_impl(model, model.nhfield() as usize, MjrContext::upload_hfield)
494    }
495
496    /// Set the camera used for rendering.
497    pub fn set_camera(&mut self, camera: MjvCamera)  {
498        self.camera = camera;
499    }
500
501    /// Enables/disables RGB rendering.
502    pub fn set_rgb_rendering(&mut self, enable: bool) {
503        self.flags.set(RendererFlags::RENDER_RGB, enable);
504        self.rgb = if enable { Some(vec![0; 3 * self.width * self.height].into_boxed_slice()) } else { None } ;
505    }
506
507    /// Enables/disables depth rendering.
508    pub fn set_depth_rendering(&mut self, enable: bool) {
509        self.flags.set(RendererFlags::RENDER_DEPTH, enable);
510        self.depth = if enable { Some(vec![0.0; self.width * self.height].into_boxed_slice()) } else { None } ;
511    }
512
513    /// Sets the font scale. To be used on construction.
514    ///
515    /// # Errors
516    /// Returns [`RendererError::GlutinError`] if the OpenGL context cannot be made current.
517    pub fn with_font_scale(mut self, font_scale: MjtFontScale) -> Result<Self, RendererError> {
518        self.set_font_scale(font_scale)?;
519        Ok(self)
520    }
521
522    /// Update the visualization options and return `self`. To be used on construction.
523    pub fn with_opts(mut self, options: MjvOption) -> Self {
524        self.set_opts(options);
525        self
526    }
527
528    /// Set the camera used for rendering. To be used on construction.
529    pub fn with_camera(mut self, camera: MjvCamera) -> Self  {
530        self.set_camera(camera);
531        self
532    }
533
534    /// Enables/disables RGB rendering. To be used on construction.
535    pub fn with_rgb_rendering(mut self, enable: bool) -> Self {
536        self.set_rgb_rendering(enable);
537        self
538    }
539
540    /// Enables/disables depth rendering. To be used on construction.
541    pub fn with_depth_rendering(mut self, enable: bool) -> Self {
542        self.set_depth_rendering(enable);
543        self
544    }
545
546    /// Set the PNG compression level used by [`MjRenderer::save_rgb`] and [`MjRenderer::save_depth`].
547    pub fn set_png_compression(&mut self, compression: png::Compression) {
548        self.png_compression = compression;
549    }
550
551    /// Set the PNG compression level. To be used on construction.
552    pub fn with_png_compression(mut self, compression: png::Compression) -> Self {
553        self.set_png_compression(compression);
554        self
555    }
556
557    /// Update the scene with new data from data.
558    /// When `data` was created from a different model than the renderer,
559    /// the internal scene is automatically recreated for the new model.
560    ///
561    /// # Errors
562    /// - [`RendererError::GlutinError`] if the OpenGL context could not be made current
563    ///   (only when the [`MjModel`] in `data` differs from the model that created the internal [`MjvScene`]).
564    pub fn sync_data<M: Deref<Target = MjModel>>(&mut self, data: &mut MjData<M>) -> Result<(), RendererError> {
565        if data.model().signature() != self.scene.signature() {
566            /* Model changed: preserve the extra-geom headroom and user-geom
567             * capacity, only substitute the per-model ngeom base count. */
568            // Ensure the GL context is current before dropping old GPU resources
569            // (MjrContext::drop calls mjr_freeContext) and creating new ones.
570            self.gl_state.make_current().map_err(RendererError::GlutinError)?;
571
572            let user_geom_cap = self.user_scene.maxgeom() as usize;
573            let new_ngeom = data.model().ffi().ngeom as usize;
574            self.scene = MjvScene::new(data.model(), new_ngeom + self.extra_geom);
575            self.user_scene = MjvScene::new(data.model(), user_geom_cap);
576            (self.near, self.far) = model_near_far(data.model());
577
578            // Recreate the rendering context so that GPU resources (textures,
579            // meshes, heightfields, skins) match the new model.
580            // SAFETY: the GL context was made current above.
581            self.context = unsafe { MjrContext::new(data.model()) };
582            self.context.offscreen();
583            self.context.change_font(self.font_scale);
584        }
585
586        self.scene.update(data, &self.option, &MjvPerturb::default(), &mut self.camera);
587        Ok(())
588    }
589
590    /// Update the scene with new data and render it.
591    ///
592    /// # Panics
593    /// Panics if syncing the data or performing the render step fails (e.g. the OpenGL
594    /// context cannot be made current). Use [`MjRenderer::sync_data`] + [`MjRenderer::render`]
595    /// instead.
596    #[deprecated(note = "replaced with sync_data + render", since = "3.0.0")]
597    pub fn sync<M: Deref<Target = MjModel>>(&mut self, data: &mut MjData<M>) {
598        self.sync_data(data).unwrap();
599        self.render().unwrap();
600    }
601
602    /// Return a flattened RGB image of the scene.
603    pub fn rgb_flat(&self) -> Option<&[u8]> {
604        self.rgb.as_deref()
605    }
606
607    /// Return an RGB image of the scene. This method accepts two generic parameters <WIDTH, HEIGHT>
608    /// that define the shape of the output slice.
609    ///
610    /// # Panics
611    /// - If the image size doesn't match the required dimensions.
612    /// - If RGB rendering is disabled.
613    ///
614    /// Use [`MjRenderer::try_rgb`] for a fallible alternative.
615    pub fn rgb<const WIDTH: usize, const HEIGHT: usize>(&self) -> &[[[u8; 3]; WIDTH]; HEIGHT] {
616        self.try_rgb::<WIDTH, HEIGHT>().unwrap()
617    }
618
619    /// Fallible version of [`MjRenderer::rgb`].
620    /// # Returns
621    /// On success, returns [`Ok`] variant containing the rendered RGB image.
622    /// # Errors
623    /// - [`RendererError::DimensionMismatch`] if the image size doesn't match the required dimensions.
624    /// - [`RendererError::RgbDisabled`] if RGB rendering is disabled.
625    pub fn try_rgb<const WIDTH: usize, const HEIGHT: usize>(&self) -> Result<&[[[u8; 3]; WIDTH]; HEIGHT], RendererError> {
626        if let Some(flat) = self.rgb_flat() {
627            bytemuck::try_from_bytes(flat)
628                .map_err(|_| RendererError::DimensionMismatch)
629        }
630        else {
631            Err(RendererError::RgbDisabled)
632        }
633    }
634
635    /// Return a flattened depth image of the scene.
636    pub fn depth_flat(&self) -> Option<&[f32]> {
637        self.depth.as_deref()
638    }
639
640    /// Return a depth image of the scene. This method accepts two generic parameters <WIDTH, HEIGHT>
641    /// that define the shape of the output slice.
642    ///
643    /// # Panics
644    /// - If the image size doesn't match the required dimensions.
645    /// - If depth rendering is disabled.
646    ///
647    /// Use [`MjRenderer::try_depth`] for a fallible alternative.
648    pub fn depth<const WIDTH: usize, const HEIGHT: usize>(&self) -> &[[f32; WIDTH]; HEIGHT] {
649        self.try_depth::<WIDTH, HEIGHT>().unwrap()
650    }
651
652    /// Fallible version of [`MjRenderer::depth`].
653    /// # Returns
654    /// On success, returns [`Ok`] variant containing the rendered depth image.
655    /// # Errors
656    /// - [`RendererError::DimensionMismatch`] if the image size doesn't match the required dimensions.
657    /// - [`RendererError::DepthDisabled`] if depth rendering is disabled.
658    pub fn try_depth<const WIDTH: usize, const HEIGHT: usize>(&self) -> Result<&[[f32; WIDTH]; HEIGHT], RendererError> {
659        if let Some(flat) = self.depth_flat() {
660            let bytes: &[u8] = bytemuck::cast_slice(flat);
661            bytemuck::try_from_bytes(bytes)
662                .map_err(|_| RendererError::DimensionMismatch)
663        }
664        else {
665            Err(RendererError::DepthDisabled)
666        }
667    }
668
669    /// Save an RGB image of the scene to a path.
670    /// # Returns
671    /// `Ok(())` on success.
672    /// # Errors
673    /// - [`RendererError::RgbDisabled`] when RGB rendering is disabled.
674    /// - [`RendererError::IoError`] if a file I/O operation fails.
675    pub fn save_rgb<T: AsRef<Path>>(&self, path: T) -> Result<(), RendererError> {
676        if let Some(rgb) = &self.rgb {
677            write_png(
678                path,
679                rgb,
680                self.width as u32,
681                self.height as u32,
682                png::ColorType::Rgb,
683                png::BitDepth::Eight,
684                self.png_compression
685            )?;
686            Ok(())
687        }
688        else {
689            Err(RendererError::RgbDisabled)
690        }
691    }
692
693    /// Save a depth image of the scene to a path. The image is 16-bit PNG, which
694    /// can be converted into depth (distance) data by dividing the grayscale values by
695    /// 65535.0 and applying inverse normalization: `depth = min + (grayscale / 65535.0) * (max - min)`.
696    ///
697    /// If `normalize` is `true`, then the data is normalized with per-frame min-max normalization.
698    /// When `normalize` is `false`, the depth values are mapped using the model's camera
699    /// near/far clip planes as the range, providing a fixed (frame-independent) mapping.
700    ///
701    /// Use of [`MjRenderer::save_depth_raw`] is recommended if performance is critical, as
702    /// it skips PNG encoding and also saves the true depth values directly.
703    /// # Returns
704    /// An [`Ok`]`((min, max))` is returned, where min and max represent the normalization parameters.
705    /// # Errors
706    /// - [`RendererError::DepthDisabled`] when depth rendering is disabled.
707    /// - [`RendererError::IoError`] if a file I/O operation fails.
708    pub fn save_depth<T: AsRef<Path>>(&self, path: T, normalize: bool) -> Result<(f32, f32), RendererError> {
709        if let Some(depth) = &self.depth {
710            let (norm, min, max) =
711            if normalize {
712                let max = depth.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
713                let min = depth.iter().cloned().fold(f32::INFINITY, f32::min);
714                let range = max - min;
715                if range == 0.0 {
716                    (vec![0u8; depth.len() * 2].into_boxed_slice(), min, max)
717                } else {
718                    (depth.iter().flat_map(|&x| (((x - min) / range * DEPTH_U16_SCALE).clamp(0.0, DEPTH_U16_SCALE) as u16).to_be_bytes()).collect::<Box<_>>(), min, max)
719                }
720            }
721            else {
722                // Use model's camera near/far clip planes as the fixed normalization range.
723                // After linearization (in render()), depth values are in meters within [near, far].
724                let near = self.near;
725                let far = self.far;
726                (depth.iter().flat_map(|&x| (((x - near) / (far - near) * DEPTH_U16_SCALE).clamp(0.0, DEPTH_U16_SCALE) as u16).to_be_bytes()).collect::<Box<_>>(), near, far)
727            };
728
729            write_png(
730                path,
731                &norm,
732                self.width as u32,
733                self.height as u32,
734                png::ColorType::Grayscale,
735                png::BitDepth::Sixteen,
736                self.png_compression
737            )?;
738            Ok((min, max))
739        }
740        else {
741            Err(RendererError::DepthDisabled)
742        }
743    }
744
745    /// Save the raw depth data to the `path`. The data is encoded
746    /// as a sequence of bytes, where groups of four represent a single f32 value
747    /// in native byte order.
748    /// # Returns
749    /// `Ok(())` on success.
750    /// # Errors
751    /// - [`RendererError::DepthDisabled`] when depth rendering is disabled.
752    /// - [`RendererError::IoError`] if a file I/O operation fails.
753    pub fn save_depth_raw<T: AsRef<Path>>(&self, path: T) -> Result<(), RendererError> {
754        if let Some(depth) = &self.depth {
755            let file = File::create(path.as_ref())?;
756            let mut writer = BufWriter::new(file);
757
758            let bytes: &[u8] = bytemuck::cast_slice(depth);
759            writer.write_all(bytes)?;
760            Ok(())
761        }
762        else {
763            Err(RendererError::DepthDisabled)
764        }
765    }
766
767    /// Draws the scene to internal arrays.
768    /// Use [`MjRenderer::rgb`] or [`MjRenderer::depth`] to obtain the rendered image.
769    ///
770    /// # Errors
771    /// - [`RendererError::GlutinError`] if the OpenGL context could not be made current.
772    /// - [`RendererError::SceneError`] if the user-scene sync overflows the geom buffer.
773    /// - [`RendererError::ContextError`] if reading pixels from the framebuffer fails.
774    pub fn render(&mut self) -> Result<(), RendererError> {
775        /* Sync user scene geoms into the main scene before rendering */
776        sync_geoms(&self.user_scene, &mut self.scene).map_err(RendererError::SceneError)?;
777
778        self.gl_state.make_current().map_err(RendererError::GlutinError)?;
779        let vp = MjrRectangle::new(0, 0, self.width as i32, self.height as i32);
780        self.scene.render(&vp, &self.context);
781
782        /* Fully flatten everything */
783        let flat_rgb = self.rgb.as_deref_mut();
784        let flat_depth = self.depth.as_deref_mut();
785
786        /* Read to whatever is enabled */
787        self.context.read_pixels(
788            flat_rgb,
789            flat_depth,
790            &vp
791        ).map_err(RendererError::ContextError)?;
792
793        /* Flip the read pixels vertically, as OpenGL reads bottom-up */
794        if let Some(rgb) = self.rgb.as_deref_mut() {
795            flip_image_vertically(rgb, self.height, self.width * 3);
796        }
797
798        /* Make depth values be the actual distance in meters and flip them vertically */
799        if let Some(depth) = self.depth.as_deref_mut() {
800            flip_image_vertically(depth, self.height, self.width);
801
802            let near = self.near;
803            let far = self.far;
804            for value in depth {
805                *value = near / (1.0 - *value * (1.0 - near / far));
806            }
807        }
808
809        Ok(())
810    }
811}
812
813/// Errors that can occur during renderer operations.
814#[derive(Debug)]
815#[non_exhaustive]
816pub enum RendererError {
817    /// The event loop failed to initialize.
818    #[cfg(feature = "renderer-winit-fallback")]
819    EventLoopError(winit::error::EventLoopError),
820    /// A glutin operation failed.
821    GlutinError(glutin::error::Error),
822    /// The supplied width or height was zero; MuJoCo requires positive dimensions.
823    ZeroDimension,
824    /// OpenGL / window initialization failed.
825    #[cfg(feature = "renderer-winit-fallback")]
826    GlInitFailed(crate::error::GlInitError),
827    /// RGB rendering was not enabled.
828    RgbDisabled,
829    /// Depth rendering was not enabled.
830    DepthDisabled,
831    /// The requested `WIDTH`/`HEIGHT` do not match the renderer's dimensions.
832    DimensionMismatch,
833    /// An I/O error occurred (e.g. while saving to a file).
834    IoError(io::Error),
835    /// A scene operation failed (e.g. user-scene sync overflowed the geom buffer).
836    SceneError(crate::error::MjSceneError),
837    /// A rendering-context operation failed (e.g. an asset or aux-buffer ID is out of range).
838    ContextError(crate::error::MjrContextError),
839    /// The model's structure signature does not match the renderer's scene.
840    /// Call [`MjRenderer::sync_data`] first.
841    SignatureMismatch,
842}
843
844/// Formats a human-readable description of the renderer error.
845impl Display for RendererError {
846    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
847        match self {
848            #[cfg(feature = "renderer-winit-fallback")]
849            Self::EventLoopError(e) => write!(f, "event loop failed to initialize: {e}"),
850            Self::GlutinError(e) => write!(f, "glutin error: {e}"),
851            Self::ZeroDimension => write!(f, "renderer width and height must both be greater than zero"),
852            #[cfg(feature = "renderer-winit-fallback")]
853            Self::GlInitFailed(e) => write!(f, "GL initialization failed: {e}"),
854            Self::RgbDisabled => write!(f, "RGB rendering is not enabled (renderer.with_rgb_rendering(true))"),
855            Self::DepthDisabled => write!(f, "depth rendering is not enabled (renderer.with_depth_rendering(true))"),
856            Self::DimensionMismatch => write!(f, "the input width and height don't match the renderer's configuration"),
857            Self::IoError(e) => write!(f, "I/O error: {e}"),
858            Self::SceneError(e) => write!(f, "scene error: {e}"),
859            Self::ContextError(e) => write!(f, "rendering context error: {e}"),
860            Self::SignatureMismatch => write!(f, "model signature mismatch: call sync_data first"),
861        }
862    }
863}
864
865/// Provides the underlying error source, if any.
866impl Error for RendererError {
867    fn source(&self) -> Option<&(dyn Error + 'static)> {
868        match self {
869            #[cfg(feature = "renderer-winit-fallback")]
870            Self::EventLoopError(e) => Some(e),
871            #[cfg(feature = "renderer-winit-fallback")]
872            Self::GlInitFailed(e) => Some(e),
873            Self::GlutinError(e) => Some(e),
874            Self::IoError(e) => Some(e),
875            Self::SceneError(e) => Some(e),
876            Self::ContextError(e) => Some(e),
877            _ => None,
878        }
879    }
880}
881
882/// Converts an [`io::Error`] into [`RendererError::IoError`].
883impl From<io::Error> for RendererError {
884    fn from(e: io::Error) -> Self {
885        Self::IoError(e)
886    }
887}
888
889/// Converts a [`png::EncodingError`] into [`RendererError::IoError`].
890impl From<png::EncodingError> for RendererError {
891    fn from(e: png::EncodingError) -> Self {
892        Self::IoError(io::Error::from(e))
893    }
894}
895
896/// Converts an [`MjSceneError`] into [`RendererError::SceneError`].
897impl From<crate::error::MjSceneError> for RendererError {
898    fn from(e: crate::error::MjSceneError) -> Self {
899        Self::SceneError(e)
900    }
901}
902
903/// Converts an [`crate::error::MjrContextError`] into [`RendererError::ContextError`].
904impl From<crate::error::MjrContextError> for RendererError {
905    fn from(e: crate::error::MjrContextError) -> Self {
906        Self::ContextError(e)
907    }
908}
909
910/// Converts a [`GlInitError`] into [`RendererError::GlInitFailed`].
911#[cfg(feature = "renderer-winit-fallback")]
912impl From<crate::error::GlInitError> for RendererError {
913    fn from(e: crate::error::GlInitError) -> Self {
914        Self::GlInitFailed(e)
915    }
916}
917
918bitflags! {
919    /// Flags that enable features of the renderer.
920    #[derive(Debug)]
921    struct RendererFlags: u8 {
922        const RENDER_RGB = 1 << 0;
923        const RENDER_DEPTH = 1 << 1;
924    }
925}
926
927
928
929/// Ensures the OpenGL context is current before GPU resources are freed.
930impl Drop for MjRenderer {
931    fn drop(&mut self) {
932        // Ensure the GL context is current before the implicit field drops
933        // (MjrContext's Drop calls mjr_freeContext which requires an active GL context).
934        let _ = self.gl_state.make_current();
935    }
936}
937
938/*
939** Don't run any tests as OpenGL hates if anything
940** runs outside the main thread.
941*/
942
943#[cfg(test)]
944mod test {
945    use crate::assert_relative_eq;
946
947    use super::*;
948
949    const MODEL: &str = stringify!(
950        <mujoco>
951
952        <visual>
953            <global offwidth="1280" offheight="720"/>
954        </visual>
955
956        <worldbody>
957            <geom name="floor" type="plane" size="10 10 1" euler="0 0 0"/>
958            <geom type="box" size="1 10 10" pos="-1 0 0" euler="0 0 0"/>
959
960            <camera name="depth_test" euler="90 90 0" pos="2.25 0 1"/>
961
962        </worldbody>
963        </mujoco>
964    );
965
966    /// Depth calculation test.
967    /// This is only run on Linux due to EGL requirements (winit cannot be used on multiple threads).
968    #[test]
969    #[cfg(target_os = "linux")]
970    fn test_depth() {
971        let model = MjModel::from_xml_string(MODEL).expect("could not load the model");
972        let mut data = MjData::new(&model);
973        data.step();
974
975        let mut renderer = MjRenderer::builder()
976            .rgb(false)
977            .depth(true)
978            .camera(MjvCamera::new_fixed(model.name_to_id(MjtObj::mjOBJ_CAMERA, "depth_test").unwrap()))
979            .build(&model)
980            .unwrap();
981
982        renderer.sync_data(&mut data).unwrap();
983        renderer.render().unwrap();
984        let min = renderer.depth_flat().unwrap().iter().fold(f32::INFINITY, |a , &b| a.min(b));
985        let max = renderer.depth_flat().unwrap().iter().fold(f32::NEG_INFINITY, |a , &b| a.max(b));
986
987        assert_relative_eq!(min, max, epsilon = 1e-4);
988        assert_relative_eq!(min, 2.25, epsilon = 1e-4);
989    }
990
991    /// Decode a PNG file and return the raw pixel bytes.
992    #[cfg(target_os = "linux")]
993    fn decode_png_pixels(path: &std::path::Path) -> Vec<u8> {
994        let decoder = png::Decoder::new(std::io::BufReader::new(std::fs::File::open(path).unwrap()));
995        let mut reader = decoder.read_info().unwrap();
996        let mut buf = vec![0u8; reader.output_buffer_size().unwrap()];
997        reader.next_frame(&mut buf).unwrap();
998        buf
999    }
1000
1001    /// Verify that all compression levels produce identical pixel data.
1002    /// This is only run on Linux due to EGL requirements (winit cannot be used on multiple threads).
1003    #[test]
1004    #[cfg(target_os = "linux")]
1005    fn test_png_compression_lossless() {
1006        let model = MjModel::from_xml_string(MODEL).expect("could not load the model");
1007        let mut data = MjData::new(&model);
1008        data.step();
1009
1010        let mut renderer = MjRenderer::builder()
1011            .rgb(true)
1012            .depth(false)
1013            .build(&model)
1014            .unwrap();
1015
1016        renderer.sync_data(&mut data).unwrap();
1017        renderer.render().unwrap();
1018
1019        let tmp = std::env::temp_dir();
1020        let path_none = tmp.join("mujoco_rs_test_none.png");
1021        let path_fast = tmp.join("mujoco_rs_test_fast.png");
1022        let path_high = tmp.join("mujoco_rs_test_high.png");
1023
1024        renderer.save_rgb(&path_none).unwrap();
1025        renderer.set_png_compression(png::Compression::Fast);
1026        renderer.save_rgb(&path_fast).unwrap();
1027        renderer.set_png_compression(png::Compression::High);
1028        renderer.save_rgb(&path_high).unwrap();
1029
1030        let pixels_none = decode_png_pixels(&path_none);
1031        let pixels_fast = decode_png_pixels(&path_fast);
1032        let pixels_high = decode_png_pixels(&path_high);
1033
1034        assert_eq!(pixels_none, pixels_fast, "Fast compression must produce identical pixels");
1035        assert_eq!(pixels_none, pixels_high, "High compression must produce identical pixels");
1036
1037        let _ = std::fs::remove_file(&path_none);
1038        let _ = std::fs::remove_file(&path_fast);
1039        let _ = std::fs::remove_file(&path_high);
1040    }
1041}