Skip to content

renderer

renderer

Modules:

Name Description
abstract_renderer
filament_rendering
opengl_context
opengl_rendering

abstract_renderer

Classes:

Name Description
MjAbstractRenderer

Attributes:

Name Type Description
RENDERING_COMPLETE

RENDERING_COMPLETE module-attribute 

RENDERING_COMPLETE = 'RENDERING_COMPLETE'

MjAbstractRenderer 

MjAbstractRenderer(model: MjModel, device_id: int | None = None)

Bases: ABC

Methods:

Name Description
close
disable_depth_rendering
disable_segmentation_rendering
enable_depth_rendering
enable_segmentation_rendering
render
update

Attributes:

Name Type Description
device_id int | None
model MjModel
render_outputs list[Any]
scene MjvScene
Source code in molmo_spaces/renderer/abstract_renderer.py 
12
13
14
15
16
def __init__(self, model: mj.MjModel, device_id: int | None = None) -> None:
    self._model = model
    self._device_id = device_id

    self._scene: mj.MjvScene | None = None
device_id property 
device_id: int | None
model property 
model: MjModel
render_outputs instance-attribute 
render_outputs: list[Any]
scene property 
scene: MjvScene
close abstractmethod 
close() -> None
Source code in molmo_spaces/renderer/abstract_renderer.py 
43
44
@abc.abstractmethod
def close(self) -> None: ...
disable_depth_rendering abstractmethod 
disable_depth_rendering() -> None
Source code in molmo_spaces/renderer/abstract_renderer.py 
34
35
@abc.abstractmethod
def disable_depth_rendering(self) -> None: ...
disable_segmentation_rendering abstractmethod 
disable_segmentation_rendering() -> None
Source code in molmo_spaces/renderer/abstract_renderer.py 
40
41
@abc.abstractmethod
def disable_segmentation_rendering(self) -> None: ...
enable_depth_rendering abstractmethod 
enable_depth_rendering() -> None
Source code in molmo_spaces/renderer/abstract_renderer.py 
31
32
@abc.abstractmethod
def enable_depth_rendering(self) -> None: ...
enable_segmentation_rendering abstractmethod 
enable_segmentation_rendering() -> None
Source code in molmo_spaces/renderer/abstract_renderer.py 
37
38
@abc.abstractmethod
def enable_segmentation_rendering(self) -> None: ...
render abstractmethod 
render(*args, **kwargs) -> Any
Source code in molmo_spaces/renderer/abstract_renderer.py 
46
47
@abc.abstractmethod
def render(self, *args, **kwargs) -> Any: ...
update abstractmethod 
update(*args, **kwargs) -> None
Source code in molmo_spaces/renderer/abstract_renderer.py 
49
50
@abc.abstractmethod
def update(self, *args, **kwargs) -> None: ...

filament_rendering

Classes:

Name Description
Args
MjFilamentRenderer

Attributes:

Name Type Description
args
data
image
model
pil_image
renderer

args module-attribute 

args = tyro.cli(Args)

data module-attribute 

data = mj.MjData(model)

image module-attribute 

image = renderer.render()

model module-attribute 

model = mj.MjModel.from_xml_path(args.model.as_posix())

pil_image module-attribute 

pil_image = Image.fromarray(image)

renderer module-attribute 

renderer = MjFilamentRenderer(model)

Args dataclass 

Args(model: Path)

Attributes:

Name Type Description
model Path
model instance-attribute 
model: Path

MjFilamentRenderer 

MjFilamentRenderer(model: MjModel, device_id: int | None = None, height: int = 720, width: int = 1280, max_geom: int = 10000)

Bases: MjAbstractRenderer

Methods:

Name Description
close
disable_depth_rendering
disable_segmentation_rendering
enable_depth_rendering
enable_segmentation_rendering
geomid_to_bodyid
mark_textures_dirty
render
render_rgb
update
upload_textures

Attributes:

Name Type Description
device_id int | None
height
model MjModel
render_outputs list[Any]
scene MjvScene
width
Source code in molmo_spaces/renderer/filament_rendering.py 
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
def __init__(
    self,
    model: mj.MjModel,
    device_id: int | None = None,
    height: int = 720,
    width: int = 1280,
    max_geom: int = 10000,
) -> None:
    super().__init__(model, device_id)

    self._width = width
    self._height = height

    self._model = model

    self._scene = mj.MjvScene(model=model, maxgeom=max_geom)
    self._scene_option = mj.MjvOption()

    self._scene_option.sitegroup *= 0
    self._scene.flags[mj.mjtRndFlag.mjRND_SHADOW] = True

    self._depth_rendering = False
    self._segmentation_rendering = False

    self._mjr_context = mj.MjrContext(model, mj.mjtFontScale.mjFONTSCALE_150.value)
    # mj.mjr_resizeOffscreen(width, height, self._mjr_context)
    mj.mjr_setBuffer(mj.mjtFramebuffer.mjFB_OFFSCREEN.value, self._mjr_context)
    self._mjr_context.readDepthMap = mj.mjtDepthMap.mjDEPTH_ZEROFAR.value

    self._textures_need_upload = False
device_id property 
device_id: int | None
height property 
height
model property 
model: MjModel
render_outputs instance-attribute 
render_outputs: list[Any]
scene property 
scene: MjvScene
width property 
width
close 
close() -> None
Source code in molmo_spaces/renderer/filament_rendering.py 
307
308
309
310
def close(self) -> None:
    if hasattr(self, "_mjr_context") and self._mjr_context:
        self._mjr_context.free()
    self._mjr_context = None
disable_depth_rendering 
disable_depth_rendering() -> None
Source code in molmo_spaces/renderer/filament_rendering.py 
51
52
def disable_depth_rendering(self) -> None:
    self._depth_rendering = False
disable_segmentation_rendering 
disable_segmentation_rendering() -> None
Source code in molmo_spaces/renderer/filament_rendering.py 
58
59
def disable_segmentation_rendering(self) -> None:
    self._segmentation_rendering = False
enable_depth_rendering 
enable_depth_rendering() -> None
Source code in molmo_spaces/renderer/filament_rendering.py 
47
48
49
def enable_depth_rendering(self) -> None:
    self._segmentation_rendering = False
    self._depth_rendering = True
enable_segmentation_rendering 
enable_segmentation_rendering() -> None
Source code in molmo_spaces/renderer/filament_rendering.py 
54
55
56
def enable_segmentation_rendering(self) -> None:
    self._segmentation_rendering = True
    self._depth_rendering = False
geomid_to_bodyid 
geomid_to_bodyid(geomid)
Source code in molmo_spaces/renderer/filament_rendering.py 
61
62
def geomid_to_bodyid(self, geomid):
    return self.model.geom_bodyid[geomid]
mark_textures_dirty 
mark_textures_dirty() -> None
Source code in molmo_spaces/renderer/filament_rendering.py 
265
266
def mark_textures_dirty(self) -> None:
    self._textures_need_upload = True
render 
render(*, out: ndarray | None = None, width: int | None = None, height: int | None = None) -> ndarray
Source code in molmo_spaces/renderer/filament_rendering.py 
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
def render(
    self,
    *,
    out: np.ndarray | None = None,
    width: int | None = None,
    height: int | None = None,
) -> np.ndarray:
    assert self._scene is not None, "Internal scene:MjvScene must be initialized by now"

    height = height or self._height
    width = width or self._width
    rect = mj.MjrRect(0, 0, width, height)

    original_flags = self._scene.flags.copy()

    # Enable shadow rendering (required for shadows to appear in rendered images)
    # Shadows are controlled by lights with castshadow enabled
    self._scene.flags[mj.mjtRndFlag.mjRND_SHADOW] = True

    # Using segmented rendering for depth makes the calculated depth more
    # accurate at far distances.
    if self._depth_rendering or self._segmentation_rendering:
        self._scene.flags[mj.mjtRndFlag.mjRND_SEGMENT] = True
        self._scene.flags[mj.mjtRndFlag.mjRND_IDCOLOR] = True

    # Upload textures to GPU before rendering if textures have been modified
    # This is necessary when textures are modified via model.tex_data
    # Only upload when needed to avoid performance overhead
    if self._textures_need_upload:
        self.upload_textures()
        self._textures_need_upload = False

    if self._depth_rendering:
        out_shape = (rect.height, rect.width)
        out_dtype = np.float32
    else:
        out_shape = (rect.height, rect.width, 3)
        out_dtype = np.uint8

    if out is None:
        out = np.empty(out_shape, dtype=out_dtype)
    else:
        if out.shape != out_shape:
            raise ValueError(
                f"Expected `out.shape == {out_shape}`. Got `out.shape={out.shape}`"
                " instead. When using depth rendering, the out array should be of"
                " shape `(width, height)` and otherwise (width, height, 3)."
                f" Got `(self.height, self.width)={(self.height, self.width)}` and"
                f" `self._depth_rendering={self._depth_rendering}`."
            )

    assert self._mjr_context, "MjrContext must be created by now, but it's None"
    mj.mjr_render(rect, self._scene, self._mjr_context)

    if self._depth_rendering:
        mj.mjr_readPixels(rgb=None, depth=out, viewport=rect, con=self._mjr_context)

        # Get the distances to the near and far clipping planes.
        extent = self.model.stat.extent
        near = self.model.vis.map.znear * extent
        far = self.model.vis.map.zfar * extent

        # Calculate OpenGL perspective matrix values in float32 precision
        # so they are close to what glFrustum returns
        # https://registry.khronos.org/OpenGL-Refpages/gl2.1/xhtml/glFrustum.xml
        zfar = np.float32(far)
        znear = np.float32(near)
        c_coef = -(zfar + znear) / (zfar - znear)
        d_coef = -(np.float32(2) * zfar * znear) / (zfar - znear)

        # In reverse Z mode the perspective matrix is transformed by the following
        c_coef = np.float32(-0.5) * c_coef - np.float32(0.5)
        d_coef = np.float32(-0.5) * d_coef

        # We need 64 bits to convert Z from ndc to metric depth without noticeable
        # losses in precision
        out_64 = out.astype(np.float64)

        # Undo OpenGL projection
        # Note: We do not need to take action to convert from window coordinates
        # to normalized device coordinates because in reversed Z mode the mapping
        # is identity
        out_64 = d_coef / (out_64 + c_coef)

        # Cast result back to float32 for backwards compatibility
        # This has a small accuracy cost
        out[:] = out_64.astype(np.float32)

        # Reset scene flags.
        np.copyto(self._scene.flags, original_flags)
    elif self._segmentation_rendering:
        mj.mjr_readPixels(rgb=out, depth=None, viewport=rect, con=self._mjr_context)

        # Convert 3-channel uint8 to 1-channel uint32.
        image3 = out.astype(np.uint32)
        segimage = image3[:, :, 0] + image3[:, :, 1] * (2**8) + image3[:, :, 2] * (2**16)
        # Remap segid to 3-channel (object ID, object type, body ID) triplet
        # Seg ID 0 is background -- will be remapped to (-1, -1, -1).

        # Find the maximum segment ID in the image to size the output array correctly
        max_segid = np.max(segimage) if segimage.size > 0 else 0

        # Create output array with size to accommodate all possible segment IDs
        # Add 1 to account for 0-based indexing and ensure we have enough space
        segid2output = np.full((max_segid + 1, 3), fill_value=-1, dtype=np.int32)

        visible_geoms = [g for g in self._scene.geoms[: self._scene.ngeom] if g.segid != -1]
        visible_segids = np.array([g.segid + 1 for g in visible_geoms], np.int32)
        visible_objid = np.array([g.objid for g in visible_geoms], np.int32)
        visible_objtype = np.array([g.objtype for g in visible_geoms], np.int32)
        visible_bodyid = np.array(
            [self.geomid_to_bodyid(g.objid) for g in visible_geoms], np.int32
        )

        # Only set values for valid segment IDs that are within bounds
        valid_mask = (visible_segids >= 0) & (visible_segids < segid2output.shape[0])
        if np.any(valid_mask):
            segid2output[visible_segids[valid_mask], 0] = visible_objid[valid_mask]
            segid2output[visible_segids[valid_mask], 1] = visible_objtype[valid_mask]
            segid2output[visible_segids[valid_mask], 2] = visible_bodyid[valid_mask]

        out = segid2output[segimage]

        # Reset scene flags.
        np.copyto(self._scene.flags, original_flags)
    else:
        mj.mjr_readPixels(rgb=out, depth=None, viewport=rect, con=self._mjr_context)
        mj.mjr_readPixels(rgb=out, depth=None, viewport=rect, con=self._mjr_context)

    return out
