bigger critic with LN to check if this help with expresivity

BartekCupial · BartekCupial · commit ab5f8b47c59a · 2024-05-21T15:37:45.000+02:00
diff --git a/mrunner_exps/layer_norm/2024_05_21_monk-APPO-T-layernorm-MLP_critic.py b/mrunner_exps/layer_norm/2024_05_21_monk-APPO-T-layernorm-MLP_critic.py
@@ -0,0 +1,101 @@
+from mrunner.helpers.specification_helper import create_experiments_helper
+
+name = globals()["script"][:-3]
+
+# params for all exps
+config = {
+    "env": "challenge",
+    "exp_tags": [name],
+    "exp_point": "monk-APPO-T",
+    "train_for_env_steps": 500_000_000,
+    "group": "monk-APPO-T",
+    "character": "mon-hum-neu-mal",
+    "num_workers": 16,
+    "num_envs_per_worker": 16,
+    "worker_num_splits": 2,
+    "rollout": 32,
+    "batch_size": 4096,  # this equals bs = 128, 128 * 32 = 4096
+    "async_rl": True,
+    "serial_mode": False,
+    "wandb_user": "bartekcupial",
+    "wandb_project": "sf2_nethack",
+    "wandb_group": "gmum",
+    "with_wandb": True,
+    "use_pretrained_checkpoint": True,
+    "model_path": "/net/pr2/projects/plgrid/plgggmum_crl/bcupial/sf_checkpoints/amzn-AA-BC_pretrained",
+    "use_prev_action": True,
+    "model": "ScaledNet",
+    "use_resnet": True,
+    "learning_rate": 0.0001,
+    "rnn_size": 1738,
+    "h_dim": 1738,
+    "gamma": 1.0,
+    "skip_train": 25_000_000,
+    "lr_schedule": "linear_decay",
+    "save_milestones_ith": 25_000_000,
+}
+
+params_grid = []
+expected_batch_size = 4096
+
+for rollout in [128]:
+    for target_batch_size in [128]:
+        batch_size = min(expected_batch_size, min(target_batch_size * rollout, expected_batch_size * 8))
+        batches_to_accumulate = max(1, (rollout * target_batch_size) // expected_batch_size)
+        optim_step_every_ith = max(1, batches_to_accumulate // 8)
+        params_grid.append(
+            {
+                "seed": list(range(3)),
+                "learning_rate": [0.0001],
+                "freeze": [{"actor_encoder": 0}],
+                "rollout": [rollout],
+                "batch_size": [batch_size],  # 32 * 512, 64 * 256, 128 * 128
+                "num_batches_per_epoch": [min(8, batches_to_accumulate)],
+                "optim_step_every_ith": [optim_step_every_ith],
+                "target_batch_size": [target_batch_size],
+                "actor_critic_share_weights": [False],
+                "critic_add_layernorm": [True],
+                "critic_replace_bn_with_ln": [True, False],
+                "critic_mlp_layers": [[512], [512, 512], [512, 512, 512]],
+            }
+        )
+        params_grid.append(
+            {
+                "seed": list(range(3)),
+                "learning_rate": [0.0001],
+                "freeze": [{"actor_encoder": 0}],
+                "rollout": [rollout],
+                "batch_size": [batch_size],  # 32 * 512, 64 * 256, 128 * 128
+                "num_batches_per_epoch": [min(8, batches_to_accumulate)],
+                "optim_step_every_ith": [optim_step_every_ith],
+                "target_batch_size": [target_batch_size],
+                "actor_critic_share_weights": [False],
+                "critic_mlp_layers": [[512], [512, 512], [512, 512, 512]],
+            }
+        )
+        params_grid.append(
+            {
+                "seed": list(range(3)),
+                "learning_rate": [0.0001],
+                "freeze": [{"encoder": 0}],
+                "rollout": [rollout],
+                "batch_size": [batch_size],  # 32 * 512, 64 * 256, 128 * 128
+                "num_batches_per_epoch": [min(8, batches_to_accumulate)],
+                "optim_step_every_ith": [optim_step_every_ith],
+                "target_batch_size": [target_batch_size],
+                "critic_mlp_layers": [[512], [512, 512], [512, 512, 512]],
+            }
+        )
+
+
+experiments_list = create_experiments_helper(
+    experiment_name=name,
+    project_name="sf2_nethack",
+    with_neptune=False,
+    script="python3 mrunner_run.py",
+    python_path=".",
+    tags=[name],
+    base_config=config,
+    params_grid=params_grid,
+    mrunner_ignore=".mrunnerignore",
+)
diff --git a/sample_factory/cfg/cfg.py b/sample_factory/cfg/cfg.py
@@ -1,3 +1,4 @@
+import ast
 import multiprocessing
 import os
 from argparse import ArgumentParser
@@ -565,10 +566,15 @@ def add_model_args(p: ArgumentParser):
     p.add_argument(
         "--decoder_mlp_layers",
         default=[],
-        type=int,
-        nargs="*",
+        type=ast.literal_eval,
         help="Optional decoder MLP layers after the policy core. If empty (default) decoder is identity function.",
     )
+    p.add_argument(
+        "--critic_mlp_layers",
+        default=[],
+        type=ast.literal_eval,
+        help="Optional critic MLP layers after the policy core. If empty (default) critic is a linear function.",
+    )
 
     p.add_argument(
         "--nonlinearity", default="elu", choices=["elu", "relu", "tanh"], type=str, help="Type of nonlinearity to use."
diff --git a/sample_factory/model/actor_critic.py b/sample_factory/model/actor_critic.py
@@ -150,7 +150,7 @@ def __init__(
 
         decoder_out_size: int = self.decoder.get_out_size()
 
-        self.critic_linear = nn.Linear(decoder_out_size, 1)
+        self.critic = model_factory.make_model_critic_func(cfg, self.decoder.get_out_size())
         self.action_parameterization = self.get_action_parameterization(decoder_out_size)
 
         self.apply(self.initialize_weights)
@@ -165,7 +165,7 @@ def forward_core(self, head_output: Tensor, rnn_states):
 
     def forward_tail(self, core_output, values_only: bool, sample_actions: bool) -> TensorDict:
         decoder_output = self.decoder(core_output)
-        values = self.critic_linear(decoder_output).squeeze()
+        values = self.critic(decoder_output).squeeze()
 
         result = TensorDict(values=values)
         if values_only:
@@ -212,7 +212,7 @@ def __init__(
         self.critic_decoder = model_factory.make_model_decoder_func(cfg, self.critic_core.get_out_size())
         self.decoders = [self.actor_decoder, self.critic_decoder]
 
-        self.critic_linear = nn.Linear(self.critic_decoder.get_out_size(), 1)
+        self.critic = model_factory.make_model_critic_func(cfg, self.critic_decoder.get_out_size())
         self.action_parameterization = self.get_action_parameterization(self.critic_decoder.get_out_size())
 
         self.apply(self.initialize_weights)
@@ -284,7 +284,7 @@ def forward_tail(self, core_output, values_only: bool, sample_actions: bool) ->
 
         # second core output corresponds to the critic
         critic_decoder_output = self.critic_decoder(core_outputs[1])
-        values = self.critic_linear(critic_decoder_output).squeeze()
+        values = self.critic(critic_decoder_output).squeeze()
 
         result = TensorDict(values=values)
         if values_only:
diff --git a/sample_factory/model/critic.py b/sample_factory/model/critic.py
@@ -0,0 +1,81 @@
+import math
+from abc import ABC
+from typing import List
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+from sample_factory.algo.utils.action_distributions import ContinuousActionDistribution
+from sample_factory.algo.utils.torch_utils import calc_num_elements
+from sample_factory.model.model_utils import ModelModule, create_mlp, nonlinearity
+from sample_factory.utils.typing import Config
+
+
+class Critic(ModelModule, ABC):
+    pass
+
+
+class MlpCritic(Critic):
+    def __init__(self, cfg: Config, critic_input_size: int):
+        super().__init__(cfg)
+        self.critic_input_size = critic_input_size
+        self.critic_out_size = 1
+        critic_layers: List[int] = cfg.critic_mlp_layers
+        activation = nonlinearity(cfg)
+        self.mlp = create_mlp(critic_layers, critic_input_size, activation)
+        if len(critic_layers) > 0:
+            self.mlp = torch.jit.script(self.mlp)
+
+        mlp_out_size = calc_num_elements(self.mlp, (critic_input_size,))
+        self.critic_linear = nn.Linear(mlp_out_size, self.critic_out_size)
+
+    def forward(self, core_output):
+        return self.critic_linear(self.mlp(core_output))
+
+
+class ValueParameterizationContinuousNonAdaptiveStddev(nn.Module):
+    """Use a single learned parameter for action stddevs."""
+
+    def __init__(self, cfg, core_out_size):
+        super().__init__()
+        self.cfg = cfg
+
+        # calculate only value means using the critic neural network
+        self.distribution_linear = nn.Linear(core_out_size, 1)
+        # stddev is a single learned parameter
+        initial_stddev = torch.empty([1])
+        initial_stddev.fill_(math.log(self.cfg.initial_stddev))
+        self.learned_stddev = nn.Parameter(initial_stddev, requires_grad=True)
+
+    def forward(self, actor_core_output: Tensor):
+        value_means = self.distribution_linear(actor_core_output)
+        batch_size = value_means.shape[0]
+        value_stddevs = self.learned_stddev.repeat(batch_size, 1)
+        value_distribution_params = torch.cat((value_means, value_stddevs), dim=1)
+        value_distribution = ContinuousActionDistribution(params=value_distribution_params)
+        return value_distribution_params, value_distribution
+
+
+class ParametrizedCritic(Critic):
+    def __init__(self, cfg: Config, critic_input_size: int):
+        super().__init__(cfg)
+        self.critic_input_size = critic_input_size
+        self.critic_out_size = 2
+        critic_layers: List[int] = cfg.critic_mlp_layers
+        activation = nonlinearity(cfg)
+        self.mlp = create_mlp(critic_layers, critic_input_size, activation)
+        if len(critic_layers) > 0:
+            self.mlp = torch.jit.script(self.mlp)
+
+        mlp_out_size = calc_num_elements(self.mlp, (critic_input_size,))
+        self.critic_parametrization = ValueParameterizationContinuousNonAdaptiveStddev(cfg, mlp_out_size)
+
+    def forward(self, core_output):
+        value_distribution_params, self.last_value_distribution = self.critic_parametrization(self.mlp(core_output))
+        values = self.last_value_distribution.sample()
+        return values
+
+
+def default_make_critic_func(cfg: Config, critic_input_size: int) -> Critic:
+    return MlpCritic(cfg, critic_input_size)
diff --git a/sample_factory/model/model_factory.py b/sample_factory/model/model_factory.py
@@ -2,6 +2,7 @@
 
 from sample_factory.model.actor_critic import ActorCritic, default_make_actor_critic_func
 from sample_factory.model.core import ModelCore, default_make_core_func
+from sample_factory.model.critic import Critic, default_make_critic_func
 from sample_factory.model.decoder import Decoder, default_make_decoder_func
 from sample_factory.model.encoder import Encoder, default_make_encoder_func
 from sample_factory.utils.typing import ActionSpace, Config, ObsSpace
@@ -11,6 +12,7 @@
 MakeEncoderFunc = Callable[[Config, ObsSpace], Encoder]
 MakeCoreFunc = Callable[[Config, int], ModelCore]
 MakeDecoderFunc = Callable[[Config, int], Decoder]
+MakeCriticFunc = Callable[[Config, int], Critic]
 
 
 class ModelFactory:
@@ -28,6 +30,7 @@ def __init__(self):
         self.make_model_encoder_func: MakeEncoderFunc = default_make_encoder_func
         self.make_model_core_func: MakeCoreFunc = default_make_core_func
         self.make_model_decoder_func: MakeDecoderFunc = default_make_decoder_func
+        self.make_model_critic_func: MakeCriticFunc = default_make_critic_func
 
     def register_actor_critic_factory(self, make_actor_critic_func: MakeActorCriticFunc):
         """
@@ -59,3 +62,11 @@ def register_decoder_factory(self, make_model_decoder_func: MakeDecoderFunc):
         """
         log.debug(f"register_decoder_factory: {make_model_decoder_func}")
         self.make_model_decoder_func = make_model_decoder_func
+
+    def register_critic_factory(self, make_model_critic_func: MakeCriticFunc):
+        """
+        Override the default decoder with a custom model.
+        The computational graph structure is: observations -> encoder -> core -> decoder -> actions
+        """
+        log.debug(f"register_critic_factory: {make_model_critic_func}")
+        self.make_model_critic_func = make_model_critic_func
diff --git a/sf_examples/nethack/train_nethack.py b/sf_examples/nethack/train_nethack.py
@@ -3,6 +3,7 @@
 from os.path import join
 from typing import Callable
 
+import torch
 import torch.nn as nn
 
 from sample_factory.algo.learning.learner import Learner
@@ -12,6 +13,7 @@
 from sample_factory.envs.env_utils import register_env
 from sample_factory.model.actor_critic import ActorCritic, default_make_actor_critic_func
 from sample_factory.model.encoder import Encoder
+from sample_factory.model.model_utils import get_rnn_size
 from sample_factory.train import run_rl
 from sample_factory.utils.typing import ActionSpace, Config, ObsSpace
 from sample_factory.utils.utils import log
@@ -71,7 +73,8 @@ def load_pretrained_checkpoint(model, checkpoint_dir: str, checkpoint_kind: str,
         del checkpoint_dict["model"]["returns_normalizer.running_var"]
         del checkpoint_dict["model"]["returns_normalizer.count"]
 
-    model.load_state_dict(checkpoint_dict["model"])
+    incompatibile = model.load_state_dict(checkpoint_dict["model"], strict=False)
+    log.debug(incompatibile)
 
 
 def load_pretrained_checkpoint_from_shared_weights(
@@ -126,17 +129,19 @@ def hook(module, input, output):
                 else:
                     register_hooks(child)
 
-        register_hooks(model.critic_encoder)
+        register_hooks(model)
 
         tmp_env = make_env_func_batched(cfg, env_config=None)
         obs, info = tmp_env.reset()
-        model.critic_encoder(obs)
+        rnn_states = torch.zeros([1, get_rnn_size(cfg)], dtype=torch.float32)
+        model(obs, rnn_states)
 
         if cfg.critic_replace_bn_with_ln:
             replace_batchnorm_with_layernorm(model.critic_encoder)
         inject_layernorm_before_activation(model.critic_encoder)
 
-        model.critic_linear = linear_layernorm(model.critic_linear)
+        inject_layernorm_before_activation(model.critic)
+        model.critic.critic_linear = linear_layernorm(model.critic.critic_linear)
 
         for handle in handles:
             handle.remove()
