Major modeling refactoring

docs/_content/api/algos/search.md

@@ -1,7 +1,7 @@
 # Search
 
 ```{eval-rst}
-.. automodule:: rl4co.models.zoo.common.search.base
+.. automodule:: rl4co.models.common.search.base
    :members:
    :undoc-members:
 ```

docs/_content/api/models/base.md

@@ -10,23 +10,23 @@ Autoregressive models are models that generate sequences one element at a time,
 ### Policy
 
 ```{eval-rst}
-.. automodule:: rl4co.models.zoo.common.autoregressive.policy
+.. automodule:: rl4co.models.common.constructive.autoregressive.policy
    :members:
    :undoc-members:
 ```
 
 ### Encoder
 
 ```{eval-rst}
-.. automodule:: rl4co.models.zoo.common.autoregressive.encoder
+.. automodule:: rl4co.models.common.constructive.autoregressive.encoder
    :members:
    :undoc-members:
 ```
 
 ### Decoder
 
 ```{eval-rst}
-.. automodule:: rl4co.models.zoo.common.autoregressive.decoder
+.. automodule:: rl4co.models.common.constructive.autoregressive.decoder
    :members:
    :undoc-members:
 ```
@@ -38,15 +38,15 @@ Non-autoregressive models generate a heatmap of probabilities from one node to a
 ### Policy
 
 ```{eval-rst}
-.. automodule:: rl4co.models.zoo.common.nonautoregressive.policy
+.. automodule:: rl4co.models.common.nonautoregressive.policy
    :members:
    :undoc-members:
 ```
 
 ### Encoder
 
 ```{eval-rst}
-.. automodule:: rl4co.models.zoo.common.nonautoregressive.encoder
+.. automodule:: rl4co.models.common.nonautoregressive.encoder
    :members:
    :undoc-members:
 ```
@@ -57,7 +57,7 @@ Note that we still need a decoding class for the heatmap (for example, to mask o
 
 
 ```{eval-rst}
-.. automodule:: rl4co.models.zoo.common.nonautoregressive.decoder
+.. automodule:: rl4co.models.common.nonautoregressive.decoder
    :members:
    :undoc-members:
 ```

examples/3-creating-new-env-model.ipynb

@@ -507,10 +507,10 @@
     "        - locs: x, y coordinates of the cities\n",
     "    \"\"\"\n",
     "\n",
-    "    def __init__(self, embedding_dim, linear_bias=True):\n",
+    "    def __init__(self, embed_dim, linear_bias=True):\n",
     "        super(TSPInitEmbedding, self).__init__()\n",
     "        node_dim = 2  # x, y\n",
-    "        self.init_embed = nn.Linear(node_dim, embedding_dim, linear_bias)\n",
+    "        self.init_embed = nn.Linear(node_dim, embed_dim, linear_bias)\n",
     "\n",
     "    def forward(self, td):\n",
     "        out = self.init_embed(td[\"locs\"])\n",
@@ -539,13 +539,13 @@
     "        - current node embedding\n",
     "    \"\"\"\n",
     "\n",
-    "    def __init__(self, embedding_dim,  linear_bias=True):\n",
+    "    def __init__(self, embed_dim,  linear_bias=True):\n",
     "        super(TSPContext, self).__init__()\n",
     "        self.W_placeholder = nn.Parameter(\n",
-    "            torch.Tensor(2 * embedding_dim).uniform_(-1, 1)\n",
+    "            torch.Tensor(2 * embed_dim).uniform_(-1, 1)\n",
     "        )\n",
     "        self.project_context = nn.Linear(\n",
-    "            embedding_dim*2, embedding_dim, bias=linear_bias\n",
+    "            embed_dim*2, embed_dim, bias=linear_bias\n",
     "        )\n",
     "\n",
     "    def forward(self, embeddings, td):\n",
@@ -620,7 +620,7 @@
     "# Instantiate policy with the embeddings we created above\n",
     "emb_dim = 128\n",
     "policy = AutoregressivePolicy(env,\n",
-    "                              embedding_dim=emb_dim,\n",
+    "                              embed_dim=emb_dim,\n",
     "                              init_embedding=TSPInitEmbedding(emb_dim),\n",
     "                              context_embedding=TSPContext(emb_dim),\n",
     "                              dynamic_embedding=StaticEmbedding(emb_dim)\n",

examples/advanced/2-flash-attention-2.ipynb

@@ -76,7 +76,7 @@
     "from rl4co.envs import TSPEnv\n",
     "from rl4co.models.zoo.am import AttentionModel\n",
     "from rl4co.utils.trainer import RL4COTrainer\n",
-    "from rl4co.models.zoo.common.autoregressive import GraphAttentionEncoder\n",
+    "from rl4co.models.common.constructive.autoregressive import GraphAttentionEncoder\n",
     "\n"
    ]
   },
