allenneuraldynamics / aind-behavior-curriculum

Starter repository for behavior base primitives.

aind-behavior-curriculum

A core problem in mice training is accurately keeping track of each mouse's training stage and accurately setting the corresponding rig parameters. As the number of behavior studies, research assistants, and mice increase, manual tracking and parameter input is prone to human error. This library provides a flexible framework for defining mice curriculum enabling mouse training to be automated.

Installation

pip install aind-behavior-curriculum

Documentation

Understanding a Curriculum

A Curriculum is structured as a graph of training Stages. Each Stage is associated with a Task, which is a set of rig parameters. Stages are connected by Stage Transitions, which are directed edges associated with a trigger condition.

Stages and Stage Transitions form the nodes and edges of a Curriculum graph, respectively. With this structure alone, a user can define a basic curriculum with the flexibility of defining skip connections and regressions. For nodes with multiple ongoing edges, edges are labelled by priority, set by the user.

An example curriculum consisting of purely stages and stage transitions. This Curriculum consists of a skip connection between Stage 'StageA' and Stage 'Graduated'. Stage Transitions are triggered on a parameter 't2' and the skip transition is ordered before the transition going to Stage StageB.

$~$

This library also supports Curriculum hypergraphs.

Conceptually, a user may want to change the rig parameters associated with a stage, but this set of rig parameters would be unnatural to classify as a new training stage altogether. In this situation, the user may define a graph of Policies and Policy Transitions within a Stage. A Policy, changes the task parameters of a Stage, as described above. A Policy Transition acts just like a Stage Transition, and defines transitions between Policies on a trigger condition. Like Stage Transitions, Policy Transitions can connect any two arbitrary Policies and are ordered by priority set by the user.

An example Curriculum consisting of Stage and Policy graphs. Left: The high level policy graph. Right: Internal policy graphs.

Policies are more nuanced than Stages.

Yellow Policies in the example indicate 'Start Policies'. To initialize the rig parameters of a Stage, the user must specify which Policy/Policies in the Stage policy graph to start with.

Unlike Stages, a mouse can occupy multiple active Policies within a Stage. As described later, the Trainer will record the net combination of rig parameters.

$~$

Any hypergraph is supported!

Here are some examples of the possibilities. The high-level stage graph are shown to the left and the inidividual policy graphs are shown to the right.

A 'Tree' Curriculum

A 'Train Track' Curriculum

A 'Policy Triangle' Curriculum

A 'Stage Triangle' Curriculum

$~$

Understanding the Trainer

The Trainer is responsible for recording where a mouse is in its associated curriculum hypergraph. The Trainer contains 4 primary functions:

1. Registration: This is the entry point where the mice enter the system. Here, the user provides the Trainer with a mouse and associates the mouse with a curriculum, a start stage, and start policies as a starting place for evaluation.

2. Evaluation: For each registered mouse, the Trainer looks at the mouse's current position in its hypergraph curriculum. The Trainer collects all the current outgoing transitions and checks which evaluate to True. The Trainer determines the updated hypergraph position and associated Task parameters according to the following simple rules:

   • Trainer takes the outgoing Stage Transition with the highest priority. If multiple Stage Transitions evaluate to True, then the Stage Transition with the highest priority is chosen. Priority is set by the user.
   • Trainer takes the outgoing Policy Transition with the highest priority. If multiple Policy Transitions evaluate to True, then the Policy Transition with the highest priority is chosen. Priority is set by the user.
   • Stage Transitions override Policy Transitions. If a Stage Transition and Policy Transition both evaluate to True, the Trainer jumps directly to the next Stage .
   • If no transitions are True, the mouse stays in place.
   • For multiple active Policies that evaluate to True, Trainer sets the current Task parameters to the net combination of incident Policies.

3. Mouse Override: This allows the user to update a mouse's position manually to any position in its curriculum. Future evaluation occurs from this new position. Due to this feature, it is possible to design a Curriculum of 'floating stages' and 'floating policies'.

4. Mouse Eject: This allows the user to remove a mouse from its curriculum entirely. The position of the mouse is recorded as 'None' and stays at 'None' on future evaluation unless the mouse is overrides back onto curriculum.

