diff --git a/rllib/agents/dqn/simple_q_policy.py b/rllib/agents/dqn/simple_q_policy.py index 297e3b54f..996a62786 100644 --- a/rllib/agents/dqn/simple_q_policy.py +++ b/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): diff --git a/rllib/agents/ppo/appo_policy.py b/rllib/agents/ppo/appo_policy.py index eb3f5ad7f..ce439a99e 100644 --- a/rllib/agents/ppo/appo_policy.py +++ b/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"]) diff --git a/rllib/agents/sac/sac_policy.py b/rllib/agents/sac/sac_policy.py index 8d5aa770b..f1ac1224e 100644 --- a/rllib/agents/sac/sac_policy.py +++ b/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) diff --git a/rllib/policy/eager_tf_policy.py b/rllib/policy/eager_tf_policy.py index fc5490778..3cb58ec46 100644 --- a/rllib/policy/eager_tf_policy.py +++ b/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 diff --git a/rllib/policy/tf_policy.py b/rllib/policy/tf_policy.py index ad030ec73..8fa4ecb19 100644 --- a/rllib/policy/tf_policy.py +++ b/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): diff --git a/rllib/tests/test_checkpoint_restore.py b/rllib/tests/test_checkpoint_restore.py index c2a1d62b2..7d00e5423 100644 --- a/rllib/tests/test_checkpoint_restore.py +++ b/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" } }