Arena¶

An arena is a physical environment in which the simulated fly is placed. Arenas can have rugged surfaces, visual features, odor features, or any combination thereof. The user can implement their own arenas by inheriting from the flygym.mujoco.arena.BaseArena class.

This page provides the API reference for the BaseArena abstract class as well as the preprogrammed arenas.

Base arena¶

class flygym.mujoco.arena.BaseArena(*args: List, **kwargs: Dict)¶

Bases: ABC

Base class for all arenas.

Attributes:

arenaAny: The arena object that the terrain is built on. Exactly what it is depends on the physics simulator.
frictionTuple [float]: Default sliding, torsional, and rolling friction coefficients of surfaces. This is provided for the user’s convinience but can be overriden for either all or some surfaces.

friction = (100.0, 0.005, 0.0001)¶

get_olfaction(sensor_pos: ndarray) → ndarray¶

Get the odor intensity readings from the environment.

Parameters:

sensor_posnp.ndarray: The Cartesian coordinates of the antennae of the fly as a (n, 3) NumPy array where n is the number of sensors (usually n=4: 2 antennae + 2 maxillary palps), and the second dimension gives the corrdinates in (x, y, z).

Returns:

np.ndarray: The odor intensity readings from the environment as a (k, n) NumPy array where k is the dimension of the odor signal and n is the number of odor sensors (usally n=4: 2 antennae + 2 maxillary palps).

abstract get_spawn_position(rel_pos: ndarray, rel_angle: ndarray) → Tuple[ndarray, ndarray]¶

Given a relative entity spawn position and orientation (as if it was a simple flat terrain), return the abjusted position and orientation. This is useful for environments that have complex terrain (eg. with obstacles) where the entity’s spawn position needs to be shifted accordingly.

For example, if the arena has flat terrain, this method can simply return rel_pos, rel_angle unchanged (as is the case by default). If there is are featues on the ground that are 0.1 mm in height, then this method should return rel_pos + [0, 0, 0.1], rel_angle.

Parameters:

rel_posnp.ndarray: (x, y, z) position of the entity in mm as supplied by the user (before any transformation).
rel_anglenp.ndarray: Euler angle (rotation along x, y, z in radian) of the fly’s orientation as supplied by the user (before any transformation).
*args: User defined arguments and keyword arguments.
**kwargs: User defined arguments and keyword arguments.

Returns:

np.ndarray: Adjusted (x, y, z) position of the entity.
np.ndarray: Adjusted euler angle (rotation along x, y, z in raidan) of the fly’s oreintation.

property odor_dimensions: int¶

The dimension of the odor signal. This can be used to emulate multiple monomolecular chemical concentrations or multiple composite ordor intensities.

Returns:

int: The dimension of the odor space.

post_visual_render_hook(physics: dm_control.mjcf.Physics, *args, **kwargs) → None¶: Make necessary changes (eg. make certain visualization markers opaque) after rendering the visual inputs. By default, this does nothing.

pre_visual_render_hook(physics: dm_control.mjcf.Physics, *args, **kwargs) → None¶: Make necessary changes (eg. make certain visualization markers transparent) before rendering the visual inputs. By default, this does nothing.

spawn_entity(entity: Any, rel_pos: ndarray, rel_angle: ndarray) → None¶

Add the fly to the arena.

Parameters:

entitymjcf.RootElement: The entity to be added to the arena (this should be the fly).
rel_posnp.ndarray: (x, y, z) position of the entity.
rel_anglenp.ndarray: euler angle representation (rot around x, y, z) of the entity’s orientation if it were spawned on a simple flat terrain.

step(dt: float, physics: dm_control.mjcf.Physics, *args, **kwargs) → None¶

Advance the arena by one step. This is useful for interactive environments (eg. moving object). Typically, this method is called from the core simulation class (eg. NeuroMechFly).

Parameters:

dtfloat: The time step in seconds since the last update. Typically, this is the same as the time step of the physics simulation (provided that this method is called by the core simulation every time the simulation steps).
physicsmjcf.Physics: The physics object of the simulation. This is typically provided by the core simulation class (eg. NeuroMechFly.physics) when the core simulation calls this method.
*args: User defined arguments and keyword arguments.
**kwargs: User defined arguments and keyword arguments.

class flygym.mujoco.arena.FlatTerrain(size: Tuple[float, float] = (50, 50), friction: Tuple[float, float, float] = (1, 0.005, 0.0001), ground_alpha: float = 1.0, scale_bar_pos: Tuple[float, float, float] | None = None)¶

Flat terrain with no obstacles.

Parameters:

sizeTuple[float, float], optional: The size of the arena in mm, by default (50, 50).
frictionTuple[float, float, float]: The sliding, torsional, and rolling friction coefficients of the ground, by default (1, 0.005, 0.0001).
ground_alphafloat: Opacity of the ground, by default 1 (fully opaque).
scale_bar_posTuple[float, float, float], optional: If supplied, a 1 mm scale bar will be placed at this location.

Attributes:

root_elementmjcf.RootElement: The root MJCF element of the arena.
frictionTuple[float, float, float]: The sliding, torsional, and rolling friction coefficients of the ground, by default (1, 0.005, 0.0001).

Preprogrammed complex terrain arenas¶

class flygym.mujoco.arena.GappedTerrain(x_range: Tuple[float, float] = (-10, 20), y_range: Tuple[float, float] = (-10, 10), friction: Tuple[float, float, float] = (1, 0.005, 0.0001), gap_width: float = 0.4, block_width: float = 1.0, gap_depth: float = 2.0, ground_alpha: float = 1.0, scale_bar_pos: Tuple[float, float, float] | None = None)¶

Terrain with horizontal gaps.

Parameters:

x_rangeTuple[float, float], optional: Range of the arena in the x direction (anterior-posterior axis of the fly) over which the block-gap pattern should span, by default (-10, 20)
y_rangeTuple[float, float], optional: Same as above in y, by default (-10, 10).
frictionTuple[float, float, float], optional: Sliding, torsional, and rolling friction coefficients, by default (1, 0.005, 0.0001).
gap_widthfloat, optional: Width of each gap, by default 0.4.
block_widthfloat, optional: Width of each block (piece of floor), by default 1.
gap_depthfloat, optional: Height of the gaps, by default 2.
ground_alphafloat, optional: Opacity of the ground, by default 1 (fully opaque).
scale_bar_posTuple[float, float, float], optional: If supplied, a 1 mm scale bar will be placed at this location.

Attributes:

root_elementmjcf.RootElement: The root MJCF element of the arena.
frictionTuple[float, float, float]: The sliding, torsional, and rolling friction coefficients of the ground, by default (1, 0.005, 0.0001).
x_rangeTuple[float, float]: Range of the arena in the x direction (anterior-posterior axis of the fly) over which the block-gap pattern should span, by default (-10, 20).
y_rangeTuple[float, float]Í: Same as above in y, by default (-10, 10)
gap_widthfloat: Width of each gap, by default 0.4
block_widthfloat: Width of each block (piece of floor), by default 1
gap_depthfloat: Height of the gaps, by default 2

class flygym.mujoco.arena.BlocksTerrain(x_range: Tuple[float, float] = (-10, 20), y_range: Tuple[float, float] = (-10, 10), friction: Tuple[float, float, float] = (1, 0.005, 0.0001), block_size: float = 1.3, height_range: Tuple[float, float] = (0.45, 0.45), ground_alpha: float = 1.0, rand_seed: int = 0, scale_bar_pos: Tuple[float, float, float] | None = None)¶

Terrain formed by blocks at random heights.

Parameters:

x_rangeTuple[float, float], optional: Range of the arena in the x direction (anterior-posterior axis of the fly) over which the block-gap pattern should span, by default (-10, 20).
y_rangeTuple[float, float], optional: Same as above in y, by default (-10, 10).
frictionTuple[float, float, float]: Sliding, torsional, and rolling friction coefficients, by default (1, 0.005, 0.0001).
block_sizefloat, optional: The side length of the rectangular blocks forming the terrain in mm, by default 1.3.
height_rangeTuple[float, float], optional: Range from which the height of the extruding blocks should be sampled. Only half of the blocks arranged in a diagonal pattern are extruded, by default (0.45, 0.45).
ground_alphafloat, optional: Opacity of the ground, by default 1 (fully opaque).
rand_seedint, optional: Seed for generating random block heights, by default 0.
scale_bar_posTuple[float, float, float], optional: If supplied, a 1 mm scale bar will be placed at this location.

Attributes:

root_elementmjcf.RootElement: The root MJCF element of the arena.
x_rangeTuple[float, float], optional: Range of the arena in the x direction (anterior-posterior axis of the fly) over which the block-gap pattern should span, by default (-10, 20).
y_rangeTuple[float, float], optional: Same as above in y, by default (-10, 10).
frictionTuple[float, float, float]: Sliding, torsional, and rolling friction coefficients, by default (1, 0.005, 0.0001).
block_sizefloat: The side length of the rectangular blocks forming the terrain in mm, by default 1.3.
height_rangeTuple[float, float]: Range from which the height of the extruding blocks should be sampled. Only half of the blocks arranged in a diagonal pattern are extruded, by default (0.45, 0.45).

