Support multiscene, multiview rendering.

Update Renderable so that it can switch its own draw states (material
instance, shadow casting, reflection textures, etc.) per render request.
This allows us to use the same Renderable with different effects from
different views within the same frame.

PiperOrigin-RevId: 922102817
Change-Id: I1cf02fddb7a00345a741db1bee88e716517765f4

src/experimental/filament/compat/scene_bridge.cc

@@ -129,8 +129,8 @@ void SceneBridge::Update(const mjrRect& viewport, const mjvScene* scene) {
       model_objects_->CreateSkinFlexMesh(scene, *geom);
     }
 
-    UniquePtr<mjrRenderable> renderable = CreateGeomRenderable(
-        *geom, ctx_, model_objects_.get(), headpos, scene->flags);
+    UniquePtr<mjrRenderable> renderable =
+        CreateGeomRenderable(*geom, ctx_, model_objects_.get(), scene->flags);
 
     mjrf_addRenderableToScene(scene_.get(), renderable.get());
     renderables_.push_back(std::move(renderable));

src/experimental/filament/compat/scene_geom_util.cc

@@ -40,12 +40,6 @@ static float GetPlaneTileSize(const mjModel* model, int matid,
   }
 }
 
-static bool IsBehind(const float* headpos, const float* pos, const float* mat) {
-  return ((headpos[0] - pos[0]) * mat[2] + (headpos[1] - pos[1]) * mat[5] +
-              (headpos[2] - pos[2]) * mat[8] <
-          0.0f);
-}
-
 static void PrepareGeomMeshes(mjrRenderable* renderable, const mjvGeom& geom,
                               ModelObjects* model_objs) {
   const mjModel* model = model_objs->GetModel();
@@ -69,19 +63,8 @@ static void PrepareGeomMeshes(mjrRenderable* renderable, const mjvGeom& geom,
       break;
     case mjGEOM_PLANE: {
       mjrf_setRenderableGeomMesh(renderable, geom_type, nstack, nslice, nquad);
-
-      float size[3];
-      std::memcpy(size, &geom.size, 3 * sizeof(float));
-      const bool is_infinite = !(size[0] > 0 && size[1] > 0);
-      if (is_infinite) {
-        // Infinite planes are scaled to match the tile size used by
-        // re-centering in engine_vis_visualize.c.
-        const float plane_scale = static_cast<float>(mjMAXPLANEGRID) / 2.0f;
-        size[0] = plane_scale;
-        size[1] = plane_scale;
-      }
       // Planes only define an xy size, so set the z-dimension to 1.0f.
-      size[2] = 1.0f;
+      const float size[3] = {geom.size[0], geom.size[1], 1.0f};
       mjrf_setRenderableSize(renderable, size);
       break;
     }
@@ -160,7 +143,7 @@ static void PrepareGeomMeshes(mjrRenderable* renderable, const mjvGeom& geom,
 }
 
 static void UpdateGeomMaterial(mjrRenderable* renderable, const mjvGeom& geom,
-                               ModelObjects* model_objs, const float headpos[3],
+                               ModelObjects* model_objs,
                                const mjtByte render_flags[mjNRNDFLAG]) {
   const mjModel* model = model_objs->GetModel();
 
@@ -175,21 +158,8 @@ static void UpdateGeomMaterial(mjrRenderable* renderable, const mjvGeom& geom,
   material.color[1] = geom.rgba[1];
   material.color[2] = geom.rgba[2];
   material.color[3] = geom.rgba[3];
-  if (geom.type == mjGEOM_PLANE) {
-    if (IsBehind(headpos, geom.pos, geom.mat)) {
-      material.color[3] *= 0.3;
-      mjrf_setRenderableReceiveShadows(renderable, false);
-      material.reflective = false;
-    } else {
-      mjrf_setRenderableReceiveShadows(renderable, true);
-      material.reflective = geom.reflectance > 0 && material.color[3] == 1.0f;
-    }
-  }
+
   mjrf_setRenderableLayerMask(renderable, geom.category);
-  if (geom.category == mjCAT_DECOR) {
-    mjrf_setRenderableCastShadows(renderable, false);
-    mjrf_setRenderableReceiveShadows(renderable, false);
-  }
 
   if (geom.matid >= 0 && geom.matid < model->nmat) {
     auto get_texture = [&](int role) -> const mjrTexture* {
@@ -310,12 +280,12 @@ static void UpdateGeomMaterial(mjrRenderable* renderable, const mjvGeom& geom,
 
 UniquePtr<mjrRenderable> CreateGeomRenderable(
     const mjvGeom& geom, mjrfContext* ctx, ModelObjects* model_objs,
-    const float headpos[3], const mjtByte render_flags[mjNRNDFLAG]) {
+    const mjtByte render_flags[mjNRNDFLAG]) {
   mjrRenderableParams params;
   mjr_defaultRenderableParams(&params);
   auto renderable = CreateRenderable(ctx, params);
   PrepareGeomMeshes(renderable.get(), geom, model_objs);
-  UpdateGeomMaterial(renderable.get(), geom, model_objs, headpos, render_flags);
+  UpdateGeomMaterial(renderable.get(), geom, model_objs, render_flags);
   return renderable;
 }
 }  // namespace mujoco

src/experimental/filament/compat/scene_geom_util.h

@@ -26,7 +26,7 @@ namespace mujoco {
 // Creates a Renderable from the given mjvGeom.
 UniquePtr<mjrRenderable> CreateGeomRenderable(
     const mjvGeom& geom, mjrfContext* ctx, ModelObjects* model_objs,
-    const float headpos[3], const mjtByte render_flags[mjNRNDFLAG]);
+    const mjtByte render_flags[mjNRNDFLAG]);
 
 }  // namespace mujoco
 

src/experimental/filament/filament/filament_context.cc

@@ -18,7 +18,9 @@
 #include <cstdint>
 #include <memory>
 #include <span>
+#include <unordered_map>
 #include <utility>
+#include <vector>
 
 #include <filament/Engine.h>
 #include <filament/IndexBuffer.h>
@@ -94,6 +96,14 @@ mjrFrameHandle FilamentContext::Render(
 
   ValidateSwapChains(requests);
 
+  std::unordered_map<mjrScene*, std::vector<const mjrRenderRequest*>> scene_to_requests;
+  for (const mjrRenderRequest& request : requests) {
+    scene_to_requests[request.scene].push_back(&request);
+  }
+  for (auto& [scene, requests] : scene_to_requests) {
+    SceneView::downcast(scene)->PrepareToRender(requests);
+  }
+
   bool render_began = false;
   mjrRenderTarget* current_target = nullptr;
   for (const mjrRenderRequest& request : requests) {

src/experimental/filament/filament/material.cc

@@ -13,6 +13,9 @@
 // limitations under the License.
 
 #include "experimental/filament/filament/material.h"
+#include <cstddef>
+#include <cstdint>
+#include <functional>
 
 #include <filament/Color.h>
 #include <filament/Material.h>
@@ -29,6 +32,51 @@
 
 namespace mujoco {
 
+template <class T>
+static void Combine(uint64_t& seed, const T* v) {
+  seed ^= std::hash<T>()(*v) + 0x9e3779b9 + (seed << 6) + (seed >> 2);
+}
+
+template <typename T>
+static uint64_t hash(const T& obj) {
+  static_assert(std::is_trivially_copyable_v<T>,
+                "Only trivially copyable types are hashable.");
+  int num_bytes = sizeof(T);
+  const std::byte* ptr = reinterpret_cast<const std::byte*>(&obj);
+  uint64_t seed = 0;
+  while (num_bytes >= sizeof(uint64_t)) {
+    Combine(seed, reinterpret_cast<const uint64_t*>(ptr));
+    ptr += sizeof(uint64_t);
+    num_bytes -= sizeof(uint64_t);
+  }
+  while (num_bytes >= sizeof(uint32_t)) {
+    Combine(seed, reinterpret_cast<const uint32_t*>(ptr));
+    ptr += sizeof(uint32_t);
+    num_bytes -= sizeof(uint32_t);
+  }
+  while (num_bytes >= sizeof(uint16_t)) {
+    Combine(seed, reinterpret_cast<const uint16_t*>(ptr));
+    ptr += sizeof(uint16_t);
+    num_bytes -= sizeof(uint16_t);
+  }
+  while (num_bytes >= sizeof(uint8_t)) {
+    Combine(seed, reinterpret_cast<const uint8_t*>(ptr));
+    ptr += sizeof(uint8_t);
+    num_bytes -= sizeof(uint8_t);
+  }
+  return seed;
+}
+
+uint64_t BuildMaterialKey(ObjectManager::MaterialType material_type,
+                         const mjrMaterial& material) {
+  // Normally, hashing the struct by memory would be a problem because of
+  // padding and other uninitialized data. However, we do a memset(0) on the
+  // entire structure in mjr_defaultMaterial so this should be safe.
+  uint64_t key = hash(material);
+  Combine(key, &material_type);
+  return key;
+}
+
 ObjectManager::MaterialType GetMaterialType(const mjrMaterial& material,
                                             const Mesh* mesh) {
   if (material.decor_ux) {
@@ -80,7 +128,7 @@ ObjectManager::MaterialType GetMaterialType(const mjrMaterial& material,
   if (color_texture == nullptr) {
     if (material.color[3] < 1.0f) {
       return ObjectManager::kPhongColorFade;
-    } else if (material.reflective) {
+    } else if (material.reflectance > 0) {
       return ObjectManager::kPhongColorReflect;
     } else {
       return ObjectManager::kPhongColor;
@@ -88,7 +136,7 @@ ObjectManager::MaterialType GetMaterialType(const mjrMaterial& material,
   } else if (color_texture->GetSamplerType() == mjTEXTURE_CUBE) {
     if (material.color[3] < 1.0f) {
       return ObjectManager::kPhongCubeFade;
-    } else if (material.reflective) {
+    } else if (material.reflectance > 0) {
       return ObjectManager::kPhongCubeReflect;
     } else {
       return ObjectManager::kPhongCube;
@@ -96,7 +144,7 @@ ObjectManager::MaterialType GetMaterialType(const mjrMaterial& material,
   } else if (has_texcoords) {
     if (material.color[3] < 1.0f) {
       return ObjectManager::kPhong2dUvFade;
-    } else if (material.reflective) {
+    } else if (material.reflectance > 0) {
       return ObjectManager::kPhong2dUvReflect;
     } else {
       return ObjectManager::kPhong2dUv;
@@ -104,7 +152,7 @@ ObjectManager::MaterialType GetMaterialType(const mjrMaterial& material,
   } else {
     if (material.color[3] < 1.0f) {
       return ObjectManager::kPhong2dFade;
-    } else if (material.reflective) {
+    } else if (material.reflectance > 0) {
       return ObjectManager::kPhong2dReflect;
     } else {
       return ObjectManager::kPhong2d;

src/experimental/filament/filament/material.h

@@ -15,6 +15,7 @@
 #ifndef MUJOCO_SRC_EXPERIMENTAL_FILAMENT_FILAMENT_MATERIAL_H_
 #define MUJOCO_SRC_EXPERIMENTAL_FILAMENT_FILAMENT_MATERIAL_H_
 
+#include <cstdint>
 #include <filament/Engine.h>
 #include <filament/MaterialInstance.h>
 #include "experimental/filament/filament/mesh.h"
@@ -23,6 +24,11 @@
 
 namespace mujoco {
 
+// Generates a hash for the given material (type) and material parameters.
+using MaterialKey = uint64_t;
+MaterialKey BuildMaterialKey(ObjectManager::MaterialType type,
+                             const mjrMaterial& material);
+
 // Returns a MaterialType that best matches the given material data and mesh.
 ObjectManager::MaterialType GetMaterialType(
     const mjrMaterial& material, const Mesh* mesh);

src/experimental/filament/filament/renderable.cc

@@ -15,13 +15,18 @@
 #include "experimental/filament/filament/renderable.h"
 
 #include <cstdint>
+#include <functional>
+#include <memory>
 #include <numbers>
+#include <span>
 
 #include <filament/Engine.h>
 #include <filament/Material.h>
 #include <filament/RenderableManager.h>
 #include <filament/Scene.h>
 #include <filament/TransformManager.h>
+#include <math/TMatHelpers.h>
+#include <math/TVecHelpers.h>
 #include <math/mat3.h>
 #include <math/mat4.h>
 #include <math/vec2.h>
@@ -29,10 +34,12 @@
 #include <math/vec4.h>
 #include <utils/EntityManager.h>
 #include <mujoco/mujoco.h>
+#include "experimental/filament/filament_util.h"
 #include "experimental/filament/filament/builtins.h"
 #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/render_context_filament.h"
 
 namespace mujoco {
@@ -57,7 +64,6 @@ Renderable::Renderable(filament::Engine* engine,
 Renderable::~Renderable() noexcept {
   filament::Engine* engine = GetEngine();
   utils::EntityManager& em = utils::EntityManager::get();
-
   for (Part& part : parts_) {
     if (assigned_scene_) {
       assigned_scene_->remove(part.entity);
@@ -65,11 +71,8 @@ Renderable::~Renderable() noexcept {
     engine->destroy(part.entity);
     em.destroy(part.entity);
   }
-  for (int i = 0; i < mjNUM_DRAW_MODES; ++i) {
-    if (instances_[i] != nullptr) {
-      engine->destroy(instances_[i]);
-      instances_[i] = nullptr;
-    }
+  for (auto& instance : instances_) {
+    engine->destroy(instance.second);
   }
 }
 
@@ -134,9 +137,6 @@ void Renderable::InitPartEntity(Part& part) {
   } else {
     builder.culling(false);
   }
-  if (instances_[static_cast<int>(draw_mode_)] != nullptr) {
-    builder.material(0, instances_[static_cast<int>(draw_mode_)]);
-  }
   builder.castShadows(params_.cast_shadows);
   builder.receiveShadows(params_.receive_shadows);
   builder.layerMask(0xff, params_.layer_mask);
@@ -168,7 +168,14 @@ void Renderable::UpdateTransform() {
   }
 
   filament::TransformManager& tm = GetEngine()->getTransformManager();
-  if (get_transform_fn_) {
+  if (geom_type_ == mjGEOM_PLANE && (trs_.size[0] <= 0 || trs_.size[1] <= 0)) {
+    infinite_plane_ = true;
+    const mat4f transform =
+        filament::math::mat4f(trs_.rotation, trs_.translation);
+    for (Part& part : parts_) {
+      tm.setTransform(tm.getInstance(part.entity), transform);
+    }
+  } else if (get_transform_fn_) {
     for (int i = 0; i < parts_.size(); ++i) {
       const mat4f transform = get_transform_fn_(i, trs_);
       tm.setTransform(tm.getInstance(parts_[i].entity), transform);
@@ -220,68 +227,171 @@ void Renderable::RemoveFromScene(filament::Scene* scene) {
 
 void Renderable::UpdateMaterial(const mjrMaterial& material) {
   material_ = material;
-
-  const Mesh* mesh = !parts_.empty() ? parts_[0].mesh : nullptr;
-
-  const ObjectManager::MaterialType type = GetMaterialType(material, mesh);
-  AssignMaterial(mjDRAW_MODE_COLOR, type);
-  if (!material_.decor_ux) {
-    AssignMaterial(mjDRAW_MODE_WIREFRAME, type);
-    AssignMaterial(mjDRAW_MODE_DEPTH, ObjectManager::kUnlitDepth);
-    AssignMaterial(mjDRAW_MODE_SEGMENTATION, ObjectManager::kUnlitSegmentation);
-  }
-
-  for (int i = 0; i < mjNUM_DRAW_MODES; ++i) {
-    if (instances_[i]) {
-      UpdateMaterialInstance(instances_[i], material_, object_mgr_);
-    }
-  }
-}
-
-void Renderable::AssignMaterial(mjrDrawMode mode,
-                                ObjectManager::MaterialType material_type) {
-  const int index = static_cast<int>(mode);
-
-  filament::Material* material = object_mgr_->GetMaterial(material_type);
-  if (instances_[index]) {
-    if (instances_[index]->getMaterial() == material) {
-      // The correct material is already assigned, do nothing.
-      return;
-    } else {
-      GetEngine()->destroy(instances_[index]);
-      instances_[index] = nullptr;
-    }
-  }
-  if (material) {
-    instances_[index] = material->createInstance();
-    if (geom_type_ == mjGEOM_PLANE || geom_type_ == mjGEOM_TRIANGLE) {
-      instances_[index]->setCullingMode(
-          filament::MaterialInstance::CullingMode::NONE);
-    }
-  }
 }
 
 const mjrMaterial& Renderable::GetMaterial() const {
   return material_;
 }
 
-void Renderable::SetDrawMode(mjrDrawMode mode) {
-  // Only SceneObjects support non-color draw modes.
-  if (material_.decor_ux) {
-    mode = mjDRAW_MODE_COLOR;
+void Renderable::Prepare(std::span<const mjrRenderRequest*> requests) {
+  // 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.
+  if (!draw_queue_.empty()) {
+    mju_error(
+        "Expected all previous draw calls to be bound before preparing new "
+        "draw calls.");
+  }
+  curr_state_ = DrawState();
+
+  int num_reflections = 0;
+  for (const mjrRenderRequest* request : requests) {
+    DrawState draw_state;
+    mjrMaterial material = material_;
+
+    draw_state.wireframe = (request->draw_mode == mjDRAW_MODE_WIREFRAME);
+    if (material.decor_ux) {
+      draw_state.cast_shadows = false;
+      draw_state.receive_shadows = false;
+    }
+
+    if (geom_type_ == mjGEOM_PLANE) {
+      const float3 camera_pos = ReadFloat3(request->camera.pos);
+      const float3 position = transform_[3].xyz;
+      const float3 forward = transform_[2].xyz;
+      const bool is_behind_camera = dot(camera_pos - position, forward) < 0;
+      if (is_behind_camera) {
+        material.color[3] *= 0.3;
+        draw_state.receive_shadows = false;
+      }
+    }
+
+    const bool reflective = request->draw_mode == mjDRAW_MODE_COLOR &&
+                            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;
+    }
+
+    draw_state.material_key =
+        PrepareMaterialInstance(material, request->draw_mode);
+    draw_queue_.push_back(draw_state);
+  }
+}
+
+MaterialKey Renderable::PrepareMaterialInstance(const mjrMaterial& material,
+                                                mjrDrawMode draw_mode) {
+  ObjectManager::MaterialType type;
+  if (draw_mode == mjDRAW_MODE_COLOR || draw_mode == mjDRAW_MODE_WIREFRAME) {
+    const Mesh* mesh = !parts_.empty() ? parts_[0].mesh : nullptr;
+    type = GetMaterialType(material, mesh);
+  } else if (draw_mode == mjDRAW_MODE_DEPTH) {
+    type = ObjectManager::kUnlitDepth;
+  } else if (draw_mode == mjDRAW_MODE_SEGMENTATION) {
+    type = ObjectManager::kUnlitSegmentation;
+  } else {
+    mju_error("Invalid draw mode: %d", draw_mode);
+    return 0;
   }
 
-  SetWireframe(mode == mjDRAW_MODE_WIREFRAME);
+  const MaterialKey key = BuildMaterialKey(type, material);
 
-  filament::MaterialInstance* instance = instances_[static_cast<int>(mode)];
-  if (instance) {
+  auto it = instances_.find(key);
+  if (it == instances_.end()) {
+    filament::MaterialInstance* instance =
+        object_mgr_->GetMaterial(type)->createInstance();
+    UpdateMaterialInstance(instance, material, object_mgr_);
+    if (geom_type_ == mjGEOM_PLANE || geom_type_ == mjGEOM_TRIANGLE) {
+      instance->setCullingMode(filament::MaterialInstance::CullingMode::NONE);
+    }
+    instances_[key] = instance;
+  }
+  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.");
+  }
+
+  if (geom_type_ == mjGEOM_PLANE && infinite_plane_) {
+    // Emulate an infinite plane by recentering a large quad in world space
+    // relative to the camera. We use the shared mjMAXPLANEGRID value as the
+    // size of the quad to ensure the texture scaling matches.
+    static constexpr float kInfiniteScale = 0.5f * mjMAXPLANEGRID;
+    const mat4f scaling =
+        mat4f::scaling(float3{kInfiniteScale, kInfiniteScale, 1.0f});
+
+    const float3 camera_pos = ReadFloat3(request.camera.pos);
+    const float3 plane_origin = transform_[3].xyz;
+    const mat3f plane_rotation = transform_.upperLeft();
+
+    const float3 vec = camera_pos - plane_origin;
+    const float3 plane_x = normalize(plane_rotation[0]);
+    const float3 plane_y = normalize(plane_rotation[1]);
+
+    // Project camera position onto the plane's local XY axes.
+    float dx = dot(vec, plane_x);
+    float dy = dot(vec, plane_y);
+
+    // Quantize based on uv_scale.
+    const float tile_size[] = {kInfiniteScale / material_.uv_scale[0],
+                               kInfiniteScale / material_.uv_scale[1]};
+    dx = tile_size[0] * mju_round(dx / tile_size[0]);
+    dy = tile_size[1] * mju_round(dy / tile_size[1]);
+
+    // Calculate the new center quad as a displacement from the plane origin.
+    const float3 displacement = dx * plane_x + dy * plane_y;
+    const float3 center = plane_origin + displacement;
+
+    const mat4f transform = mat4f(plane_rotation, center) * scaling;
+    filament::TransformManager& tm = GetEngine()->getTransformManager();
+    for (Part& part : parts_) {
+      tm.setTransform(tm.getInstance(part.entity), transform);
+    }
+  }
+
+  const DrawState& state = draw_queue_.front();
+  SetCastShadows(state.cast_shadows);
+  SetReceiveShadows(state.receive_shadows);
+  SetWireframe(state.wireframe);
+  SetMaterialInstance(state.material_key);
+  curr_state_ = state;
+  draw_queue_.pop_front();
+}
+
+void Renderable::SetMaterialInstance(MaterialKey key) {
+  if (key != curr_state_.material_key) {
+    filament::MaterialInstance* instance = instances_[key];
     filament::RenderableManager& rm = GetEngine()->getRenderableManager();
     for (Part& part : parts_) {
       filament::RenderableManager::Instance ri = rm.getInstance(part.entity);
       rm.setMaterialInstanceAt(ri, 0, instance);
     }
+    curr_state_.material_key = key;
   }
-  draw_mode_ = mode;
 }
 
 std::uint8_t Renderable::SetLayerMask(std::uint8_t mask) {

src/experimental/filament/filament/renderable.h

@@ -16,7 +16,11 @@
 #define MUJOCO_SRC_EXPERIMENTAL_FILAMENT_FILAMENT_RENDERABLE_H_
 
 #include <cstdint>
+#include <deque>
 #include <functional>
+#include <memory>
+#include <span>
+#include <unordered_map>
 #include <vector>
 
 #include <filament/Engine.h>
@@ -26,8 +30,10 @@
 #include <math/vec3.h>
 #include <utils/Entity.h>
 #include <mujoco/mujoco.h>
+#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/render_context_filament.h"
 
 namespace mujoco {
@@ -94,15 +100,23 @@ class Renderable : public mjrRenderable {
   // Removes this renderable from the filament Scene.
   void RemoveFromScene(filament::Scene* scene);
 
-  // Determines how this renderable will be drawn. See mjrDrawMode for details.
-  void SetDrawMode(mjrDrawMode mode);
-
   // Updates the parameters and textures of the material for this renderable.
   void UpdateMaterial(const mjrMaterial& material);
 
   // Returns this renderable's current material.
   const mjrMaterial& GetMaterial() const;
 
+  // Prepares the material instances for the given render requests.
+  void Prepare(std::span<const mjrRenderRequest*> requests);
+
+  // 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);
   }
@@ -132,6 +146,21 @@ 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.
+  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;
+  };
+
   // When composing a multi-part renderable, each Entity will have its own
   // transform offset based on the transform of the Renderable itself.
   using GetTransformFn = std::function<filament::math::mat4f(int, const Trs&)>;
@@ -140,22 +169,26 @@ class Renderable : public mjrRenderable {
   void InitPartEntity(Part& part);
   void UpdateTransform();
 
-  void AssignMaterial(mjrDrawMode mode,
-                      ObjectManager::MaterialType material_type);
+  MaterialKey PrepareMaterialInstance(const mjrMaterial& material,
+                                      mjrDrawMode draw_mode);
+  void SetMaterialInstance(MaterialKey key);
 
   filament::Engine* GetEngine();
 
   ObjectManager* object_mgr_;
   mjrRenderableParams params_;
-  filament::MaterialInstance* instances_[mjNUM_DRAW_MODES] = {nullptr};
   mjtGeom geom_type_ = mjGEOM_NONE;
   mjrMaterial material_;
-  mjrDrawMode draw_mode_ = mjDRAW_MODE_COLOR;
+  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;
   std::vector<Part> parts_;
   filament::math::mat4f transform_;
   GetTransformFn get_transform_fn_;
   Trs trs_;
+  bool infinite_plane_ = false;
   bool wireframe_ = false;
 };
 

src/experimental/filament/filament/scene_view.cc

@@ -14,11 +14,9 @@
 
 #include "experimental/filament/filament/scene_view.h"
 
-#include <algorithm>
-#include <array>
-#include <cstddef>
 #include <cstdint>
 #include <memory>
+#include <span>
 #include <string_view>
 
 #include <filament/ColorGrading.h>
@@ -145,6 +143,7 @@ SceneView::SceneView(filament::Engine* engine, const mjrSceneParams& params)
   reflect_view_->setPostProcessingEnabled(false);
   reflect_view_->setFrontFaceWindingInverted(true);
   reflect_view_->setVisibleLayers(0xff, params.reflection_layer_mask);
+  reflect_view_->setMultiSampleAntiAliasingOptions({.enabled = false});
 
   // Rotate the fog to align with mujoco's +Z up space.
   auto fog = main_view_->getFogEntity();
@@ -167,7 +166,6 @@ SceneView::~SceneView() {
   }
   lights_.clear();
   renderables_.clear();
-  reflect_targets_.clear();
   engine_->destroyCameraComponent(reflect_camera_->getEntity());
   engine_->destroy(reflect_view_);
   engine_->destroyCameraComponent(camera_->getEntity());
@@ -194,18 +192,11 @@ void SceneView::RemoveFromScene(Light* light) {
 void SceneView::AddToScene(Renderable* renderable) {
   if (renderables_.insert(renderable).second) {
     renderable->AddToScene(scene_);
-    if (renderable->GetMaterial().reflective) {
-      AddReflectiveRenderable(renderable);
-    }
   }
 }
 
 void SceneView::RemoveFromScene(Renderable* renderable) {
   if (renderables_.erase(renderable)) {
-    auto it = std::find(reflectives_.begin(), reflectives_.end(), renderable);
-    if (it != reflectives_.end()) {
-      reflectives_.erase(it);
-    }
     renderable->RemoveFromScene(scene_);
   }
 }
@@ -224,34 +215,30 @@ void SceneView::SetSkybox(const Texture* skybox_texture) {
   }
 }
 
+void SceneView::PrepareToRender(std::span<const mjrRenderRequest*> requests) {
+  for (const mjrRenderRequest* request : requests) {
+    if (request->scene != this) {
+      mju_error("Invalid scene for SceneView::PrepareToRender.");
+    }
+  }
+  for (Renderable* renderable : renderables_) {
+    renderable->Prepare(requests);
+  }
+}
+
 void SceneView::Render(filament::Renderer* renderer, const mjrRenderRequest& request) {
   if (request.scene != this) {
     mju_error("Invalid scene for SceneView::Render.");
   }
 
-  if (request.enable_reflections) {
-    EnableReflections();
-  } else {
-    DisableReflections();
-  }
-
   filament::Viewport viewport(request.viewport.left, request.viewport.bottom,
                               request.viewport.width, request.viewport.height);
-  main_view_->setViewport(viewport);
-  depth_segment_view_->setViewport(viewport);
-  reflect_view_->setViewport(viewport);
-
-  SetupCamera(request.camera, viewport, camera_);
-
-  for (auto& iter : renderables_) {
-    iter->SetDrawMode(request.draw_mode);
-  }
-
   filament::View* view = main_view_;
   if (request.draw_mode == mjDRAW_MODE_DEPTH ||
       request.draw_mode == mjDRAW_MODE_SEGMENTATION) {
     view = depth_segment_view_;
   }
+  view->setViewport(viewport);
   view->setShadowingEnabled(request.enable_shadows);
   view->setPostProcessingEnabled(request.enable_post_processing);
 
@@ -264,25 +251,30 @@ void SceneView::Render(filament::Renderer* renderer, const mjrRenderRequest& req
     view->setMultiSampleAntiAliasingOptions({.enabled = false});
   }
 
-  // Render reflection passes.
-  if (request.draw_mode == mjDRAW_MODE_COLOR && reflections_enabled_) {
-    for (size_t i = 0; i < reflectives_.size(); ++i) {
-      Renderable* renderable = reflectives_[i];
+  SetupCamera(request.camera, viewport, camera_);
+
+  for (auto& iter : renderables_) {
+    iter->BindMaterialInstance(request);
+
+    if (RenderTarget* target = iter->GetReflectionTarget()) {
+      viewport.left = 0;
+      viewport.bottom = 0;
+      reflect_view_->setViewport(viewport);
 
       // We assume the 0th entity is the reflective entity.
-      mat4 transform(renderable->GetTransform());
+      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);
+      std::uint8_t previous_layer_mask = iter->SetLayerMask(0x00);
 
       // Render the reflection to its render target.
-      reflect_view_->setRenderTarget(
-          reflect_targets_[i]->GetFilamentRenderTarget());
+      reflect_view_->setRenderTarget(target->GetFilamentRenderTarget());
       renderer->render(reflect_view_);
+      reflect_view_->setRenderTarget(nullptr);
 
       // Unhide the reflective surface.
-      renderable->SetLayerMask(previous_layer_mask);
+      iter->SetLayerMask(previous_layer_mask);
     }
   }
 
@@ -296,33 +288,6 @@ void SceneView::Render(filament::Renderer* renderer, const mjrRenderRequest& req
   }
 }
 
-void SceneView::AddReflectiveRenderable(Renderable* renderable) {
-  const int index = reflectives_.size();
-  reflectives_.push_back(renderable);
-
-  // Ensure we have the same number of render targets as we do reflective
-  // renderables.
-  while (reflect_targets_.size() < reflectives_.size()) {
-    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>(engine_, config));
-  }
-
-  // Prepare a render target for the reflective renderable.
-  auto viewport = reflect_view_->getViewport();
-  auto& target = reflect_targets_[index];
-  target->Prepare(viewport.width, viewport.height);
-
-  if (reflections_enabled_) {
-    mjrMaterial material = renderable->GetMaterial();
-    material.reflection_texture = target->GetColorTexture();
-    renderable->UpdateMaterial(material);
-  }
-}
-
 void SceneView::SetColorGradingOptions(const ColorGradingOptions& opts) {
   auto tone_mapper = CreateToneMapper(opts.tone_mapper);
   auto color_grading = ToBuilder(color_grading_options_)
@@ -336,33 +301,6 @@ void SceneView::SetColorGradingOptions(const ColorGradingOptions& opts) {
   color_grading_options_ = opts;
 }
 
-void SceneView::EnableReflections() {
-  if (reflections_enabled_) {
-    return;
-  }
-
-  reflections_enabled_ = true;
-  for (int i = 0; i < reflectives_.size(); ++i) {
-    Renderable* renderable = reflectives_[i];
-    mjrMaterial material = renderable->GetMaterial();
-    material.reflection_texture = reflect_targets_[i]->GetColorTexture();
-    renderable->UpdateMaterial(material);
-  }
-}
-
-void SceneView::DisableReflections() {
-  if (!reflections_enabled_) {
-    return;
-  }
-
-  reflections_enabled_ = false;
-  for (Renderable* renderable : reflectives_) {
-    mjrMaterial material = renderable->GetMaterial();
-    material.reflection_texture = nullptr;
-    renderable->UpdateMaterial(material);
-  }
-}
-
 filament::View* SceneView::GetDefaultRenderView() {
   return main_view_;
 }

src/experimental/filament/filament/scene_view.h

@@ -15,10 +15,8 @@
 #ifndef MUJOCO_SRC_EXPERIMENTAL_FILAMENT_FILAMENT_SCENE_VIEW_H_
 #define MUJOCO_SRC_EXPERIMENTAL_FILAMENT_FILAMENT_SCENE_VIEW_H_
 
-#include <array>
-#include <memory>
+#include <span>
 #include <unordered_set>
-#include <vector>
 
 #include <filament/Camera.h>
 #include <filament/ColorGrading.h>
@@ -29,7 +27,6 @@
 #include "experimental/filament/filament/color_grading_options.h"
 #include "experimental/filament/filament/light.h"
 #include "experimental/filament/filament/renderable.h"
-#include "experimental/filament/filament/render_target.h"
 #include "experimental/filament/filament/texture.h"
 #include "experimental/filament/render_context_filament.h"
 
@@ -55,6 +52,12 @@ class SceneView : public mjrScene {
   void RemoveFromScene(Renderable* renderable);
   void SetSkybox(const Texture* skybox_texture);
 
+  // Performs necessary preparations in order to render the given requests.
+  // Assumes that the Render() function will be called the same number of times
+  // and in the same order with the given requests.
+  void PrepareToRender(std::span<const mjrRenderRequest*> requests);
+
+  // Fulfills the given render request using the renderer.
   void Render(filament::Renderer* renderer, const mjrRenderRequest& request);
 
   // Returns the filament Engine managing the scene.
@@ -80,13 +83,6 @@ class SceneView : public mjrScene {
   }
 
  private:
-  void EnableReflections();
-  void DisableReflections();
-
-  // Marks a renderable as reflective. Reflective renderables have to be
-  // rendered in their own passes to create the reflective texture.
-  void AddReflectiveRenderable(Renderable* renderable);
-
   filament::Engine* engine_ = nullptr;
   filament::Scene* scene_ = nullptr;
   filament::Camera* camera_ = nullptr;
@@ -103,11 +99,6 @@ class SceneView : public mjrScene {
   // Custom view and camera for reflective surfaces.
   filament::View* reflect_view_ = nullptr;
   filament::Camera* reflect_camera_ = nullptr;
-
-  // The list of reflective renderables and their corresponding render targets.
-  bool reflections_enabled_ = true;
-  std::vector<Renderable*> reflectives_;
-  std::vector<std::unique_ptr<RenderTarget>> reflect_targets_;
 };
 }  // namespace mujoco
 

src/experimental/filament/filament_util.h

@@ -52,9 +52,9 @@ template <typename T>
 inline filament::math::mat3f ReadMat3(const T* arr, int index = 0) {
   // clang-format off
   const T* ptr = arr + (9 * index);
-  return filament::math::mat3f(ptr[0], ptr[3], ptr[6],
-                               ptr[1], ptr[4], ptr[7],
-                               ptr[2], ptr[5], ptr[8]);
+  return filament::math::mat3f(ptr[0], ptr[1], ptr[2],
+                               ptr[3], ptr[4], ptr[5],
+                               ptr[6], ptr[7], ptr[8]);
   // clang-format on
 }
 

src/experimental/filament/render_context_filament.h

@@ -139,8 +139,6 @@ typedef enum mjrDrawMode_ {
   mjDRAW_MODE_SEGMENTATION,
 } mjrDrawMode;
 
-enum { mjNUM_DRAW_MODES = 4 };  // Number of modes in `mjrDrawMode`.
-
 // Parameters describing the camera to use for rendering an image.
 typedef mjvGLCamera mjrCamera;
 
@@ -614,8 +612,6 @@ struct mjrMaterial {
   // The emissive (glow) factor of the object.
   float emissive;
 
-  // Whether or not the object is a reflective surface. Only applies to planes.
-  mjtByte reflective;
   // The blend factor to use for reflective surfaces. A value of 1.0 means that
   // the surface is fully reflective (i.e. a mirror).
   float reflectance;