[RLlib] Add testing Policy.compute_single_action() for all agents. (#8903)

rllib/agents/a3c/tests/test_a2c.py

@@ -2,7 +2,8 @@ import unittest
 
 import ray
 import ray.rllib.agents.a3c as a3c
-from ray.rllib.utils.test_utils import check_compute_action, framework_iterator
+from ray.rllib.utils.test_utils import check_compute_single_action, \
+    framework_iterator
 
 
 class TestA2C(unittest.TestCase):
@@ -30,14 +31,14 @@ class TestA2C(unittest.TestCase):
                 for i in range(num_iterations):
                     results = trainer.train()
                     print(results)
-                check_compute_action(trainer)
+                check_compute_single_action(trainer)
 
     def test_a2c_exec_impl(ray_start_regular):
         config = {"min_iter_time_s": 0}
         for _ in framework_iterator(config, ("tf", "torch")):
             trainer = a3c.A2CTrainer(env="CartPole-v0", config=config)
             assert isinstance(trainer.train(), dict)
-            check_compute_action(trainer)
+            check_compute_single_action(trainer)
 
     def test_a2c_exec_impl_microbatch(ray_start_regular):
         config = {
@@ -47,7 +48,7 @@ class TestA2C(unittest.TestCase):
         for _ in framework_iterator(config, ("tf", "torch")):
             trainer = a3c.A2CTrainer(env="CartPole-v0", config=config)
             assert isinstance(trainer.train(), dict)
-            check_compute_action(trainer)
+            check_compute_single_action(trainer)
 
 
 if __name__ == "__main__":

rllib/agents/a3c/tests/test_a3c.py

@@ -2,7 +2,8 @@ import unittest
 
 import ray
 import ray.rllib.agents.a3c as a3c
-from ray.rllib.utils.test_utils import check_compute_action, framework_iterator
+from ray.rllib.utils.test_utils import check_compute_single_action, \
+    framework_iterator
 
 
 class TestA3C(unittest.TestCase):
@@ -30,7 +31,7 @@ class TestA3C(unittest.TestCase):
                 for i in range(num_iterations):
                     results = trainer.train()
                     print(results)
-                check_compute_action(trainer)
+                check_compute_single_action(trainer)
 
 
 if __name__ == "__main__":

rllib/agents/ars/ars_tf_policy.py

@@ -54,7 +54,14 @@ class ARSTFPolicy:
             for _, variable in self.variables.variables.items())
         self.sess.run(tf.global_variables_initializer())
 
-    def compute_actions(self, observation, add_noise=False, update=True):
+    def compute_actions(self,
+                        observation,
+                        add_noise=False,
+                        update=True,
+                        **kwargs):
+        # Batch is given as list of one.
+        if isinstance(observation, list) and len(observation) == 1:
+            observation = observation[0]
         observation = self.preprocessor.transform(observation)
         observation = self.observation_filter(observation[None], update=update)
         action = self.sess.run(
@@ -64,6 +71,15 @@ class ARSTFPolicy:
             action += np.random.randn(*action.shape) * self.action_noise_std
         return action
 
+    def compute_single_action(self,
+                              observation,
+                              add_noise=False,
+                              update=True,
+                              **kwargs):
+        action = self.compute_actions(
+            [observation], add_noise=add_noise, update=update, **kwargs)
+        return action[0], [], {}
+
     def get_state(self):
         return {"state": self.get_flat_weights()}
 

rllib/agents/ars/tests/test_ars.py

@@ -2,7 +2,8 @@ import unittest
 
 import ray
 import ray.rllib.agents.ars as ars
-from ray.rllib.utils.test_utils import framework_iterator, check_compute_action
+from ray.rllib.utils.test_utils import framework_iterator, \
+    check_compute_single_action
 
 
 class TestARS(unittest.TestCase):
@@ -16,14 +17,14 @@ class TestARS(unittest.TestCase):
 
         num_iterations = 2
 
-        for _ in framework_iterator(config, ("torch", "tf")):
+        for _ in framework_iterator(config, ("tf", "torch")):
             plain_config = config.copy()
             trainer = ars.ARSTrainer(config=plain_config, env="CartPole-v0")
             for i in range(num_iterations):
                 results = trainer.train()
                 print(results)
 
-            check_compute_action(trainer)
+            check_compute_single_action(trainer)
 
 
 if __name__ == "__main__":

rllib/agents/ddpg/tests/test_apex_ddpg.py

@@ -3,8 +3,8 @@ import unittest
 
 import ray
 import ray.rllib.agents.ddpg.apex as apex_ddpg
-from ray.rllib.utils.test_utils import check, framework_iterator, \
-    check_compute_action
+from ray.rllib.utils.test_utils import check, check_compute_single_action, \
+    framework_iterator
 
 
 class TestApexDDPG(unittest.TestCase):
@@ -41,7 +41,7 @@ class TestApexDDPG(unittest.TestCase):
 
             for _ in range(num_iterations):
                 print(trainer.train())
-            check_compute_action(trainer)
+            check_compute_single_action(trainer)
 
             # Test again per-worker scale distribution
             # (should not have changed).

rllib/agents/ddpg/tests/test_ddpg.py

@@ -11,8 +11,8 @@ from ray.rllib.execution.replay_buffer import LocalReplayBuffer
 from ray.rllib.policy.sample_batch import SampleBatch
 from ray.rllib.utils.framework import try_import_tf, try_import_torch
 from ray.rllib.utils.numpy import fc, huber_loss, l2_loss, relu, sigmoid
-from ray.rllib.utils.test_utils import check, framework_iterator, \
-    check_compute_action
+from ray.rllib.utils.test_utils import check, check_compute_single_action, \
+    framework_iterator
 from ray.rllib.utils.torch_ops import convert_to_torch_tensor
 
 tf = try_import_tf()
@@ -44,7 +44,7 @@ class TestDDPG(unittest.TestCase):
             for i in range(num_iterations):
                 results = trainer.train()
                 print(results)
-            check_compute_action(trainer)
+            check_compute_single_action(trainer)
 
     def test_ddpg_exploration_and_with_random_prerun(self):
         """Tests DDPG's Exploration (w/ random actions for n timesteps)."""

rllib/agents/ddpg/tests/test_td3.py

@@ -3,8 +3,8 @@ import unittest
 
 import ray.rllib.agents.ddpg.td3 as td3
 from ray.rllib.utils.framework import try_import_tf
-from ray.rllib.utils.test_utils import check, framework_iterator, \
-    check_compute_action
+from ray.rllib.utils.test_utils import check, check_compute_single_action, \
+    framework_iterator
 
 tf = try_import_tf()
 
@@ -22,7 +22,7 @@ class TestTD3(unittest.TestCase):
             for i in range(num_iterations):
                 results = trainer.train()
                 print(results)
-            check_compute_action(trainer)
+            check_compute_single_action(trainer)
 
     def test_td3_exploration_and_with_random_prerun(self):
         """Tests TD3's Exploration (w/ random actions for n timesteps)."""

rllib/agents/dqn/tests/test_apex_dqn.py

@@ -3,8 +3,8 @@ import unittest
 
 import ray
 import ray.rllib.agents.dqn.apex as apex
-from ray.rllib.utils.test_utils import check, framework_iterator, \
-    check_compute_action
+from ray.rllib.utils.test_utils import check, check_compute_single_action, \
+    framework_iterator
 
 
 class TestApexDQN(unittest.TestCase):
@@ -47,7 +47,7 @@ class TestApexDQN(unittest.TestCase):
             expected = [0.4, 0.016190862, 0.00065536]
             check([i["cur_epsilon"] for i in infos], [0.0] + expected)
 
-            check_compute_action(trainer)
+            check_compute_single_action(trainer)
 
             # TODO(ekl) fix iterator metrics bugs w/multiple trainers.
             #            for i in range(1):

rllib/agents/dqn/tests/test_dqn.py

@@ -4,8 +4,8 @@ import unittest
 import ray
 import ray.rllib.agents.dqn as dqn
 from ray.rllib.utils.framework import try_import_tf
-from ray.rllib.utils.test_utils import check, framework_iterator, \
-    check_compute_action
+from ray.rllib.utils.test_utils import check, check_compute_single_action, \
+    framework_iterator
 
 tf = try_import_tf()
 
@@ -33,7 +33,7 @@ class TestDQN(unittest.TestCase):
                 results = trainer.train()
                 print(results)
 
-            check_compute_action(trainer)
+            check_compute_single_action(trainer)
 
             # Rainbow.
             # TODO(sven): Add torch once DQN-torch supports distributional-Q.
@@ -50,7 +50,7 @@ class TestDQN(unittest.TestCase):
                 results = trainer.train()
                 print(results)
 
-            check_compute_action(trainer)
+            check_compute_single_action(trainer)
 
     def test_dqn_exploration_and_soft_q_config(self):
         """Tests, whether a DQN Agent outputs exploration/softmaxed actions."""

rllib/agents/dqn/tests/test_simple_q.py

@@ -8,8 +8,8 @@ from ray.rllib.agents.dqn.simple_q_torch_policy import build_q_losses as \
 from ray.rllib.policy.sample_batch import SampleBatch
 from ray.rllib.utils.framework import try_import_tf
 from ray.rllib.utils.numpy import fc, one_hot, huber_loss
-from ray.rllib.utils.test_utils import check, framework_iterator, \
-    check_compute_action
+from ray.rllib.utils.test_utils import check, check_compute_single_action, \
+    framework_iterator
 
 tf = try_import_tf()
 
@@ -27,7 +27,7 @@ class TestSimpleQ(unittest.TestCase):
                 results = trainer.train()
                 print(results)
 
-            check_compute_action(trainer)
+            check_compute_single_action(trainer)
 
     def test_simple_q_loss_function(self):
         """Tests the Simple-Q loss function results on all frameworks."""

rllib/agents/es/es_tf_policy.py

@@ -100,7 +100,14 @@ class ESTFPolicy:
             for _, variable in self.variables.variables.items())
         self.sess.run(tf.global_variables_initializer())
 
-    def compute_actions(self, observation, add_noise=False, update=True):
+    def compute_actions(self,
+                        observation,
+                        add_noise=False,
+                        update=True,
+                        **kwargs):
+        # Batch is given as list of one.
+        if isinstance(observation, list) and len(observation) == 1:
+            observation = observation[0]
         observation = self.preprocessor.transform(observation)
         observation = self.observation_filter(observation[None], update=update)
         # `actions` is a list of (component) batches.
@@ -114,6 +121,15 @@ class ESTFPolicy:
         actions = unbatch(actions)
         return actions
 
+    def compute_single_action(self,
+                              observation,
+                              add_noise=False,
+                              update=True,
+                              **kwargs):
+        action = self.compute_actions(
+            [observation], add_noise=add_noise, update=update, **kwargs)
+        return action[0], [], {}
+
     def _add_noise(self, single_action, single_action_space):
         if isinstance(single_action_space, gym.spaces.Box):
             single_action += np.random.randn(*single_action.shape) * \

rllib/agents/es/es_torch_policy.py

@@ -54,7 +54,14 @@ def before_init(policy, observation_space, action_space, config):
     type(policy).set_flat_weights = _set_flat_weights
     type(policy).get_flat_weights = _get_flat_weights
 
-    def _compute_actions(policy, obs_batch, add_noise=False, update=True):
+    def _compute_actions(policy,
+                         obs_batch,
+                         add_noise=False,
+                         update=True,
+                         **kwargs):
+        # Batch is given as list -> Try converting to numpy first.
+        if isinstance(obs_batch, list) and len(obs_batch) == 1:
+            obs_batch = obs_batch[0]
         observation = policy.preprocessor.transform(obs_batch)
         observation = policy.observation_filter(
             observation[None], update=update)

rllib/agents/es/tests/test_es.py

@@ -2,7 +2,8 @@ import unittest
 
 import ray
 import ray.rllib.agents.es as es
-from ray.rllib.utils.test_utils import framework_iterator, check_compute_action
+from ray.rllib.utils.test_utils import check_compute_single_action, \
+    framework_iterator
 
 
 class TestES(unittest.TestCase):
@@ -17,14 +18,14 @@ class TestES(unittest.TestCase):
 
         num_iterations = 2
 
-        for _ in framework_iterator(config, ("torch", "tf")):
+        for _ in framework_iterator(config, ("tf", "torch")):
             plain_config = config.copy()
             trainer = es.ESTrainer(config=plain_config, env="CartPole-v0")
             for i in range(num_iterations):
                 results = trainer.train()
                 print(results)
 
-            check_compute_action(trainer)
+            check_compute_single_action(trainer)
 
 
 if __name__ == "__main__":

rllib/agents/impala/tests/test_impala.py

@@ -3,7 +3,8 @@ import unittest
 import ray
 import ray.rllib.agents.impala as impala
 from ray.rllib.utils.framework import try_import_tf
-from ray.rllib.utils.test_utils import framework_iterator, check_compute_action
+from ray.rllib.utils.test_utils import check_compute_single_action, \
+    framework_iterator
 
 tf = try_import_tf()
 
@@ -11,7 +12,7 @@ tf = try_import_tf()
 class TestIMPALA(unittest.TestCase):
     @classmethod
     def setUpClass(cls):
-        ray.init()
+        ray.init(local_mode=True)
 
     @classmethod
     def tearDownClass(cls):
@@ -22,7 +23,7 @@ class TestIMPALA(unittest.TestCase):
         config = impala.DEFAULT_CONFIG.copy()
         num_iterations = 1
 
-        for _ in framework_iterator(config, frameworks=("torch", "tf")):
+        for _ in framework_iterator(config, frameworks=("tf", "torch")):
             local_cfg = config.copy()
             for env in ["Pendulum-v0", "CartPole-v0"]:
                 print("Env={}".format(env))
@@ -33,7 +34,7 @@ class TestIMPALA(unittest.TestCase):
                 trainer = impala.ImpalaTrainer(config=local_cfg, env=env)
                 for i in range(num_iterations):
                     print(trainer.train())
-                check_compute_action(trainer)
+                check_compute_single_action(trainer)
                 trainer.stop()
 
                 # Test w/ LSTM.
@@ -43,7 +44,7 @@ class TestIMPALA(unittest.TestCase):
                 trainer = impala.ImpalaTrainer(config=local_cfg, env=env)
                 for i in range(num_iterations):
                     print(trainer.train())
-                check_compute_action(trainer, include_state=True)
+                check_compute_single_action(trainer, include_state=True)
                 trainer.stop()
 
 

rllib/agents/marwil/tests/test_marwil.py

@@ -3,7 +3,8 @@ import unittest
 import ray
 import ray.rllib.agents.marwil as marwil
 from ray.rllib.utils.framework import try_import_tf
-from ray.rllib.utils.test_utils import framework_iterator, check_compute_action
+from ray.rllib.utils.test_utils import check_compute_single_action, \
+    framework_iterator
 
 tf = try_import_tf()
 
@@ -28,7 +29,8 @@ class TestMARWIL(unittest.TestCase):
             trainer = marwil.MARWILTrainer(config=config, env="CartPole-v0")
             for i in range(num_iterations):
                 trainer.train()
-            check_compute_action(trainer, include_prev_action_reward=True)
+            check_compute_single_action(
+                trainer, include_prev_action_reward=True)
             trainer.stop()
 
 

rllib/agents/pg/tests/test_pg.py

@@ -7,7 +7,8 @@ from ray.rllib.evaluation.postprocessing import Postprocessing
 from ray.rllib.models.tf.tf_action_dist import Categorical
 from ray.rllib.models.torch.torch_action_dist import TorchCategorical
 from ray.rllib.policy.sample_batch import SampleBatch
-from ray.rllib.utils import check, fc, framework_iterator, check_compute_action
+from ray.rllib.utils import check, check_compute_single_action, fc, \
+    framework_iterator
 
 
 class TestPG(unittest.TestCase):
@@ -27,7 +28,8 @@ class TestPG(unittest.TestCase):
             trainer = pg.PGTrainer(config=config, env="CartPole-v0")
             for i in range(num_iterations):
                 trainer.train()
-            check_compute_action(trainer, include_prev_action_reward=True)
+            check_compute_single_action(
+                trainer, include_prev_action_reward=True)
 
     def test_pg_loss_functions(self):
         """Tests the PG loss function math."""

rllib/agents/ppo/tests/test_appo.py

@@ -3,7 +3,8 @@ import unittest
 import ray
 import ray.rllib.agents.ppo as ppo
 from ray.rllib.utils.framework import try_import_tf
-from ray.rllib.utils.test_utils import framework_iterator, check_compute_action
+from ray.rllib.utils.test_utils import check_compute_single_action, \
+    framework_iterator
 
 tf = try_import_tf()
 
@@ -28,14 +29,14 @@ class TestAPPO(unittest.TestCase):
             trainer = ppo.APPOTrainer(config=_config, env="CartPole-v0")
             for i in range(num_iterations):
                 print(trainer.train())
-            check_compute_action(trainer)
+            check_compute_single_action(trainer)
 
             _config = config.copy()
             _config["vtrace"] = True
             trainer = ppo.APPOTrainer(config=_config, env="CartPole-v0")
             for i in range(num_iterations):
                 print(trainer.train())
-            check_compute_action(trainer)
+            check_compute_single_action(trainer)
 
 
 if __name__ == "__main__":

rllib/agents/ppo/tests/test_ddppo.py

@@ -3,7 +3,8 @@ import unittest
 import ray
 import ray.rllib.agents.ppo as ppo
 from ray.rllib.utils.framework import try_import_tf
-from ray.rllib.utils.test_utils import framework_iterator, check_compute_action
+from ray.rllib.utils.test_utils import check_compute_single_action, \
+    framework_iterator
 
 tf = try_import_tf()
 
@@ -27,7 +28,7 @@ class TestDDPPO(unittest.TestCase):
             trainer = ppo.ddppo.DDPPOTrainer(config=config, env="CartPole-v0")
             for i in range(num_iterations):
                 trainer.train()
-            check_compute_action(trainer)
+            check_compute_single_action(trainer)
             trainer.stop()
 
 

rllib/agents/ppo/tests/test_ppo.py

@@ -16,7 +16,7 @@ from ray.rllib.policy.sample_batch import SampleBatch
 from ray.rllib.utils.framework import try_import_tf
 from ray.rllib.utils.numpy import fc
 from ray.rllib.utils.test_utils import check, framework_iterator, \
-    check_compute_action
+    check_compute_single_action
 
 tf = try_import_tf()
 
@@ -56,7 +56,8 @@ class TestPPO(unittest.TestCase):
             trainer = ppo.PPOTrainer(config=config, env="CartPole-v0")
             for i in range(num_iterations):
                 trainer.train()
-            check_compute_action(trainer, include_prev_action_reward=True)
+            check_compute_single_action(
+                trainer, include_prev_action_reward=True)
 
     def test_ppo_fake_multi_gpu_learning(self):
         """Test whether PPOTrainer can learn CartPole w/ faked multi-GPU."""

rllib/agents/sac/tests/test_sac.py

@@ -13,8 +13,8 @@ from ray.rllib.execution.replay_buffer import LocalReplayBuffer
 from ray.rllib.policy.sample_batch import SampleBatch
 from ray.rllib.utils.framework import try_import_tf, try_import_torch
 from ray.rllib.utils.numpy import fc, relu
-from ray.rllib.utils.test_utils import check, framework_iterator, \
-    check_compute_action
+from ray.rllib.utils.test_utils import check, check_compute_single_action, \
+    framework_iterator
 from ray.rllib.utils.torch_ops import convert_to_torch_tensor
 
 tf = try_import_tf()
@@ -67,7 +67,7 @@ class TestSAC(unittest.TestCase):
                 for i in range(num_iterations):
                     results = trainer.train()
                     print(results)
-                check_compute_action(trainer)
+                check_compute_single_action(trainer)
 
     def test_sac_loss_function(self):
         """Tests SAC loss function results across all frameworks."""

rllib/policy/policy.py

@@ -165,7 +165,7 @@ class Policy(metaclass=ABCMeta):
                 for s in state
             ]
 
