[RLlib] Minor fixes (torch GPU bugs + some cleanup). (#11609)

python/ray/cloudpickle/cloudpickle_fast.py

@@ -23,7 +23,7 @@ import typing
 from .cloudpickle import (
     _is_dynamic, _extract_code_globals, _BUILTIN_TYPE_NAMES, DEFAULT_PROTOCOL,
     _find_imported_submodules, _get_cell_contents, _is_importable_by_name, _builtin_type,
-    Enum, _get_or_create_tracker_id,  _make_skeleton_class, _make_skeleton_enum,
+    Enum, _get_or_create_tracker_id, _make_skeleton_class, _make_skeleton_enum,
     _extract_class_dict, dynamic_subimport, subimport, _typevar_reduce, _get_bases,
     cell_set, _make_empty_cell,
 )

rllib/agents/dqn/dqn_torch_model.py

@@ -142,7 +142,9 @@ class DQNTorchModel(TorchModelV2, nn.Module):
         if self.num_atoms > 1:
             # Distributional Q-learning uses a discrete support z
             # to represent the action value distribution
-            z = torch.range(0.0, self.num_atoms - 1, dtype=torch.float32)
+            z = torch.range(
+                0.0, self.num_atoms - 1,
+                dtype=torch.float32).to(action_scores.device)
             z = self.v_min + \
                 z * (self.v_max - self.v_min) / float(self.num_atoms - 1)
 

rllib/agents/dqn/dqn_torch_policy.py

@@ -48,7 +48,8 @@ class QLoss:
 
         if num_atoms > 1:
             # Distributional Q-learning which corresponds to an entropy loss
-            z = torch.range(0.0, num_atoms - 1, dtype=torch.float32)
+            z = torch.range(
+                0.0, num_atoms - 1, dtype=torch.float32).to(rewards.device)
             z = v_min + z * (v_max - v_min) / float(num_atoms - 1)
 
             # (batch_size, 1) * (1, num_atoms) = (batch_size, num_atoms)
@@ -321,7 +322,7 @@ def after_init(policy: Policy, obs_space: gym.spaces.Space,
                config: TrainerConfigDict) -> None:
     ComputeTDErrorMixin.__init__(policy)
     TargetNetworkMixin.__init__(policy, obs_space, action_space, config)
-    # Move target net to device (this is done autoatically for the
+    # Move target net to device (this is done automatically for the
     # policy.model, but not for any other models the policy has).
     policy.target_q_model = policy.target_q_model.to(policy.device)
 

rllib/agents/impala/tests/test_impala.py

@@ -18,25 +18,6 @@ class TestIMPALA(unittest.TestCase):
     def tearDownClass(cls) -> None:
         ray.shutdown()
 
-    def test_impala_lr_schedule(self):
-        config = impala.DEFAULT_CONFIG.copy()
-        config["lr_schedule"] = [
-            [0, 0.0005],
-            [10000, 0.000001],
-        ]
-        local_cfg = config.copy()
-        trainer = impala.ImpalaTrainer(config=local_cfg, env="CartPole-v0")
-
-        def get_lr(result):
-            return result["info"]["learner"]["default_policy"]["cur_lr"]
-
-        try:
-            r1 = trainer.train()
-            r2 = trainer.train()
-            assert get_lr(r2) < get_lr(r1), (r1, r2)
-        finally:
-            trainer.stop()
-
     def test_impala_compilation(self):
         """Test whether an ImpalaTrainer can be built with both frameworks."""
         config = impala.DEFAULT_CONFIG.copy()
@@ -70,6 +51,25 @@ class TestIMPALA(unittest.TestCase):
                     include_prev_action_reward=True)
                 trainer.stop()
 
+    def test_impala_lr_schedule(self):
+        config = impala.DEFAULT_CONFIG.copy()
+        config["lr_schedule"] = [
+            [0, 0.0005],
+            [10000, 0.000001],
+        ]
+        local_cfg = config.copy()
+        trainer = impala.ImpalaTrainer(config=local_cfg, env="CartPole-v0")
+
+        def get_lr(result):
+            return result["info"]["learner"]["default_policy"]["cur_lr"]
+
+        try:
+            r1 = trainer.train()
+            r2 = trainer.train()
+            assert get_lr(r2) < get_lr(r1), (r1, r2)
+        finally:
+            trainer.stop()
+
 
 if __name__ == "__main__":
     import pytest

