<commit_before># python3
# 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.

"""Recurrent Replay Distributed DQN (R2D2) learner implementation."""

import functools
import time
from typing import Dict, Iterator, List, Mapping, Union, Optional

import acme
from acme import specs
from acme.adders import reverb as adders
from acme.tf import losses
from acme.tf import networks
from acme.tf import savers as tf2_savers
from acme.tf import utils as tf2_utils
from acme.utils import counting
from acme.utils import loggers
import numpy as np
import reverb
import sonnet as snt
import tensorflow as tf
import tree

Variables = List[np.ndarray]


class R2D2Learner(acme.Learner, tf2_savers.TFSaveable):
  """R2D2 learner.

  This is the learning component of the R2D2 agent. It takes a dataset as input
  and implements update functionality to learn from this dataset.
  """

  def __init__(
      self,
      environment_spec: specs.EnvironmentSpec,
      network: Union[networks.RNNCore, snt.RNNCore],
      target_network: Union[networks.RNNCore, snt.RNNCore],
      burn_in_length: int,
      sequence_length: int,
      dataset: tf.data.Dataset,
      reverb_client: Optional[reverb.TFClient] = None,
      counter: Optional[counting.Counter] = None,
      logger: Optional[loggers.Logger] = None,
      discount: float = 0.99,
      target_update_period: int = 100,
      importance_sampling_exponent: float = 0.2,
      max_replay_size: int = 1_000_000,
      learning_rate: float = 1e-3,
      # TODO(sergomez): rename to use_core_state for consistency with JAX agent.
      store_lstm_state: bool = True,
      max_priority_weight: float = 0.9,
      n_step: int = 5,
      clip_grad_norm: Optional[float] = None,
  ):

    if not isinstance(network, networks.RNNCore):
      network.unroll = functools.partial(snt.static_unroll, network)
      target_network.unroll = functools.partial(snt.static_unroll,
                                                target_network)

    # Internalise agent components (replay buffer, networks, optimizer).
    # TODO(b/155086959): Fix type stubs and remove.
    self._iterator: Iterator[reverb.ReplaySample] = iter(dataset)  # pytype: disable=wrong-arg-types
    self._network = network
    self._target_network = target_network
    self._optimizer = snt.optimizers.Adam(learning_rate, epsilon=1e-3)
    self._reverb_client = reverb_client

    # Internalise the hyperparameters.
    self._store_lstm_state = store_lstm_state
    self._burn_in_length = burn_in_length
    self._discount = discount
    self._max_replay_size = max_replay_size
    self._importance_sampling_exponent = importance_sampling_exponent
    self._max_priority_weight = max_priority_weight
    self._target_update_period = target_update_period
    self._num_actions = environment_spec.actions.num_values
    self._sequence_length = sequence_length
    self._n_step = n_step
    self._clip_grad_norm = clip_grad_norm

    if burn_in_length:
      self._burn_in = lambda o, s: self._network.unroll(o, s, burn_in_length)
    else:
      self._burn_in = lambda o, s: (o, s)  # pylint: disable=unnecessary-lambda

    # Learner state.
    self._variables = network.variables
    self._num_steps = tf.Variable(
        0., dtype=tf.float32, trainable=False, name='step')

    # Internalise logging/counting objects.
    self._counter = counting.Counter(counter, 'learner')
    self._logger = logger or loggers.TerminalLogger('learner', time_delta=100.)

    # Do not record timestamps until after the first learning step is done.
    # This is to avoid including the time it takes for actors to come online and
    # fill the replay buffer.
    self._timestamp = None

  # @tf.function
  def _step(self) -> Dict[str, tf.Tensor]:

    # Draw a batch of data from replay.
    sample: reverb.ReplaySample = next(self._iterator)

    data = tf2_utils.batch_to_sequence(sample.data)
    observations, actions, rewards, discounts, extra = (data.observation,
                                                        data.action,
                                                        data.reward,
                                                        data.discount,
                                                        data.extras)
    unused_sequence_length, batch_size = actions.shape

    # Get initial state for the LSTM, either from replay or simply use zeros.<commit_msg>Wrap tf-r2d2's step in tf.function.<commit_after># python3
# 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.

"""Recurrent Replay Distributed DQN (R2D2) learner implementation."""

