Wrappers

block

ScheduleAttr

ScheduleAttr(env: TrialEnv, schedule, attr_list)

Bases: TrialWrapper

Schedule attributes.

Parameters:

Name	Type	Description	Default
env	TrialEnv	TrialEnv object	required
schedule			required

Source code in neurogym/wrappers/block.py

def __init__(self, env: TrialEnv, schedule, attr_list) -> None:
    super().__init__(env)
    self.schedule = schedule
    self.attr_list = attr_list

MultiEnvs

MultiEnvs(envs, env_input=False)

Bases: TrialWrapper

Wrap multiple environments.

Parameters:

Name	Type	Description	Default
envs		list of env object	required
env_input		bool, if True, add scalar inputs indicating current environment. default False.	False

Source code in neurogym/wrappers/block.py

def __init__(self, envs, env_input=False) -> None:
    super().__init__(envs[0])
    for env in envs:
        env.unwrapped.set_top(self)
    self.envs = envs
    self.i_env = 0

    self.env_input = env_input
    if env_input:
        env_shape = envs[0].observation_space.shape
        if len(env_shape) > 1:
            msg = f"Env must have 1-D Box shape but got {env_shape}."
            raise ValueError(msg)
        _have_equal_shape(envs)
        self.observation_space = spaces.Box(
            -np.inf,
            np.inf,
            shape=(env_shape[0] + len(self.envs),),
            dtype=self.envs[0].observation_space.dtype,
        )

set_i

set_i(i) -> None

Set the i-th environment.

Source code in neurogym/wrappers/block.py

def set_i(self, i) -> None:
    """Set the i-th environment."""
    self.i_env = i
    self.env = self.envs[self.i_env]

ScheduleEnvs

ScheduleEnvs(envs, schedule, env_input=False)

Bases: TrialWrapper

Schedule environments.

Parameters:

Name	Type	Description	Default
envs		list of env object	required
schedule		utils.scheduler.BaseSchedule object	required
env_input		bool, if True, add scalar inputs indicating current environment. default False.	False

Source code in neurogym/wrappers/block.py

def __init__(self, envs, schedule, env_input=False) -> None:
    super().__init__(envs[0])
    for env in envs:
        env.unwrapped.set_top(self)
    self.envs = envs
    self.schedule = schedule
    self.i_env = self.next_i_env = 0

    self.env_input = env_input
    if env_input:
        env_shape = envs[0].observation_space.shape
        if len(env_shape) > 1:
            msg = f"Env must have 1-D Box shape but got {env_shape}."
            raise ValueError(msg)
        _have_equal_shape(envs)
        self.observation_space = spaces.Box(
            -np.inf,
            np.inf,
            shape=(env_shape[0] + len(self.envs),),
            dtype=np.float32,
        )

reset

reset(**kwargs)

Resets environments.

Reset each environment in self.envs and use the scheduler to select the environment returning the initial observation. This environment is also used to set the current environment self.env.

Source code in neurogym/wrappers/block.py

def reset(self, **kwargs):
    # TODO: kwargs to specify the condition for new_trial
    """Resets environments.

    Reset each environment in self.envs and use the scheduler to select the environment returning
    the initial observation. This environment is also used to set the current environment self.env.
    """
    self.schedule.reset()
    return_i_env = self.schedule()

    # first reset all the env excepted return_i_env
    for i, env in enumerate(self.envs):
        if i == return_i_env:
            continue

        # change the current env so that calling _top.new_trial() in env.reset() will generate a trial for the env
        # being currently reset (and not an env that is not yet reset)
        self.set_i(i)
        # same env used here and in the first call to new_trial()
        self.next_i_env = self.i_env

        env.reset(**kwargs)

    # then reset return_i_env and return the result
    self.set_i(return_i_env)
    self.next_i_env = self.i_env
    return self.env.reset()

set_i

set_i(i) -> None

Set the current environment to the i-th environment in the list envs.

Source code in neurogym/wrappers/block.py

def set_i(self, i) -> None:
    """Set the current environment to the i-th environment in the list envs."""
    self.i_env = i
    self.env = self.envs[self.i_env]
    self.schedule.i = i

TrialHistoryV2

TrialHistoryV2(env: TrialEnv, probs: ndarray | None = None)

Bases: TrialWrapper

Change ground truth probability based on previous outcome.

Parameters:

Name	Type	Description	Default
probs	ndarray | None	matrix of probabilities of the current choice conditioned on the previous. Shape, num-choices x num-choices	None

Source code in neurogym/wrappers/block.py

def __init__(
    self,
    env: TrialEnv,
    probs: np.ndarray | None = None,
) -> None:
    super().__init__(env)
    try:
        self.n_ch = len(self.choices)  # max num of choices
    except AttributeError as e:
        msg = "TrialHistory requires task to have attribute choices."
        raise AttributeError(msg) from e
    if probs is None:
        probs = np.ones((self.n_ch, self.n_ch)) / self.n_ch  # uniform
    self.probs = probs
    if self.probs.shape != (self.n_ch, self.n_ch):
        msg = f"{self.probs.shape=} should be {self.n_ch, self.n_ch=}."
        raise ValueError(msg)
    self.prev_trial = self.rng.choice(self.n_ch)  # random initialization