rllib/agents/impala/vtrace_tf_policy.py

@@ -80,7 +80,8 @@ class VTraceLoss:
                 behaviour_policy_logits=behaviour_logits,
                 target_policy_logits=target_logits,
                 actions=tf.unstack(actions, axis=2),
-                discounts=tf.cast(~dones, tf.float32) * discount,
+                discounts=tf.cast(~tf.cast(dones, tf.bool), tf.float32) *
+                discount,
                 rewards=rewards,
                 values=values,
                 bootstrap_value=bootstrap_value,
@@ -154,8 +155,7 @@ def build_vtrace_loss(policy, model, dist_class, train_batch):
     if isinstance(policy.action_space, gym.spaces.Discrete):
         is_multidiscrete = False
         output_hidden_shape = [policy.action_space.n]
-    elif isinstance(policy.action_space,
-                    gym.spaces.multi_discrete.MultiDiscrete):
+    elif isinstance(policy.action_space, gym.spaces.MultiDiscrete):
         is_multidiscrete = True
         output_hidden_shape = policy.action_space.nvec.astype(np.int32)
     else:

rllib/agents/impala/vtrace_torch_policy.py

@@ -118,8 +118,7 @@ def build_vtrace_loss(policy, model, dist_class, train_batch):
     if isinstance(policy.action_space, gym.spaces.Discrete):
         is_multidiscrete = False
         output_hidden_shape = [policy.action_space.n]
-    elif isinstance(policy.action_space,
-                    gym.spaces.multi_discrete.MultiDiscrete):
+    elif isinstance(policy.action_space, gym.spaces.MultiDiscrete):
         is_multidiscrete = True
         output_hidden_shape = policy.action_space.nvec.astype(np.int32)
     else:

rllib/agents/pg/tests/test_pg.py

@@ -105,14 +105,15 @@ class TestPG(unittest.TestCase):
                     framework=fw)
             expected_logp = dist_cls(expected_logits, policy.model).logp(
                 train_batch[SampleBatch.ACTIONS])
+            adv = train_batch[Postprocessing.ADVANTAGES]
             if sess:
                 expected_logp = sess.run(expected_logp)
+            elif fw == "torch":
+                expected_logp = expected_logp.detach().cpu().numpy()
+                adv = adv.detach().cpu().numpy()
             else:
                 expected_logp = expected_logp.numpy()
-            expected_loss = -np.mean(
-                expected_logp *
-                (train_batch[Postprocessing.ADVANTAGES] if fw != "torch" else
-                 train_batch[Postprocessing.ADVANTAGES].numpy()))
+            expected_loss = -np.mean(expected_logp * adv)
             check(results, expected_loss, decimals=4)
 
 

rllib/agents/ppo/appo_tf_policy.py

@@ -170,8 +170,9 @@ def appo_surrogate_loss(
                     unpacked_old_policy_behaviour_logits, drop_last=True),
                 actions=tf.unstack(
                     make_time_major(loss_actions, drop_last=True), axis=2),
-                discounts=tf.cast(~make_time_major(dones, drop_last=True),
-                                  tf.float32) * policy.config["gamma"],
+                discounts=tf.cast(
+                    ~make_time_major(tf.cast(dones, tf.bool), drop_last=True),
+                    tf.float32) * policy.config["gamma"],
                 rewards=make_time_major(rewards, drop_last=True),
                 values=values_time_major[:-1],  # drop-last=True
                 bootstrap_value=values_time_major[-1],

rllib/agents/ppo/appo_torch_policy.py

