Changing structure of Environment Embeddings

[fedebotu]

Env Embeddings

At the moment, the environment embeddings have been modularized such that they can be easily switched for different environments. The idea is that they should transform data from feature space to the hidden space ready to be processed. (docs)

The current structure is as follows:

rl4co/
└── models/
    └── nn/
        └── env_embeddings/
            ├── context.py
            ├── dynamic.py
            └── init.py

Problem

However, this can easily get cluttered, especially with the addition of several new environments (each env has its own class inside of each file).
Moreover, this categorization can be suboptimal for models such as Non-autoregressive approaches (in which we mostly just have a init embedding, #122 ); also, for new multi-task learning problems such as Fei Liu and @Xi-L 's paper, it might be better to have separate files to handle environments (or their attribute compositions).

Proposal

Divide and conquer the environment embeddings as follows:

rl4co/
└── models/
    └── nn/
        └── env/
            ├── eda/
            │   ├── dpp.py
            │   └── ...
            ├── routing/
            │   ├── cvrp.py
            │   ├── multitask.py
            │   ├── tsp.py
            │   └── ...
            ├── scheduling/
            │   ├── jsp.py
            │   └── ...
            └── other_envs/
                └── ...

In which each file has all the related env embeddings (i.e. context, dynamic and init).

Expected benefits:

• Easier for newcomers. If one wants to solve a new problem, simply make related embeddings all in one file
• Clarity. Previous approach was starting to become too cluttered and harder to follow
• Modularity. More in line with RL4CO's modularity - we have a problem as environment and similarly, we have model embedding divided by each problem

What do you think? cc @cbhua @ngastzepeda

[cbhua]

As we discussed before, I like this refactoring! Let's try to keep the rl4co/envs and the rl4co/models/nn/envs having the same file structure:

rl4co/
├── envs/
│   ├── routing/
│   │   ├── tsp.py (reset, step, reward, ...)
│   │   └── ...
│   ├── scheduling/
│   │   ├── ffsp.py
│   │   └── ...
│   └── ...
└── models/
    └── nn/
        └── envs/ (Q: will we use envs or env?)
            ├── routing/
            │   ├── tsp.py (init, context, dynamic, ...)
            │   └── ...
            ├── scheduling/
            │   ├── ffsp.py
            │   └── ...
            └── ...

This could make it easier for users to understand and clearer for them to contribute.

Besides, there is a radical design: we could even put environments and their neural networks together inside the rl4co/envs folder.

rl4co/
├── envs/
│   ├── routing/
│   │   ├── tsp/
│   │   │   ├── env.py (reset, step, reward, ...)
│   │   │   └── nn.py (init, context, dynamic, ...)
│   │   └── ...
│   ├── scheduling/
│   │   ├── ffsp/
│   │   │   ├── env.py (reset, step, reward, ...)
│   │   │   └── nn.py (init, context, dynamic, ...)
│   │   └── ...
│   └── ...
└── models/
    └── nn/
        └── ...

This could be too much for refactoring in the current moment since envs relys quite a few functions from RL4CO. But if we want to split environments as an independent package in the future, we may think about this.

[fedebotu]

Honestly I think envs should be separate from the nn part, with the reason being that theoretically the user should not need to touch the environment logic but they could touch the network and modeling parts

[cbhua]

Make sense, let's choose the first design.

[fedebotu]

Note: we should also update the documentation accordingly