monitor

Monitor

Monitor(
    env: TrialEnv,
    config: Config | str | Path | None = None,
    name: str | None = None,
    trigger: str = "trial",
    interval: int = 1000,
    verbose: bool = True,
    plot_create: bool = False,
    plot_steps: int = 1000,
    ext: str = "png",
    step_fn: Callable | None = None,
    save_dir: str | Path | None = None,
)

Bases: Wrapper

Monitor class to log, visualize, and evaluate NeuroGym environment behavior.

Wraps a NeuroGym TrialEnv to track actions, rewards, and performance metrics, save them to disk, and optionally generate trial visualizations. Supports logging at trial or step level, with configurable frequency and verbosity.

Parameters:

Name	Type	Description	Default
env	TrialEnv	The NeuroGym environment to wrap.	required
config	Config | str | Path | None	Optional configuration source (Config object, TOML file path, or dictionary).	None
name	str | None	Optional monitor name; defaults to the environment class name.	None
trigger	str	When to save data ("trial" or "step").	'trial'
interval	int	How often to save data, in number of trials or steps.	1000
plot_create	bool	Whether to generate and save visualizations of environment behavior.	False
plot_steps	int	Number of steps to visualize in each plot.	1000
ext	str	Image file extension for saved plots (e.g., "png").	'png'
step_fn	Callable | None	Optional custom step function to override the environment's.	None
verbose	bool	Whether to log information when logging or saving data.	True
level		Logging verbosity level (e.g., "INFO", "DEBUG").	required
log_trigger		When to log progress ("trial" or "step").	required
log_interval		How often to log, in trials or steps.	required

Attributes:

Name	Type	Description
config	Config	Final validated configuration object.
data	dict[str, list]	Collected behavioral data for each completed trial.
data_eval	dict[str, Any]	Evaluation data collected during policy evaluation runs.
cum_reward		Cumulative reward for the current trial.
num_tr		Number of completed trials.
t		Step counter (used when trigger is "step").
save_dir		Directory where data and plots are saved.

Source code in neurogym/wrappers/monitor.py

def __init__(
    self,
    env: TrialEnv,
    config: Config | str | Path | None = None,
    name: str | None = None,
    trigger: str = "trial",
    interval: int = 1000,
    verbose: bool = True,
    plot_create: bool = False,
    plot_steps: int = 1000,
    ext: str = "png",
    step_fn: Callable | None = None,
    save_dir: str | Path | None = None,
) -> None:
    super().__init__(env)
    self.env = env
    self.step_fn = step_fn

    cfg: Config
    # TODO: This is very unwieldy.
    # The respective config classes should just have sane defaults.
    if config is None:
        config_dict = {
            "env": {"name": env.unwrapped.__class__.__name__},
            "monitor": {
                "name": name or "Monitor",
                "trigger": trigger,
                "interval": interval,
                "verbose": verbose,
                "save_dir": save_dir,
                "plot": {
                    "create": plot_create,
                    "step": plot_steps,
                    "title": env.unwrapped.__class__.__name__,
                    "ext": ext,
                },
            },
            "local_dir": LOCAL_DIR,
        }
        cfg = Config.model_validate(config_dict)
    elif isinstance(config, (str, Path)):
        cfg = Config(config_file=config)
    else:
        cfg = config  # type: ignore[arg-type]

    self.config: Config = cfg

    # Assign names for the environment and/or the monitor if they are empty
    if len(self.config.env.name) == 0:
        self.config.env.name = self.env.unwrapped.__class__.__name__
    if len(self.config.monitor.name) == 0:
        self.config.monitor.name = self.__class__.__name__

    # Set the save_dir
    if save_dir is None:
        save_dir = self.config.local_dir / f"{self.config.env.name}/{iso_timestamp()}"
    if self.config.monitor.save_dir is None:
        self.config.monitor.save_dir = save_dir

    # data to save
    self.data: dict[str, list] = {
        "action": [],
        "reward": [],
        "cum_reward": [],
        "performance": [],
    }
    # Evaluation-only data collected during the policy run (not saved to disk)
    self.data_eval: dict[str, Any] = {}
    self.cum_reward = 0.0
    if self.config.monitor.trigger == "step":
        self.t = 0
    self.num_tr = 0

    # Directory for saving plots
    self.save_dir = ensure_dir(self.config.monitor.save_dir)

    # Figures
    if self.config.monitor.plot.create:
        self.stp_counter = 0
        self.ob_mat: list = []
        self.act_mat: list = []
        self.rew_mat: list = []
        self.gt_mat: list = []
        self.perf_mat: list = []

    # Ensure the reset function is called
    initial_ob, *_ = env.reset()
    self.initial_ob = initial_ob

    # Network layers whose parameters are to be monitored
    self.activations: dict[str, ActivationProbe] = {}

