fix(mjlab): settle CGL draws after scene changes

Author: Trinkle23897

docs/env/mjlab.rst

@@ -262,6 +262,13 @@ truncation, physics state, and public rendering at multiple steps.
 The oracle uses the upstream ``auto_reset=False`` option to retain terminal
 observations; EnvPool resets a completed slot on its next step.
 
+On macOS, native and oracle rendering settle each new frame with a second GPU
+draw. A captured scene with identical camera, geometry, and lighting bytes
+reproduces a one-level color difference on CGL/Metal's first draw; repeated
+readback alone does not remove it. This does not advance physics or change the
+scene, and pixel comparisons remain exact. Other families and platforms keep
+their existing rendering path.
+
 Independent tests observe native resets without oracle synchronization. They
 check different seeds, consecutive resets, parallel slots, and replay of the
 whole reset sequence. Independently randomized goals, poses, model properties,

envpool/mujoco/mjlab/mjlab_env.h

@@ -146,7 +146,7 @@ class MjlabEnv : public Env<MjlabEnvSpec>, public RenderableEnv {
     }
     renderer_->Render(simulation_.physics.Model(),
                       simulation_.physics.RenderData(), width, height, camera,
-                      output, &camera_, &option_, true);
+                      output, &camera_, &option_, true, true);
   }
 
 #ifdef ENVPOOL_TEST

envpool/mujoco/mjlab/mjlab_test.py

@@ -261,7 +261,7 @@ def test_rollout_and_render(self, task: str) -> None:
                     self.assertTrue(
                         left.observation_space[name].contains(lo[name][0]), name
                     )
-                if step in (0, 32, 97, 194):
+                if step in (0, 1, 32, 97, 194):
                     a, b = (
                         left.render(env_ids=[1, 0]),
                         right.render(env_ids=[1, 0]),
@@ -275,6 +275,9 @@ def test_rollout_and_render(self, task: str) -> None:
                     np.testing.assert_array_equal(
                         a[:1], left.render(env_ids=[1])
                     )
+                    np.testing.assert_array_equal(
+                        a, left.render(env_ids=[1, 0])
+                    )
                     self.assertGreater(int(a.max()) - int(a.min()), 20)
                 control = np.repeat(control[None], 2, axis=0)
                 lo, lr, lt, lx, li = left.step(control, env_id=li["env_id"])

envpool/mujoco/mjlab/oracle_probe.py

@@ -251,9 +251,14 @@ def record(
                     # geometry without advancing or synchronizing physics.
                     renderer.update(env.sim.data)
                     frame = renderer.render()
-                    if step == 0 and platform.system() == "Darwin":
-                        for _ in range(4):
-                            renderer.update(env.sim.data)
+                    if platform.system() == "Darwin":
+                        # Match the native CGL warmup and per-frame draw
+                        # settling, without advancing or synchronizing physics.
+                        if step == 0:
+                            for _ in range(4):
+                                renderer.update(env.sim.data)
+                                frame = renderer.render()
+                        else:
                             frame = renderer.render()
                     frames.append(frame.copy())
                     frame_steps.append(step)

envpool/mujoco/offscreen_renderer.cc

@@ -817,7 +817,8 @@ void OffscreenRenderer::Render(const mjModel* model, mjData* data, int width,
                                int height, int camera_id, unsigned char* rgb,
                                const mjvCamera* camera_override,
                                const mjvOption* option_override,
-                               bool update_camera_first) {
+                               bool update_camera_first,
+                               bool settle_cgl_frame) {
   if (!initialized_) {
     Initialize(model);
   }
@@ -873,9 +874,18 @@ void OffscreenRenderer::Render(const mjModel* model, mjData* data, int width,
     }
     cgl_first_frame_settled_ = true;
   } else {
+    if (settle_cgl_frame) {
+      // CGL/Metal can leave the first draw of a new scene unsettled, even
+      // after context warmup. Replaying identical camera/geom/light bytes
+      // reproduces one-quantum color changes; another readback does not help.
+      // Only callers requiring per-frame settling pay for this second draw.
+      mjr_finish();
+      mjr_render(viewport, &scene_, &context_);
+    }
     read_pixels();
   }
 #else
+  (void)settle_cgl_frame;
   read_pixels();
 #endif
 

envpool/mujoco/offscreen_renderer.h

@@ -55,7 +55,7 @@ class OffscreenRenderer {
               int camera_id, unsigned char* rgb,
               const mjvCamera* camera_override = nullptr,
               const mjvOption* option_override = nullptr,
-              bool update_camera_first = false);
+              bool update_camera_first = false, bool settle_cgl_frame = false);
 
  private:
   void Initialize(const mjModel* model);