@@ -114,11 +114,11 @@ def appo_surrogate_loss(policy: Policy, model: ModelV2,
             unpacked_old_policy_behaviour_logits = torch.chunk(
                 old_policy_behaviour_logits, output_hidden_shape, dim=1)
 
-        # Prepare actions for loss
+        # Prepare actions for loss.
         loss_actions = actions if is_multidiscrete else torch.unsqueeze(
             actions, dim=1)
 
-        # Prepare KL for Loss
+        # Prepare KL for loss.
         action_kl = _make_time_major(
             old_policy_action_dist.kl(action_dist), drop_last=True)
 
@@ -152,7 +152,7 @@ def appo_surrogate_loss(policy: Policy, model: ModelV2,
         logp_ratio = is_ratio * torch.exp(actions_logp - prev_actions_logp)
         policy._is_ratio = is_ratio
 
-        advantages = vtrace_returns.pg_advantages
+        advantages = vtrace_returns.pg_advantages.to(policy.device)
         surrogate_loss = torch.min(
             advantages * logp_ratio,
             advantages *
@@ -163,8 +163,8 @@ def appo_surrogate_loss(policy: Policy, model: ModelV2,
         mean_policy_loss = -reduce_mean_valid(surrogate_loss)
 
         # The value function loss.
-        delta = values_time_major[:-1] - vtrace_returns.vs
-        value_targets = vtrace_returns.vs
+        value_targets = vtrace_returns.vs.to(policy.device)
+        delta = values_time_major[:-1] - value_targets
         mean_vf_loss = 0.5 * reduce_mean_valid(torch.pow(delta, 2.0))
 
         # The entropy loss.
@@ -315,6 +315,9 @@ def setup_late_mixins(policy: Policy, obs_space: gym.spaces.Space,
     KLCoeffMixin.__init__(policy, config)
     ValueNetworkMixin.__init__(policy, obs_space, action_space, config)
     TargetNetworkMixin.__init__(policy, obs_space, action_space, config)
+    # Move target net to device (this is done automatically for the
+    # policy.model, but not for any other models the policy has).
+    policy.target_model = policy.target_model.to(policy.device)
 
 
 # Build a child class of `TorchPolicy`, given the custom functions defined

rllib/agents/ppo/tests/test_appo.py

@@ -22,6 +22,7 @@ class TestAPPO(unittest.TestCase):
         num_iterations = 2
 
         for _ in framework_iterator(config):
+            print("w/o v-trace")
             _config = config.copy()
             trainer = ppo.APPOTrainer(config=_config, env="CartPole-v0")
             for i in range(num_iterations):
@@ -29,6 +30,7 @@ class TestAPPO(unittest.TestCase):
             check_compute_single_action(trainer)
             trainer.stop()
 
+            print("w/ v-trace")
             _config = config.copy()
             _config["vtrace"] = True
             trainer = ppo.APPOTrainer(config=_config, env="CartPole-v0")

rllib/agents/qmix/qmix_policy.py

@@ -143,7 +143,7 @@ class QMixLoss(nn.Module):
         return loss, mask, masked_td_error, chosen_action_qvals, targets
 
 
-# TODO(sven): Make this a TorchPolicy child.
+# TODO(sven): Make this a TorchPolicy child via `build_torch_policy`.
 class QMixTorchPolicy(Policy):
     """QMix impl. Assumes homogeneous agents for now.
 

rllib/agents/sac/tests/test_sac.py

@@ -4,6 +4,7 @@ import numpy as np
 import re
 import unittest
 
+import ray
 import ray.rllib.agents.sac as sac
 from ray.rllib.agents.sac.sac_tf_policy import sac_actor_critic_loss as tf_loss
 from ray.rllib.agents.sac.sac_torch_policy import actor_critic_loss as \
@@ -45,6 +46,14 @@ class SimpleEnv(Env):
 
 
 class TestSAC(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls) -> None:
+        ray.init()
+
+    @classmethod
+    def tearDownClass(cls) -> None:
+        ray.shutdown()
+
     def test_sac_compilation(self):
         """Tests whether an SACTrainer can be built with all frameworks."""
         config = sac.DEFAULT_CONFIG.copy()

rllib/contrib/maddpg/maddpg.py

@@ -58,9 +58,9 @@ DEFAULT_CONFIG = with_common_config({
     "critic_hidden_activation": "relu",
     # N-step Q learning
     "n_step": 1,
-    # Algorithm for good policies
+    # Algorithm for good policies.
     "good_policy": "maddpg",
-    # Algorithm for adversary policies
+    # Algorithm for adversary policies.
     "adv_policy": "maddpg",
 
     # === Replay buffer ===

rllib/contrib/maddpg/maddpg_policy.py

@@ -29,8 +29,8 @@ class MADDPGPostprocessing:
                                episode=None):
         # FIXME: Get done from info is required since agentwise done is not
         # supported now.
-        sample_batch.data["dones"] = self.get_done_from_info(
-            sample_batch.data["infos"])
+        sample_batch.data[SampleBatch.DONES] = self.get_done_from_info(
+            sample_batch.data[SampleBatch.INFOS])
 
         # N-step Q adjustments
         if self.config["n_step"] > 1:
@@ -94,9 +94,9 @@ class MADDPGTFPolicy(MADDPGPostprocessing, TFPolicy):
                     name=name + "_%d" % i) for i, space in enumerate(space_n)
             ]
 
-        obs_ph_n = _make_ph_n(obs_space_n, "obs")
-        act_ph_n = _make_ph_n(act_space_n, "actions")
-        new_obs_ph_n = _make_ph_n(obs_space_n, "new_obs")
+        obs_ph_n = _make_ph_n(obs_space_n, SampleBatch.OBS)
+        act_ph_n = _make_ph_n(act_space_n, SampleBatch.ACTIONS)
+        new_obs_ph_n = _make_ph_n(obs_space_n, SampleBatch.NEXT_OBS)
         new_act_ph_n = _make_ph_n(act_space_n, "new_actions")
         rew_ph = tf1.placeholder(
             tf.float32, shape=None, name="rewards_{}".format(agent_id))
@@ -328,7 +328,7 @@ class MADDPGTFPolicy(MADDPGPostprocessing, TFPolicy):
             if use_state_preprocessor:
                 model_n = [
                     ModelCatalog.get_model({
-                        "obs": obs,
+                        SampleBatch.OBS: obs,
                         "is_training": self._get_is_training_placeholder(),
                     }, obs_space, act_space, 1, self.config["model"])
                     for obs, obs_space, act_space in zip(
@@ -359,7 +359,7 @@ class MADDPGTFPolicy(MADDPGPostprocessing, TFPolicy):
         with tf1.variable_scope(scope, reuse=tf1.AUTO_REUSE) as scope:
             if use_state_preprocessor:
                 model = ModelCatalog.get_model({
-                    "obs": obs,
+                    SampleBatch.OBS: obs,
                     "is_training": self._get_is_training_placeholder(),
                 }, obs_space, act_space, 1, self.config["model"])
                 out = model.last_layer

rllib/evaluation/collectors/simple_list_collector.py

@@ -14,7 +14,6 @@ from ray.rllib.utils.debug import summarize
 from ray.rllib.utils.typing import AgentID, EpisodeID, EnvID, PolicyID, \
     TensorType
 from ray.rllib.utils.framework import try_import_tf, try_import_torch
-from ray.rllib.utils.torch_ops import convert_to_non_torch_type
 from ray.util.debug import log_once
 
 _, tf, _ = try_import_tf()
@@ -27,9 +26,8 @@ logger = logging.getLogger(__name__)
 
 
 def to_float_np_array(v: List[Any]) -> np.ndarray:
-    if torch.is_tensor(v[0]):
+    if torch and torch.is_tensor(v[0]):
         raise ValueError
-        v = convert_to_non_torch_type(v)
     arr = np.array(v)
     if arr.dtype == np.float64:
         return arr.astype(np.float32)  # save some memory
@@ -172,8 +170,8 @@ class _AgentCollector:
             if col in self.buffers:
                 continue
             shift = self.shift_before - (1 if col == SampleBatch.OBS else 0)
-            # Python primitive.
-            if isinstance(data, (int, float, bool, str)):
+            # Python primitive or dict (e.g. INFOs).
+            if isinstance(data, (int, float, bool, str, dict)):
                 self.buffers[col] = [0 for _ in range(shift)]
             # np.ndarray, torch.Tensor, or tf.Tensor.
             else:

rllib/models/catalog.py

@@ -247,6 +247,7 @@ class ModelCatalog:
             action_space (Space): Action space of the target gym env.
             name (str): An optional string to name the placeholder by.
                 Default: "action".
+
         Returns:
             action_placeholder (Tensor): A placeholder for the actions
         """

rllib/models/modelv2.py

@@ -60,7 +60,7 @@ class ModelV2:
         self._last_output = None
         self.time_major = self.model_config.get("_time_major")
         self.inference_view_requirements = {
-            SampleBatch.OBS: ViewRequirement(shift=0),
+            SampleBatch.OBS: ViewRequirement(shift=0, space=self.obs_space),
         }
 
     # TODO: (sven): Get rid of `get_initial_state` once Trajectory

rllib/models/torch/modules/noisy_layer.py

@@ -1,7 +1,6 @@
 import numpy as np
 
 from ray.rllib.utils.framework import get_activation_fn, try_import_torch
-from ray.rllib.utils.framework import get_variable
 
 torch, nn = try_import_torch()
 
@@ -36,47 +35,39 @@ class NoisyLayer(nn.Module):
         if self.activation is not None:
             self.activation = self.activation()
 
-        self.sigma_w = get_variable(
-            np.random.uniform(
-                low=-1.0 / np.sqrt(float(self.in_size)),
-                high=1.0 / np.sqrt(float(self.in_size)),
-                size=[self.in_size, out_size]),
-            framework="torch",
-            dtype=torch.float32,
-            torch_tensor=True,
-            trainable=True)
-        self.sigma_b = get_variable(
-            np.full(
-                shape=[out_size],
-                fill_value=sigma0 / np.sqrt(float(self.in_size))),
-            framework="torch",
-            dtype=torch.float32,
-            torch_tensor=True,
-            trainable=True)
-        self.w = get_variable(
-            np.full(
-                shape=[self.in_size, self.out_size],
-                fill_value=6 / np.sqrt(float(in_size) + float(out_size))),
-            framework="torch",
-            dtype=torch.float32,
-            torch_tensor=True,
-            trainable=True)
-        self.b = get_variable(
-            np.zeros([out_size]),
-            framework="torch",
-            dtype=torch.float32,
-            torch_tensor=True,
-            trainable=True)
+        sigma_w = nn.Parameter(
+            torch.from_numpy(
+                np.random.uniform(
+                    low=-1.0 / np.sqrt(float(self.in_size)),
+                    high=1.0 / np.sqrt(float(self.in_size)),
+                    size=[self.in_size, out_size])).float())
+        self.register_parameter("sigma_w", sigma_w)
+        sigma_b = nn.Parameter(
+            torch.from_numpy(
+                np.full(
+                    shape=[out_size],
+                    fill_value=sigma0 / np.sqrt(float(self.in_size)))).float())
+        self.register_parameter("sigma_b", sigma_b)
+
+        w = nn.Parameter(
+            torch.from_numpy(
+                np.full(
+                    shape=[self.in_size, self.out_size],
+                    fill_value=6 /
+                    np.sqrt(float(in_size) + float(out_size)))).float())
+        self.register_parameter("w", w)
+        b = nn.Parameter(torch.from_numpy(np.zeros([out_size])).float())
+        self.register_parameter("b", b)
 
     def forward(self, inputs):
         epsilon_in = self._f_epsilon(
             torch.normal(
                 mean=torch.zeros([self.in_size]),
-                std=torch.ones([self.in_size])))
+                std=torch.ones([self.in_size])).to(inputs.device))
         epsilon_out = self._f_epsilon(
             torch.normal(
                 mean=torch.zeros([self.out_size]),
-                std=torch.ones([self.out_size])))
+                std=torch.ones([self.out_size])).to(inputs.device))
         epsilon_w = torch.matmul(
             torch.unsqueeze(epsilon_in, -1),
             other=torch.unsqueeze(epsilon_out, 0))

rllib/models/torch/torch_action_dist.py

@@ -4,7 +4,7 @@ import numpy as np
 import tree
 
 from ray.rllib.models.action_dist import ActionDistribution
-from ray.rllib.models.modelv2 import ModelV2
+from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
 from ray.rllib.utils.annotations import override
 from ray.rllib.utils.framework import try_import_torch
 from ray.rllib.utils.numpy import SMALL_NUMBER, MIN_LOG_NN_OUTPUT, \
@@ -20,9 +20,13 @@ class TorchDistributionWrapper(ActionDistribution):
     """Wrapper class for torch.distributions."""
 
     @override(ActionDistribution)
-    def __init__(self, inputs: List[TensorType], model: ModelV2):
+    def __init__(self, inputs: List[TensorType], model: TorchModelV2):
+        # If inputs are not a torch Tensor, make them one and make sure they
+        # are on the correct device.
         if not isinstance(inputs, torch.Tensor):
-            inputs = torch.Tensor(inputs)
+            inputs = torch.from_numpy(inputs)
+            if isinstance(model, TorchModelV2):
+                inputs = inputs.to(next(model.parameters()).device)
         super().__init__(inputs, model)
         # Store the last sample here.
         self.last_sample = None
@@ -332,8 +336,8 @@ class TorchMultiActionDistribution(TorchDistributionWrapper):
 
         Args:
             inputs (torch.Tensor): A single tensor of shape [BATCH, size].
-            model (ModelV2): The ModelV2 object used to produce inputs for this
-                distribution.
+            model (TorchModelV2): The TorchModelV2 object used to produce
+                inputs for this distribution.
             child_distributions (any[torch.Tensor]): Any struct
                 that contains the child distribution classes to use to
                 instantiate the child distributions from `inputs`. This could
@@ -345,7 +349,9 @@ class TorchMultiActionDistribution(TorchDistributionWrapper):
                 and possibly nested action space.
         """
         if not isinstance(inputs, torch.Tensor):
-            inputs = torch.Tensor(inputs)
+            inputs = torch.from_numpy(inputs)
+            if isinstance(model, TorchModelV2):
+                inputs = inputs.to(next(model.parameters()).device)
         super().__init__(inputs, model)
 
         self.action_space_struct = get_base_struct_from_space(action_space)

rllib/policy/tf_policy.py

@@ -137,7 +137,12 @@ class TFPolicy(Policy):
         """
         self.framework = "tf"
         super().__init__(observation_space, action_space, config)
+
+        assert model is None or isinstance(model, ModelV2), \
+            "Model classes for TFPolicy other than `ModelV2` not allowed! " \
+            "You passed in {}.".format(model)
         self.model = model
+
         self.exploration = self._create_exploration()
         self._sess = sess
         self._obs_input = obs_input
@@ -270,13 +275,9 @@ class TFPolicy(Policy):
         ]
         self._grads = [g for (g, v) in self._grads_and_vars]
 