reset

reset(seed=None)

Reset the environment.

Parameters:

Name	Type	Description	Default
seed		Random seed for the environment	None

Returns:

Type	Description
	The initial observation from the environment reset

Source code in neurogym/wrappers/monitor.py

def reset(self, seed=None):
    """Reset the environment.

    Args:
        seed: Random seed for the environment

    Returns:
        The initial observation from the environment reset
    """
    self.cum_reward = 0
    return super().reset(seed=seed)

step

step(
    action: Any,
    collect_data: bool = True,
    record: bool = True,
) -> tuple[Any, float, bool, bool, dict[str, Any]]

Execute one environment step.

This method: 1. Takes a step in the environment 2. Collects data if sv_fig is enabled 3. Saves data when a trial completes and saving conditions are met

Parameters:

Name	Type	Description	Default
action	Any	The action to take in the environment	required
collect_data	bool	If True, collect and save data	True
record	bool	A toggle for recording activation traces.	True

Returns:

Type	Description
tuple[Any, float, bool, bool, dict[str, Any]]	Tuple of (observation, reward, terminated, truncated, info)

Source code in neurogym/wrappers/monitor.py

def step(
    self, action: Any, collect_data: bool = True, record: bool = True
) -> tuple[Any, float, bool, bool, dict[str, Any]]:
    """Execute one environment step.

    This method:
    1. Takes a step in the environment
    2. Collects data if sv_fig is enabled
    3. Saves data when a trial completes and saving conditions are met

    Args:
        action: The action to take in the environment
        collect_data: If True, collect and save data
        record: A toggle for recording activation traces.

    Returns:
        Tuple of (observation, reward, terminated, truncated, info)
    """
    if self.step_fn is not None:
        obs, rew, terminated, truncated, info = self.step_fn(action)
    else:
        obs, rew, terminated, truncated, info = self.env.step(action)
    self.cum_reward += rew
    if self.config.monitor.plot.create:
        self.store_data(obs, action, rew, info)
    if self.config.monitor.trigger == "step":
        self.t += 1
    if info.get("new_trial", False):
        self.num_tr += 1
        self.data["action"].append(action)
        self.data["reward"].append(rew)
        self.data["cum_reward"].append(self.cum_reward)
        self.cum_reward = 0
        for key in info:
            if key not in self.data:
                self.data[key] = [info[key]]
            else:
                self.data[key].append(info[key])

        # save data
        save = (
            self.t >= self.config.monitor.interval
            if self.config.monitor.trigger == "step"
            else self.num_tr % self.config.monitor.interval == 0
        )
        if save and collect_data:
            # Create save path with pathlib for cross-platform compatibility
            save_path = self.save_dir / f"trial_{self.num_tr}.npz"
            np.savez(save_path, **self.data)  # type: ignore[arg-type]

            if self.config.monitor.verbose:
                logger.info("--------------------")
                logger.info(f"Data saved to: {save_path}")
                logger.info(f"Number of trials: {self.num_tr}")
                logger.info(f"Average reward: {np.mean(self.data['reward'])}")
                logger.info(f"Average performance: {np.mean(self.data['performance'])}")
                logger.info("--------------------")
            self.reset_data()
            if self.config.monitor.plot.create:
                self.stp_counter = 0
            if self.config.monitor.trigger == "step":
                self.t = 0

        # Start a new trial for all activation monitors.
        if record:
            for am in self.activations.values():
                am.init_new_trial = True
    return obs, rew, terminated, truncated, info

reset_data

reset_data() -> None

Reset all data containers to empty lists.

Source code in neurogym/wrappers/monitor.py

def reset_data(self) -> None:
    """Reset all data containers to empty lists."""
    for key in self.data:
        self.data[key] = []

store_data

store_data(
    obs: Any, action: Any, rew: float, info: dict[str, Any]
) -> None

Store data for visualization figures.

Parameters:

Name	Type	Description	Default
obs	Any	Current observation	required
action	Any	Current action	required
rew	float	Current reward	required
info	dict[str, Any]	Info dictionary from environment	required

Source code in neurogym/wrappers/monitor.py

