Limit number of worlds shown in Newton visualizers

docs/guide/visualization.rst

@@ -3,6 +3,16 @@ Visualization
 
 Newton provides multiple viewer backends for different visualization needs, from real-time rendering to offline recording and external integrations.
 
+All viewer backends share a common interface with ``set_model()``, ``begin_frame()``, ``log_state()``, and ``end_frame()`` methods.
+
+**Limiting rendered worlds**: When training with many parallel environments, rendering all worlds can impact performance. 
+All viewers support the ``max_worlds`` parameter to limit visualization to a subset of environments:
+
+.. code-block:: python
+
+    # Only render the first 4 environments
+    viewer.set_model(model, max_worlds=4)
+
 Real-time Viewers
 -----------------
 

newton/_src/viewer/viewer.py

@@ -65,6 +65,7 @@ def __init__(self):
 
         # World offset support
         self.world_offsets = None  # Array of vec3 offsets per world
+        self.max_worlds = None  # Limit on worlds to render (None = all)
 
         # Display options as individual boolean attributes
         self.show_joints = False
@@ -117,11 +118,20 @@ def is_key_down(self, key) -> bool:
         """
         return False
 
-    def set_model(self, model):
+    def set_model(self, model, max_worlds: int | None = None):
+        """
+        Set the model to be visualized.
+
+        Args:
+            model: The Newton model to visualize.
+            max_worlds: Maximum number of worlds to render (None = all).
+                        Useful for performance when training with many environments.
+        """
         if self.model is not None:
             raise RuntimeError("Viewer set_model() can be called only once.")
 
         self.model = model
+        self.max_worlds = max_worlds
 
         if model is not None:
             self.device = model.device
@@ -131,6 +141,22 @@ def set_model(self, model):
             if self.world_offsets is None:
                 self._auto_compute_world_offsets()
 
+    def _should_render_world(self, world_idx: int) -> bool:
+        """Check if a world should be rendered based on max_worlds limit."""
+        if world_idx == -1:  # Global entities always rendered
+            return True
+        if self.max_worlds is None:
+            return True
+        return world_idx < self.max_worlds
+
+    def _get_render_world_count(self) -> int:
+        """Get the number of worlds to render."""
+        if self.model is None:
+            return 0
+        if self.max_worlds is None:
+            return self.model.num_worlds
+        return min(self.max_worlds, self.model.num_worlds)
+
     def _get_shape_isomesh(self, shape_idx: int):
         """Get the isomesh for a collision shape with an SDF volume.
 