render_rgb 
render_rgb(*, out: ndarray | None = None, width: int | None = None, height: int | None = None) -> ndarray
Source code in molmo_spaces/renderer/filament_rendering.py 
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
def render_rgb(
    self,
    *,
    out: np.ndarray | None = None,
    width: int | None = None,
    height: int | None = None,
) -> np.ndarray:
    assert self._scene is not None, "Internal scene:MjvScene must be initialized by now"

    height = height or self._height
    width = width or self._width
    rect = mj.MjrRect(0, 0, width, height)

    # Enable shadow rendering (required for shadows to appear in rendered images)
    # Shadows are controlled by lights with castshadow enabled
    self._scene.flags[mj.mjtRndFlag.mjRND_SHADOW] = True

    # Using segmented rendering for depth makes the calculated depth more
    # accurate at far distances.
    if self._depth_rendering or self._segmentation_rendering:
        self._scene.flags[mj.mjtRndFlag.mjRND_SEGMENT] = True
        self._scene.flags[mj.mjtRndFlag.mjRND_IDCOLOR] = True

    # Upload textures to GPU before rendering if textures have been modified
    # This is necessary when textures are modified via model.tex_data
    # Only upload when needed to avoid performance overhead
    if self._textures_need_upload:
        self.upload_textures()
        self._textures_need_upload = False

    if self._depth_rendering:
        out_shape = (rect.height, rect.width)
        out_dtype = np.float32
    else:
        out_shape = (rect.height, rect.width, 3)
        out_dtype = np.uint8

    if out is None:
        out = np.empty(out_shape, dtype=out_dtype)
    else:
        if out.shape != out_shape:
            raise ValueError(
                f"Expected `out.shape == {out_shape}`. Got `out.shape={out.shape}`"
                " instead. When using depth rendering, the out array should be of"
                " shape `(width, height)` and otherwise (width, height, 3)."
                f" Got `(self.height, self.width)={(self.height, self.width)}` and"
                f" `self._depth_rendering={self._depth_rendering}`."
            )

    assert self._mjr_context, "MjrContext must be created by now, but it's None"
    mj.mjr_render(rect, self._scene, self._mjr_context)

    if self._depth_rendering:
        mj.mjr_readPixels(rgb=None, depth=out, viewport=rect, con=self._mjr_context)
    elif self._segmentation_rendering:
        mj.mjr_readPixels(rgb=out, depth=None, viewport=rect, con=self._mjr_context)
    else:
        mj.mjr_readPixels(rgb=out, depth=None, viewport=rect, con=self._mjr_context)
        mj.mjr_readPixels(rgb=out, depth=None, viewport=rect, con=self._mjr_context)

    return out
