Repository: ray-project/ray

Issue:
[RLlib] Add informative error message when bad Conv2D stack is used with fixed `num_outputs` (no flattening at end).
We should add an informative error message to visionnet.py (torch and tf), when a bad Conv2D stack (`conv_filters` in model config) is configured by the user, with a fixed `num_outputs` (no flattening at end). In this case, the output of the last Conv2D layer must always be [B, num_outputs, 1, 1].

```
import gym
import numpy as np
import ray
from ray import tune
from ray.rllib.models import ModelCatalog
from ray.tune.registry import register_env
from ray.rllib.agents.ppo import PPOTrainer, ppo
from ray.tune import CLIReporter
from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
from ray.rllib.models.torch.fcnet import FullyConnectedNetwork
from ray.rllib.models.torch.visionnet import VisionNetwork

from ray.rllib.utils.framework import try_import_torch


ray.init()

torch, nn = try_import_torch()

#### ENV

class JSS(gym.Env):

    def __init__(self, env_config):
        self.problem_size = 16
        self.action_space = gym.spaces.Discrete(self.problem_size + 1)
        self.observation_space = gym.spaces.Dict({
            'real_obs': gym.spaces.Box(low=0.0, high=1.0, shape=(self.problem_size, 8, 1), dtype=np.float),
            'action_mask': gym.spaces.Box(low=0, high=1, shape=(self.action_space.n, ), dtype=np.int),
        })
        self.i = 0

    def reset(self):
        self.i = 0
        legal_action = np.ones((self.action_space.n,), dtype=np.int)
        legal_action[0] = 0
        output = {
            'real_obs': np.zeros((self.problem_size, 8, 1), dtype=np.float),
            'action_mask': legal_action
        }
        return output

    def step(self, action: int):
        self.i += 1
        legal_action = np.ones((self.action_space.n, ), dtype=np.int)
        legal_action[self.problem_size] = 0
        legal_action[self.i % self.problem_size + 1] = 0
        output = {
            'real_obs': np.zeros((self.problem_size, 8, 1), dtype=np.float),
            'action_mask': legal_action
        }
        return output

    def render(self, mode='human'):
        pass


def env_creator(env_config):
    return JSS(env_config)


register_env("jss_env", env_creator)

## MODEL

# THIS ONE WORKS
class FCMaskedActionsModel(TorchModelV2, nn.Module):

    def __init__(self,
                 obs_space,
                 action_space,
                 num_outputs,
                 model_config,
                 name):
        nn.Module.__init__(self)
        super(FCMaskedActionsModel, self).__init__(obs_space, action_space, action_space.n, model_config, name)
        true_obs_space = gym.spaces.MultiBinary(n=obs_space.shape[0] - action_space.n)
        self.action_model = FullyConnectedNetwork(
            obs_space=true_obs_space, action_space=action_space, num_outputs=action_space.n,
            model_config=model_config, name=name + "action_model")

    def forward(self, input_dict, state, seq_lens):
        # Extract the available actions tensor from the observation.
        action_mask = input_dict["obs"]["action_mask"]
        # Compute the predicted action embedding

        raw_actions, state = self.action_model({
            "obs": input_dict["obs"]["real_obs"]
        })
        inf_mask = torch.clamp(torch.log(action_mask), min=torch.finfo(raw_actions.dtype).min)
        actions = raw_actions + inf_mask

        return actions, state

    def value_function(self):
        return self.action_model.value_function()


#THIS ONE HAS A SIZE MISMATCH PROBLEM WITH THE LOGIT LAYER
class ConvMaskedActionsModel(TorchModelV2, nn.Module):

    def __init__(self,
                 obs_space,
                 action_space,
                 num_outputs,
                 model_config,
                 name):
        nn.Module.__init__(self)
        super(ConvMaskedActionsModel, self).__init__(obs_space, action_space, action_space.n, model_config, name)
        true_obs_space = obs_space.original_space['real_obs']
        self.action_model = VisionNetwork(
            obs_space=true_obs_space, action_space=action_space, num_outputs=action_space.n,
            model_config=model_config, name=name + "action_model_conv")

    def forward(self, input_dict, state, seq_lens):
        # Extract the available actions tensor from the observation.
        action_mask = input_dict["obs"]["action_mask"]
        # Compute the predicted action embedding
        raw_actions, state = self.action_model({
            "obs": input_dict["obs"]["real_obs"]
        })
        inf_mask = torch.clamp(torch.log(action_mask), min=torch.finfo(raw_actions.dtype).min)
        actions = raw_actions + inf_mask

        return actions, state

    def value_function(self):
        return self.action_model.value_function()

ModelCatalog.register_custom_model("fc_masked_model", FCMaskedActionsModel)
ModelCatalog.register_custom_model("conv_masked_model", ConvMaskedActionsModel)


config = ppo.DEFAULT_CONFIG.copy()
config['framework'] = 'torch'
config['env'] = 'jss_env'
config['model'] = {
    "conv_filters": [[4, [4, 4], 1], [8, [4, 4], 1]],
    "custom_model": "conv_masked_model",
    # Nonlinearity for fully connected net (tanh, relu)
    "fcnet_activation": "tanh",
}
config['train_batch_size'] = 1024
config['rollout_fragment_length'] = 1000
config['sgd_minibatch_size'] = 256
config['lr'] = 1e-4
config['evaluation_interval'] = None
config['metrics_smoothing_episodes'] = 1000

stop = {
    "time_total_s": 60,
}
reporter = CLIReporter(max_progress_rows=15)
reporter.add_metric_column("episode_reward_max")
analysis = tune.run(PPOTrainer, config=config, stop=stop, progress_reporter=reporter)
print("Best config: ", analysis.get_best_config(metric="episode_reward_max"))
ray.shutdown()
```