@@ -183,7 +209,7 @@ def set_world_offsets(self, spacing: tuple[float, float, float] | list[float] |
         if self.model is None:
             raise RuntimeError("Model must be set before calling set_world_offsets()")
 
-        num_worlds = self.model.num_worlds
+        num_worlds = self._get_render_world_count()
 
         # Get up axis from model
         up_axis = self.model.up_axis
@@ -251,7 +277,7 @@ def _get_world_extents(self) -> tuple[float, float, float] | None:
     def _auto_compute_world_offsets(self):
         """Automatically compute world offsets based on model extents."""
         # If only one world or no worlds, no offsets needed
-        if self.model.num_worlds <= 1:
+        if self._get_render_world_count() <= 1:
             return
 
         max_extents = self._get_world_extents()
@@ -873,6 +899,10 @@ def _populate_shapes(self):
 
         # loop over shapes
         for s in range(shape_count):
+            # skip shapes from worlds beyond max_worlds limit
+            if not self._should_render_world(shape_world[s]):
+                continue
+
             geo_type = shape_geo_type[s]
             geo_scale = [float(v) for v in shape_geo_scale[s]]
             geo_thickness = float(shape_geo_thickness[s])
@@ -1011,6 +1041,10 @@ def _populate_sdf_isomesh_instances(self):
         shape_count = len(shape_body)
 
         for s in range(shape_count):
+            # skip shapes from worlds beyond max_worlds limit
+            if not self._should_render_world(shape_world[s]):
+                continue
+
             # Only process collision shapes with SDF volumes
             is_collision_shape = shape_flags[s] & int(newton.ShapeFlags.COLLIDE_SHAPES)
             if not is_collision_shape:
@@ -1125,8 +1159,12 @@ def _populate_inertia_boxes(self):
         shape_name = "/model/inertia_boxes"
         batch = ViewerBase.ShapeInstances(shape_name, static, flags, mesh_name, self.device)
 
-        # loop over bodys
+        # loop over bodies
         for body in range(body_count):
+            # skip bodies from worlds beyond max_worlds limit
+            if not self._should_render_world(body_world[body]):
+                continue
+
             rot, principal_inertia = wp.eig3(wp.mat33(body_inertia[body]))
             xform = wp.transform(body_com[body], wp.quat_from_matrix(rot))
 

newton/_src/viewer/viewer_file.py

@@ -56,9 +56,9 @@ def __init__(self, output_path: str, auto_save: bool = True, save_interval: int
         self._frame_count = 0
         self._model_recorded = False
 
-    def set_model(self, model):
+    def set_model(self, model, max_worlds: int | None = None):
         """Override set_model to record the model when it's set."""
-        super().set_model(model)
+        super().set_model(model, max_worlds=max_worlds)
 
         if model is not None and not self._model_recorded:
             self._recorder.record_model(model)

newton/_src/viewer/viewer_gl.py

@@ -185,14 +185,15 @@ def log_gizmo(
         self._gizmo_log[name] = transform
 
     @override
-    def set_model(self, model):
+    def set_model(self, model, max_worlds: int | None = None):
         """
         Set the Newton model to visualize.
 
         Args:
             model: The Newton model instance.
+            max_worlds: Maximum number of worlds to render (None = all).
         """
-        super().set_model(model)
+        super().set_model(model, max_worlds=max_worlds)
 
         self.picking = Picking(model, pick_stiffness=10000.0, pick_damping=1000.0, world_offsets=self.world_offsets)
         self.wind = Wind(model)

newton/examples/__init__.py

@@ -337,6 +337,12 @@ def create_parser():
         default=False,
         help="Use MuJoCo's native contact solver instead of Newton contacts (default: use Newton contacts).",
     )
+    parser.add_argument(
+        "--max-worlds",
+        type=int,
+        default=None,
+        help="Maximum number of worlds to render (for performance with many environments).",
+    )
 
     return parser
 

newton/examples/selection/example_selection_cartpole.py

@@ -20,7 +20,8 @@
 # ArticulationView. This example spawns multiple cartpole robots and applies
 # simple random control policy.
 #
-# Command: python -m newton.examples selection_cartpole
+# To limit the number of worlds to render use the max-worlds argument.
+# Command: python -m newton.examples selection_cartpole --num-worlds 16 --max-worlds 8
 #
 ###########################################################################
 
@@ -55,7 +56,7 @@ def apply_forces_kernel(joint_q: wp.array2d(dtype=float), joint_f: wp.array2d(dt
 
 
 class Example:
-    def __init__(self, viewer, num_worlds=16, verbose=True):
+    def __init__(self, viewer, num_worlds=16, max_worlds=None, verbose=True):
         self.fps = 60
         self.frame_dt = 1.0 / self.fps
 
@@ -110,7 +111,7 @@ def __init__(self, viewer, num_worlds=16, verbose=True):
         if not isinstance(self.solver, newton.solvers.SolverMuJoCo):
             self.cartpoles.eval_fk(self.state_0)
 
-        self.viewer.set_model(self.model)
+        self.viewer.set_model(self.model, max_worlds=max_worlds)
         self.viewer.set_world_offsets((2.0, 0.0, 0.0))
 
         # Ensure FK evaluation (for non-MuJoCo solvers):
@@ -235,6 +236,6 @@ def test_final(self):
 
         torch.set_device(args.device)
 
-    example = Example(viewer, num_worlds=args.num_worlds)
+    example = Example(viewer, num_worlds=args.num_worlds, max_worlds=args.max_worlds)
 
     newton.examples.run(example, args)