[RLlib] Fix issue (bug): LSTM + non-shared vf + PPO + tuple actions (#6890)

* Add `RandomEnv` example to examples folder.
Convert warning into Error message when using an LSTM in a non-shared-vf network (after the warning, the program would crash).

* LINT.

* Fix issue #6884. LSTM + non-shared vf NN + PPO crashes when using a Tuple action space.

* LINT

* Change warning message for Model: shared_vf=False, LSTM=True cases.

* Bug fix.

* Add examples/random_env.py test to Jenkins.

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

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)

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.tf_modelv2 import TFModelV2
 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 make_v1_wrapper(legacy_model_cls):
                     # 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,