mujoco_rs/
renderer.rs

1//! Module related to implementation of the [`MjRenderer`].
2use crate::wrappers::mj_visualization::MjvScene;
3use crate::wrappers::mj_rendering::MjrContext;
4use crate::builder_setters;
5use crate::prelude::*;
6
7use glfw_mjrc_fork as glfw;
8use glfw_mjrc_fork::{Context, InitError, PWindow, WindowHint};
9
10use bitflags::bitflags;
11use png::Encoder;
12
13use std::io::{self, BufWriter, ErrorKind, Write};
14use std::fmt::Display;
15use std::error::Error;
16use std::path::Path;
17use std::fs::File;
18
19const RGB_NOT_FOUND_ERR_STR: &str = "RGB rendering is not enabled (renderer.with_rgb_rendering(true))";
20const DEPTH_NOT_FOUND_ERR_STR: &str = "depth rendering is not enabled (renderer.with_depth_rendering(true))";
21const INVALID_INPUT_SIZE: &str = "the input width and height don't match the renderer's configuration";
22const EXTRA_INTERNAL_VISUAL_GEOMS: usize = 100;
23
24
25/// A builder for [`MjRenderer`].
26#[derive(Debug)]
27pub struct MjRendererBuilder {
28    width: u32,
29    height: u32,
30    num_visual_internal_geom: u32,
31    num_visual_user_geom: u32,
32    rgb: bool,
33    depth: bool,
34    font_scale: MjtFontScale,
35    camera: MjvCamera,
36    opts: MjvOption,
37}
38
39impl MjRendererBuilder {
40    /// Create a builder with default configuration.
41    /// Defaults are:
42    /// - `width` and `height`: use offwidth and offheight of MuJoCo's visual/global settings from the model,
43    /// - `num_visual_internal_geom`: 100,
44    /// - `num_visual_user_geom`: 0,
45    /// - `rgb`: true,
46    /// - `depth`: false.
47    pub fn new() -> Self {
48        Self {
49            width: 0, height: 0,
50            num_visual_internal_geom: EXTRA_INTERNAL_VISUAL_GEOMS as u32, num_visual_user_geom: 0,
51            rgb: true, depth: false, font_scale: MjtFontScale::mjFONTSCALE_100,
52            camera: MjvCamera::default(), opts: MjvOption::default()
53        }
54    }
55
56    builder_setters! {
57        width: u32; "
58image width.
59
60<div class=\"warning\">
61
62The width must be less or equal to the offscreen buffer width,
63which can be configured at the top of the model's XML like so:
64
65```xml
66<visual>
67    <global offwidth=\"1920\" .../>
68</visual>
69```
70
71</div>";
72
73        height: u32; "\
74image height.
75
76<div class=\"warning\">
77
78The height must be less or equal to the offscreen buffer height,
79which can be configured at the top of the model's XML like so:
80
81```xml
82<visual>
83    <global offheight=\"1080\" .../>
84</visual>
85```
86
87</div>";
88
89        num_visual_internal_geom: u32; "\
90            maximum number of additional visual-only internal geoms to allocate for.
91            Note that the total number of geoms in the internal scene will be
92            `num_visual_internal_geom` + `num_visual_user_geom`.";
93
94        num_visual_user_geom: u32;      "maximum number of additional visual-only user geoms (drawn by the user).";
95        rgb: bool;                      "RGB rendering enabled (true) or disabled (false).";
96        depth: bool;                    "depth rendering enabled (true) or disabled (false).";
97        font_scale: MjtFontScale;       "font scale of drawn text (with [MjrContext]).";
98        camera: MjvCamera;              "camera used for drawing.";
99        opts: MjvOption;                "visualization options.";
100
101    }
102
103    /// Builds a [`MjRenderer`].
104    pub fn build(self, model: &MjModel) -> Result<MjRenderer<'_>, RendererError> {
105        // Assume model's maximum should be used
106        let mut height = self.height;
107        let mut width = self.width;
108        if width == 0 && height == 0 {
109            let global = &model.vis().global;
110            height = global.offheight as u32;
111            width = global.offwidth as u32;
112        }
113
114        let mut glfw = glfw::init_no_callbacks()
115            .map_err(|err| RendererError::GlfwInitError(err))?;
116
117        // Create window for rendering
118        glfw.window_hint(WindowHint::Visible(false));
119        let (mut window, _) = match glfw.create_window(
120            width, height,
121            "", glfw::WindowMode::Windowed
122        ) {
123            Some(x) => Ok(x),
124            None => Err(RendererError::WindowCreationError)
125        }?;
126
127        window.make_current();
128        glfw.set_swap_interval(glfw::SwapInterval::None);
129
130        // Initialize the rendering context to render to the offscreen buffer.
131        let mut context = MjrContext::new(model);
132        context.offscreen();
133
134        // The 3D scene for visualization
135        let scene = MjvScene::new(
136            model,
137            model.ffi().ngeom as usize + self.num_visual_internal_geom as usize
138            + self.num_visual_user_geom as usize
139        );
140
141        let user_scene = MjvScene::new(
142            model,
143            self.num_visual_user_geom as usize
144        );
145
146        // Construct the renderer and create allocated buffers.
147        let renderer = MjRenderer {
148            scene, user_scene, context, window, model, camera: self.camera, option: self.opts,
149            flags: RendererFlags::empty(), rgb: None, depth: None,
150            width: width as usize, height: height as usize
151        }   // These require special care
152            .with_rgb_rendering(self.rgb)
153            .with_depth_rendering(self.depth);
154
155        Ok(renderer)
156    }
157}
158
159
160impl Default for MjRendererBuilder {
161    fn default() -> Self {
162        Self::new()
163    }
164}
165
166/// A renderer for rendering 3D scenes.
167/// By default, RGB rendering is enabled and depth rendering is disabled.
168pub struct MjRenderer<'m> {
169    scene: MjvScene<'m>,
170    user_scene: MjvScene<'m>,
171    context: MjrContext,
172    model: &'m MjModel,
173
174    /* Glfw */
175    window: PWindow,
176
177    /* Configuration */
178    camera: MjvCamera,
179    option: MjvOption,
180    flags: RendererFlags,
181
182    /* Storage */
183    // Use Box to allow less space to be used
184    // when rgb or depth rendering is disabled
185    rgb: Option<Box<[u8]>>,
186    depth: Option<Box<[f32]>>,
187
188    width: usize,
189    height: usize,
190}
191
192impl<'m> MjRenderer<'m> {
193    /// Construct a new renderer.
194    /// The `max_geom` parameter
195    /// defines how much space will be allocated for additional, user-defined visual-only geoms.
196    /// It can thus be set to 0 if no additional geoms will be drawn by the user.
197    /// # Scene allocation
198    /// The renderer uses two scenes:
199    /// - the internal scene: used by the renderer to draw the model's state.
200    /// - the user scene: used by the user to add additional geoms to the internal scene
201    /// 
202    /// The **internal scene** allocates the amount of space needed to fit every pre-existing
203    /// model geom + user visual-only geoms + additional visual-only geoms that aren't from the user (e.g., tendons).
204    /// By default, the renderer reserves 100 extra geom slots for drawing the additional visual-only geoms.
205    /// If that is not enough or it is too much, you can construct [`MjRenderer`] via its builder
206    /// ([`MjRenderer::builder`]), which allows more configuration. 
207    /// 
208    /// <div class="warning">
209    /// 
210    /// Parameters `width` and `height` must be less or equal to the offscreen buffer size,
211    /// which can be configured at the top of the model's XML like so:
212    /// 
213    /// ```xml
214    /// <visual>
215    ///    <global offwidth="1920" offheight="1080"/>
216    /// </visual>
217    /// ```
218    /// 
219    /// </div>
220    pub fn new(model: &'m MjModel, width: usize, height: usize, max_geom: usize) -> Result<Self, RendererError> {
221        let mut glfw = glfw::init_no_callbacks()
222            .map_err(|err| RendererError::GlfwInitError(err))?;
223
224        /* Create window for rendering */
225        glfw.window_hint(WindowHint::Visible(false));
226        let (mut _window, _) = match glfw.create_window(width as u32, height as u32, "", glfw::WindowMode::Windowed) {
227            Some(x) => Ok(x),
228            None => Err(RendererError::WindowCreationError)
229        }?;
230
231        _window.make_current();
232        glfw.set_swap_interval(glfw::SwapInterval::None);
233
234        /* Initialize the rendering context to render to the offscreen buffer. */
235        let mut context = MjrContext::new(model);
236        context.offscreen();
237
238        /* The 3D scene for visualization */
239        let scene = MjvScene::new(model, model.ffi().ngeom as usize + max_geom + EXTRA_INTERNAL_VISUAL_GEOMS);
240        let user_scene = MjvScene::new(model, max_geom);
241
242        let camera = MjvCamera::new_free(model);
243        let option = MjvOption::default();
244
245        let mut s = Self {
246            scene, user_scene, context, window: _window, model, camera, option,
247            flags: RendererFlags::empty(), rgb: None, depth: None,
248            width, height
249        };
250
251        s = s.with_rgb_rendering(true);
252        Ok(s)
253    }
254
255    /// Create a [`MjRendererBuilder`] to configure [`MjRenderer`].
256    pub fn builder() -> MjRendererBuilder {
257        MjRendererBuilder::new()
258    }
259
260    /// Return an immutable reference to the internal scene.
261    pub fn scene(&self) -> &MjvScene<'m>{
262        &self.scene
263    }
264
265    /// Return an immutable reference to a user scene for drawing custom visual-only geoms.
266    pub fn user_scene(&self) -> &MjvScene<'m>{
267        &self.user_scene
268    }
269
270    /// Return a mutable reference to a user scene for drawing custom visual-only geoms.
271    pub fn user_scene_mut(&mut self) -> &mut MjvScene<'m>{
272        &mut self.user_scene
273    }
274
275    /// Return an immutable reference to visualization options.
276    pub fn opts(&self) -> &MjvOption {
277        &self.option
278    }
279
280    /// Return a mutable reference to visualization options.
281    pub fn opts_mut(&mut self) -> &mut MjvOption {
282        &mut self.option
283    }
284
285    /// Return an immutable reference to the camera.
286    pub fn camera(&self) -> &MjvCamera {
287        &self.camera
288    }
289
290    /// Return a mutable reference to the camera.
291    pub fn camera_mut(&mut self) -> &mut MjvCamera {
292        &mut self.camera
293    }
294
295    /// Check if RGB rendering is enabled.
296    pub fn rgb_enabled(&self) -> bool {
297        self.flags.contains(RendererFlags::RENDER_RGB)
298    }
299
300    /// Check if depth rendering is enabled.
301    pub fn depth_enabled(&self) -> bool {
302        self.flags.contains(RendererFlags::RENDER_DEPTH)
303    }
304
305    /// Sets the font size.
306    pub fn set_font_scale(&mut self, font_scale: MjtFontScale) {
307        self.context.change_font(font_scale);
308    }
309
310    /// Update the visualization options and return a reference to self.
311    pub fn set_opts(&mut self, options: MjvOption) {
312        self.option = options;
313    }
314
315    /// Set the camera used for rendering.
316    pub fn set_camera(&mut self, camera: MjvCamera)  {
317        self.camera = camera;
318    }
319
320    /// Enables/disables RGB rendering.
321    pub fn set_rgb_rendering(&mut self, enable: bool) {
322        self.flags.set(RendererFlags::RENDER_RGB, enable);
323        self.rgb = if enable { Some(vec![0; 3 * self.width * self.height].into_boxed_slice()) } else { None } ;
324    }
325
326    /// Enables/disables depth rendering.
327    pub fn set_depth_rendering(&mut self, enable: bool) {
328        self.flags.set(RendererFlags::RENDER_DEPTH, enable);
329        self.depth = if enable { Some(vec![0.0; self.width * self.height].into_boxed_slice()) } else { None } ;
330    }
331
332    /// Sets the font size. To be used on construction.
333    pub fn with_font_scale(mut self, font_scale: MjtFontScale) -> Self {
334        self.set_font_scale(font_scale);
335        self
336    }
337
338    /// Update the visualization options and return a reference to self. To be used on construction.
339    pub fn with_opts(mut self, options: MjvOption) -> Self {
340        self.set_opts(options);
341        self
342    }
343
344    /// Set the camera used for rendering. To be used on construction.
345    pub fn with_camera(mut self, camera: MjvCamera) -> Self  {
346        self.set_camera(camera);
347        self
348    }
349
350    /// Enables/disables RGB rendering. To be used on construction.
351    pub fn with_rgb_rendering(mut self, enable: bool) -> Self {
352        self.set_rgb_rendering(enable);
353        self
354    }
355
356    /// Enables/disables depth rendering. To be used on construction.
357    pub fn with_depth_rendering(mut self, enable: bool) -> Self {
358        self.set_depth_rendering(enable);
359        self
360    }
361
362    /// Update the scene with new data from data.
363    pub fn sync(&mut self, data: &mut MjData) {
364        let model_data_ptr = unsafe {  data.model().__raw() };
365        let bound_model_ptr = unsafe { self.model.__raw() };
366        assert_eq!(model_data_ptr, bound_model_ptr, "'data' must be created from the same model as the renderer.");
367
368        self.scene.update(data, &self.option, &MjvPerturb::default(), &mut self.camera);
369
370        /* Draw user scene geoms */
371        sync_geoms(&self.user_scene, &mut self.scene)
372            .expect("could not sync the user scene with the internal scene; this is a bug, please report it.");
373
374        self.render();
375    }
376
377    /// Return a flattened RGB image of the scene.
378    pub fn rgb_flat(&self) -> Option<&[u8]> {
379        self.rgb.as_deref()
380    }
381
382    /// Return an RGB image of the scene. This methods accepts two generic parameters <WIDTH, HEIGHT>
383    /// that define the shape of the output slice.
384    pub fn rgb<const WIDTH: usize, const HEIGHT: usize>(&self) -> io::Result<&[[[u8; 3]; WIDTH]; HEIGHT]> {
385        if let Some(flat) = self.rgb_flat() {
386            if flat.len() == WIDTH * HEIGHT * 3 {
387                let p_shaped = flat.as_ptr() as *const [[[u8; 3]; WIDTH]; HEIGHT];
388
389                // SAFETY: The alignment of the output is the same as the original.
390                // The lifetime also matches  'a in &'a self, which prevents data races.
391                // Length (number of elements) matches the output's.
392                Ok(unsafe { p_shaped.as_ref().unwrap() })
393            }
394            else {
395                Err(io::Error::new(io::ErrorKind::InvalidInput, INVALID_INPUT_SIZE))
396            }
397        }
398        else {
399            Err(io::Error::new(io::ErrorKind::NotFound, RGB_NOT_FOUND_ERR_STR))
400        }
401    }
402
403    /// Return a flattened depth image of the scene.
404    pub fn depth_flat(&self) -> Option<&[f32]> {
405        self.depth.as_deref()
406    }
407
408    /// Return a depth image of the scene. This methods accepts two generic parameters <WIDTH, HEIGHT>
409    /// that define the shape of the output slice.
410    pub fn depth<const WIDTH: usize, const HEIGHT: usize>(&self) -> io::Result<&[[f32; WIDTH]; HEIGHT]> {
411        if let Some(flat) = self.depth_flat() {
412            if flat.len() == WIDTH * HEIGHT {
413                let p_shaped = flat.as_ptr() as *const [[f32; WIDTH]; HEIGHT];
414
415                // SAFETY: The alignment of the output is the same as the original.
416                // The lifetime matches  'a in &'a self, which prevents data races.
417                // Length (number of elements) matches the output's.
418                Ok(unsafe { p_shaped.as_ref().unwrap() })
419            }
420            else {
421                Err(io::Error::new(io::ErrorKind::InvalidInput, INVALID_INPUT_SIZE))
422            }
423        }
424        else {
425            Err(io::Error::new(io::ErrorKind::NotFound, DEPTH_NOT_FOUND_ERR_STR))
426        }
427    }
428
429    /// Save an RGB image of the scene to a path.
430    /// # Errors
431    /// - [`ErrorKind::NotFound`] when RGB rendering is disabled,
432    /// - other errors related to write.
433    pub fn save_rgb<T: AsRef<Path>>(&self, path: T) -> io::Result<()> {
434        if let Some(rgb) = &self.rgb {
435            let file = File::create(path.as_ref())?;
436            let w = BufWriter::new(file);
437
438            let mut encoder = Encoder::new(w, self.width as u32, self.height as u32);
439            encoder.set_color(png::ColorType::Rgb);
440            encoder.set_depth(png::BitDepth::Eight);
441            encoder.set_compression(png::Compression::NoCompression);
442
443            let mut writer = encoder.write_header()?;
444            writer.write_image_data(rgb)?;
445            Ok(())
446        }
447        else {
448            Err(io::Error::new(ErrorKind::NotFound, RGB_NOT_FOUND_ERR_STR))
449        }
450    }
451
452    /// Save a depth image of the scene to a path. The image is 16-bit PNG, which
453    /// can be converted into depth (distance) data by dividing the grayscale values by
454    /// 65535.0 and applying inverse normalization: `depth = min + (grayscale / 65535.0) * (max - min)`.
455    /// If `normalize` is `true`, then the data is normalized with min-max normalization.
456    /// Use of [`MjRenderer::save_depth_raw`] is recommended if performance is critical, as
457    /// it skips PNG encoding and also saves the true depth values directly.
458    /// # Returns
459    /// An [`Ok`]`((min, max))` is returned, where min and max represent the normalization parameters.
460    /// # Errors
461    /// - [`ErrorKind::NotFound`] when depth rendering is disabled,
462    /// - other errors related to write.
463    pub fn save_depth<T: AsRef<Path>>(&self, path: T, normalize: bool) -> io::Result<(f32, f32)> {
464        if let Some(depth) = &self.depth {
465            let file = File::create(path.as_ref())?;
466            let w = BufWriter::new(file);
467
468            let mut encoder = Encoder::new(w, self.width as u32, self.height as u32);
469            encoder.set_color(png::ColorType::Grayscale);
470            encoder.set_depth(png::BitDepth::Sixteen);
471            encoder.set_compression(png::Compression::NoCompression);
472
473            let (norm, min, max) =
474            if normalize {
475                let max = depth.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
476                let min = depth.iter().cloned().fold(f32::INFINITY, f32::min);
477                (depth.iter().flat_map(|&x| (((x - min) / (max - min) * 65535.0).min(65535.0) as u16).to_be_bytes()).collect::<Box<_>>(), min, max)
478            }
479            else {
480                (depth.iter().flat_map(|&x| ((x * 65535.0).min(65535.0) as u16).to_be_bytes()).collect::<Box<_>>(), 0.0, 1.0)
481            };
482
483            let mut writer = encoder.write_header()?;
484            writer.write_image_data(&norm)?;
485            Ok((min, max))
486        }
487        else {
488            Err(io::Error::new(ErrorKind::NotFound, DEPTH_NOT_FOUND_ERR_STR))
489        }
490    }
491
492    /// Save the raw depth data to the `path`. The data is encoded
493    /// as a sequence of bytes, where groups of four represent a single f32 value.
494    /// The lower bytes of individual f32 appear first (low-endianness).
495    /// # Errors
496    /// - [`ErrorKind::NotFound`] when depth rendering is disabled,
497    /// - other errors related to write.
498    pub fn save_depth_raw<T: AsRef<Path>>(&self, path: T) -> io::Result<()> {
499        if let Some(depth) = &self.depth {
500            let file = File::create(path.as_ref())?;
501            let mut writer = BufWriter::new(file);
502
503            /* Fast conversion to a byte slice to prioritize performance */
504            let p = unsafe { std::slice::from_raw_parts(
505                depth.as_ptr() as *const u8,
506                std::mem::size_of::<f32>() * depth.len()
507            ) };
508
509            writer.write_all(p)?;
510            Ok(())
511        }
512        else {
513            Err(io::Error::new(ErrorKind::NotFound, DEPTH_NOT_FOUND_ERR_STR))
514        }
515    }
516
517    /// Draws the scene to internal arrays.
518    /// Use [`MjRenderer::rgb`] or [`MjRenderer::depth`] to obtain the rendered image.
519    fn render(&mut self) {
520        self.window.make_current();
521        let vp = MjrRectangle::new(0, 0, self.width as i32, self.height as i32);
522        self.scene.render(&vp, &self.context);
523
524        /* Fully flatten everything */
525        let flat_rgb = self.rgb.as_deref_mut();
526        let flat_depth = self.depth.as_deref_mut();
527
528        /* Read to whatever is enabled */
529        self.context.read_pixels(
530            flat_rgb,
531            flat_depth,
532            &vp
533        );
534
535        /* Make depth values be the actual distance in meters */
536        if let Some(depth) = self.depth.as_deref_mut() {
537            let map = &self.model.vis().map;
538            let near = map.znear;
539            let far = map.zfar;
540            for value in depth {
541                let z_ndc = 2.0 * *value - 1.0;
542                *value = 2.0 * near * far / (far + near - z_ndc * (far - near));
543            }
544        }
545    }
546}
547
548
549#[derive(Debug)]
550pub enum RendererError {
551    GlfwInitError(InitError),
552    WindowCreationError,
553}
554
555impl Display for RendererError {
556    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
557        match self {
558            Self::GlfwInitError(e) => write!(f, "glfw failed to initialize: {}", e),
559            Self::WindowCreationError => write!(f, "failed to create window"),
560        }
561    }
562}
563
564impl Error for RendererError {
565    fn source(&self) -> Option<&(dyn Error + 'static)> {
566        match self {
567            Self::GlfwInitError(e) => Some(e),
568            _ => None
569        }
570    }
571}
572
573bitflags! { 
574    /// Flags that enable features of the renderer.
575    struct RendererFlags: u8 {
576        const RENDER_RGB = 1 << 0;
577        const RENDER_DEPTH = 1 << 1;
578    }
579}
580
581
582
583/* 
584** Don't run any tests as OpenGL hates if anything
585** runs outside the main thread.
586*/
587
588// #[cfg(test)]
589// mod test {
590//     use super::*;
591
592//     const MODEL: &str = "
593//         <mujoco>
594//             <worldbody>
595//             </worldbody>
596//         </mujoco>
597//     ";
598
599//     // #[test]
600//     // fn test_update_normal() {
601//     //     // let model = MjModel::from_xml_string(MODEL).unwrap();
602//     //     // let model2 = MjModel::from_xml_string(MODEL).unwrap();
603//     //     // let mut data = model.make_data();
604//     //     // let mut renderer = Renderer::new(&model, 720, 1280).unwrap();
605        
606//     //     // /* Check if scene updates without errors. */
607//     //     // renderer.update_scene(&mut data);
608//     // }
609// }
610