update 
update(data: MjData, camera: int | str | MjvCamera = -1, scene_option: MjvOption | None = None) -> None
Source code in molmo_spaces/renderer/filament_rendering.py 
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
def update(
    self,
    data: mj.MjData,
    camera: int | str | mj.MjvCamera = -1,
    scene_option: mj.MjvOption | None = None,
) -> None:
    assert self._scene is not None, "Internal scene:MjvScene must be initialized by now"

    if not isinstance(camera, mj.MjvCamera):
        camera_id = camera
        if isinstance(camera_id, str):
            camera_id = mj.mj_name2id(self.model, mj.mjtObj.mjOBJ_CAMERA.value, camera_id)
            if camera_id == -1:
                raise ValueError(f'The camera "{camera}" does not exist.')
        if camera_id < -1 or camera_id >= self.model.ncam:
            raise ValueError(
                f"The camera id {camera_id} is out of range [-1, {self.model.ncam})."
            )

        camera = mj.MjvCamera()
        camera.fixedcamid = camera_id

        if camera_id == -1:
            camera.type = mj.mjtCamera.mjCAMERA_FREE
            mj.mjv_defaultFreeCamera(self.model, camera)
        else:
            camera.type = mj.mjtCamera.mjCAMERA_FIXED

    scene_option = scene_option or self._scene_option
    mj.mjv_updateScene(
        self.model,
        data,
        scene_option,
        None,
        camera,
        mj.mjtCatBit.mjCAT_ALL.value,
        self._scene,
    )
upload_textures 
upload_textures() -> None
Source code in molmo_spaces/renderer/filament_rendering.py 
257
258
259
260
261
262
263
def upload_textures(self) -> None:
    assert self._mjr_context, "MjrContext must be created by now, but it's None"
    if self.model.ntex == 0:
        return

    for tex_id in range(self.model.ntex):
        mj.mjr_uploadTexture(self.model, self._mjr_context, tex_id)

opengl_context

Classes:

Name Description
EGLGLContext

An EGL context for headless accelerated OpenGL rendering on GPU devices.

Functions:

Name Description
create_initialized_egl_device_display

Creates an initialized EGL display directly on a device.

Attributes:

Name Type Description
EGL_ATTRIBUTES
EGL_DISPLAY
EGL_DISPLAY_INITIALIZED
PYOPENGL_PLATFORM
xla_flags

EGL_ATTRIBUTES module-attribute 

EGL_ATTRIBUTES = (EGL.EGL_RED_SIZE, 8, EGL.EGL_GREEN_SIZE, 8, EGL.EGL_BLUE_SIZE, 8, EGL.EGL_ALPHA_SIZE, 8, EGL.EGL_DEPTH_SIZE, 24, EGL.EGL_STENCIL_SIZE, 8, EGL.EGL_COLOR_BUFFER_TYPE, EGL.EGL_RGB_BUFFER, EGL.EGL_SURFACE_TYPE, EGL.EGL_PBUFFER_BIT, EGL.EGL_RENDERABLE_TYPE, EGL.EGL_OPENGL_BIT, EGL.EGL_NONE)

EGL_DISPLAY module-attribute 

EGL_DISPLAY = None

EGL_DISPLAY_INITIALIZED module-attribute 

EGL_DISPLAY_INITIALIZED = False

PYOPENGL_PLATFORM module-attribute 

PYOPENGL_PLATFORM = os.environ.get('PYOPENGL_PLATFORM')

xla_flags module-attribute 

xla_flags = os.environ.get('XLA_FLAGS', '')

EGLGLContext 

EGLGLContext(max_width, max_height, device_id=0)

An EGL context for headless accelerated OpenGL rendering on GPU devices.

Methods:

Name Description
__del__
free

Frees resources associated with this context.
make_current

Attributes:

Name Type Description
device_id
Source code in molmo_spaces/renderer/opengl_context.py 
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
def __init__(self, max_width, max_height, device_id=0) -> None:
    global EGL_DISPLAY, EGL_DISPLAY_INITIALIZED
    del max_width, max_height  # unused
    self.device_id = device_id
    num_configs = ctypes.c_long()
    config_size = 1
    config = EGL.EGLConfig()
    EGL.eglReleaseThread()

    if not EGL_DISPLAY_INITIALIZED:
        # only initialize for the first time
        EGL_DISPLAY = create_initialized_egl_device_display(device_id=device_id)
        if EGL_DISPLAY == EGL.EGL_NO_DISPLAY:
            raise ImportError(
                "Cannot initialize a EGL device display. This likely means that your EGL "
                "driver does not support the PLATFORM_DEVICE extension, which is "
                "required for creating a headless rendering context."
            )
        atexit.register(EGL.eglTerminate, EGL_DISPLAY)
        EGL_DISPLAY_INITIALIZED = True
    EGL.eglChooseConfig(
        EGL_DISPLAY, EGL_ATTRIBUTES, ctypes.byref(config), config_size, num_configs
    )
    if num_configs.value < 1:
        raise RuntimeError(
            "EGL failed to find a framebuffer configuration that matches the "
            f"desired attributes: {EGL_ATTRIBUTES}"
        )
    EGL.eglBindAPI(EGL.EGL_OPENGL_API)
    self._context = EGL.eglCreateContext(EGL_DISPLAY, config, EGL.EGL_NO_CONTEXT, None)
    if not self._context:
        raise RuntimeError("Cannot create an EGL context.")
device_id instance-attribute 
device_id = device_id
__del__ 
__del__() -> None
Source code in molmo_spaces/renderer/opengl_context.py 
184
185
def __del__(self) -> None:
    self.free()
free 
free() -> None

Frees resources associated with this context.

