Repository: ray-project/ray

Issue:
[RLlib] Using LSTM model raises ValueError (Tuple obs_space, similar to #3367)

### System Information
- Fedora 7.7 (Maipo) server
- Ray from pip
- Ray version 0.8.0
- Python version 3.6.8
### Problem
I have a custom environment with a Tuple observation space, using the default model in rllib that is trained using PPO.  Everything works when I don't turn on the LSTM option.  However, when I do turn it on, I get the following stack trace:
`
File "/afs/ece.cmu.edu/usr/charlieh/.local/lib/python3.6/site-packages/ray/rllib/agents/trainer_template.py", line 90, in __init__
    Trainer.__init__(self, config, env, logger_creator)
  File "/afs/ece.cmu.edu/usr/charlieh/.local/lib/python3.6/site-packages/ray/rllib/agents/trainer.py", line 398, in __init__
    Trainable.__init__(self, config, logger_creator)
  File "/afs/ece.cmu.edu/usr/charlieh/.local/lib/python3.6/site-packages/ray/tune/trainable.py", line 96, in __init__
    self._setup(copy.deepcopy(self.config))
  File "/afs/ece.cmu.edu/usr/charlieh/.local/lib/python3.6/site-packages/ray/rllib/agents/trainer.py", line 523, in _setup
    self._init(self.config, self.env_creator)
  File "/afs/ece.cmu.edu/usr/charlieh/.local/lib/python3.6/site-packages/ray/rllib/agents/trainer_template.py", line 109, in _init
    self.config["num_workers"])
  File "/afs/ece.cmu.edu/usr/charlieh/.local/lib/python3.6/site-packages/ray/rllib/agents/trainer.py", line 568, in _make_workers
    logdir=self.logdir)
  File "/afs/ece.cmu.edu/usr/charlieh/.local/lib/python3.6/site-packages/ray/rllib/evaluation/worker_set.py", line 64, in __init__
    RolloutWorker, env_creator, policy, 0, self._local_config)
  File "/afs/ece.cmu.edu/usr/charlieh/.local/lib/python3.6/site-packages/ray/rllib/evaluation/worker_set.py", line 220, in _make_worker
    _fake_sampler=config.get("_fake_sampler", False))
  File "/afs/ece.cmu.edu/usr/charlieh/.local/lib/python3.6/site-packages/ray/rllib/evaluation/rollout_worker.py", line 350, in __init__
    self._build_policy_map(policy_dict, policy_config)
  File "/afs/ece.cmu.edu/usr/charlieh/.local/lib/python3.6/site-packages/ray/rllib/evaluation/rollout_worker.py", line 766, in _build_policy_map
    policy_map[name] = cls(obs_space, act_space, merged_conf)
  File "/afs/ece.cmu.edu/usr/charlieh/.local/lib/python3.6/site-packages/ray/rllib/policy/tf_policy_template.py", line 143, in __init__
    obs_include_prev_action_reward=obs_include_prev_action_reward)
  File "/afs/ece.cmu.edu/usr/charlieh/.local/lib/python3.6/site-packages/ray/rllib/policy/dynamic_tf_policy.py", line 198, in __init__
    before_loss_init(self, obs_space, action_space, config)
  File "/afs/ece.cmu.edu/usr/charlieh/.local/lib/python3.6/site-packages/ray/rllib/policy/tf_policy_template.py", line 127, in before_loss_init_wrapper
    self._extra_action_fetches = extra_action_fetches_fn(self)
  File "/afs/ece.cmu.edu/usr/charlieh/.local/lib/python3.6/site-packages/ray/rllib/agents/ppo/ppo_policy.py", line 170, in vf_preds_and_logits_fetches
    SampleBatch.VF_PREDS: policy.model.value_function(),
  File "/afs/ece.cmu.edu/usr/charlieh/.local/lib/python3.6/site-packages/ray/rllib/models/tf/modelv1_compat.py", line 148, in value_function
    seq_lens=None)
  File "/afs/ece.cmu.edu/usr/charlieh/.local/lib/python3.6/site-packages/ray/rllib/models/catalog.py", line 481, in get_model
    seq_lens)
  File "/afs/ece.cmu.edu/usr/charlieh/.local/lib/python3.6/site-packages/ray/rllib/models/catalog.py", line 521, in _get_model
    num_outputs, options)
  File "/afs/ece.cmu.edu/usr/charlieh/.local/lib/python3.6/site-packages/ray/rllib/models/model.py", line 57, in __init__
    input_dict["obs"], obs_space)
  File "/afs/ece.cmu.edu/usr/charlieh/.local/lib/python3.6/site-packages/ray/rllib/models/model.py", line 230, in restore_original_dimensions
    return _unpack_obs(obs, obs_space.original_space, tensorlib=tensorlib)
  File "/afs/ece.cmu.edu/usr/charlieh/.local/lib/python3.6/site-packages/ray/rllib/models/model.py", line 260, in _unpack_obs
    prep.shape[0], obs.shape))
ValueError: Expected flattened obs shape of [None, 143], got (?, 256)
`
My state space is
`obs_space = spaces.Tuple(spaces.Discrete(35),spaces.Discrete(35),spaces.Discrete(35), spaces.Discrete(3), spaces.Discrete(35))`
and if I run 
`prep = get_preprocessor(spy_state_space)(spy_state_space)
print(prep)`
I get

> (143,)

As expected.  I believe that this is similar to issue #3367 (except this time the observation space is a Tuple rather than a Dict.  Am I doing something wrong or is this a bug?  If requested, I can also try posting my code.


Discussion:
Cc @sven1977 
Will try to reproduce. ...
My minimal example works fine. Could you check, to see what I might do differently?
```
import gym
from gym.spaces import Tuple, Discrete
import numpy as np

from ray.rllib.agents.ppo import PPOTrainer
from ray.rllib.utils import try_import_tf

tf = try_import_tf()


class RandomEnv(gym.Env):
    """
    A randomly acting environment that can be instantiated with arbitrary
    action and observation spaces.
    """
    def __init__(self, config):
        # Action space.
        self.action_space = config["action_space"]
        # Observation space from which to sample.
        self.observation_space = config["observation_space"]
        # Reward space from which to sample.
        self.reward_space = config.get(
            "reward_space",
            gym.spaces.Box(low=-1.0, high=1.0, shape=(), dtype=np.float32)
        )
        # Chance that an episode ends at any step.
        self.p_done = config.get("p_done", 0.1)

    def reset(self):
        return self.observation_space.sample()

    def step(self, action):
        return self.observation_space.sample(), float(self.reward_space.sample()), \
            bool(np.random.choice(
                [True, False], p=[self.p_done, 1.0 - self.p_done]
            )), {}


if __name__ == "__main__":
    trainer = PPOTrainer(
        config={
            "model": {
                "use_lstm": True,
            },
            "vf_share_layers": True,
            "num_workers": 0,  # no parallelism
            "env_config": {
                "action_space": Discrete(2),
                # Test a simple Tuple observation space.
                "observation_space": Tuple([Discrete(2), Discrete(2)])
            }
        },
        env=RandomEnv,
    )
    for _ in range(2):
        results = trainer.train()
        print(results)
```
Ah, got it! When I set `vf_share_layers` to False (in my example above), I get the same error as you do. Ok, this seems to be a bug.
So RLlib sets `use_lstm` automatically to False if you do `vf_share_layers == False` (separate policy and vf networks) and outputs a warning: `2020-01-22 13:26:40,266	WARNING modelv1_compat.py:131 -- It is not recommended to use a LSTM model with vf_share_layers=False (consider setting it to True). If you want to not share layers, you can implement a custom LSTM model that overrides the value_function() method.`

I'll fix the error message (it shouldn't be thrown at all), but in the meantime, could you set your config to `vf_share_layers=True` and model->`use_lstm=True`? Then it should work.

