Manage reflections at the scene level.

Renderables may be created/destroyed frequently causing issues with
the underlying RenderTargets used for reflections. Instead, use a
ReflectionManager (associated with a Scene) to manage the RenderTargets.

PiperOrigin-RevId: 926072766
Change-Id: I56c70000183417a92ab9297bc6fec5779f484cbd

src/experimental/filament/filament/reflection_manager.cc

@@ -0,0 +1,61 @@
+// Copyright 2026 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/filament/reflection_manager.h"
+#include <memory>
+
+#include <filament/Engine.h>
+#include "experimental/filament/filament/render_target.h"
+#include "experimental/filament/render_context_filament.h"
+
+namespace mujoco {
+
+ReflectionManager::ReflectionManager(filament::Engine* engine)
+    : engine_(engine) {}
+
+ReflectionManager::~ReflectionManager() {
+}
+
+void ReflectionManager::ClearRenderables() {
+  renderables_.clear();
+}
+
+mjrTexture* ReflectionManager::Register(mjrRenderable* renderable, int width,
+                                        int height) {
+  if (targets_.size() == renderables_.size()) {
+    mjrRenderTargetConfig config;
+    mjr_defaultRenderTargetConfig(&config);
+    config.color_format = mjPIXEL_FORMAT_RGBA8;
+    config.depth_format = mjPIXEL_FORMAT_DEPTH32F;
+    targets_.push_back(std::make_unique<RenderTarget>(engine_, config));
+  }
+  RenderTarget* target = targets_[renderables_.size()].get();
+  target->Prepare(width, height);
+  renderables_.push_back(renderable);
+  return target->GetColorTexture();
+}
+
+int ReflectionManager::GetNumRenderables() const {
+  return renderables_.size();
+}
+
+mjrRenderable* ReflectionManager::GetRenderable(int index) const {
+  return renderables_[index];
+}
+
+const RenderTarget* ReflectionManager::GetRenderTarget(int index) const {
+  return targets_[index].get();
+}
+
+}  // namespace mujoco

src/experimental/filament/filament/reflection_manager.h

@@ -0,0 +1,61 @@
+// Copyright 2026 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_FILAMENT_REFLECTION_MANAGER_H_
+#define MUJOCO_SRC_EXPERIMENTAL_FILAMENT_FILAMENT_REFLECTION_MANAGER_H_
+
+#include <memory>
+#include <vector>
+
+#include <filament/Engine.h>
+#include "experimental/filament/filament/render_target.h"
+#include "experimental/filament/render_context_filament.h"
+
+namespace mujoco {
+
+// Manages the allocation of RenderTargets for reflective surfaces.
+class ReflectionManager {
+ public:
+  ReflectionManager(filament::Engine* engine);
+  ~ReflectionManager();
+
+  ReflectionManager(const ReflectionManager&) = delete;
+  ReflectionManager& operator=(const ReflectionManager&) = delete;
+
+  // Registers a Renderable as being reflective. Internally, this function will
+  // create a RenderTarget of the given size and return the texture that the
+  // Renderable can use as its reflection.
+  mjrTexture* Register(mjrRenderable* renderable, int width, int height);
+
+  // Clears all previously registered renderables. This should be called at the
+  // beginning of a frame.
+  void ClearRenderables();
+
+  // Returns the number of registered renderables
+  int GetNumRenderables() const;
+
+  // Returns the renderable at the given index.
+  mjrRenderable* GetRenderable(int index) const;
+
+  // Returns the RenderTarget at the given index.
+  const RenderTarget* GetRenderTarget(int index) const;
+
+ private:
+  filament::Engine* engine_;
+  std::vector<mjrRenderable*> renderables_;
+  std::vector<std::unique_ptr<RenderTarget>> targets_;
+};
+}  // namespace mujoco
+
+#endif  // MUJOCO_SRC_EXPERIMENTAL_FILAMENT_FILAMENT_REFLECTION_MANAGER_H_

src/experimental/filament/filament/renderable.cc

@@ -16,7 +16,6 @@
 
 #include <cstdint>
 #include <functional>
-#include <memory>
 #include <numbers>
 #include <span>
 
@@ -39,7 +38,7 @@
 #include "experimental/filament/filament/material.h"
 #include "experimental/filament/filament/mesh.h"
 #include "experimental/filament/filament/object_manager.h"
