[RLlib] Fix use_lstm flag for ModelV2 (w/o ModelV1 wrapping) and add it for PyTorch. (#8734)

rllib/utils/test_utils.py

@@ -248,7 +248,9 @@ def check_learning_achieved(tune_results, min_reward):
     print("ok")
 
 
-def check_compute_action(trainer, include_prev_action_reward=False):
+def check_compute_action(trainer,
+                         include_state=False,
+                         include_prev_action_reward=False):
     """Tests different combinations of arguments for trainer.compute_action.
 
     Args:
@@ -269,17 +271,27 @@ def check_compute_action(trainer, include_prev_action_reward=False):
     for explore in [True, False]:
         for full_fetch in [True, False]:
             obs = np.clip(obs_space.sample(), -1.0, 1.0)
+            state_in = None
+            if include_state:
+                state_in = pol.model.get_initial_state()
             action_in = action_space.sample() \
                 if include_prev_action_reward else None
             reward_in = 1.0 if include_prev_action_reward else None
-            action = trainer.compute_action(
+            out = trainer.compute_action(
                 obs,
+                state=state_in,
                 prev_action=action_in,
                 prev_reward=reward_in,
                 explore=explore,
                 full_fetch=full_fetch)
-            if full_fetch:
-                action, _, _ = action
+
+            state_out = None
+            if state_in or full_fetch:
+                action, state_out, _ = out
+            if state_out:
+                for si, so in zip(state_in, state_out):
+                    check(list(si.shape), so.shape)
+
             if not action_space.contains(action):
                 raise ValueError(
                     "Returned action ({}) of trainer {} not in Env's "