def store_data(self, obs: Any, action: Any, rew: float, info: dict[str, Any]) -> None:
    """Store data for visualization figures.

    Args:
        obs: Current observation
        action: Current action
        rew: Current reward
        info: Info dictionary from environment
    """
    if self.stp_counter <= self.config.monitor.plot.step:
        self.ob_mat.append(obs)
        self.act_mat.append(action)
        self.rew_mat.append(rew)
        if "gt" in info:
            self.gt_mat.append(info["gt"])
        else:
            self.gt_mat.append(-1)
        if "performance" in info:
            self.perf_mat.append(info["performance"])
        else:
            self.perf_mat.append(-1)
        self.stp_counter += 1
    elif len(self.rew_mat) > 0:
        fname = self.save_dir / f"task_{self.num_tr:06d}.{self.config.monitor.plot.ext}"
        obs_mat = np.array(self.ob_mat)
        act_mat = np.array(self.act_mat)
        visualize_run(
            ob=obs_mat,
            actions=act_mat,
            gt=self.gt_mat,
            rewards=self.rew_mat,
            performance=self.perf_mat,
            fname=fname,
            name=self.config.monitor.plot.title,
            env=self.env,
            initial_ob=self.initial_ob,
        )
        self.ob_mat = []
        self.act_mat = []
        self.rew_mat = []
        self.gt_mat = []
        self.perf_mat = []

evaluate_policy

evaluate_policy(
    num_trials: int = 100,
    model: Any | None = None,
    verbose: bool = True,
) -> dict[str, float | list[float]]

Evaluates the average performance of the RL agent in the environment.

This method runs the given model (or random policy if None) on the environment for a specified number of trials and collects performance metrics.

Parameters:

Name	Type	Description	Default
num_trials	int	Number of trials to run for evaluation	100
model	Any | None	The policy model to evaluate (if None, uses random actions)	None
verbose	bool	If True, prints progress information	True

Returns: dict: Dictionary containing performance metrics: - mean_performance: Average performance (if reported by environment) - mean_reward: Proportion of positive rewards - performances: List of performance values for each trial - rewards: List of rewards for each trial.

Source code in neurogym/wrappers/monitor.py

def evaluate_policy(
    self,
    num_trials: int = 100,
    model: Any | None = None,
    verbose: bool = True,
) -> dict[str, float | list[float]]:
    """Evaluates the average performance of the RL agent in the environment.

    This method runs the given model (or random policy if None) on the
    environment for a specified number of trials and collects performance
    metrics.

    Args:
        num_trials: Number of trials to run for evaluation
        model: The policy model to evaluate (if None, uses random actions)
        verbose: If True, prints progress information
    Returns:
        dict: Dictionary containing performance metrics:
            - mean_performance: Average performance (if reported by environment)
            - mean_reward: Proportion of positive rewards
            - performances: List of performance values for each trial
            - rewards: List of rewards for each trial.
    """
    # Reset environment
    obs, _ = self.env.reset()

    # Initialize hidden states
    states = None
    episode_starts = np.array([True])

    # Tracking variables
    rewards = []
    cum_reward = 0.0
    cum_rewards = []
    performances = []
    self.data_eval = {"action": [], "reward": []}
    # Initialize trial count
    trial_count = 0

    # Run trials
    while trial_count < num_trials + 1:
        if model is not None:
            if _SB3_INSTALLED and isinstance(model.policy, RecurrentActorCriticPolicy):
                action, states = model.predict(
                    obs,
                    state=states,
                    episode_start=episode_starts,
                    deterministic=True,
                )
            else:
                action, _ = model.predict(obs, deterministic=True)
        else:
            action = self.env.action_space.sample()

        # Use collect_data=False to avoid saving evaluation data
        obs, reward, _, _, info = self.step(action, collect_data=False, record=trial_count < num_trials - 1)
        # Update episode_starts after each step
        episode_starts = np.array([False])
        cum_reward += reward

        if info.get("new_trial", False):
            trial_count += 1
            rewards.append(reward)
            cum_rewards.append(cum_reward)
            cum_reward = 0.0
            if "performance" in info:
                performances.append(info["performance"])

            if verbose and trial_count % 1000 == 0:  # FIXME: why is this value hardcoded?
                logger.info(f"Completed {trial_count}/{num_trials} trials")

            self.data_eval["action"].append(action)
            self.data_eval["reward"].append(reward)
            for key in info:
                if key not in self.data_eval:
                    self.data_eval[key] = [info[key]]
                else:
                    self.data_eval[key].append(info[key])

            # Reset states at the end of each trial
            states = None
            episode_starts = np.array([True])

    # Calculate metrics
    performances = performances[:-1]
    performance_array = np.array([p for p in performances if p != -1])
    rewards = rewards[:-1]
    reward_array = np.array(rewards)
    cum_rewards = cum_rewards[:-1]
    cum_reward_array = np.array(cum_rewards)

    # Convert lists to numpy arrays for easier downstream analysis.
    # NOTE:
    # The 'trial' entry is misaligned by 1, so we need to discard
    # the last trial dictionary and the first value from all other elements.
    for key in self.data_eval:
        if key != "trial":
            self.data_eval[key] = np.array(self.data_eval[key][1:])
        else:
            self.data_eval[key] = self.data_eval[key][:-1]

    return {
        "rewards": rewards,
        "mean_reward": (float(np.mean(reward_array > 0)) if len(reward_array) > 0 else 0),
        "cum_rewards": cum_rewards,
        "mean_cum_reward": (float(np.mean(cum_reward_array)) if len(cum_reward_array) > 0 else 0),
        "performances": performances,
        "mean_performance": (float(np.mean(performance_array)) if len(performance_array) > 0 else 0),
    }