Discussion:

Resolving patch:
diff --git a/rllib/models/tf/visionnet.py b/rllib/models/tf/visionnet.py
--- a/rllib/models/tf/visionnet.py
+++ b/rllib/models/tf/visionnet.py
@@ -11,15 +11,17 @@ class VisionNetwork(TFModelV2):
 
     def __init__(self, obs_space, action_space, num_outputs, model_config,
                  name):
+        if not model_config.get("conv_filters"):
+            model_config["conv_filters"] = _get_filter_config(obs_space.shape)
+
         super(VisionNetwork, self).__init__(obs_space, action_space,
                                             num_outputs, model_config, name)
 
-        activation = get_activation_fn(model_config.get("conv_activation"))
-        filters = model_config.get("conv_filters")
-        if not filters:
-            filters = _get_filter_config(obs_space.shape)
-        no_final_linear = model_config.get("no_final_linear")
-        vf_share_layers = model_config.get("vf_share_layers")
+        activation = get_activation_fn(
+            self.model_config.get("conv_activation"), framework="tf")
+        filters = self.model_config["conv_filters"]
+        no_final_linear = self.model_config.get("no_final_linear")
+        vf_share_layers = self.model_config.get("vf_share_layers")
 
         inputs = tf.keras.layers.Input(
             shape=obs_space.shape, name="observations")
@@ -73,6 +75,16 @@ def __init__(self, obs_space, action_space, num_outputs, model_config,
                     padding="same",
                     data_format="channels_last",
                     name="conv_out")(last_layer)
+
+                if conv_out.shape[1] != 1 or conv_out.shape[2] != 1:
+                    raise ValueError(
+                        "Given `conv_filters` ({}) do not result in a [B, 1, "
+                        "1, {} (`num_outputs`)] shape (but in {})! Please "
+                        "adjust your Conv2D stack such that the dims 1 and 2 "
+                        "are both 1.".format(
+                            self.model_config["conv_filters"],
+                            self.num_outputs, list(conv_out.shape)))
+
             # num_outputs not known -> Flatten, then set self.num_outputs
             # to the resulting number of nodes.
             else:
diff --git a/rllib/models/torch/visionnet.py b/rllib/models/torch/visionnet.py
--- a/rllib/models/torch/visionnet.py
+++ b/rllib/models/torch/visionnet.py
@@ -15,17 +15,18 @@ class VisionNetwork(TorchModelV2, nn.Module):
 
     def __init__(self, obs_space, action_space, num_outputs, model_config,
                  name):
+        if not model_config.get("conv_filters"):
+            model_config["conv_filters"] = _get_filter_config(obs_space.shape)
+
         TorchModelV2.__init__(self, obs_space, action_space, num_outputs,
                               model_config, name)
         nn.Module.__init__(self)
 
         activation = get_activation_fn(
-            model_config.get("conv_activation"), framework="torch")
-        filters = model_config.get("conv_filters")
-        if not filters:
-            filters = _get_filter_config(obs_space.shape)
-        no_final_linear = model_config.get("no_final_linear")
-        vf_share_layers = model_config.get("vf_share_layers")
+            self.model_config.get("conv_activation"), framework="torch")
+        filters = self.model_config["conv_filters"]
+        no_final_linear = self.model_config.get("no_final_linear")
+        vf_share_layers = self.model_config.get("vf_share_layers")
 
         # Whether the last layer is the output of a Flattened (rather than
         # a n x (1,1) Conv2D).
@@ -152,8 +153,17 @@ def forward(self, input_dict, state, seq_lens):
         if not self.last_layer_is_flattened:
             if self._logits:
                 conv_out = self._logits(conv_out)
+            if conv_out.shape[2] != 1 or conv_out.shape[3] != 1:
+                raise ValueError(
+                    "Given `conv_filters` ({}) do not result in a [B, {} "
+                    "(`num_outputs`), 1, 1] shape (but in {})! Please adjust "
+                    "your Conv2D stack such that the last 2 dims are both "
+                    "1.".format(
+                        self.model_config["conv_filters"], self.num_outputs,
+                        list(conv_out.shape)))
             logits = conv_out.squeeze(3)
             logits = logits.squeeze(2)
+
             return logits, state
         else:
             return conv_out, state


Test patch: