SofaEnv Base Class

SofaEnv is the abstract class for SOFA simulation environments. When you create a new scene, you should use SofaEnv as the parent class of your task specific environment. SofaEnv handles calling the createScene function of your scene description, setting up rendering if needed, acting as the interface to the actual SOFA simulation, and exposing Gym functions.

Note

Reading the RGB buffer from OpenGL through pyglet does not work on WSL -> set the rendermode to WSL if you want to retrieve image observations.

Classes:

RenderFramework(value)

RenderFramework options for SofaEnv.

RenderMode(value)

RenderMode options for SofaEnv.

SofaEnv(scene_path[, time_step, frame_skip, ...])

Abstract class for SOFA simulation environments.

class sofa_env.base.RenderFramework(value)

Bases: Enum

RenderFramework options for SofaEnv.

This enum specifies which library will be used between pyglet and pygame

Attributes:

PYGAME

PYGLET

PYGAME = 1
PYGLET = 0
class sofa_env.base.RenderMode(value)

Bases: Enum

RenderMode options for SofaEnv.

This enum specifies if you want to simulate

  • state based without rendering anything (NONE),

  • generate image observations headlessly without creating a window (HEADLESS),

  • generate image observations headlessly, but only when env.update_rgb_buffer() or env.update_rgb_buffer_remote() are called manually (MANUAL).

  • create a window to observe the simulation for a remote workstation or WSL (REMOTE).

  • or create a window to observe the simulation (HUMAN).

Attributes:

HEADLESS

HUMAN

MANUAL

NONE

REMOTE

HEADLESS = 1
HUMAN = 0
MANUAL = 4
NONE = 3
REMOTE = 2
class sofa_env.base.SofaEnv(scene_path: str | Path, time_step: float = 0.01, frame_skip: int = 1, render_mode: RenderMode = RenderMode.NONE, create_scene_kwargs: dict | None = None, render_framework: RenderFramework = RenderFramework.PYGLET)

Bases: Env

Abstract class for SOFA simulation environments. Specific simulations can either subclass this class or use it as a template and implement everything themselves.

In either case, there are a few requirements on SofaEnv classes:

  • no SOFA code should be called in __init__(). Instead, all initialization should be done on the first call to reset(). This class does this by checking the status of a boolean on each call to reset().

  • the class must possess members called observation_space and action_space, which are subclasses of gymnasium.spaces.Space.

  • getting rendered images from SOFA assumes, that you created a camera in your scene description. Your createScene function should thus return a dictionary, that has a key value pair of 'camera': Union[Sofa.Core.Object, sofa_env.sofa_templates.camera.Camera].

  • the scene_description script has to contain a createScene function that takes at least a Sofa.Core.Node as its root node. It can optionally accept keyword arguments that you pass through create_scene_kwargs.

  • you are responsible for implementing the functions step, reset, and _do_action.

Notes

  • if using vectorized environments, MUST use a subprocess wrapper because only one SOFA simulation can exist per process. It will not throw an error, but the simulations will be invalid.

  • for rendering or generating image observations, you have the options to set render_mode to one of the RenderMode enum cases. (1) HUMAN will create a pyglet window and render the images into that window. This classe’s step and reset function will also return this image as a numpy array. (2) HEADLESS will do the same thing, but without a window. Pyglet will use EGL to create a render context, that does not need an actual window. (3) REMOTE Create and show a pyglet Window (similar to HUMAN render mode) for a remote workstation or when working under WSL. The exported display has to be adjusted and a display server like Xming or Mobaxterm is needed for the visualization. (4) NONE is the case where you are not interested in visual observations. Use this, if you are only interested in the states of the simulation.

Parameters:
  • scene_path (Union[Path, str]) – absolute path to the scene file (.py) to load.

  • time_step (float) – size of simulation time step in seconds (default: 0.01).

  • frame_skip (int) – number of simulation time steps taken (call _do_action and advance simulation) each time step is called (default: 1).

  • render_mode (RenderMode) – create a window (RenderMode.HUMAN) or run headless (RenderMode.HEADLESS).

  • create_scene_kwargs (Optional[dict]) – a dictionary to pass additional keyword arguments to the createScene function.

  • render_framework (RenderFramework) – choose between pyglet and pygame for rendering

Attributes:

action_space

observation_space

Methods:

close()

