SVG-0 agent with entropy and behavior prior.

The stochastic value gradients (SVG)-0 agent bears similarity to DPG but uses the reparameterization trick to learn stochastic policies. This version uses a continuous version of RETRACE to learn the value function across multiple steps and an added entropy bonus for the actor loss. Additionally the code supports learning behavior priors with an example on running with the locomotion go to target task.

PiperOrigin-RevId: 390159206
Change-Id: Ide9d8fd3dca9d6513ac571becd748a701e3c0173

acme/agents/tf/svg0_prior/README.md

@@ -0,0 +1,26 @@
+# Stochastic Value Gradients (SVG) with Behavior Prior.
+
+This folder contains a version of the SVG-0 agent introduced in
+([Heess et al., 2015]) that has been extended with an entropy bonus, RETRACE
+([Munos et al., 2016]) for off-policy correction and code to learn behavior
+priors ([Tirumala et al., 2019], [Galashov et al., 2019]).
+
+The base SVG-0 algorithm is similar to DPG and DDPG
+([Silver et al., 2015], [Lillicrap et al., 2015]) but uses the
+reparameterization trick to learn stochastic and not deterministic policies. In
+addition, the RETRACE algorithm is used to learn value functions using multiple
+timesteps of data with importance sampling for off policy correction.
+
+In addition an optional Behavior Prior can be learnt using this setup with an
+information asymmetry that has shown to boost performance in some domains.
+Example code to run with and without behavior priors on the DeepMind Control
+Suite and Locomotion tasks are provided in the `examples` folder.
+
+
+[Heess et al., 2015]: https://arxiv.org/abs/1510.09142
+[Munos et al., 2016]: https://arxiv.org/abs/1606.02647
+[Lillicrap et al., 2015]: https://arxiv.org/abs/1509.02971
+[Silver et al., 2014]: http://proceedings.mlr.press/v32/silver14
+[Tirumala et al., 2020]: https://arxiv.org/abs/2010.14274
+[Galashov et al., 2019]: https://arxiv.org/abs/1905.01240
+

acme/agents/tf/svg0_prior/__init__.py

@@ -0,0 +1,21 @@
+# Copyright 2018 DeepMind Technologies Limited. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Implementations of a SVG0 agent with prior."""
+
+from acme.agents.tf.svg0_prior.agent import SVG0
+from acme.agents.tf.svg0_prior.agent_distributed import DistributedSVG0
+from acme.agents.tf.svg0_prior.learning import SVG0Learner
+from acme.agents.tf.svg0_prior.networks import make_default_networks
+from acme.agents.tf.svg0_prior.networks import make_network_with_prior

acme/agents/tf/svg0_prior/acting.py

@@ -0,0 +1,67 @@
+# Copyright 2018 DeepMind Technologies Limited. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""SVG0 actor implementation."""
+
+from typing import Optional
+
+from acme import adders
+from acme import types
+
+from acme.agents.tf import actors
+from acme.tf import utils as tf2_utils
+from acme.tf import variable_utils as tf2_variable_utils
+
+import dm_env
+import sonnet as snt
+
+
+class SVG0Actor(actors.FeedForwardActor):
+  """An actor that also returns `log_prob`."""
+
+  def __init__(
+      self,
+      policy_network: snt.Module,
+      adder: Optional[adders.Adder] = None,
+      variable_client: Optional[tf2_variable_utils.VariableClient] = None,
+      deterministic_policy: Optional[bool] = False,
+  ):
+    super().__init__(policy_network, adder, variable_client)
+    self._log_prob = None
+    self._deterministic_policy = deterministic_policy
+
+  def select_action(self, observation: types.NestedArray) -> types.NestedArray:
+    # Add a dummy batch dimension and as a side effect convert numpy to TF.
+    batched_observation = tf2_utils.add_batch_dim(observation)
+
+    # Compute the policy, conditioned on the observation.
+    policy = self._policy_network(batched_observation)
+    if self._deterministic_policy:
+      action = policy.mean()
+    else:
+      action = policy.sample()
+    self._log_prob = policy.log_prob(action)
+    return tf2_utils.to_numpy_squeeze(action)
+
+  def observe(
+      self,
+      action: types.NestedArray,
+      next_timestep: dm_env.TimeStep,
+  ):
+    if not self._adder:
+      return
+
+    extras = {'log_prob': self._log_prob}
+    extras = tf2_utils.to_numpy_squeeze(extras)
+    self._adder.add(action, next_timestep, extras)