Advanced Vision

Arenas

class flygym.examples.vision.arena.MovingObjArena(size=(300, 300), friction=(1, 0.005, 0.0001), obj_radius=1, init_ball_pos=(5, 0), move_speed=10, move_direction='right', lateral_magnitude=2)

Bases: BaseArena

Flat terrain with a hovering moving object.

Parameters:

sizetuple[int, int]

The size of the terrain in (x, y) dimensions.

frictiontuple[float, float, float]

Sliding, torsional, and rolling friction coefficients, by default (1, 0.005, 0.0001)

obj_radiusfloat

Radius of the spherical floating object in mm.

init_ball_postuple[float,float]

Initial position of the object, by default (5, 0).

move_speedfloat

Speed of the moving object. By default 10.

move_directionstr

Which way the ball moves toward first. Can be “left”, “right”, or “random”. By default “right”.

lateral_magnitudefloat

Magnitude of the lateral movement of the object as a multiplier of forward velocity. For example, when lateral_magnitude is 1, the object moves at a heading (1, 1) when its movement is the most lateral. By default 2.

Attributes:

ball_postuple[float,float,float]

The position of the floating object in the arena.

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 coordinates 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 (usually n=4: 2 antennae + 2 maxillary palps).

get_spawn_position(rel_pos, rel_angle)

Given a relative entity spawn position and orientation (as if it was a simple flat terrain), return the adjusted position and orientation. This is useful for environments that have complex terrain (e.g. 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 features 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).

Returns:

np.ndarray

Adjusted (x, y, z) position of the entity.

np.ndarray

Adjusted euler angles (rotations along x, y, z in radian) of the fly’s orientation.

init_lights()

property odor_dimensions: int

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

Returns:

int

The dimension of the odor space.

post_visual_render_hook(physics: dm_control.mjcf.Physics, *args, **kwargs) → None

Make necessary changes (e.g. 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 (e.g. make certain visualization markers transparent) before rendering the visual inputs. By default, this does nothing.

reset(physics)

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, physics)

Advance the arena by one step. This is useful for interactive environments (e.g. moving object). Typically, this method is called from the core simulation class (e.g. 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 (e.g. 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.examples.vision.arena.MovingFlyArena(terrain_type: str = 'flat', x_range: tuple[float, float] | None = (-10, 20), y_range: tuple[float, float] | None = (-20, 20), block_size: float | None = 1.3, height_range: tuple[float, float] | None = (0.2, 0.2), rand_seed: int = 0, ground_alpha: float = 1, friction=(1, 0.005, 0.0001), leading_fly_height=0.5, init_fly_pos=(5, 0), move_speed=10, radius=10)

Bases: BaseArena

Flat terrain with a hovering moving fly.

Parameters:

terrain_typestr

Type of terrain. Can be “flat” or “blocks”. By default “flat”.

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, 35).

y_rangetuple[float, float], optional

Same as above in y, by default (-20, 20).

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.2, 0.2).

rand_seedint, optional

Seed for generating random block heights, by default 0.

ground_alphafloat, optional

Opacity of the ground, by default 1 (fully opaque).

frictiontuple[float, float, float]

Sliding, torsional, and rolling friction coefficients, by default (1, 0.005, 0.0001)

init_fly_postuple[float,float]

Initial position of the fly, by default (5, 0).

move_speedfloat

Speed of the moving fly. By default 10.

move_directionstr

Which way the fly moves toward first. Can be “left”, “right”, or “random”. By default “right”.

lateral_magnitudefloat

Magnitude of the lateral movement of the fly as a multiplier of forward velocity. For example, when lateral_magnitude is 1, the fly moves at a heading (1, 1) when its movement is the most lateral. By default 2.

Attributes:

fly_postuple[float,float,float]

The position of the floating fly in the arena.

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 coordinates 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 (usually n=4: 2 antennae + 2 maxillary palps).

get_spawn_position(rel_pos, rel_angle)

Given a relative entity spawn position and orientation (as if it was a simple flat terrain), return the adjusted position and orientation. This is useful for environments that have complex terrain (e.g. 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 features 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).

Returns:

np.ndarray

Adjusted (x, y, z) position of the entity.

np.ndarray

Adjusted euler angles (rotations along x, y, z in radian) of the fly’s orientation.

init_lights()

property odor_dimensions: int

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

Returns:

int

The dimension of the odor space.

post_visual_render_hook(physics: dm_control.mjcf.Physics, *args, **kwargs) → None

Make necessary changes (e.g. 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 (e.g. make certain visualization markers transparent) before rendering the visual inputs. By default, this does nothing.

reset(physics)

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, physics)

Advance the arena by one step. This is useful for interactive environments (e.g. moving object). Typically, this method is called from the core simulation class (e.g. 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 (e.g. 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.examples.vision.arena.MovingBarArena(azimuth_func: Callable[[float], float], visual_angle=(10, 60), distance=12, rgba=(0, 0, 0, 1), **kwargs)

Bases: Tethered

Flat or blocks terrain with a moving cylinder to simulate a moving bar on a circular screen.

Parameters:

azimuth_funcCallable[[float], float]

Function that takes time as input and returns the azimuth angle of the cylinder.

visual_angletuple[float, float]

Width and height of the cylinder in degrees.

distancefloat

Distance from the center of the arena to the center of the cylinders.

rgbatuple[float, float, float, float]

Color of the cylinder.

kwargsdict

Additional arguments to passed to the superclass.

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 coordinates 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 (usually n=4: 2 antennae + 2 maxillary palps).

get_pos(t)

Returns the position of the cylinder at time t.

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 adjusted position and orientation. This is useful for environments that have complex terrain (e.g. 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 features 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).

Returns:

np.ndarray

Adjusted (x, y, z) position of the entity.

np.ndarray

Adjusted euler angles (rotations along x, y, z in radian) of the fly’s orientation.

init_lights()

property odor_dimensions: int

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

Returns:

int

The dimension of the odor space.

post_visual_render_hook(physics: dm_control.mjcf.Physics, *args, **kwargs) → None

Make necessary changes (e.g. 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 (e.g. make certain visualization markers transparent) before rendering the visual inputs. By default, this does nothing.

reset(physics)

Resets the position of the cylinder.

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

Add an entity (e.g. the fly) to the arena.

Parameters:

entitymjcf.RootElement

The entity to be added to the arena.

rel_posnp.ndarray

(x, y, z) position of the entity if it were spawned on a simple flat environment.

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, physics)

Updates the position of the cylinder.

class flygym.examples.vision.arena.ObstacleOdorArena(terrain: BaseArena, obstacle_positions: ndarray = np.array([(7.5, 0), (12.5, 5), (17.5, -5)]), obstacle_colors: ndarray | tuple = (0, 0, 0, 1), obstacle_radius: float = 1, obstacle_height: float = 4, odor_source: ndarray = np.array([[25, 0, 2]]), peak_odor_intensity: ndarray = np.array([[1]]), diffuse_func: Callable = lambda x: ..., marker_colors: list[tuple[float, float, float, float]] | None = None, marker_size: float = 0.1, user_camera_settings: tuple[tuple[float, float, float], tuple[float, float, float], float] | None = None)

Bases: BaseArena

friction = (100.0, 0.005, 0.0001)

get_olfaction(antennae_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 coordinates 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 (usually n=4: 2 antennae + 2 maxillary palps).

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 adjusted position and orientation. This is useful for environments that have complex terrain (e.g. 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 features 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).

Returns:

np.ndarray

Adjusted (x, y, z) position of the entity.

np.ndarray

Adjusted euler angles (rotations along x, y, z in radian) of the fly’s orientation.

init_lights()

num_sensors = 4

property odor_dimensions: int

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

Returns:

int

The dimension of the odor space.

post_visual_render_hook(physics)

Make necessary changes (e.g. make certain visualization markers opaque) after rendering the visual inputs. By default, this does nothing.

pre_visual_render_hook(physics)

Make necessary changes (e.g. 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 (e.g. moving object). Typically, this method is called from the core simulation class (e.g. 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 (e.g. NeuroMechFly.physics) when the core simulation calls this method.

*args

User defined arguments and keyword arguments.

**kwargs

User defined arguments and keyword arguments.

Simple visual taxis

class flygym.examples.vision.simple_visual_taxis.VisualTaxis(camera: Camera, obj_threshold=0.15, decision_interval=0.05, **kwargs)

Bases: HybridTurningController

A simple visual taxis task where the fly has to follow a moving object.

Parameters:

cameraCamera

The camera to be used for rendering.

obj_thresholdfloat

The threshold for object detection. Minimum and maximum brightness values are 0 and 1. If an ommatidium’s intensity reading is below this value, then it is considered that this ommatidium is seeing the object.

decision_intervalfloat

The interval between updates of descending drives, in seconds.

kwargs

Additional keyword arguments to be passed to HybridTurningController.__init__.

Notes

Please refer to the “MPD Task Specifications” page of the API references for the detailed specifications of the action space, the observation space, the reward, the “terminated” and “truncated” flags, and the “info” dictionary.

reset(seed=0, **kwargs)

See HybridTurningController.reset.

step(control_signal)

Step the simulation forward in time. Note that this method is to be called every time the descending steering signals are updated. This typically includes many forward steps of the physics simulation.

Parameters:

control_signalarray_like

The control signal to apply to the simulation.

Returns:

visual_featuresarray_like

The preprocessed visual features extracted from the observation.

rewardfloat

The reward obtained from the current step.

terminatedbool

Whether the episode is terminated or not. Always False.

truncatedbool

Whether the episode is truncated or not. Always False.

infodict

Additional information about the step.

Connectome-constrained vision model

class flygym.examples.vision.RealTimeVisionNetwork(connectome: Namespace = Namespace(file='fib25-fib19_v2.2.json', extent=15, n_syn_fill=1), dynamics: Namespace = Namespace(type='PPNeuronIGRSynapses', activation=Namespace(type='relu')), node_config: Namespace = Namespace(bias=Namespace(type='RestingPotential', groupby=['type'], initial_dist='Normal', mode='sample', requires_grad=True, mean=0.5, std=0.05, penalize=Namespace(activity=True), seed=0), time_const=Namespace(type='TimeConstant', groupby=['type'], initial_dist='Value', value=0.05, requires_grad=True)), edge_config: Namespace = Namespace(sign=Namespace(type='SynapseSign', initial_dist='Value', requires_grad=False, groupby=['source_type', 'target_type']), syn_count=Namespace(type='SynapseCount', initial_dist='Lognormal', mode='mean', requires_grad=False, std=1.0, groupby=['source_type', 'target_type', 'dv', 'du']), syn_strength=Namespace(type='SynapseCountScaling', initial_dist='Value', requires_grad=True, scale_elec=0.01, scale_chem=0.01, clamp='non_negative', groupby=['source_type', 'target_type'])))

Bases: Network

This class extends flyvision.network.Network. The main difference is that flyvision.network.Network receives the entire history of visual input as a block, which enables more efficient computation on the GPU. In contrast, RealTimeVisionNetwork receives visual input one frame at a time, allowing for deployment in closed-loop simulations. See flyvision and Lappalainen et al., 2024 for more details.

cleanup_step_by_step_simulation() → None

Clean up the network after the simulation ends. This clears the parameters that were set up for the step-by-step simulation and resets the gradient tracking / training flags.

forward_one_step(curr_visual_input: Tensor) → AutoDeref | Tensor

Simulate the network one step forward.

Parameters:

curr_visual_inputTensor

Raw visual input experienced by the fly (i.e., intensity reading from each ommatidium). This is a tensor of shape (num_samples, num_ommatidia).

Returns:

Union[AutoDeref, Tensor]

If as_states is set to True in __init__, this returns the network state after stepping the network simulation by one step. The return value is of type flyvision.utils.tensor_utils.AutoDeref. Otherwise, the actual activities of the nodes are returned as a torch tensor (i.e., agnostic to the flyvision state representation).

setup_step_by_step_simulation(dt: float, initial_state: str | AutoDeref | None = 'auto', as_states: bool = False, num_samples: int = 1) → None

Set up the network for step-by-step simulation.

Parameters:

dtfloat

Integration time step for the visual system neural network simulation. Note that this is typically different from (larger than) the time step of the physics simulation.

initial_stateUnion[str, AutoDeref, None], optional

Initial state of the network. The default is “auto”, which establishes a steady state after 1s of gray input. See RealTimeVisionNetwork.steady_state for more details.

as_statesbool, optional

Whether to return the network state or just the activities of the nodes. The default is False.

num_samplesint, optional

Number of samples to simulate in parallel. The default is 1. The user might want to change this to 2 since there are two compound eyes in the fly.

class flygym.examples.vision.RealTimeVisionNetworkView(network_dir: PathLike | NetworkDir)

Bases: NetworkView

This class extends flyvision.network.NetworkView to work with our extended RealTimeVisionNetwork. In brief, it is used as a handle to set up the RealTimeVisionNetwork from saved checkpoint. See flyvision and Lappalainen et al., 2024 for more details.

init_network(chkpt='best_chkpt', network: RealTimeVisionNetwork | None = None) → Network

Initialize the pretrained network.

Parameters:

chkpt: str

Checkpoint to load. Default: “best_chkpt”.

network: RealTimeVisionNetwork, optional

Network instance to initialize. If None, a new instance will be created.

Returns:

RealTimeVisionNetwork

The network instance.

class flygym.examples.vision.RetinaMapper(retina: Retina | None = None, boxeye: BoxEye | None = None)

Bases: object

Both flyvision and flygym use a hexagonal grid of ommatidia to model the compound eyes of the fly. To approximate the the correct number of ommatidia per eye (about 700-800), the two libraries even share the same size of the grid. However, the two libraries use different indexing conventions for the ommatidia. This class provides methods to convert stimuli between the coordinate systems of flyvision’s BoxEye representation and flygym’s Retina representation.

flygym_to_flyvis(flygym_stimulus: ndarray) → ndarray

Convert a stimulus from flygym’s Retina representation to flyvision’s BoxEye representation.

Parameters:

flygym_stimulusnp.ndarray

Any value (e.g., intensities, neural activities) associated with the ommatidia in flygym’s Retina ordering. The shape is (…, num_ommatidia): in other words, this method works as long as the size along the last dimension is the same as the number of ommatidia (721).

Returns:

np.ndarray

The same values, but now in flyvision’s BoxEye ordering.

flyvis_to_flygym(flyvis_stimulus: ndarray) → ndarray

Convert a stimulus from flyvision’s BoxEye representation to flygym’s Retina representation.

Parameters:

flyvis_stimulusnp.ndarray

Any value (e.g., intensities, neural activities) associated with the ommatidia in flyvision’s BoxEye ordering. The shape is (…, num_ommatidia): in other words, this method works as long as the size along the last dimension is the same as the number of ommatidia (721).

Returns:

np.ndarray

The same values, but now in flygym’s Retina ordering.

class flygym.examples.vision.RealisticVisionFly(vision_network_dir=None, *args, **kwargs)

Bases: HybridTurningFly

This class extends the HybridTurningFly to couple it with the visual system network from Lappalainen et al., 2024. This allows the user to receive, as a part of the observation, the activities of visual system neurons.

Parameters:

vision_network_dirstr, optional

Path to the directory containing the vision network checkpoint. If not provided, model 000 from Lappalainen et al., 2024 will be used.

Notes

Please refer to the “MDP Task Specifications” page of the API references for the detailed specifications of the action space, the observation space, the reward, the “terminated” and “truncated” flags, and the “info” dictionary.

close()

Close the fly. See HybridTurningFly.close.

post_step(sim: Simulation)

Same as HybridTurningController, except the additional nn_activities key in the info dictionary, which contains the activities of the visual system neurons as a flyvision.LayerActivity object, and the nn_activities_arr key in the observation dictionary, which contains the activities of the visual system neurons, represented as a numpy array of shape (2, num_cells_per_eye). The 0th dimension corresponds to the eyes in the order (left, right).

reset(*args, **kwargs)

Reset the fly. See HybridTurningFly.reset.