plot_training_history

plot_training_history(
    figsize: tuple[int, int] = (12, 6),
    save_fig: bool = True,
    plot_performance: bool = True,
) -> Figure | None

Plot rewards and performance training history from saved data files with one data point per trial.

Parameters:

Name	Type	Description	Default
figsize	tuple[int, int]	Figure size as (width, height) tuple	(12, 6)
save_fig	bool	Whether to save the figure to disk	True
plot_performance	bool	Whether to plot performance in a separate plot	True

Returns: matplotlib figure object

Source code in neurogym/wrappers/monitor.py

def plot_training_history(
    self,
    figsize: tuple[int, int] = (12, 6),
    save_fig: bool = True,
    plot_performance: bool = True,
) -> plt.Figure | None:
    """Plot rewards and performance training history from saved data files with one data point per trial.

    Args:
        figsize: Figure size as (width, height) tuple
        save_fig: Whether to save the figure to disk
        plot_performance: Whether to plot performance in a separate plot
    Returns:
        matplotlib figure object
    """
    files = natsorted(self.save_dir.glob("*.npz"))

    if not files:
        logger.warning("No data files found matching pattern: *.npz")
        return None

    logger.info(f"Found {len(files)} data files")

    # Arrays to hold average values
    avg_rewards_per_file = []
    avg_cum_rewards_per_file = []
    avg_performances_per_file = []
    file_indices = []
    total_trials = 0

    for file in files:
        data = np.load(file, allow_pickle=True)

        if "reward" in data:
            rewards = data["reward"]
            if len(rewards) > 0:
                avg_rewards_per_file.append(np.mean(rewards))
                total_trials += len(rewards)
                file_indices.append(total_trials)

        if "cum_reward" in data:
            cum_rewards = data["cum_reward"]
            if len(cum_rewards) > 0:
                avg_cum_rewards_per_file.append(np.mean(cum_rewards))

        if "performance" in data:
            perfs = data["performance"]
            if len(perfs) > 0:
                avg_performances_per_file.append(np.mean(perfs))

    file_indices = [0, *file_indices]
    avg_rewards_per_file = [0, *avg_rewards_per_file]
    avg_cum_rewards_per_file = [0, *avg_cum_rewards_per_file]
    avg_performances_per_file = [0, *avg_performances_per_file]

    fig, axes = plt.subplots(1, 2 if plot_performance else 1, figsize=figsize)
    if not isinstance(axes, np.ndarray):
        axes = [axes]

    # 1. Rewards and Cumulative Rewards plot
    ax1 = axes[0]

    if len(avg_rewards_per_file) == len(file_indices):
        ax1.plot(
            file_indices,
            avg_rewards_per_file,
            "o-",
            color="blue",
            label="Avg Reward",
            linewidth=2,
        )
    if len(avg_cum_rewards_per_file) == len(file_indices):
        ax1.plot(
            file_indices,
            avg_cum_rewards_per_file,
            "s--",
            color="red",
            label="Avg Cum Reward",
            linewidth=2,
        )

    ax1.set_xlabel("Trials")
    ax1.set_ylabel("Reward / Cumulative Reward")
    common_ylim = (-0.05, 1.05)
    ax1.set_ylim(common_ylim)
    ax1.set_title("Reward and Cumulative Reward per File")

    ax1.grid(True, which="both", axis="y", linestyle="--", alpha=0.7)
    ax1.legend(loc="center right", bbox_to_anchor=(1, 0.2))

    # 2. Optional: Performances plot
    if plot_performance and len(axes) > 1:
        ax2 = axes[1]
        if len(avg_performances_per_file) == len(file_indices):
            ax2.plot(
                file_indices,
                avg_performances_per_file,
                "o-",
                color="green",
                linewidth=2,
            )
        ax2.set_xlabel("Trials")
        ax2.set_ylabel("Average Performance (0-1)")
        ax2.set_ylim(common_ylim)
        ax2.set_title("Average Performance per File")

        ax2.grid(True, which="both", axis="y", linestyle="--", alpha=0.7)
    plt.tight_layout()
    fig.subplots_adjust(top=0.8)

    for ax in fig.axes:
        ax.spines["top"].set_visible(False)
        ax.spines["right"].set_visible(False)

    plt.suptitle(
        f"Training History for {self.config.env.name}\n({len(files)} data files, {total_trials} total trials)",
        fontsize=12,
    )

    if save_fig:
        # TODO: Add an option to provide a custom path.
        # TODO: Add an option to select the file type (PNG, PDF, etc.).
        save_path = self.save_dir / "training_history.pdf"
        plt.savefig(save_path, dpi=300, bbox_inches="tight")
        logger.info(f"Figure saved to {save_path}")

    return fig