Resolving patch:
diff --git a/rllib/examples/random_env.py b/rllib/examples/random_env.py
new file mode 100644
--- /dev/null
+++ b/rllib/examples/random_env.py
@@ -0,0 +1,69 @@
+"""
+Example of a custom gym environment and model. Run this for a demo.
+
+This example shows:
+  - using a custom environment
+  - using a custom model
+  - using Tune for grid search
+
+You can visualize experiment results in ~/ray_results using TensorBoard.
+"""
+
+import gym
+from gym.spaces import Tuple, Discrete
+import numpy as np
+
+from ray.rllib.agents.ppo import PPOTrainer
+from ray.rllib.utils import try_import_tf
+
+tf = try_import_tf()
+
+
+class RandomEnv(gym.Env):
+    """
+    A randomly acting environment that can be instantiated with arbitrary
+    action and observation spaces.
+    """
+
+    def __init__(self, config):
+        # Action space.
+        self.action_space = config["action_space"]
+        # Observation space from which to sample.
+        self.observation_space = config["observation_space"]
+        # Reward space from which to sample.
+        self.reward_space = config.get(
+            "reward_space",
+            gym.spaces.Box(low=-1.0, high=1.0, shape=(), dtype=np.float32))
+        # Chance that an episode ends at any step.
+        self.p_done = config.get("p_done", 0.1)
+
+    def reset(self):
+        return self.observation_space.sample()
+
+    def step(self, action):
+        return self.observation_space.sample(), \
+            float(self.reward_space.sample()), \
+            bool(np.random.choice(
+                [True, False], p=[self.p_done, 1.0 - self.p_done]
+            )), {}
+
+
+if __name__ == "__main__":
+    trainer = PPOTrainer(
+        config={
+            "model": {
+                "use_lstm": True,
+            },
+            "vf_share_layers": False,
+            "num_workers": 0,  # no parallelism
+            "env_config": {
+                "action_space": Discrete(2),
+                # Test a simple Tuple observation space.
+                "observation_space": Tuple([Discrete(3),
+                                            Discrete(2)])
+            }
+        },
+        env=RandomEnv,
+    )
+    results = trainer.train()
+    print(results)
diff --git a/rllib/models/tf/modelv1_compat.py b/rllib/models/tf/modelv1_compat.py
--- a/rllib/models/tf/modelv1_compat.py
+++ b/rllib/models/tf/modelv1_compat.py
@@ -1,3 +1,4 @@
+import copy
 import logging
 import numpy as np
 