-        # TODO(sven/ekl): Deprecate support for v1 models.
-        if hasattr(self, "model") and isinstance(self.model, ModelV2):
+        if self.model:
             self._variables = ray.experimental.tf_utils.TensorFlowVariables(
                 [], self._sess, self.variables())
-        else:
-            self._variables = ray.experimental.tf_utils.TensorFlowVariables(
-                self._loss, self._sess)
 
         # gather update ops for any batch norm layers
         if not self._update_ops:
@@ -331,7 +332,8 @@ class TFPolicy(Policy):
         fetched = builder.get(to_fetch)
 
         # Update our global timestep by the batch size.
-        self.global_timestep += fetched[0].shape[0]
+        self.global_timestep += len(obs_batch) if isinstance(obs_batch, list) \
+            else obs_batch.shape[0]
 
         return fetched
 

rllib/tests/test_local.py

@@ -15,6 +15,7 @@ class LocalModeTest(unittest.TestCase):
     def test_local(self):
         cf = DEFAULT_CONFIG.copy()
         cf["model"]["fcnet_hiddens"] = [10]
+        cf["num_workers"] = 2
 
         for _ in framework_iterator(cf):
             agent = PGTrainer(cf, "CartPole-v0")

rllib/tests/test_multi_agent_pendulum.py