Source code in molmo_spaces/renderer/opengl_context.py 
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
def free(self) -> None:
    """Frees resources associated with this context."""
    global EGL_DISPLAY, EGL_DISPLAY_INITIALIZED
    if self._context and EGL_DISPLAY_INITIALIZED:
        try:
            current_context = EGL.eglGetCurrentContext()
            if current_context and self._context.address == current_context.address:
                EGL.eglMakeCurrent(
                    EGL_DISPLAY,
                    EGL.EGL_NO_SURFACE,
                    EGL.EGL_NO_SURFACE,
                    EGL.EGL_NO_CONTEXT,
                )
            EGL.eglDestroyContext(EGL_DISPLAY, self._context)
            EGL.eglReleaseThread()
        except EGLError:
            # Display may have already been terminated by atexit handler
            # during exception cleanup. Nothing we can do here.
            pass
    self._context = None
make_current 
make_current() -> None
Source code in molmo_spaces/renderer/opengl_context.py 
155
156
157
158
159
160
161
def make_current(self) -> None:
    global EGL_DISPLAY
    if not EGL.eglMakeCurrent(
        EGL_DISPLAY, EGL.EGL_NO_SURFACE, EGL.EGL_NO_SURFACE, self._context
    ):
        error = EGL.eglGetError()
        raise RuntimeError(f"Failed to make the EGL context current. EGL error: {error}")

create_initialized_egl_device_display 

create_initialized_egl_device_display(device_id=0)

Creates an initialized EGL display directly on a device.

Source code in molmo_spaces/renderer/opengl_context.py 
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
def create_initialized_egl_device_display(device_id=0):
    """Creates an initialized EGL display directly on a device."""
    all_devices = EGL.eglQueryDevicesEXT()
    selected_device = (
        os.environ.get("CUDA_VISIBLE_DEVICES", None)
        if os.environ.get("MUJOCO_EGL_DEVICE_ID", None) is None
        else os.environ.get("MUJOCO_EGL_DEVICE_ID", None)
    )
    if selected_device is None:
        candidates = all_devices
        device_idx = 0 if device_id == -1 else device_id
    else:
        if not selected_device.isdigit():
            device_inds = [int(x) for x in selected_device.split(",")]
            if device_id == -1:
                device_idx = device_inds[0]
            else:
                assert device_id in device_inds, (
                    "specified device id is not made visible in environment variables."
                )
                device_idx = device_id
        else:
            device_idx = int(selected_device)
        if not 0 <= device_idx < len(all_devices):
            raise RuntimeError(
                f"The MUJOCO_EGL_DEVICE_ID environment variable must be an integer "
                f"between 0 and {len(all_devices) - 1} (inclusive), got {device_idx}."
            )
    candidates = all_devices[device_idx : device_idx + 1]
    for device in candidates:
        display = EGL.eglGetPlatformDisplayEXT(EGL.EGL_PLATFORM_DEVICE_EXT, device, None)
        if display != EGL.EGL_NO_DISPLAY and EGL.eglGetError() == EGL.EGL_SUCCESS:
            # `eglInitialize` may or may not raise an exception on failure depending
            # on how PyOpenGL is configured. We therefore catch a `GLError` and also
            # manually check the output of `eglGetError()` here.
            try:
                initialized = EGL.eglInitialize(display, None, None)
            except error.GLError:
                pass
            else:
                if initialized == EGL.EGL_TRUE and EGL.eglGetError() == EGL.EGL_SUCCESS:
                    return display
    return EGL.EGL_NO_DISPLAY

opengl_rendering

Classes:

Name Description
Args
MjOpenGLRenderer

Attributes:

Name Type Description
args
data
image
model
pil_image
renderer

args module-attribute 

args = tyro.cli(Args)

data module-attribute 

data = mj.MjData(model)

image module-attribute 

image = renderer.render()

model module-attribute 

model = mj.MjModel.from_xml_path(args.model.as_posix())

pil_image module-attribute 

pil_image = Image.fromarray(image)

renderer module-attribute 

renderer = MjOpenGLRenderer(model)

Args dataclass 

Args(model: Path)

Attributes:

Name Type Description
model Path
model instance-attribute 
model: Path

MjOpenGLRenderer 

MjOpenGLRenderer(model: MjModel, device_id: int | None = None, height: int = 720, width: int = 1280, max_geom: int = 10000)

Bases: MjAbstractRenderer

Methods:

Name Description
__del__
__enter__
__exit__
close
disable_depth_rendering
disable_segmentation_rendering
enable_depth_rendering
enable_segmentation_rendering
geomid_to_bodyid
mark_textures_dirty

Mark that textures have been modified and need to be uploaded.
render

Renders the scene as a numpy array of pixel values.
update

Updates geometry used for rendering.
upload_textures

Upload all textures to the GPU render context.

Attributes:

Name Type Description
device_id int | None
height
model MjModel
render_outputs list[Any]
scene MjvScene
width
Source code in molmo_spaces/renderer/opengl_rendering.py 
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
def __init__(
    self,
    model: mj.MjModel,
    device_id: int | None = None,
    height: int = 720,
    width: int = 1280,
    max_geom: int = 10000,
) -> None:
    # TODO(wilbert): remove this an use gpustat instead of full torch, and only
    # if using linux. On MacOS we don't have multi-gpu so makes no sense to try to
    # pass a device_id, right?
    if device_id is None:
        try:
            import torch

            if torch.cuda.is_available():
                device_id = 0
        except ImportError:
            pass

    super().__init__(model, device_id)

    self._width = width
    self._height = height

    self._model = model

    self._scene = mj.MjvScene(model=model, maxgeom=max_geom)
    self._scene_option = mj.MjvOption()

    self._scene_option.sitegroup *= 0
    self._scene.flags[mj.mjtRndFlag.mjRND_SHADOW] = True

    self._depth_rendering = False
    self._segmentation_rendering = False

    # TODO(nimrod): Figure out why pytype doesn't like gl_context.GLContext
    self._context_is_cgl = False
    if device_id is None:
        from mujoco import gl_context

        self._gl_context = gl_context.GLContext(width, height)
        self._context_is_cgl = sys.platform == "darwin"
    else:
        from molmo_spaces.renderer.opengl_context import EGLGLContext

        self._gl_context = EGLGLContext(width, height, device_id)
    self._gl_context.make_current()
    self._mjr_context = mj.MjrContext(model, mj.mjtFontScale.mjFONTSCALE_150.value)
    mj.mjr_resizeOffscreen(width, height, self._mjr_context)
    mj.mjr_setBuffer(mj.mjtFramebuffer.mjFB_OFFSCREEN.value, self._mjr_context)
    self._mjr_context.readDepthMap = mj.mjtDepthMap.mjDEPTH_ZEROFAR.value

    # TODO In MacOS, keeping the context locked seems to preclude others to progress,
    #  so it doesn't look like we can achieve true parallelism through multi threading?
    #  This also happens at the end of render()
    if self._context_is_cgl:
        from mujoco.cgl import cgl  # ty: ignore[unresolved-import]

        cgl.CGLUnlockContext(self._gl_context._context)  # pyright: ignore[reportAttributeAccessIssue]

    self._textures_need_upload = False
device_id property 
device_id: int | None
height property 
height
model property 
model: MjModel
render_outputs instance-attribute 
render_outputs: list[Any]
scene property 
scene: MjvScene
width property 
width
__del__ 
__del__() -> None
Source code in molmo_spaces/renderer/opengl_rendering.py 
373
374
def __del__(self) -> None:
    self.close()
__enter__ 
__enter__()
Source code in molmo_spaces/renderer/opengl_rendering.py 
366
367
def __enter__(self):
    return self
__exit__ 
__exit__(exc_type, exc_value, traceback)
Source code in molmo_spaces/renderer/opengl_rendering.py 
369
370
371
def __exit__(self, exc_type, exc_value, traceback):
    del exc_type, exc_value, traceback
    self.close()
close 
close() -> None
Source code in molmo_spaces/renderer/opengl_rendering.py 
358
359
360
361
362
363
364
def close(self) -> None:
    if hasattr(self, "_gl_context") and self._gl_context:
        self._gl_context.free()
    self._gl_context = None
    if hasattr(self, "_mjr_context") and self._mjr_context:
        self._mjr_context.free()
    self._mjr_context = None
disable_depth_rendering 
disable_depth_rendering() -> None
Source code in molmo_spaces/renderer/opengl_rendering.py 
85
86
def disable_depth_rendering(self) -> None:
    self._depth_rendering = False
disable_segmentation_rendering 
disable_segmentation_rendering() -> None
Source code in molmo_spaces/renderer/opengl_rendering.py 
92
93
def disable_segmentation_rendering(self) -> None:
    self._segmentation_rendering = False
enable_depth_rendering 
enable_depth_rendering() -> None
Source code in molmo_spaces/renderer/opengl_rendering.py 
81
82
83
def enable_depth_rendering(self) -> None:
    self._segmentation_rendering = False
    self._depth_rendering = True
enable_segmentation_rendering 
enable_segmentation_rendering() -> None
Source code in molmo_spaces/renderer/opengl_rendering.py 
88
89
90
def enable_segmentation_rendering(self) -> None:
    self._segmentation_rendering = True
    self._depth_rendering = False
geomid_to_bodyid 
geomid_to_bodyid(geomid)
Source code in molmo_spaces/renderer/opengl_rendering.py 
95
96
def geomid_to_bodyid(self, geomid):
    return self.model.geom_bodyid[geomid]
mark_textures_dirty 
mark_textures_dirty() -> None

Mark that textures have been modified and need to be uploaded.

Call this after modifying texture data in model.tex_data to ensure the changes will be uploaded before the next render.

Source code in molmo_spaces/renderer/opengl_rendering.py 
299
300
301
302
303
304
305
def mark_textures_dirty(self) -> None:
    """Mark that textures have been modified and need to be uploaded.

    Call this after modifying texture data in model.tex_data to ensure
    the changes will be uploaded before the next render.
    """
    self._textures_need_upload = True
render 
render(*, out: ndarray | None = None, width: int | None = None, height: int | None = None) -> ndarray

Renders the scene as a numpy array of pixel values.

Parameters:

Name Type Description Default
out ndarray | None

Alternative output array in which to place the resulting pixels. It must have the same shape as the expected output but the type will be cast if necessary. The expted shape depends on the value of self._depth_rendering: when True, we expect out.shape == (width, height), and out.shape == (width, height, 3) when False.

 None

Returns:

Type Description
ndarray

A new numpy array holding the pixels with shape (H, W) or (H, W, 3),
ndarray

depending on the value of self._depth_rendering unless
ndarray

out is None, in which case a reference to out is returned.

Raises:

Type Description
RuntimeError