record_activations

record_activations(
    layer: Module,
    name: str | None = None,
    steps: int | None = None,
) -> ActivationProbe

Record the output activations of a layer over a trial.

Parameters:

Name	Type	Description	Default
layer	Module	The layer whose activations are being monitored.	required
name	str | None	The name to use for the activation monitor. This can be useful for retrieving activation monitors at a later stage.	None
steps	int | None	The steps to record for. This could be less than the total number of steps in a trial.	None

Returns:

Type	Description
ActivationProbe	An ActivationMonitor instance.

Source code in neurogym/wrappers/monitor.py

def record_activations(
    self,
    layer: nn.Module,
    name: str | None = None,
    steps: int | None = None,
) -> ActivationProbe:
    """Record the output activations of a layer over a trial.

    Args:
        layer: The layer whose activations are being monitored.
        name: The name to use for the activation monitor.
            This can be useful for retrieving activation monitors
            at a later stage.
        steps: The steps to record for.
            This could be less than the total number of steps in a trial.

    Returns:
        An ActivationMonitor instance.
    """
    # Ensure that we have a name for this monitor
    if name is None:
        name = layer.__class__.__name__

    # The total number of steps to record during each trial.
    total_steps = int(self.unwrapped.tmax / self.unwrapped.dt)  # type: ignore[attr-defined]
    steps = total_steps if steps is None else min(steps, total_steps)

    # Create an activation monitor
    am = ActivationProbe(layer, steps, name)

    # Check if we are replacing an existing monitor
    if name in self.activations:
        logger.warning(f"Replacing an existing activation monitor '{name}'.")

    # Store and return the monitor
    self.activations[name] = am

    return am

plot_activations

plot_activations(
    name: str,
    population: str,
    figsize: tuple[int, ...] | None = None,
    neurons: int | list[int] | None = None,
    mean: bool = False,
    batch_dim: int | None = None,
) -> tuple[Figure, Axes]

Plot the neuron activations.

Parameters:

Name	Type	Description	Default
name	str	Name of the layer.	required
population	str	The neuron population to plot.	required
figsize	tuple[int, ...] | None	The size of the figure.	None
neurons	int | list[int] | None	List of neuron ids to plot. If None, all neurons will be plotted.	None
mean	bool	If set, plot the mean activation over all trials rather than each separate trial.	False
batch_dim	int | None	The batch dimension, if it exists.	None

Raises:

Type	Description
ValueError	Raised if there is no such layer in the history.
KeyError	Raised if activations have not been recorded for the requested neuron population.
ValueError	Raised if the requested neuron IDs are outside the layer range.

Returns:

Type	Description
tuple[Figure, Axes]	A plot of all neuron activations.

Source code in neurogym/wrappers/monitor.py

def plot_activations(
    self,
    name: str,
    population: str,
    figsize: tuple[int, ...] | None = None,
    neurons: int | list[int] | None = None,
    mean: bool = False,
    batch_dim: int | None = None,
) -> tuple[plt.Figure, plt.Axes]:
    """Plot the neuron activations.

    Args:
        name: Name of the layer.
        population: The neuron population to plot.
        figsize: The size of the figure.
        neurons: List of neuron ids to plot. If None, all neurons will be plotted.
        mean: If set, plot the mean activation over all trials rather than each separate trial.
        batch_dim: The batch dimension, if it exists.

    Raises:
        ValueError: Raised if there is no such layer in the history.
        KeyError: Raised if activations have not been recorded for the requested neuron population.
        ValueError: Raised if the requested neuron IDs are outside the layer range.

    Returns:
        A plot of all neuron activations.
    """
    # First, get the right layer
    am = self.activations.get(name)
    if am is None:
        msg = f"Invalid layer name '{name}'"
        raise ValueError(msg)

    # Get the activations for the specified population
    activations = am.history.get(population)
    if activations is None:
        msg = f"Invalid population '{population}' (available populations: {am.probe.populations})"
        raise KeyError(msg)

    if len(activations) == 0:
        msg = f"The history for population '{population}' is empty."
        raise KeyError(msg)

    lengths = [len(a) for a in activations]
    masked = np.ma.masked_all((len(activations), max(lengths), *activations[0][0].shape), dtype=np.float32)
    for trace_idx, trace in enumerate(activations):
        masked[trace_idx, : len(trace)] = np.array(trace)

    # Neuron selection
    if neurons is None:
        neurons = list(range(activations[0][0].shape[-1]))
    elif isinstance(neurons, int):
        neurons = [neurons]

    if max(neurons) > masked[0].shape[-1] - 1:
        msg = f"The neuron IDs must be between 0 and {activations[0][0].shape[-1] - 1}."
        raise ValueError(msg)
    neuron_count = len(neurons)

    # Compute the number of rows and columns
    rows = int(np.floor(np.sqrt(neuron_count)))
    cols = neuron_count // rows
    if neuron_count > rows * cols:
        rows += 1

    # Figure and axes
    if figsize is None:
        figsize = (3 * cols, 2 * rows)
    fig, axes = plt.subplots(rows, cols, figsize=figsize)
    if neuron_count == 1:
        axes = [axes]

    # Plot the activations for the selected neurons
    for neuron_idx, (row, col) in enumerate(product(range(rows), range(cols))):
        ax = axes[col] if rows == 1 else axes[row][col]
        if neuron_idx >= neuron_count:
            ax.axis("off")
            continue
        neuron = neurons[neuron_idx]
        self._plot_activations_single(neuron, ax, mean, masked, batch_dim)

    fig.suptitle(f"{population.capitalize()} neuron activations", fontsize=18, y=1)
    fig.tight_layout()

    return fig, axes