class flygym.mujoco.arena.MixedTerrain(friction: Tuple[float, float, float] = (1, 0.005, 0.0001), gap_width: float = 0.5, gapped_block_width: float = 1.0, gap_depth: float = 2, block_size: float = 1.3, height_range: Tuple[float, float] = (0.45, 0.45), ground_alpha: float = 1.0, rand_seed: int = 0, scale_bar_pos: Tuple[float, float, float] | None = None)¶

A mixture of flat, blocks, and gaps terrains.

Parameters:

frictionTuple[float, float, float], optional: Sliding, torsional, and rolling friction coefficients, by default (1, 0.005, 0.0001)
gap_widthfloat, optional: Width of each gap, by default 0.4.
gapped_block_widthfloat, optional: Width of each block (piece of floor), by default 1.
gap_depthfloat, optional: Height of the gaps, by default 2.
block_sizefloat, optional: The side length of the rectangular blocks forming the terrain in mm, by default 1.3.
height_rangeTuple[float, float], optional: Range from which the height of the extruding blocks should be sampled. Only half of the blocks arranged in a diagonal pattern are extruded, by default (0.45, 0.45).
ground_alphafloat, optional: Opacity of the ground, by default 1 (fully opaque).
rand_seedint, optional: Seed for generating random block heights, by default 0.
scale_bar_posTuple[float, float, float], optional: If supplied, a 1 mm scale bar will be placed at this location.

Attributes:

root_elementmjcf.RootElement: The root MJCF element of the arena.
frictionTuple[float, float, float]: Sliding, torsional, and rolling friction coefficients, by default (1, 0.005, 0.0001).

Preprogrammed arenas with sensory features¶

class flygym.mujoco.arena.OdorArena(size: Tuple[float, float] = (300, 300), friction: Tuple[float, float, float] = (1, 0.005, 0.0001), num_sensors: int = 4, odor_source: ndarray = np.array([[10, 0, 0]]), peak_intensity: ndarray = np.array([[1]]), diffuse_func: Callable = lambda x: ..., marker_colors: List[Tuple[float, float, float, float]] | None = None, marker_size: float = 0.25)¶

Flat terrain with an odor source.

Parameters:

sizeTuple[float, float], optional: The size of the arena in mm, by default (300, 300).
frictionTuple[float, float, float], optional: The sliding, torsional, and rolling friction coefficients of the ground, by default (1, 0.005, 0.0001).
num_sensorsint, optional: The number of odor sensors, by default 4: 2 antennae + 2 maxillary palps.
odor_sourcenp.ndarray, optional: The position of the odor source in (x, y, z) coordinates. The shape of the array is (n_sources, 3).
peak_intensitynp.ndarray, optional: The peak intensity of the odor source. The shape of the array is (n_sources, n_dimensions). Note that the odor intensity can be multidimensional.
diffuse_funcCallable, optional: The function that, given a distance from the odor source, returns the relative intensity of the odor. By default, this is a inverse square relationship.
marker_colorsList[Tuple[float, float, float, float]], optional: A list of n_sources RGBA values (each as a tuple) indicating the colors of the markers indicating the positions of the odor sources. The RGBA values should be given in the range [0, 1]. By default, the matplotlib color cycle is used.
marker_sizefloat, optional: The size of the odor source markers, by default 0.25.

Attributes:

root_elementmjcf.RootElement: The root MJCF element of the arena.
frictionTuple[float, float, float]: The sliding, torsional, and rolling friction coefficients of the ground, by default (1, 0.005, 0.0001)
num_sensorsint: The number of odor sensors, by default 4: 2 antennae + 2 maxillary palps.
odor_sourcenp.ndarray: The position of the odor source in (x, y, z) coordinates. The shape of the array is (n_sources, 3).
peak_intensitynp.ndarray: The peak intensity of the odor source. The shape of the array is (n_sources, n_dimensions). Note that the odor intensity can be multidimensional.
num_odor_sourcesint: Number of odor sources.
odor_dimensionsint: The dimension of the odor signal.
diffuse_funcCallable: The function that, given a distance from the odor source, returns the relative intensity of the odor. By default, this is a inverse square relationship.
birdeye_camdm_control.mujoco.Camera: MuJoCo camera that gives a birdeye view of the arena.
birdeye_cam_zoomdm_control.mujoco.Camera: MuJoCo camera that gives a birdeye view of the arena, zoomed in toward the fly.