Implement reflective surfaces.

Reflective surfaces (planes) use a custom shader to blend their
texture with a texture that contains the reflected image. The
reflected image is rendered to a texture as a separate render
pass in the SceneView.

Because reflective textures are generated after the preparation
phase, additional functions are needed to reapply the texture
to the drawable.

PiperOrigin-RevId: 889213413
Change-Id: I7bad39d62e3e5ba38cf6ac667fec5b87bfa7b2a2

src/experimental/filament/CMakeLists.txt

@@ -43,6 +43,7 @@ target_sources(${MUJOCO_FILAMENT_TARGET_NAME}
         filament/light.h
         filament/material.cc
         filament/material.h
+        filament/math_util.cc
         filament/math_util.h
         filament/model_util.cc
         filament/model_util.h

src/experimental/filament/filament/drawable.cc

@@ -181,7 +181,8 @@ void Drawable::Update(const mjModel* model, const mjvScene* scene,
   mjv_cameraInModel(head_pos, nullptr, nullptr, scene);
 
   SetTransform(geom);
-  UpdateMaterial(geom, scene->flags[mjRND_IDCOLOR], head_pos);
+  UpdateMaterial(geom, scene->flags[mjRND_IDCOLOR],
+                 scene->flags[mjRND_REFLECTION], head_pos);
   renderables_.SetWireframe(scene->flags[mjRND_WIREFRAME]);
 }
 
@@ -224,13 +225,21 @@ void Drawable::SetDrawMode(Material::DrawMode mode) {
   renderables_.SetMaterialInstance(material_.GetMaterialInstance(mode));
 }
 
+void Drawable::UpdateReflectionTexture(const filament::Texture* tex) {
+  material_.UpdateReflectionTexture(tex);
+}
+
+void Drawable::SetLayerMask(std::uint8_t mask) {
+  renderables_.SetLayerMask(mask);
+}
+
 void Drawable::SetTransform(const mjvGeom& geom) {
   // Flex and skin geometries are in global space.
   if (geom.type == mjGEOM_FLEX || geom.type == mjGEOM_SKIN) {
     return;
   }
 
-  const mat4 transform(ReadMat3(geom.mat), ReadFloat3(geom.pos));
+  transform_ = mat4(ReadMat3(geom.mat), ReadFloat3(geom.pos));
 
   float3 size = ReadFloat3(geom.size);
   filament::TransformManager& tm =
@@ -239,7 +248,7 @@ void Drawable::SetTransform(const mjvGeom& geom) {
     const utils::Entity& entity = renderables_[j];
 
     // Update object transform.
-    mat4 entity_transform = transform;
+    mat4 entity_transform = transform_;
 
     // Some built-in drawables are composed of multiple entities. For example,
     // capsules are a combination of a open tube and two dome end caps.
@@ -342,7 +351,7 @@ void Drawable::SetTransform(const mjvGeom& geom) {
 }
 
 void Drawable::UpdateMaterial(const mjvGeom& geom, bool use_segid_color,
-                              const mjtNum* headpos) {
+                              bool enable_reflection, const mjtNum* headpos) {
   ObjectManager* object_mgr = material_.GetObjectManager();
   const mjModel* model = object_mgr->GetModel();
 
@@ -351,8 +360,11 @@ void Drawable::UpdateMaterial(const mjvGeom& geom, bool use_segid_color,
     if (IsBehind(headpos, geom.pos, geom.mat)) {
       color[3] *= 0.3;
       renderables_.SetReceiveShadows(false);
+      reflective_ = false;
     } else {
       renderables_.SetReceiveShadows(true);
+      reflective_ =
+          enable_reflection && geom.reflectance > 0 && color.a == 1.0f;
     }
   }
 
@@ -405,6 +417,8 @@ void Drawable::UpdateMaterial(const mjvGeom& geom, bool use_segid_color,
       if (textures.color == nullptr) {
         if (color.a < 1.0f) {
           material_.SetNormalMaterialType(ObjectManager::kPhongColorFade);
+        } else if (reflective_) {
+          material_.SetNormalMaterialType(ObjectManager::kPhongColorReflect);
         } else {
           material_.SetNormalMaterialType(ObjectManager::kPhongColor);
         }
@@ -412,18 +426,24 @@ void Drawable::UpdateMaterial(const mjvGeom& geom, bool use_segid_color,
                 filament::Texture::Sampler::SAMPLER_CUBEMAP) {
         if (color.a < 1.0f) {
           material_.SetNormalMaterialType(ObjectManager::kPhongCubeFade);
+        } else if (reflective_) {
+          material_.SetNormalMaterialType(ObjectManager::kPhongCubeReflect);
         } else {
           material_.SetNormalMaterialType(ObjectManager::kPhongCube);
         }
       } else if (has_texcoords) {
         if (color.a < 1.0f) {
           material_.SetNormalMaterialType(ObjectManager::kPhong2dUvFade);
+        } else if (reflective_) {
+          material_.SetNormalMaterialType(ObjectManager::kPhong2dUvReflect);
         } else {
           material_.SetNormalMaterialType(ObjectManager::kPhong2dUv);
         }
       } else {
         if (color.a < 1.0f) {
           material_.SetNormalMaterialType(ObjectManager::kPhong2dFade);
+        } else if (reflective_) {
+          material_.SetNormalMaterialType(ObjectManager::kPhong2dReflect);
         } else {
           material_.SetNormalMaterialType(ObjectManager::kPhong2d);
         }
@@ -433,6 +453,7 @@ void Drawable::UpdateMaterial(const mjvGeom& geom, bool use_segid_color,
 
   Material::Params params;
   params.color = color;
+  params.reflectance = geom.reflectance;
   params.emissive = geom.emission;
   params.specular = geom.specular;
   params.glossiness = geom.shininess;

src/experimental/filament/filament/drawable.h

@@ -15,8 +15,10 @@
 #ifndef MUJOCO_SRC_EXPERIMENTAL_FILAMENT_FILAMENT_DRAWABLE_H_
 #define MUJOCO_SRC_EXPERIMENTAL_FILAMENT_FILAMENT_DRAWABLE_H_
 
+#include <cstdint>
 #include <filament/Engine.h>
 #include <filament/Scene.h>
+#include <math/mat4.h>
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjtnum.h>
 #include <mujoco/mjvisualize.h>
@@ -46,11 +48,26 @@ class Drawable {
   // transform, material, etc.) of the geom.
   void Update(const mjModel* model, const mjvScene* scene, const mjvGeom& geom);
 
+  // Returns the transform of the drawable.
+  const filament::math::mat4& GetTransform() const { return transform_; }
+
   // Swaps the MaterialInstance that will be used to render the Drawable (e.g.
   // normal, depth, segmentation, etc.). This must be called before the filament
   // beginFrame/endFrame.
   void SetDrawMode(Material::DrawMode mode);
 
+  // Sets the layer mask for all managed entities. This can be used to show
+  // or hide the drawable from specific passes. The default layer mask is 0x01.
+  void SetLayerMask(std::uint8_t mask);
+
+  // Returns true if the drawable is reflective.
+  bool IsReflective() const { return reflective_; }
+
+  // Sets the reflection texture for the drawable. We have a separate setter
+  // because we need to render the reflection texture before it can be applied
+  // to the material.
+  void UpdateReflectionTexture(const filament::Texture* tex);
+
  private:
   void AddMesh(int data_id);
   void AddHeightField(int hfield_id);
@@ -61,10 +78,12 @@ class Drawable {
 
   // Updates the material parameters of the drawable for rendering.
   void UpdateMaterial(const mjvGeom& geom, bool use_segid_color,
-                      const mjtNum* headpos);
+                      bool enable_reflection, const mjtNum* headpos);
 
   Material material_;
   Renderables renderables_;
+  bool reflective_ = false;
+  filament::math::mat4 transform_;
 };
 
 }  // namespace mujoco

src/experimental/filament/filament/material.cc

@@ -71,6 +71,11 @@ void Material::UpdateTextures(const Textures& textures) {
   UpdateMaterialInstances();
 }
 
+void Material::UpdateReflectionTexture(const filament::Texture* tex) {
+  textures_.reflection = tex;
+  UpdateMaterialInstances();
+}
+
 void Material::UpdateMaterialInstances() {
   filament::MaterialInstance* instance = instances_[DrawMode::kNormal];
   if (instance == nullptr) {
@@ -108,6 +113,9 @@ void Material::UpdateMaterialInstances() {
   if (material->hasParameter("UvOffset")) {
     instance->setParameter("UvOffset", params_.uv_offset);
   }
+  if (material->hasParameter("Reflectance")) {
+    instance->setParameter("Reflectance", params_.reflectance);
+  }
 
   if (instances_[DrawMode::kSegmentation]) {
     instances_[DrawMode::kSegmentation]->setParameter(
@@ -179,6 +187,14 @@ void Material::UpdateMaterialInstances() {
       instance->setParameter("Emissive", fallback, sampler);
     }
   }
+  if (material->hasParameter("Reflection")) {
+    if (textures_.reflection) {
+      instance->setParameter("Reflection", textures_.reflection, sampler);
+    } else {
+      auto* fallback = object_mgr_->GetFallbackTexture(mjTEXROLE_USER);
+      instance->setParameter("Reflection", fallback, sampler);
+    }
+  }
 }
 
 }  // namespace mujoco

src/experimental/filament/filament/material.h

@@ -46,6 +46,7 @@ class Material {
     const filament::Texture* occlusion = nullptr;
     const filament::Texture* orm = nullptr;
     const filament::Texture* emissive = nullptr;
+    const filament::Texture* reflection = nullptr;
   };
 
   // The parameters that can be applied to the drawable's material.
@@ -60,6 +61,7 @@ class Material {
     float metallic = -1.0f;
     float roughness = -1.0f;
     float emissive = -1.0f;
+    float reflectance = 0.0f;
     bool tex_uniform = false;
   };
 
@@ -78,6 +80,10 @@ class Material {
   // Updates the material textures of the drawable for rendering.
   void UpdateTextures(const Textures& textures);
 
+  // Update the reflection texture. We do this separately since the reflection
+  // texture needs to be rendered before it can be applied to the material.
+  void UpdateReflectionTexture(const filament::Texture* tex);
+
   // Returns the material instance assigned to the draw mode.
   filament::MaterialInstance* GetMaterialInstance(DrawMode mode) {
     return instances_[mode];

src/experimental/filament/filament/math_util.cc

@@ -0,0 +1,84 @@
+// 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/math_util.h"
+
+#include <math/mat4.h>
+#include <math/vec3.h>
+#include <math/vec4.h>
+#include <math/TVecHelpers.h>
+
+namespace mujoco {
+
+using filament::math::float3;
+using filament::math::float4;
+using filament::math::mat4;
+
+mat4 ToReflectionMatrix(const mat4& xform) {
+  const float3 normal = xform[2].xyz;
+  const float dist = dot(xform[3].xyz, normal);
+  // clang-format off
+  return mat4(
+      1.0f - 2.0f * normal.x * normal.x,
+      0.0f - 2.0f * normal.y * normal.x,
+      0.0f - 2.0f * normal.z * normal.x,
+      0.0f,
+      0.0f - 2.0f * normal.x * normal.y,
+      1.0f - 2.0f * normal.y * normal.y,
+      0.0f - 2.0f * normal.z * normal.y,
+      0.0f,
+      0.0f - 2.0f * normal.x * normal.z,
+      0.0f - 2.0f * normal.y * normal.z,
+      1.0f - 2.0f * normal.z * normal.z,
+      0.0f,
+      2.0f * dist * normal.x,
+      2.0f * dist * normal.y,
+      2.0f * dist * normal.z,
+      1.0f
+  );
+  // clang-format on
+}
+
+mat4 CalculateObliqueProjection(const mat4& projection, const float4& plane) {
+  mat4 res = projection;
+  auto sgn = [](float x) {
+    return (x > 0.0f) ? 1.0f : x < 0.0f ? -1.0f : 0.0f;
+  };
+
+  // The plane should be oriented such that the side to be kept is positive.
+  // The camera is at (0,0,0) in camera space. The value of the plane equation
+  // at the camera origin is plane.w. The reflected scene is on the opposite
+  // side of the plane from the camera. If plane.w is positive, the camera is
+  // on the positive side, so the reflected scene is on the negative side.
+  // We need to flip the plane in this case.
+  float4 clip_plane = plane;
+  if (plane.w > 0) {
+    clip_plane = -plane;
+  }
+
+  float4 q;
+  q.x = (sgn(clip_plane.x) + res[2][0]) / res[0][0];
+  q.y = (sgn(clip_plane.y) + res[2][1]) / res[1][1];
+  q.z = -1.0f;
+  q.w = (1.0f + res[2][2]) / res[3][2];
+
+  const float4 c = clip_plane * (2.0f / dot(clip_plane, q));
+  res[0][2] = c.x;
+  res[1][2] = c.y;
+  res[2][2] = c.z;
+  res[3][2] = c.w;
+  return res;
+}
+
+}  // namespace mujoco

src/experimental/filament/filament/math_util.h

@@ -16,6 +16,7 @@
 #define MUJOCO_SRC_EXPERIMENTAL_FILAMENT_FILAMENT_MATH_UTIL_H_
 
 #include <math/mat3.h>
+#include <math/mat4.h>
 #include <math/vec2.h>
 #include <math/vec3.h>
 #include <math/vec4.h>
@@ -54,6 +55,15 @@ inline filament::math::mat3 ReadMat3(const T* arr, int index = 0) {
   // clang-format on
 }
 
+// Calculates a reflection matrix for a plane defined by its transform.
+filament::math::mat4 ToReflectionMatrix(const filament::math::mat4& xform);
+
+// Modifies a projection matrix so its near plane coincides with an arbitrary
+// plane defined in camera space.
+filament::math::mat4 CalculateObliqueProjection(
+    const filament::math::mat4& projection,
+    const filament::math::float4& plane);
+
 }  // namespace mujoco
 
 #endif  // MUJOCO_SRC_EXPERIMENTAL_FILAMENT_FILAMENT_MATH_UTIL_H_

src/experimental/filament/filament/scene_view.cc

@@ -15,6 +15,7 @@
 #include "experimental/filament/filament/scene_view.h"
 
 #include <array>
+#include <cstddef>
 #include <memory>
 #include <optional>
 #include <string_view>
@@ -25,8 +26,9 @@
 #include <filament/LightManager.h>
 #include <filament/Material.h>
 #include <filament/Options.h>
-#include <filament/RenderableManager.h>
 #include <filament/Renderer.h>
+#include <filament/RenderableManager.h>
+#include <filament/RenderTarget.h>
 #include <filament/Skybox.h>
 #include <filament/TransformManager.h>
 #include <filament/View.h>
@@ -36,6 +38,7 @@
 #include <math/scalar.h>
 #include <math/vec3.h>
 #include <math/vec4.h>
+#include <math/TVecHelpers.h>
 #include <utils/EntityManager.h>
 #include <mujoco/mujoco.h>
 #include "experimental/filament/filament/color_grading_options.h"
@@ -45,11 +48,13 @@
 #include "experimental/filament/filament/math_util.h"
 #include "experimental/filament/filament/model_util.h"
 #include "experimental/filament/filament/object_manager.h"
+#include "experimental/filament/filament/render_target_util.h"
 
 namespace mujoco {
 
 using filament::math::float3;
 using filament::math::float4;
+using filament::math::mat3;
 using filament::math::mat4;
 
 static constexpr int kNormalIndex =
@@ -82,10 +87,37 @@ filament::ColorGrading::Builder ToBuilder(const ColorGradingOptions& opts) {
       .curves(opts.shadow_gamma, opts.mid_point, opts.highlight_scale);
 }
 
+// Sets up the `reflection_camera`'s projection matrix so that it is a
+// reflection of the `src_camera` across the plane defined by the
+// `surface_xform`. The generated projection is an oblique projection so that
+// the texture can be applied directly to the plane with screenspace uvs.
+static void SetupReflectionCamera(const mat4& surface_xform,
+                                  const filament::Camera* src_camera,
+                                  filament::Camera* reflection_camera,
+                                  float near = 0.01f, float far = 100.0f) {
+  const mat4 src_model_matrix = src_camera->getModelMatrix();
+  const mat4 src_view_matrix = src_camera->getViewMatrix();
+  const mat4 src_projection = src_camera->getProjectionMatrix();
+
+  reflection_camera->setModelMatrix(ToReflectionMatrix(surface_xform) *
+                                    src_model_matrix);
+
+  const float3 normal = surface_xform[2].xyz;
+  const float3 view_pos = (src_view_matrix * surface_xform[3]).xyz;
+  const float3 view_normal = (src_view_matrix * float4(normal, 0.0f)).xyz;
+
+  const float3 plane_normal_camera = view_normal;
+  const float plane_dist_camera = -dot(plane_normal_camera, view_pos);
+  const float4 oblique_plane(plane_normal_camera, plane_dist_camera);
+  const mat4 oblique = CalculateObliqueProjection(src_projection, oblique_plane);
+  reflection_camera->setCustomProjection(oblique, near, far);
+}
+
 SceneView::SceneView(filament::Engine* engine, ObjectManager* object_mgr)
     : object_mgr_(object_mgr), engine_(engine) {
   scene_ = engine_->createScene();
   camera_ = engine_->createCamera(utils::EntityManager::get().create());
+  reflect_camera_ = engine_->createCamera(utils::EntityManager::get().create());
 
   for (auto& view : views_) {
     view = engine_->createView();
@@ -93,6 +125,12 @@ SceneView::SceneView(filament::Engine* engine, ObjectManager* object_mgr)
     view->setCamera(camera_);
   }
 
+  reflect_view_ = engine_->createView();
+  reflect_view_->setScene(scene_);
+  reflect_view_->setCamera(reflect_camera_);
+  reflect_view_->setShadowingEnabled(false);
+  reflect_view_->setPostProcessingEnabled(false);
+
   const mjModel* m = object_mgr_->GetModel();
 
   // Configure options for the normal view.
@@ -197,6 +235,11 @@ SceneView::SceneView(filament::Engine* engine, ObjectManager* object_mgr)
 SceneView::~SceneView() {
   lights_.clear();
   drawables_.clear();
+  reflect_targets_.clear();
+
+  engine_->destroyCameraComponent(reflect_camera_->getEntity());
+  engine_->destroy(reflect_view_);
+
   engine_->destroyCameraComponent(camera_->getEntity());
   engine_->destroy(views_[kNormalIndex]->getColorGrading());
   for (auto& view : views_) {
@@ -216,6 +259,25 @@ void SceneView::Render(filament::Renderer* renderer, DrawMode draw_mode,
     view->setMultiSampleAntiAliasingOptions({.enabled = false});
   }
 
+  // Render reflection passes.
+  if (draw_mode == DrawMode::kNormal) {
+    for (size_t i = 0; i < reflectives_.size(); ++i) {
+      Drawable* drawable = reflectives_[i];
+
+      SetupReflectionCamera(drawable->GetTransform(), camera_, reflect_camera_);
+
+      // Hide reflective surface from its own reflection pass.
+      drawable->SetLayerMask(0x00);
+
+      // Render the reflection to its render target.
+      reflect_view_->setRenderTarget(reflect_targets_[i]->GetRenderTarget());
+      renderer->render(reflect_view_);
+
+      // Unhide the reflective surface.
+      drawable->SetLayerMask(0x01);
+    }
+  }
+
   view->setRenderTarget(target);
   renderer->render(view);
   view->setRenderTarget(nullptr);
@@ -233,10 +295,12 @@ filament::View* SceneView::PrepareRenderView(DrawMode mode) {
 }
 
 void SceneView::SetViewport(mjrRect viewport) {
+  auto filament_viewport = ReadViewport(viewport);
   aspect_ratio_ = (float)viewport.width / (float)viewport.height;
   for (auto& view : views_) {
-    view->setViewport(ReadViewport(viewport));
+    view->setViewport(filament_viewport);
   }
+  reflect_view_->setViewport(filament_viewport);
 }
 
 void SceneView::SetColorGradingOptions(const ColorGradingOptions& opts) {
@@ -382,6 +446,7 @@ void SceneView::UpdateScene(const mjrContext* context, const mjvScene* scene) {
     iter->RemoveFromScene(scene_);
   }
   drawables_.clear();
+  reflectives_.clear();
   for (int i = 0; i < scene->ngeom; ++i) {
     const mjvGeom* geom = scene->geoms + i;
 
@@ -394,6 +459,9 @@ void SceneView::UpdateScene(const mjrContext* context, const mjvScene* scene) {
     auto drawable = std::make_unique<Drawable>(object_mgr_, *geom);
     drawable->AddToScene(scene_);
     drawable->Update(object_mgr_->GetModel(), scene, *geom);
+    if (drawable->IsReflective()) {
+      AddReflectiveDrawable(drawable.get());
+    }
     drawables_.push_back(std::move(drawable));
   }
 
@@ -433,6 +501,24 @@ void SceneView::UpdateScene(const mjrContext* context, const mjvScene* scene) {
   }
 }
 
+void SceneView::AddReflectiveDrawable(Drawable* drawable) {
+  const int index = reflectives_.size();
+  reflectives_.push_back(drawable);
+
+  // Ensure we have the same number of render targets as we do reflective
+  // drawables.
+  while (reflect_targets_.size() < reflectives_.size()) {
+    reflect_targets_.push_back(std::make_unique<RenderTargetAndTextures>(
+        engine_, kRenderTargetReflectionColor, kRenderTargetDepth));
+  }
+
+  // Prepare a render target for the reflective drawable.
+  auto viewport = reflect_view_->getViewport();
+  auto& target = reflect_targets_[index];
+  target->Prepare(viewport.width, viewport.height);
+  drawable->UpdateReflectionTexture(target->GetColorTexture());
+}
+
 filament::Engine* SceneView::GetEngine() const { return engine_; }
 
 filament::View* SceneView::GetDefaultRenderView() {

src/experimental/filament/filament/scene_view.h

@@ -36,6 +36,7 @@
 #include "experimental/filament/filament/light.h"
 #include "experimental/filament/filament/material.h"
 #include "experimental/filament/filament/object_manager.h"
+#include "experimental/filament/filament/render_target_util.h"
 
 namespace mujoco {
 
@@ -86,6 +87,9 @@ class SceneView {
 
   void PrepareLights();
 
+  // Registers the given drawable as a reflective surface.
+  void AddReflectiveDrawable(Drawable* drawable);
+
   // 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(
@@ -108,6 +112,17 @@ class SceneView {
   float fallback_head_light_intensity_ = 0.f;
   float fallback_scene_light_intensity_ = 80'000.f;
   float fallback_environment_light_intensity_ = 5'000.f;
+
+  // Custom view and camera for reflective surfaces.
+  filament::View* reflect_view_ = nullptr;
+  filament::Camera* reflect_camera_ = nullptr;
+
+  // The list of drawables that are reflective.
+  std::vector<Drawable*> reflectives_;
+
+  // Each reflective drawable has its own render target which is used to render
+  // the reflected image.
+  std::vector<std::unique_ptr<RenderTargetAndTextures>> reflect_targets_;
 };
 
 }  // namespace mujoco