rtgym.agent package

Subpackages

Submodules

rtgym.agent.agent module

Agent interface following grid_and_place’s control and neuron API.

class Agent(gym, device=None)[source]

Bases: object

add_neuron_group(neuron_profile)[source]
property arena
get_response(agent_data, return_format='tensor', keys=None, str_filter=None, type_filter=None, device=None)[source]
init_control(control_profile)[source]

Initialize trajectory_generator (EMA) or the original random_walk.

init_neurons(neuron_profiles)[source]

Initialize named neuron groups; missing type selects diffusion_cell.

property neuron_groups
num_neurons(keys=None, str_filter=None, type_filter=None)[source]
random_spawn(batch_size)[source]
random_traverse(duration_ts, batch_size, init_state=None, pause_prob=0, **kwargs)[source]

Generate a tensor trajectory of duration_ts samples, continuing state.

spawn(init_state=None)[source]
property state
step(mv_dir, spd, head_dir)[source]
step_state(state, mv_dir, spd, head_dir)[source]

rtgym.agent.control module

EMA navigation adapted from grasp-lyrl/grid_and_place’s rtgym.

The motion equations and random draw order match its eager implementation. The control device is configurable; no implicit compilation occurs.

class TrajectoryGenerator(gym, bhv_device=None)[source]

Bases: object

Smoothed speed/direction targets with tangential boundary avoidance.

spd_mean and spd_sd are in grid cells/second, as in grid_and_place. Use generate_trajectory for an explicit number of samples.

generate_trajectory(duration_ts, batch_size, init_state=None, init_pos=None)[source]
init_from_profile(profile)[source]
reset()[source]
step(state, mv_dir, spd, head_dir)[source]

Apply commanded grid-cell displacement with the same wall rejection.

rtgym.agent.neurons module

The sensory module contains the Neurons class which is responsible for creating and managing spatially and movement-modulated sensory cells of the agent.

class Neurons(gym)[source]

Bases: object

The class object that manages the sensory system of the agent. More broadly, it handles all simulated neuronal responses of the agent.

When the gym is initialized, an Agent object is automatically created, which in turn creates a Neurons object. The Neurons object is initially just a placeholder and must be initialized with a sensory profile (a dictionary) that defines the simulated neuronal groups and their parameters.

During spatial traversal, RatatouGym separates the concerns of trajectory generation and neuronal response computation. Once a trajectory is generated, RatatouGym calls the get_response method of the Neurons object. This method takes the trajectory as input and computes the corresponding neuronal responses using the defined tuning curves.

This class should not be initialized directly. The RatatouGym class will automatically manage it.

Parameters:

gym (RatatouGym) – Parent RatatouGym object.

add_neuron_group(sensory_profile: Dict[str, Any])[source]

Add a sensory cell to the sensory system.

Parameters:

sensory_profile – Dictionary containing the sensory profile.

aggregate_res_maps(keys=None, str_filter=None, type_filter=None)[source]

Aggregate sensory response maps from spatial modalities.

Combines response maps from multiple spatial sensory modalities into a single array for analysis or decoding purposes.

Parameters:

  • keys (list, optional) – Keys to filter the sensories.

  • str_filter (str, optional) – String filter for sensory names.

  • type_filter (str, optional) – Type filter for sensory modalities.

Returns:

Aggregated sensory response maps of shape (n_cells, H, W).

Return type:

np.ndarray

Raises:

AssertionError – If non-spatial modulated sensory cells are included.

property arena
property common_params
compute_res()[source]
decode_response(response: numpy.ndarray, res_maps=None, keys=None, str_filter=None, type_filter=None, use_torch=True, device=None, method='euclidean', **kwargs)[source]

Decode sensory response into spatial coordinates using various optimization methods.

This method converts high-dimensional sensory responses (e.g., from place cells, grid cells) back to spatial coordinates. Multiple algorithms are available, ranging from exact brute-force search to fast approximate methods.

Parameters:

  • response (np.ndarray) – Neurons response array of shape: - (B, T, D) for trajectory decoding - (B, D) for single state decoding where B=batch size, T=time steps, D=feature dimensions

  • res_maps (np.ndarray, optional) – Precomputed response template maps. If None, computed from filtered sensory modalities.

  • keys (list, optional) – Specific sensory keys to include in decoding.

  • str_filter (str, optional) – String filter for sensory names.

  • type_filter (str, optional) – Type filter for sensory modalities.

  • use_torch (bool, optional) – Enable PyTorch acceleration (default: True).

  • device (str or torch.device, optional) – Computation device for PyTorch.

  • method (str, optional) – Decoding algorithm to use: - “euclidean”: Brute-force exact search (default) - “torch_euclidean”: GPU-accelerated exact search - “kdtree”: K-d tree for fast exact search - “faiss”: FAISS library for very fast approximate search - “interpolation”: Spatial interpolation with anchor points

  • **kwargs – Additional parameters passed to specific methods.

Returns:

Decoded coordinates wrapped in appropriate

dataclass. Shape matches input: (B,T,2) for trajectories, (B,2) for states.

Return type:

Union[Trajectory, AgentState]

Examples

>>> # Decode place cell responses to trajectory
>>> trajectory = sensory.decode_response(responses, method="kdtree")
>>>
>>> # Fast approximate decoding with FAISS
>>> trajectory = sensory.decode_response(responses, method="faiss", n_clusters=50)
filter_neurons(keys=None, str_filter=None, type_filter=None)[source]

This helps to find the keys of the sensory cells that match the given criteria.

It will prioritize the most specific filter. The specificity from most to least is:

keys > str_filter > type_filter

get_response(agent_data, return_format='tensor', keys=None, str_filter=None, type_filter=None, device=None)[source]

Return selected responses as a dictionary, NumPy array, or tensor.

Movement-modulated cells retain their NumPy equations; tensor requests convert those small outputs. Spatial cells query cached fields directly.

init_from_profile(sensory_profile)[source]
list_all()[source]

List all the sensory cells.

load(file_path)[source]

Load the sensory cells from a file.

Parameters:

file_path – Path to the file where the sensory cells are saved.

load_from_state_dict(state_dict, append=True)[source]

Load the sensory cells from a state dictionary.

Parameters:

  • state_dict – State dictionary of the sensory cells.

  • append – If True, append the sensory cells to the existing sensory cells. If False, replace the existing sensory cells.

num_neurons(keys=None, str_filter=None, type_filter=None)[source]
property s_res
save(file_path)[source]

Save the sensory cells to a file.

Parameters:

file_path – Path to the file where the sensory cells will be saved.

property t_res
to(device, dtype=None)[source]

Prepare spatial fields for repeated tensor queries.

Module contents