Every Trainer function keeps a record of mouse history in SubjectHistory which can be referenced or exported for rig automation and further analysis.

$~$

Building a Curriculum

For examples of how to build a Curriculum, please reference examples/example_project and examples/example_project_2 within the project files and their associated diagrams, examples/example_project/diagrams and examples/example_project_2/diagrams.

Tips for building your own Curriculum:

• Focus on one graph at a time. Define all the Tasks/Stages/Stage Transitions associated with the higher level graph, and then move onto defining the Policies/Policy Transitions associated with each Stage.

• Metrics contains all the variables that trigger conditions associated with Stage Transitions and Policy Transitions. Progressively add to Metrics as needed.

• Keep Stage Transitions and Policy Transitions simple. A typical transition will only trigger on one metric variable. This makes transitions much easier to name.

• Validate Stage Transition and Policy Transition priority with the Curriculum.export_digram(...) utility, which labels edges with its rank. Use Curriculum.set_stage_transition_priority(...) and Stage.set_policy_transition_priority(...) to reorder priority.

Common mistakes:

• Every Stage needs a set of start policies, see Curriculum.set_start_policies(...). If a stage with no policies is desired, use curriculum_utils.create_empty_stage(...). This is a common pattern for the final stage of a Curriculum, so the library also offers a prebuilt final stage curriculum_utils.GRADUATED.

• The callables in Policy and Policy Transition/Stage Transition have different input signatures. Please reference Policy.validate_rule(...) and PolicyTransition.validate_rule(...)/StageTransition.validate_rule(...)

$~$

Building a Trainer

The 4 primary functions of the Trainer described above are decoupled from any database. To use the Trainer in practice, the user must define Trainer.load_data(...) and Trainer.write_data(...) which connect to a user's databases for mice curriculum, mice history, and mice metrics. Please see examples/example_project/trainer.py for an example.

$~$

Inside Allen Institute of Neural Dynamics

Allen Institute of Neural Dynamics offers an internal repository template that automatically uploads the repository's curriculum to a central bucket available here: https://github.com/AllenNeuralDynamics/aind-behavior-curriculum-template This way, curriculums can be accessed across rig computers and reused/modified similar to Github commits.

As of (5/9/2024), a Metrics database has yet to be defined, therefore a Trainer cannot be defined.

aind-behavior-curriculum's Issues

Description of a full ecossystem pipeline deployed in AIND

Bootstrapping an animal the first time:

1. Register the Animal in the Curriculum (Trainer.register_subject)
2. A first session will be suggested (by updating Task.TaskParameters via Trainer.evaluate_subjects)
   A suggestion will be generated by updating the current settings (see _get_net_parameter_update function). We might need to wrap this function a bit differently.
3. Write this data to SLIMS (Mouse Content -> Behavior Session Event:
   a. Attach the json file that deserializes from Task object)
   b. Fill in the other metadata fields associated with the Event (eg. experimenter, scheduled date, etc...)

Running a session:

1. Find a task to run
   a. From a Rig computer list, query SLIMS for Behavior Session Event for a given mouse:
   b. Filter/query the list as you see fit (e.g. "Is there a session scheduled for today for user X?")
   c. Grab the latest JSON attached to it.
   e. This Task object will be used to run the task (the user can override the full step 1, and I will refer to it as 1*)
   (1*). The user ignores SLIMS and instead creates their own Task object via other means.
2. Run a task and collect data
   a. Deserialize the JSON from (2) back to a Task object.
   b. Run it
3. Calculate metrics on the rig and save session.json
   a. If needed, query for historical sessions from the document DB. This will be a list of session dictionaries, which contain historical computed metrics and task parameters.
   b. Generate a session.json (aind-data-schema) and add Task and Metrics to the fields StimulusEpoch.script.parameters, StimulusEpoch.output_parameters, respectively.
   c. TODO: From the Trainer class, deserialize the currently active Stage and Policies to JSON. This output should go somewhere near the Metrics class as it is the suggested output of the Trainer. (see response below)
4. Upload data
   a. copy data to VAST (prefer to use aind-watchdog-service)
   b. upload data to cloud (taken care of if using aind-watchdog-service, otherwise use aind-data-transfer-service)
5. Save recommendation for the next session
   a. upload to SLIMS as Behavior Content Event
