Merge branch 'google-deepmind:main' into mjx-warp-segmentation

doc/changelog.rst

@@ -2,6 +2,16 @@
 Changelog
 =========
 
+Upcoming version (not yet released)
+-----------------------------------
+
+- Added island support for the :ref:`PGS solver<soAlgorithms>`.
+
+Python
+^^^^^^
+
+- Added ``MjSpec.encode`` method, wrapping :ref:`mj_encode`.
+
 Version 3.8.0 (April 24, 2026)
 ------------------------------
 

doc/computation/index.rst

@@ -1424,12 +1424,6 @@ While islanding is not free (see implementation in `engine_island.c
 - Unconstrained DOFs are completely untouched by the solver, which otherwise needs to discover that they are unaffected.
 - Solving separate islands can be multi-threaded.
 
-.. admonition:: Known issues
-   :class: note
-
-   Islanding is not yet supported by the PGS solver.
-
-
 .. _soParameters:
 
 Parameters

doc/python.rst

@@ -547,6 +547,12 @@ Compiled ``MjSpec`` objects can be saved to XML string with the ``to_xml()`` met
      </worldbody>
    </mujoco>
 
+Alternatively, the spec can be saved directly to a file using ``encode()``:
+
+.. code-block:: python
+
+   spec.encode('model.xml', model)
+
 Attachment
 ----------
 

mjx/mujoco/mjx/_src/types.py

@@ -775,6 +775,7 @@ class Model(PyTreeNode):
   mesh_texcoord: np.ndarray
   flex_vertadr: np.ndarray
   flex_vertnum: np.ndarray
+  flex_interp: np.ndarray
   flex_vert0: np.ndarray
   flex_nodeadr: np.ndarray
   flex_nodenum: np.ndarray

python/mujoco/specs.cc

@@ -431,6 +431,51 @@ PYBIND11_MODULE(_specs, m) {
       throw FatalError(std::string(err.data()));
     }
   });
+  mjSpec.def(
+      "encode",
+      [](MjSpec& self, std::string filename,
+         std::optional<py::object> model,
+         std::optional<std::string> content_type) -> int {
+        raw::MjModel* m = nullptr;
+        if (model.has_value() && !model->is_none()) {
+          auto& wrapper =
+              py::cast<_impl::MjModelWrapper&>(*model);
+          m = wrapper.get();
+        }
+
+        mjVFS vfs;
+        mjVFS* vfs_ptr = nullptr;
+        if (!self.assets.empty()) {
+          mj_defaultVFS(&vfs);
+          vfs_ptr = &vfs;
+          for (const auto& asset : self.assets) {
+            std::string buffer_name =
+                py::cast<std::string>(asset.first);
+            std::string buffer =
+                py::cast<std::string>(asset.second);
+            mj_addBufferVFS(vfs_ptr, buffer_name.c_str(),
+                            buffer.c_str(), buffer.size());
+          }
+        }
+
+        std::array<char, 1024> err;
+        err[0] = '\0';
+        const char* ct =
+            content_type.has_value() ? content_type->c_str() : nullptr;
+        int nbytes = mj_encode(self.ptr, m, filename.c_str(), ct,
+                               vfs_ptr, err.data(), err.size());
+
+        if (vfs_ptr) {
+          mj_deleteVFS(vfs_ptr);
+        }
+
+        if (nbytes < 0) {
+          throw FatalError(std::string(err.data()));
+        }
+        return nbytes;
+      },
+      py::arg("filename"), py::arg("model") = py::none(),
+      py::arg("content_type") = py::none());
   mjSpec.def(
       "add_default",
       [](MjSpec* spec, std::string& classname,

python/mujoco/specs_test.py

@@ -22,6 +22,7 @@ import textwrap
 import typing
 import zipfile  # pylint: disable=unused-import
 
+from absl import flags
 from absl.testing import absltest
 from etils import epath
 import mujoco
@@ -34,6 +35,11 @@ def get_linenumber():
 
 
 class SpecsTest(absltest.TestCase):
+  def setUp(self):
+    super().setUp()
+    # Mark flags as parsed to avoid pytest errors about unparsed flags.
+    # This is needed for `create_tempdir()` calls below.
+    flags.FLAGS.mark_as_parsed()
 
   def test_typing(self):
     spec = mujoco.MjSpec()
@@ -1964,5 +1970,49 @@ class SpecsTest(absltest.TestCase):
     self.assertGreater(cam_sd[0], cam_sd[1])  # dist > depth
     self.assertAlmostEqual(cam_sd[1], 2.0, places=6)  # depth is still 2.0
 
+  def test_encode_xml(self):
+    # Create a simple spec and compile.
+    spec = mujoco.MjSpec()
+    body = spec.worldbody.add_body()
+    geom = body.add_geom()
+    geom.size[0] = 1
+    model = spec.compile()
+
+    # Encode to XML.
+    filename = os.path.join(self.create_tempdir().full_path, 'output.xml')
+    nbytes = spec.encode(filename, model)
+    self.assertGreater(nbytes, 0)
+
+    # Verify the output is valid XML that can be loaded.
+    reloaded = mujoco.MjSpec.from_file(filename)
+    reloaded_model = reloaded.compile()
+    self.assertEqual(reloaded_model.ngeom, model.ngeom)
+
+  def test_encode_xml_without_model(self):
+    # Create a simple spec and compile so XML can be written.
+    spec = mujoco.MjSpec()
+    body = spec.worldbody.add_body()
+    geom = body.add_geom()
+    geom.size[0] = 1
+    spec.compile()
+
+    # Encode to XML without passing a model explicitly.
+    filename = os.path.join(self.create_tempdir().full_path, 'output.xml')
+    nbytes = spec.encode(filename)
+    self.assertGreater(nbytes, 0)
+
+  def test_encode_no_encoder_raises(self):
+    # Create a simple spec and compile.
+    spec = mujoco.MjSpec()
+    body = spec.worldbody.add_body()
+    geom = body.add_geom()
+    geom.size[0] = 1
+    model = spec.compile()
+
+    # Encode with an unknown extension should fail.
+    filename = os.path.join(self.create_tempdir().full_path, 'output.unknown')
+    with self.assertRaises(mujoco.FatalError):
+      spec.encode(filename, model)
+
 if __name__ == '__main__':
   absltest.main()

simulate/simulate.cc

@@ -349,7 +349,7 @@ void UpdateProfiler(mj::Simulate* sim, const mjModel* m, const mjData* d) {
       sim->figcost.linepnt[start + 2] = 0;
     }
 
-    for (int i=0; i<sim->figcost.linepnt[0]; i++) {
+    for (int i=0; i<npoints; i++) {
       // x
       sim->figcost.linedata[start + 0][2*i] = i;
       sim->figcost.linedata[start + 1][2*i] = i;

src/engine/engine_forward.c

@@ -953,7 +953,7 @@ static void solve_threaded(const mjModel* m, mjData* d, int flg_Newton) {
 // compute efc_b, efc_force, qfrc_constraint; update qacc
 void mj_fwdConstraint(const mjModel* m, mjData* d) {
   TM_START;
-  int nv = m->nv, nefc = d->nefc, nisland = d->nisland;
+  int nv = m->nv, nefc = d->nefc, nisland = d->nisland, nidof;
 
   // always clear qfrc_constraint
   mju_zero(d->qfrc_constraint, nv);
@@ -970,50 +970,69 @@ void mj_fwdConstraint(const mjModel* m, mjData* d) {
   mj_mulJacVec(m, d, d->efc_b, d->qacc_smooth);
   mju_subFrom(d->efc_b, d->efc_aref, nefc);
 
+  // check for invalid solver type
+  if (m->opt.solver != mjSOL_PGS && m->opt.solver != mjSOL_CG && m->opt.solver != mjSOL_NEWTON) {
+    mjERROR("unknown solver type %d", m->opt.solver);
+  }
+
   // warmstart solver
   warmstart(m, d);
   mju_zeroInt(d->solver_niter, mjNISLAND);
 
   // check if islands are supported
-  int islands_supported = !mjDISABLED(mjDSBL_ISLAND)    &&
+  int islands_supported = !mjDISABLED(mjDSBL_ISLAND) && nisland > 0;
-                          nisland > 0                   &&
-                          m->opt.noslip_iterations == 0 &&
-                          (m->opt.solver == mjSOL_CG || m->opt.solver == mjSOL_NEWTON);
 
   // run solver over constraint islands
   if (islands_supported) {
-    int nidof = d->nidof;
+    switch ((mjtSolver) m->opt.solver) {
-
+    case mjSOL_PGS:
-    // copy inputs to islands (vel+acc deps, pos-dependent already copied in mj_island)
-    mju_gather(d->ifrc_smooth,     d->qfrc_smooth,     d->map_idof2dof, nidof);
-    mju_gather(d->ifrc_constraint, d->qfrc_constraint, d->map_idof2dof, nidof);
-    mju_gather(d->iacc_smooth,     d->qacc_smooth,     d->map_idof2dof, nidof);
-    mju_gather(d->iacc,            d->qacc,            d->map_idof2dof, nidof);
-    mju_gather(d->iefc_force,      d->efc_force,       d->map_iefc2efc, nefc);
-    mju_gather(d->iefc_aref,       d->efc_aref,        d->map_iefc2efc, nefc);
-
-    // solve per island, with or without threads
-    if (!d->threadpool) {
-      // no threadpool, loop over islands
       for (int island=0; island < nisland; island++) {
-        if (m->opt.solver == mjSOL_NEWTON) {
+        mj_solPGS_island(m, d, island, m->opt.iterations);
-          mj_solNewton_island(m, d, island, m->opt.iterations);
-        } else {
-          mj_solCG_island(m, d, island, m->opt.iterations);
-        }
       }
-    } else {
+      break;
-      // have threadpool, solve using threads
+
-      solve_threaded(m, d, m->opt.solver == mjSOL_NEWTON);
+    case mjSOL_CG:
+    case mjSOL_NEWTON:
+      // copy inputs to islands (vel+acc deps, pos-dependent already copied in mj_island)
+      nidof = d->nidof;
+      mju_gather(d->ifrc_smooth,     d->qfrc_smooth,     d->map_idof2dof, nidof);
+      mju_gather(d->ifrc_constraint, d->qfrc_constraint, d->map_idof2dof, nidof);
+      mju_gather(d->iacc_smooth,     d->qacc_smooth,     d->map_idof2dof, nidof);
+      mju_gather(d->iacc,            d->qacc,            d->map_idof2dof, nidof);
+      mju_gather(d->iefc_force,      d->efc_force,       d->map_iefc2efc, nefc);
+      mju_gather(d->iefc_aref,       d->efc_aref,        d->map_iefc2efc, nefc);
+
+      // solve per island, with or without threads
+      if (!d->threadpool) {
+        // no threadpool, loop over islands
+        for (int island=0; island < nisland; island++) {
+          if (m->opt.solver == mjSOL_NEWTON) {
+            mj_solNewton_island(m, d, island, m->opt.iterations);
+          } else {
+            mj_solCG_island(m, d, island, m->opt.iterations);
+          }
+        }
+      } else {
+        // have threadpool, solve using threads
+        solve_threaded(m, d, m->opt.solver == mjSOL_NEWTON);
+      }
+
+      // copy back solver outputs (scatter dofs since ni <= nv)
+      mju_scatter(d->qacc,            d->iacc,            d->map_idof2dof, nidof);
+      mju_scatter(d->qfrc_constraint, d->ifrc_constraint, d->map_idof2dof, nidof);
+      mju_gather(d->efc_force, d->iefc_force, d->map_efc2iefc, nefc);
+      break;
     }
 
-    // copy back solver outputs (scatter dofs since ni <= nv)
+    // run noslip solver per island if enabled
-    mju_scatter(d->qacc,            d->iacc,            d->map_idof2dof, nidof);
+    if (m->opt.noslip_iterations > 0) {
-    mju_scatter(d->qfrc_constraint, d->ifrc_constraint, d->map_idof2dof, nidof);
+      for (int island=0; island < nisland; island++) {
-    mju_gather(d->efc_force, d->iefc_force, d->map_efc2iefc, nefc);
+        mj_solNoSlip_island(m, d, island, m->opt.noslip_iterations);
+      }
+    }
   }
 
-  // run solver over all constraints
+  // run solver over all constraints (monolithic)
   else {
     switch ((mjtSolver) m->opt.solver) {
     case mjSOL_PGS:                     // PGS
@@ -1027,15 +1046,17 @@ void mj_fwdConstraint(const mjModel* m, mjData* d) {
     case mjSOL_NEWTON:                  // Newton
       mj_solNewton(m, d, m->opt.iterations);
       break;
+    }
 
-    default:
+    // run noslip solver if enabled
-      mjERROR("unknown solver type %d", m->opt.solver);
+    if (m->opt.noslip_iterations > 0) {
+      mj_solNoSlip(m, d, m->opt.noslip_iterations);
     }
   }
 
-  // run noslip solver if enabled
+  // dual solvers: map efc_force to joint space (always monolithic)
-  if (m->opt.noslip_iterations > 0) {
+  if (m->opt.solver == mjSOL_PGS || m->opt.noslip_iterations > 0) {
-    mj_solNoSlip(m, d, m->opt.noslip_iterations);
+    mj_dualFinish(m, d);
   }
 
   TM_END(mjTIMER_CONSTRAINT);

src/engine/engine_solver.c

@@ -66,7 +66,6 @@ static void saveStats(const mjModel* m, mjData* d, int island, int iter,
 
 
 // finalize dual solver: map to joint space
-// TODO: b/295296178 - add island support to Dual solvers
 static void dualFinish(const mjModel* m, mjData* d) {
   // map constraint force to joint space
   mj_mulJacTVec(m, d, d->qfrc_constraint, d->efc_force);
@@ -77,10 +76,17 @@ static void dualFinish(const mjModel* m, mjData* d) {
 }
 
 
+// PGS: map efc_force to joint space
+void mj_dualFinish(const mjModel* m, mjData* d) {
+  dualFinish(m, d);
+}
+
+
 // compute 1/diag(AR)
-// TODO: b/295296178 - add island support to Dual solvers
+//   res[c] = 1 / AR[efclist[c], efclist[c]] for c = 0..nefc-1
-static void ARdiaginv(const mjModel* m, const mjData* d, mjtNum* res, int flg_subR) {
+//   efclist is NULL for monolithic (sequential) iteration
-  int nefc = d->nefc;
+static void ARdiaginv(const mjModel* m, const mjData* d, mjtNum* res,
+                      int nefc, const int* efclist, int flg_subR) {
   const mjtNum *AR = d->efc_AR;
   const mjtNum *R = d->efc_R;
 
@@ -90,12 +96,13 @@ static void ARdiaginv(const mjModel* m, const mjData* d, mjtNum* res, int flg_su
     const int *rownnz = d->efc_AR_rownnz;
     const int *colind = d->efc_AR_colind;
 
-    for (int i=0; i < nefc; i++) {
+    for (int c=0; c < nefc; c++) {
+      int i = efclist ? efclist[c] : c;
       int nnz = rownnz[i];
       for (int j=0; j < nnz; j++) {
         int adr = rowadr[i] + j;
         if (i == colind[adr]) {
-          res[i] = 1 / (flg_subR ? mju_max(mjMINVAL, AR[adr] - R[i]) : AR[adr]);
+          res[c] = 1 / (flg_subR ? mju_max(mjMINVAL, AR[adr] - R[i]) : AR[adr]);
           break;
         }
       }
@@ -104,16 +111,17 @@ static void ARdiaginv(const mjModel* m, const mjData* d, mjtNum* res, int flg_su
 
   // dense
   else {
-    for (int i=0; i < nefc; i++) {
+    int d_nefc = d->nefc;  // global nefc
-      int adr = i * (nefc + 1);
+    for (int c=0; c < nefc; c++) {
-      res[i] = 1 / (flg_subR ? mju_max(mjMINVAL, AR[adr] - R[i]) : AR[adr]);
+      int i = efclist ? efclist[c] : c;
+      int adr = i * (d_nefc + 1);
+      res[c] = 1 / (flg_subR ? mju_max(mjMINVAL, AR[adr] - R[i]) : AR[adr]);
     }
   }
 }
 
 
 // extract diagonal block from AR, clamp diag to 1e-10 if flg_subR
-// TODO: b/295296178 - add island support to Dual solvers
 static void extractBlock(const mjModel* m, const mjData* d, mjtNum* Ac,
                          int start, int n, int flg_subR) {
   int nefc = d->nefc;
@@ -173,7 +181,6 @@ static void extractBlock(const mjModel* m, const mjData* d, mjtNum* Ac,
 
 
 // compute residual for one block
-// TODO: b/295296178 - add island support to Dual solvers
 static void residual(const mjModel* m, const mjData* d, mjtNum* res, int i, int dim, int flg_subR) {
   int nefc = d->nefc;
 
@@ -203,7 +210,6 @@ static void residual(const mjModel* m, const mjData* d, mjtNum* res, int i, int
 
 
 // compute cost change
-// TODO: b/295296178 - add island support to Dual solvers
 static mjtNum costChange(const mjtNum* A, mjtNum* force, const mjtNum* oldforce,
                          const mjtNum* res, int dim) {
   mjtNum change;
@@ -229,9 +235,9 @@ static mjtNum costChange(const mjtNum* A, mjtNum* force, const mjtNum* oldforce,
 
 
 // set efc_state to dual constraint state; return nactive
-// TODO: b/295296178 - add island support to Dual solvers
+//   iterates over efclist (or sequentially if NULL), classifies by ne/nf ranges
-static int dualState(const mjModel* m, const mjData* d, int* state) {
+static int dualState(const mjData* d, int* state,
-  int ne = d->ne, nf = d->nf, nefc = d->nefc;
+                     int ne, int nf, int nefc, const int* efclist) {
   const mjtNum* force = d->efc_force;
   const mjtNum* floss = d->efc_frictionloss;
 
@@ -239,10 +245,14 @@ static int dualState(const mjModel* m, const mjData* d, int* state) {
   int nactive = ne + nf;
 
   // equality
-  mju_fillInt(state, mjCNSTRSTATE_QUADRATIC, ne);
+  for (int c=0; c < ne; c++) {
+    int i = efclist ? efclist[c] : c;
+    state[i] = mjCNSTRSTATE_QUADRATIC;
+  }
 
   // friction
-  for (int i=ne; i < ne+nf; i++) {
+  for (int c=ne; c < ne+nf; c++) {
+    int i = efclist ? efclist[c] : c;
     if (force[i] <= -floss[i]) {
       state[i] = mjCNSTRSTATE_LINEARPOS;  // opposite of primal
     } else if (force[i] >= floss[i]) {
@@ -253,7 +263,9 @@ static int dualState(const mjModel* m, const mjData* d, int* state) {
   }
 
   // limit and contact
-  for (int i=ne+nf; i < nefc; i++) {
+  for (int c=ne+nf; c < nefc; c++) {
+    int i = efclist ? efclist[c] : c;
+
     // non-negative
     if (d->efc_type[i] != mjCNSTR_CONTACT_ELLIPTIC) {
       if (force[i] <= 0) {
@@ -302,7 +314,7 @@ static int dualState(const mjModel* m, const mjData* d, int* state) {
       mju_fillInt(state+i, result, dim);
 
       // advance
-      i += (dim-1);
+      c += (dim-1);
     }
   }
 
@@ -310,26 +322,85 @@ static int dualState(const mjModel* m, const mjData* d, int* state) {
 }
 
 
+// update constraint state, return nactive and nchange
+static int dualStateChange(const mjData* d, int* state, int* oldstate,
+                           int ne, int nf, int nefc,
+                           const int* efclist, int* nchange) {
+  // save old state
+  for (int c=0; c < nefc; c++) {
+    int i = efclist ? efclist[c] : c;
+    oldstate[c] = state[i];
+  }
+
+  // update state
+  int nactive = dualState(d, state, ne, nf, nefc, efclist);
+
+  // count state changes
+  *nchange = 0;
+  for (int c=0; c < nefc; c++) {
+    int i = efclist ? efclist[c] : c;
+    *nchange += (oldstate[c] != state[i]);
+  }
+
+  return nactive;
+}
+
+
+// solve QCQP and project onto friction ellipsoid, write to force[i+1..i+dim-1]
+static void solveQCQP(mjtNum* force, int i, int dim,
+                      mjtNum* Ac, mjtNum* bc, const mjtNum* mu) {
+  int flg_active;
+  mjtNum v[6];
+
+  // solve
+  if (dim == 3) {
+    flg_active = mju_QCQP2(v, Ac, bc, mu, force[i]);
+  } else if (dim == 4) {
+    flg_active = mju_QCQP3(v, Ac, bc, mu, force[i]);
+  } else {  // dim == 5
+    flg_active = mju_QCQP(v, Ac, bc, mu, force[i], dim-1);
+  }
+
+  // on constraint: put v on ellipsoid, in case QCQP is approximate
+  if (flg_active) {
+    mjtNum s = 0;
+    for (int j=0; j < dim-1; j++) {
+      s += v[j]*v[j] / (mu[j]*mu[j]);
+    }
+    s = mju_sqrt(force[i]*force[i] / mju_max(mjMINVAL, s));
+    for (int j=0; j < dim-1; j++) {
+      v[j] *= s;
+    }
+  }
+
+  // assign
+  mju_copy(force+i+1, v, dim-1);
+}
+
+
 //---------------------------- PGS solver ----------------------------------------------------------
 
-// TODO: b/295296178 - add island support to Dual solvers
+// core PGS solver: iterates over constraints specified by efclist
-void mj_solPGS(const mjModel* m, mjData* d, int maxiter) {
+//   island: island index for stats (use -1 for monolithic, mapped to 0)
-  int ne = d->ne, nf = d->nf, nefc = d->nefc;
+//   ne, nf, nefc: constraint type counts
+//   efclist: maps list position c to monolithic efc index (NULL for sequential)
+static void solPGS(const mjModel* m, mjData* d, int island,
+                   int ne, int nf, int nefc,
+                   const int* efclist, int maxiter) {
   const mjtNum *floss = d->efc_frictionloss;
   mjtNum *force = d->efc_force;
   mj_markStack(d);
   mjtNum* ARinv = mjSTACKALLOC(d, nefc, mjtNum);
   int* oldstate = mjSTACKALLOC(d, nefc, int);
 
-  // TODO: b/295296178 - Use island index (currently hardcoded to 0)
+  int island_stat = mjMAX(0, island);  // island index for diagnostic stats
-  int island = 0;
   mjtNum scale = 1 / (m->stat.meaninertia * mjMAX(1, m->nv));
 
   // precompute inverse diagonal of AR
-  ARdiaginv(m, d, ARinv, 0);
+  ARdiaginv(m, d, ARinv, nefc, efclist, 0);
 
   // initial constraint state
-  dualState(m, d, d->efc_state);
+  dualState(d, d->efc_state, ne, nf, nefc, efclist);
 
   // main iteration
   int iter = 0;
@@ -338,7 +409,9 @@ void mj_solPGS(const mjModel* m, mjData* d, int maxiter) {
     mjtNum improvement = 0;
 
     // perform one sweep
-    for (int i=0; i < nefc; i++) {
+    for (int c=0; c < nefc; c++) {
+      int i = efclist ? efclist[c] : c;
+
       // get constraint dimensionality
       int dim;
       if (d->efc_type[i] == mjCNSTR_CONTACT_ELLIPTIC) {
@@ -361,16 +434,16 @@ void mj_solPGS(const mjModel* m, mjData* d, int maxiter) {
       // simple constraint
       if (d->efc_type[i] != mjCNSTR_CONTACT_ELLIPTIC) {
         // unconstrained minimum
-        force[i] -= res[0]*ARinv[i];
+        force[i] -= res[0]*ARinv[c];
 
         // impose interval and inequality constraints
-        if (i >= ne && i < ne+nf) {
+        if (c >= ne && c < ne+nf) {
           if (force[i] < -floss[i]) {
             force[i] = -floss[i];
           } else if (force[i] > floss[i]) {
             force[i] = floss[i];
           }
-        } else if (i >= ne+nf) {
+        } else if (c >= ne+nf) {
           if (force[i] < 0) {
             force[i] = 0;
           }
@@ -380,7 +453,7 @@ void mj_solPGS(const mjModel* m, mjData* d, int maxiter) {
       // elliptic cone constraint
       else {
         // get friction
-        mjtNum *mu =  d->contact[d->efc_id[i]].friction;
+        mjtNum *mu = d->contact[d->efc_id[i]].friction;
 
         //-------------------- perform normal or ray update
 
@@ -390,7 +463,7 @@ void mj_solPGS(const mjModel* m, mjData* d, int maxiter) {
         // normal force too small: normal update
         if (force[i] < mjMINVAL) {
           // unconstrained minimum
-          force[i] -= res[0]*ARinv[i];
+          force[i] -= res[0]*ARinv[c];
 
           // clamp
           if (force[i] < 0) {
@@ -447,61 +520,31 @@ void mj_solPGS(const mjModel* m, mjData* d, int maxiter) {
 
         // QCQP
         else {
-          int flg_active;
+          solveQCQP(force, i, dim, Ac, bc, mu);
-          mjtNum v[6];
-
-          // solve
-          if (dim == 3) {
-            flg_active = mju_QCQP2(v, Ac, bc, mu, force[i]);
-          } else if (dim == 4) {
-            flg_active = mju_QCQP3(v, Ac, bc, mu, force[i]);
-          } else {
-            flg_active = mju_QCQP(v, Ac, bc, mu, force[i], dim-1);
-          }
-
-          // on constraint: put v on ellipsoid, in case QCQP is approximate
-          if (flg_active) {
-            mjtNum s = 0;
-            for (int j=0; j < dim-1; j++) {
-              s += v[j]*v[j] / (mu[j]*mu[j]);
-            }
-            s = mju_sqrt(force[i]*force[i] / mju_max(mjMINVAL, s));
-            for (int j=0; j < dim-1; j++) {
-              v[j] *= s;
-            }
-          }
-
-          // assign
-          mju_copy(force+i+1, v, dim-1);
         }
       }
 
       // accumulate improvement
       if (dim == 1) {
-        Athis[0] = 1/ARinv[i];
+        Athis[0] = 1/ARinv[c];
       }
       improvement -= costChange(Athis, force+i, oldforce, res, dim);
 
       // skip the rest of this constraint
-      i += (dim-1);
+      c += (dim-1);
     }
 
-    // process state
+    // update constraint state
-    mju_copyInt(oldstate, d->efc_state, nefc);
+    int nchange;
-    int nactive = dualState(m, d, d->efc_state);
+    int nactive = dualStateChange(d, d->efc_state, oldstate, ne, nf, nefc, efclist, &nchange);
-    int nchange = 0;
-    for (int i=0; i < nefc; i++) {
-      nchange += (oldstate[i] != d->efc_state[i]);
-    }
 
     // scale improvement, save stats
     improvement *= scale;
-    saveStats(m, d, island, iter, improvement, 0, 0, nactive, nchange, 0, 0);
+    saveStats(m, d, island_stat, iter, improvement, 0, 0, nactive, nchange, 0, 0);
 
     // increment iteration count
     iter++;
 
-
     // terminate
     if (improvement < m->opt.tolerance) {
       break;
@@ -509,51 +552,69 @@ void mj_solPGS(const mjModel* m, mjData* d, int maxiter) {
   }
 
   // finalize statistics
-  if (island < mjNISLAND) {
+  if (island_stat < mjNISLAND) {
     // update solver iterations
-    d->solver_niter[island] += iter;
+    d->solver_niter[island_stat] += iter;
 
     // set nnz
     if (mj_isSparse(m)) {
-      d->solver_nnz[island] = 0;
+      d->solver_nnz[island_stat] = 0;
-      for (int i=0; i < nefc; i++) {
+      for (int c=0; c < nefc; c++) {
-        d->solver_nnz[island] += d->efc_AR_rownnz[i];
+        d->solver_nnz[island_stat] += d->efc_AR_rownnz[efclist ? efclist[c] : c];
       }
     } else {
-      d->solver_nnz[island] = nefc*nefc;
+      d->solver_nnz[island_stat] = nefc*nefc;
     }
   }
 
-  // map to joint space
-  dualFinish(m, d);
-
   mj_freeStack(d);
 }
 
 
+// PGS entry point (monolithic, no dualFinish — caller handles it)
+void mj_solPGS(const mjModel* m, mjData* d, int maxiter) {
+  solPGS(m, d, /*island=*/-1, d->ne, d->nf, d->nefc, /*efclist=*/NULL, maxiter);
+}
+
+
+// PGS entry point (one island)
+void mj_solPGS_island(const mjModel* m, mjData* d, int island, int maxiter) {
+  int ne = d->island_ne[island];
+  int nf = d->island_nf[island];
+  int nefc = d->island_nefc[island];
+  int iefcadr = d->island_iefcadr[island];
+
+  solPGS(m, d, island, ne, nf, nefc, d->map_iefc2efc + iefcadr, maxiter);
+}
+
+
 //---------------------------- NoSlip solver -------------------------------------------------------
 
-// TODO: b/295296178 - add island support to Dual solvers
+// core NoSlip solver: iterates over constraints specified by efclist
-void mj_solNoSlip(const mjModel* m, mjData* d, int maxiter) {
+//   island: island index for stats (use -1 for monolithic, mapped to 0)
-  int dim, iter = 0, ne = d->ne, nf = d->nf, nefc = d->nefc;
+//   ne, nf, nefc: constraint type counts
+//   efclist: maps list position c to monolithic efc index (NULL for sequential)
+static void solNoSlip(const mjModel* m, mjData* d, int island,
+                      int ne, int nf, int nefc,
+                      const int* efclist, int maxiter) {
+  int dim, iter = 0;
   const mjtNum *floss = d->efc_frictionloss;
   mjtNum *force = d->efc_force;
   mjtNum *mu, improvement;
-  mjtNum v[5], Ac[25], bc[5], res[5], oldforce[5], delta[5], mid, y, K0, K1;
+  mjtNum Ac[25], bc[5], res[5], oldforce[5], delta[5], mid, y, K0, K1;
   mjContact* con;
   mj_markStack(d);
   mjtNum* ARinv = mjSTACKALLOC(d, nefc, mjtNum);
   int* oldstate = mjSTACKALLOC(d, nefc, int);
 
-  // TODO: b/295296178 - Use island index (currently hardcoded to 0)
+  int island_stat = mjMAX(0, island);
-  int island = 0;
   mjtNum scale = 1 / (m->stat.meaninertia * mjMAX(1, m->nv));
 
   // precompute inverse diagonal of A
-  ARdiaginv(m, d, ARinv, 1);
+  ARdiaginv(m, d, ARinv, nefc, efclist, 1);
 
   // initial constraint state
-  dualState(m, d, d->efc_state);
+  dualState(d, d->efc_state, ne, nf, nefc, efclist);
 
   // main iteration
   while (iter < maxiter) {
@@ -562,19 +623,22 @@ void mj_solNoSlip(const mjModel* m, mjData* d, int maxiter) {
 
     // correct for cost change at iter 0
     if (iter == 0) {
-      for (int i=0; i < nefc; i++) {
+      for (int c=0; c < nefc; c++) {
+        int i = efclist ? efclist[c] : c;
         improvement += 0.5*force[i]*force[i]*d->efc_R[i];
       }
     }
 
     // perform one sweep: dry friction
-    for (int i=ne; i < ne+nf; i++) {
+    for (int c=ne; c < ne+nf; c++) {
+      int i = efclist ? efclist[c] : c;
+
       // compute residual, save old
       residual(m, d, res, i, 1, 1);
       oldforce[0] = force[i];
 
       // unconstrained minimum
-      force[i] -= res[0]*ARinv[i];
+      force[i] -= res[0]*ARinv[c];
 
       // impose interval constraints
       if (force[i] < -floss[i]) {
@@ -585,11 +649,13 @@ void mj_solNoSlip(const mjModel* m, mjData* d, int maxiter) {
 
       // add to improvement
       delta[0] = force[i] - oldforce[0];
-      improvement -= 0.5*delta[0]*delta[0]/ARinv[i] + delta[0]*res[0];
+      improvement -= 0.5*delta[0]*delta[0]/ARinv[c] + delta[0]*res[0];
     }
 
     // perform one sweep: contact friction
-    for (int i=ne+nf; i < nefc; i++) {
+    for (int c=ne+nf; c < nefc; c++) {
+      int i = efclist ? efclist[c] : c;
+
       // pyramidal contact
       if (d->efc_type[i] == mjCNSTR_CONTACT_PYRAMIDAL) {
         // get contact info
@@ -648,7 +714,7 @@ void mj_solNoSlip(const mjModel* m, mjData* d, int maxiter) {
         }
 
         // skip the rest of this contact
-        i += 2*(dim-1)-1;
+        c += 2*(dim-1)-1;
       }
 
       // elliptic contact
@@ -678,55 +744,29 @@ void mj_solNoSlip(const mjModel* m, mjData* d, int maxiter) {
 
         // QCQP
         else {
-          int flg_active = 0;
+          solveQCQP(force, i, dim, Ac, bc, mu);
-
-          // solve
-          if (dim == 3) {
-            flg_active = mju_QCQP2(v, Ac, bc, mu, force[i]);
-          } else if (dim == 4) {
-            flg_active = mju_QCQP3(v, Ac, bc, mu, force[i]);
-          } else {
-            flg_active = mju_QCQP(v, Ac, bc, mu, force[i], dim-1);
-          }
-
-          // on constraint: put v on ellipsoid, in case QCQP is approximate
-          if (flg_active) {
-            mjtNum s = 0;
-            for (int j=0; j < dim-1; j++) {
-              s += v[j]*v[j]/(mu[j]*mu[j]);
-            }
-            s = mju_sqrt(force[i]*force[i] / mju_max(mjMINVAL, s));
-            for (int j=0; j < dim-1; j++) {
-              v[j] *= s;
-            }
-          }
-
-          // assign
-          mju_copy(force+i+1, v, dim-1);
         }
 
         // accumulate improvement
         improvement -= costChange(Ac, force+i+1, oldforce, res, dim-1);
 
         // skip the rest of this contact
-        i += (dim-1);
+        c += (dim-1);
       }
     }
 
-    // process state
+    // update constraint state
-    mju_copyInt(oldstate, d->efc_state, nefc);
+    int nchange;
-    int nactive = dualState(m, d, d->efc_state);
+    int nactive = dualStateChange(d, d->efc_state, oldstate, ne, nf, nefc, efclist, &nchange);
-    int nchange = 0;
-    for (int i=0; i < nefc; i++) {
-      nchange += (oldstate[i] != d->efc_state[i]);
-    }
 
     // scale improvement, save stats
     improvement *= scale;
 
     // save noslip stats after all the entries from regular solver
-    int stats_iter = iter + d->solver_niter[island];
+    if (island_stat < mjNISLAND) {
-    saveStats(m, d, island, stats_iter, improvement, 0, 0, nactive, nchange, 0, 0);
+      int stats_iter = iter + d->solver_niter[island_stat];
+      saveStats(m, d, island_stat, stats_iter, improvement, 0, 0, nactive, nchange, 0, 0);
+    }
 
     // increment iteration count
     iter++;
@@ -738,15 +778,31 @@ void mj_solNoSlip(const mjModel* m, mjData* d, int maxiter) {
   }
 
   // update solver iterations
-  d->solver_niter[island] += iter;
+  if (island_stat < mjNISLAND) {
-
+    d->solver_niter[island_stat] += iter;
-  // map to joint space
+  }
-  dualFinish(m, d);
 
   mj_freeStack(d);
 }
 
 
+// NoSlip entry point (monolithic, no dualFinish — caller handles it)
+void mj_solNoSlip(const mjModel* m, mjData* d, int maxiter) {
+  solNoSlip(m, d, /*island=*/-1, d->ne, d->nf, d->nefc, /*efclist=*/NULL, maxiter);
+}
+
+
+// NoSlip entry point (one island)
+void mj_solNoSlip_island(const mjModel* m, mjData* d, int island, int maxiter) {
+  int ne = d->island_ne[island];
+  int nf = d->island_nf[island];
+  int nefc = d->island_nefc[island];
+  int iefcadr = d->island_iefcadr[island];
+
+  solNoSlip(m, d, island, ne, nf, nefc, d->map_iefc2efc + iefcadr, maxiter);
+}
+
+
 //------------------------- Primal solvers ---------------------------------------------------------
 
 // Primal context

src/engine/engine_solver.h

@@ -35,10 +35,19 @@ void mj_solNewton(const mjModel* m, mjData* d, int maxiter);
 
 //------------------------------ per-island solvers ------------------------------------------------
 
+// PGS solver (one island, no dualFinish — caller handles it)
+void mj_solPGS_island(const mjModel* m, mjData* d, int island, int maxiter);
+
+// NoSlip solver (one island, no dualFinish — caller handles it)
+void mj_solNoSlip_island(const mjModel* m, mjData* d, int island, int maxiter);
+
 // CG solver
 void mj_solCG_island(const mjModel* m, mjData* d, int island, int maxiter);
 
 // Newton entry point
 void mj_solNewton_island(const mjModel* m, mjData* d, int island, int maxiter);
 
+// map efc_force to joint space (used after dual island dispatch)
+void mj_dualFinish(const mjModel* m, mjData* d);
+
 #endif  // MUJOCO_SRC_ENGINE_ENGINE_SOLVER_H_

src/experimental/filament/CMakeLists.txt

@@ -27,7 +27,6 @@ target_sources(${MUJOCO_FILAMENT_TARGET_NAME}
         filament/builtins.h
         filament/color_grading_options.cc
         filament/color_grading_options.h
-        filament/draw_mode.h
         filament/filament_context.cc
         filament/filament_context.h
         filament/filament_platform_factory.cc

src/experimental/filament/compat/imgui_bridge.cc

@@ -72,7 +72,7 @@ uintptr_t ImguiBridge::UploadImage(uintptr_t tex_id, const uint8_t* pixels,
 
   // Assign a new texture ID.
   if (tex_id == 0) {
-    tex_id = textures_.size() + 1;
+    tex_id = next_tex_id_++;
   }
 
   std::unique_ptr<Texture>& texture = textures_[tex_id];
@@ -123,7 +123,7 @@ void ImguiBridge::CreateTexture(ImTextureData* data) {
   config.format = mjPIXEL_FORMAT_RGBA8;
   config.color_space = mjCOLORSPACE_LINEAR;
 
-  const uintptr_t tex_id = textures_.size() + 1;
+  const uintptr_t tex_id = next_tex_id_++;
   textures_[tex_id] =
       std::make_unique<Texture>(scene_view_->GetEngine(), config);
   data->SetTexID((ImTextureID)tex_id);

src/experimental/filament/compat/imgui_bridge.h

@@ -64,6 +64,7 @@ class ImguiBridge {
   std::vector<std::unique_ptr<Renderable>> renderables_;
   std::vector<std::unique_ptr<Mesh>> meshes_;
   std::unordered_map<uintptr_t, std::unique_ptr<Texture>> textures_;
+  uintptr_t next_tex_id_ = 1;
 };
 
 // Draws text at the given screen coordinates in clip space (i.e. [-1,-1,-1] to

src/experimental/filament/compat/mjr_filament_renderer.cc

@@ -25,7 +25,6 @@
 #include "experimental/filament/compat/imgui_bridge.h"
 #include "experimental/filament/compat/imgui_editor.h"
 #include "experimental/filament/compat/scene_bridge.h"
-#include "experimental/filament/filament/draw_mode.h"
 #include "experimental/filament/filament/filament_context.h"
 #include "experimental/filament/filament/model_util.h"
 #include "experimental/filament/filament/render_target.h"
@@ -42,11 +41,14 @@ void MjrFilamentRenderer::Init(const mjModel* model) {
   scene_bridge_ = std::make_unique<SceneBridge>(GetObjectManager(), model);
   imgui_bridge_ = std::make_unique<ImguiBridge>(GetObjectManager());
 
+  mjr_defaultRenderRequest(&render_requests_[0]);
+  mjr_defaultRenderRequest(&render_requests_[1]);
+
   render_requests_[0].scene = scene_bridge_->GetSceneView();
-  render_requests_[0].draw_mode = DrawMode::Color;
+  render_requests_[0].draw_mode = mjDRAW_MODE_COLOR;
 
   render_requests_[1].scene = imgui_bridge_->GetSceneView();
-  render_requests_[1].draw_mode = DrawMode::Color;
+  render_requests_[1].draw_mode = mjDRAW_MODE_COLOR;
 
   // The UX camera is a fixed orthographic camera. We only need to change the
   // width/height based on the viewport per frame.
@@ -78,11 +80,11 @@ void MjrFilamentRenderer::Render(const mjrRect& viewport, const mjvScene* scene)
   }
 
   if (scene->flags[mjRND_SEGMENT]) {
-    render_requests_[0].draw_mode  = DrawMode::Segmentation;
+    render_requests_[0].draw_mode = mjDRAW_MODE_SEGMENTATION;
   } else if (scene->flags[mjRND_DEPTH]) {
-    render_requests_[0].draw_mode  = DrawMode::Depth;
+    render_requests_[0].draw_mode = mjDRAW_MODE_DEPTH;
   } else {
-    render_requests_[0].draw_mode  = DrawMode::Color;
+    render_requests_[0].draw_mode = mjDRAW_MODE_COLOR;
   }
 
   render_requests_[0].width = viewport.width;
@@ -142,10 +144,11 @@ void MjrFilamentRenderer::ReadPixels(mjrRect viewport, unsigned char* rgb,
     const size_t num_requests =
         (mode_ == FrameBufferMode::OffScreenWithGui) ? 2 : 1;
 
-    ReadPixelsRequest read_request;
+    mjrReadPixelsRequest read_request;
+    mjr_defaultReadPixelsRequest(&read_request);
     read_request.output = rgb;
     read_request.num_bytes = viewport.width * viewport.height * 3;
-    const FrameHandle frame = FilamentContext::Render(
+    const mjrFrameHandle frame = FilamentContext::Render(
         {&render_requests_[0], num_requests}, {&read_request, 1});
     FilamentContext::WaitForFrame(frame);
 
@@ -163,13 +166,14 @@ void MjrFilamentRenderer::ReadPixels(mjrRect viewport, unsigned char* rgb,
     render_requests_[0].target = target.get();
     render_requests_[1].target = target.get();
 
-    DrawMode last_draw_mode = render_requests_[0].draw_mode;
+    mjrDrawMode last_draw_mode = render_requests_[0].draw_mode;
-    render_requests_[0].draw_mode = DrawMode::Depth;
+    render_requests_[0].draw_mode = mjDRAW_MODE_DEPTH;
 
-    ReadPixelsRequest read_request;
+    mjrReadPixelsRequest read_request;
+    mjr_defaultReadPixelsRequest(&read_request);
     read_request.output = reinterpret_cast<uint8_t*>(depth);
     read_request.num_bytes = viewport.width * viewport.height * sizeof(float);
-    const FrameHandle frame =
+    const mjrFrameHandle frame =
         FilamentContext::Render({&render_requests_[0], 1}, {&read_request, 1});
     FilamentContext::WaitForFrame(frame);
 

src/experimental/filament/compat/mjr_filament_renderer.h

@@ -76,7 +76,7 @@ class MjrFilamentRenderer : public FilamentContext {
   };
 
   FrameBufferMode mode_ = FrameBufferMode::Window;
-  RenderRequest render_requests_[2];
+  mjrRenderRequest render_requests_[2];
   std::unique_ptr<SceneBridge> scene_bridge_;
   std::unique_ptr<ImguiBridge> imgui_bridge_;
 };

src/experimental/filament/compat/model_objects.cc

@@ -26,12 +26,7 @@
 #include <vector>
 
 #include <filament/Engine.h>
-#include <filament/IndirectLight.h>
-#include <filament/Material.h>
-#include <filament/Skybox.h>
 #include <math/TVecHelpers.h>
-#include <math/mat3.h>
-#include <math/scalar.h>
 #include <math/vec2.h>
 #include <math/vec3.h>
 #include <math/vec4.h>
@@ -47,7 +42,6 @@ namespace mujoco {
 using filament::math::float2;
 using filament::math::float3;
 using filament::math::float4;
-using filament::math::mat3f;
 
 enum class MeshType {
   kNormal,
@@ -534,12 +528,6 @@ ModelObjects::ModelObjects(const mjModel* model, filament::Engine* engine)
 }
 
 ModelObjects::~ModelObjects() {
-  for (auto& iter : skyboxes_) {
-    engine_->destroy(iter);
-  }
-  for (auto& iter : indirect_lights_) {
-    engine_->destroy(iter);
-  }
   meshes_.clear();
   textures_.clear();
 }
@@ -681,48 +669,13 @@ const Texture* ModelObjects::GetTexture(int mat_id, int role) const {
   return GetTexture(tex_id);
 }
 
-filament::IndirectLight* ModelObjects::CreateIndirectLight(int tex_id,
+const Texture* ModelObjects::GetSkyboxTexture() const {
-                                                           float intensity) {
-  filament::Texture* texture = nullptr;
-  const Texture::SphericalHarmonics* spherical_harmonics = nullptr;
-  auto texture_iter = textures_.find(tex_id);
-  if (texture_iter != textures_.end()) {
-    texture = texture_iter->second->GetFilamentTexture();
-    spherical_harmonics = texture_iter->second->GetSphericalHarmonics();
-  }
-
-  filament::IndirectLight::Builder builder;
-  builder.reflections(texture);
-  if (spherical_harmonics != nullptr) {
-    builder.irradiance(3, *spherical_harmonics);
-  }
-  builder.intensity(intensity);
-  // Rotate the light to match mujoco's Z-up convention.
-  builder.rotation(mat3f::rotation(filament::math::f::PI / 2, float3{1, 0, 0}));
-  filament::IndirectLight* indirect_light = builder.build(*engine_);
-  indirect_lights_.push_back(indirect_light);
-  return indirect_light;
-}
-
-filament::Skybox* ModelObjects::CreateSkybox() {
-  filament::Texture* skybox_texture = nullptr;
   for (auto& iter : textures_) {
-    const int texture_type = model_->tex_type[iter.first];
+    if (model_->tex_type[iter.first] == mjTEXTURE_SKYBOX) {
-    if (texture_type == mjTEXTURE_SKYBOX) {
+      return iter.second.get();
-      skybox_texture = iter.second->GetFilamentTexture();
-      break;
     }
   }
-
+  return nullptr;
-  if (skybox_texture == nullptr) {
-    return nullptr;
-  }
-
-  filament::Skybox::Builder builder;
-  builder.environment(skybox_texture);
-  filament::Skybox* skybox = builder.build(*engine_);
-  skyboxes_.push_back(skybox);
-  return skybox;
 }
 
 }  // namespace mujoco

src/experimental/filament/compat/model_objects.h

@@ -18,11 +18,8 @@
 #include <array>
 #include <memory>
 #include <unordered_map>
-#include <vector>
 
 #include <filament/Engine.h>
-#include <filament/IndirectLight.h>
-#include <filament/Skybox.h>
 #include <mujoco/mjmodel.h>
 #include <mujoco/mujoco.h>
 #include "experimental/filament/filament/mesh.h"
@@ -30,7 +27,7 @@
 
 namespace mujoco {
 
-// Creates and owns various filament objects based on the data in a mjrContext.
+// Creates and owns various filament objects based on the mjModel.
 class ModelObjects {
  public:
   ModelObjects(const mjModel* model, filament::Engine* engine);
@@ -58,10 +55,6 @@ class ModelObjects {
 
   void CreateSkinFlexMesh(const mjvScene* scene, const mjvGeom& geom);
 
-  // Returns the filament engine used by the ModelObjects to create filament
-  // objects.
-  filament::Engine* GetEngine() const { return engine_; }
-
   // Returns the cached instance of a filament object created from the mjModel.
   const Mesh* GetShapeBuffer(ShapeType shape) const;
   const Mesh* GetMeshBuffer(int data_id) const;
@@ -69,9 +62,7 @@ class ModelObjects {
   const Mesh* GetFlexSkinGeomMesh(int geom_id) const;
   const Texture* GetTexture(int tex_id) const;
   const Texture* GetTexture(int mat_id, int role) const;
-
+  const Texture* GetSkyboxTexture() const;
-  filament::Skybox* CreateSkybox();
-  filament::IndirectLight* CreateIndirectLight(int tex_id, float intensity);
 
   float GetSpecularMultiplier() const { return specular_multiplier_; }
   float GetShininessMultiplier() const { return shininess_multiplier_; }
@@ -85,8 +76,6 @@ class ModelObjects {
  private:
   const mjModel* model_ = nullptr;
   filament::Engine* engine_ = nullptr;
-  std::vector<filament::Skybox*> skyboxes_;
-  std::vector<filament::IndirectLight*> indirect_lights_;
   std::array<std::unique_ptr<Mesh>, kNumShapes> shapes_;
   std::unordered_map<int, std::unique_ptr<Mesh>> meshes_;
   std::unordered_map<int, std::unique_ptr<Mesh>> convex_hulls_;

src/experimental/filament/compat/scene_bridge.cc

@@ -322,10 +322,7 @@ void SceneBridge::PrepareLights() {
     }
   }
 
-  filament::Skybox* skybox = model_objects_->CreateSkybox();
+  scene_view_->SetSkybox(model_objects_->GetSkyboxTexture());
-  if (skybox) {
-    scene_view_->AddToScene(skybox);
-  }
 }
 
 filament::math::mat4 CalculateClipFromWorld(const mjrRect& viewport,

src/experimental/filament/compat/scene_geom_util.cc

@@ -30,12 +30,12 @@
 #include <mujoco/mjvisualize.h>
 #include <mujoco/mujoco.h>
 #include "experimental/filament/compat/model_objects.h"
-#include "experimental/filament/filament/material.h"
 #include "experimental/filament/filament/math_util.h"
 #include "experimental/filament/filament/mesh.h"
 #include "experimental/filament/filament/object_manager.h"
 #include "experimental/filament/filament/renderable.h"
 #include "experimental/filament/filament/texture.h"
+#include "experimental/filament/render_context_filament.h"
 
 namespace mujoco {
 
@@ -424,7 +424,7 @@ static void UpdateGeomMaterial(Renderable& renderable, const mjvGeom& geom,
   // the programmatic UVs.
 
   if (textures.color) {
-    if (textures.color->GetFilamentTexture()->getTarget() ==
+    if (Texture::downcast(textures.color)->GetFilamentTexture()->getTarget() ==
         filament::Texture::Sampler::SAMPLER_2D) {
       // For 2D textures, `tex_repeat` specifies how many times the texture
       // image is repeated. The `tex_uniform` flag determines if the repetition

src/experimental/filament/filament/draw_mode.h

@@ -1,36 +0,0 @@
-// 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_DRAW_MODE_H_
-#define MUJOCO_SRC_EXPERIMENTAL_FILAMENT_FILAMENT_DRAW_MODE_H_
-
-namespace mujoco {
-
-// The different modes that can be used to render the scene.
-enum class DrawMode {
-  // Render the scene with "normal" colors and lighting.
-  Color,
-  // Render the scene as a grayscale depth map.
-  Depth,
-  // Render each object with a unique, uniform (flat) color regardless of
-  // lighting and texture.
-  Segmentation,
-};
-
-static constexpr int kNumDrawModes = 3;
-
-}  // namespace mujoco
-
-
-#endif  // MUJOCO_SRC_EXPERIMENTAL_FILAMENT_FILAMENT_DRAW_MODE_H_

src/experimental/filament/filament/filament_context.cc

@@ -79,16 +79,16 @@ FilamentContext::~FilamentContext() {
   filament::Engine::destroy(engine_);
 }
 
-FilamentContext::FrameHandle FilamentContext::Render(
+mjrFrameHandle FilamentContext::Render(
-    std::span<const RenderRequest> requests,
+    std::span<const mjrRenderRequest> requests,
-    std::span<const ReadPixelsRequest> read_requests) {
+    std::span<const mjrReadPixelsRequest> read_requests) {
   if (read_requests.size() > 1) {
     mju_error("Only one read request is supported for now.");
   }
 
   bool render_began = false;
-  RenderTarget* current_target = nullptr;
+  mjrRenderTarget* current_target = nullptr;
-  for (const RenderRequest& request : requests) {
+  for (const mjrRenderRequest& request : requests) {
     if (request.target != current_target && render_began) {
       renderer_->endFrame();
       render_began = false;
@@ -126,7 +126,8 @@ FilamentContext::FrameHandle FilamentContext::Render(
         scene_view_request.draw_mode = request.draw_mode;
         scene_view_request.viewport = {0, 0, request.width, request.height};
         scene_view_request.camera = request.camera;
-        request.scene->Render(renderer_, scene_view_request);
+        SceneView* scene_view = SceneView::downcast(request.scene);
+        scene_view->Render(renderer_, scene_view_request);
       }
     } else {
       if (read_requests.empty()) {
@@ -134,7 +135,7 @@ FilamentContext::FrameHandle FilamentContext::Render(
             "Rendering to a render target without a read request is pointless.");
       }
 
-      const ReadPixelsRequest& read_request = read_requests[0];
+      const mjrReadPixelsRequest& read_request = read_requests[0];
       if (read_request.num_bytes == 0) {
         mju_error("Output buffer size is zero.");
       }
@@ -146,13 +147,16 @@ FilamentContext::FrameHandle FilamentContext::Render(
         break;
       }
       if (render_began) {
+        RenderTarget* render_target = RenderTarget::downcast(request.target);
+
         SceneView::RenderRequest scene_view_request;
         scene_view_request.draw_mode = request.draw_mode;
         scene_view_request.viewport = {0, 0, request.width, request.height};
         scene_view_request.camera = request.camera;
-        scene_view_request.target = request.target;
+        scene_view_request.target = render_target;
-        request.scene->Render(renderer_, scene_view_request);
+        SceneView* scene_view = SceneView::downcast(request.scene);
-        request.target->ReadColorPixels(renderer_, read_request.output,
+        scene_view->Render(renderer_, scene_view_request);
+        render_target->ReadColorPixels(renderer_, (uint8_t*)read_request.output,
                                         read_request.num_bytes);
       }
     }
@@ -175,7 +179,7 @@ FilamentContext::FrameHandle FilamentContext::Render(
   return ++frame_counter_;
 }
 
-void FilamentContext::WaitForFrame(FrameHandle frame_handle) {
+void FilamentContext::WaitForFrame(mjrFrameHandle frame_handle) {
   if (frame_counter_ < frame_handle) {
     engine_->flushAndWait();
   }

src/experimental/filament/filament/filament_context.h

@@ -15,7 +15,6 @@
 #ifndef MUJOCO_SRC_EXPERIMENTAL_FILAMENT_FILAMENT_FILAMENT_CONTEXT_H_
 #define MUJOCO_SRC_EXPERIMENTAL_FILAMENT_FILAMENT_FILAMENT_CONTEXT_H_
 
-#include <cstddef>
 #include <cstdint>
 #include <span>
 #include <memory>
@@ -25,77 +24,30 @@
 #include <filament/Renderer.h>
 #include <filament/SwapChain.h>
 #include <math/vec4.h>
-#include <mujoco/mjvisualize.h>
-#include "experimental/filament/filament/draw_mode.h"
-#include "experimental/filament/filament/scene_view.h"
 #include "experimental/filament/filament/object_manager.h"
-#include "experimental/filament/filament/render_target.h"
 #include "experimental/filament/render_context_filament.h"
 
 namespace mujoco {
 
 // Manages the filament renderer and provides APIs for rendering scenes.
-class FilamentContext {
+class FilamentContext : public mjrfContext {
  public:
   explicit FilamentContext(const mjrFilamentConfig* config);
   ~FilamentContext();
 
-  // Information needed to render a single image of a scene.
+  FilamentContext(const FilamentContext&) = delete;
-  struct RenderRequest {
+  FilamentContext& operator=(const FilamentContext&) = delete;
-    // The scene to render.
-    SceneView* scene = nullptr;
-
-    // The method (e.g. Color, Depth, Segmentation, etc.) to use for rendering.
-    DrawMode draw_mode = DrawMode::Color;
-
-    // The camera from which to render the scene.
-    mjvGLCamera camera;
-
-    // The dimensions of the output image.
-    int width = 0;
-    int height = 0;
-
-    // The render target into which to render the image. If nullptr, the image
-    // will be rendered to the window (as previously configured in
-    // mjrFilamentConfig::native_window).
-    RenderTarget* target = nullptr;
-  };
-
-  // Information needed to read pixels from a render target.
-  struct ReadPixelsRequest {
-    RenderTarget* target = nullptr;
-
-    // The buffer into which the read pixels will be written.
-    uint8_t* output = nullptr;
-
-    // The number of bytes in the output buffer. This should match the size of
-    // the render target texture.
-    std::size_t num_bytes = 0;
-
-    // Callback when the read pixels operation is complete. This will be called
-    // during WaitForFrame() or in a subsequent call to Render(). This function
-    // can optionally be used to free the output buffer if needed.
-    void (*read_completed_callback)(void* user_data) = nullptr;
-
-    // User data to pass to the completion callback.
-    void* user_data = nullptr;
-  };
-
-  // Rendering is asynchronous by nature. Each render request is assigned a
-  // unique Handle which can be used to query the status of the request. The
-  // Handle can also be used to block until the request is completed.
-  using FrameHandle = std::uint64_t;
 
   // Queues the given render requests for rendering. This function copies the
   // necessary data from the requests into the renderer thread and returns
   // immediately afterwards. The renderer thread will then perform the actual
   // rendering on the GPU. Callers can use WaitForFrame to block until the
   // rendering is complete.
-  FrameHandle Render(std::span<const RenderRequest> render_requests,
+  mjrFrameHandle Render(std::span<const mjrRenderRequest> render_requests,
-                     std::span<const ReadPixelsRequest> read_requests = {});
+                     std::span<const mjrReadPixelsRequest> read_requests = {});
 
   // Blocks until the given frame has completed rendering.
-  void WaitForFrame(FrameHandle frame_handle);
+  void WaitForFrame(mjrFrameHandle frame_handle);
 
   // Sets the clear color for the renderer.
   void SetClearColor(const filament::math::float4& color);
@@ -107,8 +59,12 @@ class FilamentContext {
 
   ObjectManager* GetObjectManager() const { return object_manager_.get(); }
 
-  FilamentContext(const FilamentContext&) = delete;
+  static FilamentContext* downcast(mjrfContext* context) {
-  FilamentContext& operator=(const FilamentContext&) = delete;
+    return static_cast<FilamentContext*>(context);
+  }
+  static const FilamentContext* downcast(const mjrfContext* context) {
+    return static_cast<const FilamentContext*>(context);
+  }
 
  private:
   mjrFilamentConfig config_;

src/experimental/filament/filament/light.cc

@@ -27,27 +27,13 @@
 #include <mujoco/mujoco.h>
 #include "experimental/filament/filament/math_util.h"
 #include "experimental/filament/filament/texture.h"
+#include "experimental/filament/render_context_filament.h"
 
 namespace mujoco {
 
 using filament::math::float3;
 using filament::math::mat3f;
 
-void mjr_defaultLightParams(mjrLightParams* params) {
-  params->type = mjLIGHT_POINT;
-  params->texture = nullptr;
-  params->color[0] = 0;
-  params->color[1] = 0;
-  params->color[2] = 0;
-  params->intensity = 0.0f;
-  params->cast_shadows = true;
-  params->range = 10.0f;
-  params->spot_cone_angle = 180.f;
-  params->bulb_radius = 0.0f;
-  params->shadow_map_size = 2048;
-  params->vsm_blur_width = 0.0f;
-}
-
 Light::Light(filament::Engine* engine, const mjrLightParams& params)
     : engine_(engine), params_(params) {
   // Filament treats image-based lights (IBLs) as separate objects (i.e.
@@ -56,9 +42,10 @@ Light::Light(filament::Engine* engine, const mjrLightParams& params)
     filament::IndirectLight::Builder builder;
     if (params.texture) {
       // Allow null textures for fallback lights.
-      builder.reflections(params.texture->GetFilamentTexture());
+      const Texture* texture = Texture::downcast(params.texture);
+      builder.reflections(texture->GetFilamentTexture());
       const Texture::SphericalHarmonics* spherical_harmonics =
-          params.texture->GetSphericalHarmonics();
+          texture->GetSphericalHarmonics();
       if (spherical_harmonics != nullptr) {
         builder.irradiance(3, *spherical_harmonics);
       }

src/experimental/filament/filament/light.h

@@ -20,40 +20,12 @@
 #include <math/vec3.h>
 #include <utils/Entity.h>
 #include <mujoco/mujoco.h>
-#include "experimental/filament/filament/texture.h"
+#include "experimental/filament/render_context_filament.h"
 
 namespace mujoco {
 
-typedef mjtLightType mjrLightType;
-
-// Configuration parameters for a light.
-struct mjrLightParams {
-  // The type of light (e.g. spot, point, directional, etc.)
-  mjrLightType type;
-  // The texture to use for image lights.
-  const Texture* texture;
-  // The color of the light.
-  float color[3];
-  // The intensity of the light, in candela.
-  float intensity;
-  // Whether or not the light casts shadows.
-  mjtByte cast_shadows;
-  // The range/distance in which the light is effective, in meters.
-  float range;
-  // The angle of the spot light cone, in degrees.
-  float spot_cone_angle;
-  // The radius of the bulb used for soft shadows.
-  float bulb_radius;
-  // The size of the shadow map.
-  int shadow_map_size;
-  // Blur width for EL VSM.
-  float vsm_blur_width;
-};
-
-void mjr_defaultLightParams(mjrLightParams* params);
-
 // Manages the filament Entities for a single mjvLight.
-class Light {
+class Light : public mjrLight {
  public:
   Light(filament::Engine* engine, const mjrLightParams& params);
   ~Light() noexcept;
@@ -84,6 +56,13 @@ class Light {
   void Enable();
   void Disable();
 
+  static Light* downcast(mjrLight* light) {
+    return static_cast<Light*>(light);
+  }
+  static const Light* downcast(const mjrLight* light) {
+    return static_cast<const Light*>(light);
+  }
+
  private:
   filament::Engine* engine_ = nullptr;
   filament::IndirectLight* ibl_ = nullptr;

src/experimental/filament/filament/material.cc

@@ -14,8 +14,6 @@
 
 #include "experimental/filament/filament/material.h"
 
-#include <array>
-
 #include <filament/Color.h>
 #include <filament/Material.h>
 #include <filament/MaterialInstance.h>
@@ -25,45 +23,10 @@
 #include "experimental/filament/filament/math_util.h"
 #include "experimental/filament/filament/texture.h"
 #include "experimental/filament/filament/object_manager.h"
+#include "experimental/filament/render_context_filament.h"
 
 namespace mujoco {
 
-template <int N>
-static void setf(float (&arr)[N], const std::array<float, N>& values) {
-  for (int i = 0; i < N; ++i) {
-    arr[i] = values[i];
-  }
-}
-
-void mjr_defaultMaterialTextures(mjrMaterialTextures* textures) {
-  textures->color = nullptr;
-  textures->normal = nullptr;
-  textures->metallic = nullptr;
-  textures->roughness = nullptr;
-  textures->occlusion = nullptr;
-  textures->orm = nullptr;
-  textures->emissive = nullptr;
-  textures->reflection = nullptr;
-}
-
-void mjr_defaultMaterialParams(mjrMaterialParams* params) {
-  setf(params->color, {1.f, 1.f, 1.f, 1.f});
-  setf(params->segmentation_color, {1, 1, 1, 1});
-  setf(params->uv_scale, {1, 1, 1});
-  setf(params->uv_offset, {0, 0, 0});
-  setf(params->scissor, {0, 0, 0, 0});
-
-  params->emissive = -1.0f;
-  params->specular = -1.0f;
-  params->glossiness = -1.0f;
-  params->metallic = -1.0f;
-  params->roughness = -1.0f;
-  params->reflectance = 0.0f;
-  params->tex_uniform = false;
-  params->reflective = false;
-}
-
-
 void UpdateMaterialInstance(filament::MaterialInstance* instance,
                             const mjrMaterialParams& params,
                             const mjrMaterialTextures& textures,
@@ -118,11 +81,12 @@ void UpdateMaterialInstance(filament::MaterialInstance* instance,
   sampler.setMinFilter(
       filament::TextureSampler::MinFilter::LINEAR_MIPMAP_LINEAR);
 
-  auto TrySetTexture = [&](const char* name, const Texture* texture,
+  auto TrySetTexture = [&](const char* name, const mjrTexture* texture,
                            mjtTextureRole role) {
     if (material->hasParameter(name)) {
       if (texture != nullptr) {
-        instance->setParameter(name, texture->GetFilamentTexture(), sampler);
+        instance->setParameter(
+            name, Texture::downcast(texture)->GetFilamentTexture(), sampler);
       } else {
         instance->setParameter(name, object_mgr->GetFallbackTexture(role),
                                sampler);

src/experimental/filament/filament/material.h

@@ -17,46 +17,11 @@
 
 #include <filament/Engine.h>
 #include <filament/MaterialInstance.h>
-#include <mujoco/mujoco.h>
-#include "experimental/filament/filament/texture.h"
 #include "experimental/filament/filament/object_manager.h"
+#include "experimental/filament/render_context_filament.h"
 
 namespace mujoco {
 
-// The textures that can be assigned to the drawable's material.
-struct mjrMaterialTextures {
-  const Texture* color;
-  const Texture* normal;
-  const Texture* metallic;
-  const Texture* roughness;
-  const Texture* occlusion;
-  const Texture* orm;
-  const Texture* emissive;
-  const Texture* reflection;
-};
-
-void mjr_defaultMaterialTextures(mjrMaterialTextures* textures);
-
-// The parameters that can be applied to the drawable's material.
-struct mjrMaterialParams {
-  float color[4];
-  float segmentation_color[4];
-  float tex_repeat[2];
-  float uv_scale[3];
-  float uv_offset[3];
-  float scissor[4];
-  float specular;
-  float glossiness;
-  float metallic;
-  float roughness;
-  float emissive;
-  float reflectance;
-  mjtByte tex_uniform;
-  mjtByte reflective;
-};
-
-void mjr_defaultMaterialParams(mjrMaterialParams* params);
-
 // Updates the material instances based on the currently set parameters and
 // textures.
 void UpdateMaterialInstance(filament::MaterialInstance* instance,

src/experimental/filament/filament/mesh.cc

@@ -33,6 +33,7 @@
 #include <math/vec4.h>
 #include <mujoco/mujoco.h>
 #include "experimental/filament/filament/math_util.h"
+#include "experimental/filament/render_context_filament.h"
 
 namespace mujoco {
 
@@ -101,11 +102,6 @@ int FillSequence(std::byte* buffer, std::size_t num_bytes) {
   return num;
 }
 
-// Initializes the mjrMeshData to default values.
-void mjr_defaultMeshData(mjrMeshData* data) {
-  std::memset(data, 0, sizeof(mjrMeshData));
-}
-
 Mesh::Mesh(filament::Engine* engine, const mjrMeshData& data)
     : engine_(engine), shared_state_(std::make_shared<SharedState>()) {
   type_ = data.primitive_type == mjMESH_PRIMITIVE_TYPE_TRIANGLES

src/experimental/filament/filament/mesh.h

@@ -16,7 +16,6 @@
 #define MUJOCO_SRC_EXPERIMENTAL_FILAMENT_FILAMENT_MESH_H_
 
 #include <array>
-#include <cstddef>
 #include <functional>
 #include <memory>
 #include <mutex>
@@ -30,121 +29,22 @@
 #include <filament/RenderableManager.h>
 #include <filament/VertexBuffer.h>
 #include <math/vec4.h>
-#include <mujoco/mujoco.h>
+#include "experimental/filament/render_context_filament.h"
 
 // Functions for creating filament vertex and index buffers.
 namespace mujoco {
 
-// The type of data stored in an index buffer.
-typedef enum mjrIndexType_ {
-  mjINDEX_TYPE_U16 = 0,
-  mjINDEX_TYPE_U32 = 1,
-} mjrIndexType;
-
-// The type of primitive to be drawn by vertex data.
-typedef enum mjrMeshPrimitiveType_ {
-  mjMESH_PRIMITIVE_TYPE_TRIANGLES = 0,
-  mjMESH_PRIMITIVE_TYPE_LINES = 1,
-} mjrMeshPrimitiveType;
-
-// The usage/purpose of an attribute of a vertex.
-typedef enum mjrVertexAttributeUsage_ {
-  mjVERTEX_ATTRIBUTE_USAGE_POSITION = 0,
-  mjVERTEX_ATTRIBUTE_USAGE_NORMAL = 1,
-  mjVERTEX_ATTRIBUTE_USAGE_TANGENTS = 2,
-  mjVERTEX_ATTRIBUTE_USAGE_UV = 3,
-  mjVERTEX_ATTRIBUTE_USAGE_COLOR = 4,
-} mjrVertexAttributeUsage;
-
-// The data format of an attribute of a vertex.
-typedef enum mjrVertexAttributeType_ {
-  mjVERTEX_ATTRIBUTE_TYPE_FLOAT2 = 0,
-  mjVERTEX_ATTRIBUTE_TYPE_FLOAT3 = 1,
-  mjVERTEX_ATTRIBUTE_TYPE_FLOAT4 = 2,
-  mjVERTEX_ATTRIBUTE_TYPE_UBYTE4 = 3,
-} mjrVertexAttributeType;
-
-// Maximum number of vertex attributes that can be used by a mesh.
-enum { mjMAX_VERTEX_ATTRIBUTES = 16 };
-
-// Information about a single attribute of a vertex.
-struct mjrVertexAttribute {
-  // The data for the attribute.
-  const void* bytes;
-
-  // The usage/purpose of the attribute.
-  mjrVertexAttributeUsage usage;
-
-  // The data format of the attribute.
-  mjrVertexAttributeType type;
-};
-
-// The binary contents of a mesh.
-struct mjrMeshData {
-  // The number of vertices in the mesh. Each of the vertex arrays below is
-  // assumed to have this number of elements.
-  mjtSize nvertices;
-
-  // The number of attributes for each vertex in the mesh.
-  int nattributes;
-
-  // Information about each attribute of a vertex in the mesh. See `interleaved`
-  // for more details.
-  mjrVertexAttribute attributes[mjMAX_VERTEX_ATTRIBUTES];
-
-  // Whether the vertex attributes are interleaved or not.
-  //
-  // If true, assumes that the attributes are packed in the order specified in
-  // the attributes array, with no padding in-between. Additionally, the
-  // `data` pointer for each attribute is assumed to point to the first element
-  // of that type.
-  //
-  // If false, assume each attribute is stored in a separate array as defined
-  // by the `data` field of the attribute.
-  mjtByte interleaved;
-
-  // The number of indices in the mesh. The indices array is assumed to have
-  // this number of elements.
-  mjtSize nindices;
-
-  // The indices of the mesh, stored as either ushort or uint depending on the
-  // index type.
-  const void* indices;
-
-  // The type of data stored in the indices array.
-  mjrIndexType index_type;
-
-  // The type of primitive to be drawn by vertex data.
-  mjrMeshPrimitiveType primitive_type;
-
-  // Whether to compute the bounds of the mesh using the vertex positions.
-  mjtByte compute_bounds;
-
-  // The bounds of the mesh. If bounds_min == bounds_max, then we assume that
-  // that the bounds are not set (i.e. the bounds is empty).
-  float bounds_min[3];
-  float bounds_max[3];
-
-  // Because rendering may be multithreaded, we cannot make assumptions about
-  // when the mesh data will finish uploading to the GPU. As such, we will use
-  // this callback to notify callers when it is safe to free the mesh data.
-  void (*release_callback)(void* user_data);
-
-  // User data to pass to the release callback.
-  void* user_data;
-};
-
-// Initializes the MeshData to default values.
-void mjr_defaultMeshData(mjrMeshData* data);
-
 // Owns a Vertex and Index buffer representing a geometry mesh.
-class Mesh {
+class Mesh : public mjrMesh {
  public:
   // Creates a Mesh from the given MeshData.
   Mesh(filament::Engine* engine, const mjrMeshData& data);
 
   ~Mesh();
 
+  Mesh(const Mesh&) = delete;
+  Mesh& operator=(const Mesh&) = delete;
+
   // Returns the filament IndexBuffer for the mesh.
   filament::IndexBuffer* GetFilamentIndexBuffer() const;
 
@@ -163,8 +63,12 @@ class Mesh {
   // Returns the bounds of the mesh.
   filament::Box GetBounds() const;
 
-  Mesh(const Mesh&) = delete;
+  static Mesh* downcast(mjrMesh* mesh) {
-  Mesh& operator=(const Mesh&) = delete;
+    return static_cast<Mesh*>(mesh);
+  }
+  static const Mesh* downcast(const mjrMesh* mesh) {
+    return static_cast<const Mesh*>(mesh);
+  }
 
  private:
   void BuildVertexBuffer(const mjrMeshData& data);

src/experimental/filament/filament/render_target.cc

@@ -30,11 +30,6 @@
 
 namespace mujoco {
 
-void mjr_defaultRenderTargetConfig(mjrRenderTargetConfig* config) {
-  config->color_format = mjPIXEL_FORMAT_RGBA8;
-  config->depth_format = mjPIXEL_FORMAT_DEPTH32F;
-}
-
 RenderTarget::RenderTarget(filament::Engine* engine,
                            const mjrRenderTargetConfig& config)
     : engine_(engine), config_(config) {}

src/experimental/filament/filament/render_target.h

@@ -22,20 +22,12 @@
 #include <filament/Engine.h>
 #include <filament/Texture.h>
 #include "experimental/filament/filament/texture.h"
+#include "experimental/filament/render_context_filament.h"
 
 namespace mujoco {
 
-// Defines the basic properties of a render target.
-struct mjrRenderTargetConfig {
-  mjrPixelFormat color_format;
-  mjrPixelFormat depth_format;
-};
-
-// Initializes the RenderTargetConfig to default values.
-void mjr_defaultRenderTargetConfig(mjrRenderTargetConfig* config);
-
 // Manages a filament RenderTarget and the textures which are bound to it.
-class RenderTarget {
+class RenderTarget : public mjrRenderTarget {
  public:
   // Defines the types of textures to create for the color and depth
   // attachments.
@@ -62,6 +54,13 @@ class RenderTarget {
   // Returns the underlying filament render target.
   filament::RenderTarget* GetFilamentRenderTarget() const;
 
+  static RenderTarget* downcast(mjrRenderTarget* render_target) {
+    return static_cast<RenderTarget*>(render_target);
+  }
+  static const RenderTarget* downcast(const mjrRenderTarget* render_target) {
+    return static_cast<const RenderTarget*>(render_target);
+  }
+
  private:
   void Destroy();
 

src/experimental/filament/filament/renderable.cc

@@ -26,20 +26,17 @@
 #include <math/mat4.h>
 #include <utils/EntityManager.h>
 #include <mujoco/mujoco.h>
-#include "experimental/filament/filament/draw_mode.h"
 #include "experimental/filament/filament/material.h"
 #include "experimental/filament/filament/math_util.h"
 #include "experimental/filament/filament/mesh.h"
 #include "experimental/filament/filament/object_manager.h"
+#include "experimental/filament/filament/texture.h"
+#include "experimental/filament/render_context_filament.h"
 
 namespace mujoco {
 
 using filament::math::mat4f;
 
-void mjr_defaultRenderableParams(mjrRenderableParams* params) {
-  params->shading_model = mjSHADING_MODEL_SCENE_OBJECT;
-}
-
 Renderable::Renderable(ObjectManager* object_mgr, const mjrRenderableParams& params)
     : object_mgr_(object_mgr), params_(params) {
   mjr_defaultMaterialParams(&material_params_);
@@ -57,7 +54,7 @@ Renderable::~Renderable() noexcept {
     engine->destroy(part.entity);
     em.destroy(part.entity);
   }
-  for (int i = 0; i < kNumDrawModes; ++i) {
+  for (int i = 0; i < mjNUM_DRAW_MODES; ++i) {
     if (instances_[i] != nullptr) {
       engine->destroy(instances_[i]);
       instances_[i] = nullptr;
@@ -206,13 +203,13 @@ void Renderable::UpdateMaterial(const mjrMaterialParams& params,
   material_params_ = params;
   material_textures_ = textures;
 
-  AssignMaterial(DrawMode::Color, GetColorMaterialType());
+  AssignMaterial(mjDRAW_MODE_COLOR, GetColorMaterialType());
   if (params_.shading_model == mjSHADING_MODEL_SCENE_OBJECT) {
-    AssignMaterial(DrawMode::Depth, ObjectManager::kUnlitDepth);
+    AssignMaterial(mjDRAW_MODE_DEPTH, ObjectManager::kUnlitDepth);
-    AssignMaterial(DrawMode::Segmentation, ObjectManager::kUnlitSegmentation);
+    AssignMaterial(mjDRAW_MODE_SEGMENTATION, ObjectManager::kUnlitSegmentation);
   }
 
-  for (int i = 0; i < kNumDrawModes; ++i) {
+  for (int i = 0; i < mjNUM_DRAW_MODES; ++i) {
     if (instances_[i]) {
       UpdateMaterialInstance(instances_[i], material_params_,
                              material_textures_, object_mgr_);
@@ -221,7 +218,7 @@ void Renderable::UpdateMaterial(const mjrMaterialParams& params,
   SetDrawMode(draw_mode_);
 }
 
-void Renderable::AssignMaterial(DrawMode mode,
+void Renderable::AssignMaterial(mjrDrawMode mode,
                                 ObjectManager::MaterialType material_type) {
   const int index = static_cast<int>(mode);
 
@@ -248,10 +245,10 @@ const mjrMaterialTextures& Renderable::GetMaterialTextures() const {
   return material_textures_;
 }
 
-void Renderable::SetDrawMode(DrawMode mode) {
+void Renderable::SetDrawMode(mjrDrawMode mode) {
   // Only SceneObjects support non-color draw modes.
   if (params_.shading_model != mjSHADING_MODEL_SCENE_OBJECT) {
-    mode = DrawMode::Color;
+    mode = mjDRAW_MODE_COLOR;
   }
 
   filament::MaterialInstance* instance = instances_[static_cast<int>(mode)];
@@ -369,6 +366,7 @@ ObjectManager::MaterialType Renderable::GetColorMaterialType() const {
   // geometry) and `mesh_texcoordadr` stores the address of the mesh uvs if
   // it has them.
   bool has_texcoords = false;
+  const Texture* color_texture = Texture::downcast(material_textures_.color);
   if (!parts_.empty()) {
     const auto attribs = parts_[0].mesh->GetVertexAttributes();
     auto it = std::find(attribs.begin(), attribs.end(),
@@ -376,7 +374,7 @@ ObjectManager::MaterialType Renderable::GetColorMaterialType() const {
     has_texcoords = (it != attribs.end());
   }
 
-  if (material_textures_.color == nullptr) {
+  if (color_texture == nullptr) {
     if (material_params_.color[3] < 1.0f) {
       return ObjectManager::kPhongColorFade;
     } else if (material_params_.reflective) {
@@ -384,7 +382,7 @@ ObjectManager::MaterialType Renderable::GetColorMaterialType() const {
     } else {
       return ObjectManager::kPhongColor;
     }
-  } else if (material_textures_.color->GetFilamentTexture()->getTarget() ==
+  } else if (color_texture->GetFilamentTexture()->getTarget() ==
               filament::Texture::Sampler::SAMPLER_CUBEMAP) {
     if (material_params_.color[3] < 1.0f) {
       return ObjectManager::kPhongCubeFade;

src/experimental/filament/filament/renderable.h

@@ -24,29 +24,14 @@
 #include <filament/Scene.h>
 #include <math/mat4.h>
 #include <utils/Entity.h>
-#include "experimental/filament/filament/draw_mode.h"
 #include "experimental/filament/filament/material.h"
 #include "experimental/filament/filament/math_util.h"
 #include "experimental/filament/filament/mesh.h"
 #include "experimental/filament/filament/object_manager.h"
+#include "experimental/filament/render_context_filament.h"
 
 namespace mujoco {
 
-// The shading model (material) for a Renderable.
-typedef enum mjrShadingModel_ {
-  mjSHADING_MODEL_SCENE_OBJECT,
-  mjSHADING_MODEL_DECOR,
-  mjSHADING_MODEL_DECOR_LINES,
-  mjSHADING_MODEL_UX,
-} mjrShadingModel;
-
-// Configuration parameters for a Renderable.
-struct mjrRenderableParams {
-  mjrShadingModel shading_model;
-};
-
-void mjr_defaultRenderableParams(mjrRenderableParams* params);
-
 // A Renderable is effectively two things: a mesh and a material.
 //
 // The mesh describes the surface geometry of the object and the material
@@ -61,7 +46,7 @@ void mjr_defaultRenderableParams(mjrRenderableParams* params);
 // the user specifies the MaterialParams and MaterialTextures to use with the
 // ShadingModel. Its these properties that ultimately define the actual material
 // of the Renderable.
-class Renderable {
+class Renderable : public mjrRenderable {
  public:
   // Default filament values for priority and layer mask.
   static constexpr std::uint8_t kDefaultPriority = 4;
@@ -124,7 +109,7 @@ class Renderable {
 
   // Further defines the material of the renderable. Only applies to renderables
   // with a SceneObject shading model.
-  void SetDrawMode(DrawMode mode);
+  void SetDrawMode(mjrDrawMode mode);
 
   // Updates the parameters for the material.
   void UpdateMaterial(const mjrMaterialParams& params,
@@ -139,6 +124,13 @@ class Renderable {
   // Returns the filament Engine managing the renderables.
   filament::Engine* GetEngine();
 
+  static Renderable* downcast(mjrRenderable* renderable) {
+    return static_cast<Renderable*>(renderable);
+  }
+  static const Renderable* downcast(const mjrRenderable* renderable) {
+    return static_cast<const Renderable*>(renderable);
+  }
+
  private:
   struct Part {
     utils::Entity entity;
@@ -149,16 +141,16 @@ class Renderable {
 
   void InitPartEntity(Part& part);
 
-  void AssignMaterial(DrawMode mode, ObjectManager::MaterialType material_type);
+  void AssignMaterial(mjrDrawMode mode, ObjectManager::MaterialType material_type);
 
   ObjectManager::MaterialType GetColorMaterialType() const;
 
   ObjectManager* object_mgr_;
   mjrRenderableParams params_;
-  filament::MaterialInstance* instances_[kNumDrawModes] = {nullptr};
+  filament::MaterialInstance* instances_[mjNUM_DRAW_MODES] = {nullptr};
   mjrMaterialParams material_params_;
   mjrMaterialTextures material_textures_;
-  DrawMode draw_mode_ = DrawMode::Color;
+  mjrDrawMode draw_mode_ = mjDRAW_MODE_COLOR;
   filament::Scene* assigned_scene_ = nullptr;
   std::vector<Part> parts_;
   filament::math::mat4f transform_;

src/experimental/filament/filament/scene_view.cc

@@ -40,13 +40,12 @@
 #include <utils/EntityManager.h>
 #include <mujoco/mujoco.h>
 #include "experimental/filament/filament/color_grading_options.h"
-#include "experimental/filament/filament/draw_mode.h"
 #include "experimental/filament/filament/light.h"
-#include "experimental/filament/filament/material.h"
 #include "experimental/filament/filament/math_util.h"
 #include "experimental/filament/filament/render_target.h"
 #include "experimental/filament/filament/renderable.h"
 #include "experimental/filament/filament/texture.h"
+#include "experimental/filament/render_context_filament.h"
 
 namespace mujoco {
 
@@ -54,10 +53,6 @@ using filament::math::float3;
 using filament::math::float4;
 using filament::math::mat4;
 
-static constexpr int kNormalIndex = static_cast<int>(DrawMode::Color);
-static constexpr int kDepthIndex = static_cast<int>(DrawMode::Depth);
-static constexpr int kSegmentIndex = static_cast<int>(DrawMode::Segmentation);
-
 static filament::ColorGrading::Builder ToBuilder(
     const ColorGradingOptions& opts) {
   return filament::ColorGrading::Builder()
@@ -146,11 +141,11 @@ SceneView::SceneView(filament::Engine* engine) : engine_(engine) {
 
   // Disable post processing for the depth and segmentation views to preserve
   // the values.
-  views_[kDepthIndex]->setPostProcessingEnabled(false);
+  views_[mjDRAW_MODE_DEPTH]->setPostProcessingEnabled(false);
-  views_[kSegmentIndex]->setPostProcessingEnabled(false);
+  views_[mjDRAW_MODE_SEGMENTATION]->setPostProcessingEnabled(false);
 
   // Rotate the fog to align with mujoco's +Z up space.
-  auto fog = views_[kNormalIndex]->getFogEntity();
+  auto fog = views_[mjDRAW_MODE_COLOR]->getFogEntity();
   auto& tm = engine->getTransformManager();
   tm.create(fog);
   tm.setTransform(tm.getInstance(fog),
@@ -158,6 +153,10 @@ SceneView::SceneView(filament::Engine* engine) : engine_(engine) {
 }
 
 SceneView::~SceneView() {
+  if (skybox_) {
+    scene_->setSkybox(nullptr);
+    engine_->destroy(skybox_);
+  }
   for (auto& light : lights_) {
     light->RemoveFromScene(scene_);
   }
@@ -210,15 +209,17 @@ void SceneView::RemoveFromScene(Renderable* renderable) {
   }
 }
 
-void SceneView::AddToScene(filament::Skybox* skybox) {
+void SceneView::SetSkybox(const Texture* skybox_texture) {
-  skybox_ = skybox;
+  if (skybox_) {
-  scene_->setSkybox(skybox);
-}
-
-void SceneView::RemoveFromScene(filament::Skybox* skybox) {
-  if (skybox_ == skybox) {
-    skybox_ = nullptr;
     scene_->setSkybox(nullptr);
+    engine_->destroy(skybox_);
+    skybox_ = nullptr;
+  }
+  if (skybox_texture) {
+    filament::Skybox::Builder builder;
+    builder.environment(skybox_texture->GetFilamentTexture());
+    skybox_ = builder.build(*engine_);
+    scene_->setSkybox(skybox_);
   }
 }
 
@@ -249,7 +250,7 @@ void SceneView::Render(filament::Renderer* renderer,
   }
 
   // Render reflection passes.
-  if (request.draw_mode == DrawMode::Color && reflections_enabled_) {
+  if (request.draw_mode == mjDRAW_MODE_COLOR && reflections_enabled_) {
     for (size_t i = 0; i < reflectives_.size(); ++i) {
       Renderable* renderable = reflectives_[i];
 
@@ -311,7 +312,7 @@ void SceneView::SetColorGradingOptions(const ColorGradingOptions& opts) {
   auto color_grading = ToBuilder(color_grading_options_)
                            .toneMapper(tone_mapper.get())
                            .build(*engine_);
-  views_[kNormalIndex]->setColorGrading(color_grading);
+  views_[mjDRAW_MODE_COLOR]->setColorGrading(color_grading);
   if (color_grading_) {
     engine_->destroy(color_grading_);
   }
@@ -320,11 +321,11 @@ void SceneView::SetColorGradingOptions(const ColorGradingOptions& opts) {
 }
 
 void SceneView::EnableShadows() {
-  views_[kNormalIndex]->setShadowingEnabled(true);
+  views_[mjDRAW_MODE_COLOR]->setShadowingEnabled(true);
 }
 
 void SceneView::DisableShadows() {
-  views_[kNormalIndex]->setShadowingEnabled(false);
+  views_[mjDRAW_MODE_COLOR]->setShadowingEnabled(false);
 }
 
 void SceneView::EnableReflections() {
@@ -345,19 +346,18 @@ void SceneView::DisableReflections() {
     textures.reflection = nullptr;
     renderable->UpdateMaterial(renderable->GetMaterialParams(), textures);
   }
-
 }
 
 void SceneView::EnablePostProcessing() {
-  views_[kNormalIndex]->setPostProcessingEnabled(true);
+  views_[mjDRAW_MODE_COLOR]->setPostProcessingEnabled(true);
 }
 
 void SceneView::DisablePostProcessing() {
-  views_[kNormalIndex]->setPostProcessingEnabled(false);
+  views_[mjDRAW_MODE_COLOR]->setPostProcessingEnabled(false);
 }
 
 filament::View* SceneView::GetDefaultRenderView() {
-  return views_[kNormalIndex];
+  return views_[mjDRAW_MODE_COLOR];
 }
 
 ColorGradingOptions SceneView::GetColorGradingOptions() const {

src/experimental/filament/filament/scene_view.h

@@ -27,10 +27,11 @@
 #include <filament/View.h>
 #include <mujoco/mujoco.h>
 #include "experimental/filament/filament/color_grading_options.h"
-#include "experimental/filament/filament/draw_mode.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"
 
 namespace mujoco {
 
@@ -39,23 +40,25 @@ namespace mujoco {
 // The filament Scene is populated with the objects (e.g. lights, renderables,
 // skybox, etc.). It manages multiple views to support a variety of draw modes
 // (e.g. normal, depth, segmentation, etc.) as well as reflective surfaces.
-class SceneView {
+class SceneView : public mjrScene {
  public:
   SceneView(filament::Engine* engine);
   ~SceneView();
 
+  SceneView(const SceneView&) = delete;
+  SceneView& operator=(const SceneView&) = delete;
+
   // Adds/removes entities from the scene.
   void AddToScene(Light* light);
   void RemoveFromScene(Light* light);
   void AddToScene(Renderable* renderable);
   void RemoveFromScene(Renderable* renderable);
-  void AddToScene(filament::Skybox* skybox);
+  void SetSkybox(const Texture* skybox_texture);
-  void RemoveFromScene(filament::Skybox* skybox);
 
   // Parameters for rendering the scene.
   struct RenderRequest {
     // The draw mode (e.g. normal, depth, segmentation) to render.
-    DrawMode draw_mode = DrawMode::Color;
+    mjrDrawMode draw_mode = mjDRAW_MODE_COLOR;
     // The target viewport for the rendered image.
     mjrRect viewport;
     // The camera from which to render the scene.
@@ -90,8 +93,12 @@ class SceneView {
   ColorGradingOptions GetColorGradingOptions() const;
   void SetColorGradingOptions(const ColorGradingOptions& opts);
 
-  SceneView(const SceneView&) = delete;
+  static SceneView* downcast(mjrScene* scene) {
-  SceneView& operator=(const SceneView&) = delete;
+    return static_cast<SceneView*>(scene);
+  }
+  static const SceneView* downcast(const mjrScene* scene) {
+    return static_cast<const SceneView*>(scene);
+  }
 
  private:
   // Marks a renderable as reflective. Reflective renderables have to be
@@ -103,7 +110,7 @@ class SceneView {
   filament::Camera* camera_ = nullptr;
   filament::ColorGrading* color_grading_ = nullptr;
   ColorGradingOptions color_grading_options_;
-  std::array<filament::View*, kNumDrawModes> views_;
+  std::array<filament::View*, mjNUM_DRAW_MODES> views_;
 
   // Scene objects.
   std::unordered_set<Light*> lights_;

src/experimental/filament/filament/texture.cc

@@ -110,14 +110,6 @@ static filament::Texture::InternalFormat GetTextureInternalFormat(
   }
 }
 
-void mjr_defaultTextureData(mjrTextureData* data) {
-  std::memset(data, 0, sizeof(mjrTextureData));
-}
-
-void mjr_defaultTextureConfig(mjrTextureConfig* config) {
-  std::memset(config, 0, sizeof(mjrTextureConfig));
-}
-
 Texture::Texture(filament::Engine* engine, const mjrTextureConfig& config,
                  InternalFlags flags)
     : engine_(engine), config_(config) {

src/experimental/filament/filament/texture.h

@@ -15,76 +15,18 @@
 #ifndef MUJOCO_SRC_EXPERIMENTAL_FILAMENT_FILAMENT_TEXTURE_H_
 #define MUJOCO_SRC_EXPERIMENTAL_FILAMENT_FILAMENT_TEXTURE_H_
 
-#include <cstddef>
-
 #include <filament/Engine.h>
 #include <filament/Texture.h>
 #include <math/vec3.h>
 #include <mujoco/mjmodel.h>
 #include <mujoco/mujoco.h>
+#include "experimental/filament/render_context_filament.h"
 
 // Functions for creating filament textures.
 namespace mujoco {
 
-// Pixel formats for textures.
-typedef enum mjrPixelFormat_ {
-  mjPIXEL_FORMAT_UNKNOWN = 0,
-  mjPIXEL_FORMAT_R8,
-  mjPIXEL_FORMAT_RGB8,
-  mjPIXEL_FORMAT_RGBA8,
-  mjPIXEL_FORMAT_R32F,
-  mjPIXEL_FORMAT_DEPTH32F,
-  mjPIXEL_FORMAT_KTX,
-} mjrPixelFormat;
-
-typedef mjtTexture mjrTextureTarget;
-typedef mjtColorSpace mjrColorSpace;
-
-// The binary contents of a texture.
-struct mjrTextureData {
-  // Pointer to the image data. If null, an empty texture will be created.
-  const void* bytes;
-
-  // The number of bytes in the image data.
-  mjtSize nbytes;
-
-  // Because rendering may be multithreaded, we cannot make assumptions about
-  // when the image data will finish uploading to the GPU. As such, we will use
-  // this callback to notify callers when it is safe to free the image data.
-  void (*release_callback)(void* user_data);
-
-  // User data to pass to the release callback.
-  void* user_data;
-};
-
-// Initializes the TextureData to default values.
-void mjr_defaultTextureData(mjrTextureData* data);
-
-// Defines the basic properties of a texture.
-struct mjrTextureConfig {
-  // The width of the texture. For compressed textures (e.g. KTX), this is the
-  // number of bytes in the compressed data.
-  int width;
-
-  // The height of the texture. For compressed textures (e.g. KTX), this should
-  // be 0.
-  int height;
-
-  // The target of the texture (e.g. 2D, cube, etc.)
-  mjrTextureTarget target;
-
-  // The format of the pixels in the texture (e.g. RGB8, RGBA8, KTX, etc.)
-  mjrPixelFormat format;
-
-  // The color space of the texture (e.g. LINEAR, sRGB, etc.)
-  mjrColorSpace color_space;
-};
-
-// Initializes the TextureConfig to default values.
-void mjr_defaultTextureConfig(mjrTextureConfig* config);
-
 // Wrapper around a filament::Texture.
-class Texture {
+class Texture : public mjrTexture {
  public:
   // Flags for internal use.
   struct InternalFlags {
@@ -99,6 +41,9 @@ class Texture {
 
   ~Texture();
 
+  Texture(const Texture&) = delete;
+  Texture& operator=(const Texture&) = delete;
+
   // Uploads the given data to the texture.
   void Upload(const mjrTextureData& data);
 
@@ -117,8 +62,12 @@ class Texture {
     return has_spherical_harmonics_ ? &spherical_harmonics_ : nullptr;
   }
 
-  Texture(const Texture&) = delete;
+  static Texture* downcast(mjrTexture* texture) {
-  Texture& operator=(const Texture&) = delete;
+    return static_cast<Texture*>(texture);
+  }
+  static const Texture* downcast(const mjrTexture* texture) {
+    return static_cast<const Texture*>(texture);
+  }
 
  private:
   void ReleaseData();

src/experimental/filament/render_context_filament.cc

@@ -14,6 +14,7 @@
 
 #include "experimental/filament/render_context_filament.h"
 
+#include <array>
 #include <cstdint>
 #include <cstring>
 
@@ -36,12 +37,91 @@ static void CheckFilamentContext() {
   }
 }
 
+template <int N>
+static void setf(float (&arr)[N], const std::array<float, N>& values) {
+  for (int i = 0; i < N; ++i) {
+    arr[i] = values[i];
+  }
+}
+
+
 extern "C" {
 
 void mjrf_defaultFilamentConfig(mjrFilamentConfig* config) {
   memset(config, 0, sizeof(mjrFilamentConfig));
 }
 
+void mjr_defaultTextureData(mjrTextureData* data) {
+  memset(data, 0, sizeof(mjrTextureData));
+}
+
+void mjr_defaultTextureConfig(mjrTextureConfig* config) {
+  memset(config, 0, sizeof(mjrTextureConfig));
+}
+
+void mjr_defaultMeshData(mjrMeshData* data) {
+  std::memset(data, 0, sizeof(mjrMeshData));
+}
+
+void mjr_defaultLightParams(mjrLightParams* params) {
+  params->type = mjLIGHT_POINT;
+  params->texture = nullptr;
+  params->color[0] = 0;
+  params->color[1] = 0;
+  params->color[2] = 0;
+  params->intensity = 0.0f;
+  params->cast_shadows = true;
+  params->range = 10.0f;
+  params->spot_cone_angle = 180.f;
+  params->bulb_radius = 0.0f;
+  params->shadow_map_size = 2048;
+  params->vsm_blur_width = 0.0f;
+}
+
+void mjr_defaultMaterialTextures(mjrMaterialTextures* textures) {
+  textures->color = nullptr;
+  textures->normal = nullptr;
+  textures->metallic = nullptr;
+  textures->roughness = nullptr;
+  textures->occlusion = nullptr;
+  textures->orm = nullptr;
+  textures->emissive = nullptr;
+  textures->reflection = nullptr;
+}
+
+void mjr_defaultMaterialParams(mjrMaterialParams* params) {
+  setf(params->color, {1.f, 1.f, 1.f, 1.f});
+  setf(params->segmentation_color, {1, 1, 1, 1});
+  setf(params->uv_scale, {1, 1, 1});
+  setf(params->uv_offset, {0, 0, 0});
+  setf(params->scissor, {0, 0, 0, 0});
+  params->emissive = -1.0f;
+  params->specular = -1.0f;
+  params->glossiness = -1.0f;
+  params->metallic = -1.0f;
+  params->roughness = -1.0f;
+  params->reflectance = 0.0f;
+  params->tex_uniform = false;
+  params->reflective = false;
+}
+
+void mjr_defaultRenderableParams(mjrRenderableParams* params) {
+  params->shading_model = mjSHADING_MODEL_SCENE_OBJECT;
+}
+
+void mjr_defaultRenderTargetConfig(mjrRenderTargetConfig* config) {
+  config->color_format = mjPIXEL_FORMAT_RGBA8;
+  config->depth_format = mjPIXEL_FORMAT_DEPTH32F;
+}
+
+void mjr_defaultRenderRequest(mjrRenderRequest* request) {
+  memset(request, 0, sizeof(mjrRenderRequest));
+}
+
+void mjr_defaultReadPixelsRequest(mjrReadPixelsRequest* request) {
+  memset(request, 0, sizeof(mjrReadPixelsRequest));
+}
+
 void mjrf_makeFilamentContext(const mjModel* m, mjrContext* con,
                               const mjrFilamentConfig* config) {
   // TODO: Support multiple contexts and multiple threads. For now, we'll just

src/experimental/filament/render_context_filament.h

@@ -29,12 +29,345 @@ extern "C" {
 // IMPORTANT: This API should still be considered experimental and is likely
 // change frequently.
 
+// Opaque types.
+struct mjrTexture {};
+struct mjrMesh {};
+struct mjrScene {};
+struct mjrLight {};
+struct mjrRenderable {};
+struct mjrRenderTarget {};
+
+// Opaque type for the filament rendering context.
+struct mjrfContext {};
+
+// The different modes that can be used to render a scene.
+typedef enum mjrDrawMode_ {
+  // Render the scene with "normal" colors and lighting.
+  mjDRAW_MODE_COLOR,
+  // Render the scene as a grayscale depth map.
+  mjDRAW_MODE_DEPTH,
+  // Render each object with a unique, uniform (flat) color regardless of
+  // lighting and texture.
+  mjDRAW_MODE_SEGMENTATION,
+} mjrDrawMode;
+
+enum { mjNUM_DRAW_MODES = 3 };
+
+// The shading model (material) for a Renderable.
+typedef enum mjrShadingModel_ {
+  // For renderables in the main 3D scene.
+  mjSHADING_MODEL_SCENE_OBJECT = 0,
+  // For UX renderables.
+  mjSHADING_MODEL_UX,
+  // For decorative elements in a Scene (e.g. contact points, force vectors,
+  // etc.). These objects will not be affected by lighting.
+  mjSHADING_MODEL_DECOR,
+  // As above, but uses a line primitives for drawing.
+  mjSHADING_MODEL_DECOR_LINES,
+} mjrShadingModel;
+
+// The type of data stored in an index buffer.
+typedef enum mjrIndexType_ {
+  mjINDEX_TYPE_U16 = 0,
+  mjINDEX_TYPE_U32,
+} mjrIndexType;
+
+// The type of primitive to be drawn by vertex data.
+typedef enum mjrMeshPrimitiveType_ {
+  mjMESH_PRIMITIVE_TYPE_TRIANGLES = 0,
+  mjMESH_PRIMITIVE_TYPE_LINES,
+} mjrMeshPrimitiveType;
+
+// The usage/purpose of an attribute of a vertex.
+typedef enum mjrVertexAttributeUsage_ {
+  mjVERTEX_ATTRIBUTE_USAGE_POSITION = 0,
+  mjVERTEX_ATTRIBUTE_USAGE_NORMAL,
+  mjVERTEX_ATTRIBUTE_USAGE_TANGENTS,
+  mjVERTEX_ATTRIBUTE_USAGE_UV,
+  mjVERTEX_ATTRIBUTE_USAGE_COLOR,
+} mjrVertexAttributeUsage;
+
+// The data format of an attribute of a vertex.
+typedef enum mjrVertexAttributeType_ {
+  mjVERTEX_ATTRIBUTE_TYPE_FLOAT2 = 0,
+  mjVERTEX_ATTRIBUTE_TYPE_FLOAT3,
+  mjVERTEX_ATTRIBUTE_TYPE_FLOAT4,
+  mjVERTEX_ATTRIBUTE_TYPE_UBYTE4,
+} mjrVertexAttributeType;
+
+// Pixel formats for textures.
+typedef enum mjrPixelFormat_ {
+  mjPIXEL_FORMAT_UNKNOWN = 0,
+  mjPIXEL_FORMAT_R8,
+  mjPIXEL_FORMAT_RGB8,
+  mjPIXEL_FORMAT_RGBA8,
+  mjPIXEL_FORMAT_R32F,
+  mjPIXEL_FORMAT_DEPTH32F,
+  mjPIXEL_FORMAT_KTX,
+} mjrPixelFormat;
+
 typedef enum mjrGraphicsApi_ {  // backend graphics API to use
   mjGRAPHICS_API_DEFAULT = 0,   // default based on platform
   mjGRAPHICS_API_OPENGL,        // OpenGL (desktop) / WebGL
   mjGRAPHICS_API_VULKAN         // Vulkan
 } mjrGraphicsApi;
 
+
+// Rendering is asynchronous by nature. Each render request is assigned a
+// unique Handle which can be used to query the status of the request. The
+// Handle can also be used to block until the request is completed.
+typedef std::uint64_t mjrFrameHandle;
+
+// Bring some legacy mjt types into the mjr namespace.
+typedef mjtTexture mjrTextureTarget;
+typedef mjtColorSpace mjrColorSpace;
+typedef mjtLightType mjrLightType;
+
+// The textures that can be assigned to the drawable's material.
+struct mjrMaterialTextures {
+  const mjrTexture* color;
+  const mjrTexture* normal;
+  const mjrTexture* metallic;
+  const mjrTexture* roughness;
+  const mjrTexture* occlusion;
+  const mjrTexture* orm;
+  const mjrTexture* emissive;
+  const mjrTexture* reflection;
+};
+
+// Initializes the mjrMaterialTextures to default values.
+void mjr_defaultMaterialTextures(mjrMaterialTextures* textures);
+
+// The parameters that can be applied to the drawable's material.
+struct mjrMaterialParams {
+  float color[4];
+  float segmentation_color[4];
+  float tex_repeat[2];
+  float uv_scale[3];
+  float uv_offset[3];
+  float scissor[4];
+  float specular;
+  float glossiness;
+  float metallic;
+  float roughness;
+  float emissive;
+  float reflectance;
+  mjtByte tex_uniform;
+  mjtByte reflective;
+};
+
+// Initializes the mjrMaterialParams to default values.
+void mjr_defaultMaterialParams(mjrMaterialParams* params);
+
+// The binary contents of a texture.
+struct mjrTextureData {
+  // Pointer to the image data. If null, an empty texture will be created.
+  const void* bytes;
+
+  // The number of bytes in the image data.
+  mjtSize nbytes;
+
+  // Because rendering may be multithreaded, we cannot make assumptions about
+  // when the image data will finish uploading to the GPU. As such, we will use
+  // this callback to notify callers when it is safe to free the image data.
+  void (*release_callback)(void* user_data);
+
+  // User data to pass to the release callback.
+  void* user_data;
+};
+
+// Initializes the mjrTextureData to default values.
+void mjr_defaultTextureData(mjrTextureData* data);
+
+// Defines the basic properties of a texture.
+struct mjrTextureConfig {
+  // The width of the texture. For compressed textures (e.g. KTX), this is the
+  // number of bytes in the compressed data.
+  int width;
+
+  // The height of the texture. For compressed textures (e.g. KTX), this should
+  // be 0.
+  int height;
+
+  // The target of the texture (e.g. 2D, cube, etc.)
+  mjrTextureTarget target;
+
+  // The format of the pixels in the texture (e.g. RGB8, RGBA8, KTX, etc.)
+  mjrPixelFormat format;
+
+  // The color space of the texture (e.g. LINEAR, sRGB, etc.)
+  mjrColorSpace color_space;
+};
+
+// Initializes the mjrTextureConfig to default values.
+void mjr_defaultTextureConfig(mjrTextureConfig* config);
+
+// Configuration parameters for a Renderable.
+struct mjrRenderableParams {
+  // The shading model to use for the Renderable.
+  mjrShadingModel shading_model;
+};
+
+// Initializes the mjrRenderableParams to default values.
+void mjr_defaultRenderableParams(mjrRenderableParams* params);
+
+// Information about a single attribute of a vertex.
+struct mjrVertexAttribute {
+  // The data for the attribute.
+  const void* bytes;
+
+  // The usage/purpose of the attribute.
+  mjrVertexAttributeUsage usage;
+
+  // The data format of the attribute.
+  mjrVertexAttributeType type;
+};
+
+// Maximum number of vertex attributes in a mesh.
+enum { mjMAX_VERTEX_ATTRIBUTES = 16 };
+
+// The binary contents of a mesh.
+struct mjrMeshData {
+  // The number of vertices in the mesh. Each of the vertex arrays below is
+  // assumed to have this number of elements.
+  mjtSize nvertices;
+
+  // The number of attributes for each vertex in the mesh.
+  int nattributes;
+
+  // Information about each attribute of a vertex in the mesh. See `interleaved`
+  // for more details.
+  mjrVertexAttribute attributes[mjMAX_VERTEX_ATTRIBUTES];
+
+  // Whether the vertex attributes are interleaved or not.
+  //
+  // If true, assumes that the attributes are packed in the order specified in
+  // the attributes array, with no padding in-between. Additionally, the
+  // `data` pointer for each attribute is assumed to point to the first element
+  // of that type.
+  //
+  // If false, assume each attribute is stored in a separate array as defined
+  // by the `data` field of the attribute.
+  mjtByte interleaved;
+
+  // The number of indices in the mesh. The indices array is assumed to have
+  // this number of elements.
+  mjtSize nindices;
+
+  // The indices of the mesh, stored as either ushort or uint depending on the
+  // index type.
+  const void* indices;
+
+  // The type of data stored in the indices array.
+  mjrIndexType index_type;
+
+  // The type of primitive to be drawn by vertex data.
+  mjrMeshPrimitiveType primitive_type;
+
+  // Whether to compute the bounds of the mesh using the vertex positions.
+  mjtByte compute_bounds;
+
+  // The bounds of the mesh. If bounds_min == bounds_max, then we assume that
+  // that the bounds are not set (i.e. the bounds is empty).
+  float bounds_min[3];
+  float bounds_max[3];
+
+  // Because rendering may be multithreaded, we cannot make assumptions about
+  // when the mesh data will finish uploading to the GPU. As such, we will use
+  // this callback to notify callers when it is safe to free the mesh data.
+  void (*release_callback)(void* user_data);
+
+  // User data to pass to the release callback.
+  void* user_data;
+};
+
+// Initializes the mjrMeshData to default values.
+void mjr_defaultMeshData(mjrMeshData* data);
+
+// Configuration parameters for a light.
+struct mjrLightParams {
+  // The type of light (e.g. spot, point, directional, etc.)
+  mjrLightType type;
+  // The texture to use for image lights.
+  const mjrTexture* texture;
+  // The color of the light.
+  float color[3];
+  // The intensity of the light, in candela.
+  float intensity;
+  // Whether or not the light casts shadows.
+  mjtByte cast_shadows;
+  // The range/distance in which the light is effective, in meters.
+  float range;
+  // The angle of the spot light cone, in degrees.
+  float spot_cone_angle;
+  // The radius of the bulb used for soft shadows.
+  float bulb_radius;
+  // The size of the shadow map.
+  int shadow_map_size;
+  // Blur width for EL VSM.
+  float vsm_blur_width;
+};
+
+// Initializes the mjrLightParams to default values.
+void mjr_defaultLightParams(mjrLightParams* params);
+
+// Defines the basic properties of a render target.
+struct mjrRenderTargetConfig {
+  mjrPixelFormat color_format;
+  mjrPixelFormat depth_format;
+};
+
+// Initializes the RenderTargetConfig to default values.
+void mjr_defaultRenderTargetConfig(mjrRenderTargetConfig* config);
+
+// Information needed to render a single image of a scene.
+struct mjrRenderRequest {
+  // The scene to render.
+  mjrScene* scene;
+
+  // The method (e.g. Color, Depth, Segmentation, etc.) to use for rendering.
+  mjrDrawMode draw_mode;
+
+  // The camera from which to render the scene.
+  mjvGLCamera camera;
+
+  // The dimensions of the output image.
+  int width;
+  int height;
+
+  // The render target into which to render the image. If nullptr, the image
+  // will be rendered to the window (as previously configured in
+  // mjrFilamentConfig::native_window).
+  mjrRenderTarget* target;
+};
+
+// Initializes the mjrRenderRequest to default values.
+void mjr_defaultRenderRequest(mjrRenderRequest* request);
+
+// Information needed to read pixels from a render target.
+struct mjrReadPixelsRequest {
+  mjrRenderTarget* target;
+
+  // The buffer into which the read pixels will be written.
+  void* output;
+
+  // The number of bytes in the output buffer. This should match the size of
+  // the render target texture.
+  mjtSize num_bytes;
+
+  // Callback when the read pixels operation is complete. This will be called
+  // during WaitForFrame() or in a subsequent call to Render(). This function
+  // can optionally be used to free the output buffer if needed.
+  void (*read_completed_callback)(void* user_data);
+
+  // User data to pass to the completion callback.
+  void* user_data;
+};
+
+// Initializes the mjrReadPixelsRequest to default values.
+void mjr_defaultReadPixelsRequest(mjrReadPixelsRequest* request);
+
+// Configuration parameters for the filament rendering context.
 struct mjrFilamentConfig {
   // The native window handle into which we can render directly.
   void* native_window;

src/experimental/platform/sim/sim_profiler.cc

@@ -17,12 +17,11 @@
 #include <imgui.h>
 #include <implot.h>
 #include <mujoco/mujoco.h>
+#include "experimental/platform/ux/imgui_widgets.h"
 
 namespace mujoco::platform {
 
-SimProfiler::SimProfiler() {
+SimProfiler::SimProfiler() { Clear(); }
-  Clear();
-}
 
 void SimProfiler::Clear() {
   constexpr int kProfilerMaxFrames = 200;
@@ -88,21 +87,21 @@ void SimProfiler::Update(const mjModel* model, const mjData* data) {
   // Solver diagnostics.
   mjtNum sqrt_nnz = 0;
   int solver_niter = 0;
-  const int nisland = data->nefc ? mjMAX(1, mjMIN(data->nisland, mjNISLAND)) : 0;
+  const int nisland =
-  for (int island=0; island < nisland; island++) {
+      data->nefc ? mjMAX(1, mjMIN(data->nisland, mjNISLAND)) : 0;
+  for (int island = 0; island < nisland; island++) {
     sqrt_nnz += data->solver_nnz[island];
     solver_niter += data->solver_niter[island];
   }
   sqrt_nnz = mju_sqrt(sqrt_nnz);
 
   dim_dof_.erase(dim_dof_.begin());
-  int nv = (model->opt.enableflags & mjENBL_SLEEP) ? data->nv_awake
+  int nv = (model->opt.enableflags & mjENBL_SLEEP) ? data->nv_awake : model->nv;
-                                                     : model->nv;
   dim_dof_.push_back(nv);
 
   dim_body_.erase(dim_body_.begin());
   int nbody = (model->opt.enableflags & mjENBL_SLEEP) ? data->nbody_awake
-                                                        : model->nbody;
+                                                      : model->nbody;
   dim_body_.push_back(nbody);
 
   dim_constraint_.erase(dim_constraint_.begin());
@@ -115,16 +114,17 @@ void SimProfiler::Update(const mjModel* model, const mjData* data) {
   dim_contact_.push_back(data->ncon);
 
   dim_iteration_.erase(dim_iteration_.begin());
-  dim_iteration_.push_back(static_cast<float>(solver_niter) / nisland);
+  dim_iteration_.push_back(static_cast<float>(solver_niter) /
+                           mjMAX(1, nisland));
 }
 
-
+void SimProfiler::CpuTimeGraph(ImVec2 plot_size) {
-void SimProfiler::CpuTimeGraph() {
+  ImPlotFlags flags =
-  if (ImPlot::BeginPlot("CPU Time", ImVec2(-1, 0), ImPlotFlags_NoMouseText)) {
+      ImPlot_SetupPlotFlags(plot_size) | ImPlotFlags_NoMouseText;
+  if (ImPlot::BeginPlot("CPU msec vs frame", plot_size, flags)) {
     ImPlot::PushStyleVar(ImPlotStyleVar_LineWeight, 2.0f);
-    ImPlot::SetupAxis(ImAxis_X1, "frame", ImPlotAxisFlags_AutoFit);
+    ImPlot_SetupTimeAxis(plot_size, "");
-    ImPlot::SetupAxis(ImAxis_Y1, "msec", ImPlotAxisFlags_AutoFit);
+    ImPlot_SetupValueAxis(plot_size, "", "%.2f");
-    ImPlot::SetupAxisFormat(ImAxis_Y1, "%.2f");
     ImPlot::SetupLegend(ImPlotLocation_NorthEast);
     ImPlot::SetupFinish();
 
@@ -143,12 +143,13 @@ void SimProfiler::CpuTimeGraph() {
   }
 }
 
-void SimProfiler::DimensionsGraph() {
+void SimProfiler::DimensionsGraph(ImVec2 plot_size) {
-  if (ImPlot::BeginPlot("Dimensions", ImVec2(-1, 0), ImPlotFlags_NoMouseText)) {
+  ImPlotFlags flags =
+      ImPlot_SetupPlotFlags(plot_size) | ImPlotFlags_NoMouseText;
+  if (ImPlot::BeginPlot("Dimensions vs frame", plot_size, flags)) {
     ImPlot::PushStyleVar(ImPlotStyleVar_LineWeight, 2.0f);
-    ImPlot::SetupAxis(ImAxis_X1, "frame", ImPlotAxisFlags_AutoFit);
+    ImPlot_SetupTimeAxis(plot_size, "");
-    ImPlot::SetupAxis(ImAxis_Y1, "count", ImPlotAxisFlags_AutoFit);
+    ImPlot_SetupValueAxis(plot_size, "", "%.0f");
-    ImPlot::SetupAxisFormat(ImAxis_Y1, "%.0f");
     ImPlot::SetupLegend(ImPlotLocation_NorthEast);
     ImPlot::SetupFinish();
 

src/experimental/platform/sim/sim_profiler.h

@@ -17,6 +17,7 @@
 
 #include <vector>
 
+#include <imgui.h>
 #include <mujoco/mujoco.h>
 
 namespace mujoco::platform {
@@ -33,8 +34,8 @@ class SimProfiler {
   void Update(const mjModel* model, const mjData* data);
 
   // Displays the profiling data using ImPlot.
-  void CpuTimeGraph();
+  void CpuTimeGraph(ImVec2 plot_size = ImVec2(-1, 0));
-  void DimensionsGraph();
+  void DimensionsGraph(ImVec2 plot_size = ImVec2(-1, 0));
 
  private:
   std::vector<float> cpu_total_;

src/experimental/platform/ux/gui.cc

@@ -97,15 +97,15 @@ void SetupTheme(GuiTheme theme) {
     c[ImGuiCol_TextSelectedBg] = ImVec4(0.73, 0.73, 0.73, 0.35);
     c[ImGuiCol_ModalWindowDimBg] = ImVec4(0.80, 0.80, 0.80, 0.35);
     c[ImGuiCol_DragDropTarget] = ImVec4(1.00, 1.00, 0.00, 0.90);
-    c[ImGuiCol_NavHighlight] = ImVec4(0.26, 0.59, 0.98, 1.00);
+    c[ImGuiCol_NavCursor] = ImVec4(0.26, 0.59, 0.98, 1.00);
     c[ImGuiCol_NavWindowingHighlight] = ImVec4(1.00, 1.00, 1.00, 0.70);
     c[ImGuiCol_NavWindowingDimBg] = ImVec4(0.80, 0.80, 0.80, 0.20);
     c[ImGuiCol_DockingEmptyBg] = ImVec4(0.38, 0.38, 0.38, 1.00);
     c[ImGuiCol_Tab] = ImVec4(0.25, 0.25, 0.25, 1.00);
     c[ImGuiCol_TabHovered] = ImVec4(0.40, 0.40, 0.40, 1.00);
-    c[ImGuiCol_TabActive] = ImVec4(0.33, 0.33, 0.33, 1.00);
+    c[ImGuiCol_TabSelected] = ImVec4(0.33, 0.33, 0.33, 1.00);
-    c[ImGuiCol_TabUnfocused] = ImVec4(0.25, 0.25, 0.25, 1.00);
+    c[ImGuiCol_TabDimmed] = ImVec4(0.25, 0.25, 0.25, 1.00);
-    c[ImGuiCol_TabUnfocusedActive] = ImVec4(0.33, 0.33, 0.33, 1.00);
+    c[ImGuiCol_TabDimmedSelected] = ImVec4(0.33, 0.33, 0.33, 1.00);
     c[ImGuiCol_DockingPreview] = ImVec4(0.85, 0.85, 0.85, 0.28);
     c[ImGuiCol_WindowBg].w = 1.0f;
   } else if (theme == GuiTheme::kLight) {
@@ -192,7 +192,7 @@ void SetupTheme(GuiTheme theme) {
   float rounding = 4.0f;
   s.DisplaySafeAreaPadding = ImVec2(0, 0);
   s.WindowPadding = ImVec2(hspacing, vspacing);
-  s.FramePadding = ImVec2(hspacing, vspacing);
+  s.FramePadding = ImVec2(hspacing, 2);
   s.ItemSpacing = ImVec2(hspacing, vspacing);
   s.ItemInnerSpacing = ImVec2(hspacing, vspacing);
   s.WindowRounding = rounding;
@@ -205,7 +205,7 @@ void SetupTheme(GuiTheme theme) {
   s.WindowBorderSize = 0.0f;
   s.FrameBorderSize = 1.0f;
   s.PopupBorderSize = 1.0f;
-  s.IndentSpacing = 20.0f;
+  s.IndentSpacing = 6.0f;
   s.ScrollbarSize = 12.0f;
   s.GrabMinSize = 5.0f;
   s.WindowMenuButtonPosition = ImGuiDir_None;
@@ -213,15 +213,32 @@ void SetupTheme(GuiTheme theme) {
   s.DockingNodeHasCloseButton = false;
 }
 
+void RescaleDock(float ratio) {
+  if (ratio == 1) return;
+  ImGuiID root = ImGui::GetID("Root");
+  ImGuiDockNode* root_node = ImGui::DockBuilderGetNode(root);
+  if (root_node) {
+    struct ScaleNodes {
+      static void Apply(ImGuiDockNode* node, float r) {
+        node->SizeRef.x *= r;
+        if (node->ChildNodes[0]) Apply(node->ChildNodes[0], r);
+        if (node->ChildNodes[1]) Apply(node->ChildNodes[1], r);
+      }
+    };
+    ScaleNodes::Apply(root_node, ratio);
+  }
+}
+
 ImVec4 ConfigureDockingLayout() {
   ImGuiViewport* viewport = ImGui::GetMainViewport();
   const float scale = ImGui::GetWindowDpiScale();
+  const float font_scale = ImGui::GetIO().FontGlobalScale;
 
   const float kOptionsRelWidth = 0.22f;
   const float kInspectorRelWidth = 0.22f;
   const float kStatsRelHeight = 0.3f;
-  const float kToolsBarHeight = 48.f * scale;
+  const float kToolsBarHeight = 36.f * scale * font_scale;
-  const float kStatusBarHeight = 32.f * scale;
+  const float kStatusBarHeight = 32.f * scale * font_scale;
 
   const ImVec2 dockspace_pos{viewport->WorkPos.x,
                              viewport->WorkPos.y + kToolsBarHeight};
@@ -303,6 +320,10 @@ ImVec4 ConfigureDockingLayout() {
     platform::ScopedStyle style;
     style.Var(ImGuiStyleVar_WindowBorderSize, 1.0f);
     style.Var(ImGuiStyleVar_WindowRounding, 0.0f);
+    style.Var(ImGuiStyleVar_WindowMinSize, ImVec2(1, 1));
+    const float toolbar_vpad =
+        std::max(0.f, (kToolsBarHeight - ImGui::GetFrameHeight()) * 0.5f);
+    style.Var(ImGuiStyleVar_WindowPadding, ImVec2(4, toolbar_vpad));
     ImGui::SetNextWindowPos(viewport->WorkPos, ImGuiCond_Always);
     ImGui::SetNextWindowSize(ImVec2(viewport->Size.x, kToolsBarHeight),
                              ImGuiCond_Always);
@@ -315,6 +336,7 @@ ImVec4 ConfigureDockingLayout() {
     platform::ScopedStyle style;
     style.Var(ImGuiStyleVar_WindowBorderSize, 1.0f);
     style.Var(ImGuiStyleVar_WindowRounding, 0.0f);
+    style.Var(ImGuiStyleVar_WindowMinSize, ImVec2(1, 1));
     ImGui::SetNextWindowPos(ImVec2(0, viewport->Size.y - kStatusBarHeight),
                             ImGuiCond_Always);
     ImGui::SetNextWindowSize(ImVec2(viewport->Size.x, kStatusBarHeight),
@@ -335,18 +357,23 @@ ImVec4 ConfigureDockingLayout() {
 void StepControlGui(const mjModel* model, StepControl* step_control,
                     int& speed_index) {
   platform::ScopedStyle style;
-  style.Var(ImGuiStyleVar_FrameRounding, 2.f);
+  style.Var(ImGuiStyleVar_FrameRounding, 8.f);
 
   const ImColor yellow(255, 215, 0, 255);
   const ImColor green(40, 180, 40, 255);
-  const float scale = ImGui::GetWindowDpiScale();
-  ImVec2 button_size(48.f * scale, 32.f * scale);
 
   auto make_button = [&](const char* icon, StepControl::PauseState target_state,
-                         ImColor color, const char* tooltip = "",
+                         ImColor color, ImDrawFlags corners,
-                         float hover_alpha = 1.f) {
+                         const char* tooltip = "",
+                         float hover_alpha = 1.f, float width_scale = 1.f) {
+    ImVec2 size(0, 0);
+    if (width_scale != 1.f) {
+      const ImGuiStyle& s = ImGui::GetStyle();
+      const float w = ImGui::CalcTextSize(icon).x + s.FramePadding.x * 2;
+      size.x = w * width_scale;
+    }
     bool active = step_control->GetPauseState() == target_state;
-    if (ImGui_ColorButton(icon, active, color, button_size, hover_alpha)) {
+    if (ImGui_ColorButtonEx(icon, active, color, corners, size, hover_alpha)) {
       step_control->SetPauseState(target_state);
     }
     if (!std::string_view(tooltip).empty()) {
@@ -355,26 +382,28 @@ void StepControlGui(const mjModel* model, StepControl* step_control,
   };
 
   make_button(ICON_FA_PAUSE, StepControl::PauseState::kNormalPaused, yellow,
-              "Pause");
+              ImDrawFlags_RoundCornersLeft, "Pause", .3f, 1.6f);
   ImGui::SameLine(0.f, 0.f);
   make_button(ICON_FA_MAGIC, StepControl::PauseState::kViscousPaused, yellow,
-              "Viscous Pause");
+              ImDrawFlags_RoundCornersNone, "Viscous Pause", .3f, 1.3f);
   ImGui::SameLine(0.f, 0.f);
-  make_button(ICON_FA_PLAY, StepControl::PauseState::kUnpaused, green, "", .6f);
+  make_button(ICON_FA_PLAY, StepControl::PauseState::kUnpaused, green,
+              ImDrawFlags_RoundCornersRight, "", .3f, 1.6f);
 
   // Speed selection.
-  ImGui::SameLine();
+  style.Reset();
-  const float pad_y = (button_size.y - ImGui::GetFontSize()) * .5f;
+  ImGui::SameLine(0, ImGui::GetFrameHeight() * .6f);
   ImGui::PushStyleVar(ImGuiStyleVar_FramePadding,
-                      ImVec2(ImGui::GetStyle().FramePadding.x + 5.f, pad_y));
+                      ImVec2(ImGui::GetStyle().FramePadding.x + 5.f,
+                             ImGui::GetStyle().FramePadding.y));
 
   const auto [misaligned, measured] = IsSpeedMisaligned(*step_control);
   char speed_preview[64];
   if (misaligned) {
-    snprintf(speed_preview, sizeof(speed_preview), "%s%s (%-4.1f%%)",
+    snprintf(speed_preview, sizeof(speed_preview), "%s %s (%-4.1f%%)",
              ICON_FA_TACHOMETER, kPercentRealTime[speed_index], measured);
   } else {
-    snprintf(speed_preview, sizeof(speed_preview), "%s%s", ICON_FA_TACHOMETER,
+    snprintf(speed_preview, sizeof(speed_preview), "%s %s", ICON_FA_TACHOMETER,
              kPercentRealTime[speed_index]);
   }
 
@@ -400,39 +429,22 @@ void StepControlGui(const mjModel* model, StepControl* step_control,
   }
 }
 
-bool ThemeSelectGui(GuiTheme* theme) {
+bool ThemeSelectGui(GuiTheme* theme, const ImVec2& size) {
   static constexpr const char* ICON_DARKMODE = ICON_FA_CIRCLE;
   static constexpr const char* ICON_LIGHTMODE = ICON_FA_CIRCLE_O;
   static constexpr const char* ICON_CLASSICMODE = ICON_FA_ADJUST;
   const char* theme_icons[] = {ICON_LIGHTMODE, ICON_DARKMODE, ICON_CLASSICMODE};
   const char* theme_tooltips[] = {"Light Mode", "Dark Mode", "Classic Mode"};
-  const GuiTheme theme_values[] = {
-      GuiTheme::kLight,
-      GuiTheme::kDark,
-      GuiTheme::kClassic,
-  };
-
-  bool changed = false;
 
   int theme_idx = static_cast<int>(*theme);
-  ImGui::SetNextItemWidth(ImGui::CalcTextSize(theme_icons[0]).x +
+  if (ImGui::Button(theme_icons[theme_idx], size)) {
-                          ImGui::GetStyle().FramePadding.x * 2);
+    theme_idx = (theme_idx + 1) % IM_ARRAYSIZE(theme_icons);
-  if (ImGui::BeginCombo("##Theme", theme_icons[theme_idx],
+    *theme = static_cast<GuiTheme>(theme_idx);
-                        ImGuiComboFlags_NoArrowButton)) {
+    return true;
-    for (int n = 0; n < IM_ARRAYSIZE(theme_icons); n++) {
-      if (ImGui::Selectable(theme_icons[n], (theme_idx == n))) {
-        *theme = theme_values[n];
-        changed = true;
-      }
-      if (ImGui::IsItemHovered()) {
-        ImGui::SetTooltip("%s", theme_tooltips[n]);
-      }
-    }
-    ImGui::EndCombo();
   }
-  ImGui::SetItemTooltip("%s", "Theme");
+  ImGui::SetItemTooltip("%s", theme_tooltips[theme_idx]);
 
-  return changed;
+  return false;
 }
 
 bool LabelSelectionGui(mjvOption* opts) {
@@ -506,7 +518,7 @@ bool CameraSelectionGui(const mjModel* model, mjData* data, mjvCamera& camera,
 
   auto select = [&](int type, int idx) {
     if (ImGui::Selectable(GetCameraName(model, camera, type).c_str(),
-            (type == idx))) {
+                          (type == idx))) {
       return true;
     }
     return false;
@@ -817,22 +829,6 @@ void PhysicsGui(mjModel* model, float min_width) {
     ImGui::TreePop();
   }
 
-  if (ImGui::TreeNodeEx("Actuator Groups")) {
-    if (ImGui::BeginTable("##ActuatorGroupsTable", num_cols)) {
-      const ImVec2 size = GetFlexElementSize(num_cols);
-      for (int i = 0; i < 6; ++i) {
-        char label[64];
-        std::snprintf(label, sizeof(label), "Act Group %d", i);
-        ImGui::TableNextColumn();
-        int flipped = ~opt.disableactuator;
-        ImGui_BitToggle(label, &flipped, 1 << i, size);
-        opt.disableactuator = ~flipped;
-      }
-      ImGui::EndTable();
-    }
-    ImGui::TreePop();
-  };
-
   if (ImGui::TreeNodeEx("Algorithmic Parameters")) {
     ImGui_Input("Timestep", &opt.timestep, {0, 1, 0.01, 0.1});
     ImGui_Input("Iterations", &opt.iterations, {0, 1000, 1, 10});
@@ -867,6 +863,22 @@ void PhysicsGui(mjModel* model, float min_width) {
     ImGui::TreePop();
   }
 
+  if (ImGui::TreeNodeEx("Actuator Groups")) {
+    if (ImGui::BeginTable("##ActuatorGroupsTable", num_cols)) {
+      const ImVec2 size = GetFlexElementSize(num_cols);
+      for (int i = 0; i < 6; ++i) {
+        char label[64];
+        std::snprintf(label, sizeof(label), "Act Group %d", i);
+        ImGui::TableNextColumn();
+        int flipped = ~opt.disableactuator;
+        ImGui_BitToggle(label, &flipped, 1 << i, size);
+        opt.disableactuator = ~flipped;
+      }
+      ImGui::EndTable();
+    }
+    ImGui::TreePop();
+  }
+
   ImGui::PopItemWidth();
 }
 
@@ -986,8 +998,6 @@ void RenderingGui(const mjModel* model, mjvOption* vis_options,
       static_cast<int>(std::floor(available_width / min_width)), 1, 6);
 
   if (ImGui::TreeNodeEx("Model Elements", ImGuiTreeNodeFlags_DefaultOpen)) {
-    ImGui::Unindent(ImGui::GetTreeNodeToLabelSpacing() / 2);
-
     if (ImGui::BeginTable("##ModelElementsTable", num_cols)) {
       const ImVec2 size = GetFlexElementSize(num_cols);
       for (int i = 0; i < mjNVISFLAG; ++i) {
@@ -996,14 +1006,10 @@ void RenderingGui(const mjModel* model, mjvOption* vis_options,
       }
       ImGui::EndTable();
     }
-
-    ImGui::Indent(ImGui::GetTreeNodeToLabelSpacing() / 2);
     ImGui::TreePop();
   }
 
   if (ImGui::TreeNodeEx("Render Flags", ImGuiTreeNodeFlags_DefaultOpen)) {
-    ImGui::Unindent(ImGui::GetTreeNodeToLabelSpacing() / 2);
-
     if (ImGui::BeginTable("##RenderFlagsTable", num_cols)) {
       const ImVec2 size = GetFlexElementSize(num_cols);
       for (int i = 0; i < mjNRNDFLAG; ++i) {
@@ -1012,8 +1018,6 @@ void RenderingGui(const mjModel* model, mjvOption* vis_options,
       }
       ImGui::EndTable();
     }
-
-    ImGui::Indent(ImGui::GetTreeNodeToLabelSpacing() / 2);
     ImGui::TreePop();
   }
 }
@@ -1030,8 +1034,6 @@ void GroupsGui(const mjModel* model, mjvOption* vis_options, float min_width) {
 
   auto GroupGui = [&](const char* name, mjtByte* group) {
     if (ImGui::TreeNodeEx(name, ImGuiTreeNodeFlags_DefaultOpen)) {
-      ImGui::Unindent(ImGui::GetTreeNodeToLabelSpacing() / 2);
-
       char label[64];
       std::snprintf(label, sizeof(label), "##%s", name);
       if (ImGui::BeginTable(label, num_cols)) {
@@ -1044,8 +1046,6 @@ void GroupsGui(const mjModel* model, mjvOption* vis_options, float min_width) {
 
         ImGui::EndTable();
       }
-
-      ImGui::Indent(ImGui::GetTreeNodeToLabelSpacing() / 2);
       ImGui::TreePop();
     }
   };
@@ -1150,12 +1150,24 @@ void ControlsGui(const mjModel* model, const mjData* data,
   ImGui::PopItemWidth();
 }
 
-void ConvergenceGui(const mjModel* model, mjData* data) {
+static int GetPlotXLimit(const mjData* data) {
-  if (ImPlot::BeginPlot("Convergence (log 10)", ImVec2(-1, 0),
+  int max_niter = 0;
-                        ImPlotFlags_NoMouseText)) {
+  const int nisland0 =
+      data->nefc ? mjMAX(1, mjMIN(data->nisland, mjNISLAND)) : 0;
+  for (int k = 0; k < nisland0; k++) {
+    max_niter = mjMAX(max_niter, data->solver_niter[k]);
+  }
+  return mjMAX(10, ((max_niter + 9) / 10) * 10);
+}
+
+void ConvergenceGui(const mjModel* model, mjData* data, ImVec2 plot_size) {
+  int xlim = GetPlotXLimit(data);
+  ImPlotFlags flags =
+      ImPlot_SetupPlotFlags(plot_size) | ImPlotFlags_NoMouseText;
+  if (ImPlot::BeginPlot("Convergence (log 10) vs iter", plot_size, flags)) {
     ImPlot::PushStyleVar(ImPlotStyleVar_LineWeight, 2.0f);
-    ImPlot::SetupAxis(ImAxis_X1, "iteration", ImPlotAxisFlags_AutoFit);
+    ImPlot::SetupAxis(ImAxis_X1, "", ImPlotAxisFlags_AutoFit);
-    ImPlot::SetupAxisLimits(ImAxis_X1, 0, 20, ImPlotCond_Always);
+    ImPlot::SetupAxisLimits(ImAxis_X1, 0, xlim, ImPlotCond_Always);
     ImPlot::SetupAxisFormat(ImAxis_Y1, "%.1f");
     ImPlot::SetupAxisLimits(ImAxis_Y1, -20, 5, ImPlotCond_Always);
     ImPlot::SetupLegend(ImPlotLocation_NorthEast);
@@ -1211,11 +1223,14 @@ void ConvergenceGui(const mjModel* model, mjData* data) {
   }
 }
 
-void CountsGui(const mjModel* model, mjData* data) {
+void CountsGui(const mjModel* model, mjData* data, ImVec2 plot_size) {
-  if (ImPlot::BeginPlot("Counts", ImVec2(-1, 0), ImPlotFlags_NoMouseText)) {
+  int xlim = GetPlotXLimit(data);
+  ImPlotFlags flags =
+      ImPlot_SetupPlotFlags(plot_size) | ImPlotFlags_NoMouseText;
+  if (ImPlot::BeginPlot("Counts vs iter", plot_size, flags)) {
     ImPlot::PushStyleVar(ImPlotStyleVar_LineWeight, 2.0f);
-    ImPlot::SetupAxis(ImAxis_X1, "iteration", ImPlotAxisFlags_AutoFit);
+    ImPlot::SetupAxis(ImAxis_X1, "", ImPlotAxisFlags_AutoFit);
-    ImPlot::SetupAxisLimits(ImAxis_X1, 0, 20, ImPlotCond_Always);
+    ImPlot::SetupAxisLimits(ImAxis_X1, 0, xlim, ImPlotCond_Always);
     ImPlot::SetupAxisFormat(ImAxis_Y1, "%.0f");
     ImPlot::SetupAxisLimits(ImAxis_Y1, 0, 80, ImPlotCond_Always);
     ImPlot::SetupLegend(ImPlotLocation_NorthEast);

src/experimental/platform/ux/gui.h

@@ -42,6 +42,9 @@ enum class GuiTheme {
 // Updates the ImGui internal style state to match the requested theme.
 void SetupTheme(GuiTheme theme);
 
+// Rescales all dock node widths by the given ratio.
+void RescaleDock(float ratio);
+
 // Configures the ImGui docking module to the standard layout used by Studio.
 // This includes the following named sections:
 //   "ToolBar": fixed size bar spanning the top of the window; for placing
@@ -81,7 +84,7 @@ void StepControlGui(const mjModel* model, StepControl* step_control,
                     int& speed_index);
 
 // UX for selecting the GUI theme.
-bool ThemeSelectGui(GuiTheme* theme);
+bool ThemeSelectGui(GuiTheme* theme, const ImVec2& size = ImVec2(0, 0));
 
 // UX for selecting the visualization label option.
 bool LabelSelectionGui(mjvOption* opts);
@@ -138,10 +141,12 @@ void NoiseGui(const mjModel* model, const mjData* data, float& noise_scale,
               float& noise_rate);
 
 // UX for the solver convergence chart.
-void ConvergenceGui(const mjModel* model, mjData* data);
+void ConvergenceGui(const mjModel* model, mjData* data,
+                    ImVec2 plot_size = ImVec2(-1, 0));
 
 // UX for the solver counts chart.
-void CountsGui(const mjModel* model, mjData* data);
+void CountsGui(const mjModel* model, mjData* data,
+               ImVec2 plot_size = ImVec2(-1, 0));
 
 // UX for displaying basic simulation information. Note that the pause state and
 // FPS needs to be tracked by the caller and passed here to be displayed.

src/experimental/platform/ux/imgui_widgets.cc

@@ -14,6 +14,7 @@
 
 #include "experimental/platform/ux/imgui_widgets.h"
 
+#include <algorithm>
 #include <cstdint>
 #include <cstring>
 #include <sstream>
@@ -23,6 +24,7 @@
 #include <vector>
 
 #include <imgui.h>
+#include <implot.h>
 #include <mujoco/mujoco.h>
 
 namespace mujoco::platform {
@@ -373,9 +375,75 @@ void ImGui_EndHSplit(bool open) {
 }
 
 void MaybeSaveToClipboard(const std::string& contents) {
-  if (ImGui::GetIO().SetClipboardTextFn) {
+  ImGui::SetClipboardText(contents.c_str());
-    ImGui::GetIO().SetClipboardTextFn(nullptr, contents.c_str());
+}
+
+ImPlotFlags ImPlot_SetupPlotFlags(ImVec2 plot_size) {
+  ImPlotFlags flags = ImPlotFlags_None;
+  if (plot_size.x > 0 && plot_size.y > 0) {
+    const float min_dim = std::min(plot_size.x, plot_size.y);
+    if (min_dim < 300) {
+      flags |= ImPlotFlags_NoTitle;
+    }
+    if (min_dim < 200) {
+      flags |= ImPlotFlags_NoLegend;
+    }
+  }
+  return flags;
+}
+
+void ImPlot_SetupTimeAxis(ImVec2 plot_size, const char* label,
+                          ImPlotAxisFlags extra_flags) {
+  ImPlotAxisFlags flags = extra_flags;
+  if (plot_size.x > 0 && plot_size.x < 300) {
+    flags |= ImPlotAxisFlags_NoTickLabels;
+  }
+  ImPlot::SetupAxis(ImAxis_X1, label, flags);
+}
+
+void ImPlot_SetupValueAxis(ImVec2 plot_size, const char* label,
+                           const char* format, ImPlotAxisFlags extra_flags) {
+  ImPlotAxisFlags flags = extra_flags;
+  if (plot_size.y > 0 && plot_size.y < 150) {
+    flags |= ImPlotAxisFlags_NoTickLabels;
+  }
+  ImPlot::SetupAxis(ImAxis_Y1, label, flags);
+  if (format) {
+    ImPlot::SetupAxisFormat(ImAxis_Y1, format);
   }
 }
 
+void ImPlot_SetupFixedAxis(ImVec2 plot_size, double y_min, double y_max,
+                           const char* label, const char* format,
+                           const double* tick_values,
+                           const char* const* tick_labels, int n_ticks) {
+  ImPlotAxisFlags flags = ImPlotAxisFlags_None;
+  if (plot_size.y > 0 && plot_size.y < 150) {
+    flags |= ImPlotAxisFlags_NoTickLabels;
+  }
+  ImPlot::SetupAxis(ImAxis_Y1, label, flags);
+  ImPlot::SetupAxisLimits(ImAxis_Y1, y_min, y_max, ImPlotCond_Always);
+  if (format) {
+    ImPlot::SetupAxisFormat(ImAxis_Y1, format);
+  }
+  if (tick_values && n_ticks > 0) {
+    ImPlot::SetupAxisTicks(ImAxis_Y1, tick_values, n_ticks, tick_labels);
+  }
+}
+
+ImPlotPairLayout ImPlot_ComputePairLayout() {
+  ImVec2 avail = ImGui::GetContentRegionAvail();
+  bool is_wide = avail.x > avail.y;
+
+  const float item_spacing = ImGui::GetStyle().ItemSpacing.y;
+  ImVec2 plot_size(is_wide ? (avail.x - item_spacing) * 0.5f : avail.x,
+                   is_wide ? avail.y : (avail.y - item_spacing) * 0.5f);
+
+  return {
+      plot_size,
+      is_wide ? ImPlotLayoutDirection::kHorizontal
+              : ImPlotLayoutDirection::kVertical,
+  };
+}
+
 }  // namespace mujoco::platform

src/experimental/platform/ux/imgui_widgets.h

@@ -27,6 +27,7 @@
 #include <imgui.h>
 #include <imgui_internal.h>
 #include <misc/cpp/imgui_stdlib.h>
+#include <implot.h>
 #include <mujoco/mujoco.h>
 #include "experimental/platform/ux/enum_utils.h"
 
@@ -550,6 +551,55 @@ inline bool ImGui_ColorButton(const char* label, bool active, ImColor color,
   return ImGui::Button(label, size);
 }
 
+// Like ImGui_ColorButton, but with per-corner rounding control via ImDrawFlags.
+// Use ImDrawFlags_RoundCornersLeft, ImDrawFlags_RoundCornersRight,
+// ImDrawFlags_RoundCornersNone, ImDrawFlags_RoundCornersAll, etc.
+inline bool ImGui_ColorButtonEx(const char* label, bool active, ImColor color,
+                                ImDrawFlags corners,
+                                const ImVec2& size = ImVec2(0, 0),
+                                float hover_alpha = 0.5f) {
+  const ImGuiStyle& s = ImGui::GetStyle();
+  const ImVec2 label_size = ImGui::CalcTextSize(label, nullptr, true);
+  const ImVec2 btn_size(
+      size.x > 0 ? size.x : label_size.x + s.FramePadding.x * 2,
+      size.y > 0 ? size.y : label_size.y + s.FramePadding.y * 2);
+
+  const ImVec2 pos = ImGui::GetCursorScreenPos();
+  ImGui::InvisibleButton(label, btn_size);
+  const bool clicked = ImGui::IsItemClicked();
+  const bool hovered = ImGui::IsItemHovered();
+
+  // Determine background color.
+  const ImColor hover_color(color.Value.x, color.Value.y, color.Value.z,
+                            color.Value.w * hover_alpha);
+  ImColor bg;
+  if (active) {
+    bg = color;
+  } else if (hovered) {
+    bg = hover_color;
+  } else {
+    bg = ImGui::GetColorU32(ImGuiCol_Button);
+  }
+
+  // Draw background with per-corner rounding.
+  ImDrawList* dl = ImGui::GetWindowDrawList();
+  const ImVec2 max(pos.x + btn_size.x, pos.y + btn_size.y);
+  dl->AddRectFilled(pos, max, bg, s.FrameRounding, corners);
+
+  // Draw border.
+  if (s.FrameBorderSize > 0) {
+    dl->AddRect(pos, max, ImGui::GetColorU32(ImGuiCol_Border),
+                s.FrameRounding, corners, s.FrameBorderSize);
+  }
+
+  // Draw label centered.
+  const ImVec2 text_pos(pos.x + (btn_size.x - label_size.x) * 0.5f,
+                        pos.y + (btn_size.y - label_size.y) * 0.5f);
+  dl->AddText(text_pos, ImGui::GetColorU32(ImGuiCol_Text), label);
+
+  return clicked;
+}
+
 // Begin a boxed section with outer borders - use EndBoxSection to close.
 inline bool BeginBoxSection(const char* id, ImGuiTableFlags extra_flags = 0) {
   ImGuiTableFlags flags = ImGuiTableFlags_BordersOuter | extra_flags;
@@ -566,6 +616,46 @@ inline void EndBoxSection() { ImGui::EndTable(); }
 // Saves the given contents to the clipboard if the clipboard is available.
 void MaybeSaveToClipboard(const std::string& contents);
 
+// Returns plot flags with title/legend conditionally hidden when the plot
+// area is too small.  `plot_size` is the final rendered size of the plot.
+ImPlotFlags ImPlot_SetupPlotFlags(ImVec2 plot_size);
+
+// Sets up the X axis as a "time/frame" axis.
+// Hides tick labels when the plot is narrow.
+// Uses `label` as the axis label (empty string to hide) and auto-fit limits.
+void ImPlot_SetupTimeAxis(
+    ImVec2 plot_size, const char* label = "",
+    ImPlotAxisFlags extra_flags = ImPlotAxisFlags_AutoFit);
+
+// Sets up a Y axis with auto-fit limits.
+// Hides tick labels when the plot is short.
+void ImPlot_SetupValueAxis(
+    ImVec2 plot_size, const char* label = "", const char* format = nullptr,
+    ImPlotAxisFlags extra_flags = ImPlotAxisFlags_AutoFit);
+
+// Sets up a Y axis with fixed limits and optional explicit ticks.
+// Hides tick labels when the plot is short.
+void ImPlot_SetupFixedAxis(ImVec2 plot_size, double y_min, double y_max,
+                           const char* label = "", const char* format = nullptr,
+                           const double* tick_values = nullptr,
+                           const char* const* tick_labels = nullptr,
+                           int n_ticks = 0);
+
+enum class ImPlotLayoutDirection {
+  kHorizontal,
+  kVertical,
+};
+
+struct ImPlotPairLayout {
+  ImVec2 plot_size;  // Size for each individual plot.
+  ImPlotLayoutDirection direction;
+};
+
+// Computes a responsive layout for two plots that share the available
+// content region.  When the region is wider than tall, the plots are placed
+// side-by-side; otherwise they are stacked vertically.
+ImPlotPairLayout ImPlot_ComputePairLayout();
+
 }  // namespace mujoco::platform
 
 #endif  // MUJOCO_SRC_EXPERIMENTAL_PLATFORM_UX_IMGUI_WIDGETS_H_

src/experimental/studio/app.cc

@@ -72,13 +72,11 @@ static void SelectParentPerturb(const mjModel* model, mjvPerturb& perturb) {
 }
 
 static constexpr const char* ICON_COPY_CAMERA = platform::ICON_FA_COPY;
-static constexpr const char* ICON_UNLOAD_MODEL = platform::ICON_FA_EJECT;
 static constexpr const char* ICON_RELOAD_MODEL = platform::ICON_FA_REFRESH;
 static constexpr const char* ICON_RESET_MODEL = platform::ICON_FA_UNDO;
 static constexpr const char* ICON_PREV_FRAME = platform::ICON_FA_CARET_LEFT;
 static constexpr const char* ICON_NEXT_FRAME = platform::ICON_FA_CARET_RIGHT;
 static constexpr const char* ICON_CURR_FRAME = platform::ICON_FA_FAST_FORWARD;
-static constexpr const char* ICON_RELOAD_SPEC = platform::ICON_FA_REFRESH;
 static constexpr const char* ICON_UNDO_SPEC = platform::ICON_FA_UNDO;
 static constexpr const char* ICON_REDO_SPEC = platform::ICON_FA_REPEAT;
 
@@ -580,6 +578,14 @@ void App::HandleKeyboardEvents() {
     tmp_.inspector_panel = !tmp_.inspector_panel;
   } else if (ImGui_IsChordJustPressed(ImGuiKey_Tab)) {
     tmp_.options_panel = !tmp_.options_panel;
+  } else if (ImGui_IsChordJustPressed(ImGuiKey_Minus | ImGuiMod_Ctrl)) {
+    float old_scale = ui_.font_scale;
+    ui_.font_scale = std::clamp(ui_.font_scale - 0.1f, 0.5f, 3.0f);
+    platform::RescaleDock(ui_.font_scale / old_scale);
+  } else if (ImGui_IsChordJustPressed(ImGuiKey_Equal | ImGuiMod_Ctrl)) {
+    float old_scale = ui_.font_scale;
+    ui_.font_scale = std::clamp(ui_.font_scale + 0.1f, 0.5f, 3.0f);
+    platform::RescaleDock(ui_.font_scale / old_scale);
   } else if (ImGui_IsChordJustPressed(ImGuiKey_Minus)) {
     SetSpeedIndex(tmp_.speed_index + 1);
   } else if (ImGui_IsChordJustPressed(ImGuiKey_Equal)) {
@@ -825,6 +831,8 @@ void App::BuildGui() {
     platform::SetupTheme(ui_.theme);
   }
 
+  ImGui::GetIO().FontGlobalScale = ui_.font_scale;
+
   const ImVec4 workspace_rect = platform::ConfigureDockingLayout();
 
   // Place charts in bottom right corner of the workspace.
@@ -884,8 +892,12 @@ void App::BuildGui() {
     ImGui::SetNextWindowPos(chart_pos, ImGuiCond_FirstUseEver);
     ImGui::SetNextWindowSize(chart_size, ImGuiCond_FirstUseEver);
     if (ImGui::Begin("Performance", &tmp_.chart_performance)) {
-      profiler_.CpuTimeGraph();
+      auto layout = platform::ImPlot_ComputePairLayout();
-      profiler_.DimensionsGraph();
+      profiler_.CpuTimeGraph(layout.plot_size);
+      if (layout.direction == platform::ImPlotLayoutDirection::kHorizontal) {
+        ImGui::SameLine();
+      }
+      profiler_.DimensionsGraph(layout.plot_size);
     }
     ImGui::End();
   }
@@ -894,8 +906,12 @@ void App::BuildGui() {
     ImGui::SetNextWindowPos(chart_pos, ImGuiCond_FirstUseEver);
     ImGui::SetNextWindowSize(chart_size, ImGuiCond_FirstUseEver);
     if (ImGui::Begin("Solver", &tmp_.chart_solver)) {
-      platform::CountsGui(model(), data());
+      auto layout = platform::ImPlot_ComputePairLayout();
-      platform::ConvergenceGui(model(), data());
+      platform::CountsGui(model(), data(), layout.plot_size);
+      if (layout.direction == platform::ImPlotLayoutDirection::kHorizontal) {
+        ImGui::SameLine();
+      }
+      platform::ConvergenceGui(model(), data(), layout.plot_size);
     }
     ImGui::End();
   }
@@ -1011,14 +1027,14 @@ void App::ModelOptionsGui() {
       ImGuiTreeNodeFlags_SpanAvailWidth | ImGuiTreeNodeFlags_Framed;
 
   ImGui::BeginChild("PhysicsGui", {0, 0}, child_flags);
-  if (ImGui::TreeNodeEx("Physics Settings", node_flags)) {
+  if (ImGui::TreeNodeEx("Physics", node_flags)) {
     platform::PhysicsGui(model(), min_width);
     ImGui::TreePop();
   }
   ImGui::EndChild();
 
   ImGui::BeginChild("RenderingGui", {0, 0}, child_flags);
-  if (ImGui::TreeNodeEx("Rendering Settings", node_flags)) {
+  if (ImGui::TreeNodeEx("Rendering", node_flags)) {
     platform::RenderingGui(model(), &vis_options_, renderer_->GetRenderFlags(),
                            min_width);
     ImGui::TreePop();
@@ -1394,13 +1410,12 @@ void App::HelpGui() {
 }
 
 void App::ToolBarGui() {
+  ImGui::PushStyleVar(ImGuiStyleVar_CellPadding, ImVec2(0, 0));
   if (ImGui::BeginTable("##ToolBarTable", 2)) {
     platform::ScopedStyle style;
-    const ImColor red(220, 40, 40, 255);
+    style.Var(ImGuiStyleVar_ItemSpacing,
-
+              ImVec2(ImGui::GetStyle().ItemSpacing.x * 2.0f,
-    const float scale = ImGui::GetWindowDpiScale();
+                     ImGui::GetStyle().ItemSpacing.y));
-    const ImVec2 button_size(48.f * scale, 32.f * scale);
-    const ImVec2 play_button_size(80.f * scale, 32.f * scale);
 
     const float label_width = GetExpectedLabelWidth();
     const float copy_btn_width = ImGui::CalcTextSize(ICON_COPY_CAMERA).x +
@@ -1412,56 +1427,40 @@ void App::ToolBarGui() {
     const float right_width = label_width + sp + label_width + sp +
                               label_width + sp + copy_btn_width + sp +
                               theme_width;
-    const float separator_width = .2f * button_size.x;
+    const float separator_width = ImGui::GetFrameHeight() * .6f;
 
     ImGui::TableSetupColumn("", ImGuiTableColumnFlags_WidthStretch);
     ImGui::TableSetupColumn("", ImGuiTableColumnFlags_WidthFixed, right_width);
 
     ImGui::TableNextColumn();
+    ImGui::SetCursorPosX(ImGui::GetCursorPosX() + ImGui::GetStyle().WindowPadding.x);
 
-    // Combined (Unload, Reload) widget
+    const float btn_size = ImGui::GetFrameHeight();
+    const ImVec2 square_size(btn_size, btn_size);
+
+    // Reload button.
     {
       style.Var(ImGuiStyleVar_FrameRounding, 2.0f);
-
+      if (ImGui::Button(ICON_RELOAD_MODEL, square_size)) {
-      // Unload button.
-      {
-        const ImColor a = red;
-        const ImColor h(a.Value.x, a.Value.y, a.Value.z, a.Value.w * 0.6f);
-        style.Color(ImGuiCol_ButtonHovered, h);
-        style.Color(ImGuiCol_ButtonActive, a);
-
-        if (ImGui::Button(ICON_UNLOAD_MODEL, button_size)) {
-          InitEmptyModel();
-        }
-        ImGui::SetItemTooltip("%s", "Unload");
-        style.Reset();
-      }
-
-      // Reload button.
-      ImGui::SameLine(0, 0);
-      if (ImGui::Button(ICON_RELOAD_MODEL, button_size)) {
         RequestModelReload();
       }
       ImGui::SetItemTooltip("%s", "Reload");
     }
 
     // Reset button.
-    ImGui::SameLine(0, separator_width);
+    ImGui::SameLine(0, 0.5 * separator_width);
-    if (ImGui::Button(ICON_RESET_MODEL, button_size)) {
+    if (ImGui::Button(ICON_RESET_MODEL, square_size)) {
       ResetPhysics();
     }
     ImGui::SetItemTooltip("%s", "Reset");
 
-
     // Combined (Normal Pause, Viscous Pause, Play) widget and Speed selection.
     ImGui::SameLine(0, separator_width);
     platform::StepControlGui(model(), &step_control_, tmp_.speed_index);
 
     ImGui::TableNextColumn();
-    ImGui::SetCursorPosY(ImGui::GetCursorPosY() +
-                         (button_size.y - ImGui::GetFrameHeight()) * 0.5f);
 
-    if (ImGui::Button(ICON_COPY_CAMERA)) {
+    if (ImGui::Button(ICON_COPY_CAMERA, square_size)) {
       std::string camera_string = platform::CameraToString(data(), &camera_);
       platform::MaybeSaveToClipboard(camera_string);
     }
@@ -1481,13 +1480,14 @@ void App::ToolBarGui() {
 
     ImGui::SameLine();
     ImGui::SetNextItemWidth(GetExpectedLabelWidth());
-    if (platform::ThemeSelectGui(&ui_.theme)) {
+    if (platform::ThemeSelectGui(&ui_.theme, square_size)) {
       platform::SetupTheme(ui_.theme);
       ImGui::GetIO().WantSaveIniSettings = true;
     }
 
     ImGui::EndTable();
   }
+  ImGui::PopStyleVar();
 }
 
 void App::StatusBarGui() {
@@ -1851,11 +1851,13 @@ float App::GetExpectedLabelWidth() {
 App::UiState::Dict App::UiState::ToDict() const {
   return {
       {"theme", std::to_string(static_cast<int>(theme))},
+      {"font_scale", std::to_string(font_scale)},
   };
 }
 
 void App::UiState::FromDict(const Dict& dict) {
   *this = UiState();
   theme = ReadIniValue(dict, "theme", theme);
+  font_scale = platform::ReadIniValue(dict, "font_scale", font_scale);
 }
 }  // namespace mujoco::studio

src/experimental/studio/app.h

@@ -105,6 +105,7 @@ class App {
     int camera_idx = platform::kTumbleCameraIdx;
     int key_idx = 0;
     platform::GuiTheme theme = platform::GuiTheme::kLight;
+    float font_scale = 1.0f;
 
     using Dict = std::unordered_map<std::string, std::string>;
     Dict ToDict() const;

src/user/user_api.cc

@@ -20,6 +20,7 @@
 #include <cstdio>
 #include <cstdlib>
 #include <cstring>
+#include <filesystem>
 #include <functional>
 #include <iterator>
 #include <map>
@@ -143,6 +144,18 @@ mjSpec* mj_parse(const char* filename, const char* content_type,
 int mj_encode(const mjSpec* s, const mjModel* m, const char* filename,
               const char* content_type, const mjVFS* vfs, char* error,
               int error_sz) {
+  // TODO(shaves) Move MJCF and URDF to encoders/decoders.
+  auto filepath = mujoco::user::FilePath(filename);
+  if (filepath.Ext() == ".xml" ||
+      (content_type && std::strcmp(content_type, "text/xml") == 0)) {
+    int result = mj_saveXML(s, filename, error, error_sz);
+    if (result < 0) {
+      return -1;
+    }
+
+    return std::filesystem::file_size(filename);
+  }
+
   const mjpEncoder* encoder = mjp_findEncoder(filename, content_type);
   if (!encoder) {
     if (error) {

test/engine/engine_island_test.cc

@@ -15,6 +15,7 @@
 // Tests for engine/engine_island.c.
 
 #include <string>
+#include <vector>
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
@@ -602,5 +603,41 @@ TEST_F(IslandTest, EqualityConstraintOfTendons) {
   mj_deleteModel(model);
 }
 
+TEST_F(IslandTest, PGSIslandExact) {
+  const std::string xml_path = GetTestDataFilePath(kIlslandEfcPath);
+  char error[1024];
+  mjModel* m = mj_loadXML(xml_path.c_str(), nullptr, error, sizeof(error));
+  ASSERT_THAT(m, NotNull()) << error;
+  mjData* d = mj_makeData(m);
+
+  // simulate to get a non-trivial state
+  while (d->time < 0.5) {
+    mj_step(m, d);
+  }
+
+  // switch to PGS, disable early termination
+  m->opt.solver = mjSOL_PGS;
+  m->opt.tolerance = 0;
+
+  // solve with islands
+  m->opt.disableflags &= ~mjDSBL_ISLAND;
+  mj_forward(m, d);
+  ASSERT_GT(d->nisland, 1);
+  std::vector<mjtNum> qfrc_island(d->qfrc_constraint,
+                                  d->qfrc_constraint + m->nv);
+
+  // solve without islands
+  m->opt.disableflags |= mjDSBL_ISLAND;
+  mj_forward(m, d);
+  std::vector<mjtNum> qfrc_mono(d->qfrc_constraint,
+                                d->qfrc_constraint + m->nv);
+
+  // expect exact match
+  EXPECT_EQ(qfrc_island, qfrc_mono);
+
+  mj_deleteData(d);
+  mj_deleteModel(m);
+}
+
 }  // namespace
 }  // namespace mujoco