@@ -225,15 +225,15 @@
     "env = TSPEnv(num_loc=1000)\n",
     "\n",
     "num_heads = 8\n",
-    "embedding_dim = 128\n",
+    "embed_dim = 128\n",
     "num_layers = 3\n",
-    "enc_simple = GraphAttentionEncoder(env, num_heads=num_heads, embedding_dim=embedding_dim, num_layers=num_layers,\n",
+    "enc_simple = GraphAttentionEncoder(env, num_heads=num_heads, embed_dim=embed_dim, num_layers=num_layers,\n",
     "                            sdpa_fn=scaled_dot_product_attention_simple)\n",
     "\n",
-    "enc_fa1 = GraphAttentionEncoder(env, num_heads=num_heads, embedding_dim=embedding_dim, num_layers=num_layers,\n",
+    "enc_fa1 = GraphAttentionEncoder(env, num_heads=num_heads, embed_dim=embed_dim, num_layers=num_layers,\n",
     "                            sdpa_fn=scaled_dot_product_attention)\n",
     "\n",
-    "enc_fa2 = GraphAttentionEncoder(env, num_heads=num_heads, embedding_dim=embedding_dim, num_layers=num_layers,\n",
+    "enc_fa2 = GraphAttentionEncoder(env, num_heads=num_heads, embed_dim=embed_dim, num_layers=num_layers,\n",
     "                            sdpa_fn=scaled_dot_product_attention_flash_attn)\n",
     "\n",
     "# Flash Attention supports only FP16 and BFloat16\n",
@@ -248,14 +248,14 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "def build_models(num_heads=8, embedding_dim=128, num_layers=3):\n",
-    "    enc_simple = GraphAttentionEncoder(env, num_heads=num_heads, embedding_dim=embedding_dim, num_layers=num_layers,\n",
+    "def build_models(num_heads=8, embed_dim=128, num_layers=3):\n",
+    "    enc_simple = GraphAttentionEncoder(env, num_heads=num_heads, embed_dim=embed_dim, num_layers=num_layers,\n",
     "                                sdpa_fn=scaled_dot_product_attention_simple)\n",
     "\n",
-    "    enc_fa1 = GraphAttentionEncoder(env, num_heads=num_heads, embedding_dim=embedding_dim, num_layers=num_layers,\n",
+    "    enc_fa1 = GraphAttentionEncoder(env, num_heads=num_heads, embed_dim=embed_dim, num_layers=num_layers,\n",
     "                                sdpa_fn=scaled_dot_product_attention)\n",
     "\n",
-    "    enc_fa2 = GraphAttentionEncoder(env, num_heads=num_heads, embedding_dim=embedding_dim, num_layers=num_layers,\n",
+    "    enc_fa2 = GraphAttentionEncoder(env, num_heads=num_heads, embed_dim=embed_dim, num_layers=num_layers,\n",
     "                                sdpa_fn=scaled_dot_product_attention_flash_attn)\n",
     "\n",
     "    # Flash Attention supports only FP16 and BFloat16\n",
@@ -295,10 +295,10 @@
     "times_fa1 = []\n",
     "times_fa2 = []\n",
     "\n",
-    "# for embedding_dim in [64, 128, 256]:\n",
-    "for embedding_dim in [128]:\n",
+    "# for embed_dim in [64, 128, 256]:\n",
+    "for embed_dim in [128]:\n",
     "    # Get models\n",
-    "    enc_simple, enc_fa1, enc_fa2 = build_models(embedding_dim=embedding_dim)\n",
+    "    enc_simple, enc_fa1, enc_fa2 = build_models(embed_dim=embed_dim)\n",
     "\n",
     "    for problem_size in sizes:\n",
     "\n",

examples/modeling/1-decoding-strategies.ipynb

@@ -70,7 +70,7 @@
     "\n",
     "# Policy: neural network, in this case with encoder-decoder architecture\n",
     "policy = AttentionModelPolicy(env.name, \n",
-    "                              embedding_dim=128,\n",
+    "                              embed_dim=128,\n",
     "                              num_encoder_layers=3,\n",
     "                              num_heads=8,\n",
     "                            )\n",