-        batched_action, state_out, info = self.compute_actions(
+        out = self.compute_actions(
             [obs],
             state_batch,
             prev_action_batch=prev_action_batch,
@@ -175,7 +175,16 @@ class Policy(metaclass=ABCMeta):
             explore=explore,
             timestep=timestep)
 
-        single_action = unbatch(batched_action)
+        # Some policies don't return a tuple, but always just a single action.
+        # E.g. ES and ARS.
+        if not isinstance(out, tuple):
+            single_action = out
+            state_out = []
+            info = {}
+        # Normal case: Policy should return (action, state, info) tuple.
+        else:
+            batched_action, state_out, info = out
+            single_action = unbatch(batched_action)
         assert len(single_action) == 1
         single_action = single_action[0]
 

rllib/tests/test_rollout.py

@@ -66,9 +66,6 @@ class TestRollout(unittest.TestCase):
     def test_a3c(self):
         rollout_test("A3C")
 
-    def test_ars(self):
-        rollout_test("ARS")
-
     def test_ddpg(self):
         rollout_test("DDPG", env="Pendulum-v0")
 

rllib/utils/__init__.py

@@ -13,8 +13,8 @@ from ray.rllib.utils.policy_client import PolicyClient
 from ray.rllib.utils.policy_server import PolicyServer
 from ray.rllib.utils.schedules import LinearSchedule, PiecewiseSchedule, \
     PolynomialSchedule, ExponentialSchedule, ConstantSchedule
-from ray.rllib.utils.test_utils import check, framework_iterator, \
-    check_compute_action
+from ray.rllib.utils.test_utils import check, check_compute_single_action, \
+    framework_iterator
 from ray.tune.utils import merge_dicts, deep_update
 
 
@@ -71,7 +71,7 @@ def try_import_tree():
 __all__ = [
     "add_mixins",
     "check",
-    "check_compute_action",
+    "check_compute_single_action",
     "deprecation_warning",
     "fc",
     "force_list",

rllib/utils/test_utils.py

@@ -28,9 +28,9 @@ def framework_iterator(config=None,
         config (Optional[dict]): An optional config dict to alter in place
             depending on the iteration.
         frameworks (Tuple[str]): A list/tuple of the frameworks to be tested.
-            Allowed are: "tf", "tfe", and "torch".
-        session (bool): If True, enter a tf.Session() and yield that as
-            well in the tf-case (otherwise, yield (fw, None)).
+            Allowed are: "tf", "tfe", "torch", and None.
+        session (bool): If True and only in the tf-case: Enter a tf.Session()
+            and yield that as second return value (otherwise yield (fw, None)).
 
     Yields:
         str: If enter_session is False:
@@ -95,7 +95,7 @@ def check(x, y, decimals=5, atol=None, rtol=None, false=False):
         x (any): The value to be compared (to the expectation: `y`). This
             may be a Tensor.
         y (any): The expected value to be compared to `x`. This must not
-            be a Tensor.
+            be a tf-Tensor, but may be a tfe/torch-Tensor.
         decimals (int): The number of digits after the floating point up to
             which all numeric values have to match.
         atol (float): Absolute tolerance of the difference between x and y
@@ -244,13 +244,13 @@ def check_learning_achieved(tune_results, min_reward):
     print("ok")
 
 
-def check_compute_action(trainer,
-                         include_state=False,
-                         include_prev_action_reward=False):
+def check_compute_single_action(trainer,
+                                include_state=False,
+                                include_prev_action_reward=False):
     """Tests different combinations of arguments for trainer.compute_action.
 
     Args:
-        trainer (Trainer): The trainer object to test.
+        trainer (Trainer): The Trainer object to test.
         include_prev_action_reward (bool): Whether to include the prev-action
             and -reward in the `compute_action` call.
 
@@ -264,31 +264,44 @@ def check_compute_action(trainer,
 
     obs_space = pol.observation_space
     action_space = pol.action_space
-    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
-            out = trainer.compute_action(
-                obs,
-                state=state_in,
-                prev_action=action_in,
-                prev_reward=reward_in,
-                explore=explore,
-                full_fetch=full_fetch)
 
-            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)
+    for what in [pol, trainer]:
+        print("what={}".format(what))
+        method_to_test = trainer.compute_action if what is trainer else \
+            pol.compute_single_action
 
-            if not action_space.contains(action):
-                raise ValueError(
-                    "Returned action ({}) of trainer {} not in Env's "
-                    "action_space ({})!".format(action, trainer, action_space))
+        for explore in [True, False]:
+            print("explore={}".format(explore))
+            for full_fetch in ([False, True] if what is trainer else [False]):
+                print("full-fetch={}".format(full_fetch))
+                call_kwargs = {}
+                if what is trainer:
+                    call_kwargs["full_fetch"] = full_fetch
+
+                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 = method_to_test(
+                    obs,
+                    state_in,
+                    prev_action=action_in,
+                    prev_reward=reward_in,
+                    explore=explore,
+                    **call_kwargs)
+
+                state_out = None
+                if state_in or full_fetch or what is pol:
+                    action, state_out, _ = action
+                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/policy {} not in "
+                        "Env's action_space "
+                        "({})!".format(action, what, action_space))