if this method is called after the close method.
Source code in molmo_spaces/renderer/opengl_rendering.py 
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
def render(
    self,
    *,
    out: np.ndarray | None = None,
    width: int | None = None,
    height: int | None = None,
) -> np.ndarray:
    """Renders the scene as a numpy array of pixel values.

    Args:
      out: Alternative output array in which to place the resulting pixels. It
        must have the same shape as the expected output but the type will be
        cast if necessary. The expted shape depends on the value of
        `self._depth_rendering`: when `True`, we expect `out.shape == (width,
        height)`, and `out.shape == (width, height, 3)` when `False`.

    Returns:
      A new numpy array holding the pixels with shape `(H, W)` or `(H, W, 3)`,
      depending on the value of `self._depth_rendering` unless
      `out is None`, in which case a reference to `out` is returned.

    Raises:
      RuntimeError: if this method is called after the close method.
    """
    assert self._scene is not None, "Internal scene:MjvScene must be initialized by now"

    height = height or self._height
    width = width or self._width
    rect = mj.MjrRect(0, 0, width, height)

    original_flags = self._scene.flags.copy()

    # Enable shadow rendering (required for shadows to appear in rendered images)
    # Shadows are controlled by lights with castshadow enabled
    self._scene.flags[mj.mjtRndFlag.mjRND_SHADOW] = True

    # Using segmented rendering for depth makes the calculated depth more
    # accurate at far distances.
    if self._depth_rendering or self._segmentation_rendering:
        self._scene.flags[mj.mjtRndFlag.mjRND_SEGMENT] = True
        self._scene.flags[mj.mjtRndFlag.mjRND_IDCOLOR] = True

    if self._gl_context is None:
        raise RuntimeError("render cannot be called after close.")

    self._gl_context.make_current()

    # Upload textures to GPU before rendering if textures have been modified
    # This is necessary when textures are modified via model.tex_data
    # Only upload when needed to avoid performance overhead
    if self._textures_need_upload:
        self.upload_textures()
        self._textures_need_upload = False

    if self._depth_rendering:
        out_shape = (rect.height, rect.width)
        out_dtype = np.float32
    else:
        out_shape = (rect.height, rect.width, 3)
        out_dtype = np.uint8

    if out is None:
        out = np.empty(out_shape, dtype=out_dtype)
    else:
        if out.shape != out_shape:
            raise ValueError(
                f"Expected `out.shape == {out_shape}`. Got `out.shape={out.shape}`"
                " instead. When using depth rendering, the out array should be of"
                " shape `(width, height)` and otherwise (width, height, 3)."
                f" Got `(self.height, self.width)={(self.height, self.width)}` and"
                f" `self._depth_rendering={self._depth_rendering}`."
            )

    assert self._mjr_context, "MjrContext must be created by now, but it's None"
    mj.mjr_render(rect, self._scene, self._mjr_context)

    if self._depth_rendering:
        mj.mjr_readPixels(rgb=None, depth=out, viewport=rect, con=self._mjr_context)

        # Get the distances to the near and far clipping planes.
        extent = self.model.stat.extent
        near = self.model.vis.map.znear * extent
        far = self.model.vis.map.zfar * extent

        # Calculate OpenGL perspective matrix values in float32 precision
        # so they are close to what glFrustum returns
        # https://registry.khronos.org/OpenGL-Refpages/gl2.1/xhtml/glFrustum.xml
        zfar = np.float32(far)
        znear = np.float32(near)
        c_coef = -(zfar + znear) / (zfar - znear)
        d_coef = -(np.float32(2) * zfar * znear) / (zfar - znear)

        # In reverse Z mode the perspective matrix is transformed by the following
        c_coef = np.float32(-0.5) * c_coef - np.float32(0.5)
        d_coef = np.float32(-0.5) * d_coef

        # We need 64 bits to convert Z from ndc to metric depth without noticeable
        # losses in precision
        out_64 = out.astype(np.float64)

        # Undo OpenGL projection
        # Note: We do not need to take action to convert from window coordinates
        # to normalized device coordinates because in reversed Z mode the mapping
        # is identity
        out_64 = d_coef / (out_64 + c_coef)

        # Cast result back to float32 for backwards compatibility
        # This has a small accuracy cost
        out[:] = out_64.astype(np.float32)

        # Reset scene flags.
        np.copyto(self._scene.flags, original_flags)
    elif self._segmentation_rendering:
        mj.mjr_readPixels(rgb=out, depth=None, viewport=rect, con=self._mjr_context)

        # Convert 3-channel uint8 to 1-channel uint32.
        image3 = out.astype(np.uint32)
        segimage = image3[:, :, 0] + image3[:, :, 1] * (2**8) + image3[:, :, 2] * (2**16)
        # Remap segid to 3-channel (object ID, object type, body ID) triplet
        # Seg ID 0 is background -- will be remapped to (-1, -1, -1).

        # Find the maximum segment ID in the image to size the output array correctly
        max_segid = np.max(segimage) if segimage.size > 0 else 0

        # Create output array with size to accommodate all possible segment IDs
        # Add 1 to account for 0-based indexing and ensure we have enough space
        segid2output = np.full((max_segid + 1, 3), fill_value=-1, dtype=np.int32)

        visible_geoms = [g for g in self._scene.geoms[: self._scene.ngeom] if g.segid != -1]
        visible_segids = np.array([g.segid + 1 for g in visible_geoms], np.int32)
        visible_objid = np.array([g.objid for g in visible_geoms], np.int32)
        visible_objtype = np.array([g.objtype for g in visible_geoms], np.int32)
        visible_bodyid = np.array(
            [self.geomid_to_bodyid(g.objid) for g in visible_geoms], np.int32
        )

        # Only set values for valid segment IDs that are within bounds
        valid_mask = (visible_segids >= 0) & (visible_segids < segid2output.shape[0])
        if np.any(valid_mask):
            segid2output[visible_segids[valid_mask], 0] = visible_objid[valid_mask]
            segid2output[visible_segids[valid_mask], 1] = visible_objtype[valid_mask]
            segid2output[visible_segids[valid_mask], 2] = visible_bodyid[valid_mask]

        out = segid2output[segimage]

        # Reset scene flags.
        np.copyto(self._scene.flags, original_flags)
    else:
        mj.mjr_readPixels(rgb=out, depth=None, viewport=rect, con=self._mjr_context)
        mj.mjr_readPixels(rgb=out, depth=None, viewport=rect, con=self._mjr_context)

    out[:] = np.flipud(out)

    # TODO In MacOS, keeping the context locked seems to preclude others to progress,
    #  so it doesn't look like we can achieve true parallelism through multi threading?
    #  This also happens at the end of __init__()
    if self._context_is_cgl:
        from mujoco.cgl import cgl  # ty: ignore[unresolved-import]

        cgl.CGLUnlockContext(self._gl_context._context)  # pyright: ignore[reportAttributeAccessIssue] # ty: ignore

    return out