6. Go back to 1. Repeat until mouse graduates.

Manual intervention edge-cases:

User overrides suggestion

This results from taking the 1* approach. There is no really good way to guarantee that the mouse can go back to the curriculum automatically as the pipeline loses "control" over what task and parameters ran at the rig. This situation should be handled by an application-specific routine and left to the criteria of the developer

User manually moves the animal to a different stage.

This will reset the curriculum to a stage akin to "Bootstrapping an animal the first time 2". A first session will be suggested from no prior Metrics and the curriculum can continue.

User manually removes animal from curriculum and resumes later

Worst case scenario, this defaults back to "### User manually moves animal to a different stage.". Best case, the user resumes the animal in the same stage. If this is the case, the pipeline resumes from "Running a session. 5."

Export requires graphviz

For some operating systems graphviz may not already be installed. It might be helpful for future users to have a note that they need that in the README as the error raised is informative but doesnt explicitly say "graphviz" anywhere in it and may confuse future users.

Clarify distinction between `Stage` and `Policy`

Policy and Stage remain one of the most subtle concepts in the present library. We should think about a better way to distinguish them. Here's I am currently seeing these two concepts/implementations, feel free to use this thread to discuss and hopefully spin a pull-request to the docs.

On the curriculum

As stage on the current docs:

A Curriculum is structured as a graph of training Stage. Each Stage is associated with a Task, which defines a set of configuration parameters via TaskParameters. Stages are connected by StageTransition, which are directed edges associated with a trigger condition.

In other words, a Curriculum, in its simplest form, can be seen as a container of Stages + any logic associated to the transition between them (StageTransitions).

Stage is, in turn, a "container" of a Task instance. The concept of an "instance" is critical here. While the Task object defines a class, Stage works by wrapping an instance. What this means is that two distinct Stages A and B, could actually implement the same Task type BUT different instances of the same Task (Think about two different training stages of the same task implementation).

Let's step back a bit here and think about what affordances we have at this point.

1 - In each Stage we can define a set of parameters for a specific Task
2 - We can express transition logic that transitions between Stages
3 - We can have Stages of different Tasks or the same Task

Lets look at two ways this could be used (I am going to use pseudo-code for brevity)

Curriculum within the same Task

Let's consider a Task named VrForaging. The experimenter has 2 distinct and discrete operation modes that the animal transitions through training. Let's call them: ForagingForApples and ForageForCheese. The animal starts training on ForagingForApples and after it fulfills a transition criteria (say collects 100 rewards in a behavior session), graduates to ForageForCheese.

Curriculum across Tasks

Let's consider a curriculum where the animal is first trained in the previous task VrForaging and Stage ForageApples. Now, instead of graduating to ForageForcheese, we want to teach the mouse a completely new task, say DynamicForaging. This new task also has stages, say BaitedStage.

As you can see, a Stage is always required. It is the materialization of a Task class/type.

But what about policies?

Consider the previous scenario where the animal is in Task VrForaging and Stage ForagingForApples. Additionally, consider a variable inside the task that depicts the amount of reward the animal receives. Say the experimenter wants to automatically update this value based on the amount of water the animal drank in the previous session.

We can express this logic using Stages by considering:

ForagingForApples(5ul) -> ForagingForApples(6ul) if WaterDrank > 1000 & WaterDrank< 1200;
ForagingForApples(5ul) -> ForagingForApples(7ul) if WaterDrank > 800  & WaterDrank < 1000;
ForagingForApples(5ul) -> ForagingForApples(8ul) if WaterDrank > 600  & WaterDrank < 800;
ForagingForApples(5ul) -> ForagingForApples(9ul) if WaterDrank > 400  & WaterDrank < 600;

While possible this is unnecessarily annoying as we would need to code a large number of stage transitions between all the stages.

To meet this need, we can instead use a Policy. A Policy can be thought of as a set of functions that run on top of the session outcome. For instance, the above scenario could be instead coded in the continuous domain by:

ForagingForApples(x), where x = WaterDrank / 2000

In this example, only one policy is active. However, multiple policies can be active simultaneously that are used to update different set of parameters.

