Description of a full ecossystem pipeline deployed in AIND

[bruno-f-cruz]

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."

[mochic]

This looks like a good initial plan. We will likely use np_codeocean for 4 but may use watchdog or switch at a later date. @corbennett does this look good to you?

[bruno-f-cruz]

For now, lets pack both Metrics and Curriculum state in the output_parameters field. I suggest simply having 2 properties at the top level: -metrics (Deserializes the Metrics class) -curriculum (Deserializes the state for the curriculum. Still need a Model for this, I will try to come up with something over the weekend)

[mochic]

@bruno-f-cruz Thank you for your consistent feedback on how you expect us to use this codebase.

[bruno-f-cruz]

I finally had some time to sit down and think about the Trainer class. These are mostly my thoughts and may be a biased interpretation of what Jonathan wrote, but please correct me if I am wrong.

---

Curriculum

A curriculum instance is a static object that fully defines the possible progression through a set of stages. It can be seen as a hierarchical structure:

• Curriculum
  • List of Stages of different Tasks
    • Policies available in each Stage
    • Initial TaskParameters (i.e., initial settings)
    • Policy transitions
    • List of Stage transitions

Trainer

A trainer takes the previous Curriculum object and uses the output of a specific session (Metrics) to suggest the next Task, TaskParameters, and Policies to use in the following session. It is important to note that these are merely suggestions that may or may not be accepted by the user (more on this later). To provide this suggestion, the Trainer must be able to index a specific location/state in the full curriculum. While the curriculum has a static definition, the state of a Trainer, given a Curriculum, is defined by:

• Current Stage
• Current active Policy

Let's call these "TrainerState". We can then write:

TaskParameters_T+1 (i.e., Suggestion) = Trainer(Curriculum, TrainerState_T, TaskParameters_T, Metrics_T)

where Metrics and TaskParameters_T must match the Stage type.

Importantly, as an artifact of the Trainer, a new TrainerState_T+1 is also generated. This new TrainerState is responsible for encoding the new Current Stage and new Active Policy, under the assumption that the next acquired session will run with the suggested TaskParameters_T+1.

Where do all pieces fit in the pipeline?

TaskParameters_T

These are the settings that were USED during an experiment. They are part of the session.json and will thus be indexed in a documentdb. They are obtained by serializing the Task subclass and stored in the StimulusEpoch.script.parameters field as per @dyf's instructions.

Curriculum

This is a static definition serialized to JSON. It can be recovered as long as it is saved somewhere. @dyf suggested it would be ideal to keep a common library of these curricula (not sure if this already exists?). An alternative would be to save them in the GitHub repo that implements a given curriculum. The JSON should only need to be updated if the Curriculum changes, thus bumping its version.

TaskParameter_T+1

This is a suggestion that has yet to materialize into real data. We decided to use SLIMS as a way to store the Task object that contains these parameters as a JSON file. If the suggestion is accepted, TaskParameters_T+1 will trigger an identical pipeline to TaskParameters_T.

Metrics_T

Similarly to TaskParameters, these should be serialized to the StimulusEpoch.output_parameters at the time of creation of session.json.

TrainerState_T+1

A recent discussion (see above) suggests that this class should be serialized next to Metrics. This would be "ok" assuming that only "on-curriculum" training will occur. However, this has a few problems should users want to override a suggestion and/or a specific state of the trainer. The bottom line is: I think we will need a modifiable object, somewhere, that can be used to store and, if needed, manually update the state of the curriculum. While the session.json might solve the "storing" issue, it seems like a clunky solution if the state needs to be modified, as it would require modifying the session.json itself. Maybe I am just missing something on how session.jsons are used and indexed? It may be the case that a different behavior event is needed in SLIMS perhaps? This would be simply responsible for keeping the state of the trainer for a given animal.

---

How to recover from manual changes in TrainerState?

While thinking about this subject I also went through some simulated scenarios of the pipeline that made me write down some notes on how to recover the Trainer from a manual intervention.

A simple heuristic could be:

1. Is the incoming TaskParameters, Metrics, and Curriculum compatible with the current TrainerState (i.e., Stage and Policies)?

2. Does the user want to try to "coerce" the previous TaskParameters into the trainer? (Remember, the user might have manually set parameters that were not suggested by the trainer, OR it may have manually changed the curriculum state.)

if (1 & 2) == True run Trainer else throw exception and require manual intervention from the user.

It is important to note that manual intervention at the level of TaskParameters or TrainerState might lead to deadlocks that the experimenter did not anticipate due to, for instance, policy transitions applied on top of Policies that would never be co-active otherwise.

Alternatively, we could also only allow a "strict" mode to be used where only stage transitions are allowed. This would essentially remove any dependency on the previous session (since each stage starts from an initial set of parameters and active policies).

Anyway happy to discuss this in person, just wanted to make sure I captured these thoughts somewhere.

[bruno-f-cruz]

After today's meeting, I believe that the consensus is that the TrainerState information should be put in a container together with a suggestion for the next session and go into SLIMS.

If a manual override is performed by the user (see "How to recover from manual changes in TrainerState?" above), it is expected to lead to an update of both task settings and curriculum state. As a result, SLIMS should be seen as both a database for scheduled experiments as well as a database for ephemeral curriculum state.

Remaining questions:

1. What should this information look like? Come up with a Container[Stage, ActivePolicies] where the TaskParameters object can be extracted from the Stage (i.e. Stage.Task.TaskParameters) object. While this may appear appealing at first, it also makes any rig software that runs a specific task dependent on the subclassed curriculum, and not the other way around. So suddenly, even if the software only runs TaskA, it also needs to be aware of TasksB and C. Not to mention that Task definition would have to be centralized at the level of the curriculum or we would otherwise risk circular dependencies. I think we should try to avoid this if possible.

I think this can easily be solved, however, by dict-indexing the json object (i.e not deserialize onto the python object), accessing the json["Container"]["Stage"]["Task"] and deserialize to the proper task at this level.

2. Future need for a UI

This should be made easy by following and sharing a common structure. Since we are adopting pydantic as our lib to build models, we should look into FastUI to procedurally generate UIs