update 
update(data: MjData, camera: int | str | MjvCamera = -1, scene_option: MjvOption | None = None) -> None

Updates geometry used for rendering.

Parameters:

Name Type Description Default
data MjData

An instance of MjData.

 required
camera int | str | MjvCamera

An instance of MjvCamera, a string or an integer

 -1
scene_option MjvOption | None

A custom MjvOption instance to use to render the scene instead of the default.

 None

Raises:

Type Description
ValueError

If camera_id is outside the valid range, or if camera does not exist.
Source code in molmo_spaces/renderer/opengl_rendering.py 
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
def update(
    self,
    data: mj.MjData,
    camera: int | str | mj.MjvCamera = -1,
    scene_option: mj.MjvOption | None = None,
) -> None:
    """Updates geometry used for rendering.

    Args:
      data: An instance of `MjData`.
      camera: An instance of `MjvCamera`, a string or an integer
      scene_option: A custom `MjvOption` instance to use to render
        the scene instead of the default.

    Raises:
      ValueError: If `camera_id` is outside the valid range, or if camera does
        not exist.
    """
    assert self._scene is not None, "Internal scene:MjvScene must be initialized by now"

    if not isinstance(camera, mj.MjvCamera):
        camera_id = camera
        if isinstance(camera_id, str):
            camera_id = mj.mj_name2id(self.model, mj.mjtObj.mjOBJ_CAMERA.value, camera_id)
            if camera_id == -1:
                raise ValueError(f'The camera "{camera}" does not exist.')
        if camera_id < -1 or camera_id >= self.model.ncam:
            raise ValueError(
                f"The camera id {camera_id} is out of range [-1, {self.model.ncam})."
            )

        camera = mj.MjvCamera()
        camera.fixedcamid = camera_id

        if camera_id == -1:
            camera.type = mj.mjtCamera.mjCAMERA_FREE
            mj.mjv_defaultFreeCamera(self.model, camera)
        else:
            camera.type = mj.mjtCamera.mjCAMERA_FIXED

    scene_option = scene_option or self._scene_option
    mj.mjv_updateScene(
        self.model,
        data,
        scene_option,
        None,
        camera,
        mj.mjtCatBit.mjCAT_ALL.value,
        self._scene,
    )
upload_textures 
upload_textures() -> None

Upload all textures to the GPU render context.

This should be called after modifying texture data in model.tex_data to ensure the changes are visible in rendered images.

NOTE: This only uploads textures to THIS renderer's context (MjOpenGLRenderer). The passive viewer has its own separate renderer context and won't see these updates.

Parameters:

Name Type Description Default
data

Optional MjData to use for updating the scene after texture upload. If provided, will call mjv_updateScene() to refresh the scene.

 required
Source code in molmo_spaces/renderer/opengl_rendering.py 
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
def upload_textures(self) -> None:
    """Upload all textures to the GPU render context.

    This should be called after modifying texture data in model.tex_data
    to ensure the changes are visible in rendered images.

    NOTE: This only uploads textures to THIS renderer's context (MjOpenGLRenderer).
    The passive viewer has its own separate renderer context and won't see these updates.

    Args:
        data: Optional MjData to use for updating the scene after texture upload.
              If provided, will call mjv_updateScene() to refresh the scene.
    """
    import logging

    log = logging.getLogger(__name__)

    if self._gl_context is None or self._mjr_context is None:
        log.debug("upload_textures(): Skipping - GL context or Mjr context is None")
        return

    # Skip if no textures exist
    if self.model.ntex == 0:
        log.debug("upload_textures(): Skipping - no textures in model (ntex == 0)")
        return

    log.debug(f"upload_textures(): Uploading {self.model.ntex} textures to GPU render context")
    self._gl_context.make_current()
    # Upload all textures to the render context
    for tex_id in range(self.model.ntex):
        mj.mjr_uploadTexture(self.model, self._mjr_context, tex_id)

    # Unlock context if needed (for macOS)
    if self._context_is_cgl:
        from mujoco.cgl import cgl  # ty: ignore[unresolved-import]

        cgl.CGLUnlockContext(self._gl_context._context)  # pyright: ignore[reportAttributeAccessIssue] # ty: ignore