Currently, the base class in [base.py] implements an MLP.
There are two more classes implementing tree architectures that I am not sure if they are well maintained (Michał might know)
There is also a class for the State Flow model in state_flow.py used by the DB loss.
An inconvenient aspect of the current implementation is how the combination of loss and architecture is handled in the configuration files. Currently, there needs to be a configuration file for every architecture/loss combination.
The new implementation should satisfy, in my opinion, the following desiderata:
It should be easy to create new architectures for the policy models, say a CNN, GNN, Transformer, etc.
The use of policy should not be entangled with the loss.
There could be generic architectures (MLP, CNN, etc.) whose specific main properties can be specified via the config. For example, the number of layers and units of an MLP.
Additionally, it should be easy to add a new class with a specific architecture that would be available to use with any loss.