In other words, a Stage, on top of being a simple container for a Task instance, also defines a set of policies. As a result, two stages can define the same Task, but differ solely on the underlying policies. (For instance, one Stage could update the WaterDrank by a factor of 1/2000 whereas the other may do it by a different factor of 1/4000)

Why Policy transitions?

The previous architecture already affords vast flexibility. However, one thing that becomes very difficult is to independently control active policies. Imagine that you have two concurrently active policies (UpdateWater and UpdateDistance). At some point, you want to change the UpdateWater to UpdateWaterByAlot. This could be done by coding an extra stage, as mentioned before. Unfortunately, we need to account for UpdateDistance too!
Couldn't we just add this Policy to the new stage? You could, but what happens if UpdateDistance has a corresponding UpdateDistanceByAlot too? Now we suddenly need to expand our number of stages to account for all possible pair-wise combinations of active policies.

To solve this, each Stage can instead also code its own set of PolicyTransitions that define how different policies transition into others. This allows different policies to be concurrently active within the Stage while simultaneously allow them to be updated to different ones independently.

Consider explicitly making classes abstract

As discussed in #32 and #29 I think we should consider making most of the classes to be inherited by users explicitly abstract (i.e. inheriting from the std ABC module).

These would include, at least:
Task, TaskParameters, Metrics, Curriculum, PolicyGraph, StageGraph

Consider adding `stage_name` descriptor to `Task` object

Is your feature request related to a problem? Please describe.
From the point of view of the Task schema, no field is available to keep track of stage or experiment type. This would be interesting as the same Task could be used, via different parameters, to instantiate different experiments (or stages).

As I was thinking about this, it felt like violating a bit the separation of concerns principle but I think it may be warranted for a few reasons:

1. If a user wants to buy into the Task specification but does not want to use the full Curriculum (Curriculum actually deserializes the stage information but only with the full curriculum, not the suggestion)
2. For instances o Task kept in a remote database (e.g. slims) it would be nice to have a human readable field that indicated with stage or experiment it is being suggested by the Task json instance.

Describe the solution you'd like

Add stage_name to Task with the following signature:

`stage_name`: Optional[str] = Field(default=None)

The following should be added:

• When a Stage is instantiated if no name is passed to Stage (i.e. = None), a validator should attempt to use Task.stage_name automatically. If stage_name is null, it should throw, since name is a required field.

Describe alternatives you've considered

Leave it up to the user to add this field to TaskParameters. In which case it would not be the concern of the curriculum. But it would also be much harder to generate a common API that uses this field

I am a bit afraid of the possible duplication of metadata between Stage.name and this suggested property. If that is a problem, we may just want to change the name of this property for something else to prevent confusion. Perhaps task_descriptor?

Additional context
This field should only be used for generating human-readable metadata, I would not rely on it for any functional implication in the curriculum. Instead we should rely on Stage.name

Typo in `Trainer` description

It should mention Stage instead of Trainer I believe?

aind-behavior-curriculum/docs/source/index.rst

Line 128 in bb28161

- :py:class:`~aind_behavior_curriculum.trainer.Trainer` takes the outgoing :py:class:`~aind_behavior_curriculum.trainer.Trainer` with

Add curriculum visualization

Implement Task schema

• add a basic Task schema, based on various examples
• add an example Task in ./examples/
• add a unit test that constructs the example Task

@bruno-f-cruz @hanhou please add links to your repos here.

update readme

Improve docs cross-referencing

We should consider refactoring the docs website to cross-reference with the API documentation of the package. For instance, Metrics should link to the respective technical reference instead of just highlighting the word.
Additionally, we should consider splitting the technical documentation by module instead of flattening the whole thing. As an example:

The callables in **Policy** and ...

should be replaced by:

The callables in :py:class:`~aind_behavior_curriculum.curriculum.Policy` and...

Improve `Exception` patterns

There is a lot of assert patterns in the code. While ok for debugging, as we move to production, we should favor raise Exception over assert patterns.

Review type hinting for correct json schema generation

https://github.com/AllenNeuralDynamics/aind-behavior-curriculum/blob/bb28161608b4544e5f71e23c2190a6682f17dcaf/src/aind_behavior_curriculum/trainer.py#L20C26-L20C38

I dont think this is generating a valid json schema. Mostly a reminder for me to check later.