@@ -19,6 +19,11 @@ class TestMultiAgentPendulum(unittest.TestCase):
         register_env("multi_agent_pendulum",
                      lambda _: MultiAgentPendulum({"num_agents": 1}))
 
+        stop = {
+            "timesteps_total": 500000,
+            "episode_reward_mean": -400.0,
+        }
+
         # Test for both torch and tf.
         for fw in framework_iterator(frameworks=["torch", "tf"]):
             trials = run_experiments(
@@ -26,10 +31,7 @@ class TestMultiAgentPendulum(unittest.TestCase):
                     "test": {
                         "run": "PPO",
                         "env": "multi_agent_pendulum",
-                        "stop": {
-                            "timesteps_total": 500000,
-                            "episode_reward_mean": -300.0,
-                        },
+                        "stop": stop,
                         "config": {
                             "train_batch_size": 2048,
                             "vf_clip_param": 10.0,
@@ -49,9 +51,11 @@ class TestMultiAgentPendulum(unittest.TestCase):
                     }
                 },
                 verbose=1)
-            if trials[0].last_result["episode_reward_mean"] < -300.0:
-                raise ValueError("Did not get to -200 reward",
-                                 trials[0].last_result)
+            if trials[0].last_result["episode_reward_mean"] <= \
+                    stop["episode_reward_mean"]:
+                raise ValueError(
+                    "Did not get to {} reward".format(
+                        stop["episode_reward_mean"]), trials[0].last_result)
 
 
 if __name__ == "__main__":