noise

Noise

Noise(env: TrialEnv, std_noise: float = 0.1)

Bases: Wrapper

Add Gaussian noise to the observations.

Parameters:

Name	Type	Description	Default
env	TrialEnv	The NeuroGym environment to wrap.	required
std_noise	float	Standard deviation of noise. (def: 0.1)	0.1
perf_th		If != None, the wrapper will adjust the noise so the mean performance is not larger than perf_th. (def: None, float)	required
w		Window used to compute the mean performance. (def: 100, int)	required
step_noise		Step used to increment/decrease std. (def: 0.001, float)	required

Source code in neurogym/wrappers/noise.py

def __init__(
    self,
    env: TrialEnv,
    std_noise: float = 0.1,
) -> None:
    super().__init__(env)
    self.env = env
    self.std_noise = std_noise

pass_action

PassAction

PassAction(env: TrialEnv)

Bases: Wrapper

Modifies observation by adding the previous action.

Source code in neurogym/wrappers/pass_action.py

def __init__(self, env: TrialEnv) -> None:
    super().__init__(env)
    self.env = env
    if isinstance(env.observation_space, spaces.Discrete):
        env_oss = int(env.observation_space.n)  # Number of discrete states
        self.observation_space = spaces.Discrete(n=env_oss + 1)
    elif env.observation_space.shape is not None:
        env_oss = env.observation_space.shape[0]
        self.observation_space = spaces.Box(
            -np.inf,
            np.inf,
            shape=(env_oss + 1,),
            dtype=np.float32,
        )
    else:
        msg = "env.observation_space.shape for PassAction may not be None"
        raise ValueError(msg)

pass_reward

PassReward

PassReward(env: TrialEnv)

Bases: Wrapper

Modifies observation by adding the previous reward.

Source code in neurogym/wrappers/pass_reward.py

def __init__(self, env: TrialEnv) -> None:
    """Modifies observation by adding the previous reward."""
    super().__init__(env)
    if isinstance(env.observation_space, spaces.Discrete):
        env_oss = int(env.observation_space.n)  # Number of discrete states
        self.observation_space = spaces.Discrete(n=env_oss + 1)
    elif env.observation_space.shape is not None:
        env_oss = env.observation_space.shape[0]
        self.observation_space = spaces.Box(
            -np.inf,
            np.inf,
            shape=(env_oss + 1,),
            dtype=np.float32,
        )
    else:
        msg = "env.observation_space.shape for PassReward may not be None"
        raise ValueError(msg)

reaction_time

ReactionTime

ReactionTime(
    env: TrialEnv,
    urgency: float = 0.0,
    end_on_stimulus: bool = True,
)

Bases: Wrapper

Allow reaction time response.

Modifies a given environment by allowing the network to act at any time after the fixation period. By default, the trial ends when the stimulus period ends. Optionally, the original trial structure can be preserved while still allowing early responses.

Parameters:

Name	Type	Description	Default
env	TrialEnv	The environment to wrap	required
urgency	float	Urgency signal added to reward at each timestep	0.0
end_on_stimulus	bool	If True (default), trial ends when stimulus ends. If False, preserves original trial timing while allowing early responses during stimulus period.	True

Source code in neurogym/wrappers/reaction_time.py

def __init__(
    self,
    env: TrialEnv,
    urgency: float = 0.0,
    end_on_stimulus: bool = True,
) -> None:
    super().__init__(env)
    self.urgency = urgency
    self.end_on_stimulus = end_on_stimulus
    self.tr_dur = 0

side_bias

SideBias

SideBias(
    env: TrialEnv,
    probs: list[list[float]],
    block_dur: float | tuple[int, int] = 200,
)

Bases: TrialWrapper

Changes the probability of ground truth with block-wise biases.

Parameters:

