Add per-environment geom visibility mask for heterogeneous entities

src/bridge/bindings.cpp

@@ -51,6 +51,7 @@ NB_MODULE(_gs_madrona_batch_renderer, m) {
             nb::ndarray<const float, nb::shape<-1>, nb::device::cpu> cam_zfar,
             nb::ndarray<const uint32_t, nb::shape<-1>, nb::device::cpu> cam_proj_type,
             nb::ndarray<const int32_t, nb::shape<-1>, nb::device::cpu> enabled_geom_groups,
+            nb::ndarray<const int32_t, nb::shape<-1, -1>, nb::device::cpu> geom_env_mask,
             bool add_cam_debug_geo,
             bool use_rt,
             VisualizerGPUHandles *viz_gpu_hdls)
@@ -81,6 +82,9 @@ NB_MODULE(_gs_madrona_batch_renderer, m) {
                 .geomDataIDs = ptr_geom_data_ids,     // (int32_t *)geom_data_ids.data(),
                 .geomMatIDs = ptr_geom_mat_ids,       // (int32_t *)geom_mat_ids.data(),
                 .enabledGeomGroups = (int32_t *)enabled_geom_groups.data(),
+                // [num_worlds, num_geoms] visibility mask; None -> all visible.
+                .geomEnvMask = geom_env_mask.is_valid()
+                    ? (const int32_t *)geom_env_mask.data() : nullptr,
                 .geomSizes = (math::Vector3 *)geom_sizes.data(),
                 .matRGBA = (math::Vector4 *)mat_rgba.data(),
                 .matTexIDs = (int32_t *)mat_tex_ids.data(),
@@ -147,10 +151,11 @@ NB_MODULE(_gs_madrona_batch_renderer, m) {
            nb::arg("cam_zfar"),
            nb::arg("cam_proj_type"),
            nb::arg("enabled_geom_groups"),
+           nb::arg("geom_env_mask") = nb::none(),
            nb::arg("add_cam_debug_geo") = false,
            nb::arg("use_rt") = false,
            nb::arg("visualizer_gpu_handles") = nb::none(),
-           nb::keep_alive<1, 31>())
+           nb::keep_alive<1, 32>())
         .def("init", [](Manager &mgr,
             nb::ndarray<nb::pytorch, const float, nb::shape<-1, -1, 3>> geom_pos,
             nb::ndarray<nb::pytorch, const float, nb::shape<-1, -1, 4>> geom_rot,

src/bridge/gs_madrona/renderer_gs.py

@@ -15,6 +15,12 @@
 
 class GeomRetriever:
     def retrieve_rigid_meshes_static(self) -> dict:
+        """Return the static geometry kwargs forwarded to MadronaBatchRenderer.
+
+        May include an optional ``geom_env_mask``: an int32 ``[num_worlds, num_geoms]``
+        array (1 = geom visible in that world, 0 = hidden) for heterogeneous entities.
+        Omitting it keeps every geom visible everywhere (legacy behavior).
+        """
         raise NotImplementedError()
 
     def retrieve_rigid_property_torch(self, num_worlds) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

src/bridge/mgr.cpp

@@ -695,9 +695,19 @@ Manager::Impl * Manager::Impl::make(
     REQ_CUDA(cudaMemcpy(cam_zfar, gs_model.camZFar, sizeof(float) * gs_model.numCams, cudaMemcpyHostToDevice));
     REQ_CUDA(cudaMemcpy(cam_proj_type, gs_model.camProjType, sizeof(uint32_t) * gs_model.numCams, cudaMemcpyHostToDevice));
 
+    // Optional per-(world, geom) visibility mask for heterogeneous entities.
+    int32_t *geom_env_mask_gpu = nullptr;
+    if (gs_model.geomEnvMask != nullptr) {
+        size_t mask_count = (size_t)mgr_cfg.numWorlds * gs_model.numGeoms;
+        geom_env_mask_gpu = (int32_t *)cu::allocGPU(sizeof(int32_t) * mask_count);
+        REQ_CUDA(cudaMemcpy(geom_env_mask_gpu, gs_model.geomEnvMask,
+            sizeof(int32_t) * mask_count, cudaMemcpyHostToDevice));
+    }
+
     sim_cfg.geomTypes = geom_types_gpu;
     sim_cfg.geomDataIDs = geom_data_ids_gpu;
     sim_cfg.geomSizes = geom_sizes_gpu;
+    sim_cfg.geomEnvMask = geom_env_mask_gpu;
     sim_cfg.camFovy = cam_fovy;
     sim_cfg.camZNear = cam_znear;
     sim_cfg.camZFar = cam_zfar;
@@ -754,6 +764,9 @@ Manager::Impl * Manager::Impl::make(
     cu::deallocGPU(geom_types_gpu);
     cu::deallocGPU(geom_data_ids_gpu);
     cu::deallocGPU(geom_sizes_gpu);
+    if (geom_env_mask_gpu != nullptr) {
+        cu::deallocGPU(geom_env_mask_gpu);
+    }
     cu::deallocGPU(cam_fovy);
     cu::deallocGPU(cam_znear);
     cu::deallocGPU(cam_zfar);

src/bridge/mgr.hpp

@@ -38,6 +38,9 @@ struct GSModel {
     int32_t *geomDataIDs;
     int32_t *geomMatIDs;
     int32_t *enabledGeomGroups;
+    // Optional per-(world, geom) visibility mask [numWorlds, numGeoms], row-major:
+    // 0 hides that geom in that world; nullptr = all geoms visible everywhere.
+    const int32_t *geomEnvMask;
     madrona::math::Vector3 *geomSizes;
     madrona::math::Vector4 *matRGBA;
     int32_t *matTexIDs;

src/bridge/sim.cpp

@@ -142,8 +142,15 @@ Sim::Sim(Engine &ctx,
 
     for (CountT geom_idx = 0; geom_idx < (CountT)cfg.numGeoms; geom_idx++) {
 
+        // Hide geoms not active in this world (heterogeneous entities). The
+        // instance is still created so the exported [numWorlds, numGeoms]
+        // columns stay aligned, but left non-renderable.
+        bool env_visible = (cfg.geomEnvMask == nullptr) ||
+            (cfg.geomEnvMask[(CountT)ctx.worldID().idx * (CountT)cfg.numGeoms +
+                             geom_idx] != 0);
+
         Entity instance;
-        if (cfg.geomDataIDs[geom_idx] == -1) {
+        if (cfg.geomDataIDs[geom_idx] == -1 || !env_visible) {
             instance = ctx.makeEntity<RenderEntity>();
             render::RenderingSystem::disableEntityRenderable(ctx, instance);
             ctx.get<ObjectID>(instance) = ObjectID {0};

src/bridge/sim.hpp

@@ -54,6 +54,8 @@ struct Sim : public madrona::WorldBase {
         int32_t *geomTypes;
         int32_t *geomDataIDs;
         Vector3 *geomSizes;
+        // Optional device-side [numWorlds, numGeoms] visibility mask (see GSModel::geomEnvMask).
+        const int32_t *geomEnvMask;
         uint32_t numGeoms;
         uint32_t numCams;
         uint32_t numLights;