rtgym package

Subpackages

• rtgym.agent package
  • Subpackages
    • rtgym.agent.behavior package
      • Submodules
      • rtgym.agent.behavior.autonomous_behavior module
      • rtgym.agent.behavior.base_behavior module
      • rtgym.agent.behavior.behavior module
      • rtgym.agent.behavior.behavior_config module
      • rtgym.agent.behavior.controllable_behavior module
      • Module contents
    • rtgym.agent.sensory package
      • Subpackages
      • Submodules
      • rtgym.agent.sensory.sensory module
      • Module contents
  • Submodules
  • rtgym.agent.agent module
    • Agent
      • Agent.add_sensory()
      • Agent.arena
      • Agent.autonomous
      • Agent.controllable
      • Agent.get_response()
      • Agent.random_spawn()
      • Agent.random_traverse()
      • Agent.sensories
      • Agent.set_behavior()
      • Agent.set_sensory()
      • Agent.set_sensory_manually()
      • Agent.spatial_resolution
      • Agent.spawn()
      • Agent.state
      • Agent.step()
      • Agent.temporal_resolution
  • Module contents
• rtgym.arena package
  • Subpackages
    • rtgym.arena.arena_shapes package
      • Submodules
      • rtgym.arena.arena_shapes.circle module
      • rtgym.arena.arena_shapes.cornered_rectangle module
      • rtgym.arena.arena_shapes.maze_0 module
      • rtgym.arena.arena_shapes.maze_1 module
      • rtgym.arena.arena_shapes.maze_2 module
      • rtgym.arena.arena_shapes.rectangle module
      • rtgym.arena.arena_shapes.trainer_0 module
      • rtgym.arena.arena_shapes.triangle module
      • rtgym.arena.arena_shapes.two_rooms module
      • Module contents
  • Submodules
  • rtgym.arena.arena module
    • Arena
      • Arena.spatial_resolution
      • Arena.dimensions
      • Arena.free_space
      • Arena.subscribers
      • Arena.arena_height
      • Arena.arena_map
      • Arena.arena_width
      • Arena.generate_random_pos()
      • Arena.get_specs()
      • Arena.init_arena_map()
      • Arena.inv_arena_map
      • Arena.load()
      • Arena.map
      • Arena.notify_subscribers()
      • Arena.print_specs()
      • Arena.save()
      • Arena.set_arena_map()
      • Arena.subscribe()
      • Arena.validate_index()
      • Arena.vis()
  • rtgym.arena.room_editor module
  • Module contents
• rtgym.dataclass package
  • Submodules
  • rtgym.dataclass.agent_state module
    • AgentState
      • AgentState.float_coord
      • AgentState.int_coord
  • rtgym.dataclass.raw_agent_state module
    • RawAgentState
      • RawAgentState.batch_size
      • RawAgentState.copy()
      • RawAgentState.float_coord
      • RawAgentState.int_coord
      • RawAgentState.reset()
      • RawAgentState.to_agent_state()
  • rtgym.dataclass.raw_trajectory module
    • RawTrajectory
      • RawTrajectory.append()
      • RawTrajectory.to_trajectory()
      • RawTrajectory.update_state()
  • rtgym.dataclass.trajectory module
    • Trajectory
      • Trajectory.coord
      • Trajectory.copy()
      • Trajectory.float_coord
      • Trajectory.from_dict()
      • Trajectory.int_coord
      • Trajectory.load()
      • Trajectory.n_steps
      • Trajectory.reshape()
      • Trajectory.slice()
      • Trajectory.state_dict()
      • Trajectory.t_range()
  • Module contents
