Separate Light management into a separate class.

PiperOrigin-RevId: 918508291
Change-Id: I9cadc1b612006a0450f8a2fd9088d56cc2e610a3

src/experimental/filament/CMakeLists.txt

@@ -51,6 +51,8 @@ target_sources(${MUJOCO_FILAMENT_TARGET_NAME}
         filament/scene_view.h
         filament/texture.cc
         filament/texture.h
+        compat/light_manager.cc
+        compat/light_manager.h
         compat/model_objects.cc
         compat/model_objects.h
         compat/scene_bridge.cc

src/experimental/filament/compat/light_manager.cc

@@ -0,0 +1,213 @@
+// Copyright 2025 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "experimental/filament/compat/light_manager.h"
+
+#include <memory>
+#include <string>
+#include <utility>
+
+#include <math/mat4.h>
+#include <math/mathfwd.h>
+#include <math/vec3.h>
+#include <math/vec4.h>
+#include <mujoco/mujoco.h>
+#include "experimental/filament/compat/model_objects.h"
+#include "experimental/filament/filament_util.h"
+#include "experimental/filament/render_context_filament_cpp.h"
+#include "experimental/filament/render_context_filament.h"
+
+namespace mujoco {
+
+using filament::math::float3;
+using filament::math::float4;
+using filament::math::mat3;
+using filament::math::mat4;
+
+static UniquePtr<mjrTexture> CreateFallbackIndirectLightTexture(
+    mjrfContext* ctx) {
+  const std::string filename = ResolveFilamentAssetPath("ibl.ktx");
+  mjResource* resource =
+      mju_openResource("", filename.c_str(), nullptr, nullptr, 0);
+  if (!resource) {
+    mju_error("Failed to open resource: %s", filename.c_str());
+  }
+  const void* bytes = nullptr;
+  const int nbytes = mju_readResource(resource, &bytes);
+  if (bytes == nullptr || nbytes <= 0) {
+    mju_error("Failed to read resource: %s", filename.c_str());
+  }
+
+  mjrTextureConfig config;
+  mjr_defaultTextureConfig(&config);
+  config.width = 1;
+  config.height = 1;
+  config.sampler_type = mjTEXTURE_CUBE;
+  config.format = mjPIXEL_FORMAT_KTX;
+  config.color_space = mjCOLORSPACE_AUTO;
+
+  auto texture = CreateTexture(ctx, config);
+
+  mjrTextureData payload;
+  mjr_defaultTextureData(&payload);
+  payload.bytes = bytes;
+  payload.nbytes = nbytes;
+  payload.release_callback = +[](void* user_data) {
+    mju_closeResource((mjResource*)user_data);
+  };
+  payload.user_data = resource;
+
+  mjrf_setTextureData(texture.get(), &payload);
+  return texture;
+}
+
+LightManager::LightManager(mjrfContext* ctx, mjrScene* scene,
+                           ModelObjects* model_objects)
+    : ctx_(ctx), scene_(scene) {
+  const mjModel* model = model_objects->GetModel();
+  default_shadow_map_size_ = ReadElement(
+      model, "filament.shadows.map_size", default_shadow_map_size_);
+  default_vsm_blur_width_ = ReadElement(
+      model, "filament.shadows.vsm_blur_width", default_vsm_blur_width_);
+  fallback_head_light_intensity_ =
+      ReadElement(model, "filament.fallback.head_light_intensity",
+                  fallback_head_light_intensity_);
+  fallback_scene_light_intensity_ =
+      ReadElement(model, "filament.fallback.scene_light_intensity",
+                  fallback_scene_light_intensity_);
+  fallback_environment_light_intensity_ =
+      ReadElement(model, "filament.fallback.environment_light_intensity",
+                  fallback_environment_light_intensity_);
+  Prepare(model_objects);
+}
+
+LightManager::~LightManager() {
+  for (auto& iter : lights_) {
+    mjrf_removeLightFromScene(scene_, iter.get());
+  }
+  lights_.clear();
+  if (fallback_ibl_) {
+    mjrf_removeLightFromScene(scene_, fallback_ibl_.get());
+  }
+  fallback_ibl_.reset();
+}
+
+void LightManager::Prepare(ModelObjects* model_objects) {
+  const mjModel* model = model_objects->GetModel();
+
+  bool has_image_based_light = false;
+  float total_light_intensity = 0.0f;
+  for (int i = 0; i < model->nlight; ++i) {
+    total_light_intensity += model->light_intensity[i];
+
+    if (model->light_type[i] == mjLIGHT_IMAGE) {
+      mjrLightParams params;
+      mjr_defaultLightParams(&params);
+      params.type = mjLIGHT_IMAGE;
+      params.texture = model_objects->GetTexture(model->light_texid[i]);
+      params.intensity = model->light_intensity[i];
+      auto light_obj = CreateLight(ctx_, params);
+      mjrf_addLightToScene(scene_, light_obj.get());
+      lights_.emplace_back(std::move(light_obj));
+      has_image_based_light = true;
+    } else {
+      mjrLightParams params;
+      mjr_defaultLightParams(&params);
+      params.color[0] = model->light_diffuse[0];
+      params.color[1] = model->light_diffuse[1];
+      params.color[2] = model->light_diffuse[2];
+      params.type = (mjtLightType)model->light_type[i];
+      params.cast_shadows = model->light_castshadow[i];
+      params.bulb_radius = model->light_bulbradius[i];
+      params.range = model->light_range[i];
+      params.intensity = model->light_intensity[i];
+      params.shadow_map_size = default_shadow_map_size_;
+      params.vsm_blur_width = default_vsm_blur_width_;
+      if (params.type == mjLIGHT_SPOT) {
+        params.spot_cone_angle = model->light_cutoff[i];
+      }
+
+      auto light_obj = CreateLight(ctx_, params);
+      mjrf_addLightToScene(scene_, light_obj.get());
+      lights_.emplace_back(std::move(light_obj));
+    }
+  }
+
+  // Add a placeholder (black) headlight as our last light. Going forward, we'll
+  // assume lights_.back() is always the headlight.
+  {
+    mjrLightParams params;
+    mjr_defaultLightParams(&params);
+    // We break with the spec here slightly and use a spot light for the head
+    // light instead of a directional params. This is because filament only
+    // supports a single directional light, and we'd rather allow a scene
+    // light to be that directional params. It's also a bit odd for a
+    // directional light to move with the camera.
+    params.type = mjLIGHT_SPOT;
+    params.cast_shadows = 0;
+    params.intensity = 0.0f;
+    params.spot_cone_angle = 90.0f;
+    auto light_obj = CreateLight(ctx_, params);
+    mjrf_addLightToScene(scene_, light_obj.get());
+    lights_.emplace_back(std::move(light_obj));
+  }
+
+  if (!has_image_based_light && total_light_intensity > 0.0f) {
+    // Create a black indirect light to ensure that the skybox is
+    // oriented to respect mujoco's Z-up convention.
+    mjrLightParams params;
+    mjr_defaultLightParams(&params);
+    params.type = mjLIGHT_IMAGE;
+    params.intensity = 10.0f;
+    fallback_ibl_ = CreateLight(ctx_, params);
+    mjrf_addLightToScene(scene_, fallback_ibl_.get());
+  }
+
+  // There are no "physical" lights in the scene which means we're likely
+  // dealing with a "classic renderer" scene. In this case, let's add a
+  // default environment light and set the light intensity ourselves.
+  if (total_light_intensity == 0.0f) {
+    // Create a fallback environment light.
+    fallback_ibl_texture_ = CreateFallbackIndirectLightTexture(ctx_);
+
+    mjrLightParams params;
+    mjr_defaultLightParams(&params);
+    params.type = mjLIGHT_IMAGE;
+    params.texture = fallback_ibl_texture_.get();
+    params.intensity = fallback_environment_light_intensity_;
+    fallback_ibl_ = CreateLight(ctx_, params);
+    mjrf_addLightToScene(scene_, fallback_ibl_.get());
+
+    // Distribute the fallback scene light intensity among the lights.
+    const float intensity = fallback_scene_light_intensity_ / lights_.size();
+    for (auto& light : lights_) {
+      if (light) {
+        const bool is_headlight = (light == lights_.back());
+        mjrf_setLightIntensity(light.get(),
+                               is_headlight ? fallback_head_light_intensity_
+                                            : intensity);
+      }
+    }
+  }
+
+  mjrf_setSceneSkybox(scene_, model_objects->GetSkyboxTexture());
+}
+
+mjrLight* LightManager::GetLight(int index) {
+  if (index < 0 || index >= lights_.size()) {
+    return nullptr;
+  }
+  return lights_[index].get();
+}
+}  // namespace mujoco

src/experimental/filament/compat/light_manager.h

@@ -0,0 +1,56 @@
+// Copyright 2025 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef MUJOCO_SRC_EXPERIMENTAL_FILAMENT_COMPAT_LIGHT_MANAGER_H_
+#define MUJOCO_SRC_EXPERIMENTAL_FILAMENT_COMPAT_LIGHT_MANAGER_H_
+
+#include <vector>
+
+#include "experimental/filament/compat/model_objects.h"
+#include "experimental/filament/render_context_filament.h"
+#include "experimental/filament/render_context_filament_cpp.h"
+
+namespace mujoco {
+
+// Manages Light entities for an mjrScene using data from an mjvScene.
+class LightManager {
+ public:
+  LightManager(mjrfContext* ctx, mjrScene* scene, ModelObjects* model_objects);
+  ~LightManager();
+
+  // Returns the light with the given index in the mjModel. Note that an extra
+  // headlight is assigned of the index `nlight`.
+  mjrLight* GetLight(int index);
+
+  LightManager(const LightManager&) = delete;
+  LightManager& operator=(const LightManager&) = delete;
+
+ private:
+  void Prepare(ModelObjects* model_objects);
+
+  mjrfContext* ctx_ = nullptr;
+  mjrScene* scene_ = nullptr;
+  UniquePtr<mjrLight> fallback_ibl_{nullptr, nullptr};
+  UniquePtr<mjrTexture> fallback_ibl_texture_{nullptr, nullptr};
+  std::vector<UniquePtr<mjrLight>> lights_;
+  int default_shadow_map_size_ = 2048;
+  float default_vsm_blur_width_ = 0.0f;
+  float fallback_head_light_intensity_ = 0.f;
+  float fallback_scene_light_intensity_ = 80'000.f;
+  float fallback_environment_light_intensity_ = 5'000.f;
+};
+
+}  // namespace mujoco
+
+#endif  // MUJOCO_SRC_EXPERIMENTAL_FILAMENT_COMPAT_LIGHT_MANAGER_H_

src/experimental/filament/compat/scene_bridge.cc

@@ -16,8 +16,6 @@
 
 #include <memory>
 #include <optional>
-#include <string>
-#include <string_view>
 #include <utility>
 
 #include <math/TMatHelpers.h>
@@ -27,6 +25,7 @@
 #include <math/vec3.h>
 #include <math/vec4.h>
 #include <mujoco/mujoco.h>
+#include "experimental/filament/compat/light_manager.h"
 #include "experimental/filament/compat/model_objects.h"
 #include "experimental/filament/compat/scene_geom_util.h"
 #include "experimental/filament/filament_util.h"
@@ -40,43 +39,6 @@ using filament::math::float4;
 using filament::math::mat3;
 using filament::math::mat4;
 
-static UniquePtr<mjrTexture> CreateFallbackIndirectLightTexture(
-    mjrfContext* ctx) {
-  const std::string filename = ResolveFilamentAssetPath("ibl.ktx");
-  mjResource* resource =
-      mju_openResource("", filename.c_str(), nullptr, nullptr, 0);
-  if (!resource) {
-    mju_error("Failed to open resource: %s", filename.c_str());
-  }
-  const void* bytes = nullptr;
-  const int nbytes = mju_readResource(resource, &bytes);
-  if (bytes == nullptr || nbytes <= 0) {
-    mju_error("Failed to read resource: %s", filename.c_str());
-  }
-
-  mjrTextureConfig config;
-  mjr_defaultTextureConfig(&config);
-  config.width = 1;
-  config.height = 1;
-  config.sampler_type = mjTEXTURE_CUBE;
-  config.format = mjPIXEL_FORMAT_KTX;
-  config.color_space = mjCOLORSPACE_AUTO;
-
-  auto texture = CreateTexture(ctx, config);
-
-  mjrTextureData payload;
-  mjr_defaultTextureData(&payload);
-  payload.bytes = bytes;
-  payload.nbytes = nbytes;
-  payload.release_callback = +[](void* user_data) {
-    mju_closeResource((mjResource*)user_data);
-  };
-  payload.user_data = resource;
-
-  mjrf_setTextureData(texture.get(), &payload);
-  return texture;
-}
-
 SceneBridge::SceneBridge(mjrfContext* ctx, const mjModel* model)
     : ctx_(ctx) {
   mjrSceneParams params;
@@ -92,32 +54,12 @@ SceneBridge::SceneBridge(mjrfContext* ctx, const mjModel* model)
                                  filament::math::float4(0, 0, 0, 1));
   mjrf_setClearColor(ctx_, &clear_color[0]);
 
-  default_shadow_map_size_ = ReadElement(
-      model, "filament.shadows.map_size", default_shadow_map_size_);
-  default_vsm_blur_width_ = ReadElement(
-      model, "filament.shadows.vsm_blur_width", default_vsm_blur_width_);
-  fallback_head_light_intensity_ =
-      ReadElement(model, "filament.fallback.head_light_intensity",
-                  fallback_head_light_intensity_);
-  fallback_scene_light_intensity_ =
-      ReadElement(model, "filament.fallback.scene_light_intensity",
-                  fallback_scene_light_intensity_);
-  fallback_environment_light_intensity_ =
-      ReadElement(model, "filament.fallback.environment_light_intensity",
-                  fallback_environment_light_intensity_);
-
-  PrepareLights();
+  light_manager_ =
+      std::make_unique<LightManager>(ctx_, scene_.get(), model_objects_.get());
 }
 
 SceneBridge::~SceneBridge() {
-  for (auto& iter : lights_) {
-    mjrf_removeLightFromScene(scene_.get(), iter.get());
-  }
-  lights_.clear();
-  if (fallback_ibl_) {
-    mjrf_removeLightFromScene(scene_.get(), fallback_ibl_.get());
-  }
-  fallback_ibl_.reset();
+  light_manager_.reset();
   for (auto& iter : renderables_) {
     mjrf_removeRenderableFromScene(scene_.get(), iter.get());
   }
@@ -132,107 +74,6 @@ std::optional<float3> SceneBridge::ClipFromWorld(const float3& pos) const{
   return clip_pos.xyz / clip_pos.w;
 }
 
-void SceneBridge::PrepareLights() {
-  const mjModel* model = model_objects_->GetModel();
-
-  bool has_image_based_light = false;
-  float total_light_intensity = 0.0f;
-  for (int i = 0; i < model->nlight; ++i) {
-    total_light_intensity += model->light_intensity[i];
-
-    if (model->light_type[i] == mjLIGHT_IMAGE) {
-      mjrLightParams params;
-      mjr_defaultLightParams(&params);
-      params.type = mjLIGHT_IMAGE;
-      params.texture = model_objects_->GetTexture(model->light_texid[i]);
-      params.intensity = model->light_intensity[i];
-      auto light_obj = CreateLight(ctx_, params);
-      mjrf_addLightToScene(scene_.get(), light_obj.get());
-      lights_.emplace_back(std::move(light_obj));
-      has_image_based_light = true;
-    } else {
-      mjrLightParams params;
-      mjr_defaultLightParams(&params);
-      params.color[0] = model->light_diffuse[0];
-      params.color[1] = model->light_diffuse[1];
-      params.color[2] = model->light_diffuse[2];
-      params.type = (mjtLightType)model->light_type[i];
-      params.cast_shadows = model->light_castshadow[i];
-      params.bulb_radius = model->light_bulbradius[i];
-      params.range = model->light_range[i];
-      params.intensity = model->light_intensity[i];
-      params.shadow_map_size = default_shadow_map_size_;
-      params.vsm_blur_width = default_vsm_blur_width_;
-      if (params.type == mjLIGHT_SPOT) {
-        params.spot_cone_angle = model->light_cutoff[i];
-      }
-
-      auto light_obj = CreateLight(ctx_, params);
-      mjrf_addLightToScene(scene_.get(), light_obj.get());
-      lights_.emplace_back(std::move(light_obj));
-    }
-  }
-
-  // Add a placeholder (black) headlight as our last light. Going forward, we'll
-  // assume lights_.back() is always the headlight.
-  {
-    mjrLightParams params;
-    mjr_defaultLightParams(&params);
-    // We break with the spec here slightly and use a spot light for the head
-    // light instead of a directional params. This is because filament only
-    // supports a single directional light, and we'd rather allow a scene
-    // light to be that directional params. It's also a bit odd for a
-    // directional light to move with the camera.
-    params.type = mjLIGHT_SPOT;
-    params.cast_shadows = 0;
-    params.intensity = 0.0f;
-    params.spot_cone_angle = 90.0f;
-    auto light_obj = CreateLight(ctx_, params);
-    mjrf_addLightToScene(scene_.get(), light_obj.get());
-    lights_.emplace_back(std::move(light_obj));
-  }
-
-  if (!has_image_based_light && total_light_intensity > 0.0f) {
-    // Create a black indirect light to ensure that the skybox is
-    // oriented to respect mujoco's Z-up convention.
-    mjrLightParams params;
-    mjr_defaultLightParams(&params);
-    params.type = mjLIGHT_IMAGE;
-    params.intensity = 10.0f;
-    fallback_ibl_ = CreateLight(ctx_, params);
-    mjrf_addLightToScene(scene_.get(), fallback_ibl_.get());
-  }
-
-  // There are no "physical" lights in the scene which means we're likely
-  // dealing with a "classic renderer" scene. In this case, let's add a
-  // default environment light and set the light intensity ourselves.
-  if (total_light_intensity == 0.0f) {
-    // Create a fallback environment light.
-    fallback_ibl_texture_ = CreateFallbackIndirectLightTexture(ctx_);
-
-    mjrLightParams params;
-    mjr_defaultLightParams(&params);
-    params.type = mjLIGHT_IMAGE;
-    params.texture = fallback_ibl_texture_.get();
-    params.intensity = fallback_environment_light_intensity_;
-    fallback_ibl_ = CreateLight(ctx_, params);
-    mjrf_addLightToScene(scene_.get(), fallback_ibl_.get());
-
-    // Distribute the fallback scene light intensity among the lights.
-    const float intensity = fallback_scene_light_intensity_ / lights_.size();
-    for (auto& light : lights_) {
-      if (light) {
-        const bool is_headlight = (light == lights_.back());
-        mjrf_setLightIntensity(light.get(),
-                               is_headlight ? fallback_head_light_intensity_
-                                            : intensity);
-      }
-    }
-  }
-
-  mjrf_setSceneSkybox(scene_.get(), model_objects_->GetSkyboxTexture());
-}
-
 mat4 CalculateClipFromWorld(const mjrRect& viewport, const mjrCamera& cam) {
   const float3 cam_pos(cam.pos[0], cam.pos[1], cam.pos[2]);
   const float3 cam_fwd(cam.forward[0], cam.forward[1], cam.forward[2]);
@@ -295,10 +136,16 @@ void SceneBridge::Update(const mjrRect& viewport, const mjvScene* scene) {
     renderables_.push_back(std::move(renderable));
   }
 
+  mjrLight* headlight = nullptr;
   bool headlight_enabled = false;
   for (int i = 0; i < scene->nlight; ++i) {
     const mjvLight& scene_light = scene->lights[i];
     if (scene_light.id < 0 && scene_light.headlight) {
+      // The headlight, if it exists, is assigned the id `scene->nlight`.
+      headlight = light_manager_->GetLight(scene->nlight);
+      if (!headlight) {
+        continue;
+      }
       // We position the headlight slightly behind the camera to avoid some
       // odd clipping issues.
       headlight_enabled = true;
@@ -306,24 +153,21 @@ void SceneBridge::Update(const mjrRect& viewport, const mjvScene* scene) {
       headpos[1] -= gazedir[1] * 0.05f;
       headpos[2] -= gazedir[2] * 0.05f;
 
-      // The headlight is always the "back" light.
-      UniquePtr<mjrLight>& light = lights_.back();
-      mjrf_setLightColor(light.get(), scene_light.diffuse);
-      mjrf_setLightTransform(light.get(), headpos, gazedir);
+      mjrf_setLightColor(headlight, scene_light.diffuse);
+      mjrf_setLightTransform(headlight, headpos, gazedir);
       continue;
-    } else if (scene_light.id < lights_.size() - 1) {
-      UniquePtr<mjrLight>& light = lights_[scene_light.id];
-      if (light) {
-        mjrf_setLightColor(light.get(), scene_light.diffuse);
-        mjrf_setLightTransform(light.get(), scene_light.pos, scene_light.dir);
-      }
+    } else if (mjrLight* light = light_manager_->GetLight(scene_light.id)) {
+      mjrf_setLightColor(light, scene_light.diffuse);
+      mjrf_setLightTransform(light, scene_light.pos, scene_light.dir);
     } else {
       mju_error("Unexpected light id: %d", scene_light.id);
     }
   }
 
   // Enable/disable the headlight based on whether or not it's in the scene.
-  mjrf_setLightEnabled(lights_.back().get(), headlight_enabled);
+  if (headlight) {
+    mjrf_setLightEnabled(headlight, headlight_enabled);
+  }
 }
 
 void SceneBridge::UploadMesh(const mjModel* model, int id) {

src/experimental/filament/compat/scene_bridge.h

@@ -18,13 +18,13 @@
 #include <functional>
 #include <memory>
 #include <optional>
-#include <string_view>
 #include <vector>
 
 #include <math/mat4.h>
 #include <math/vec3.h>
 #include <mujoco/mjvisualize.h>
 #include <mujoco/mujoco.h>
+#include "experimental/filament/compat/light_manager.h"
 #include "experimental/filament/compat/model_objects.h"
 #include "experimental/filament/render_context_filament.h"
 #include "experimental/filament/render_context_filament_cpp.h"
@@ -57,8 +57,6 @@ class SceneBridge {
   SceneBridge& operator=(const SceneBridge&) = delete;
 
  private:
-  void PrepareLights();
-
   // Converts a point in world space to clip space, eg. in the range [-1,-1, 0]
   // to [1, 1, 1]. Returns std::nullopt if the point is behind the camera.
   std::optional<filament::math::float3> ClipFromWorld(
@@ -66,19 +64,12 @@ class SceneBridge {
 
   mjrfContext* ctx_ = nullptr;
   std::unique_ptr<ModelObjects> model_objects_;
+  std::unique_ptr<LightManager> light_manager_;
   mjrCamera camera_;
   DrawTextAtFn draw_text_callback_;
   UniquePtr<mjrScene> scene_{nullptr, nullptr};
-  UniquePtr<mjrLight> fallback_ibl_{nullptr, nullptr};
-  UniquePtr<mjrTexture> fallback_ibl_texture_{nullptr, nullptr};
-  std::vector<UniquePtr<mjrLight>> lights_;
   std::vector<UniquePtr<mjrRenderable>> renderables_;
   filament::math::mat4 clip_from_world_;
-  int default_shadow_map_size_ = 2048;
-  float default_vsm_blur_width_ = 0.0f;
-  float fallback_head_light_intensity_ = 0.f;
-  float fallback_scene_light_intensity_ = 80'000.f;
-  float fallback_environment_light_intensity_ = 5'000.f;
 };
 
 }  // namespace mujoco