Performs necessary cleanup when environment is no longer needed.

get_depth_from_open_gl()

Reads the depth buffer from OpenGL and returns a depth array with absolute distance to the camera values.

get_depth_from_pygame()

Reads the rgb buffer from openGl for pygame and returns a copy.

get_depth_from_pyglet()

Reads the depth buffer from pyglet and returns a copy.

get_rgb_from_open_gl()

Reads the rgb buffer from OpenGL and returns a copy.

get_rgb_from_pygame()

Reads the rgb buffer from pygame and returns a copy.

get_rgb_from_pyglet()

Reads the rgb buffer from pyglet and returns a copy.

render([mode])

Returns the rgb observation from the simulation.

reset([seed, options])

Resets the SOFA simulation.

step(action)

Runs #frame_skip timesteps of the environment's dynamics.

action_space: spaces.Space[ActType]
close() None

Performs necessary cleanup when environment is no longer needed.

get_depth_from_open_gl() ndarray

Reads the depth buffer from OpenGL and returns a depth array with absolute distance to the camera values.

get_depth_from_pygame() ndarray

Reads the rgb buffer from openGl for pygame and returns a copy.

get_depth_from_pyglet() ndarray

Reads the depth buffer from pyglet and returns a copy.

Note

This function differs from get_depth_from_open_gl in that this function returns the depth buffer as a uint8 value in [0, 255]. get_depth_from_open_gl returns the depth information as float values in absolute distances to the camera.

get_rgb_from_open_gl() ndarray

Reads the rgb buffer from OpenGL and returns a copy.

get_rgb_from_pygame() ndarray

Reads the rgb buffer from pygame and returns a copy.

get_rgb_from_pyglet() ndarray

Reads the rgb buffer from pyglet and returns a copy.

Note

Pyglet widows have a front and back buffer that are exchanged when window.flip() is called. The content of the back buffer is undefined, which is why we return a copy. Without the copy, we could get out a bit more performance from the environment.

observation_space: spaces.Space[ObsType]
render(mode: str | None = None) ndarray

Returns the rgb observation from the simulation.

reset(seed: int | SeedSequence | None = None, options: Dict[str, Any] | None = None) Tuple[ndarray | None, Dict]

Resets the SOFA simulation. If this is the first env.reset(), the simulation is initialized.

Parameters:
  • seed (Union[int, np.random.SeedSequence, None]) – the seed to use for the environment.

  • options (Optional[Dict[str, Any]]) – a dictionary of options.

Returns:

the current visual observation from the env, if the render mode is not RenderMode.NONE

Return type:

rgb_buffer (Union[np.ndarray, None])

Note

Your implementation of this function should 1. call super().reset(), 2. manually reset any (SOFA) components that need resetting (e.g. by setting a new pose), and 3. return the initial observation.

Most environments have SOFA objects that each have their own random number generator. To correctly seed the environment, you need to set the seed of each SOFA object. However, we can only do that, if the environment is already initialized, so reset() was called at least once. See reset() in DeflectSpheresEnv for an example.

rng: np.random.Generator | None
seed_sequence: np.random.SeedSequence | None
step(action: Any) ndarray | None

Runs #frame_skip timesteps of the environment’s dynamics. The action will be applied each time. When the end of an episode is reached, you are responsible for calling reset() to reset this environment’s state.

Parameters:

action (instance of self.action_space) – an action that should be applied frame_skip times as defined in _do_action.

Returns:

The current visual observation from the env, if the render mode is not RenderMode.NONE

Return type:

rgb_buffer (Union[np.ndarray, None])

Note

In order for this function to work properly, please implement env._do_action(action) to describe how an action should be applied to the SOFA simulation. The function will be called frame_skip times.

Your implementation of env.step(action) should

  1. call rgb_array = super().step(action), (rgb_array will be None if render_mode is RenderMode.NONE)

  2. calculate and return next observation, reward, done, and info.

  • observation (instance of self.observation_space): the observation from the env

  • reward (float) : amount of reward for the current state, state action pair, or state transition

  • done (bool): whether the episode has ended

  • info (dict): auxiliary diagnostic information for logging and debugging

For full control of how steps are applied to the simulation, you can also reimplement the complete step function and do not call the super().step(action). Just make sure you also call self._maybe_update_rgb_buffer() to get valid images.