Name	Type	Description	Default
env	TrialEnv	The task environment to wrap (must expose choices).	required
probs	list[list[float]]	Explicit probability matrix with shape (n_blocks, n_choices). Each row defines the choice probabilities for one block and must sum to 1.0. The number of columns (n_choices) must match the number of choices defined in the task (i.e., len(env.choices)), so that each probability maps to a valid choice.	required
block_dur	float | tuple[int, int]	Duration of each block, with behavior depending on the type: - int (≥ 1): Use a fixed number of trials per block (e.g., block_dur=20 means each block has exactly 20 trials). - float (0 < value < 1): Specify a per-trial probability of switching to a new block (e.g., block_dur=0.1 means there's a 10% chance of switching blocks after each trial). - tuple (low, high): Draw the number of trials per block randomly from a uniform distribution over the integer range [low, high] (inclusive).	200

Examples:

• probs=[[0.8, 0.2], [0.2, 0.8], [0.4, 0.6]], block_dur=200

  Stay 200 trials per block, randomly switch to new block after;

• probs=[[0.8, 0.2], [0.2, 0.8], [0.4, 0.6]], block_dur=0.1

  10% probability per trial to switch to new random block;

• probs=[[0.8, 0.2], [0.2, 0.8], [0.4, 0.6]], block_dur=(200, 400)

  Random trials in [200, 400] range per block, then switch.

Source code in neurogym/wrappers/side_bias.py

def __init__(
    self,
    env: TrialEnv,
    probs: list[list[float]],
    block_dur: float | tuple[int, int] = 200,
) -> None:
    super().__init__(env)

    # Validate environment
    if not isinstance(self.task, TrialEnv):
        msg = "SideBias requires the wrapped task to be a TrialEnv."
        raise TypeError(msg)

    try:
        self.choices = self.task.choices  # type: ignore[attr-defined]
    except AttributeError as e:
        msg = "SideBias requires task to have attribute choices."
        raise AttributeError(msg) from e

    # Reject non-matrix types (no automatic matrix generation)
    if (
        not isinstance(probs, list)
        or not all(isinstance(row, list) for row in probs)
        or not all(isinstance(prob, (float, int)) for row in probs for prob in row)
    ):
        msg = (
            "probs must be a 2D list of lists (matrix) with shape (n_blocks, n_choices),"
            "e.g., probs = [[0.5, 0.5], [0.2, 0.8], [0.8, 0.2]] for n_blocks = 3 and n_choices = 2."
        )
        raise TypeError(msg)

    # Convert to numpy array and validate
    self.choice_prob = np.array(probs, dtype=float)

    # Validate matrix dimensions
    if self.choice_prob.ndim != 2:
        msg = f"probs must be a 2D matrix, got {self.choice_prob.ndim}D array."
        raise ValueError(msg)

    n_blocks, n_choices = self.choice_prob.shape

    if n_choices != len(self.choices):
        msg = (
            f"The number of choices inferred from the probability matrix ({n_choices}) "
            f"does not match the number of choices defined in the task ({len(self.choices)}). "
            "These must be equal to ensure each probability corresponds to a valid choice."
        )
        raise ValueError(msg)

    # Validate that each row sums to 1 (with tolerance for floating point)
    row_sums = np.sum(self.choice_prob, axis=1)
    tolerance = 1e-10
    if not np.allclose(row_sums, 1.0, atol=tolerance):
        invalid_rows = np.where(~np.isclose(row_sums, 1.0, atol=tolerance))[0]
        msg = (
            f"Each row in probs must sum to 1.0. "
            f"Rows {invalid_rows.tolist()} have sums {row_sums[invalid_rows].tolist()}"
        )
        raise ValueError(msg)

    # Validate probabilities are non-negative
    if np.any(self.choice_prob < 0):
        msg = "All probabilities in probs must be non-negative."
        raise ValueError(msg)

    self.n_block = n_blocks
    self.curr_block = self.task.rng.choice(self.n_block)

    # Validate and set block_dur
    self._validate_and_set_block_dur(block_dur)

    # Initialize block state
    self._remaining_trials, self._p_switch = self._new_block_duration()

new_trial

new_trial(**kwargs)

Generate new trial with block-based probability biases.

Source code in neurogym/wrappers/side_bias.py

def new_trial(self, **kwargs):
    """Generate new trial with block-based probability biases."""
    # Determine if we should switch blocks
    if self._p_switch is not None:
        # Probability-based switching
        switch_block = self.task.rng.random() < self._p_switch
    else:
        # Counter-based switching
        self._remaining_trials -= 1
        switch_block = self._remaining_trials <= 0

    if switch_block:
        # Switch to a different random block (never same as current)
        if self.n_block > 1:
            # Multiple blocks available - choose different one
            available_blocks = [b for b in range(self.n_block) if b != self.curr_block]
            self.curr_block = self.task.rng.choice(available_blocks)
        # If only 1 block, stay in same block (edge case)

        # Generate new duration/probability for the new block
        self._remaining_trials, self._p_switch = self._new_block_duration()

    # Set trial parameters based on current block
    probs = self.choice_prob[self.curr_block]
    kwargs["ground_truth"] = self.task.rng.choice(self.choices, p=probs)
    kwargs["probs"] = probs

    return self.env.new_trial(**kwargs)