import os
import pickle
import sys
import time
import weakref
from typing import TYPE_CHECKING, Callable
import numpy as np
import torch
import gstaichi as ti
from numpy.typing import ArrayLike
import genesis as gs
import genesis.utils.geom as gu
from genesis.utils.misc import ALLOCATE_TENSOR_WARNING
from genesis.engine.force_fields import ForceField
from genesis.engine.materials.base import Material
from genesis.engine.states.solvers import SimState
from genesis.options import (
AvatarOptions,
BaseCouplerOptions,
LegacyCouplerOptions,
FEMOptions,
MPMOptions,
PBDOptions,
ProfilingOptions,
RigidOptions,
SFOptions,
SimOptions,
SPHOptions,
ToolOptions,
ViewerOptions,
VisOptions,
)
from genesis.options.morphs import Morph
from genesis.options.surfaces import Surface
from genesis.options.renderers import Rasterizer, RendererOptions
from genesis.options.sensors import SensorOptions
from genesis.options.recorders import RecorderOptions
from genesis.recorders import RecorderManager
from genesis.repr_base import RBC
from genesis.utils.tools import FPSTracker
from genesis.utils.misc import redirect_libc_stderr, tensor_to_array
from genesis.vis import Visualizer
from genesis.utils.warnings import warn_once
if TYPE_CHECKING:
from genesis.engine.entities.base_entity import Entity
from genesis.recorders import Recorder
[docs]
@gs.assert_initialized
class Scene(RBC):
"""
A ``genesis.Scene`` object wraps all components in a simulation environment, including a simulator (containing multiple physics solvers), entities, and a visualizer (controlling both the viewer and all the cameras).
Basically, everything happens inside a scene.
Parameters
----------
sim_options : gs.options.SimOptions
The options configuring the overarching `simulator`, which in turn manages all the solvers.
coupler_options : gs.options.CouplerOptions
The options configuring the `coupler` between different solvers.
tool_options : gs.options.ToolOptions
The options configuring the tool_solver (``scene.sim.ToolSolver``).
rigid_options : gs.options.RigidOptions
The options configuring the rigid_solver (``scene.sim.RigidSolver``).
avatar_options : gs.options.AvatarOptions
The options configuring the avatar_solver (``scene.sim.AvatarSolver``).
mpm_options : gs.options.MPMOptions
The options configuring the mpm_solver (``scene.sim.MPMSolver``).
sph_options : gs.options.SPHOptions
The options configuring the sph_solver (``scene.sim.SPHSolver``).
fem_options : gs.options.FEMOptions
The options configuring the fem_solver (``scene.sim.FEMSolver``).
sf_options : gs.options.SFOptions
The options configuring the sf_solver (``scene.sim.SFSolver``).
pbd_options : gs.options.PBDOptions
The options configuring the pbd_solver (``scene.sim.PBDSolver``).
vis_options : gs.options.VisOptions
The options configuring the visualization system (``scene.visualizer``). Visualizer controls both the interactive viewer and the cameras.
viewer_options : gs.options.ViewerOptions
The options configuring the viewer (``scene.visualizer.viewer``).
renderer : gs.renderers.RendererOptions
The renderer options used by `camera` for rendering images. This doesn't affect the behavior of the interactive viewer.
show_viewer : bool
Whether to show the interactive viewer. Set it to False if you only need headless rendering.
show_FPS : bool
Whether to show the FPS in the terminal.
"""
def __init__(
self,
sim_options: SimOptions | None = None,
coupler_options: BaseCouplerOptions | None = None,
tool_options: ToolOptions | None = None,
rigid_options: RigidOptions | None = None,
avatar_options: AvatarOptions | None = None,
mpm_options: MPMOptions | None = None,
sph_options: SPHOptions | None = None,
fem_options: FEMOptions | None = None,
sf_options: SFOptions | None = None,
pbd_options: PBDOptions | None = None,
vis_options: VisOptions | None = None,
viewer_options: ViewerOptions | None = None,
profiling_options: ProfilingOptions | None = None,
renderer: RendererOptions | None = None,
show_viewer: bool | None = None,
show_FPS: bool | None = None, # deprecated, use profiling_options.show_FPS instead
):
# Delay simulator import to allow specifying gstaichi array type at init
from genesis.engine.simulator import Simulator
# Handling of default arguments
sim_options = sim_options or SimOptions()
coupler_options = coupler_options or LegacyCouplerOptions()
tool_options = tool_options or ToolOptions()
rigid_options = rigid_options or RigidOptions()
avatar_options = avatar_options or AvatarOptions()
mpm_options = mpm_options or MPMOptions()
sph_options = sph_options or SPHOptions()
fem_options = fem_options or FEMOptions()
sf_options = sf_options or SFOptions()
pbd_options = pbd_options or PBDOptions()
vis_options = vis_options or VisOptions()
viewer_options = viewer_options or ViewerOptions()
profiling_options = profiling_options or ProfilingOptions()
renderer = renderer or Rasterizer()
if show_FPS is not None:
warn_once("Scene.show_FPS is deprecated. Please use Scene.profiling_options.show_FPS")
profiling_options.show_FPS = show_FPS
# validate options
self._validate_options(
sim_options,
coupler_options,
tool_options,
rigid_options,
avatar_options,
mpm_options,
sph_options,
fem_options,
sf_options,
pbd_options,
vis_options,
viewer_options,
profiling_options,
renderer,
)
self.sim_options = sim_options
self.coupler_options = coupler_options
self.tool_options = tool_options
self.rigid_options = rigid_options
self.avatar_options = avatar_options
self.mpm_options = mpm_options
self.sph_options = sph_options
self.fem_options = fem_options
self.sf_options = sf_options
self.pbd_options = pbd_options
self.profiling_options = profiling_options
self.vis_options = vis_options
self.viewer_options = viewer_options
self.renderer_options = renderer
# merge options
self.tool_options.copy_attributes_from(self.sim_options)
self.rigid_options.copy_attributes_from(self.sim_options)
self.avatar_options.copy_attributes_from(self.sim_options)
self.mpm_options.copy_attributes_from(self.sim_options)
self.sph_options.copy_attributes_from(self.sim_options)
self.fem_options.copy_attributes_from(self.sim_options)
self.sf_options.copy_attributes_from(self.sim_options)
self.pbd_options.copy_attributes_from(self.sim_options)
# simulator
self._sim = Simulator(
scene=self,
options=self.sim_options,
coupler_options=self.coupler_options,
tool_options=self.tool_options,
rigid_options=self.rigid_options,
avatar_options=self.avatar_options,
mpm_options=self.mpm_options,
sph_options=self.sph_options,
fem_options=self.fem_options,
sf_options=self.sf_options,
pbd_options=self.pbd_options,
)
# visualizer
self._visualizer = Visualizer(
scene=self,
show_viewer=show_viewer,
vis_options=vis_options,
viewer_options=viewer_options,
renderer_options=renderer,
)
# recorders
self._recorder_manager = RecorderManager(self._sim.dt)
# emitters
self._emitters = gs.List()
self._backward_ready = False
self._forward_ready = False
self._uid = gs.UID()
self._t = 0
self._is_built = False
gs.logger.info(f"Scene ~~~<{self._uid}>~~~ created.")
def __del__(self):
self.destroy()
def _validate_options(
self,
sim_options: SimOptions,
coupler_options: BaseCouplerOptions,
tool_options: ToolOptions,
rigid_options: RigidOptions,
avatar_options: AvatarOptions,
mpm_options: MPMOptions,
sph_options: SPHOptions,
fem_options: FEMOptions,
sf_options: SFOptions,
pbd_options: PBDOptions,
vis_options: VisOptions,
viewer_options: ViewerOptions,
profiling_options: ProfilingOptions,
renderer_options: RendererOptions,
):
if not isinstance(sim_options, SimOptions):
gs.raise_exception("`sim_options` should be an instance of `SimOptions`.")
if not isinstance(coupler_options, BaseCouplerOptions):
gs.raise_exception("`coupler_options` should be an instance of `BaseCouplerOptions`.")
if not isinstance(tool_options, ToolOptions):
gs.raise_exception("`tool_options` should be an instance of `ToolOptions`.")
if not isinstance(rigid_options, RigidOptions):
gs.raise_exception("`rigid_options` should be an instance of `RigidOptions`.")
if not isinstance(avatar_options, AvatarOptions):
gs.raise_exception("`avatar_options` should be an instance of `AvatarOptions`.")
if not isinstance(mpm_options, MPMOptions):
gs.raise_exception("`mpm_options` should be an instance of `MPMOptions`.")
if not isinstance(sph_options, SPHOptions):
gs.raise_exception("`sph_options` should be an instance of `SPHOptions`.")
if not isinstance(fem_options, FEMOptions):
gs.raise_exception("`fem_options` should be an instance of `FEMOptions`.")
if not isinstance(sf_options, SFOptions):
gs.raise_exception("`sf_options` should be an instance of `SFOptions`.")
if not isinstance(pbd_options, PBDOptions):
gs.raise_exception("`pbd_options` should be an instance of `PBDOptions`.")
if not isinstance(vis_options, VisOptions):
gs.raise_exception("`vis_options` should be an instance of `VisOptions`.")
if not isinstance(viewer_options, ViewerOptions):
gs.raise_exception("`viewer_options` should be an instance of `ViewerOptions`.")
if not isinstance(profiling_options, ProfilingOptions):
gs.raise_exception("`profiling_options` should be an instance of `ProfilingOptions`.")
if not isinstance(renderer_options, RendererOptions):
gs.raise_exception("`renderer` should be an instance of `gs.renderers.Renderer`.")
# Validate rigid_options against sim_options
if rigid_options.box_box_detection is None:
rigid_options.box_box_detection = not sim_options.requires_grad
elif rigid_options.box_box_detection and sim_options.requires_grad:
gs.raise_exception(
"`rigid_options.box_box_detection` cannot be True when `sim_options.requires_grad` is True."
)
if not rigid_options.use_gjk_collision and sim_options.requires_grad:
gs.raise_exception(
"`rigid_options.use_gjk_collision` cannot be False when `sim_options.requires_grad` is True."
)
if rigid_options.enable_mujoco_compatibility and sim_options.requires_grad:
gs.raise_exception(
"`rigid_options.enable_mujoco_compatibility` cannot be True when `sim_options.requires_grad` is True."
)
[docs]
def destroy(self):
if getattr(self, "_recorder_manager", None) is not None:
if self._recorder_manager.is_recording:
self._recorder_manager.stop()
self._recorder_manager = None
if getattr(self, "_visualizer", None) is not None:
self._visualizer.destroy()
self._visualizer = None
# Stop tracking this scene
try:
gs._scene_registry.remove(weakref.ref(self))
except ValueError:
# This scene may have been destroyed previously
pass
[docs]
@gs.assert_unbuilt
def add_entity(
self,
morph: Morph,
material: Material | None = None,
surface: Surface | None = None,
visualize_contact: bool = False,
vis_mode: str | None = None,
):
"""
Add an entity to the scene.
Parameters
----------
morph : gs.morphs.Morph
The morph of the entity.
material : gs.materials.Material | None, optional
The material of the entity. If None, use ``gs.materials.Rigid()``.
surface : gs.surfaces.Surface | None, optional
The surface of the entity. If None, use ``gs.surfaces.Default()``.
visualize_contact : bool
Whether to visualize contact forces applied to this entity as arrows in the viewer and rendered images. Note that this will not be displayed in images rendered by camera using the `RayTracer` renderer.
vis_mode : str | None, optional
The visualization mode of the entity. This is a handy shortcut for setting `surface.vis_mode` without explicitly creating a surface object.
Returns
-------
entity : genesis.Entity
The created entity.
"""
if material is None:
material = gs.materials.Rigid()
if surface is None:
# assign a local surface, otherwise modification will apply on global default surface
surface = gs.surfaces.Default()
if isinstance(material, gs.materials.Rigid):
# small sdf res is sufficient for primitives regardless of size
if isinstance(morph, gs.morphs.Primitive):
material._sdf_max_res = 32
# some morph should not smooth surface normal
if isinstance(morph, (gs.morphs.Box, gs.morphs.Cylinder, gs.morphs.Terrain)):
surface.smooth = False
if isinstance(morph, (gs.morphs.URDF, gs.morphs.MJCF, gs.morphs.Terrain)):
if not isinstance(material, (gs.materials.Rigid, gs.materials.Avatar, gs.materials.Hybrid)):
gs.raise_exception(f"Unsupported material for morph: {material} and {morph}.")
if surface.double_sided is None:
if isinstance(material, gs.materials.PBD.Cloth):
surface.double_sided = True
else:
surface.double_sided = False
if vis_mode is not None:
surface.vis_mode = vis_mode
# validate and populate default surface.vis_mode considering morph type
if isinstance(material, (gs.materials.Rigid, gs.materials.Avatar, gs.materials.Tool)):
if surface.vis_mode is None:
surface.vis_mode = "visual"
if surface.vis_mode not in ("visual", "collision", "sdf"):
gs.raise_exception(
f"Unsupported `surface.vis_mode` for material {material}: '{surface.vis_mode}'. Expected one of: ['visual', 'collision', 'sdf']."
)
elif isinstance(
material,
(
gs.materials.PBD.Liquid,
gs.materials.PBD.Particle,
gs.materials.MPM.Liquid,
gs.materials.MPM.Sand,
gs.materials.MPM.Snow,
gs.materials.SPH.Liquid,
),
):
if surface.vis_mode is None:
surface.vis_mode = "particle"
if surface.vis_mode not in ("particle", "recon"):
gs.raise_exception(
f"Unsupported `surface.vis_mode` for material {material}: '{surface.vis_mode}'. Expected one of: ['particle', 'recon']."
)
elif isinstance(material, (gs.materials.SF.Smoke)):
if surface.vis_mode is None:
surface.vis_mode = "particle"
if surface.vis_mode not in ["particle"]:
gs.raise_exception(
f"Unsupported `surface.vis_mode` for material {material}: '{surface.vis_mode}'. Expected one of: ['particle', 'recon']."
)
elif isinstance(material, (gs.materials.PBD.Base, gs.materials.MPM.Base, gs.materials.SPH.Base)):
if surface.vis_mode is None:
surface.vis_mode = "visual"
if surface.vis_mode not in ("visual", "particle", "recon"):
gs.raise_exception(
f"Unsupported `surface.vis_mode` for material {material}: '{surface.vis_mode}'. Expected one of: ['visual', 'particle', 'recon']."
)
elif isinstance(material, (gs.materials.FEM.Base)):
if surface.vis_mode is None:
surface.vis_mode = "visual"
if surface.vis_mode not in ["visual"]:
gs.raise_exception(
f"Unsupported `surface.vis_mode` for material {material}: '{surface.vis_mode}'. Expected one of: ['visual']."
)
elif isinstance(material, (gs.materials.Hybrid)): # determine the visual of the outer soft part
if surface.vis_mode is None:
surface.vis_mode = "particle"
if surface.vis_mode not in ["particle", "visual"]:
gs.raise_exception(
f"Unsupported `surface.vis_mode` for material {material}: '{surface.vis_mode}'. Expected one of: ['particle', 'visual']."
)
else:
gs.raise_exception()
# Set material-dependent default options
if isinstance(morph, gs.morphs.FileMorph):
# Rigid entities will convexify geom by default
if morph.convexify is None:
morph.convexify = isinstance(material, (gs.materials.Rigid, gs.materials.Avatar))
entity = self._sim._add_entity(morph, material, surface, visualize_contact)
return entity
[docs]
@gs.assert_unbuilt
def link_entities(
self,
parent_entity: "Entity",
child_entity: "Entity",
parent_link_name="",
child_link_name="",
):
"""
links two entities to act as single entity.
Parameters
----------
parent_entity : genesis.Entity
The entity in the scene that will be a parent of kinematic tree.
child_entity : genesis.Entity
The entity in the scene that will be a child of kinematic tree.
parent_link_name : str
The name of the link in the parent entity to be linked.
child_link_name : str
The name of the link in the child entity to be linked.
"""
if not isinstance(parent_entity, gs.engine.entities.RigidEntity):
gs.raise_exception("Currently only rigid entities are supported for merging.")
if not isinstance(child_entity, gs.engine.entities.RigidEntity):
gs.raise_exception("Currently only rigid entities are supported for merging.")
if not child_link_name:
for link in child_entity._links:
if link.parent_idx == -1:
child_link = link
break
else:
child_link = child_entity.get_link(child_link_name)
parent_link = parent_entity.get_link(parent_link_name)
if child_link._parent_idx != -1:
gs.logger.warning(
"Child entity already has a parent link. This may cause the entity to break into parts. Make sure "
"this operation is intended."
)
child_link._parent_idx = parent_link.idx
parent_link._child_idxs.append(child_link.idx)
[docs]
@gs.assert_unbuilt
def add_mesh_light(
self,
morph: Morph | None = None,
color: ArrayLike | None = (1.0, 1.0, 1.0, 1.0),
intensity: float = 20.0,
revert_dir: bool | None = False,
double_sided: bool | None = False,
cutoff: float | None = 180.0,
):
"""
Add a mesh light to the scene. Only supported by RayTracer.
Parameters
----------
morph : gs.morphs.Morph
The morph of the light. Must be an instance of `gs.morphs.Primitive` or `gs.morphs.Mesh`.
color : tuple of float, shape (3,)
The color of the light, specified as (r, g, b).
intensity : float
The intensity of the light.
revert_dir : bool
Whether to revert the direction of the light. If True, the light will be emitted towards the mesh's inside.
double_sided : bool
Whether to emit light from both sides of surface.
cutoff : float
The cutoff angle of the light in degrees. Range: [0.0, 180.0].
"""
if not isinstance(self.renderer_options, gs.renderers.RayTracer):
gs.raise_exception(
"This method is only supported by RayTracer. Please use 'add_light' when using BatchRenderer."
)
if not isinstance(morph, (gs.morphs.Primitive, gs.morphs.Mesh)):
gs.raise_exception("Light morph only supports `gs.morphs.Primitive` or `gs.morphs.Mesh`.")
mesh = gs.Mesh.from_morph_surface(morph, gs.surfaces.Plastic(smooth=False))
self._visualizer.add_mesh_light(mesh, color, intensity, morph.pos, morph.quat, revert_dir, double_sided, cutoff)
[docs]
@gs.assert_unbuilt
def add_light(
self,
pos: ArrayLike,
dir: ArrayLike,
color: ArrayLike = (1.0, 1.0, 1.0),
intensity: float = 1.0,
directional: bool = False,
castshadow: bool = True,
cutoff: float = 45.0,
attenuation: float = 0.0,
):
"""
Add a light to the scene for batch renderer.
Parameters
----------
pos : tuple of float, shape (3,)
The position of the light, specified as (x, y, z).
dir : tuple of float, shape (3,)
The direction of the light, specified as (x, y, z).
color : tuple of float, shape (3,)
The color of the light, specified as (r, g, b).
intensity : float
The intensity of the light.
directional : bool
Whether the light is directional.
castshadow : bool
Whether the light casts shadows.
cutoff : float
The cutoff angle of the light in degrees. Range: (0.0, 90.0).
attenuation : float
The attenuation factor of the light.
Light intensity will attenuate by distance with (1 / (1 + attenuation * distance ^ 2))
"""
if not isinstance(self.renderer_options, gs.renderers.BatchRenderer):
gs.raise_exception(
"This method is only supported by BatchRenderer. Please use 'add_mesh_light' when using RayTracer."
)
self.visualizer.add_light(pos, dir, color, intensity, directional, castshadow, cutoff, attenuation)
[docs]
@gs.assert_unbuilt
def add_sensor(self, sensor_options: "SensorOptions"):
"""
Add a sensor to the scene.
Sensors extract information from the scene without modifying the physics simulation.
Parameters
----------
sensor_options : SensorOptions
The options for the sensor.
"""
return self._sim._sensor_manager.create_sensor(sensor_options)
[docs]
@gs.assert_unbuilt
def start_recording(self, data_func: Callable, rec_options: "RecorderOptions") -> "Recorder":
"""
Automatically read and process data. See RecorderOptions for more details.
Data from `data_func` is automatically read and processed using the recorder at the
frequency `rec_options.hz` (or every step if not specified) as the scene is stepped.
Parameters
----------
data_func: Callable
A function with no arguments that returns the data to be recorded.
rec_options : RecorderOptions
The options for the recording.
Returns
-------
recorder : Recorder
The created recorder object.
"""
return self._recorder_manager.add_recorder(data_func, rec_options)
[docs]
@gs.assert_unbuilt
def add_camera(
self,
model="pinhole",
res=(320, 320),
pos=(0.5, 2.5, 3.5),
lookat=(0.5, 0.5, 0.5),
up=(0.0, 0.0, 1.0),
fov=30,
aperture=2.0,
focus_dist=None,
GUI=False,
spp=256,
denoise=None,
near=0.1,
far=20.0,
env_idx=None,
debug=False,
):
"""
Add a camera to the scene.
The camera model can be either 'pinhole' or 'thinlens'. The 'pinhole' model is a simple camera model that
captures light rays from a single point in space. The 'thinlens' model is a more complex camera model that
simulates a lens with a finite aperture size, allowing for depth of field effects.
Warning
-------
When 'pinhole' is used, the `aperture` and `focal_len` parameters are ignored.
Parameters
----------
model : str
Specifies the camera model. Options are 'pinhole' or 'thinlens'.
res : tuple of int, shape (2,)
The resolution of the camera, specified as a tuple (width, height).
pos : tuple of float, shape (3,)
The position of the camera in the scene, specified as (x, y, z).
lookat : tuple of float, shape (3,)
The point in the scene that the camera is looking at, specified as (x, y, z).
up : tuple of float, shape (3,)
The up vector of the camera, defining its orientation, specified as (x, y, z).
fov : float
The vertical field of view of the camera in degrees.
aperture : float
The aperture size of the camera, controlling depth of field.
focus_dist : float | None
The focus distance of the camera. If None, it will be auto-computed using `pos` and `lookat`.
GUI : bool
Whether to display the camera's rendered image in a separate GUI window.
spp : int, optional
Samples per pixel. Only available when using RayTracer renderer. Defaults to 256.
denoise : bool
Whether to denoise the camera's rendered image. Only available when using the RayTracer renderer. Defaults
to True on Linux, otherwise False. If OptiX denoiser is not available in your platform, consider enabling
the OIDN denoiser option when building the RayTracer.
near : float
Distance from camera center to near plane in meters.
Only available when using rasterizer in Rasterizer and BatchRender renderer. Defaults to 0.1.
far : float
Distance from camera center to far plane in meters.
Only available when using rasterizer in Rasterizer and BatchRender renderer. Defaults to 20.0.
env_idx : int, optional
The specific environment index to bind to the camera. This option must be specified if and only if a
non-batched renderer is being used. If provided, only this environment will be taken into account when
following a rigid entity via 'follow_entity' and when being attached to some rigid link via 'attach'. Note
that this option is unrelated to which environment is being rendering on the scene. Default to None for
batched renderers (ie BatchRender), 'rendered_envs_idx[0]' otherwise (ie Raytracer or Rasterizer).
debug : bool
Whether to use the debug camera. It enables to create cameras that can used to monitor / debug the
simulation without being part of the "sensors". Their output is rendered by the usual simple Rasterizer
systematically, no matter if BatchRender and RayTracer is enabled. This way, it is possible to record the
simulation with arbitrary resolution and camera pose, without interfering with what robots can perceive
from their environment. Defaults to False.
Returns
-------
camera : genesis.Camera
The created camera object.
"""
if denoise is None:
denoise = sys.platform != "darwin"
return self._visualizer.add_camera(
res, pos, lookat, up, model, fov, aperture, focus_dist, GUI, spp, denoise, near, far, env_idx, debug
)
[docs]
@gs.assert_unbuilt
def add_emitter(
self,
material: Material,
max_particles=20000,
surface: Surface | None = None,
):
"""
Add a fluid emitter to the scene.
Parameters
----------
material : gs.materials.Material
The material of the fluid to be emitted. Must be an instance of `gs.materials.MPM.Base`,
`gs.materials.SPH.Base`, `gs.materials.PBD.Particle` or `gs.materials.PBD.Liquid`.
max_particles : int
The maximum number of particles that can be emitted by the emitter. Particles will be recycled once this
limit is reached.
surface : gs.surfaces.Surface | None, optional
The surface of the emitter. If None, use `gs.surfaces.Default(color=(0.6, 0.8, 1.0, 1.0))`.
Returns
-------
emitter : genesis.Emitter
The created emitter object.
"""
from genesis.engine.entities import Emitter
if self.requires_grad:
gs.raise_exception("Emitter is not supported in differentiable mode.")
if not isinstance(
material, (gs.materials.MPM.Base, gs.materials.SPH.Base, gs.materials.PBD.Particle, gs.materials.PBD.Liquid)
):
gs.raise_exception(
"Non-supported material for emitter. Supported materials are: `gs.materials.MPM.Base`, "
"`gs.materials.SPH.Base`, `gs.materials.PBD.Particle`, `gs.materials.PBD.Liquid`."
)
if surface is None:
surface = gs.surfaces.Default(color=(0.6, 0.8, 1.0, 1.0))
if surface.vis_mode is None:
surface.vis_mode = "particle"
if surface.vis_mode == "visual":
gs.raise_exception("surface.vis_mode='visual' is not supported for fluid emitters.")
emitter = Emitter(max_particles)
entity = self.add_entity(
morph=gs.morphs.Nowhere(n_particles=max_particles),
material=material,
surface=surface,
)
emitter.set_entity(entity)
self._emitters.append(emitter)
return emitter
[docs]
@gs.assert_unbuilt
def add_force_field(self, force_field: ForceField):
"""
Add a force field to the scene.
Parameters
----------
force_field : gs.force_fields.ForceField
The force field to add to the scene.
Returns
-------
force_field : gs.force_fields.ForceField
The added force field.
"""
force_field.scene = self
self._sim._add_force_field(force_field)
return force_field
[docs]
@gs.assert_unbuilt
def build(
self,
n_envs=0,
env_spacing=(0.0, 0.0),
n_envs_per_row: int | None = None,
center_envs_at_origin=True,
compile_kernels=None,
):
"""
Builds the scene once all entities have been added. This operation is required before running the simulation.
Parameters
----------
n_envs : int
Number of parallel environments to create.
If `n_envs` is 0, the scene will not have a batching dimension. When greater than 0, the first dimension of
all the input and returned states will be the batch dimension.
env_spacing : tuple of float, shape (2,)
The spacing between adjacent environments in the scene.
This is for visualization purposes only and does not change simulation-related poses.
n_envs_per_row : int
The number of environments per row for visualization. If None, it will be set to `sqrt(n_envs)`.
center_envs_at_origin : bool
Whether to put the center of all the environments at the origin (for visualization only).
compile_kernels : bool, optional
This parameter is deprecated and will be removed in future release.
"""
if compile_kernels is not None:
warn_once("`compile_kernels` is deprecated and will be removed in future release.")
compile_kernels = True
with gs.logger.timer(f"Building scene ~~~<{self._uid}>~~~..."):
self._parallelize(n_envs, env_spacing, n_envs_per_row, center_envs_at_origin)
# simulator
self._sim.build()
# reset state
self._reset()
self._is_built = True
with gs.logger.timer("Compiling simulation kernels..."):
self._sim.step()
self._reset()
# visualizer
with gs.logger.timer("Building visualizer..."):
self._visualizer.build()
if self.profiling_options.show_FPS:
self.FPS_tracker = FPSTracker(self.n_envs, alpha=self.profiling_options.FPS_tracker_alpha)
# recorders
self._recorder_manager.build()
# Update global scene registry
def _destroy_callback(scene_ref: weakref.ReferenceType["Scene"]):
scene = scene_ref()
for i, scene_ref_i in enumerate(gs._scene_registry):
if scene is scene_ref_i():
del gs._scene_registry[i]
break
gs._scene_registry.append(weakref.ref(self, _destroy_callback))
def _parallelize(
self,
n_envs: int,
env_spacing: tuple[float, float],
n_envs_per_row: int,
center_envs_at_origin: bool,
):
self.n_envs = n_envs
self.env_spacing = env_spacing
self.n_envs_per_row = n_envs_per_row
# true batch size
self._B = max(1, self.n_envs)
self._envs_idx = torch.arange(self._B, dtype=gs.tc_int, device=gs.device)
if self.n_envs_per_row is None:
self.n_envs_per_row = np.ceil(np.sqrt(self._B)).astype(int)
# compute offset values for visualizing each env
if not isinstance(env_spacing, (list, tuple)) or len(env_spacing) != 2:
gs.raise_exception("`env_spacing` should be a tuple of length 2.")
offset_x = (np.arange(self._B) // self.n_envs_per_row) * self.env_spacing[0]
offset_y = (np.arange(self._B) % self.n_envs_per_row) * self.env_spacing[1]
offset_z = np.zeros((self._B,))
self.envs_offset = np.stack((offset_x, offset_y, offset_z), axis=-1, dtype=gs.np_float)
# move to center
if center_envs_at_origin:
center = (np.max(self.envs_offset, axis=0) + np.min(self.envs_offset, axis=0)) / 2.0
self.envs_offset -= center
"""
Notes:
- When using gpu
- for non-batched env, we only parallelize certain loops that have big loop size
- for batched env, we parallelize all loops
- When using cpu, we serialize everything.
- This is emprically as fast as parallel loops even with big batchsize (tested up to B=10000), because invoking multiple cpu processes cannot utilize all cpu usage.
- In order to exploit full cpu power, users are encouraged to launch multiple processes manually, and each will use a single cpu thred.
"""
if gs.backend == gs.cpu:
self._para_level = gs.PARA_LEVEL.NEVER
elif self.n_envs == 0:
self._para_level = gs.PARA_LEVEL.PARTIAL
else:
self._para_level = gs.PARA_LEVEL.ALL
[docs]
@gs.assert_built
def reset(self, state: SimState | None = None, envs_idx=None):
"""
Resets the scene to its initial state.
Parameters
----------
state : SimState | None
The state to reset the scene to. If None, the scene will be reset to its initial state.
If this is given, the scene's registerered initial state will be updated to this state.
envs_idx : None | array_like, optional
The indices of the environments. If None, all environments will be considered. Defaults to None.
"""
gs.logger.debug(f"Resetting Scene ~~~<{self._uid}>~~~.")
self._reset(state, envs_idx=envs_idx)
self._recorder_manager.reset(envs_idx)
def _reset(self, state: SimState | None = None, *, envs_idx=None):
if self._is_built:
if state is None:
state = self._init_state
else:
assert isinstance(state, SimState), "state must be a SimState object"
self._init_state = state
self._sim.reset(state, envs_idx)
else:
self._init_state = self._get_state()
self._t = 0
self._forward_ready = True
self._reset_grad()
# Clear the entire cache of the visualizer.
# TODO: Could be optimized to only clear cache associated the the environments being reset.
if self._visualizer.is_built:
self._visualizer.reset()
# TODO: sets _next_particle = 0; not sure this is env isolation safe
for emitter in self._emitters:
emitter.reset()
def _reset_grad(self):
self._backward_ready = True
def _get_state(self):
return self._sim.get_state()
[docs]
@gs.assert_built
def get_state(self):
"""
Returns the current state of the scene.
Returns
-------
state : genesis.SimState
The state of the scene at the current time step.
"""
return self._get_state()
[docs]
@gs.assert_built
def step(self, update_visualizer=True, refresh_visualizer=True):
"""
Runs a simulation step forward in time.
"""
if not self._forward_ready:
gs.raise_exception("Forward simulation not allowed after backward pass. Please reset scene state.")
self._sim.step()
self._t += 1
if update_visualizer:
self._visualizer.update(force=False, auto=refresh_visualizer)
# Update IPC GUI if enabled
if hasattr(self, "_ipc_gui_enabled") and self._ipc_gui_enabled:
self._sim._coupler.update_ipc_gui()
if self.profiling_options.show_FPS:
self.FPS_tracker.step()
self._recorder_manager.step(self._sim.cur_step_global)
[docs]
def stop_recording(self):
self._recorder_manager.stop()
def _step_grad(self):
self._sim.collect_output_grads()
self._sim._step_grad()
self._t -= 1
[docs]
@gs.assert_built
def draw_debug_line(self, start, end, radius=0.002, color=(1.0, 0.0, 0.0, 0.5)):
"""
Draws a line in the scene for visualization.
Parameters
----------
start : array_like, shape (3,)
The starting point of the line.
end : array_like, shape (3,)
The ending point of the line.
radius : float, optional
The radius of the line (represented as a cylinder)
color : array_like, shape (4,), optional
The color of the line in RGBA format.
Returns
-------
node : genesis.ext.pyrender.mesh.Mesh
The created debug object.
"""
with self._visualizer.viewer_lock:
return self._visualizer.context.draw_debug_line(start, end, radius, color)
[docs]
@gs.assert_built
def draw_debug_arrow(self, pos, vec=(0, 0, 1), radius=0.01, color=(1.0, 0.0, 0.0, 0.5)):
"""
Draws an arrow in the scene for visualization.
Parameters
----------
pos : array_like, shape (3,)
The starting position of the arrow.
vec : array_like, shape (3,), optional
The vector of the arrow.
radius : float, optional
The radius of the arrow body (represented as a cylinder).
color : array_like, shape (4,), optional
The color of the arrow in RGBA format.
Returns
-------
node : genesis.ext.pyrender.mesh.Mesh
The created debug object.
"""
with self._visualizer.viewer_lock:
return self._visualizer.context.draw_debug_arrow(pos, vec, radius, color)
[docs]
@gs.assert_built
def draw_debug_frame(self, T, axis_length=1.0, origin_size=0.015, axis_radius=0.01):
"""
Draws a 3-axis coordinate frame in the scene for visualization.
Parameters
----------
T : array_like, shape (4, 4)
The transformation matrix of the frame.
axis_length : float, optional
The length of the axes.
origin_size : float, optional
The size of the origin point (represented as a sphere).
axis_radius : float, optional
The radius of the axes (represented as cylinders).
Returns
-------
node : genesis.ext.pyrender.mesh.Mesh
The created debug object.
"""
with self._visualizer.viewer_lock:
return self._visualizer.context.draw_debug_frame(T, axis_length, origin_size, axis_radius)
[docs]
@gs.assert_built
def draw_debug_frames(self, Ts, axis_length=1.0, origin_size=0.015, axis_radius=0.01):
"""
Draws 3-axis coordinate frames in the scene for visualization.
Parameters
----------
Ts : array_like, shape (n, 4, 4)
The transformation matrices of frames.
axis_length : float, optional
The length of the axes.
origin_size : float, optional
The size of the origin point (represented as a sphere).
axis_radius : float, optional
The radius of the axes (represented as cylinders).
Returns
-------
node : genesis.ext.pyrender.mesh.Mesh
The created debug object.
"""
with self._visualizer.viewer_lock:
return self._visualizer.context.draw_debug_frames(Ts, axis_length, origin_size, axis_radius)
[docs]
@gs.assert_built
def draw_debug_mesh(self, mesh, pos=np.zeros(3), T=None):
"""
Draws a mesh in the scene for visualization.
Parameters
----------
mesh : trimesh.Trimesh
The mesh to be drawn.
pos : array_like, shape (3,), optional
The position of the mesh in the scene.
T : array_like, shape (4, 4) | None, optional
The transformation matrix of the mesh. If None, the mesh will be drawn at the position specified by `pos`. Otherwise, `T` has a higher priority than `pos`.
Returns
-------
node : genesis.ext.pyrender.mesh.Mesh
The created debug object.
"""
with self._visualizer.viewer_lock:
return self._visualizer.context.draw_debug_mesh(mesh, pos, T)
[docs]
@gs.assert_built
def draw_debug_sphere(self, pos, radius=0.01, color=(1.0, 0.0, 0.0, 0.5)):
"""
Draws a sphere in the scene for visualization.
Parameters
----------
pos : array_like, shape (3,)
The center position of the sphere.
radius : float, optional
radius of the sphere.
color : array_like, shape (4,), optional
The color of the sphere in RGBA format.
Returns
-------
node : genesis.ext.pyrender.mesh.Mesh
The created debug object.
"""
with self._visualizer.viewer_lock:
return self._visualizer.context.draw_debug_sphere(pos, radius, color)
[docs]
@gs.assert_built
def draw_debug_spheres(self, poss, radius=0.01, color=(1.0, 0.0, 0.0, 0.5)):
"""
Draws multiple spheres in the scene for visualization.
Parameters
----------
poss : array_like, shape (N, 3)
The positions of the spheres.
radius : float, optional
The radius of the spheres.
color : array_like, shape (4,), optional
The color of the spheres in RGBA format.
Returns
-------
node : genesis.ext.pyrender.mesh.Mesh
The created debug object.
"""
with self._visualizer.viewer_lock:
return self._visualizer.context.draw_debug_spheres(poss, radius, color)
[docs]
@gs.assert_built
def draw_debug_box(
self,
bounds,
color=(1.0, 0.0, 0.0, 1.0),
wireframe=True,
wireframe_radius=0.0015,
):
"""
Draws a box in the scene for visualization.
Parameters
----------
bounds : array_like, shape (2, 3)
The bounds of the box, specified as [[min_x, min_y, min_z], [max_x, max_y, max_z]].
color : array_like, shape (4,), optional
The color of the box in RGBA format.
wireframe : bool, optional
Whether to draw the box as a wireframe.
wireframe_radius : float, optional
The radius of the wireframe lines.
Returns
-------
node : genesis.ext.pyrender.mesh.Mesh
The created debug object.
"""
with self._visualizer.viewer_lock:
return self._visualizer.context.draw_debug_box(
bounds, color, wireframe=wireframe, wireframe_radius=wireframe_radius
)
[docs]
@gs.assert_built
def draw_debug_points(self, poss, colors=(1.0, 0.0, 0.0, 0.5)):
"""
Draws points in the scene for visualization.
Parameters
----------
poss : array_like, shape (N, 3)
The positions of the points.
colors : array_like, shape (4,), optional
The color of the points in RGBA format.
Returns
-------
node : genesis.ext.pyrender.mesh.Mesh
The created debug object.
"""
with self._visualizer.viewer_lock:
return self._visualizer.context.draw_debug_points(poss, colors)
[docs]
@gs.assert_built
def draw_debug_path(self, qposs, entity, link_idx=-1, density=0.3, frame_scaling=1.0):
"""
Draws a planned joint trajectory in the scene for visualization.
Parameters
----------
qposs : array_like, shape (N, M)
The joint positions of the planned points.
N is the number of configurations (i.e., trajectory points).
M is the number of degrees of freedom for the entity (i.e., joint dimensions).
entity : gs.engine.entities.RigidEntity
The rigid entity whose forward kinematics are used to compute the trajectory path.
link_idx : int, optional
The link id of the rigid entity to visualize. Defeault is -1.
density : float, optional
Controls the sampling density of the trajectory points to visualize. Default is 0.3.
frame_scaling : float, optional
Scaling factor for the visualization frames' size. Affects the length and thickness of the debug frames.
Default is 1.0.
Returns
-------
node : genesis.ext.pyrender.mesh.Mesh
The created debug object representing the visualized trajectory.
Notes
-----
The function uses forward kinematics (FK) to convert joint positions to Cartesian space and render debug frames.
The density parameter reduces FK computational load by sampling fewer points, with 1.0 representing the whole
trajectory.
"""
with self._visualizer.viewer_lock:
N = len(qposs)
density = np.clip(density, 0.0, 1.0)
N_new = int(N * density)
indices = torch.linspace(0, N - 2, N_new, dtype=int)
Ts = np.zeros((N_new, 4, 4))
for i in range(N_new):
pos, quat = entity.forward_kinematics(qposs[indices[i]])
Ts[i] = tensor_to_array(gu.trans_quat_to_T(pos[link_idx], quat[link_idx]))
return self._visualizer.context.draw_debug_frames(
Ts, axis_length=frame_scaling * 0.1, origin_size=0.001, axis_radius=frame_scaling * 0.005
)
[docs]
@gs.assert_built
def render_all_cameras(
self,
rgb=True,
depth=False,
segmentation=False,
colorize_seg=False,
normal=False,
antialiasing=False,
force_render=False,
):
"""
Render the scene for all cameras using the batch renderer.
Parameters
----------
rgb : bool, optional
Whether to render the rgb image.
depth : bool, optional
Whether to render the depth image.
segmentation : bool, optional
Whether to render the segmentation image.
normal : bool, optional
Whether to render the normal image.
antialiasing : bool, optional
Whether to apply anti-aliasing.
force_render : bool, optional
Whether to force render the scene.
Returns:
A tuple of tensors of shape (n_envs, H, W, 3) if rgb is not None,
otherwise a list of tensors of shape (n_envs, H, W) if depth is not None.
If n_envs == 0, the first dimension of the tensor is squeezed.
"""
if self._visualizer.batch_renderer is None:
gs.raise_exception("Method only supported by 'BatchRenderer'")
rgb_out, depth_out, seg_out, normal_out = self._visualizer.batch_renderer.render(
rgb, depth, segmentation, normal, antialiasing, force_render
)
if segmentation and colorize_seg:
seg_out = tuple(self._visualizer.batch_renderer.colorize_seg_idxc_arr(seg) for seg in seg_out)
return rgb_out, depth_out, seg_out, normal_out
[docs]
@gs.assert_built
def clear_debug_object(self, obj):
"""
Clears all the debug objects in the scene.
"""
with self._visualizer.viewer_lock:
self._visualizer.context.clear_debug_object(obj)
[docs]
@gs.assert_built
def clear_debug_objects(self):
"""
Clears all the debug objects in the scene.
"""
with self._visualizer.viewer_lock:
self._visualizer.context.clear_debug_objects()
def _backward(self):
"""
At this point, all the scene states the simulation run should have been filled with gradients.
Next, we run backward from scene state back to scene's internal taichi variables, then back through time.
"""
if not self._backward_ready:
gs.raise_exception("Multiple backward calls not allowed.")
# backward pass through time
while self._t > 0:
self._step_grad()
self._backward_ready = False
self._forward_ready = False
[docs]
def dump_ckpt_to_numpy(self) -> dict[str, np.ndarray]:
"""
Collect every Taichi field in the **scene and its active solvers** and
return them as a flat ``{key: ndarray}`` dictionary.
Returns
-------
dict[str, np.ndarray]
Mapping ``"Class.attr[.member]" -> array`` with raw field data.
"""
arrays: dict[str, np.ndarray] = {}
for name, value in self.__dict__.items():
if isinstance(value, (ti.Field, ti.Ndarray)):
arrays[".".join((self.__class__.__name__, name))] = value.to_numpy()
for solver in self.active_solvers:
arrays.update(solver.dump_ckpt_to_numpy())
return arrays
[docs]
def save_checkpoint(self, path: str | os.PathLike) -> None:
"""
Pickle the full physics state to *one* file.
Parameters
----------
path : str | os.PathLike
Destination filename.
"""
state = {
"timestamp": time.time(),
"step_index": self.t,
"arrays": self.dump_ckpt_to_numpy(),
}
with open(path, "wb") as f:
pickle.dump(state, f, protocol=pickle.HIGHEST_PROTOCOL)
[docs]
def load_checkpoint(self, path: str | os.PathLike) -> None:
"""
Restore a file produced by :py:meth:`save_checkpoint`.
Parameters
----------
path : str | os.PathLike
Path to the checkpoint pickle.
"""
with open(path, "rb") as f:
state = pickle.load(f)
arrays = state["arrays"]
for name, value in self.__dict__.items():
if isinstance(value, (ti.Field, ti.Ndarray)):
key = ".".join((self.__class__.__name__, name))
if key in arrays:
value.from_numpy(arrays[key])
for solver in self.active_solvers:
solver.load_ckpt_from_numpy(arrays)
self._t = state.get("step_index", self._t)
# ------------------------------------------------------------------------------------
# ----------------------------------- utilities --------------------------------------
# ------------------------------------------------------------------------------------
def _sanitize_envs_idx(self, envs_idx, *, unsafe=False):
# Handling default argument and special cases
if envs_idx is None:
return self._envs_idx
if self.n_envs == 0:
gs.raise_exception("`envs_idx` is not supported for non-parallelized scene.")
if isinstance(envs_idx, slice):
return self._envs_idx[envs_idx]
if isinstance(envs_idx, (int, np.integer)):
return self._envs_idx[envs_idx : envs_idx + 1]
# Early return if unsafe
if unsafe:
return envs_idx
# Perform a bunch of sanity checks
_envs_idx = torch.as_tensor(envs_idx, dtype=gs.tc_int, device=gs.device).contiguous()
if _envs_idx is not envs_idx:
gs.logger.debug(ALLOCATE_TENSOR_WARNING)
_envs_idx = torch.atleast_1d(_envs_idx)
if _envs_idx.ndim != 1:
gs.raise_exception("Expecting a 1D tensor for `envs_idx`.")
if (_envs_idx < 0).any() or (_envs_idx >= self.n_envs).any():
gs.raise_exception("`envs_idx` exceeds valid range.")
return _envs_idx
# ------------------------------------------------------------------------------------
# ----------------------------------- properties -------------------------------------
# ------------------------------------------------------------------------------------
@property
def uid(self):
"""The unique ID of the scene."""
return self._uid
@property
def dt(self):
"""The time duration for each simulation step."""
return self._sim.dt
@property
def t(self):
"""The current simulation time step."""
return self._t
@property
def substeps(self):
"""The number of substeps per simulation step."""
return self._sim.substeps
@property
def requires_grad(self):
"""Whether the scene is in differentiable mode."""
return self._sim.requires_grad
@property
def is_built(self) -> bool:
"""Whether the scene has been built."""
return self._is_built
@property
def show_FPS(self):
"""Whether to print the frames per second (FPS) in the terminal."""
warn_once("Scene.show_FPS is deprecated. Please use profiling_options.show_FPS")
return self.profiling_options.show_FPS
@property
def gravity(self):
"""The gravity in the scene."""
return self._sim.gravity
@property
def viewer(self):
"""The viewer object for the scene."""
return self._visualizer.viewer
@property
def visualizer(self):
"""The visualizer object for the scene."""
return self._visualizer
@property
def sim(self):
"""The scene's top-level simulator."""
return self._sim
@property
def cur_t(self):
"""The current simulation time."""
return self._sim.cur_t
@property
def solvers(self):
"""All the solvers managed by the scene's simulator."""
return self._sim.solvers
@property
def active_solvers(self):
"""All the active solvers managed by the scene's simulator."""
return self._sim.active_solvers
@property
def entities(self) -> list["Entity"]:
"""All the entities in the scene."""
return self._sim.entities
@property
def emitters(self):
"""All the emitters in the scene."""
return self._emitters
@property
def tool_solver(self):
"""The scene's `tool_solver`, managing all the `ToolEntity` in the scene."""
return self._sim.tool_solver
@property
def rigid_solver(self):
"""The scene's `rigid_solver`, managing all the `RigidEntity` in the scene."""
return self._sim.rigid_solver
@property
def avatar_solver(self):
"""The scene's `avatar_solver`, managing all the `AvatarEntity` in the scene."""
return self._sim.avatar_solver
@property
def mpm_solver(self):
"""The scene's `mpm_solver`, managing all the `MPMEntity` in the scene."""
return self._sim.mpm_solver
@property
def sph_solver(self):
"""The scene's `sph_solver`, managing all the `SPHEntity` in the scene."""
return self._sim.sph_solver
@property
def fem_solver(self):
"""The scene's `fem_solver`, managing all the `FEMEntity` in the scene."""
return self._sim.fem_solver
@property
def pbd_solver(self):
"""The scene's `pbd_solver`, managing all the `PBDEntity` in the scene."""
return self._sim.pbd_solver
@property
def segmentation_idx_dict(self):
"""
Returns a dictionary mapping segmentation indices to scene entities.
In the segmentation map:
- Index 0 corresponds to the background (-1).
- Indices > 0 correspond to scene elements, which may be represented as:
- `entity_id`
- `(entity_id, link_id)`
- `(entity_id, link_id, geom_id)`
depending on the material type and the configured segmentation level.
"""
return self._visualizer.segmentation_idx_dict
