mujoco

// 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.

//  Python bindings for the MuJoCo filament renderer (mjrf) API.

#include <Python.h>

#include <cstdint>
#include <cstring>
#include <memory>
#include <string>
#include <string_view>
#include <type_traits>
#include <vector>

#include <Eigen/Core>
#include <Eigen/Eigen>
#include <mujoco/mjrfilament.h>
#include <mujoco/mujoco.h>
#include "render/filament/support/model_decorations.h"
#include "render/filament/support/model_lights.h"
#include "render/filament/support/model_objects.h"
#include "render/filament/support/model_renderables.h"
#include "structs.h"
#include <pybind11/eigen.h>
#include <pybind11/eigen/matrix.h>
#include <pybind11/stl.h>

// Note: this file includes autogenerated code from:
// third_party/py/mujoco/codegen/generate_mjrf_bindings.py.
//
// The mjrf C API is composed of three types: enums, structs, and functions.
// This file generates Python wrappers/bindings for these types. Furthermore,
// it attempts to do so such that the resulting bindings are as pythonic as
// possible.
//
// Binding enums is done directly, though we remove redundant prefixes. So, e.g.
// mjrGraphicsApi becomes mjrf.GraphicsApi and mjGRAPHICS_API_DEFAULT becomes
// mjrf.GraphicsApi.GRAPHICS_API_DEFAULT.
//
// Structs are treated in two different ways. Most structs (e.g. mjrRect,
// mjrCamera, etc.) are considered to be simple POD data types. We bind these
// types directly to python and let python manage their lifetimes, creating
// copies as needed.
//
// Some structs are treated as "Objects". These map to the types that have
// explicit create/destroy functions in the C API (e.g. mjrfContext,
// mjrfTexture, mjrfLight, etc.). When created in python, their lifetimes are
// managed by a pybind11::capsule. However, when they are returned from an API
// call, they are returned such that python does not manage their lifetimes.
//
// Then, we bind the functions themselves. In almost all cases, each mjrf
// function takes a pointer to an object as its first argument. In this case,
// we actually bind the function as a method of the object itself.  So, e.g.
// `mjrf_setLightEnabled(mjrfLight*, bool)` becomes `Light.set_enabled(bool)`.
//
// Finally, to bridge the gap between Python tuple/list/numpy array types and
// the native C types (e.g. float*), we use pybind's Eigen bindings.