-#include "experimental/filament/filament/render_target.h"
+#include "experimental/filament/filament/reflection_manager.h"
 #include "experimental/filament/render_context_filament.h"
 
 namespace mujoco {
@@ -233,14 +232,14 @@ const mjrMaterial& Renderable::GetMaterial() const {
   return material_;
 }
 
-void Renderable::Prepare(std::span<const mjrRenderRequest*> requests) {
+void Renderable::Prepare(std::span<const mjrRenderRequest*> requests,
+                         ReflectionManager* reflection_mgr) {
   // We assume BindMaterialInstance will be called with the same requests in
   // the same order. As such, we'll just store the draw state in a deque rather
   // than trying to perform any kind of matching with the requests.
   draw_queue_.clear();
   curr_state_ = DrawState();
 
-  int num_reflections = 0;
   for (const mjrRenderRequest* request : requests) {
     DrawState draw_state;
     mjrMaterial material = material_;
@@ -266,21 +265,8 @@ void Renderable::Prepare(std::span<const mjrRenderRequest*> requests) {
                             request->enable_reflections &&
                             material.reflectance > 0.0;
     if (reflective) {
-      // Allocate a render target as needed and assign it to the material's
-      // reflection texture.
-      if (reflect_targets_.size() == num_reflections) {
-        mjrRenderTargetConfig config;
-        mjr_defaultRenderTargetConfig(&config);
-        config.color_format = mjPIXEL_FORMAT_RGBA8;
-        config.depth_format = mjPIXEL_FORMAT_DEPTH32F;
-        reflect_targets_.push_back(
-            std::make_unique<RenderTarget>(GetEngine(), config));
-      }
-      RenderTarget* target = reflect_targets_[num_reflections].get();
-      target->Prepare(request->viewport.width, request->viewport.height);
-      material.reflection_texture = target->GetColorTexture();
-      draw_state.reflection_idx = num_reflections;
-      ++num_reflections;
+      material.reflection_texture = reflection_mgr->Register(
+          this, request->viewport.width, request->viewport.height);
     }
 
     draw_state.material_key =
@@ -319,14 +305,6 @@ MaterialKey Renderable::PrepareMaterialInstance(const mjrMaterial& material,
   return key;
 }
 
-RenderTarget* Renderable::GetReflectionTarget() const {
-  if (curr_state_.reflection_idx < 0) {
-    return nullptr;
-  } else {
-    return reflect_targets_[curr_state_.reflection_idx].get();
-  }
-}
-
 void Renderable::BindMaterialInstance(const mjrRenderRequest& request) {
   if (draw_queue_.empty()) {
     mju_error("No material instances to bind.");

src/experimental/filament/filament/renderable.h

@@ -18,7 +18,6 @@
 #include <cstdint>
 #include <deque>
 #include <functional>
-#include <memory>
 #include <span>
 #include <unordered_map>
 #include <vector>
@@ -33,7 +32,7 @@
 #include "experimental/filament/filament/material.h"
 #include "experimental/filament/filament/mesh.h"
 #include "experimental/filament/filament/object_manager.h"
-#include "experimental/filament/filament/render_target.h"
+#include "experimental/filament/filament/reflection_manager.h"
 #include "experimental/filament/render_context_filament.h"
 
 namespace mujoco {
@@ -107,16 +106,12 @@ class Renderable : public mjrRenderable {
   const mjrMaterial& GetMaterial() const;
 
   // Prepares the material instances for the given render requests.
-  void Prepare(std::span<const mjrRenderRequest*> requests);
+  void Prepare(std::span<const mjrRenderRequest*> requests,
+               ReflectionManager* reflection_mgr);
 
   // Binds the material instance for the given render request.
   void BindMaterialInstance(const mjrRenderRequest& request);
 
-  // Returns the render target used for reflections, if any. The main render
-  // function will use this target to render a reflection for this renderable
-  // based on the current camera position.
-  RenderTarget* GetReflectionTarget() const;
-
   static Renderable* downcast(mjrRenderable* renderable) {
     return static_cast<Renderable*>(renderable);
   }
@@ -147,15 +142,13 @@ class Renderable : public mjrRenderable {
   };
 
   // For each render request in a batch, we may need to render the object
-  // differently (e.g. a different reflection texture for each camera). However,
-  // we cannot change the material instance during the actual frame rendering
-  // (i.e. between beginFrame/endFrame). Instead, we can switch out the material
-  // instance entirely for each request. We keep track of which instance, as
-  // well as other render state, to use with each render request.
+  // differently. However, we cannot change the material instance during the
+  // actual frame rendering (i.e. between beginFrame/endFrame). Instead, we can
+  // switch out the material instance entirely for each request. We keep track
+  // of which instance, as well as other render state, to use with each render
+  // request.
   struct DrawState {
     MaterialKey material_key;
-    // Index to the reflection target to use, if any.
-    int reflection_idx = -1;
     bool cast_shadows = true;
     bool receive_shadows = true;
     bool wireframe = false;
@@ -180,7 +173,6 @@ class Renderable : public mjrRenderable {
   mjtGeom geom_type_ = mjGEOM_NONE;
   mjrMaterial material_;
   std::unordered_map<MaterialKey, filament::MaterialInstance*> instances_;
-  std::vector<std::unique_ptr<RenderTarget>> reflect_targets_;
   std::deque<DrawState> draw_queue_;
   DrawState curr_state_;
   filament::Scene* assigned_scene_ = nullptr;

src/experimental/filament/filament/scene_view.cc

@@ -38,6 +38,7 @@
 #include <math/vec4.h>
 #include <utils/EntityManager.h>
 #include <mujoco/mujoco.h>
+#include "experimental/filament/filament/reflection_manager.h"
 #include "experimental/filament/filament_util.h"
 #include "experimental/filament/filament/color_grading_options.h"
 #include "experimental/filament/filament/light.h"
@@ -121,6 +122,8 @@ static void SetupReflectionCamera(const mat4& surface_xform,
 
 SceneView::SceneView(filament::Engine* engine, const mjrSceneParams& params)
     : engine_(engine) {
+  reflection_mgr_ = std::make_unique<ReflectionManager>(engine_);
+
   scene_ = engine->createScene();
   camera_ = engine->createCamera(utils::EntityManager::get().create());
   reflect_camera_ = engine->createCamera(utils::EntityManager::get().create());
@@ -221,8 +224,10 @@ void SceneView::PrepareToRender(std::span<const mjrRenderRequest*> requests) {
       mju_error("Invalid scene for SceneView::PrepareToRender.");
     }
   }
+
+  reflection_mgr_->ClearRenderables();
   for (Renderable* renderable : renderables_) {
-    renderable->Prepare(requests);
+    renderable->Prepare(requests, reflection_mgr_.get());
   }
 }
 
@@ -257,27 +262,26 @@ void SceneView::Render(filament::Renderer* renderer, const mjrRenderRequest& req
     iter->BindMaterialInstance(request);
   }
 
-  for (auto& iter : renderables_) {
-    if (RenderTarget* target = iter->GetReflectionTarget()) {
-      viewport.left = 0;
-      viewport.bottom = 0;
-      reflect_view_->setViewport(viewport);
+  for (int i = 0; i < reflection_mgr_->GetNumRenderables(); ++i) {
+    Renderable* renderable =
+        Renderable::downcast(reflection_mgr_->GetRenderable(i));
+    const mat4 transform(renderable->GetTransform());
+    SetupReflectionCamera(transform, camera_, reflect_camera_);
 
-      // We assume the 0th entity is the reflective entity.
-      mat4 transform(iter->GetTransform());
-      SetupReflectionCamera(transform, camera_, reflect_camera_);
+    // Hide reflective surface from its own reflection pass.
+    std::uint8_t previous_layer_mask = renderable->SetLayerMask(0x00);
 
-      // Hide reflective surface from its own reflection pass.
-      std::uint8_t previous_layer_mask = iter->SetLayerMask(0x00);
+    // Render the reflection to its render target.
+    viewport.left = 0;
+    viewport.bottom = 0;
+    reflect_view_->setViewport(viewport);
+    const RenderTarget* target = reflection_mgr_->GetRenderTarget(i);
+    reflect_view_->setRenderTarget(target->GetFilamentRenderTarget());
+    renderer->render(reflect_view_);
+    reflect_view_->setRenderTarget(nullptr);
 
-      // Render the reflection to its render target.
-      reflect_view_->setRenderTarget(target->GetFilamentRenderTarget());
-      renderer->render(reflect_view_);
-      reflect_view_->setRenderTarget(nullptr);
-
-      // Unhide the reflective surface.
-      iter->SetLayerMask(previous_layer_mask);
-    }
+    // Unhide the reflective surface.
+    renderable->SetLayerMask(previous_layer_mask);
   }
 
   view->setRenderTarget(render_target ? render_target->GetFilamentRenderTarget()

src/experimental/filament/filament/scene_view.h

@@ -15,6 +15,7 @@
 #ifndef MUJOCO_SRC_EXPERIMENTAL_FILAMENT_FILAMENT_SCENE_VIEW_H_
 #define MUJOCO_SRC_EXPERIMENTAL_FILAMENT_FILAMENT_SCENE_VIEW_H_
 
+#include <memory>
 #include <span>
 #include <unordered_set>
 
@@ -27,6 +28,7 @@
 #include "experimental/filament/filament/color_grading_options.h"
 #include "experimental/filament/filament/light.h"
 #include "experimental/filament/filament/renderable.h"
+#include "experimental/filament/filament/reflection_manager.h"
 #include "experimental/filament/filament/texture.h"
 #include "experimental/filament/render_context_filament.h"
 
@@ -99,6 +101,7 @@ class SceneView : public mjrScene {
   // Custom view and camera for reflective surfaces.
   filament::View* reflect_view_ = nullptr;
   filament::Camera* reflect_camera_ = nullptr;
+  std::unique_ptr<ReflectionManager> reflection_mgr_;
 };
 }  // namespace mujoco