examples/modeling/2-search-methods.ipynb

@@ -190,7 +190,7 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "INFO:rl4co.models.zoo.common.autoregressive.policy:Instantiated environment not provided; instantiating tsp\n"
+      "INFO:rl4co.models.common.constructive.autoregressive.policy:Instantiated environment not provided; instantiating tsp\n"
      ]
     },
     {

examples/modeling/3-change-encoder.ipynb

@@ -280,14 +280,14 @@
     "\n",
     "gcn_encoder = GCNEncoder(\n",
     "    env_name='cvrp', \n",
-    "    embedding_dim=128,\n",
+    "    embed_dim=128,\n",
     "    num_nodes=20, \n",
     "    num_layers=3,\n",
     ")\n",
     "\n",
     "mpnn_encoder = MessagePassingEncoder(\n",
     "    env_name='cvrp', \n",
-    "    embedding_dim=128,\n",
+    "    embed_dim=128,\n",
     "    num_nodes=20, \n",
     "    num_layers=3,\n",
     ")\n",
@@ -464,15 +464,15 @@
     "    def __init__(\n",
     "            self,\n",
     "            env_name: str,\n",
-    "            embedding_dim: int,\n",
+    "            embed_dim: int,\n",
     "            init_embedding: nn.Module = None,\n",
     "        ):\n",
     "        super(BaseEncoder, self).__init__()\n",
     "        self.env_name = env_name\n",
     "        \n",
     "        # Init embedding for each environment\n",
     "        self.init_embedding = (\n",
-    "            env_init_embedding(self.env_name, {\"embedding_dim\": embedding_dim})\n",
+    "            env_init_embedding(self.env_name, {\"embed_dim\": embed_dim})\n",
     "            if init_embedding is None\n",
     "            else init_embedding\n",
     "        )\n",

rl4co/__init__.py

@@ -1 +1 @@
-__version__ = "0.4.0dev1"
+__version__ = "0.4.0dev2"

rl4co/models/__init__.py

@@ -1,17 +1,25 @@
-from rl4co.models.zoo.active_search import ActiveSearch
-from rl4co.models.zoo.am import AttentionModel, AttentionModelPolicy
-from rl4co.models.zoo.common.autoregressive import (
+from rl4co.models.common.constructive.autoregressive import (
     AutoregressiveDecoder,
+    AutoregressiveEncoder,
     AutoregressivePolicy,
-    GraphAttentionEncoder,
 )
-from rl4co.models.zoo.common.nonautoregressive import (
+from rl4co.models.common.constructive.base import (
+    ConstructiveDecoder,
+    ConstructiveEncoder,
+    ConstructivePolicy,
+)
+from rl4co.models.common.constructive.nonautoregressive import (
     NonAutoregressiveDecoder,
     NonAutoregressiveEncoder,
-    NonAutoregressiveModel,
     NonAutoregressivePolicy,
 )
-from rl4co.models.zoo.common.search import SearchBase
+from rl4co.models.common.search import SearchBase
+from rl4co.models.rl.common.base import RL4COLitModule
+from rl4co.models.rl.ppo.ppo import PPO
+from rl4co.models.rl.reinforce.baselines import REINFORCEBaseline, get_reinforce_baseline
+from rl4co.models.rl.reinforce.reinforce import REINFORCE
+from rl4co.models.zoo.active_search import ActiveSearch
+from rl4co.models.zoo.am import AttentionModel, AttentionModelPolicy
 from rl4co.models.zoo.deepaco import DeepACO, DeepACOPolicy
 from rl4co.models.zoo.eas import EAS, EASEmb, EASLay
 from rl4co.models.zoo.ham import (
@@ -20,7 +28,7 @@
 )
 from rl4co.models.zoo.matnet import MatNet, MatNetPolicy
 from rl4co.models.zoo.mdam import MDAM, MDAMPolicy
-from rl4co.models.zoo.pomo import POMO, POMOPolicy
-from rl4co.models.zoo.ppo import PPOModel, PPOPolicy
+from rl4co.models.zoo.pomo import POMO
+from rl4co.models.zoo.ppo import PPOModel
 from rl4co.models.zoo.ptrnet import PointerNetwork, PointerNetworkPolicy
 from rl4co.models.zoo.symnco import SymNCO, SymNCOPolicy

rl4co/models/common/__init__.py

@@ -0,0 +1,15 @@
+from rl4co.models.common.constructive.autoregressive import (
+    AutoregressiveDecoder,
+    AutoregressiveEncoder,
+    AutoregressivePolicy,
+)
+from rl4co.models.common.constructive.base import (
+    ConstructiveDecoder,
+    ConstructiveEncoder,
+    ConstructivePolicy,
+)
+from rl4co.models.common.constructive.nonautoregressive import (
+    NonAutoregressiveDecoder,
+    NonAutoregressiveEncoder,
+    NonAutoregressivePolicy,
+)

rl4co/models/common/constructive/__init__.py

@@ -0,0 +1,15 @@
+from rl4co.models.common.constructive.autoregressive import (
+    AutoregressiveDecoder,
+    AutoregressiveEncoder,
+    AutoregressivePolicy,
+)
+from rl4co.models.common.constructive.base import (
+    ConstructiveDecoder,
+    ConstructiveEncoder,
+    ConstructivePolicy,
+)
+from rl4co.models.common.constructive.nonautoregressive import (
+    NonAutoregressiveDecoder,
+    NonAutoregressiveEncoder,
+    NonAutoregressivePolicy,
+)

rl4co/models/common/constructive/autoregressive/__init__.py

@@ -0,0 +1,3 @@
+from rl4co.models.common.constructive.autoregressive.decoder import AutoregressiveDecoder
+from rl4co.models.common.constructive.autoregressive.encoder import AutoregressiveEncoder
+from rl4co.models.common.constructive.autoregressive.policy import AutoregressivePolicy

rl4co/models/common/constructive/autoregressive/decoder.py

@@ -0,0 +1,11 @@
+from rl4co.models.common.constructive.base import ConstructiveDecoder
+
+
+class AutoregressiveDecoder(ConstructiveDecoder):
+    """Template class for an autoregressive decoder, simple wrapper around
+    :class: rl4co.models.common.constructive.base.ConstructiveDecoder
+
+    Tip:
+        This class will not work as it is and is just a template.
+        An example for autoregressive encoder can be found as :class: rl4co.models.zoo.am.decoder.AttentionModelDecoder.
+    """

rl4co/models/common/constructive/autoregressive/encoder.py

@@ -0,0 +1,11 @@
+from rl4co.models.common.constructive.base import ConstructiveEncoder
+
+
+class AutoregressiveEncoder(ConstructiveEncoder):
+    """Template class for an autoregressive encoder, simple wrapper around
+    :class: rl4co.models.common.constructive.base.ConstructiveEncoder
+
+    Tip:
+        This class will not work as it is and is just a template.
+        An example for autoregressive encoder can be found as :class: rl4co.models.zoo.am.encoder.AttentionModelEncoder.
+    """

rl4co/models/common/constructive/autoregressive/policy.py

@@ -0,0 +1,49 @@
+from typing import Union
+
+from rl4co.envs import RL4COEnvBase
+from rl4co.models.common.constructive.base import ConstructivePolicy
+
+from .decoder import AutoregressiveDecoder
+from .encoder import AutoregressiveEncoder
+
+
+class AutoregressivePolicy(ConstructivePolicy):
+    """Template class for an autoregressive policy, simple wrapper around
+    :class: rl4co.models.common.constructive.base.ConstructivePolicy.
+
+    Note:
+        While a decoder is required, an encoder is optional and will be initialized to
+        :class: rl4co.models.common.constructive.autoregressive.encoder.NoEncoder .
+        This can be used in decoder-only models in which at each step actions do not depend on
+        previously encoded states.
+    """
+
+    def __init__(
+        self,
+        encoder: AutoregressiveEncoder = None,
+        decoder: AutoregressiveDecoder = None,
+        env_name: Union[str, RL4COEnvBase] = "tsp",
+        temperature: float = 1.0,
+        tanh_clipping: float = 0,
+        mask_logits: bool = True,
+        train_decode_type: str = "sampling",
+        val_decode_type: str = "greedy",
+        test_decode_type: str = "greedy",
+        **unused_kw,
+    ):
+        # We raise an error for the user if no decoder was provided
+        if decoder is None:
+            raise ValueError("AutoregressivePolicy requires a decoder to be provided.")
+
+        super(AutoregressivePolicy, self).__init__(
+            encoder=encoder,
+            decoder=decoder,
+            env_name=env_name,
+            temperature=temperature,
+            tanh_clipping=tanh_clipping,
+            mask_logits=mask_logits,
+            train_decode_type=train_decode_type,
+            val_decode_type=val_decode_type,
+            test_decode_type=test_decode_type,
+            **unused_kw,
+        )