namespace mujoco::python {
namespace {

// Creates a new object of type T using one of the mjrf_create functions.
template <typename T, typename... Args>
T* Create(Args&&... args) {
  if constexpr (std::is_same_v<T, mjrfContext>) {
    return mjrf_createContext(std::forward<Args>(args)...);
  } else if constexpr (std::is_same_v<T, mjrfTexture>) {
    return mjrf_createTexture(std::forward<Args>(args)...);
  } else if constexpr (std::is_same_v<T, mjrfMesh>) {
    return mjrf_createMesh(std::forward<Args>(args)...);
  } else if constexpr (std::is_same_v<T, mjrfScene>) {
    return mjrf_createScene(std::forward<Args>(args)...);
  } else if constexpr (std::is_same_v<T, mjrfLight>) {
    return mjrf_createLight(std::forward<Args>(args)...);
  } else if constexpr (std::is_same_v<T, mjrfRenderable>) {
    return mjrf_createRenderable(std::forward<Args>(args)...);
  } else if constexpr (std::is_same_v<T, mjrfRenderTarget>) {
    return mjrf_createRenderTarget(std::forward<Args>(args)...);
  } else {
    static_assert(!sizeof(T), "Unsupported type for Create");
  }
}

// Corresponding Destroy function matching the above Create.
void Destroy(mjrfContext* obj)         { mjrf_destroyContext(obj); }
void Destroy(mjrfTexture* obj)         { mjrf_destroyTexture(obj); }
void Destroy(mjrfMesh* obj)            { mjrf_destroyMesh(obj); }
void Destroy(mjrfScene* obj)           { mjrf_destroyScene(obj); }
void Destroy(mjrfLight* obj)           { mjrf_destroyLight(obj); }
void Destroy(mjrfRenderable* obj)      { mjrf_destroyRenderable(obj); }
void Destroy(mjrfRenderTarget* obj)    { mjrf_destroyRenderTarget(obj); }

// Initializers for structs.
void Init(mjrfContextConfig& obj)      { mjrf_defaultContextConfig(&obj); }
void Init(mjrfFrameStats& obj)         { mjrf_defaultFrameStats(&obj); }
void Init(mjrfLightParams& obj)        { mjrf_defaultLightParams(&obj); }
void Init(mjrfMaterial& obj)           { mjrf_defaultMaterial(&obj); }
void Init(mjrfMeshConfig& obj)         { mjrf_defaultMeshConfig(&obj); }
void Init(mjrfMeshData& obj)           { mjrf_defaultMeshData(&obj); }
void Init(mjrfReadPixelsRequest& obj)  { mjrf_defaultReadPixelsRequest(&obj); }
void Init(mjrfRenderableParams& obj)   { mjrf_defaultRenderableParams(&obj); }
void Init(mjrfRenderRequest& obj)      { mjrf_defaultRenderRequest(&obj); }
void Init(mjrfRenderTargetConfig& obj) { mjrf_defaultRenderTargetConfig(&obj); }
void Init(mjrfSceneParams& obj)        { mjrf_defaultSceneParams(&obj); }
void Init(mjrfTextureConfig& obj)      { mjrf_defaultTextureConfig(&obj); }
void Init(mjrfTextureData& obj)        { mjrf_defaultTextureData(&obj); }
void Init(mjrRect& obj)                { std::memset(&obj, 0, sizeof(obj)); }
void Init(mjrCamera& obj)              { std::memset(&obj, 0, sizeof(obj)); }
void Init(mjrVertexAttribute& obj)     { std::memset(&obj, 0, sizeof(obj)); }

// Either owns the lifetime of the wrapped object (using a PyCapsule) or acts as
// a simple wrapper around an existing instance of the object. The former is
// used when creating objects in python code.  The latter is used when returning
// objects from the C API.
#define DEFINE_WRAPPED_OBJECT(PY_TYPE_, TYPE_)                                \
  class PY_TYPE_ : public _impl::WrapperBase<TYPE_> {                         \
    static void DestroyCapsule(PyObject* pyobj) {                             \
      void* ptr = (PyCapsule_GetPointer(pyobj, nullptr));                     \
      if (ptr) {                                                              \
        Destroy(static_cast<TYPE_*>(ptr));                                    \
      }                                                                       \
    }                                                                         \
                                                                              \
   public:                                                                    \
    template <typename... Args>                                               \
    PY_TYPE_(Args&&... args)                                                  \
        : _impl::WrapperBase<TYPE_>(                                          \
              Create<TYPE_>(std::forward<Args>(args)...), &DestroyCapsule) {} \
    PY_TYPE_(TYPE_* ptr) : _impl::WrapperBase<TYPE_>(ptr, nullptr) {}         \
    PY_TYPE_(const TYPE_* ptr)                                                \
        : _impl::WrapperBase<TYPE_>(const_cast<TYPE_*>(ptr), nullptr) {}      \
  }

DEFINE_WRAPPED_OBJECT(PyContext, mjrfContext);
DEFINE_WRAPPED_OBJECT(PyTexture, mjrfTexture);
DEFINE_WRAPPED_OBJECT(PyMesh, mjrfMesh);
DEFINE_WRAPPED_OBJECT(PyScene, mjrfScene);
DEFINE_WRAPPED_OBJECT(PyLight, mjrfLight);
DEFINE_WRAPPED_OBJECT(PyRenderable, mjrfRenderable);
DEFINE_WRAPPED_OBJECT(PyRenderTarget, mjrfRenderTarget);
#undef DEFINE_WRAPPED_OBJECT

template <typename T, std::size_t N>
struct NumArray {
  // A reference to a vector; used for binding arguments (e.g. const &).
  using Ref = Eigen::Ref<const Eigen::Vector<T, N>>;
  // A mapping to a vector; used for returning values to existing C++ arrays.
  using Map = Eigen::Map<Eigen::Vector<T, N>>;
};

using Float3 = NumArray<float, 3>;
using Float4 = NumArray<float, 4>;
using Float9 = NumArray<float, 9>;

template <typename T> struct is_object : std::false_type {};
template <> struct is_object<PyContext> : std::true_type {};
template <> struct is_object<PyTexture> : std::true_type {};
template <> struct is_object<PyMesh> : std::true_type {};
template <> struct is_object<PyScene> : std::true_type {};
template <> struct is_object<PyLight> : std::true_type {};
template <> struct is_object<PyRenderable> : std::true_type {};
template <> struct is_object<PyRenderTarget> : std::true_type {};

template <typename T> struct is_num_array : std::false_type {};
template <> struct is_num_array<Float3> : std::true_type {};
template <> struct is_num_array<Float4> : std::true_type {};
template <> struct is_num_array<Float9> : std::true_type {};

void Free(mjrfReadPixelsRequest& data) {
  data.output = nullptr;
  data.num_bytes = 0;
}

void Free(mjrfTextureData& data) {
  if (data.bytes) {
    delete[] static_cast<const char*>(data.bytes);
    data.bytes = nullptr;
    data.num_bytes = 0;
  }
}

void FreeMeshDataVertices(mjrfMeshData& data) {
  if (data.vertices[0]) {
    delete[] static_cast<const char*>(data.vertices[0]);
    for (int i = 0; i < mjMAX_VERTEX_ATTRIBUTES; ++i) {
      data.vertices[i] = nullptr;
    }
    data.num_vertices = 0;
  }
}

void FreeMeshDataIndices(mjrfMeshData& data) {
  if (data.indices) {
    delete[] static_cast<const char*>(data.indices);
    data.indices = nullptr;
    data.num_indices = 0;
  }
}

void Free(mjrfMeshData& data) {
  FreeMeshDataVertices(data);
  FreeMeshDataIndices(data);
}

template <typename T, typename U>
void Setter(T& dst, const U& src) {
  if constexpr (std::is_same_v<T, void*>) {
    dst = reinterpret_cast<T>(src);
  } else if constexpr (std::is_same_v<T, float[3]>) {
    for (int i = 0; i < 3; ++i) {
      dst[i] = src[i];
    }
  } else if constexpr (std::is_same_v<T, float[4]>) {
    for (int i = 0; i < 4; ++i) {
      dst[i] = src[i];
    }
  } else if constexpr (std::is_same_v<T, float[9]>) {
    for (int i = 0; i < 9; ++i) {
      dst[i] = src[i];
    }
  } else if constexpr (std::is_same_v<U, std::vector<mjrVertexAttribute>>) {
    for (int i = 0; i < src.size(); ++i) {
      dst[i] = src[i];
    }
  } else if constexpr (is_object<std::decay_t<U>>::value) {
    dst = const_cast<T>(src.get());
  } else {
    dst = src;
  }
}

template <typename U, typename T, typename C, typename... O>
void DefProperty(pybind11::class_<C, O...>& c, const char* name, T C::* prop) {
  if constexpr (is_num_array<U>::value) {
    c.def_property(name,
      [prop](C& c) { return typename U::Map(c.*prop); },
      [prop](C& c, const typename U::Ref& val) { Setter<T, typename U::Ref>(c.*prop, val); });
  } else if constexpr (is_object<U>::value) {
    c.def_property(name,
      [prop](C& c) { return U(c.*prop); },
      [prop](C& c, const U& val) { Setter<T, U>(c.*prop, val); });
  } else if constexpr (std::is_class_v<U>) {
    c.def_property(name,
      [prop](C& c) -> U& { return c.*prop; },
      [prop](C& c, const U& val) { Setter<T, U>(c.*prop, val); },
      py::return_value_policy::reference_internal);
  } else if constexpr (std::is_same_v<T, void*>) {
    c.def_property(name,
      [prop](C& c) -> U { return reinterpret_cast<U>(c.*prop); },
      [prop](C& c, const U& val) { Setter<T, U>(c.*prop, val); });
  } else {
    c.def_property(name,
      [prop](C& c) { return static_cast<U>(c.*prop); },
      [prop](C& c, const U& val) { Setter<T, U>(c.*prop, val); });
  }
}
}  // namespace

PYBIND11_MODULE(_render_filament, m, pybind11::mod_gil_not_used()) {
#include "render_filament_generated.cc.inc"

  // Explicit bindings for creation functions.
  mjrf_context.def(
      py::init([](mjrfContextConfig& config) {
        return std::make_unique<PyContext>(&config);
      }),
      py::arg("config"));

  mjrf_context.def(
      "create_mesh",
      [](PyContext& ctx, mjrfMeshConfig& config) {
        return std::make_unique<PyMesh>(ctx.get(), &config);
      },
      py::arg("params"), py::keep_alive<0, 1>());
  mjrf_context.def(
      "create_texture",
      [](PyContext& ctx, mjrfTextureConfig& config) {
        return std::make_unique<PyTexture>(ctx.get(), &config);
      },
      py::arg("config"), py::keep_alive<0, 1>());

  mjrf_context.def(
      "create_scene",
      [](PyContext& ctx, mjrfSceneParams& params) {
        return std::make_unique<PyScene>(ctx.get(), &params);
      },
      py::arg("params"), py::keep_alive<0, 1>());

  mjrf_context.def(
      "create_render_target",
      [](PyContext& ctx, mjrfRenderTargetConfig& config) {
        return std::make_unique<PyRenderTarget>(ctx.get(), &config);
      },
      py::arg("config"), py::keep_alive<0, 1>());

  mjrf_context.def(
      "create_light",
      [](PyContext& ctx, mjrfLightParams& params) {
        return std::make_unique<PyLight>(ctx.get(), &params);
      },
      py::arg("params"), py::keep_alive<0, 1>());

  mjrf_context.def(
      "create_renderable",
      [](PyContext& ctx, mjrfRenderableParams& params) {
        return std::make_unique<PyRenderable>(ctx.get(), &params);
      },
      py::arg("params"), py::keep_alive<0, 1>());

  // memory management
  // Note: ReadPixelsRequest stores a raw pointer to `buffer`; the caller must
  // keep the underlying Python buffer object alive until rendering finishes.
  mjrf_read_pixels_request.def(
      "set_buffer",
      [](mjrfReadPixelsRequest& req, py::buffer b) {
        py::buffer_info info = b.request(true);
        if (!PyBuffer_IsContiguous(info.view(), 'C')) {
          throw py::value_error("Buffer must be C-contiguous.");
        }
        req.output = info.ptr;
        req.num_bytes = static_cast<int>(info.size * info.itemsize);
      },
      py::arg("buffer"));

  mjrf_texture_data.def("set_image",
    [](mjrfTextureData& data, py::bytes bytes) {
      Free(data);
      std::string_view sv(bytes);
      if (sv.empty()) {
        return;
      }

      char* buf = new char[sv.size()];
      std::memcpy(buf, sv.data(), sv.size());
      data.bytes = buf;
      data.num_bytes = sv.size();
    }, py::arg("bytes"));

  mjrf_mesh_config.def(
      "attributes",
      [](mjrfMeshConfig& config, int i) {
        if (i < 0 || i >= mjMAX_VERTEX_ATTRIBUTES) {
          throw py::index_error("Index out of range");
        }
        return config.attributes[i];
      },
      py::arg("i"));

  mjrf_mesh_data.def(
      "set_vertices",
      [](mjrfMeshData& data, const mjrfMeshConfig& config, py::bytes vertices) {
        FreeMeshDataVertices(data);
        std::string_view sv(vertices);
        if (sv.empty()) {
          return;
        }

        int vertex_size = 0;
        for (int i = 0; i < config.num_attributes; ++i) {
          switch (config.attributes[i].type) {
            case mjVERTEX_ATTRIBUTE_TYPE_FLOAT2:
              vertex_size += 2 * sizeof(float);
              break;
            case mjVERTEX_ATTRIBUTE_TYPE_FLOAT3:
              vertex_size += 3 * sizeof(float);
              break;
            case mjVERTEX_ATTRIBUTE_TYPE_FLOAT4:
              vertex_size += 4 * sizeof(float);
              break;
            case mjVERTEX_ATTRIBUTE_TYPE_UBYTE4:
              vertex_size += 4 * sizeof(unsigned char);
              break;
          }
        }
        const size_t expected_size = config.max_vertices * vertex_size;
        if (sv.size() != expected_size) {
          throw py::value_error("Expected " + std::to_string(expected_size) +
                                " bytes for vertices, but got " +
                                std::to_string(sv.size()));
        }

        char* buf = new char[sv.size()];
        std::memcpy(buf, sv.data(), sv.size());

        data.num_vertices = config.max_vertices;

        int offset = 0;
        for (int i = 0; i < config.num_attributes; ++i) {
          data.vertices[i] = buf + offset;

          // Calculate the offset for the next attribute using the size of
          // the element of this attribute.
          int elem_size = 0;
          switch (config.attributes[i].type) {
            case mjVERTEX_ATTRIBUTE_TYPE_FLOAT2:
              elem_size = 2 * sizeof(float);
              break;
            case mjVERTEX_ATTRIBUTE_TYPE_FLOAT3:
              elem_size = 3 * sizeof(float);
              break;
            case mjVERTEX_ATTRIBUTE_TYPE_FLOAT4:
              elem_size = 4 * sizeof(float);
              break;
            case mjVERTEX_ATTRIBUTE_TYPE_UBYTE4:
              elem_size = 4 * sizeof(unsigned char);
              break;
          }
          offset += (elem_size * (config.interleaved ? 1 : data.num_vertices));
        }
      },
      py::arg("config"), py::arg("vertices"));

  mjrf_mesh_data.def(
      "set_indices",
      [](mjrfMeshData& data, const mjrfMeshConfig& config, py::bytes indices) {
        FreeMeshDataIndices(data);

        std::string_view sv(indices);
        if (sv.empty()) {
          return;
        }

        const size_t elem_size = config.index_type == mjINDEX_TYPE_U16
                                     ? sizeof(uint16_t)
                                     : sizeof(uint32_t);
        const size_t expected_size = config.max_indices * elem_size;
        if (sv.size() != expected_size) {
          throw py::value_error("Expected " + std::to_string(expected_size) +
                                " bytes for indices, but got " +
                                std::to_string(sv.size()));
        }

        char* buf = new char[sv.size()];
        std::memcpy(buf, sv.data(), sv.size());

        data.indices = buf;
        data.num_indices = config.max_indices;
      },
      py::arg("config"), py::arg("indices"));

  // other
  mjrf_scene.def(
      "configure_from_model",
      [](PyScene& scene, const MjModelWrapper& model) {
        return mjrf_configureSceneFromModel(scene.get(), model.get());
      },
      py::arg("model"));

  mjrf_context.def("render",
    [](PyContext& ctx, const std::vector<mjrfRenderRequest>& req,
        const std::vector<mjrfReadPixelsRequest>& read) {
      return mjrf_render(ctx.get(), req.data(), req.size(), read.data(),
                          read.size());
    },
    py::arg("req"), py::arg("read"),
    py::call_guard<py::gil_scoped_release>());

  mjrf_renderable.def("get_material", [](PyRenderable& renderable) {
    mjrfMaterial material;
    mjrf_getRenderableMaterial(renderable.get(), &material);
    return material;
  });

  using ModelObjects = mujoco::ModelObjects;
  {
    py::class_<ModelObjects> cls(m, "ModelObjects", py::module_local());
    cls.def(py::init([](PyContext& ctx, const MjModelWrapper& model) {
              return std::make_unique<ModelObjects>(model.get(), ctx.get());
            }),
            py::arg("ctx"), py::arg("model"), py::keep_alive<1, 2>(),
            py::keep_alive<1, 3>());
  }

  {
    using ModelRenderables = mujoco::ModelRenderables;

    py::class_<ModelRenderables> cls(m, "ModelRenderables", py::module_local());
    auto init = [](PyScene& scene, ModelObjects& model_objects) {
      return std::make_unique<ModelRenderables>(scene.get(), &model_objects);
    };
    cls.def(py::init(init), py::arg("scene"), py::arg("model_objects"),
            py::keep_alive<1, 2>(), py::keep_alive<1, 3>());
    cls.def("update", [](ModelRenderables& self, const MjDataWrapper& data) {
      self.Update(data.get());
    },
    py::arg("data"));
  }

  {
    using ModelLights = mujoco::ModelLights;

    py::class_<mujoco::ModelLights> cls(m, "ModelLights", py::module_local());
    auto init = [](PyScene& scene, ModelObjects& model_objects) {
      return std::make_unique<ModelLights>(scene.get(), &model_objects);
    };
    cls.def(py::init(init), py::arg("scene"), py::arg("model_objects"),
            py::keep_alive<1, 2>(), py::keep_alive<1, 3>());
    cls.def("update", [](ModelLights& self, const MjDataWrapper& data) {
      self.Update(data.get());
    },
    py::arg("data"));
  }

  {
    using ModelDecorations = mujoco::ModelDecorations;

    py::class_<mujoco::ModelDecorations> cls(m, "ModelDecorations",
                                             py::module_local());
    auto init = [](PyContext& ctx, PyScene& scene, MjModelWrapper& model,
                   int num_geoms) {
      return std::make_unique<ModelDecorations>(ctx.get(), scene.get(),
                                                model.get(), num_geoms);
    };
    cls.def(py::init(init), py::arg("ctx"), py::arg("scene"), py::arg("model"),
            py::arg("num_geoms") = 2000, py::keep_alive<1, 2>(),
            py::keep_alive<1, 3>(), py::keep_alive<1, 4>());
    cls.def(
        "update",
        [](ModelDecorations& self, MjDataWrapper& data,
           const mjvOption& vis_option, const mjvPerturb& perturb,
           mjvCamera& camera, const mjrRect& viewport,
           const std::vector<mjvGeom>& extra_geoms) {
          self.Update(data.get(), &vis_option, &perturb, &camera, viewport,
                      nullptr, extra_geoms);
        },
        py::arg("data"), py::arg("vis_option"), py::arg("perturb"),
        py::arg("camera"), py::arg("viewport"), py::arg("extra_geoms"));
  }
}
}  // namespace mujoco::python