AllenNeuralDynamics / aind-behavior-curriculum

Starter repository for behavior base primitives.
https://aind-behavior-curriculum.readthedocs.io
MIT License
1 stars 0 forks source link

Architectural class diagram #32

Open jeromelecoq opened 1 month ago

jeromelecoq commented 1 month ago

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.

jeromelecoq commented 1 month ago

@hanhou maybe you have feedbacks there?

jeromelecoq commented 1 month ago

Here is a graph of this module, as far as I understand it :

digraph "classes" {
rankdir=BT
charset="utf-8"

    subgraph cluster_0 {
        label="Base"
        BaseModel
        AindBehaviorModel
        AindBehaviorModelExtra
    }

    subgraph cluster_1 {
        label="curriculum"
        Metrics
        Rule
        Policy
        PolicyTransition
        BehaviorGraph
        PolicyGraph
        Stage
        StageTransition
        StageGraph
        Curriculum
    }

    subgraph cluster_2 {
        label="task"
        TaskParameters
        Task
    }

    subgraph cluster_3 {
        label="curriculum_utils"
        Graduated
    }

    subgraph cluster_4 {
        label="trainer"
        SubjectHistory
        Trainer
    }

BaseModel -> AindBehaviorModel
BaseModel -> AindBehaviorModelExtra
AindBehaviorModelExtra -> Metrics
AindBehaviorModel -> Policy
AindBehaviorModel -> PolicyTransition
AindBehaviorModel -> BehaviorGraph
BehaviorGraph -> PolicyGraph
AindBehaviorModel -> Stage
AindBehaviorModel -> StageTransition
BehaviorGraph -> StageGraph
AindBehaviorModel -> Curriculum

AindBehaviorModelExtra -> TaskParameters
AindBehaviorModel -> Task
Task -> Graduated
AindBehaviorModel -> SubjectHistory
}
jeromelecoq commented 1 month ago

Here is the corresponding graph : graphviz

jeromelecoq commented 1 month ago

So it has 4 layers deep and 18 classes. Is there an obvious way to simplify this tree?

bruno-f-cruz commented 1 month ago

We could simplify, for instance we could just duplicate code across the graph classes, or not use a base class for all models and instead copy-paste the configuration dict. I don't think this would be a good idea, but if we really wanted to make it horizontal and not use inheritance, we could... But tbh the use of inheritance here is warranted, as the base classes are very light-weight and serve a single purpose. Again, inheritance is just a tool like any other, it is not bad/good in a vacuum, but I think its use its warranted in this case (but happy to consider alternatives).

As I mentioned a few times before, most of those classes are abstract and users are not even expected to interact with them (looking at AindBehaviorModel and AindBehaviorModelExtra, or BehaviorGraph). In my view, the module is a single layer (In your diagram would be the 3rd layer from the bottom). This is the level users will interact with the library. Finally, as a user, I think you should be looking at the library quite differently. It's also about composition, not only inheritance.

In my mind there are three distinct modules that users are expected to interact with. These modules build on top of each other. Why? As I mentioned before, users should not be expected to buy into the full stack if they just need a part, or are not yet ready to make the full integration all at once (our current situation).

Modules

Task

This is the module you use to spec your task logic. It doesn't mean you need to buy into the curriculum, but if you want to later, you are in a really good position to do it:

task

Curriculum

This is the module you use to spec a curriculum. This is the input for a Trainer. However, the trainer could be implemented in the local machine or could be a service running remotely elsewhere.

curriculum

Trainer

This is the module that is responsible for running trainer. As opposed to the previous two modules that are mostly about specification, this module is concerned with the actual implementation of the updating logic.

As a result, depending on how much of the stack you need, you really just need to interact with <10 classes. For instance, I would probably not use any of the Policy logic, which would drop at least 3 classes :P

Anyway, this is my own opinion, I would be curious to hear @jwong-nd , @dyf and @hanhou take on it too!

jeromelecoq commented 1 month ago

Many Thanks for those details. I am also curious to hear other thoughts before discussing.

jeromelecoq commented 1 month ago

Also @alexpiet

alexpiet commented 1 month ago

Personally I think the classes could be simplified. Why is Graduated its own class, instead of just a flag on a stage? Why is there AindBehaviorModel and AindBehaviorModelExtra? I also think the type checks are excessive. I think a lot of this comes down to style, I'm a big believer in duck typing and less hierarchy. I think Bruno is the exact opposite. But - I haven't been paying attention to the details on this repo, and don't want to go through the gory details. So I'm fine with whatever works.

bruno-f-cruz commented 1 month ago

#

Re: Duck-typing

As i mentioned in #30 you can use duck-typing (provided you handle the (de)serialization logic on your own... again, I don't think you should, but you can... so use at your own risk). I will leave it at this, but happy to discuss further in person why I don't favor duck-typing for long-term projects and why I think this is necessary when using an explicit schema definition. The design of the package is such that if the user wants to do it, they can, but we still make sure that the core package doesn't rely on it, which I think is a good balance. (re-reading your question I am not sure if this is the angle you had in mind when you mentioned duck typing tho, so happy to discuss offline)

Re: AindBehaviorModel and AindBehaviorModelExtra

Once again, they are abstract classes that users should not interact with. AindBehaviorModelExtra is there to ensure that users can use duck-typing and still ensure somewhat stable (de)serialization without dropping parameters against the base class. That being said, we must validate the required fields (e.g. Task name, the existence of a task_parameters field, etc..., to ensure that any tool we use downstream is able to build on top of common interfaces. Once again, it's about allowing users the flexibility to use whatever they want but still ensuring that basic contracts are followed (the whole point of using schemas).

Re: Graduated stage

If you think making Graduated a parameter of Stage instead of a specific Stage instance, I would be happy to consider a refactoring. For me, it boils down to: If a mouse graduates, does it mean it will keep training on that graduation stage or does it exit the curriculum altogether? If the former, it does make sense to add a is_graduated_stage to Stage, if the latter, I think the current implementation makes more sense.

jeromelecoq commented 1 month ago

My personal concern is not to repeat MTRAIN path which was deemed overly complex and hard to understand by most except the initial developers. I recollect they thought it was perfectly well designed and they even gave talk on it. I learned it and I eventually thought some part of it were reasonable. So we want a solution we can train everyone to use and believe it is good. If you think this can done with this architecture, then perhaps we just need more documentation support.

jeromelecoq commented 1 month ago

Is there an argument there that the external access could be simplified while keeping the inner structure? I could propose something there I think. It seems like it is expected to access all inside objects directly. This reminds me of Tensorflow versus keras a little bit for those that are familiar.