import functools
import time
from typing import Dict, Iterator, List, Mapping, Union, Optional

import acme
from acme import specs
from acme.adders import reverb as adders
from acme.tf import losses
from acme.tf import networks
from acme.tf import savers as tf2_savers
from acme.tf import utils as tf2_utils
from acme.utils import counting
from acme.utils import loggers
import numpy as np
import reverb
import sonnet as snt
import tensorflow as tf
import tree

Variables = List[np.ndarray]


class R2D2Learner(acme.Learner, tf2_savers.TFSaveable):
  """R2D2 learner.

  This is the learning component of the R2D2 agent. It takes a dataset as input
  and implements update functionality to learn from this dataset.
  """

  def __init__(
      self,
      environment_spec: specs.EnvironmentSpec,
      network: Union[networks.RNNCore, snt.RNNCore],
      target_network: Union[networks.RNNCore, snt.RNNCore],
      burn_in_length: int,
      sequence_length: int,
      dataset: tf.data.Dataset,
      reverb_client: Optional[reverb.TFClient] = None,
      counter: Optional[counting.Counter] = None,
      logger: Optional[loggers.Logger] = None,
      discount: float = 0.99,
      target_update_period: int = 100,
      importance_sampling_exponent: float = 0.2,
      max_replay_size: int = 1_000_000,
      learning_rate: float = 1e-3,
      # TODO(sergomez): rename to use_core_state for consistency with JAX agent.
      store_lstm_state: bool = True,
      max_priority_weight: float = 0.9,
      n_step: int = 5,
      clip_grad_norm: Optional[float] = None,
  ):

    if not isinstance(network, networks.RNNCore):
      network.unroll = functools.partial(snt.static_unroll, network)
      target_network.unroll = functools.partial(snt.static_unroll,
                                                target_network)

    # Internalise agent components (replay buffer, networks, optimizer).
    # TODO(b/155086959): Fix type stubs and remove.
    self._iterator: Iterator[reverb.ReplaySample] = iter(dataset)  # pytype: disable=wrong-arg-types
    self._network = network
    self._target_network = target_network
    self._optimizer = snt.optimizers.Adam(learning_rate, epsilon=1e-3)
    self._reverb_client = reverb_client

    # Internalise the hyperparameters.
    self._store_lstm_state = store_lstm_state
    self._burn_in_length = burn_in_length
    self._discount = discount
    self._max_replay_size = max_replay_size
    self._importance_sampling_exponent = importance_sampling_exponent
    self._max_priority_weight = max_priority_weight
    self._target_update_period = target_update_period
    self._num_actions = environment_spec.actions.num_values
    self._sequence_length = sequence_length
    self._n_step = n_step
    self._clip_grad_norm = clip_grad_norm

    if burn_in_length:
      self._burn_in = lambda o, s: self._network.unroll(o, s, burn_in_length)
    else:
      self._burn_in = lambda o, s: (o, s)  # pylint: disable=unnecessary-lambda

    # Learner state.
    self._variables = network.variables
    self._num_steps = tf.Variable(
        0., dtype=tf.float32, trainable=False, name='step')

    # Internalise logging/counting objects.
    self._counter = counting.Counter(counter, 'learner')
    self._logger = logger or loggers.TerminalLogger('learner', time_delta=100.)

    # Do not record timestamps until after the first learning step is done.
    # This is to avoid including the time it takes for actors to come online and
    # fill the replay buffer.
    self._timestamp = None

  @tf.function
  def _step(self) -> Dict[str, tf.Tensor]:

    # Draw a batch of data from replay.
    sample: reverb.ReplaySample = next(self._iterator)

    data = tf2_utils.batch_to_sequence(sample.data)
    observations, actions, rewards, discounts, extra = (data.observation,
                                                        data.action,
                                                        data.reward,
                                                        data.discount,
                                                        data.extras)
    unused_sequence_length, batch_size = actions.shape

    # Get initial state for the LSTM, either from replay or simply use zeros.
{'fraction_non_alphanumeric': 0.07513110125050423, 'fraction_numerical': 0.01189995966115369, 'mean_word_length': 2.9971785570334544, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 20, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2111933', 'n_tokens_mistral': 2952, 'n_tokens_neox': 2836, 'n_words': 871}