• rtgym.utils package
  • Subpackages
    • rtgym.utils.visualization package
      • Submodules
      • rtgym.utils.visualization.plot_ratemap module
      • Module contents
  • Submodules
  • rtgym.utils.common module
    • hash_seed()
  • rtgym.utils.data_processing module
    • RatemapAggregator
      • RatemapAggregator.get_ratemap()
      • RatemapAggregator.init_counts()
      • RatemapAggregator.reset()
      • RatemapAggregator.update()
    • combine_trajectories()
    • get_gym_dimensions()
    • restrain2ff()
    • states2ff()
    • states_to_ratemap()
  • rtgym.utils.decode_response module
    • Mathematical Background
    • create_dataclass_result()
    • decode_response_euclidean()
    • decode_response_faiss()
    • decode_response_interpolation()
    • decode_response_kdtree()
    • decode_response_torch()
  • rtgym.utils.masking module
    • Masking
      • Masking.apply_gaussian_blur()
      • Masking.gaussian_kernel_1d()
      • Masking.mask()
      • Masking.new_mask()
      • Masking.set_m_max()
      • Masking.set_m_min()
      • Masking.to()
      • Masking.to_tensor()
  • rtgym.utils.verbose module
    • compute_layout()
    • generate_titles()
    • plot_dual_field()
    • plot_dual_fields()
    • plot_single_field()
    • prepare_fields()
    • print_dict()
    • title_single_field()
    • visualize_fields()
  • Module contents

Submodules

rtgym.rtgym module

RatatouGym is a python package that provides a simple interface to generate sensory responses from a virtual agent in a virtual environment.

This file contains more documentation as it serves as the main API for the package.

class RatatouGym(temporal_resolution, spatial_resolution, **kwargs)

Bases: object

The RatatouGym class is a singleton that integrates the environment, agent, behavior, and sensory streams. When initialized, the class creates an arena and an uninitialized agent. The arena is the environment the agent moves in, and the agent is the subject navigating the arena.

After initialization, both the arena and the agent must be set up. For the arena, you can either choose from predefined shapes or provide a custom 2D NumPy array as the map. For the agent, you must configure sensory and behavioral parameters. Sensory is defined by a dictionary of the sensory cell types you want to simulate. Behavior is defined by parameters such as movement speed, how frequently the agent changes direction, and how it avoids obstacles.

Parameters:

• temporal_resolution (float) – Time resolution in milliseconds. E.g. 50 will mean 50ms per timestep.

• spatial_resolution (float) – Spatial resolution in units per pixel. E.g. 1 will mean 1cm per pixel.

• **kwargs – Additional keyword arguments.

Example

>>> from rtgym import RatatouGym
>>> gym = RatatouGym(
...     temporal_resolution=50,
...     spatial_resolution=1,
... )

property arena_map

Get the arena map. The arena map represents the spatial layout of the environment, encoding walls and free space for navigation.

Returns:

Arena map of shape (H, W) where 0 represents free space

and 1 represents walls. This binary representation is used internally for collision detection and path planning.

Return type:

np.ndarray

init_arena_map(**kwargs)

Initialize the arena map with a predefined geometric shape. This method calls the corresponding function in the arena class to initialize the navigation space from a set of predefined shapes.

The method supports various predefined shapes commonly used in spatial navigation research, such as squares, rectangles, circles, and more complex geometries. Each shape type accepts specific parameters that define its dimensions and characteristics. Please refer to the rtgym.arena.arena_shapes module for more details.

Parameters:

**kwargs – Shape-specific parameters that vary depending on the selected arena geometry. Each arean shape might have different parameters. Please refer to the rtgym.arena.arena_shapes module for more details.

Example

>>> gym = RatatouGym(temporal_resolution=50.0, spatial_resolution=1.0)
>>> gym.init_arena_map(shape='rectangle', dimensions=[100, 100])

property inv_arena_map

Get the inverted arena map for visualization. The inverted map is useful for plotting and visualization purposes where free space should be highlighted rather than walls.

Returns:

Inverted arena map of shape (H, W) where 1 represents

free space and 0 represents walls. This format is optimized for visual display and matplotlib visualization.

Return type:

np.ndarray

load_arena(path)

Load an arena configuration from a file. This method allows loading previously saved arena configurations, enabling reproducible experiments and sharing of environment setups.

The loaded arena will replace the current arena configuration entirely, including its spatial layout, dimensions, and any associated metadata. The agent will be automatically notified of the arena change through the internal callback system.

Parameters:

path (str) – Path to the arena file. Should be a valid file path pointing to a previously saved arena configuration. The file format depends on the arena implementation but typically uses .npz format for numpy arrays.

Raises:

• FileNotFoundError – If the specified path does not exist.

• ValueError – If the file format is incompatible or corrupted.

Example