Refactor graph data model

Enable version autobump on merge to main via github action

Use the same action as aind-data-schema.

Fix Union[Stage] signature in inherited curriculum model

aind-behavior-curriculum/examples/example_project/curriculum.py

Line 163 in 9d77be9

Curriculum_Stages = Union[StageA, StageB, Graduated]

Worst case scenario this should be refactored to be a discriminated union (see more details here AllenNeuralDynamics/aind-data-schema#518).

Best case, see if there is a way to not have to pass the types.

Add "script" string field to `Stage`

Implement the dynamic foraging task as an example

• subclass the abstract Task schema
• save a file in examples/dynamic_foraging.py

Creating parameters in a simple way

Just going through this with an eye for simplicity.

aind-behavior-curriculum/examples/example_project/curriculum.py

Line 169 in bb28161

taskA = TaskA(task_parameters=TaskAParameters())

I appreciate it is important to validate parameters values against a schema but could external user just pass a dictionary instead of subclassing?

curriculum is misspelled in repository name

Describe the bug
curriculum, not cirriculum

publish to pypi

There are instructions for how to do this in the tag_and_publish tag_and_publish.yml.

Also need to add our service account to this repo so it has permission.

Empty classes could be creating unnecessary layers

I am going through the repo with an eye for simplicity.
why do we need this empty class?

See here :
https://github.com/AllenNeuralDynamics/aind-behavior-curriculum/blob/bb28161608b4544e5f71e23c2190a6682f17dcaf/src/aind_behavior_curriculum/task.py#L19C22-L19C44

Consider adding `version` field to `Task` class

Is your feature request related to a problem? Please describe.
After playing a bit with the current specification, I believe we need a field in Task that version of the task. This version property should follow semantic versioning. This is especially important if considering a case where users might want to validate different versions of the task to a remote trainer on deserialization of the Task object.

Describe the solution you'd like

add the following property to Task:

version: str = Field(..., description="Semantic versioned task version")

There are a few ways to ensure validation, the easiest would likely be:

class MyTask(Task):
   (...)
   version: Literal["0.0.1"] = "0.0.1"

As the field is constant, any json object created with a previous version would end up raising a validation error on deserialization.

Describe alternatives you've considered
We should also decide if this is supposed to be an optional or required field

Why do we need Stage and Task objects?

aind-behavior-curriculum/examples/example_project/curriculum.py

Line 171 in bb28161

stageA = Stage(name="StageA", task=taskA)

It looks like the parameters are directly pass to the a stage object. Could we just have a Stage object? What is the usefulness of these 2 layers?

Implement Curriculum schema

• Add the generic Curriculum schema
• Add an example that uses the basic Task from #2

fix repo name typo

• aind-behavior-cirriculum -> aind-behavior-curriculum
• fix setup.py/pyproject.toml/subdirectories, etc

Architectural class diagram

This repository seems to have a complex class architecture. I went through many of the classes and I still could not understand all dependencies and the logic behind it. They are a few empty classes with empty layers.
I think this might warrant a review of the class architecture given that it is expected from any external users to subclass components to create a new task.

Trainer registers subjects with Metric

Trainer currently initializes subjects with Metric when calling register_subject. Our starting policy for our curriculum operates on our own Metric subclass which forces us to include code that handles this edgecase in our initialization policy.

implement abstract automatic training manager

no database stuff, use pd.DataFrame as your interaction layer.

Dynamic Routing Task Regimen Curriculum reimplementation user experience

I've begun reimplementing the dynamic routing task regimen here:
https://github.com/AllenNeuralDynamics/dynamic-routing-curriculum?tab=readme-ov-file

• I've made minimal usage of policies and hypergraphs because I want to find out where our problems are in full implementation
• I can see how hypergraphs and policies can reduce the number of redundant stages especially in the cases where I have circular stage transitions to just alternate task parametrization until a certain secondary condition is met.
• I'm in the process of getting trainer to work and if I understand the documentation correctly I may be able to get automagic transitions if I have policy transitions attached that are redundant with the stage transitions.
• This is just here to document my user experience?
• We currently have a part of the regimen that was always moved to via manual intervention and I'm not sure how this will fit in but it might become more clear once I start using Trainer.