[rllib] Don't use flat weights in non-eager mode (#6001)

rllib/agents/dqn/simple_q_policy.py

@@ -72,6 +72,10 @@ class TargetNetworkMixin(object):
 
         self.update_target = do_update
 
+    @override(TFPolicy)
+    def variables(self):
+        return self.q_func_vars + self.target_q_func_vars
+
 
 def build_q_models(policy, obs_space, action_space, config):
 

rllib/agents/ppo/appo_policy.py

@@ -19,9 +19,10 @@ from ray.rllib.policy.sample_batch import SampleBatch
 from ray.rllib.evaluation.postprocessing import compute_advantages
 from ray.rllib.utils import try_import_tf
 from ray.rllib.policy.tf_policy_template import build_tf_policy
-from ray.rllib.policy.tf_policy import LearningRateSchedule
+from ray.rllib.policy.tf_policy import LearningRateSchedule, TFPolicy
 from ray.rllib.agents.ppo.ppo_policy import KLCoeffMixin, ValueNetworkMixin
 from ray.rllib.models import ModelCatalog
+from ray.rllib.utils.annotations import override
 from ray.rllib.utils.explained_variance import explained_variance
 from ray.rllib.utils.tf_ops import make_tf_callable
 
@@ -425,6 +426,10 @@ class TargetNetworkMixin(object):
 
         self.update_target = do_update
 
+    @override(TFPolicy)
+    def variables(self):
+        return self.model_vars + self.target_model_vars
+
 
 def setup_mixins(policy, obs_space, action_space, config):
     LearningRateSchedule.__init__(policy, config["lr"], config["lr_schedule"])

rllib/agents/sac/sac_policy.py

@@ -12,10 +12,12 @@ from ray.rllib.agents.sac.sac_model import SACModel
 from ray.rllib.agents.ddpg.noop_model import NoopModel
 from ray.rllib.agents.dqn.dqn_policy import _postprocess_dqn, PRIO_WEIGHTS
 from ray.rllib.policy.sample_batch import SampleBatch
+from ray.rllib.policy.tf_policy import TFPolicy
 from ray.rllib.policy.tf_policy_template import build_tf_policy
 from ray.rllib.models import ModelCatalog
 from ray.rllib.utils.error import UnsupportedSpaceException
 from ray.rllib.utils import try_import_tf, try_import_tfp
+from ray.rllib.utils.annotations import override
 from ray.rllib.utils.tf_ops import minimize_and_clip, make_tf_callable
 
 tf = try_import_tf()
@@ -330,6 +332,10 @@ class TargetNetworkMixin(object):
     def update_target(self, tau=None):
         self._do_update(np.float32(tau or self.config.get("tau")))
 
+    @override(TFPolicy)
+    def variables(self):
+        return self.model.variables() + self.target_model.variables()
+
 
 def setup_early_mixins(policy, obs_space, action_space, config):
     ExplorationStateMixin.__init__(policy, obs_space, action_space, config)

rllib/policy/eager_tf_policy.py

@@ -350,17 +350,21 @@ def build_eager_tf_policy(name,
 
         @override(Policy)
         def get_weights(self):
-            variables = self.model.variables()
+            variables = self.variables()
             return [v.numpy() for v in variables]
 
         @override(Policy)
         def set_weights(self, weights):
-            variables = self.model.variables()
+            variables = self.variables()
             assert len(weights) == len(variables), (len(weights),
                                                     len(variables))
             for v, w in zip(variables, weights):
                 v.assign(w)
 
+        def variables(self):
+            """Return the list of all savable variables for this policy."""
+            return self.model.variables()
+
         def is_recurrent(self):
             return len(self._state_in) > 0
 

rllib/policy/tf_policy.py

@@ -146,6 +146,10 @@ class TFPolicy(Policy):
             raise ValueError(
                 "seq_lens tensor must be given if state inputs are defined")
 
+    def variables(self):
+        """Return the list of all savable variables for this policy."""
+        return self.model.variables()
+
     def get_placeholder(self, name):
         """Returns the given action or loss input placeholder by name.
 
@@ -194,8 +198,13 @@ class TFPolicy(Policy):
             if g is not None
         ]
         self._grads = [g for (g, v) in self._grads_and_vars]
-        self._variables = ray.experimental.tf_utils.TensorFlowVariables(
-            self._loss, self._sess)
+        if hasattr(self, "model") and isinstance(self.model, ModelV2):
+            self._variables = ray.experimental.tf_utils.TensorFlowVariables(
+                [], self._sess, self.variables())
+        else:
+            # TODO(ekl) deprecate support for v1 models
+            self._variables = ray.experimental.tf_utils.TensorFlowVariables(
+                self._loss, self._sess)
 
         # gather update ops for any batch norm layers
         if not self._update_ops:
@@ -253,11 +262,11 @@ class TFPolicy(Policy):
 
     @override(Policy)
     def get_weights(self):
-        return self._variables.get_flat()
+        return self._variables.get_weights()
 
     @override(Policy)
     def set_weights(self, weights):
-        return self._variables.set_flat(weights)
+        return self._variables.set_weights(weights)
 
     @override(Policy)
     def export_model(self, export_dir):

rllib/tests/test_checkpoint_restore.py

@@ -21,7 +21,7 @@ def get_mean_action(alg, obs):
     return np.mean(out)
 
 
-ray.init(num_cpus=10)
+ray.init(num_cpus=10, object_store_memory=1e9)
 
 CONFIGS = {
     "SAC": {},
@@ -29,6 +29,7 @@ CONFIGS = {
         "episodes_per_batch": 10,
         "train_batch_size": 100,
         "num_workers": 2,
+        "noise_size": 2500000,
         "observation_filter": "MeanStdFilter"
     },
     "DQN": {},
@@ -56,6 +57,7 @@ CONFIGS = {
     "ARS": {
         "num_rollouts": 10,
         "num_workers": 2,
+        "noise_size": 2500000,
         "observation_filter": "MeanStdFilter"
     }
 }