@@ -6,7 +7,6 @@
 from ray.rllib.models.tf.misc import linear, normc_initializer
 from ray.rllib.utils.annotations import override
 from ray.rllib.utils import try_import_tf
-from ray.rllib.utils.debug import log_once
 from ray.rllib.utils.tf_ops import scope_vars
 
 tf = try_import_tf()
@@ -124,19 +124,29 @@ def value_function(self):
                     # Create a new separate model with no RNN state, etc.
                     branch_model_config = self.model_config.copy()
                     branch_model_config["free_log_std"] = False
+                    obs_space_vf = self.obs_space
+
                     if branch_model_config["use_lstm"]:
                         branch_model_config["use_lstm"] = False
-                        if log_once("vf_warn"):
-                            logger.warning(
-                                "It is not recommended to use a LSTM model "
-                                "with vf_share_layers=False (consider setting "
-                                "it to True). If you want to not share "
-                                "layers, you can implement a custom LSTM "
-                                "model that overrides the value_function() "
-                                "method.")
+                        logger.warning(
+                            "It is not recommended to use an LSTM model "
+                            "with the `vf_share_layers=False` option. "
+                            "If you want to use separate policy- and vf-"
+                            "networks with LSTMs, you can implement a custom "
+                            "LSTM model that overrides the value_function() "
+                            "method. "
+                            "NOTE: Your policy- and vf-NNs will use the same "
+                            "shared LSTM!")
+                        # Remove original space from obs-space not to trigger
+                        # preprocessing (input to vf-NN is already vectorized
+                        # LSTM output).
+                        obs_space_vf = copy.copy(self.obs_space)
+                        if hasattr(obs_space_vf, "original_space"):
+                            delattr(obs_space_vf, "original_space")
+
                     branch_instance = self.legacy_model_cls(
                         self.cur_instance.input_dict,
-                        self.obs_space,
+                        obs_space_vf,
                         self.action_space,
                         1,
                         branch_model_config,


Test patch:
diff --git a/ci/jenkins_tests/run_rllib_tests.sh b/ci/jenkins_tests/run_rllib_tests.sh
--- a/ci/jenkins_tests/run_rllib_tests.sh
+++ b/ci/jenkins_tests/run_rllib_tests.sh
@@ -487,3 +487,6 @@ docker run --rm --shm-size=${SHM_SIZE} --memory=${MEMORY_SIZE} $DOCKER_SHA \
 
 docker run --rm --shm-size=${SHM_SIZE} --memory=${MEMORY_SIZE} $DOCKER_SHA \
     /ray/ci/suppress_output python /ray/rllib/tests/test_env_with_subprocess.py
+
+docker run --rm --shm-size=${SHM_SIZE} --memory=${MEMORY_SIZE} $DOCKER_SHA \
+    /ray/ci/suppress_output python /ray/rllib/examples/random_env.py