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.