rllib/tests/test_rollout.py

@@ -170,7 +170,7 @@ def learn_test_multi_agent_plus_rollout(algo):
                 "policy_mapping_fn": policy_fn,
             },
         }
-        stop = {"episode_reward_mean": 190.0}
+        stop = {"episode_reward_mean": 180.0}
         tune.run(
             algo,
             config=config,

rllib/tests/test_rollout_worker.py

@@ -187,39 +187,35 @@ class TestRolloutWorker(unittest.TestCase):
     def test_global_vars_update(self):
         # Allow for Unittest run.
         ray.init(num_cpus=5, ignore_reinit_error=True)
-        for fw in framework_iterator(frameworks=()):
+        for fw in framework_iterator(frameworks=("tf2", "tf")):
             agent = A2CTrainer(
                 env="CartPole-v0",
                 config={
                     "num_workers": 1,
+                    # lr = 0.1 - [(0.1 - 0.000001) / 100000] * ts
                     "lr_schedule": [[0, 0.1], [100000, 0.000001]],
                     "framework": fw,
                 })
-            result = agent.train()
-            for i in range(10):
+            policy = agent.get_policy()
+            for i in range(3):
                 result = agent.train()
-                print("num_steps_sampled={}".format(
-                    result["info"]["num_steps_sampled"]))
                 print("num_steps_trained={}".format(
                     result["info"]["num_steps_trained"]))
                 print("num_steps_sampled={}".format(
                     result["info"]["num_steps_sampled"]))
-                print("num_steps_trained={}".format(
-                    result["info"]["num_steps_trained"]))
-                if i == 0:
-                    self.assertGreater(
-                        result["info"]["learner"]["default_policy"]["cur_lr"],
-                        0.01)
-                if result["info"]["learner"]["default_policy"]["cur_lr"] < \
-                        0.07:
-                    break
-            self.assertLess(
-                result["info"]["learner"]["default_policy"]["cur_lr"], 0.07)
+                global_timesteps = policy.global_timestep
+                print("global_timesteps={}".format(global_timesteps))
+                expected_lr = \
+                    0.1 - ((0.1 - 0.000001) / 100000) * global_timesteps
+                lr = policy.cur_lr
+                if fw == "tf":
+                    lr = policy._sess.run(lr)
+                check(lr, expected_lr, rtol=0.05)
             agent.stop()
 
     def test_no_step_on_init(self):
         register_env("fail", lambda _: FailOnStepEnv())
-        for fw in framework_iterator(frameworks=()):
+        for fw in framework_iterator():
             pg = PGTrainer(
                 env="fail", config={
                     "num_workers": 1,

rllib/utils/test_utils.py

@@ -179,7 +179,7 @@ def check(x, y, decimals=5, atol=None, rtol=None, false=False):
     else:
         if tf1 is not None:
             # y should never be a Tensor (y=expected value).
-            if isinstance(y, tf1.Tensor):
+            if isinstance(y, (tf1.Tensor, tf1.Variable)):
                 # In eager mode, numpyize tensors.
                 if tf.executing_eagerly():
                     y = y.numpy()
@@ -187,11 +187,11 @@ def check(x, y, decimals=5, atol=None, rtol=None, false=False):
                     raise ValueError(
                         "`y` (expected value) must not be a Tensor. "
                         "Use numpy.ndarray instead")
-            if isinstance(x, tf1.Tensor):
+            if isinstance(x, (tf1.Tensor, tf1.Variable)):
                 # In eager mode, numpyize tensors.
                 if tf1.executing_eagerly():
                     x = x.numpy()
-                # Otherwise, use a quick tf-session.
+                # Otherwise, use a new tf-session.
                 else:
                     with tf1.Session() as sess:
                         x = sess.run(x)