>>> gym = RatatouGym(temporal_resolution=50.0, spatial_resolution=1.0)
>>> gym.load_arena('saved_environments/maze_1.npz')

random_pos(n)

Generate random valid positions in the arena. This method creates random positions that are guaranteed to be in free space (not inside walls), making them suitable for agent placement or target positioning.

Parameters:

n (int) – Number of random positions to generate. Must be a positive integer.

Returns:

Array of shape (n, 2) containing random positions.

Return type:

np.ndarray

property s_res

Get spatial resolution of the gym. The spatial resolution controls how the arena is discretized into pixels. The units are in cm per pixel.

Returns:

Spatial resolution in cm per pixel.

Return type:

float

save_arena(path)

Save the current arena configuration to a file. This method enables persistence of arena configurations for later use, sharing, or experimental reproducibility.

The saved file will contain the complete arena state including the spatial layout, dimensions, and any associated metadata. The saved arena can be later loaded using the load_arena method.

Parameters:

path (str) – Path where the arena file will be saved. Must have .npz extension as the arena is saved in numpy’s compressed format for efficiency and compatibility.

Raises:

• AssertionError – If the file extension is not .npz.

• PermissionError – If the specified path is not writable.

• OSError – If there are issues with file system operations.

Example

>>> gym = RatatouGym(temporal_resolution=50.0, spatial_resolution=1.0)
>>> gym.init_arena_map(shape='square', width=100, height=100)
>>> gym.save_arena('my_environments/custom_maze.npz')

set_arena_map(arena_map)

Set arena map with a custom binary layout. This method allows complete control over the arena geometry by directly specifying the spatial layout as a binary array.

The custom arena map enables creation of arbitrary environments beyond the predefined shapes, allowing for complex mazes, irregular boundaries, or specialized experimental setups. The binary format uses 0 for free space and 1 for walls, following standard conventions in spatial navigation research.

Parameters:

arena_map (np.ndarray) – Custom arena map of shape (H, W) where 0 represents free space (navigable areas) and 1 represents walls (impassable barriers). The array should be 2-dimensional with consistent data types (typically bool or int).

Raises:

• ValueError – If arena_map is not a 2D array or contains invalid values.

• TypeError – If arena_map is not a numpy array or compatible type.

Example

>>> import numpy as np
>>> custom_map = np.zeros((50, 50))  # 50x50 open field
>>> custom_map[0, :] = 1  # top wall
>>> custom_map[-1, :] = 1  # bottom wall
>>> custom_map[:, 0] = 1  # left wall
>>> custom_map[:, -1] = 1  # right wall
>>> gym.set_arena_map(custom_map)

set_sensory_manually(sens_type, sensory)

property t_res

Get temporal resolution of the gym. The temporal resolution controls how continuous time in the real world is discretized into timesteps.

Returns:

Temporal resolution in milliseconds.

Return type:

float

to_coord(pos)

Convert position to coordinate indices. This method transforms continuous position values into discrete coordinate indices based on the spatial resolution of the arena.

Parameters:

pos (array-like) – Position values in the continuous space of the arena.

Returns:

Coordinate indices in the discretized arena grid.

Return type:

np.ndarray

to_sec(ts)

Convert timesteps to seconds. This is a helper function to convert timesteps in milliseconds to seconds.

Parameters:

ts (int, float, np.ndarray, or torch.Tensor) – Timestep value(s).

Returns:

Corresponding time(s) in seconds.

Return type:

int, float, np.ndarray, or torch.Tensor

Raises:

TypeError – If input type is not supported.

Example

>>> gym = RatatouGym(temporal_resolution=50, spatial_resolution=1)
>>> gym.to_sec(20)
>>> 1.0  # 20 timesteps = 1 second (50ms per timestep)
>>> gym.to_sec(np.array([20, 40, 60]))
>>> array([1., 2., 3.])
>>> gym.to_sec(torch.tensor([20, 40, 60]))
>>> tensor([1., 2., 3.])

to_ts(t)

Convert time (in seconds) to number of timesteps. This is a helper function. Note that this method will round the time to the nearest integer.

Parameters:

t (int, float, np.ndarray, or torch.Tensor) – Time value(s) in seconds

Returns:

Corresponding number of timestep(s)

Return type:

int, np.ndarray, or torch.Tensor

Raises:

TypeError – If input type is not supported.

Example

>>> gym = RatatouGym(temporal_resolution=50, spatial_resolution=1)
>>> gym.to_ts(1)
>>> 20  # 1 second = 20 timesteps
>>> gym.to_ts(np.array([1, 2, 3]))
>>> array([20 40 60])
>>> gym.to_ts(torch.tensor([1, 2, 3]))
>>> tensor([20, 40, 60], dtype=torch.int32)

vis_gif(traj, plot_w, plot_h, return_format='anim', stride=10, interval=50, max_frames=100)

Create an animated visualization of trajectory and heading direction.

This method generates an animated plot showing the agent’s path through the arena with real-time visualization of position and heading direction. The trajectory is displayed as a red line that progressively reveals the path taken, while a blue arrow indicates the agent’s current heading direction at each timestep.

The animation can be customized for performance by adjusting frame rate, sampling density, and maximum duration. For long trajectories, the stride and max_frames parameters help balance visualization quality with computational efficiency.

Parameters:

• traj – Trajectory data object containing position coordinates (int_coord, float_coord) and heading direction (hd) information. Must have shape (batch, timesteps, dims).

• plot_w (float) – Desired width of the plot in inches. Will be automatically computed based on arena aspect ratio if not specified correctly.

• plot_h (float) – Desired height of the plot in inches.

• return_format (str) – Output format for the animation. Options: - ‘anim’: Returns matplotlib animation object for further manipulation - ‘html’: Returns HTML representation for Jupyter notebook display

• stride (int) – Step size for sampling trajectory points. Higher values create faster animations by skipping frames but may lose trajectory detail. Default is 10.

• interval (int) – Time between animation frames in milliseconds. Lower values create faster playback. Default is 50ms (20 FPS).

• max_frames (int) – Maximum number of frames to generate regardless of trajectory length. Helps prevent memory issues with very long trajectories and ensures reasonable animation generation times. Default is 100.

Returns:

Animation object that can be saved,

displayed, or converted to different formats.

Return type:

matplotlib.animation.FuncAnimation

Raises:

ValueError – If trajectory data is empty or malformed (no data along expected axis).

Example

>>> # Create animation for a recorded trajectory
>>> anim = gym.vis_gif(trajectory, plot_w=8, plot_h=6,
...                    stride=5, interval=100, max_frames=50)
>>> anim.save('trajectory.gif', writer='pillow', fps=10)

vis_traj(traj, cmap='viridis', linewidth=1, return_format=None, height=3, ax=None, vis_kwargs=None)

Visualize the trajectory of the agent in the arena.

This method provides flexible visualization of agent trajectories, supporting both static plotting and animated displays. The trajectory is overlaid on the arena map, showing the path taken by the agent through the environment. The visualization automatically scales to maintain proper aspect ratios and provides options for customization through various parameters.

For static visualizations, the complete trajectory is displayed as a line plot over the arena. For animated visualizations, the trajectory is built up frame by frame, optionally showing the agent’s heading direction as it moves.

Parameters:

• traj (rtgym.dataclass.Trajectory) – Trajectory object containing the agent’s position data and other trajectory information to be visualized.

• cmap (str, optional) – Colormap for the arena background. Defaults to “viridis”.

• linewidth (float, optional) – Width of the trajectory line. Defaults to 1.

• return_format (str, optional) – Format of the returned visualization. Options are: - ‘anim’: Returns matplotlib animation object - ‘html’: Returns HTML representation of animation for Jupyter notebooks - None: Returns static plot as (fig, ax) tuple

• height (float, optional) – Height of the figure in inches. Width is automatically calculated to maintain arena aspect ratio. Defaults to 3.

• ax (matplotlib.axes.Axes, optional) – Existing axes to plot on. If provided, figure size parameters are ignored.

• vis_kwargs (dict, optional) – Additional keyword arguments passed to animation function when return_format is ‘anim’ or ‘html’.

Returns:

• If return_format is None: Returns (fig, ax) tuple for static plot

• If return_format is ‘anim’: Returns matplotlib FuncAnimation object

• If return_format is ‘html’: Returns HTML object for Jupyter display

Return type:

tuple or animation object

Module contents