Corax: Core RL in JAX
"""Corax is a framework for reinforcement learning in JAX."""

# Expose specs and types modules.
from corax import specs
from corax import types

# Expose core interfaces.
from corax.core import Actor
from corax.core import Learner
from corax.core import Saveable
from corax.core import VariableSource
from corax.core import Worker

# Expose the environment loop.
from corax.environment_loop import EnvironmentLoop

# Expose utilities for working with EnvironmentSpec
from corax.specs import make_environment_spec

corax/__init__.py
ethanluoyc/corax
"""Adders for sending data from actors to replay buffers."""
from corax.adders.base import Adder
from corax.adders.wrappers import ForkingAdder
from corax.adders.wrappers import IgnoreExtrasAdder

corax/adders/__init__.py
ethanluoyc/corax
# 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.

"""Interface for adders which transmit data to a replay buffer."""

import abc

import dm_env

from corax import types


class Adder(abc.ABC):
    """The Adder interface.

    An adder packs together data to send to the replay buffer, and potentially
    performs some reduction/transformation to this data in the process.

    All adders will use this API. Below is an illustrative example of how they
    are intended to be used in a typical RL run-loop. We assume that the
    environment conforms to the dm_env environment API.

    ```python
    # Reset the environment and add the first observation.
    timestep = env.reset()
    adder.add_first(timestep.observation)

    while not timestep.last():
      # Generate an action from the policy and step the environment.
      action = my_policy(timestep)
      timestep = env.step(action)

      # Add the action and the resulting timestep.
      adder.add(action, next_timestep=timestep)
    ```

    Note that for all adders, the `add()` method expects an action taken and the
    *resulting* timestep observed after taking this action. Note that this
    timestep is named `next_timestep` precisely to emphasize this point.
    """

    @abc.abstractmethod
    def add_first(self, timestep: dm_env.TimeStep):
        """Defines the interface for an adder's `add_first` method.

        We expect this to be called at the beginning of each episode and it will
        start a trajectory to be added to replay with an initial observation.

        Args:
          timestep: a dm_env TimeStep corresponding to the first step.
        """

    @abc.abstractmethod
    def add(
        self,
        action: types.NestedArray,
        next_timestep: dm_env.TimeStep,
        extras: types.NestedArray = (),
    ):
        """Defines the adder `add` interface.

        Args:
          action: A possibly nested structure corresponding to a_t.
          next_timestep: A dm_env Timestep object corresponding to the resulting
            data obtained by taking the given action.
          extras: A possibly nested structure of extra data to add to replay.
        """

    @abc.abstractmethod
    def reset(self):
        """Resets the adder's buffer."""

corax/adders/base.py
ethanluoyc/corax
"""Adders for Reverb replay buffers."""
from corax.adders.reverb.base import DEFAULT_PRIORITY_TABLE
from corax.adders.reverb.base import PriorityFn
from corax.adders.reverb.base import PriorityFnInput
from corax.adders.reverb.base import ReverbAdder
from corax.adders.reverb.base import Step
from corax.adders.reverb.base import Trajectory
from corax.adders.reverb.episode import EpisodeAdder
from corax.adders.reverb.sequence import EndBehavior
from corax.adders.reverb.sequence import SequenceAdder
from corax.adders.reverb.structured import StructuredAdder
from corax.adders.reverb.structured import create_n_step_transition_config
from corax.adders.reverb.structured import create_step_spec
from corax.adders.reverb.transition import NStepTransitionAdder

corax/adders/reverb/__init__.py
ethanluoyc/corax
# 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.

"""Adders that use Reverb (github.com/deepmind/reverb) as a backend."""

import abc
import time
from typing import (
    Callable,
    Iterable,
    Mapping,
    NamedTuple,
    Optional,
    Sized,
    Tuple,
    Union,
)

import dm_env
import numpy as np
import reverb
import tensorflow as tf
import tree
from absl import logging

from corax import specs
from corax import types
from corax.adders import base

DEFAULT_PRIORITY_TABLE = "priority_table"
_MIN_WRITER_LIFESPAN_SECONDS = 60
StartOfEpisodeType = Union[bool, specs.Array, tf.Tensor, tf.TensorSpec, Tuple[()]]


# TODO(b/188510142): Delete Step.
class Step(NamedTuple):
    """Step class used internally for reverb adders."""

    observation: types.NestedArray
    action: types.NestedArray
    reward: types.NestedArray
    discount: types.NestedArray
    start_of_episode: StartOfEpisodeType
    extras: types.NestedArray = ()


# TODO(b/188510142): Replace with proper Trajectory class.
Trajectory = Step


class PriorityFnInput(NamedTuple):
    """The input to a priority function consisting of stacked steps."""

    observations: types.NestedArray
    actions: types.NestedArray
    rewards: types.NestedArray
    discounts: types.NestedArray
    start_of_episode: types.NestedArray
    extras: types.NestedArray


# Define the type of a priority function and the mapping from table to function.
PriorityFn = Callable[["PriorityFnInput"], float]
PriorityFnMapping = Mapping[str, Optional[PriorityFn]]


def spec_like_to_tensor_spec(paths: Iterable[str], spec: specs.Array):
    return tf.TensorSpec.from_spec(spec, name="/".join(str(p) for p in paths))


class ReverbAdder(base.Adder):
    """Base class for Reverb adders."""

    def __init__(
        self,
        client: reverb.Client,
        max_sequence_length: int,
        max_in_flight_items: int,
        delta_encoded: bool = False,
        priority_fns: Optional[PriorityFnMapping] = None,
        validate_items: bool = True,
    ):
        """Initialize a ReverbAdder instance.

        Args:
          client: A client to the Reverb backend.
          max_sequence_length: The maximum length of sequences (corresponding to the
            number of observations) that can be added to replay.
          max_in_flight_items: The maximum number of items allowed to be "in flight"
            at the same time. See `block_until_num_items` in
            `reverb.TrajectoryWriter.flush` for more info.
          delta_encoded: If `True` (False by default) enables delta encoding, see
            `Client` for more information.
          priority_fns: A mapping from table names to priority functions; if
            omitted, all transitions/steps/sequences are given uniform priorities
            (1.0) and placed in DEFAULT_PRIORITY_TABLE.
          validate_items: Whether to validate items against the table signature
            before they are sent to the server. This requires table signature to be
            fetched from the server and cached locally.
        """
        if priority_fns:
            priority_fns = dict(priority_fns)
        else:
            priority_fns = {DEFAULT_PRIORITY_TABLE: None}

        self._client = client
        self._priority_fns = priority_fns
        self._max_sequence_length = max_sequence_length
        self._delta_encoded = delta_encoded
        # TODO(b/206629159): Remove this.
        self._max_in_flight_items = max_in_flight_items
        self._add_first_called = False

        # This is exposed as the _writer property in such a way that it will create
        # a new writer automatically whenever the internal __writer is None. Users
        # should ONLY ever interact with self._writer.
        self.__writer = None
        # Every time a new writer is created, it must fetch the signature from the
        # Reverb server. If this is set too low it can crash the adders in a
        # distributed setup where the replay may take a while to spin up.
        self._validate_items = validate_items

    def __del__(self):
        if self.__writer is not None:
            timeout_ms = 10_000
            # Try flush all appended data before closing to avoid loss of experience.
            try:
                self.__writer.flush(0, timeout_ms=timeout_ms)
            except reverb.DeadlineExceededError as e:
                logging.error(
                    "Timeout (%d ms) exceeded when flushing the writer before "
                    "deleting it. Caught Reverb exception: %s",
                    timeout_ms,
                    str(e),
                )
            self.__writer.close()
            self.__writer = None

    @property
    def _writer(self) -> reverb.TrajectoryWriter:
        if self.__writer is None:
            self.__writer = self._client.trajectory_writer(
                num_keep_alive_refs=self._max_sequence_length,
                validate_items=self._validate_items,
            )
            self._writer_created_timestamp = time.time()
        return self.__writer

    def add_priority_table(self, table_name: str, priority_fn: Optional[PriorityFn]):
        if table_name in self._priority_fns:
            raise ValueError(
                f"A priority function already exists for {table_name}. "
                f'Existing tables: {", ".join(self._priority_fns.keys())}.'
            )
        self._priority_fns[table_name] = priority_fn

    def reset(self, timeout_ms: Optional[int] = None):
        """Resets the adder's buffer."""
        if self.__writer:
            # Flush all appended data and clear the buffers.
            self.__writer.end_episode(clear_buffers=True, timeout_ms=timeout_ms)

            # Create a new writer unless the current one is too young.
            # This is to reduce the relative overhead of creating a new Reverb writer.
            if (
                time.time() - self._writer_created_timestamp
                > _MIN_WRITER_LIFESPAN_SECONDS
            ):
                self.__writer = None
        self._add_first_called = False

    def add_first(self, timestep: dm_env.TimeStep):
        """Record the first observation of a trajectory."""
        if not timestep.first():
            raise ValueError(
                "adder.add_first with an initial timestep (i.e. one for "
                "which timestep.first() is True"
            )

        # Record the next observation but leave the history buffer row open by
        # passing `partial_step=True`.
        self._writer.append(
            dict(observation=timestep.observation, start_of_episode=timestep.first()),
            partial_step=True,
        )
        self._add_first_called = True

    def add(
        self,
        action: types.NestedArray,
        next_timestep: dm_env.TimeStep,
        extras: types.NestedArray = (),
    ):
        """Record an action and the following timestep."""

        if not self._add_first_called:
            raise ValueError("adder.add_first must be called before adder.add.")

        # Add the timestep to the buffer.
        has_extras = (
            len(extras) > 0
            if isinstance(extras, Sized)  # pylint: disable=g-explicit-length-test
            else extras is not None
        )
        current_step = dict(
            # Observation was passed at the previous add call.
            action=action,
            reward=next_timestep.reward,
            discount=next_timestep.discount,
            # Start of episode indicator was passed at the previous add call.
            **({"extras": extras} if has_extras else {}),
        )
        self._writer.append(current_step)

        # Record the next observation and write.
        self._writer.append(
            dict(
                observation=next_timestep.observation,
                start_of_episode=next_timestep.first(),
            ),
            partial_step=True,
        )
        self._write()

        if next_timestep.last():
            # Complete the row by appending zeros to remaining open fields.
            # TODO(b/183945808): remove this when fields are no longer expected to be
            # of equal length on the learner side.
            dummy_step = tree.map_structure(np.zeros_like, current_step)
            self._writer.append(dummy_step)
            self._write_last()
            self.reset()

    @classmethod
    def signature(
        cls, environment_spec: specs.EnvironmentSpec, extras_spec: types.NestedSpec = ()
    ):
        """This is a helper method for generating signatures for Reverb tables.

        Signatures are useful for validating data types and shapes, see Reverb's
        documentation for details on how they are used.

        Args:
          environment_spec: A `specs.EnvironmentSpec` whose fields are nested
            structures with leaf nodes that have `.shape` and `.dtype` attributes.
            This should come from the environment that will be used to generate
            the data inserted into the Reverb table.
          extras_spec: A nested structure with leaf nodes that have `.shape` and
            `.dtype` attributes. The structure (and shapes/dtypes) of this must
            be the same as the `extras` passed into `ReverbAdder.add`.

        Returns:
          A `Step` whose leaf nodes are `tf.TensorSpec` objects.
        """
        spec_step = Step(
            observation=environment_spec.observations,
            action=environment_spec.actions,
            reward=environment_spec.rewards,
            discount=environment_spec.discounts,
            start_of_episode=specs.Array(shape=(), dtype=bool),
            extras=extras_spec,
        )
        return tree.map_structure_with_path(spec_like_to_tensor_spec, spec_step)

    @abc.abstractmethod
    def _write(self):
        """Write data to replay from the buffer."""

    @abc.abstractmethod
    def _write_last(self):
        """Write data to replay from the buffer."""

corax/adders/reverb/base.py
ethanluoyc/corax
# 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.

"""Episode adders.

This implements full episode adders, potentially with padding.
"""

from typing import Callable, Iterable, Optional, Tuple

import dm_env
import numpy as np
import reverb
import tensorflow as tf
import tree

from corax import specs
from corax import types
from corax.adders.reverb import base
from corax.adders.reverb import utils

_PaddingFn = Callable[[Tuple[int, ...], np.dtype], np.ndarray]


class EpisodeAdder(base.ReverbAdder):
    """Adder which adds entire episodes as trajectories."""

    def __init__(
        self,
        client: reverb.Client,
        max_sequence_length: int,
        delta_encoded: bool = False,
        priority_fns: Optional[base.PriorityFnMapping] = None,
        max_in_flight_items: int = 1,
        padding_fn: Optional[_PaddingFn] = None,
        # Deprecated kwargs.
        chunk_length: Optional[int] = None,
    ):
        del chunk_length

        super().__init__(
            client=client,
            max_sequence_length=max_sequence_length,
            delta_encoded=delta_encoded,
            priority_fns=priority_fns,
            max_in_flight_items=max_in_flight_items,
        )
        self._padding_fn = padding_fn

    def add(
        self,
        action: types.NestedArray,
        next_timestep: dm_env.TimeStep,
        extras: types.NestedArray = (),
    ):
        if self._writer.episode_steps >= self._max_sequence_length - 1:
            raise ValueError(
                "The number of observations within the same episode will exceed "
                "max_sequence_length with the addition of this transition."
            )

        super().add(action, next_timestep, extras)

    def _write(self):
        # This adder only writes at the end of the episode, see _write_last()
        pass

    def _write_last(self):
        if (
            self._padding_fn is not None
            and self._writer.episode_steps < self._max_sequence_length
        ):
            history = self._writer.history
            padding_step = dict(
                observation=history["observation"],
                action=history["action"],
                reward=history["reward"],
                discount=history["discount"],
                extras=history.get("extras", ()),
            )
            # Get shapes and dtypes from the last element.
            padding_step = tree.map_structure(
                lambda col: self._padding_fn(col[-1].shape, col[-1].dtype), padding_step
            )
            padding_step["start_of_episode"] = False
            while self._writer.episode_steps < self._max_sequence_length:
                self._writer.append(padding_step)

        trajectory = tree.map_structure(lambda x: x[:], self._writer.history)

        # Pack the history into a base.Step structure and get numpy converted
        # variant for priotiy computation.
        trajectory = base.Trajectory(**trajectory)

        # Calculate the priority for this episode.
        table_priorities = utils.calculate_priorities(self._priority_fns, trajectory)

        # Create a prioritized item for each table.
        for table_name, priority in table_priorities.items():
            self._writer.create_item(table_name, priority, trajectory)
            self._writer.flush(self._max_in_flight_items)

    # TODO(b/185309817): make this into a standalone method.
    @classmethod
    def signature(
        cls,
        environment_spec: specs.EnvironmentSpec,
        extras_spec: types.NestedSpec = (),
        sequence_length: Optional[int] = None,
    ):
        """This is a helper method for generating signatures for Reverb tables.

        Signatures are useful for validating data types and shapes, see Reverb's
        documentation for details on how they are used.

        Args:
          environment_spec: A `specs.EnvironmentSpec` whose fields are nested
            structures with leaf nodes that have `.shape` and `.dtype` attributes.
            This should come from the environment that will be used to generate the
            data inserted into the Reverb table.
          extras_spec: A nested structure with leaf nodes that have `.shape` and
            `.dtype` attributes. The structure (and shapes/dtypes) of this must be
            the same as the `extras` passed into `ReverbAdder.add`.
          sequence_length: An optional integer representing the expected length of
            sequences that will be added to replay.

        Returns:
          A `Step` whose leaf nodes are `tf.TensorSpec` objects.
        """

        def add_time_dim(paths: Iterable[str], spec: tf.TensorSpec):
            return tf.TensorSpec(
                shape=(sequence_length, *spec.shape),
                dtype=spec.dtype,
                name="/".join(str(p) for p in paths),
            )

        trajectory_env_spec, trajectory_extras_spec = tree.map_structure_with_path(
            add_time_dim, (environment_spec, extras_spec)
        )

        trajectory_spec = base.Trajectory(
            *trajectory_env_spec,
            start_of_episode=tf.TensorSpec(
                shape=(sequence_length,), dtype=tf.bool, name="start_of_episode"
            ),
            extras=trajectory_extras_spec,
        )

        return trajectory_spec

corax/adders/reverb/episode.py
ethanluoyc/corax
# 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.

"""Tests for Episode adders."""

import numpy as np
from absl.testing import absltest
from absl.testing import parameterized

from corax.adders.reverb import episode as adders
from corax.adders.reverb import test_utils


class EpisodeAdderTest(test_utils.AdderTestMixin, parameterized.TestCase):
    @parameterized.parameters(2, 10, 50)
    def test_adder(self, max_sequence_length):
        adder = adders.EpisodeAdder(self.client, max_sequence_length)

        # Create a simple trajectory to add.
        observations = range(max_sequence_length)
        first, steps = test_utils.make_trajectory(observations)

        expected_episode = test_utils.make_sequence(observations)
        self.run_test_adder(
            adder=adder,
            first=first,
            steps=steps,
            expected_items=[expected_episode],
            signature=adder.signature(*test_utils.get_specs(steps[0])),
        )

    @parameterized.parameters(2, 10, 50)
    def test_max_sequence_length(self, max_sequence_length):
        adder = adders.EpisodeAdder(self.client, max_sequence_length)

        first, steps = test_utils.make_trajectory(range(max_sequence_length + 1))
        adder.add_first(first)
        for action, step in steps[:-1]:
            adder.add(action, step)

        # We should have max_sequence_length-1 timesteps that have been written,
        # where the -1 is due to the dangling observation (ie we have actually
        # seen max_sequence_length observations).
        self.assertEqual(self.num_items(), 0)

        # Adding one more step should raise an error.
        with self.assertRaises(ValueError):
            action, step = steps[-1]
            adder.add(action, step)

        # Since the last insert failed it should not affect the internal state.
        self.assertEqual(self.num_items(), 0)

    @parameterized.parameters((2, 1), (10, 2), (50, 5))
    def test_padding(self, max_sequence_length, padding):
        adder = adders.EpisodeAdder(
            self.client, max_sequence_length + padding, padding_fn=np.zeros
        )

        # Create a simple trajectory to add.
        observations = range(max_sequence_length)
        first, steps = test_utils.make_trajectory(observations)

        expected_episode = test_utils.make_sequence(observations)
        for _ in range(padding):
            expected_episode.append((0, 0, 0.0, 0.0, False, ()))

        self.run_test_adder(
            adder=adder,
            first=first,
            steps=steps,
            expected_items=[expected_episode],
            signature=adder.signature(*test_utils.get_specs(steps[0])),
        )

    @parameterized.parameters((2, 1), (10, 2), (50, 5))
    def test_nonzero_padding(self, max_sequence_length, padding):
        adder = adders.EpisodeAdder(
            self.client,
            max_sequence_length + padding,
            padding_fn=lambda s, d: np.zeros(s, d) - 1,
        )

        # Create a simple trajectory to add.
        observations = range(max_sequence_length)
        first, steps = test_utils.make_trajectory(observations)

        expected_episode = test_utils.make_sequence(observations)
        for _ in range(padding):
            expected_episode.append((-1, -1, -1.0, -1.0, False, ()))

        self.run_test_adder(
            adder=adder,
            first=first,
            steps=steps,
            expected_items=[expected_episode],
            signature=adder.signature(*test_utils.get_specs(steps[0])),
        )


if __name__ == "__main__":
    absltest.main()

corax/adders/reverb/episode_test.py
ethanluoyc/corax
# 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.

"""Sequence adders.

This implements adders which add sequences or partial trajectories.
"""

import enum
import operator
from typing import Iterable, Optional

import numpy as np
import reverb
import tensorflow as tf
import tree

from corax import specs
from corax import types
from corax.adders.reverb import base
from corax.adders.reverb import utils


class EndBehavior(enum.Enum):
    """Class to enumerate available options for writing behavior at episode ends.

    Example:

      sequence_length = 3
      period = 2

    Episode steps (digits) and writing events (W):

               1 2 3 4 5 6
                   W   W

    First two sequences:

               1 2 3
               . . 3 4 5

    Written sequences for the different end of episode behaviors:
    Here are the last written sequences for each end of episode behavior:

     WRITE     . . . 4 5 6
     CONTINUE  . . . . 5 6 F
     ZERO_PAD  . . . . 5 6 0
     TRUNCATE  . . . . 5 6

    Key:
      F: First step of the next episode
      0: Zero-filled Step
    """

    WRITE = "write_buffer"
    CONTINUE = "continue_to_next_episode"
    ZERO_PAD = "zero_pad_til_next_write"
    TRUNCATE = "write_truncated_buffer"


class SequenceAdder(base.ReverbAdder):
    """An adder which adds sequences of fixed length."""

    def __init__(
        self,
        client: reverb.Client,
        sequence_length: int,
        period: int,
        *,
        delta_encoded: bool = False,
        priority_fns: Optional[base.PriorityFnMapping] = None,
        max_in_flight_items: Optional[int] = 2,
        end_of_episode_behavior: Optional[EndBehavior] = None,
        # Deprecated kwargs.
        chunk_length: Optional[int] = None,
        pad_end_of_episode: Optional[bool] = None,
        break_end_of_episode: Optional[bool] = None,
        validate_items: bool = True,
    ):
        """Makes a SequenceAdder instance.

        Args:
          client: See docstring for BaseAdder.
          sequence_length: The fixed length of sequences we wish to add.
          period: The period with which we add sequences. If less than
            sequence_length, overlapping sequences are added. If equal to
            sequence_length, sequences are exactly non-overlapping.
          delta_encoded: If `True` (False by default) enables delta encoding, see
            `Client` for more information.
          priority_fns: See docstring for BaseAdder.
          max_in_flight_items: The maximum number of items allowed to be "in flight"
            at the same time. See `block_until_num_items` in
            `reverb.TrajectoryWriter.flush` for more info.
          end_of_episode_behavior:  Determines how sequences at the end of the
            episode are handled (default `EndOfEpisodeBehavior.ZERO_PAD`). See
            the docstring for `EndOfEpisodeBehavior` for more information.
          chunk_length: Deprecated and unused.
          pad_end_of_episode: If True (default) then upon end of episode the current
            sequence will be padded (with observations, actions, etc... whose values
            are 0) until its length is `sequence_length`. If False then the last
            sequence in the episode may have length less than `sequence_length`.
          break_end_of_episode: If 'False' (True by default) does not break
            sequences on env reset. In this case 'pad_end_of_episode' is not used.
          validate_items: Whether to validate items against the table signature
            before they are sent to the server. This requires table signature to be
            fetched from the server and cached locally.
        """
        del chunk_length
        super().__init__(
            client=client,
            # We need an additional space in the buffer for the partial step the
            # base.ReverbAdder will add with the next observation.
            max_sequence_length=sequence_length + 1,
            delta_encoded=delta_encoded,
            priority_fns=priority_fns,
            max_in_flight_items=max_in_flight_items,
            validate_items=validate_items,
        )

        if pad_end_of_episode and not break_end_of_episode:
            raise ValueError(
                "Can't set pad_end_of_episode=True and break_end_of_episode=False at"
                " the same time, since those behaviors are incompatible."
            )

        self._period = period
        self._sequence_length = sequence_length

        if end_of_episode_behavior and (
            pad_end_of_episode is not None or break_end_of_episode is not None
        ):
            raise ValueError(
                "Using end_of_episode_behavior and either "
                "pad_end_of_episode or break_end_of_episode is not permitted. "
                "Please use only end_of_episode_behavior instead."
            )

        # Set pad_end_of_episode and break_end_of_episode to default values.
        if end_of_episode_behavior is None and pad_end_of_episode is None:
            pad_end_of_episode = True
        if end_of_episode_behavior is None and break_end_of_episode is None:
            break_end_of_episode = True

        self._end_of_episode_behavior = EndBehavior.ZERO_PAD
        if pad_end_of_episode is not None or break_end_of_episode is not None:
            if not break_end_of_episode:
                self._end_of_episode_behavior = EndBehavior.CONTINUE
            elif break_end_of_episode and pad_end_of_episode:
                self._end_of_episode_behavior = EndBehavior.ZERO_PAD
            elif break_end_of_episode and not pad_end_of_episode:
                self._end_of_episode_behavior = EndBehavior.TRUNCATE
            else:
                raise ValueError(
                    "Reached an unexpected configuration of the SequenceAdder "
                    f"with break_end_of_episode={break_end_of_episode} "
                    f"and pad_end_of_episode={pad_end_of_episode}."
                )
        elif isinstance(end_of_episode_behavior, EndBehavior):
            self._end_of_episode_behavior = end_of_episode_behavior
        else:
            raise ValueError(
                "end_of_episod_behavior must be an instance of "
                f"EndBehavior, received {end_of_episode_behavior}."
            )

    def reset(
        self,
    ):  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
        """Resets the adder's buffer."""
        # If we do not write on end of episode, we should not reset the writer.
        if self._end_of_episode_behavior is EndBehavior.CONTINUE:
            return

        super().reset()

    def _write(self):
        self._maybe_create_item(self._sequence_length)

    def _write_last(self):
        # Maybe determine the delta to the next time we would write a sequence.
        if self._end_of_episode_behavior in (
            EndBehavior.TRUNCATE,
            EndBehavior.ZERO_PAD,
        ):
            delta = self._sequence_length - self._writer.episode_steps
            if delta < 0:
                delta = (self._period + delta) % self._period

        # Handle various end-of-episode cases.
        if self._end_of_episode_behavior is EndBehavior.CONTINUE:
            self._maybe_create_item(self._sequence_length, end_of_episode=True)

        elif self._end_of_episode_behavior is EndBehavior.WRITE:
            # Drop episodes that are too short.
            if self._writer.episode_steps < self._sequence_length:
                return
            self._maybe_create_item(
                self._sequence_length, end_of_episode=True, force=True
            )

        elif self._end_of_episode_behavior is EndBehavior.TRUNCATE:
            self._maybe_create_item(
                self._sequence_length - delta, end_of_episode=True, force=True
            )

        elif self._end_of_episode_behavior is EndBehavior.ZERO_PAD:
            zero_step = tree.map_structure(
                lambda x: np.zeros_like(x[-2].numpy()), self._writer.history
            )
            for _ in range(delta):
                self._writer.append(zero_step)

            self._maybe_create_item(
                self._sequence_length, end_of_episode=True, force=True
            )
        else:
            raise ValueError(
                f"Unhandled end of episode behavior: {self._end_of_episode_behavior}."
                " This should never happen, please contact Acme dev team."
            )

    def _maybe_create_item(
        self, sequence_length: int, *, end_of_episode: bool = False, force: bool = False
    ):
        # Check conditions under which a new item is created.
        first_write = self._writer.episode_steps == sequence_length
        # NOTE(bshahr): the following line assumes that the only way sequence_length
        # is less than self._sequence_length, is if the episode is shorter than
        # self._sequence_length.
        period_reached = self._writer.episode_steps > self._sequence_length and (
            (self._writer.episode_steps - self._sequence_length) % self._period == 0
        )

        if not first_write and not period_reached and not force:
            return

        # TODO(b/183945808): will need to change to adhere to the new protocol.
        if not end_of_episode:
            get_traj = operator.itemgetter(slice(-sequence_length - 1, -1))
        else:
            get_traj = operator.itemgetter(slice(-sequence_length, None))

        history = self._writer.history
        trajectory = base.Trajectory(**tree.map_structure(get_traj, history))

        # Compute priorities for the buffer.
        table_priorities = utils.calculate_priorities(self._priority_fns, trajectory)

        # Create a prioritized item for each table.
        for table_name, priority in table_priorities.items():
            self._writer.create_item(table_name, priority, trajectory)
            self._writer.flush(self._max_in_flight_items)

    # TODO(bshahr): make this into a standalone method. Class methods should be
    # used as alternative constructors or when modifying some global state,
    # neither of which is done here.
    @classmethod
    def signature(
        cls,
        environment_spec: specs.EnvironmentSpec,
        extras_spec: types.NestedSpec = (),
        sequence_length: Optional[int] = None,
    ):
        """This is a helper method for generating signatures for Reverb tables.

        Signatures are useful for validating data types and shapes, see Reverb's
        documentation for details on how they are used.

        Args:
          environment_spec: A `specs.EnvironmentSpec` whose fields are nested
            structures with leaf nodes that have `.shape` and `.dtype` attributes.
            This should come from the environment that will be used to generate
            the data inserted into the Reverb table.
          extras_spec: A nested structure with leaf nodes that have `.shape` and
            `.dtype` attributes. The structure (and shapes/dtypes) of this must
            be the same as the `extras` passed into `ReverbAdder.add`.
          sequence_length: An optional integer representing the expected length of
            sequences that will be added to replay.

        Returns:
          A `Trajectory` whose leaf nodes are `tf.TensorSpec` objects.
        """

        def add_time_dim(paths: Iterable[str], spec: tf.TensorSpec):
            return tf.TensorSpec(
                shape=(sequence_length, *spec.shape),
                dtype=spec.dtype,
                name="/".join(str(p) for p in paths),
            )

        trajectory_env_spec, trajectory_extras_spec = tree.map_structure_with_path(
            add_time_dim, (environment_spec, extras_spec)
        )

        spec_step = base.Trajectory(
            *trajectory_env_spec,
            start_of_episode=tf.TensorSpec(
                shape=(sequence_length,), dtype=tf.bool, name="start_of_episode"
            ),
            extras=trajectory_extras_spec,
        )

        return spec_step

corax/adders/reverb/sequence.py
ethanluoyc/corax
# 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.

"""Tests for sequence adders."""

from absl.testing import absltest
from absl.testing import parameterized

from corax.adders.reverb import sequence as adders
from corax.adders.reverb import test_cases
from corax.adders.reverb import test_utils


class SequenceAdderTest(test_utils.AdderTestMixin, parameterized.TestCase):
    @parameterized.named_parameters(*test_cases.TEST_CASES_FOR_SEQUENCE_ADDER)
    def test_adder(
        self,
        sequence_length: int,
        period: int,
        first,
        steps,
        expected_sequences,
        end_behavior: adders.EndBehavior = adders.EndBehavior.ZERO_PAD,
        repeat_episode_times: int = 1,
    ):
        adder = adders.SequenceAdder(
            self.client,
            sequence_length=sequence_length,
            period=period,
            end_of_episode_behavior=end_behavior,
        )
        super().run_test_adder(
            adder=adder,
            first=first,
            steps=steps,
            expected_items=expected_sequences,
            repeat_episode_times=repeat_episode_times,
            end_behavior=end_behavior,
            signature=adder.signature(*test_utils.get_specs(steps[0])),
        )

    @parameterized.parameters(
        (True, True, adders.EndBehavior.ZERO_PAD),
        (False, True, adders.EndBehavior.TRUNCATE),
        (False, False, adders.EndBehavior.CONTINUE),
    )
    def test_end_of_episode_behavior_set_correctly(
        self, pad_end_of_episode, break_end_of_episode, expected_behavior
    ):
        adder = adders.SequenceAdder(
            self.client,
            sequence_length=5,
            period=3,
            pad_end_of_episode=pad_end_of_episode,
            break_end_of_episode=break_end_of_episode,
        )
        self.assertEqual(adder._end_of_episode_behavior, expected_behavior)


if __name__ == "__main__":
    absltest.main()

corax/adders/reverb/sequence_test.py
ethanluoyc/corax
# 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.

"""Generic adders that wraps Reverb's StructuredWriter."""

import itertools
import time
from typing import Callable, List, Optional, Sequence, Sized

import dm_env
import numpy as np
import reverb
import tensorflow as tf
import tree
from absl import logging
from reverb import structured_writer as sw

from corax import specs
from corax import types
from corax.adders import base as adders_base
from corax.adders.reverb import base as reverb_base
from corax.adders.reverb import sequence as sequence_adder

Step = reverb_base.Step
Trajectory = reverb_base.Trajectory
EndBehavior = sequence_adder.EndBehavior

_RESET_WRITER_EVERY_SECONDS = 60


class StructuredAdder(adders_base.Adder):
    """Generic Adder which writes to Reverb using Reverb's `StructuredWriter`.

    The StructuredAdder is a thin wrapper around Reverb's `StructuredWriter` and
    its behaviour is determined by the configs to __init__. Much of the behaviour
    provided by other Adders can be replicated using `StructuredAdder` but there
    are a few noteworthy differences:

     * The behaviour of `StructuredAdder` can be thought of as the union of all
       its configs. This means that a single adder is capable of inserting items
       of different structures into any number of tables WITHOUT any data
       duplication. Other adders are only capable of writing items of the same
       structure into multiple tables.
     * The complete structure of the step must be known at construction time when
       using the StructuredAdder. This is not the case for other Adders as they
       allow the structure of the step to become expanded over time.
     * The `StructuredAdder` assigns all items the same priority (1.0) as it does
       not currently support custom priority computations.
     * The StructuredAdder is completely generic and thus does not perform any
       preprocessing on the data (e.g. cumulative rewards as done by the
       NStepTransitionAdder) before writing it to Reverb. The user is instead
       expected to perform preprocessing in the dataset pipeline on the learner.
    """

    def __init__(
        self,
        client: reverb.Client,
        max_in_flight_items: int,
        configs: Sequence[sw.Config],
        step_spec: Step,
    ):
        """Initialize a StructuredAdder instance.

        Args:
          client: A client to the Reverb backend.
          max_in_flight_items: The maximum number of items allowed to be "in flight"
            at the same time. See `block_until_num_items` in
            `reverb.TrajectoryWriter.flush` for more info.
          configs: Configurations defining the behaviour of the wrapped Reverb
            writer.
          step_spec: spec of the step that is going to be inserted in the Adder. It
            can be created with `create_step_spec` using the environment spec and
            and the extras spec.
        """

        # We validate the configs by attempting to infer the signatures of all
        # targeted tables.
        for table, table_configs in itertools.groupby(configs, lambda c: c.table):
            try:
                sw.infer_signature(list(table_configs), step_spec)
            except ValueError as e:
                raise ValueError(
                    f"Received invalid configs for table {table}: {str(e)}"
                ) from e

        self._client = client
        self._configs = tuple(configs)
        self._none_step: Step = tree.map_structure(lambda _: None, step_spec)
        self._max_in_flight_items = max_in_flight_items

        self._writer = None
        self._writer_created_at = None

    def __del__(self):
        if self._writer is None:
            return

        # Try flush all appended data before closing to avoid loss of experience.
        try:
            self._writer.flush(0, timeout_ms=10_000)
        except reverb.DeadlineExceededError as e:
            logging.error(
                "Timeout (10 s) exceeded when flushing the writer before "
                "deleting it. Caught Reverb exception: %s",
                str(e),
            )

    def _make_step(self, **kwargs) -> Step:
        """Complete the step with None in the missing positions."""
        return Step(**{**self._none_step._asdict(), **kwargs})

    @property
    def configs(self):
        return self._configs

    def reset(self, timeout_ms: Optional[int] = None):
        """Marks the active episode as completed and flushes pending items."""
        if self._writer is not None:
            # Flush all pending items.
            self._writer.end_episode(clear_buffers=True, timeout_ms=timeout_ms)

            # Create a new writer unless the current one is too young.
            # This is to reduce the relative overhead of creating a new Reverb writer.
            if time.time() - self._writer_created_at > _RESET_WRITER_EVERY_SECONDS:
                self._writer = None

    def add_first(self, timestep: dm_env.TimeStep):
        """Record the first observation of an episode."""
        if not timestep.first():
            raise ValueError(
                "adder.add_first called with a timestep that was not the first of its"
                "episode (i.e. one for which timestep.first() is not True)"
            )

        if self._writer is None:
            self._writer = self._client.structured_writer(self._configs)
            self._writer_created_at = time.time()

        # Record the next observation but leave the history buffer row open by
        # passing `partial_step=True`.
        self._writer.append(
            data=self._make_step(
                observation=timestep.observation, start_of_episode=timestep.first()
            ),
            partial_step=True,
        )
        self._writer.flush(self._max_in_flight_items)

    def add(
        self,
        action: types.NestedArray,
        next_timestep: dm_env.TimeStep,
        extras: types.NestedArray = (),
    ):
        """Record an action and the following timestep."""

        if self._writer is None or not self._writer.step_is_open:
            raise ValueError("adder.add_first must be called before adder.add.")

        # Add the timestep to the buffer.
        has_extras = (
            len(extras) > 0
            if isinstance(extras, Sized)  # pylint: disable=g-explicit-length-test
            else extras is not None
        )

        current_step = self._make_step(
            action=action,
            reward=next_timestep.reward,
            discount=next_timestep.discount,
            extras=extras if has_extras else self._none_step.extras,
        )
        self._writer.append(current_step)

        # Record the next observation and write.
        self._writer.append(
            data=self._make_step(
                observation=next_timestep.observation,
                start_of_episode=next_timestep.first(),
            ),
            partial_step=True,
        )
        self._writer.flush(self._max_in_flight_items)

        if next_timestep.last():
            # Complete the row by appending zeros to remaining open fields.
            # TODO(b/183945808): remove this when fields are no longer expected to be
            # of equal length on the learner side.
            dummy_step = tree.map_structure(
                lambda x: None if x is None else np.zeros_like(x), current_step
            )
            self._writer.append(dummy_step)
            self.reset()


def create_step_spec(
    environment_spec: specs.EnvironmentSpec, extras_spec: types.NestedSpec = ()
) -> Step:
    return Step(
        *environment_spec,
        start_of_episode=tf.TensorSpec([], tf.bool, "start_of_episode"),
        extras=extras_spec,
    )


def _last_n(n: int, step_spec: Step) -> Trajectory:
    """Constructs a sequence with the last n elements of all the Step fields."""
    return Trajectory(*tree.map_structure(lambda x: x[-n:], step_spec))


def create_sequence_config(
    step_spec: Step,
    sequence_length: int,
    period: int,
    table: str = reverb_base.DEFAULT_PRIORITY_TABLE,
    end_of_episode_behavior: EndBehavior = EndBehavior.TRUNCATE,
    sequence_pattern: Callable[[int, Step], Trajectory] = _last_n,
) -> List[sw.Config]:
    """Generates configs that produces the same behaviour as `SequenceAdder`.

    NOTE! ZERO_PAD is not supported as the same behaviour can be achieved by
    writing with TRUNCATE and then adding padding in the dataset pipeline on the
    learner.

    Args:
      step_spec: The full structure of the data which will be appended to the
        Reverb `StructuredWriter` in each step. Please use `create_step_spec` to
        create `step_spec`.
      sequence_length: The number of steps that each trajectory should span.
      period: The period with which we add sequences. If less than
        sequence_length, overlapping sequences are added. If equal to
        sequence_length, sequences are exactly non-overlapping.
      table: Name of the Reverb table to write items to. Defaults to the default
        Acme table.
      end_of_episode_behavior: Determines how sequences at the end of the episode
        are handled (default `EndOfEpisodeBehavior.TRUNCATE`). See the docstring
        of `EndOfEpisodeBehavior` for more information.
      sequence_pattern: Transformation to obtain a sequence given the length
        and the shape of the step.

    Returns:
      A list of configs for `StructuredAdder` to produce the described behaviour.

    Raises:
      ValueError: If sequence_length is <= 0.
      NotImplementedError: If `end_of_episod_behavior` is `ZERO_PAD`.
    """
    if sequence_length <= 0:
        raise ValueError(f"sequence_length must be > 0 but got {sequence_length}.")

    if end_of_episode_behavior == EndBehavior.ZERO_PAD:
        raise NotImplementedError(
            "Zero-padding is not supported. Please use TRUNCATE instead."
        )

    if end_of_episode_behavior == EndBehavior.CONTINUE:
        raise NotImplementedError("Merging episodes is not supported.")

    def _sequence_pattern(n: int) -> sw.Pattern:
        return sw.pattern_from_transform(
            step_spec, lambda step: sequence_pattern(n, step)
        )

    # The base config is considered for all but the last step in the episode. No
    # trajectories are created for the first `sequence_step-1` steps and then a
    # new trajectory is inserted every `period` steps.
    base_config = sw.create_config(
        pattern=_sequence_pattern(sequence_length),
        table=table,
        conditions=[
            sw.Condition.step_index() >= sequence_length - 1,
            sw.Condition.step_index() % period == (sequence_length - 1) % period,
        ],
    )

    end_of_episode_configs = []
    if end_of_episode_behavior == EndBehavior.WRITE:
        # Simply write a trajectory in exactly the same way as the base config. The
        # only difference here is that we ALWAYS create a trajectory even if it
        # doesn't align with the `period`. The exceptions to the rule are episodes
        # that are shorter than `sequence_length` steps which are completely
        # ignored.
        config = sw.create_config(
            pattern=_sequence_pattern(sequence_length),
            table=table,
            conditions=[
                sw.Condition.is_end_episode(),
                sw.Condition.step_index() >= sequence_length - 1,
            ],
        )
        end_of_episode_configs.append(config)
    elif end_of_episode_behavior == EndBehavior.TRUNCATE:
        # The first trajectory is written at step index `sequence_length - 1` and
        # then written every `period` step. This means that the
        # `step_index % period` will always be equal to the below value everytime a
        # trajectory is written.
        target = (sequence_length - 1) % period

        # When the episode ends we still want to capture the steps that has been
        # appended since the last item was created. We do this by creating a config
        # for all `step_index % period`, except `target`, and condition these
        # configs so that they only are triggered when `end_episode` is called.
        for x in range(period):
            # When the last step is aligned with the period of the inserts then no
            # action is required as the item was already generated by `base_config`.
            if x == target:
                continue

            # If we were to pad the trajectory then we'll need to continue adding
            # padding until `step_index % period` is equal to `target` again. We can
            # exploit this relation by conditioning the config to only be applied for
            # a single value of `step_index % period`. This constraint means that we
            # can infer the number of padding steps required until the next write
            # would have occurred if the episode didn't end.
            #
            # Now if we assume that the padding instead is added on the dataset (or
            # the trajectory is simply truncated) then we can infer from the above
            # that the number of real steps in this padded trajectory will be the
            # difference between `sequence_length` and number of pad steps.
            num_pad_steps = (target - x) % period
            unpadded_length = sequence_length - num_pad_steps

            config = sw.create_config(
                pattern=_sequence_pattern(unpadded_length),
                table=table,
                conditions=[
                    sw.Condition.is_end_episode(),
                    sw.Condition.step_index() % period == x,
                    sw.Condition.step_index() >= sequence_length,
                ],
            )
            end_of_episode_configs.append(config)

        # The above configs will capture the "remainder" of any episode that is at
        # least `sequence_length` steps long. However, if the entire episode is
        # shorter than `sequence_length` then data might still be lost. We avoid
        # this by simply creating `sequence_length-1` configs that capture the last
        # `x` steps iff the entire episode is `x` steps long.
        for x in range(1, sequence_length):
            config = sw.create_config(
                pattern=_sequence_pattern(x),
                table=table,
                conditions=[
                    sw.Condition.is_end_episode(),
                    sw.Condition.step_index() == x - 1,
                ],
            )
            end_of_episode_configs.append(config)
    else:
        raise ValueError(
            f"Unexpected `end_of_episod_behavior`: {end_of_episode_behavior}"
        )

    return [base_config] + end_of_episode_configs


def create_n_step_transition_config(
    step_spec: Step, n_step: int, table: str = reverb_base.DEFAULT_PRIORITY_TABLE
) -> List[sw.Config]:
    """Generates configs that replicates the behaviour of NStepTransitionAdder.

    Please see the docstring of NStepTransitionAdder for more details.

    NOTE! In contrast to NStepTransitionAdder, the trajectories written by the
    `StructuredWriter` does not include the precomputed cumulative reward and
    discounts. Instead the trajectory includes the raw rewards and discounts
    required to comptute these values.

    Args:
      step_spec: The full structure of the data which will be appended to the
        Reverb `StructuredWriter` in each step. Please use `create_step_spec` to
        create `step_spec`.
      n_step: The "N" in N-step transition. See the class docstring for the
        precise definition of what an N-step transition is. `n_step` must be at
        least 1, in which case we use the standard one-step transition, i.e. (s_t,
        a_t, r_t, d_t, s_t+1, e_t).
      table: Name of the Reverb table to write items to. Defaults to the default
        Acme table.

    Returns:
      A list of configs for `StructuredAdder` to produce the described behaviour.
    """

    def _make_pattern(n: int):
        ref_step = sw.create_reference_step(step_spec)

        get_first = lambda x: x[-(n + 1) : -n]
        get_all = lambda x: x[-(n + 1) : -1]
        get_first_and_last = lambda x: x[-(n + 1) :: n]

        tmap = tree.map_structure

        # We use the exact same structure as we done when writing sequences except
        # we trim the number of steps in each sub tree. This has the benefit that
        # the postprocessing used to transform these items into N-step transition
        # structures (cumulative rewards and discounts etc.) can be applied on
        # full sequence items as well. The only difference being that the latter is
        # more wasteful than the trimmed down version we write here.
        return Trajectory(
            observation=tmap(get_first_and_last, ref_step.observation),
            action=tmap(get_first, ref_step.action),
            reward=tmap(get_all, ref_step.reward),
            discount=tmap(get_all, ref_step.discount),
            start_of_episode=tmap(get_first, ref_step.start_of_episode),
            extras=tmap(get_first, ref_step.extras),
        )

    # At the start of the episodes we'll add shorter transitions.
    start_of_episode_configs = []
    for n in range(1, n_step):
        config = sw.create_config(
            pattern=_make_pattern(n),
            table=table,
            conditions=[
                sw.Condition.step_index() == n,
            ],
        )
        start_of_episode_configs.append(config)

    # During all other steps we'll add a full N-step transition.
    base_config = sw.create_config(pattern=_make_pattern(n_step), table=table)

    # When the episode ends we'll add shorter transitions.
    end_of_episode_configs = []
    for n in range(n_step - 1, 0, -1):
        config = sw.create_config(
            pattern=_make_pattern(n),
            table=table,
            conditions=[
                sw.Condition.is_end_episode(),
                # If the entire episode is shorter than n_step then the episode
                # start configs will already create an item that covers all the
                # steps so we add this filter here to avoid adding it again.
                sw.Condition.step_index() != n,
            ],
        )
        end_of_episode_configs.append(config)

    return start_of_episode_configs + [base_config] + end_of_episode_configs

corax/adders/reverb/structured.py
ethanluoyc/corax
# 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.

"""Tests for structured."""

from typing import Sequence

import dm_env
import numpy as np
import tree
from absl.testing import absltest
from absl.testing import parameterized
from reverb import structured_writer as sw

from corax import types
from corax.adders.reverb import sequence as adders
from corax.adders.reverb import structured
from corax.adders.reverb import test_cases
from corax.adders.reverb import test_utils
from corax.utils import tree_utils


class StructuredAdderTest(test_utils.AdderTestMixin, parameterized.TestCase):
    @parameterized.named_parameters(*test_cases.BASE_TEST_CASES_FOR_SEQUENCE_ADDER)
    def test_sequence_adder(
        self,
        sequence_length: int,
        period: int,
        first,
        steps,
        expected_sequences,
        end_behavior: adders.EndBehavior,
        repeat_episode_times: int = 1,
    ):
        env_spec, extras_spec = test_utils.get_specs(steps[0])
        step_spec = structured.create_step_spec(env_spec, extras_spec)

        should_pad_trajectory = end_behavior == adders.EndBehavior.ZERO_PAD

        def _maybe_zero_pad(flat_trajectory):
            trajectory = tree.unflatten_as(step_spec, flat_trajectory)

            if not should_pad_trajectory:
                return trajectory

            padding_length = sequence_length - flat_trajectory[0].shape[0]
            if padding_length == 0:
                return trajectory

            padding = tree.map_structure(
                lambda x: np.zeros([padding_length, *x.shape[1:]], x.dtype), trajectory
            )

            return tree.map_structure(lambda *x: np.concatenate(x), trajectory, padding)

        # The StructuredAdder does not support adding padding steps as we assume
        # that the padding will be added on the learner side.
        if end_behavior == adders.EndBehavior.ZERO_PAD:
            end_behavior = adders.EndBehavior.TRUNCATE

        configs = structured.create_sequence_config(
            step_spec=step_spec,
            sequence_length=sequence_length,
            period=period,
            end_of_episode_behavior=end_behavior,
        )
        adder = structured.StructuredAdder(
            client=self.client,
            max_in_flight_items=0,
            configs=configs,
            step_spec=step_spec,
        )

        super().run_test_adder(
            adder=adder,
            first=first,
            steps=steps,
            expected_items=expected_sequences,
            repeat_episode_times=repeat_episode_times,
            end_behavior=end_behavior,
            item_transform=_maybe_zero_pad,
            signature=sw.infer_signature(configs, step_spec),
        )

    @parameterized.named_parameters(*test_cases.TEST_CASES_FOR_TRANSITION_ADDER)
    def test_transition_adder(
        self,
        n_step: int,
        additional_discount: float,
        first: dm_env.TimeStep,
        steps: Sequence[dm_env.TimeStep],
        expected_transitions: Sequence[types.Transition],
    ):
        env_spec, extras_spec = test_utils.get_specs(steps[0])
        step_spec = structured.create_step_spec(env_spec, extras_spec)

        def _as_n_step_transition(flat_trajectory):
            trajectory = tree.unflatten_as(step_spec, flat_trajectory)

            rewards, discount = _compute_cumulative_quantities(
                rewards=trajectory.reward,
                discounts=trajectory.discount,
                additional_discount=additional_discount,
                n_step=tree.flatten(trajectory.reward)[0].shape[0],
            )

            tmap = tree.map_structure
            return types.Transition(
                observation=tmap(lambda x: x[0], trajectory.observation),
                action=tmap(lambda x: x[0], trajectory.action),
                reward=rewards,
                discount=discount,
                next_observation=tmap(lambda x: x[-1], trajectory.observation),
                extras=tmap(lambda x: x[0], trajectory.extras),
            )

        configs = structured.create_n_step_transition_config(
            step_spec=step_spec, n_step=n_step
        )

        adder = structured.StructuredAdder(
            client=self.client,
            max_in_flight_items=0,
            configs=configs,
            step_spec=step_spec,
        )

        super().run_test_adder(
            adder=adder,
            first=first,
            steps=steps,
            expected_items=expected_transitions,
            stack_sequence_fields=False,
            item_transform=_as_n_step_transition,
            signature=sw.infer_signature(configs, step_spec),
        )


def _compute_cumulative_quantities(
    rewards: types.NestedArray,
    discounts: types.NestedArray,
    additional_discount: float,
    n_step: int,
):
    """Stolen from TransitionAdder."""

    # Give the same tree structure to the n-step return accumulator,
    # n-step discount accumulator, and self.discount, so that they can be
    # iterated in parallel using tree.map_structure.
    rewards, discounts, self_discount = tree_utils.broadcast_structures(
        rewards, discounts, additional_discount
    )
    flat_rewards = tree.flatten(rewards)
    flat_discounts = tree.flatten(discounts)
    flat_self_discount = tree.flatten(self_discount)

    # Copy total_discount as it is otherwise read-only.
    total_discount = [np.copy(a[0]) for a in flat_discounts]

    # Broadcast n_step_return to have the broadcasted shape of
    # reward * discount.
    n_step_return = [
        np.copy(np.broadcast_to(r[0], np.broadcast(r[0], d).shape))
        for r, d in zip(flat_rewards, total_discount)
    ]

    # NOTE: total_discount will have one less self_discount applied to it than
    # the value of self._n_step. This is so that when the learner/update uses
    # an additional discount we don't apply it twice. Inside the following loop
    # we will apply this right before summing up the n_step_return.
    for i in range(1, n_step):
        for nsr, td, r, d, sd in zip(
            n_step_return,
            total_discount,
            flat_rewards,
            flat_discounts,
            flat_self_discount,
        ):
            # Equivalent to: `total_discount *= self._discount`.
            td *= sd
            # Equivalent to: `n_step_return += reward[i] * total_discount`.
            nsr += r[i] * td
            # Equivalent to: `total_discount *= discount[i]`.
            td *= d[i]

    n_step_return = tree.unflatten_as(rewards, n_step_return)
    total_discount = tree.unflatten_as(rewards, total_discount)
    return n_step_return, total_discount


if __name__ == "__main__":
    absltest.main()

corax/adders/reverb/structured_test.py
ethanluoyc/corax
# 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.

"""Test cases used by multiple test files."""

import dm_env
import numpy as np

from corax import types
from corax.adders.reverb import sequence as sequence_adder

# Define the main set of test cases; these are given as parameterized tests to
# the test_adder method and describe a trajectory to add to replay and the
# expected transitions that should result from this trajectory. The expected
# transitions are of the form: (observation, action, reward, discount,
# next_observation, extras).
TEST_CASES_FOR_TRANSITION_ADDER = [
    dict(
        testcase_name="OneStepFinalReward",
        n_step=1,
        additional_discount=1.0,
        first=dm_env.restart(1),
        steps=(
            (0, dm_env.transition(reward=0.0, observation=2)),
            (0, dm_env.transition(reward=0.0, observation=3)),
            (0, dm_env.termination(reward=1.0, observation=4)),
        ),
        expected_transitions=(
            types.Transition(1, 0, 0.0, 1.0, 2),
            types.Transition(2, 0, 0.0, 1.0, 3),
            types.Transition(3, 0, 1.0, 0.0, 4),
        ),
    ),
    dict(
        testcase_name="OneStepDict",
        n_step=1,
        additional_discount=1.0,
        first=dm_env.restart({"foo": 1}),
        steps=(
            (0, dm_env.transition(reward=0.0, observation={"foo": 2})),
            (0, dm_env.transition(reward=0.0, observation={"foo": 3})),
            (0, dm_env.termination(reward=1.0, observation={"foo": 4})),
        ),
        expected_transitions=(
            types.Transition({"foo": 1}, 0, 0.0, 1.0, {"foo": 2}),
            types.Transition({"foo": 2}, 0, 0.0, 1.0, {"foo": 3}),
            types.Transition({"foo": 3}, 0, 1.0, 0.0, {"foo": 4}),
        ),
    ),
    dict(
        testcase_name="OneStepExtras",
        n_step=1,
        additional_discount=1.0,
        first=dm_env.restart(1),
        steps=(
            (
                0,
                dm_env.transition(reward=0.0, observation=2),
                {"state": 0},
            ),
            (
                0,
                dm_env.transition(reward=0.0, observation=3),
                {"state": 1},
            ),
            (
                0,
                dm_env.termination(reward=1.0, observation=4),
                {"state": 2},
            ),
        ),
        expected_transitions=(
            types.Transition(1, 0, 0.0, 1.0, 2, {"state": 0}),
            types.Transition(2, 0, 0.0, 1.0, 3, {"state": 1}),
            types.Transition(3, 0, 1.0, 0.0, 4, {"state": 2}),
        ),
    ),
    dict(
        testcase_name="OneStepExtrasZeroes",
        n_step=1,
        additional_discount=1.0,
        first=dm_env.restart(1),
        steps=(
            (
                0,
                dm_env.transition(reward=0.0, observation=2),
                np.zeros(1),
            ),
            (
                0,
                dm_env.transition(reward=0.0, observation=3),
                np.zeros(1),
            ),
            (
                0,
                dm_env.termination(reward=1.0, observation=4),
                np.zeros(1),
            ),
        ),
        expected_transitions=(
            types.Transition(1, 0, 0.0, 1.0, 2, np.zeros(1)),
            types.Transition(2, 0, 0.0, 1.0, 3, np.zeros(1)),
            types.Transition(3, 0, 1.0, 0.0, 4, np.zeros(1)),
        ),
    ),
    dict(
        testcase_name="TwoStep",
        n_step=2,
        additional_discount=1.0,
        first=dm_env.restart(1),
        steps=(
            (0, dm_env.transition(reward=1.0, observation=2, discount=0.5)),
            (0, dm_env.transition(reward=1.0, observation=3, discount=0.5)),
            (0, dm_env.termination(reward=1.0, observation=4)),
        ),
        expected_transitions=(
            types.Transition(1, 0, 1.0, 0.50, 2),
            types.Transition(1, 0, 1.5, 0.25, 3),
            types.Transition(2, 0, 1.5, 0.00, 4),
            types.Transition(3, 0, 1.0, 0.00, 4),
        ),
    ),
    dict(
        testcase_name="TwoStepStructuredReward",
        n_step=2,
        additional_discount=1.0,
        first=dm_env.restart(1),
        steps=(
            (0, dm_env.transition(reward=(1.0, 2.0), observation=2, discount=0.5)),
            (0, dm_env.transition(reward=(1.0, 2.0), observation=3, discount=0.5)),
            (0, dm_env.termination(reward=(1.0, 2.0), observation=4)),
        ),
        expected_transitions=(
            types.Transition(1, 0, (1.0, 2.0), (0.50, 0.50), 2),
            types.Transition(1, 0, (1.5, 3.0), (0.25, 0.25), 3),
            types.Transition(2, 0, (1.5, 3.0), (0.00, 0.00), 4),
            types.Transition(3, 0, (1.0, 2.0), (0.00, 0.00), 4),
        ),
    ),
    dict(
        testcase_name="TwoStepNDArrayReward",
        n_step=2,
        additional_discount=1.0,
        first=dm_env.restart(1),
        steps=(
            (
                0,
                dm_env.transition(
                    reward=np.array((1.0, 2.0)), observation=2, discount=0.5
                ),
            ),
            (
                0,
                dm_env.transition(
                    reward=np.array((1.0, 2.0)), observation=3, discount=0.5
                ),
            ),
            (0, dm_env.termination(reward=np.array((1.0, 2.0)), observation=4)),
        ),
        expected_transitions=(
            types.Transition(1, 0, np.array((1.0, 2.0)), np.array((0.50, 0.50)), 2),
            types.Transition(1, 0, np.array((1.5, 3.0)), np.array((0.25, 0.25)), 3),
            types.Transition(2, 0, np.array((1.5, 3.0)), np.array((0.00, 0.00)), 4),
            types.Transition(3, 0, np.array((1.0, 2.0)), np.array((0.00, 0.00)), 4),
        ),
    ),
    dict(
        testcase_name="TwoStepStructuredDiscount",
        n_step=2,
        additional_discount=1.0,
        first=dm_env.restart(1),
        steps=(
            (
                0,
                dm_env.transition(
                    reward=1.0, observation=2, discount={"a": 0.5, "b": 0.1}
                ),
            ),
            (
                0,
                dm_env.transition(
                    reward=1.0, observation=3, discount={"a": 0.5, "b": 0.1}
                ),
            ),
            (
                0,
                dm_env.termination(reward=1.0, observation=4)._replace(
                    discount={"a": 0.0, "b": 0.0}
                ),
            ),
        ),
        expected_transitions=(
            types.Transition(1, 0, {"a": 1.0, "b": 1.0}, {"a": 0.50, "b": 0.10}, 2),
            types.Transition(1, 0, {"a": 1.5, "b": 1.1}, {"a": 0.25, "b": 0.01}, 3),
            types.Transition(2, 0, {"a": 1.5, "b": 1.1}, {"a": 0.00, "b": 0.00}, 4),
            types.Transition(3, 0, {"a": 1.0, "b": 1.0}, {"a": 0.00, "b": 0.00}, 4),
        ),
    ),
    dict(
        testcase_name="TwoStepNDArrayDiscount",
        n_step=2,
        additional_discount=1.0,
        first=dm_env.restart(1),
        steps=(
            (
                0,
                dm_env.transition(
                    reward=1.0, observation=2, discount=np.array((0.5, 0.1))
                ),
            ),
            (
                0,
                dm_env.transition(
                    reward=1.0, observation=3, discount=np.array((0.5, 0.1))
                ),
            ),
            (
                0,
                dm_env.termination(reward=1.0, observation=4)._replace(
                    discount=np.array((0.0, 0.0))
                ),
            ),
        ),
        expected_transitions=(
            types.Transition(1, 0, np.array((1.0, 1.0)), np.array((0.50, 0.10)), 2),
            types.Transition(1, 0, np.array((1.5, 1.1)), np.array((0.25, 0.01)), 3),
            types.Transition(2, 0, np.array((1.5, 1.1)), np.array((0.00, 0.00)), 4),
            types.Transition(3, 0, np.array((1.0, 1.0)), np.array((0.00, 0.00)), 4),
        ),
    ),
    dict(
        testcase_name="TwoStepBroadcastedNDArrays",
        n_step=2,
        additional_discount=1.0,
        first=dm_env.restart(1),
        steps=(
            (
                0,
                dm_env.transition(
                    reward=np.array([[1.0, 2.0]]),
                    observation=2,
                    discount=np.array([[0.5], [0.1]]),
                ),
            ),
            (
                0,
                dm_env.transition(
                    reward=np.array([[1.0, 2.0]]),
                    observation=3,
                    discount=np.array([[0.5], [0.1]]),
                ),
            ),
            (
                0,
                dm_env.termination(
                    reward=np.array([[1.0, 2.0]]), observation=4
                )._replace(discount=np.array([[0.0], [0.0]])),
            ),
        ),
        expected_transitions=(
            types.Transition(
                1, 0, np.array([[1.0, 2.0], [1.0, 2.0]]), np.array([[0.50], [0.10]]), 2
            ),
            types.Transition(
                1, 0, np.array([[1.5, 3.0], [1.1, 2.2]]), np.array([[0.25], [0.01]]), 3
            ),
            types.Transition(
                2, 0, np.array([[1.5, 3.0], [1.1, 2.2]]), np.array([[0.00], [0.00]]), 4
            ),
            types.Transition(
                3, 0, np.array([[1.0, 2.0], [1.0, 2.0]]), np.array([[0.00], [0.00]]), 4
            ),
        ),
    ),
    dict(
        testcase_name="TwoStepStructuredBroadcastedNDArrays",
        n_step=2,
        additional_discount=1.0,
        first=dm_env.restart(1),
        steps=(
            (
                0,
                dm_env.transition(
                    reward={"a": np.array([[1.0, 2.0]])},
                    observation=2,
                    discount=np.array([[0.5], [0.1]]),
                ),
            ),
            (
                0,
                dm_env.transition(
                    reward={"a": np.array([[1.0, 2.0]])},
                    observation=3,
                    discount=np.array([[0.5], [0.1]]),
                ),
            ),
            (
                0,
                dm_env.termination(
                    reward={"a": np.array([[1.0, 2.0]])}, observation=4
                )._replace(discount=np.array([[0.0], [0.0]])),
            ),
        ),
        expected_transitions=(
            types.Transition(
                1,
                0,
                {"a": np.array([[1.0, 2.0], [1.0, 2.0]])},
                {"a": np.array([[0.50], [0.10]])},
                2,
            ),
            types.Transition(
                1,
                0,
                {"a": np.array([[1.5, 3.0], [1.1, 2.2]])},
                {"a": np.array([[0.25], [0.01]])},
                3,
            ),
            types.Transition(
                2,
                0,
                {"a": np.array([[1.5, 3.0], [1.1, 2.2]])},
                {"a": np.array([[0.00], [0.00]])},
                4,
            ),
            types.Transition(
                3,
                0,
                {"a": np.array([[1.0, 2.0], [1.0, 2.0]])},
                {"a": np.array([[0.00], [0.00]])},
                4,
            ),
        ),
    ),
    dict(
        testcase_name="TwoStepWithExtras",
        n_step=2,
        additional_discount=1.0,
        first=dm_env.restart(1),
        steps=(
            (
                0,
                dm_env.transition(reward=1.0, observation=2, discount=0.5),
                {"state": 0},
            ),
            (
                0,
                dm_env.transition(reward=1.0, observation=3, discount=0.5),
                {"state": 1},
            ),
            (
                0,
                dm_env.termination(reward=1.0, observation=4),
                {"state": 2},
            ),
        ),
        expected_transitions=(
            types.Transition(1, 0, 1.0, 0.50, 2, {"state": 0}),
            types.Transition(1, 0, 1.5, 0.25, 3, {"state": 0}),
            types.Transition(2, 0, 1.5, 0.00, 4, {"state": 1}),
            types.Transition(3, 0, 1.0, 0.00, 4, {"state": 2}),
        ),
    ),
    dict(
        testcase_name="ThreeStepDiscounted",
        n_step=3,
        additional_discount=0.4,
        first=dm_env.restart(1),
        steps=(
            (0, dm_env.transition(reward=1.0, observation=2, discount=0.5)),
            (0, dm_env.transition(reward=1.0, observation=3, discount=0.5)),
            (0, dm_env.termination(reward=1.0, observation=4)),
        ),
        expected_transitions=(
            types.Transition(1, 0, 1.00, 0.5, 2),
            types.Transition(1, 0, 1.20, 0.1, 3),
            types.Transition(1, 0, 1.24, 0.0, 4),
            types.Transition(2, 0, 1.20, 0.0, 4),
            types.Transition(3, 0, 1.00, 0.0, 4),
        ),
    ),
    dict(
        testcase_name="ThreeStepVaryingReward",
        n_step=3,
        additional_discount=0.5,
        first=dm_env.restart(1),
        steps=(
            (0, dm_env.transition(reward=2.0, observation=2)),
            (0, dm_env.transition(reward=3.0, observation=3)),
            (0, dm_env.transition(reward=5.0, observation=4)),
            (0, dm_env.termination(reward=7.0, observation=5)),
        ),
        expected_transitions=(
            types.Transition(1, 0, 2, 1.00, 2),
            types.Transition(1, 0, 2 + 0.5 * 3, 0.50, 3),
            types.Transition(1, 0, 2 + 0.5 * 3 + 0.25 * 5, 0.25, 4),
            types.Transition(2, 0, 3 + 0.5 * 5 + 0.25 * 7, 0.00, 5),
            types.Transition(3, 0, 5 + 0.5 * 7, 0.00, 5),
            types.Transition(4, 0, 7, 0.00, 5),
        ),
    ),
    dict(
        testcase_name="SingleTransitionEpisode",
        n_step=4,
        additional_discount=1.0,
        first=dm_env.restart(1),
        steps=((0, dm_env.termination(reward=1.0, observation=2)),),
        expected_transitions=(types.Transition(1, 0, 1.00, 0.0, 2),),
    ),
    dict(
        testcase_name="EpisodeShorterThanN",
        n_step=4,
        additional_discount=1.0,
        first=dm_env.restart(1),
        steps=(
            (0, dm_env.transition(reward=1.0, observation=2)),
            (0, dm_env.termination(reward=1.0, observation=3)),
        ),
        expected_transitions=(
            types.Transition(1, 0, 1.00, 1.0, 2),
            types.Transition(1, 0, 2.00, 0.0, 3),
            types.Transition(2, 0, 1.00, 0.0, 3),
        ),
    ),
    dict(
        testcase_name="EpisodeEqualToN",
        n_step=3,
        additional_discount=1.0,
        first=dm_env.restart(1),
        steps=(
            (0, dm_env.transition(reward=1.0, observation=2)),
            (0, dm_env.termination(reward=1.0, observation=3)),
        ),
        expected_transitions=(
            types.Transition(1, 0, 1.00, 1.0, 2),
            types.Transition(1, 0, 2.00, 0.0, 3),
            types.Transition(2, 0, 1.00, 0.0, 3),
        ),
    ),
]

BASE_TEST_CASES_FOR_SEQUENCE_ADDER = [
    dict(
        testcase_name="PeriodOne",
        sequence_length=3,
        period=1,
        first=dm_env.restart(1),
        steps=(
            (0, dm_env.transition(reward=2.0, observation=2)),
            (0, dm_env.transition(reward=3.0, observation=3)),
            (0, dm_env.transition(reward=5.0, observation=4)),
            (0, dm_env.termination(reward=7.0, observation=5)),
        ),
        expected_sequences=(
            # (observation, action, reward, discount, start_of_episode, extra)
            [
                (1, 0, 2.0, 1.0, True, ()),
                (2, 0, 3.0, 1.0, False, ()),
                (3, 0, 5.0, 1.0, False, ()),
            ],
            [
                (2, 0, 3.0, 1.0, False, ()),
                (3, 0, 5.0, 1.0, False, ()),
                (4, 0, 7.0, 0.0, False, ()),
            ],
            [
                (3, 0, 5.0, 1.0, False, ()),
                (4, 0, 7.0, 0.0, False, ()),
                (5, 0, 0.0, 0.0, False, ()),
            ],
        ),
        end_behavior=sequence_adder.EndBehavior.ZERO_PAD,
    ),
    dict(
        testcase_name="PeriodTwo",
        sequence_length=3,
        period=2,
        first=dm_env.restart(1),
        steps=(
            (0, dm_env.transition(reward=2.0, observation=2)),
            (0, dm_env.transition(reward=3.0, observation=3)),
            (0, dm_env.transition(reward=5.0, observation=4)),
            (0, dm_env.termination(reward=7.0, observation=5)),
        ),
        expected_sequences=(
            # (observation, action, reward, discount, start_of_episode, extra)
            [
                (1, 0, 2.0, 1.0, True, ()),
                (2, 0, 3.0, 1.0, False, ()),
                (3, 0, 5.0, 1.0, False, ()),
            ],
            [
                (3, 0, 5.0, 1.0, False, ()),
                (4, 0, 7.0, 0.0, False, ()),
                (5, 0, 0.0, 0.0, False, ()),
            ],
        ),
        end_behavior=sequence_adder.EndBehavior.ZERO_PAD,
    ),
    dict(
        testcase_name="EarlyTerminationPeriodOne",
        sequence_length=3,
        period=1,
        first=dm_env.restart(1),
        steps=(
            (0, dm_env.transition(reward=2.0, observation=2)),
            (0, dm_env.termination(reward=3.0, observation=3)),
        ),
        expected_sequences=(
            # (observation, action, reward, discount, start_of_episode, extra)
            [
                (1, 0, 2.0, 1.0, True, ()),
                (2, 0, 3.0, 0.0, False, ()),
                (3, 0, 0.0, 0.0, False, ()),
            ],
        ),
        end_behavior=sequence_adder.EndBehavior.ZERO_PAD,
    ),
    dict(
        testcase_name="EarlyTerminationPeriodTwo",
        sequence_length=3,
        period=2,
        first=dm_env.restart(1),
        steps=(
            (0, dm_env.transition(reward=2.0, observation=2)),
            (0, dm_env.termination(reward=3.0, observation=3)),
        ),
        expected_sequences=(
            # (observation, action, reward, discount, start_of_episode, extra)
            [
                (1, 0, 2.0, 1.0, True, ()),
                (2, 0, 3.0, 0.0, False, ()),
                (3, 0, 0.0, 0.0, False, ()),
            ],
        ),
        end_behavior=sequence_adder.EndBehavior.ZERO_PAD,
    ),
    dict(
        testcase_name="EarlyTerminationPaddingPeriodOne",
        sequence_length=4,
        period=1,
        first=dm_env.restart(1),
        steps=(
            (0, dm_env.transition(reward=2.0, observation=2)),
            (0, dm_env.termination(reward=3.0, observation=3)),
        ),
        expected_sequences=(
            # (observation, action, reward, discount, start_of_episode, extra)
            [
                (1, 0, 2.0, 1.0, True, ()),
                (2, 0, 3.0, 0.0, False, ()),
                (3, 0, 0.0, 0.0, False, ()),
                (0, 0, 0.0, 0.0, False, ()),
            ],
        ),
        end_behavior=sequence_adder.EndBehavior.ZERO_PAD,
    ),
    dict(
        testcase_name="EarlyTerminationPaddingPeriodTwo",
        sequence_length=4,
        period=2,
        first=dm_env.restart(1),
        steps=(
            (0, dm_env.transition(reward=2.0, observation=2)),
            (0, dm_env.termination(reward=3.0, observation=3)),
        ),
        expected_sequences=(
            # (observation, action, reward, discount, start_of_episode, extra)
            [
                (1, 0, 2.0, 1.0, True, ()),
                (2, 0, 3.0, 0.0, False, ()),
                (3, 0, 0.0, 0.0, False, ()),
                (0, 0, 0.0, 0.0, False, ()),
            ],
        ),
        end_behavior=sequence_adder.EndBehavior.ZERO_PAD,
    ),
    dict(
        testcase_name="EarlyTerminationNoPadding",
        sequence_length=4,
        period=1,
        first=dm_env.restart(1),
        steps=(
            (0, dm_env.transition(reward=2.0, observation=2)),
            (0, dm_env.termination(reward=3.0, observation=3)),
        ),
        expected_sequences=(
            # (observation, action, reward, discount, start_of_episode, extra)
            [
                (1, 0, 2.0, 1.0, True, ()),
                (2, 0, 3.0, 0.0, False, ()),
                (3, 0, 0.0, 0.0, False, ()),
            ],
        ),
        end_behavior=sequence_adder.EndBehavior.TRUNCATE,
    ),
    dict(
        testcase_name="LongEpisodePadding",
        sequence_length=3,
        period=3,
        first=dm_env.restart(1),
        steps=(
            (0, dm_env.transition(reward=2.0, observation=2)),
            (0, dm_env.transition(reward=3.0, observation=3)),
            (0, dm_env.transition(reward=5.0, observation=4)),
            (0, dm_env.transition(reward=7.0, observation=5)),
            (0, dm_env.transition(reward=9.0, observation=6)),
            (0, dm_env.transition(reward=11.0, observation=7)),
            (0, dm_env.termination(reward=13.0, observation=8)),
        ),
        expected_sequences=(
            # (observation, action, reward, discount, start_of_episode, extra)
            [
                (1, 0, 2.0, 1.0, True, ()),
                (2, 0, 3.0, 1.0, False, ()),
                (3, 0, 5.0, 1.0, False, ()),
            ],
            [
                (4, 0, 7.0, 1.0, False, ()),
                (5, 0, 9.0, 1.0, False, ()),
                (6, 0, 11.0, 1.0, False, ()),
            ],
            [
                (7, 0, 13.0, 0.0, False, ()),
                (8, 0, 0.0, 0.0, False, ()),
                (0, 0, 0.0, 0.0, False, ()),
            ],
        ),
        end_behavior=sequence_adder.EndBehavior.ZERO_PAD,
    ),
    dict(
        testcase_name="LongEpisodeNoPadding",
        sequence_length=3,
        period=3,
        first=dm_env.restart(1),
        steps=(
            (0, dm_env.transition(reward=2.0, observation=2)),
            (0, dm_env.transition(reward=3.0, observation=3)),
            (0, dm_env.transition(reward=5.0, observation=4)),
            (0, dm_env.transition(reward=7.0, observation=5)),
            (0, dm_env.transition(reward=9.0, observation=6)),
            (0, dm_env.transition(reward=11.0, observation=7)),
            (0, dm_env.termination(reward=13.0, observation=8)),
        ),
        expected_sequences=(
            # (observation, action, reward, discount, start_of_episode, extra)
            [
                (1, 0, 2.0, 1.0, True, ()),
                (2, 0, 3.0, 1.0, False, ()),
                (3, 0, 5.0, 1.0, False, ()),
            ],
            [
                (4, 0, 7.0, 1.0, False, ()),
                (5, 0, 9.0, 1.0, False, ()),
                (6, 0, 11.0, 1.0, False, ()),
            ],
            [
                (7, 0, 13.0, 0.0, False, ()),
                (8, 0, 0.0, 0.0, False, ()),
            ],
        ),
        end_behavior=sequence_adder.EndBehavior.TRUNCATE,
    ),
    dict(
        testcase_name="EndBehavior_WRITE",
        sequence_length=3,
        period=2,
        first=dm_env.restart(1),
        steps=(
            (0, dm_env.transition(reward=2.0, observation=2)),
            (0, dm_env.transition(reward=3.0, observation=3)),
            (0, dm_env.transition(reward=5.0, observation=4)),
            (0, dm_env.transition(reward=7.0, observation=5)),
            (0, dm_env.termination(reward=8.0, observation=6)),
        ),
        expected_sequences=(
            # (observation, action, reward, discount, start_of_episode, extra)
            [
                (1, 0, 2.0, 1.0, True, ()),
                (2, 0, 3.0, 1.0, False, ()),
                (3, 0, 5.0, 1.0, False, ()),
            ],
            [
                (3, 0, 5.0, 1.0, False, ()),
                (4, 0, 7.0, 1.0, False, ()),
                (5, 0, 8.0, 0.0, False, ()),
            ],
            [
                (4, 0, 7.0, 1.0, False, ()),
                (5, 0, 8.0, 0.0, False, ()),
                (6, 0, 0.0, 0.0, False, ()),
            ],
        ),
        end_behavior=sequence_adder.EndBehavior.WRITE,
    ),
]

TEST_CASES_FOR_SEQUENCE_ADDER = BASE_TEST_CASES_FOR_SEQUENCE_ADDER + [
    dict(
        testcase_name="NonBreakingSequenceOnEpisodeReset",
        sequence_length=3,
        period=2,
        first=dm_env.restart(1),
        steps=(
            (0, dm_env.transition(reward=2.0, observation=2)),
            (0, dm_env.transition(reward=3.0, observation=3)),
            (0, dm_env.transition(reward=5.0, observation=4)),
            (0, dm_env.transition(reward=7.0, observation=5)),
            (0, dm_env.transition(reward=9.0, observation=6)),
            (0, dm_env.transition(reward=11.0, observation=7)),
            (0, dm_env.termination(reward=13.0, observation=8)),
        ),
        expected_sequences=(
            # (observation, action, reward, discount, start_of_episode, extra)
            [
                (1, 0, 2.0, 1.0, True, ()),
                (2, 0, 3.0, 1.0, False, ()),
                (3, 0, 5.0, 1.0, False, ()),
            ],
            [
                (3, 0, 5.0, 1.0, False, ()),
                (4, 0, 7.0, 1.0, False, ()),
                (5, 0, 9.0, 1.0, False, ()),
            ],
            [
                (5, 0, 9.0, 1.0, False, ()),
                (6, 0, 11.0, 1.0, False, ()),
                (7, 0, 13.0, 0.0, False, ()),
            ],
        ),
        end_behavior=sequence_adder.EndBehavior.CONTINUE,
        repeat_episode_times=1,
    ),
    dict(
        testcase_name="NonBreakingSequenceMultipleTerminatedEpisodes",
        sequence_length=3,
        period=2,
        first=dm_env.restart(1),
        steps=(
            (0, dm_env.transition(reward=2.0, observation=2)),
            (0, dm_env.transition(reward=3.0, observation=3)),
            (0, dm_env.transition(reward=5.0, observation=4)),
            (0, dm_env.transition(reward=7.0, observation=5)),
            (0, dm_env.transition(reward=9.0, observation=6)),
            (0, dm_env.termination(reward=13.0, observation=7)),
        ),
        expected_sequences=(
            # (observation, action, reward, discount, start_of_episode, extra)
            [
                (1, 0, 2.0, 1.0, True, ()),
                (2, 0, 3.0, 1.0, False, ()),
                (3, 0, 5.0, 1.0, False, ()),
            ],
            [
                (3, 0, 5.0, 1.0, False, ()),
                (4, 0, 7.0, 1.0, False, ()),
                (5, 0, 9.0, 1.0, False, ()),
            ],
            [
                (5, 0, 9.0, 1.0, False, ()),
                (6, 0, 13.0, 0.0, False, ()),
                (7, 0, 0.0, 0.0, False, ()),
            ],
            [
                (7, 0, 0.0, 0.0, False, ()),
                (1, 0, 2.0, 1.0, True, ()),
                (2, 0, 3.0, 1.0, False, ()),
            ],
            [
                (2, 0, 3.0, 1.0, False, ()),
                (3, 0, 5.0, 1.0, False, ()),
                (4, 0, 7.0, 1.0, False, ()),
            ],
            [
                (4, 0, 7.0, 1.0, False, ()),
                (5, 0, 9.0, 1.0, False, ()),
                (6, 0, 13.0, 0.0, False, ()),
            ],
            [
                (6, 0, 13.0, 0.0, False, ()),
                (7, 0, 0.0, 0.0, False, ()),
                (1, 0, 2.0, 1.0, True, ()),
            ],
            [
                (1, 0, 2.0, 1.0, True, ()),
                (2, 0, 3.0, 1.0, False, ()),
                (3, 0, 5.0, 1.0, False, ()),
            ],
            [
                (3, 0, 5.0, 1.0, False, ()),
                (4, 0, 7.0, 1.0, False, ()),
                (5, 0, 9.0, 1.0, False, ()),
            ],
            [
                (5, 0, 9.0, 1.0, False, ()),
                (6, 0, 13.0, 0.0, False, ()),
                (7, 0, 0.0, 0.0, False, ()),
            ],
        ),
        end_behavior=sequence_adder.EndBehavior.CONTINUE,
        repeat_episode_times=3,
    ),
    dict(
        testcase_name="NonBreakingSequenceMultipleTruncatedEpisodes",
        sequence_length=3,
        period=2,
        first=dm_env.restart(1),
        steps=(
            (0, dm_env.transition(reward=2.0, observation=2)),
            (0, dm_env.transition(reward=3.0, observation=3)),
            (0, dm_env.transition(reward=5.0, observation=4)),
            (0, dm_env.transition(reward=7.0, observation=5)),
            (0, dm_env.transition(reward=9.0, observation=6)),
            (0, dm_env.truncation(reward=13.0, observation=7)),
        ),
        expected_sequences=(
            # (observation, action, reward, discount, start_of_episode, extra)
            [
                (1, 0, 2.0, 1.0, True, ()),
                (2, 0, 3.0, 1.0, False, ()),
                (3, 0, 5.0, 1.0, False, ()),
            ],
            [
                (3, 0, 5.0, 1.0, False, ()),
                (4, 0, 7.0, 1.0, False, ()),
                (5, 0, 9.0, 1.0, False, ()),
            ],
            [
                (5, 0, 9.0, 1.0, False, ()),
                (6, 0, 13.0, 1.0, False, ()),
                (7, 0, 0.0, 0.0, False, ()),
            ],
            [
                (7, 0, 0.0, 0.0, False, ()),
                (1, 0, 2.0, 1.0, True, ()),
                (2, 0, 3.0, 1.0, False, ()),
            ],
            [
                (2, 0, 3.0, 1.0, False, ()),
                (3, 0, 5.0, 1.0, False, ()),
                (4, 0, 7.0, 1.0, False, ()),
            ],
            [
                (4, 0, 7.0, 1.0, False, ()),
                (5, 0, 9.0, 1.0, False, ()),
                (6, 0, 13.0, 1.0, False, ()),
            ],
            [
                (6, 0, 13.0, 1.0, False, ()),
                (7, 0, 0.0, 0.0, False, ()),
                (1, 0, 2.0, 1.0, True, ()),
            ],
            [
                (1, 0, 2.0, 1.0, True, ()),
                (2, 0, 3.0, 1.0, False, ()),
                (3, 0, 5.0, 1.0, False, ()),
            ],
            [
                (3, 0, 5.0, 1.0, False, ()),
                (4, 0, 7.0, 1.0, False, ()),
                (5, 0, 9.0, 1.0, False, ()),
            ],
            [
                (5, 0, 9.0, 1.0, False, ()),
                (6, 0, 13.0, 1.0, False, ()),
                (7, 0, 0.0, 0.0, False, ()),
            ],
        ),
        end_behavior=sequence_adder.EndBehavior.CONTINUE,
        repeat_episode_times=3,
    ),
]

corax/adders/reverb/test_cases.py
ethanluoyc/corax
# 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.

"""Utilities for testing Reverb adders."""

from typing import Any, Callable, Optional, Sequence, Tuple, TypeVar, Union

import dm_env
import numpy as np
import reverb
import tensorflow as tf
import tree
from absl.testing import absltest

from corax import specs
from corax import types
from corax.adders import base as adders_base
from corax.adders import reverb as adders
from corax.utils import tree_utils

StepWithExtra = Tuple[Any, dm_env.TimeStep, Any]
StepWithoutExtra = Tuple[Any, dm_env.TimeStep]
Step = TypeVar("Step", StepWithExtra, StepWithoutExtra)


def make_trajectory(observations):
    """Make a simple trajectory from a sequence of observations.

    Arguments:
      observations: a sequence of observations.

    Returns:
      a tuple (first, steps) where first contains the initial dm_env.TimeStep
      object and steps contains a list of (action, step) tuples. The length of
      steps is given by episode_length.
    """
    first = dm_env.restart(observations[0])
    middle = [
        (0, dm_env.transition(reward=0.0, observation=observation))
        for observation in observations[1:-1]
    ]
    last = (0, dm_env.termination(reward=0.0, observation=observations[-1]))
    return first, middle + [last]


def make_sequence(observations):
    """Create a sequence of timesteps of the form `first, [second, ..., last]`."""
    first, steps = make_trajectory(observations)
    observation = first.observation
    sequence = []
    start_of_episode = True
    for action, timestep in steps:
        extras = ()
        sequence.append(
            (
                observation,
                action,
                timestep.reward,
                timestep.discount,
                start_of_episode,
                extras,
            )
        )
        observation = timestep.observation
        start_of_episode = False
    sequence.append((observation, 0, 0.0, 0.0, False, ()))
    return sequence


def _numeric_to_spec(x: Union[float, int, np.ndarray]):
    if isinstance(x, np.ndarray):
        return specs.Array(shape=x.shape, dtype=x.dtype)
    elif isinstance(x, (float, int)):
        return specs.Array(shape=(), dtype=type(x))
    else:
        raise ValueError(f"Unsupported numeric: {type(x)}")


def get_specs(step):
    """Infer spec from an example step."""
    env_spec = tree.map_structure(
        _numeric_to_spec,
        specs.EnvironmentSpec(
            observations=step[1].observation,
            actions=step[0],
            rewards=step[1].reward,
            discounts=step[1].discount,
        ),
    )

    has_extras = len(step) == 3
    if has_extras:
        extras_spec = tree.map_structure(_numeric_to_spec, step[2])
    else:
        extras_spec = ()

    return env_spec, extras_spec


class AdderTestMixin(absltest.TestCase):
    """A helper mixin for testing Reverb adders.

    Note that any test inheriting from this mixin must also inherit from something
    that provides the Python unittest assert methods.
    """

    server: reverb.Server
    client: reverb.Client

    @classmethod
    def setUpClass(cls):
        super().setUpClass()

        replay_table = reverb.Table.queue(adders.DEFAULT_PRIORITY_TABLE, 1000)
        cls.server = reverb.Server([replay_table])
        cls.client = reverb.Client(f"localhost:{cls.server.port}")

    def tearDown(self):
        self.client.reset(adders.DEFAULT_PRIORITY_TABLE)
        super().tearDown()

    @classmethod
    def tearDownClass(cls):
        cls.server.stop()
        super().tearDownClass()

    def num_episodes(self):
        info = self.client.server_info(1)[adders.DEFAULT_PRIORITY_TABLE]
        return info.num_episodes

    def num_items(self):
        info = self.client.server_info(1)[adders.DEFAULT_PRIORITY_TABLE]
        return info.current_size

    def items(self):
        sampler = self.client.sample(
            table=adders.DEFAULT_PRIORITY_TABLE,
            num_samples=self.num_items(),
            emit_timesteps=False,
        )
        return [sample.data for sample in sampler]  # pytype: disable=attribute-error

    def run_test_adder(
        self,
        adder: adders_base.Adder,
        first: dm_env.TimeStep,
        steps: Sequence[Step],
        expected_items: Sequence[Any],
        signature: types.NestedSpec,
        pack_expected_items: bool = False,
        stack_sequence_fields: bool = True,
        repeat_episode_times: int = 1,
        end_behavior: adders.EndBehavior = adders.EndBehavior.ZERO_PAD,
        item_transform: Optional[Callable[[Sequence[np.ndarray]], Any]] = None,
    ):
        """Runs a unit test case for the adder.

        Args:
          adder: The instance of `Adder` that is being tested.
          first: The first `dm_env.TimeStep` that is used to call
            `Adder.add_first()`.
          steps: A sequence of (action, timestep) tuples that are passed to
            `Adder.add()`.
          expected_items: The sequence of items that are expected to be created
            by calling the adder's `add_first()` method on `first` and `add()` on
            all of the elements in `steps`.
          signature: Signature that written items must be compatible with.
          pack_expected_items: Deprecated and not used. If true the expected items
            are given unpacked and need to be packed in a list before comparison.
          stack_sequence_fields: Whether to stack the sequence fields of the
            expected items before comparing to the observed items. Usually False
            for transition adders and True for both episode and sequence adders.
          repeat_episode_times: How many times to run an episode.
          end_behavior: How end of episode should be handled.
          item_transform: Transformation of item simulating the work done by the
            dataset pipeline on the learner in a real setup.
        """

        del pack_expected_items

        if not steps:
            raise ValueError("At least one step must be given.")

        has_extras = len(steps[0]) == 3
        for _ in range(repeat_episode_times):
            # Add all the data up to the final step.
            adder.add_first(first)
            for step in steps[:-1]:
                action, ts = step[0], step[1]

                if has_extras:
                    extras = step[2]
                else:
                    extras = ()

                adder.add(action, next_timestep=ts, extras=extras)

            # Add the final step.
            adder.add(*steps[-1])

        # Force run the destructor to trigger the flushing of all pending items.
        getattr(adder, "__del__", lambda: None)()

        # Ending the episode should close the writer. No new writer should yet have
        # been created as it is constructed lazily.
        if end_behavior is not adders.EndBehavior.CONTINUE:
            self.assertEqual(self.num_episodes(), repeat_episode_times)

        # Make sure our expected and observed data match.
        observed_items = self.items()

        # Check matching number of items.
        self.assertEqual(len(expected_items), len(observed_items))

        # Check items are matching according to numpy's almost_equal.
        for expected_item, observed_item in zip(expected_items, observed_items):
            if stack_sequence_fields:
                expected_item = tree_utils.stack_sequence_fields(expected_item)

            # Apply the transformation which would be done by the dataset in a real
            # setup.
            if item_transform:
                observed_item = item_transform(observed_item)

            tree.map_structure(
                np.testing.assert_array_almost_equal,
                tree.flatten(expected_item),
                tree.flatten(observed_item),
            )

        # Make sure the signature matches was is being written by Reverb.
        def _check_signature(spec: tf.TensorSpec, value: np.ndarray):
            self.assertTrue(spec.is_compatible_with(tf.convert_to_tensor(value)))

        # Check that it is possible to unpack observed using the signature.
        for item in observed_items:
            tree.map_structure(
                _check_signature, tree.flatten(signature), tree.flatten(item)
            )

corax/adders/reverb/test_utils.py
ethanluoyc/corax
# 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.

"""Transition adders.

This implements an N-step transition adder which collapses trajectory sequences
into a single transition, simplifying to a simple transition adder when N=1.
"""

import copy
from typing import Optional, Tuple

import numpy as np
import reverb
import tree

from corax import specs
from corax import types
from corax.adders.reverb import base
from corax.adders.reverb import utils
from corax.utils import tree_utils


class NStepTransitionAdder(base.ReverbAdder):
    """An N-step transition adder.

    This will buffer a sequence of N timesteps in order to form a single N-step
    transition which is added to reverb for future retrieval.

    For N=1 the data added to replay will be a standard one-step transition which
    takes the form:

          (s_t, a_t, r_t, d_t, s_{t+1}, e_t)

    where:

      s_t = state observation at time t
      a_t = the action taken from s_t
      r_t = reward ensuing from action a_t
      d_t = environment discount ensuing from action a_t. This discount is
          applied to future rewards after r_t.
      e_t [Optional] = extra data that the agent persists in replay.

    For N greater than 1, transitions are of the form:

          (s_t, a_t, R_{t:t+n}, D_{t:t+n}, s_{t+N}, e_t),

    where:

      s_t = State (observation) at time t.
      a_t = Action taken from state s_t.
      g = the additional discount, used by the agent to discount future returns.
      R_{t:t+n} = N-step discounted return, i.e. accumulated over N rewards:
            R_{t:t+n} := r_t + g * d_t * r_{t+1} + ...
                             + g^{n-1} * d_t * ... * d_{t+n-2} * r_{t+n-1}.
      D_{t:t+n}: N-step product of agent discounts g_i and environment
        "discounts" d_i.
            D_{t:t+n} := g^{n-1} * d_{t} * ... * d_{t+n-1},
        For most environments d_i is 1 for all steps except the last,
        i.e. it is the episode termination signal.
      s_{t+n}: The "arrival" state, i.e. the state at time t+n.
      e_t [Optional]: A nested structure of any 'extras' the user wishes to add.

    Notes:
      - At the beginning and end of episodes, shorter transitions are added.
        That is, at the beginning of the episode, it will add:
              (s_0 -> s_1), (s_0 -> s_2), ..., (s_0 -> s_n), (s_1 -> s_{n+1})

        And at the end of the episode, it will add:
              (s_{T-n+1} -> s_T), (s_{T-n+2} -> s_T), ... (s_{T-1} -> s_T).
      - We add the *first* `extra` of each transition, not the *last*, i.e.
          if extras are provided, we get e_t, not e_{t+n}.
    """

    def __init__(
        self,
        client: reverb.Client,
        n_step: int,
        discount: float,
        *,
        priority_fns: Optional[base.PriorityFnMapping] = None,
        max_in_flight_items: int = 5,
    ):
        """Creates an N-step transition adder.

        Args:
          client: A `reverb.Client` to send the data to replay through.
          n_step: The "N" in N-step transition. See the class docstring for the
            precise definition of what an N-step transition is. `n_step` must be at
            least 1, in which case we use the standard one-step transition, i.e.
            (s_t, a_t, r_t, d_t, s_t+1, e_t).
          discount: Discount factor to apply. This corresponds to the agent's
            discount in the class docstring.
          priority_fns: See docstring for BaseAdder.
          max_in_flight_items: The maximum number of items allowed to be "in flight"
            at the same time. See `block_until_num_items` in
            `reverb.TrajectoryWriter.flush` for more info.

        Raises:
          ValueError: If n_step is less than 1.
        """
        # Makes the additional discount a float32, which means that it will be
        # upcast if rewards/discounts are float64 and left alone otherwise.
        self.n_step = n_step
        self._discount = tree.map_structure(np.float32, discount)
        self._first_idx = 0
        self._last_idx = 0

        super().__init__(
            client=client,
            max_sequence_length=n_step + 1,
            priority_fns=priority_fns,
            max_in_flight_items=max_in_flight_items,
        )

    def add(self, *args, **kwargs):
        # Increment the indices for the start and end of the window for computing
        # n-step returns.
        if self._writer.episode_steps >= self.n_step:
            self._first_idx += 1
        self._last_idx += 1

        super().add(*args, **kwargs)

    def reset(
        self,
    ):  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
        super().reset()
        self._first_idx = 0
        self._last_idx = 0

    @property
    def _n_step(self) -> int:
        """Effective n-step, which may vary at starts and ends of episodes."""
        return self._last_idx - self._first_idx

    def _write(self):
        # Convenient getters for use in tree operations.
        get_first = lambda x: x[self._first_idx]
        get_last = lambda x: x[self._last_idx]
        # Note: this getter is meant to be used on a TrajectoryWriter.history to
        # obtain its numpy values.
        get_all_np = lambda x: x[self._first_idx : self._last_idx].numpy()

        # Get the state, action, next_state, as well as possibly extras for the
        # transition that is about to be written.
        history = self._writer.history
        s, a = tree.map_structure(
            get_first, (history["observation"], history["action"])
        )
        s_ = tree.map_structure(get_last, history["observation"])

        # Maybe get extras to add to the transition later.
        if "extras" in history:
            extras = tree.map_structure(get_first, history["extras"])

        # Note: at the beginning of an episode we will add the initial N-1
        # transitions (of size 1, 2, ...) and at the end of an episode (when
        # called from write_last) we will write the final transitions of size (N,
        # N-1, ...). See the Note in the docstring.
        # Get numpy view of the steps to be fed into the priority functions.
        reward, discount = tree.map_structure(
            get_all_np, (history["reward"], history["discount"])
        )

        # Compute discounted return and geometric discount over n steps.
        n_step_return, total_discount = self._compute_cumulative_quantities(
            reward, discount
        )

        # Append the computed n-step return and total discount.
        # Note: if this call to _write() is within a call to _write_last(), then
        # this is the only data being appended and so it is not a partial append.
        self._writer.append(
            dict(n_step_return=n_step_return, total_discount=total_discount),
            partial_step=self._writer.episode_steps <= self._last_idx,
        )
        # This should be done immediately after self._writer.append so the history
        # includes the recently appended data.
        history = self._writer.history

        # Form the n-step transition by using the following:
        # the first observation and action in the buffer, along with the cumulative
        # reward and discount computed above.
        n_step_return, total_discount = tree.map_structure(
            lambda x: x[-1], (history["n_step_return"], history["total_discount"])
        )
        transition = types.Transition(
            observation=s,
            action=a,
            reward=n_step_return,
            discount=total_discount,
            next_observation=s_,
            extras=(extras if "extras" in history else ()),
        )

        # Calculate the priority for this transition.
        table_priorities = utils.calculate_priorities(self._priority_fns, transition)

        # Insert the transition into replay along with its priority.
        for table, priority in table_priorities.items():
            self._writer.create_item(
                table=table, priority=priority, trajectory=transition
            )
            self._writer.flush(self._max_in_flight_items)

    def _write_last(self):
        # Write the remaining shorter transitions by alternating writing and
        # incrementingfirst_idx. Note that last_idx will no longer be incremented
        # once we're in this method's scope.
        self._first_idx += 1
        while self._first_idx < self._last_idx:
            self._write()
            self._first_idx += 1

    def _compute_cumulative_quantities(
        self, rewards: types.NestedArray, discounts: types.NestedArray
    ) -> Tuple[types.NestedArray, types.NestedArray]:
        # Give the same tree structure to the n-step return accumulator,
        # n-step discount accumulator, and self.discount, so that they can be
        # iterated in parallel using tree.map_structure.
        rewards, discounts, self_discount = tree_utils.broadcast_structures(
            rewards, discounts, self._discount
        )
        flat_rewards = tree.flatten(rewards)
        flat_discounts = tree.flatten(discounts)
        flat_self_discount = tree.flatten(self_discount)

        # Copy total_discount as it is otherwise read-only.
        total_discount = [np.copy(a[0]) for a in flat_discounts]

        # Broadcast n_step_return to have the broadcasted shape of
        # reward * discount.
        n_step_return = [
            np.copy(np.broadcast_to(r[0], np.broadcast(r[0], d).shape))
            for r, d in zip(flat_rewards, total_discount)
        ]

        # NOTE: total_discount will have one less self_discount applied to it than
        # the value of self._n_step. This is so that when the learner/update uses
        # an additional discount we don't apply it twice. Inside the following loop
        # we will apply this right before summing up the n_step_return.
        for i in range(1, self._n_step):
            for nsr, td, r, d, sd in zip(
                n_step_return,
                total_discount,
                flat_rewards,
                flat_discounts,
                flat_self_discount,
            ):
                # Equivalent to: `total_discount *= self._discount`.
                td *= sd
                # Equivalent to: `n_step_return += reward[i] * total_discount`.
                nsr += r[i] * td
                # Equivalent to: `total_discount *= discount[i]`.
                td *= d[i]

        n_step_return = tree.unflatten_as(rewards, n_step_return)
        total_discount = tree.unflatten_as(rewards, total_discount)
        return n_step_return, total_discount

    # TODO(bshahr): make this into a standalone method. Class methods should be
    # used as alternative constructors or when modifying some global state,
    # neither of which is done here.
    @classmethod
    def signature(
        cls, environment_spec: specs.EnvironmentSpec, extras_spec: types.NestedSpec = ()
    ):
        # This function currently assumes that self._discount is a scalar.
        # If it ever becomes a nested structure and/or a np.ndarray, this method
        # will need to know its structure / shape. This is because the signature
        # discount shape is the environment's discount shape and this adder's
        # discount shape broadcasted together. Also, the reward shape is this
        # signature discount shape broadcasted together with the environment
        # reward shape. As long as self._discount is a scalar, it will not affect
        # either the signature discount shape nor the signature reward shape, so we
        # can ignore it.

        rewards_spec, step_discounts_spec = tree_utils.broadcast_structures(
            environment_spec.rewards, environment_spec.discounts
        )
        rewards_spec = tree.map_structure(
            _broadcast_specs, rewards_spec, step_discounts_spec
        )
        step_discounts_spec = tree.map_structure(copy.deepcopy, step_discounts_spec)

        transition_spec = types.Transition(
            environment_spec.observations,
            environment_spec.actions,
            rewards_spec,
            step_discounts_spec,
            environment_spec.observations,  # next_observation
            extras_spec,
        )

        return tree.map_structure_with_path(
            base.spec_like_to_tensor_spec, transition_spec
        )


def _broadcast_specs(*args: specs.Array) -> specs.Array:
    """Like np.broadcast, but for specs.Array.

    Args:
      *args: one or more specs.Array instances.

    Returns:
      A specs.Array with the broadcasted shape and dtype of the specs in *args.
    """
    bc_info = np.broadcast(*tuple(a.generate_value() for a in args))
    dtype = np.result_type(*tuple(a.dtype for a in args))
    return specs.Array(shape=bc_info.shape, dtype=dtype)

corax/adders/reverb/transition.py
ethanluoyc/corax
# 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.

"""Tests for NStepTransition adders."""

from absl.testing import absltest
from absl.testing import parameterized

from corax.adders.reverb import test_cases
from corax.adders.reverb import test_utils
from corax.adders.reverb import transition as adders


class NStepTransitionAdderTest(test_utils.AdderTestMixin, parameterized.TestCase):
    @parameterized.named_parameters(*test_cases.TEST_CASES_FOR_TRANSITION_ADDER)
    def test_adder(
        self, n_step, additional_discount, first, steps, expected_transitions
    ):
        adder = adders.NStepTransitionAdder(self.client, n_step, additional_discount)
        super().run_test_adder(
            adder=adder,
            first=first,
            steps=steps,
            expected_items=expected_transitions,
            stack_sequence_fields=False,
            signature=adder.signature(*test_utils.get_specs(steps[0])),
        )


if __name__ == "__main__":
    absltest.main()

corax/adders/reverb/transition_test.py
ethanluoyc/corax
# 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.

"""Utilities for Reverb-based adders."""

from typing import Dict, Union

import jax
import jax.numpy as jnp
import numpy as np
import tree

from corax import types
from corax.adders.reverb import base


def zeros_like(x: Union[np.ndarray, int, float, np.number]):
    """Returns a zero-filled object of the same (d)type and shape as the input.

    The difference between this and `np.zeros_like()` is that this works well
    with `np.number`, `int`, `float`, and `jax.numpy.DeviceArray` objects without
    converting them to `np.ndarray`s.

    Args:
      x: The object to replace with 0s.

    Returns:
      A zero-filed object of the same (d)type and shape as the input.
    """
    if isinstance(x, (int, float, np.number)):
        return type(x)(0)
    elif isinstance(x, jax.Array):
        return jnp.zeros_like(x)
    elif isinstance(x, np.ndarray):
        return np.zeros_like(x)
    else:
        raise ValueError(
            f"Input ({type(x)}) must be either a numpy array, an int, or a float."
        )


def final_step_like(step: base.Step, next_observation: types.NestedArray) -> base.Step:
    """Return a list of steps with the final step zero-filled."""
    # Make zero-filled components so we can fill out the last step.
    zero_action, zero_reward, zero_discount, zero_extras = tree.map_structure(
        zeros_like, (step.action, step.reward, step.discount, step.extras)
    )

    # Return a final step that only has next_observation.
    return base.Step(
        observation=next_observation,
        action=zero_action,
        reward=zero_reward,
        discount=zero_discount,
        start_of_episode=False,
        extras=zero_extras,
    )


def calculate_priorities(
    priority_fns: base.PriorityFnMapping,
    trajectory_or_transition: Union[base.Trajectory, types.Transition],
) -> Dict[str, float]:
    """Helper used to calculate the priority of a Trajectory or Transition.

    This helper converts the leaves of the Trajectory or Transition from
    `reverb.TrajectoryColumn` objects into numpy arrays. The converted Trajectory
    or Transition is then passed into each of the functions in `priority_fns`.

    Args:
      priority_fns: a mapping from table names to priority functions (i.e. a
        callable of type PriorityFn). The given function will be used to generate
        the priority (a float) for the given table.
      trajectory_or_transition: the trajectory or transition used to compute
        priorities.

    Returns:
      A dictionary mapping from table names to the priority (a float) for the
      given collection Trajectory or Transition.
    """
    if any([priority_fn is not None for priority_fn in priority_fns.values()]):
        trajectory_or_transition = tree.map_structure(
            lambda col: col.numpy(), trajectory_or_transition
        )

    return {
        table: (priority_fn(trajectory_or_transition) if priority_fn else 1.0)
        for table, priority_fn in priority_fns.items()
    }

corax/adders/reverb/utils.py
ethanluoyc/corax
# 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.

"""A library of useful adder wrappers."""

from typing import Iterable

import dm_env

from corax import types
from corax.adders import base


class ForkingAdder(base.Adder):
    """An adder that forks data into several other adders."""

    def __init__(self, adders: Iterable[base.Adder]):
        self._adders = adders

    def reset(self):
        for adder in self._adders:
            adder.reset()

    def add_first(self, timestep: dm_env.TimeStep):
        for adder in self._adders:
            adder.add_first(timestep)

    def add(
        self,
        action: types.NestedArray,
        next_timestep: dm_env.TimeStep,
        extras: types.NestedArray = (),
    ):
        for adder in self._adders:
            adder.add(action, next_timestep, extras)


class IgnoreExtrasAdder(base.Adder):
    """An adder that ignores extras."""

    def __init__(self, adder: base.Adder):
        self._adder = adder

    def reset(self):
        self._adder.reset()

    def add_first(self, timestep: dm_env.TimeStep):
        self._adder.add_first(timestep)

    def add(
        self,
        action: types.NestedArray,
        next_timestep: dm_env.TimeStep,
        extras: types.NestedArray = (),
    ):
        self._adder.add(action, next_timestep)

corax/adders/wrappers.py
ethanluoyc/corax
# 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.

"""The base agent interface."""

import math
from typing import List, Optional, Sequence

import dm_env
import numpy as np
import reverb

from corax import core
from corax import types


def _calculate_num_learner_steps(
    num_observations: int, min_observations: int, observations_per_step: float
) -> int:
    """Calculates the number of learner steps to do at step=num_observations."""
    n = num_observations - min_observations
    if n < 0:
        # Do not do any learner steps until you have seen min_observations.
        return 0
    if observations_per_step > 1:
        # One batch every 1/obs_per_step observations, otherwise zero.
        return int(n % int(observations_per_step) == 0)
    else:
        # Always return 1/obs_per_step batches every observation.
        return int(1 / observations_per_step)


class Agent(core.Actor, core.VariableSource):
    """Agent class which combines acting and learning.

    This provides an implementation of the `Actor` interface which acts and
    learns. It takes as input instances of both `acme.Actor` and `acme.Learner`
    classes, and implements the policy, observation, and update methods which
    defer to the underlying actor and learner.

    The only real logic implemented by this class is that it controls the number
    of observations to make before running a learner step. This is done by
    passing the number of `min_observations` to use and a ratio of
    `observations_per_step` := num_actor_actions / num_learner_steps.

    Note that the number of `observations_per_step` can also be in the range[0, 1]
    in order to allow the agent to take more than 1 learner step per action.
    """

    def __init__(
        self,
        actor: core.Actor,
        learner: core.Learner,
        min_observations: Optional[int] = None,
        observations_per_step: Optional[float] = None,
        iterator: Optional[core.PrefetchingIterator] = None,
        replay_tables: Optional[List[reverb.Table]] = None,
    ):
        self._actor = actor
        self._learner = learner
        self._min_observations = min_observations
        self._observations_per_step = observations_per_step
        self._num_observations = 0
        self._iterator = iterator
        self._replay_tables = replay_tables
        self._batch_size_upper_bounds = (
            [1_000_000_000] * len(replay_tables) if replay_tables else None
        )

    def select_action(self, observation: types.NestedArray) -> types.NestedArray:
        return self._actor.select_action(observation)

    def observe_first(self, timestep: dm_env.TimeStep):
        self._actor.observe_first(timestep)

    def observe(self, action: types.NestedArray, next_timestep: dm_env.TimeStep):
        self._num_observations += 1
        self._actor.observe(action, next_timestep)

    def _has_data_for_training(self):
        if self._iterator.ready():  # type: ignore
            return True
        for table, batch_size in zip(
            self._replay_tables, self._batch_size_upper_bounds  # type: ignore
        ):
            if not table.can_sample(batch_size):
                return False
        return True

    def update(
        self,
    ):  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
        if self._iterator:
            # Perform learner steps as long as iterator has data.
            update_actor = False
            while self._has_data_for_training():
                # Run learner steps (usually means gradient steps).
                total_batches = self._iterator.retrieved_elements()
                self._learner.step()
                current_batches = self._iterator.retrieved_elements() - total_batches
                assert current_batches == 1, (
                    "Learner step must retrieve exactly one element from the iterator"
                    f" (retrieved {current_batches}). Otherwise agent can deadlock. "
                    "Example cause is that your chosen agent"
                    "s Builder has a "
                    "`make_learner` factory that prefetches the data but it "
                    "shouldn"
                    "t."
                )
                self._batch_size_upper_bounds = [
                    math.ceil(
                        t.info.rate_limiter_info.sample_stats.completed
                        / (total_batches + 1)
                    )
                    for t in self._replay_tables  # type: ignore
                ]
                update_actor = True
            if update_actor:
                # Update the actor weights only when learner was updated.
                self._actor.update()
            return

        # If dataset is not provided, follback to the old logic.
        # TODO(stanczyk): Remove when not used.
        num_steps = _calculate_num_learner_steps(
            num_observations=self._num_observations,
            min_observations=self._min_observations,  # type: ignore
            observations_per_step=self._observations_per_step,  # type: ignore
        )
        for _ in range(num_steps):
            # Run learner steps (usually means gradient steps).
            self._learner.step()
        if num_steps > 0:
            # Update the actor weights when learner updates.
            self._actor.update()

    def get_variables(self, names: Sequence[str]) -> List[List[np.ndarray]]:
        return self._learner.get_variables(names)

corax/agents/agent.py
ethanluoyc/corax
# 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.

"""ActorCore interface definition."""

import dataclasses
from typing import Callable, Generic, Mapping, Tuple, TypeVar, Union

import chex
import jax
import jax.numpy as jnp

from corax import types
from corax.jax import networks as networks_lib
from corax.jax import utils
from corax.jax.types import PRNGKey

NoneType = type(None)
# The state of the actor. This could include recurrent network state or any
# other state which needs to be propagated through the select_action calls.
State = TypeVar("State")
# The extras to be passed to the observe method.
Extras = TypeVar("Extras")
RecurrentState = TypeVar("RecurrentState")
SelectActionFn = Callable[
    [networks_lib.Params, networks_lib.Observation, State],
    Tuple[networks_lib.Action, State],
]


@dataclasses.dataclass
class ActorCore(Generic[State, Extras]):
    """Pure functions that define the algorithm-specific actor functionality."""

    init: Callable[[PRNGKey], State]
    select_action: SelectActionFn
    get_extras: Callable[[State], Extras]


# A simple feed forward policy which produces no extras and takes only an RNGKey
# as a state.
FeedForwardPolicy = Callable[
    [networks_lib.Params, PRNGKey, networks_lib.Observation], networks_lib.Action
]

FeedForwardPolicyWithExtra = Callable[
    [networks_lib.Params, PRNGKey, networks_lib.Observation],
    Tuple[networks_lib.Action, types.NestedArray],
]


Policy = Union[FeedForwardPolicy, FeedForwardPolicyWithExtra]


def batched_feed_forward_to_actor_core(
    policy: FeedForwardPolicy,
) -> ActorCore[PRNGKey, Tuple[()]]:
    """A convenience adaptor from FeedForwardPolicy to ActorCore."""

    def select_action(
        params: networks_lib.Params,
        observation: networks_lib.Observation,
        state: PRNGKey,
    ):
        rng = state
        rng1, rng2 = jax.random.split(rng)
        observation = utils.add_batch_dim(observation)
        action = utils.squeeze_batch_dim(policy(params, rng1, observation))
        return action, rng2

    def init(rng: PRNGKey) -> PRNGKey:
        return rng

    def get_extras(unused_rng: PRNGKey) -> Tuple[()]:
        return ()

    return ActorCore(init=init, select_action=select_action, get_extras=get_extras)


@chex.dataclass(frozen=True, mappable_dataclass=False)
class SimpleActorCoreStateWithExtras:
    rng: PRNGKey
    extras: Mapping[str, jnp.ndarray]


def unvectorize_select_action(actor_core: ActorCore) -> ActorCore:
    """Makes an actor core's select_action method expect unbatched arguments."""

    def unvectorized_select_action(
        params: networks_lib.Params,
        observations: networks_lib.Observation,
        state: State,
    ) -> Tuple[networks_lib.Action, State]:
        observations, state = utils.add_batch_dim((observations, state))
        actions, state = actor_core.select_action(params, observations, state)
        return utils.squeeze_batch_dim((actions, state))

    return ActorCore(
        init=actor_core.init,
        select_action=unvectorized_select_action,
        get_extras=actor_core.get_extras,
    )

corax/agents/jax/actor_core.py
ethanluoyc/corax
# 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.

"""Simple JAX actors."""

from typing import Generic, Optional

import dm_env
import jax

from corax import adders
from corax import core
from corax import types
from corax.agents.jax import actor_core
from corax.jax import networks as network_lib
from corax.jax import utils
from corax.jax import variable_utils


class GenericActor(core.Actor, Generic[actor_core.State, actor_core.Extras]):
    """A generic actor implemented on top of ActorCore.

    An actor based on a policy which takes observations and outputs actions. It
    also adds experiences to replay and updates the actor weights from the policy
    on the learner.
    """

    def __init__(
        self,
        actor: actor_core.ActorCore[actor_core.State, actor_core.Extras],
        random_key: network_lib.PRNGKey,
        variable_client: Optional[variable_utils.VariableClient],
        adder: Optional[adders.Adder] = None,
        jit: bool = True,
        backend: Optional[str] = "cpu",
        per_episode_update: bool = False,
    ):
        """Initializes a feed forward actor.

        Args:
          actor: actor core.
          random_key: Random key.
          variable_client: The variable client to get policy parameters from.
          adder: An adder to add experiences to.
          jit: Whether or not to jit the passed ActorCore's pure functions.
          backend: Which backend to use when jitting the policy.
          per_episode_update: if True, updates variable client params once at the
            beginning of each episode
        """
        self._random_key = random_key
        self._variable_client = variable_client
        self._adder = adder
        self._state = None

        # Unpack ActorCore, jitting if requested.
        if jit:
            self._init = jax.jit(actor.init, backend=backend)
            self._policy = jax.jit(actor.select_action, backend=backend)
        else:
            self._init = actor.init
            self._policy = actor.select_action
        self._get_extras = actor.get_extras
        self._per_episode_update = per_episode_update

    @property
    def _params(self):
        return self._variable_client.params if self._variable_client else []

    def select_action(self, observation: network_lib.Observation) -> types.NestedArray:
        action, self._state = self._policy(self._params, observation, self._state)
        return utils.to_numpy(action)

    def observe_first(self, timestep: dm_env.TimeStep):
        self._random_key, key = jax.random.split(self._random_key)
        self._state = self._init(key)
        if self._adder:
            self._adder.add_first(timestep)
        if self._variable_client and self._per_episode_update:
            self._variable_client.update_and_wait()

    def observe(self, action: network_lib.Action, next_timestep: dm_env.TimeStep):
        if self._adder:
            self._adder.add(action, next_timestep, extras=self._get_extras(self._state))  # type: ignore

    def update(self, wait: bool = False):
        if self._variable_client and not self._per_episode_update:
            self._variable_client.update(wait)

corax/agents/jax/actors.py
ethanluoyc/corax
# 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.

"""RL agent Builder interface."""

import abc
import dataclasses
from typing import Generic, Iterator, List, Optional, TypeVar

import reverb

from corax import adders
from corax import core
from corax import specs
from corax.jax import networks as networks_lib
from corax.utils import counting
from corax.utils import loggers

Networks = TypeVar("Networks")
Policy = TypeVar("Policy")
Sample = TypeVar("Sample")


class OfflineBuilder(abc.ABC, Generic[Networks, Policy, Sample]):
    """Interface for defining the components of an offline RL agent.

    Implementations of this interface contain a complete specification of a
    concrete offline RL agent. An instance of this class can be used to build an
    offline RL agent that operates either locally or in a distributed setup.
    """

    @abc.abstractmethod
    def make_learner(
        self,
        random_key: networks_lib.PRNGKey,
        networks: Networks,
        dataset: Iterator[Sample],
        logger_fn: loggers.LoggerFactory,
        environment_spec: specs.EnvironmentSpec,
        *,
        counter: Optional[counting.Counter] = None,
    ) -> core.Learner:
        """Creates an instance of the learner.

        Args:
          random_key: A key for random number generation.
          networks: struct describing the networks needed by the learner; this is
            specific to the learner in question.
          dataset: iterator over demonstration samples.
          logger_fn: factory providing loggers used for logging progress.
          environment_spec: A container for all relevant environment specs.
          counter: a Counter which allows for recording of counts (learner steps,
            evaluator steps, etc.) distributed throughout the agent.
        """

    @abc.abstractmethod
    def make_actor(
        self,
        random_key: networks_lib.PRNGKey,
        policy: Policy,
        environment_spec: specs.EnvironmentSpec,
        variable_source: Optional[core.VariableSource] = None,
    ) -> core.Actor:
        """Create an actor instance to be used for evaluation.

        Args:
          random_key: A key for random number generation.
          policy: Instance of a policy expected by the algorithm corresponding to
            this builder.
          environment_spec: A container for all relevant environment specs.
          variable_source: A source providing the necessary actor parameters.
        """

    @abc.abstractmethod
    def make_policy(
        self,
        networks: Networks,
        environment_spec: specs.EnvironmentSpec,
        evaluation: bool,
    ) -> Policy:
        """Creates the agent policy to be used for evaluation.

        Args:
          networks: struct describing the networks needed to generate the policy.
          environment_spec: struct describing the specs of the environment.
          evaluation: This flag is present for consistency with the
            ActorLearnerBuilder, in which case data-generating actors and evaluation
            actors can behave differently. For OfflineBuilders, this should be set
            to True.

        Returns:
          Policy to be used for evaluation. The exact form of this object may differ
          from one agent to the next; it could be a simple callable, a nest of
          callables, or an ActorCore for instance.
        """


class ActorLearnerBuilder(
    OfflineBuilder[Networks, Policy, Sample], Generic[Networks, Policy, Sample]
):
    """Defines an interface for defining the components of an RL agent.

    Implementations of this interface contain a complete specification of a
    concrete RL agent. An instance of this class can be used to build an
    RL agent which interacts with the environment either locally or in a
    distributed setup.
    """

    @abc.abstractmethod
    def make_replay_tables(
        self,
        environment_spec: specs.EnvironmentSpec,
        policy: Policy,
    ) -> List[reverb.Table]:
        """Create tables to insert data into.

        Args:
          environment_spec: A container for all relevant environment specs.
          policy: Agent's policy which can be used to extract the extras_spec.

        Returns:
          The replay tables used to store the experience the agent uses to train.
        """

    @abc.abstractmethod
    def make_dataset_iterator(
        self,
        replay_client: reverb.Client,
    ) -> Iterator[Sample]:
        """Create a dataset iterator to use for learning/updating the agent."""

    @abc.abstractmethod
    def make_adder(
        self,
        replay_client: reverb.Client,
        environment_spec: Optional[specs.EnvironmentSpec],
        policy: Optional[Policy],
    ) -> Optional[adders.Adder]:
        """Create an adder which records data generated by the actor/environment.

        Args:
          replay_client: Reverb Client which points to the replay server.
          environment_spec: specs of the environment.
          policy: Agent's policy which can be used to extract the extras_spec.
        """
        # TODO(sabela): make the parameters non-optional.

    @abc.abstractmethod
    def make_actor(
        self,
        random_key: networks_lib.PRNGKey,
        policy: Policy,
        environment_spec: specs.EnvironmentSpec,
        variable_source: Optional[core.VariableSource] = None,
        adder: Optional[adders.Adder] = None,
    ) -> core.Actor:
        """Create an actor instance.

        Args:
          random_key: A key for random number generation.
          policy: Instance of a policy expected by the algorithm corresponding to
            this builder.
          environment_spec: A container for all relevant environment specs.
          variable_source: A source providing the necessary actor parameters.
          adder: How data is recorded (e.g. added to replay).
        """

    @abc.abstractmethod
    def make_learner(
        self,
        random_key: networks_lib.PRNGKey,
        networks: Networks,
        dataset: Iterator[Sample],
        logger_fn: loggers.LoggerFactory,
        environment_spec: specs.EnvironmentSpec,
        replay_client: Optional[reverb.Client] = None,
        counter: Optional[counting.Counter] = None,
    ) -> core.Learner:
        """Creates an instance of the learner.

        Args:
          random_key: A key for random number generation.
          networks: struct describing the networks needed by the learner; this can
            be specific to the learner in question.
          dataset: iterator over samples from replay.
          logger_fn: factory providing loggers used for logging progress.
          environment_spec: A container for all relevant environment specs.
          replay_client: client which allows communication with replay. Note that
            this is only intended to be used for updating priorities. Samples should
            be obtained from `dataset`.
          counter: a Counter which allows for recording of counts (learner steps,
            actor steps, etc.) distributed throughout the agent.
        """

    def make_policy(
        self,
        networks: Networks,
        environment_spec: specs.EnvironmentSpec,
        evaluation: bool = False,
    ) -> Policy:
        """Creates the agent policy.

           Creates the agent policy given the collection of network components and
           environment spec. An optional boolean can be given to indicate if the
           policy will be used for evaluation.

        Args:
          networks: struct describing the networks needed to generate the policy.
          environment_spec: struct describing the specs of the environment.
          evaluation: when true, a version of the policy to use for evaluation
            should be returned. This is algorithm-specific so if an algorithm makes
            no distinction between behavior and evaluation policies this boolean may
            be ignored.

        Returns:
          Behavior policy or evaluation policy for the agent.
        """
        # TODO(sabela): make abstract once all agents implement it.
        del networks, environment_spec, evaluation
        raise NotImplementedError


@dataclasses.dataclass(frozen=True)
class ActorLearnerBuilderWrapper(
    ActorLearnerBuilder[Networks, Policy, Sample], Generic[Networks, Policy, Sample]
):
    """An empty wrapper for ActorLearnerBuilder."""

    wrapped: ActorLearnerBuilder[Networks, Policy, Sample]

    def make_replay_tables(
        self,
        environment_spec: specs.EnvironmentSpec,
        policy: Policy,
    ) -> List[reverb.Table]:
        return self.wrapped.make_replay_tables(environment_spec, policy)

    def make_dataset_iterator(
        self,
        replay_client: reverb.Client,
    ) -> Iterator[Sample]:
        return self.wrapped.make_dataset_iterator(replay_client)

    def make_adder(
        self,
        replay_client: reverb.Client,
        environment_spec: Optional[specs.EnvironmentSpec],
        policy: Optional[Policy],
    ) -> Optional[adders.Adder]:
        return self.wrapped.make_adder(replay_client, environment_spec, policy)

    def make_actor(
        self,
        random_key: networks_lib.PRNGKey,
        policy: Policy,
        environment_spec: specs.EnvironmentSpec,
        variable_source: Optional[core.VariableSource] = None,
        adder: Optional[adders.Adder] = None,
    ) -> core.Actor:
        return self.wrapped.make_actor(
            random_key, policy, environment_spec, variable_source, adder
        )

    def make_learner(
        self,
        random_key: networks_lib.PRNGKey,
        networks: Networks,
        dataset: Iterator[Sample],
        logger_fn: loggers.LoggerFactory,
        environment_spec: specs.EnvironmentSpec,
        replay_client: Optional[reverb.Client] = None,
        counter: Optional[counting.Counter] = None,
    ) -> core.Learner:
        return self.wrapped.make_learner(
            random_key,
            networks,
            dataset,
            logger_fn,
            environment_spec,
            replay_client,
            counter,
        )

    def make_policy(
        self,
        networks: Networks,
        environment_spec: specs.EnvironmentSpec,
        evaluation: bool = False,
    ) -> Policy:
        return self.wrapped.make_policy(networks, environment_spec, evaluation)

corax/agents/jax/builders.py
ethanluoyc/corax
from corax.agents.jax.calql.adder import CalQLAdder
from corax.agents.jax.calql.learning import CalQLLearner
from corax.agents.jax.calql.networks import CQLNetworks
from corax.agents.jax.calql.networks import apply_policy_and_sample
from corax.agents.jax.calql.networks import make_networks

corax/agents/jax/calql/__init__.py
ethanluoyc/corax
import dataclasses
from typing import Union

import dm_env
import numpy as np

from corax.adders import Adder


@dataclasses.dataclass
class SparseReward:
    reward_scale: float
    reward_bias: float

    positive_reward: float
    negative_reward: float


@dataclasses.dataclass
class DenseReward:
    reward_scale: float
    reward_bias: float


def compute_return_to_go(
    rewards,
    terminals,
    gamma,
    reward_scale,
    reward_bias,
    is_sparse_reward,
    negative_reward=0.0,
):
    """
    A config dict for getting the default high/low rewrd values for each envs
    This is used in calc_return_to_go func in sampler.py and replay_buffer.py
    """
    if len(rewards) == 0:
        return np.array([])

    reward_neg = negative_reward * reward_scale + reward_bias

    if is_sparse_reward and np.all(np.array(rewards) == reward_neg):
        """
        If the env has sparse reward and the trajectory is all negative rewards,
        we use r / (1-gamma) as return to go.
        For exapmle, if gamma = 0.99 and the rewards = [-1, -1, -1],
        then return_to_go = [-100, -100, -100]
        """
        # assuming failure reward is negative
        # use r / (1-gamma) for negative trajctory
        return_to_go = [float(reward_neg / (1 - gamma))] * len(rewards)
    else:
        return_to_go = [0] * len(rewards)
        prev_return = 0
        for i in range(len(rewards)):
            return_to_go[-i - 1] = rewards[-i - 1] + gamma * prev_return * (
                1 - terminals[-i - 1]
            )
            prev_return = return_to_go[-i - 1]

    return np.array(return_to_go, dtype=np.float32)


class CalQLAdder(Adder):
    def __init__(
        self,
        base_adder: Adder,
        discount,
        reward_config: Union[SparseReward, DenseReward],
    ):
        self._adder = base_adder
        self._discount = discount
        self._reward_config = reward_config
        self._steps = []

    def add_first(self, timestep: dm_env.TimeStep):
        assert timestep.first()
        self._steps = [(None, timestep)]
        # return self._adder.add_first(timestep)

    def add(self, action, next_timestep: dm_env.TimeStep, extras=()):
        del extras
        self._steps.append((action, next_timestep))
        if next_timestep.last():
            dones = [step.last() for (_, step) in self._steps[1:]]
            rewards = [
                (
                    step.reward * self._reward_config.reward_scale
                    + self._reward_config.reward_bias
                )
                for (_, step) in self._steps[1:]
            ]
            if isinstance(self._reward_config, SparseReward):
                mc_returns = compute_return_to_go(
                    rewards,
                    dones,
                    self._discount,
                    self._reward_config.reward_scale,
                    self._reward_config.reward_bias,
                    is_sparse_reward=True,
                    negative_reward=self._reward_config.negative_reward,
                )
            else:
                mc_returns = compute_return_to_go(
                    rewards,
                    dones,
                    self._discount,
                    self._reward_config.reward_scale,
                    self._reward_config.reward_bias,
                    is_sparse_reward=False,
                )

            initial_step = self._steps[0][1]
            self._adder.add_first(initial_step)
            for (action, step), mc_return, reward in zip(
                self._steps[1:], mc_returns, rewards
            ):
                step = step._replace(reward=reward.astype(np.float32))
                self._adder.add(action, step, extras={"mc_return": mc_return})
            self._steps = []

    def reset(self):
        self._steps = []
        self._adder.reset()

corax/agents/jax/calql/adder.py
ethanluoyc/corax
"""Tests for the CQL agent."""

import jax
import optax
from absl.testing import absltest

from corax import specs
from corax.agents.jax.calql import learning
from corax.agents.jax.calql import networks as networks_lib
from corax.testing import fakes


class CQLTest(absltest.TestCase):
    def test_train(self):
        seed = 0
        num_iterations = 6
        batch_size = 64

        # Create a fake environment to test with.
        environment = fakes.ContinuousEnvironment(
            episode_length=10, bounded=True, action_dim=6
        )
        spec = specs.make_environment_spec(environment)

        # Construct the agent.
        networks = networks_lib.make_networks(
            spec,
            policy_hidden_sizes=(8, 8),
            critic_hidden_sizes=(8, 8),
        )
        dataset = fakes.transition_iterator(environment)
        key = jax.random.PRNGKey(seed)
        learner = learning.CalQLLearner(
            batch_size,
            networks,
            key,
            demonstrations=dataset(batch_size),
            policy_optimizer=optax.adam(3e-5),
            critic_optimizer=optax.adam(3e-4),
            fixed_cql_coefficient=5.0,
            cql_lagrange_threshold=None,
            target_entropy=0.1,
            num_bc_iters=2,
            num_sgd_steps_per_step=1,
        )

        # Train the agent
        for _ in range(num_iterations):
            learner.step()


if __name__ == "__main__":
    absltest.main()

corax/agents/jax/calql/agent_test.py
ethanluoyc/corax
# type: ignore
"""CQL learner implementation."""

import time
from typing import Any, Dict, Iterator, List, NamedTuple, Optional, Tuple

import chex
import jax
import jax.numpy as jnp
import optax

import corax
from corax import types
from corax.agents.jax.calql.networks import CQLNetworks
from corax.agents.jax.calql.networks import apply_and_sample_n
from corax.jax import networks as networks_lib
from corax.jax import utils
from corax.utils import counting
from corax.utils import loggers

_CQL_COEFFICIENT_MAX_VALUE = 1e6
_CQL_GRAD_CLIPPING_VALUE = 40


class TrainingState(NamedTuple):
    """Contains training state for the learner."""

    policy_optimizer_state: optax.OptState
    critic_optimizer_state: optax.OptState
    policy_params: networks_lib.Params
    critic_params: networks_lib.Params
    target_critic_params: networks_lib.Params
    key: networks_lib.PRNGKey

    # Optimizer and value of the alpha parameter from SAC (entropy temperature).
    # These fields are only used with an adaptive coefficient (when
    # fixed_entropy_coefficeint is None in the CQLLearner)
    alpha_optimizer_state: Optional[optax.OptState] = None
    log_sac_alpha: Optional[networks_lib.Params] = None

    # Optimizer and value of the alpha parameter from CQL (regularization
    # coefficient).
    # These fields are only used with an adaptive coefficient (when
    # fixed_cql_coefficiennt is None in the CQLLearner)
    cql_optimizer_state: Optional[optax.OptState] = None
    log_cql_alpha: Optional[networks_lib.Params] = None

    steps: int = 0


class CalQLLearner(corax.Learner):
    """CQL learner.

    Learning component of the Conservative Q-Learning algorithm from
    [Kumar et al., 2020] https://arxiv.org/abs/2006.04779.
    """

    _state: TrainingState

    def __init__(
        self,
        batch_size: int,
        networks: CQLNetworks,
        random_key: networks_lib.PRNGKey,
        demonstrations: Iterator[types.Transition],
        policy_optimizer: optax.GradientTransformation,
        critic_optimizer: optax.GradientTransformation,
        tau: float = 0.005,
        fixed_cql_coefficient: Optional[float] = None,
        cql_lagrange_threshold: Optional[float] = None,
        cql_num_samples: int = 10,
        num_sgd_steps_per_step: int = 1,
        reward_scale: float = 1.0,
        discount: float = 0.99,
        fixed_entropy_coefficient: Optional[float] = None,
        target_entropy: Optional[float] = 0,
        num_bc_iters: int = 50_000,
        max_target_backup: bool = False,
        use_calql: bool = False,
        counter: Optional[counting.Counter] = None,
        logger: Optional[loggers.Logger] = None,
    ):
        """Initializes the CQL learner.

        Args:
          batch_size: batch size.
          networks: CQL networks.
          random_key: a key for random number generation.
          demonstrations: an iterator over training data.
          policy_optimizer: the policy optimizer.
          critic_optimizer: the Q-function optimizer.
          tau: target smoothing coefficient.
          fixed_cql_coefficient: the value for cql coefficient. If None, an adaptive
            coefficient will be used.
          cql_lagrange_threshold: a threshold that controls the adaptive loss for
            the cql coefficient.
          cql_num_samples: number of samples used to compute logsumexp(Q) via
            importance sampling.
          num_sgd_steps_per_step: how many gradient updated to perform per batch.
            batch is split into this many smaller batches, thus should be a multiple
            of num_sgd_steps_per_step
          reward_scale: reward scale.
          discount: discount to use for TD updates.
          fixed_entropy_coefficient: coefficient applied to the entropy bonus. If
            None, an adaptative coefficient will be used.
          target_entropy: Target entropy when using adapdative entropy bonus.
          num_bc_iters: Number of BC steps for actor initialization.
          counter: counter object used to keep track of steps.
          logger: logger object to be used by learner.
        """
        self._num_bc_iters = num_bc_iters
        adaptive_entropy_coefficient = fixed_entropy_coefficient is None
        action_spec = networks.environment_specs.actions
        if adaptive_entropy_coefficient:
            # sac_alpha is the temperature parameter that determines the relative
            # importance of the entropy term versus the reward.
            log_sac_alpha = jnp.asarray(0.0, dtype=jnp.float32)
            alpha_optimizer = optax.adam(learning_rate=3e-4)
            alpha_optimizer_state = alpha_optimizer.init(log_sac_alpha)
        else:
            if target_entropy:
                raise ValueError(
                    "target_entropy should not be set when "
                    "fixed_entropy_coefficient is provided"
                )

        adaptive_cql_coefficient = fixed_cql_coefficient is None
        if adaptive_cql_coefficient:
            log_cql_alpha = jnp.asarray(0.0, dtype=jnp.float32)
            cql_optimizer = optax.adam(learning_rate=3e-4)
            cql_optimizer_state = cql_optimizer.init(log_cql_alpha)
        else:
            if cql_lagrange_threshold:
                raise ValueError(
                    "cql_lagrange_threshold should not be set when "
                    "fixed_cql_coefficient is provided"
                )

        def alpha_loss(
            log_sac_alpha: jnp.ndarray,
            policy_params: networks_lib.Params,
            transitions: types.Transition,
            key: jnp.ndarray,
        ) -> jnp.ndarray:
            """Eq 18 from https://arxiv.org/pdf/1812.05905.pdf."""
            dist_params = networks.policy_network.apply(
                policy_params, transitions.observation
            )
            action = networks.sample(dist_params, key)
            log_prob = networks.log_prob(dist_params, action)
            sac_alpha = jnp.exp(log_sac_alpha)
            sac_alpha_loss = sac_alpha * jax.lax.stop_gradient(
                -log_prob - target_entropy
            )
            return jnp.mean(sac_alpha_loss)

        def sac_critic_loss(
            q_old_action: jnp.ndarray,
            policy_params: networks_lib.Params,
            target_critic_params: networks_lib.Params,
            transitions: types.Transition,
            key: networks_lib.PRNGKey,
        ) -> jnp.ndarray:
            """Computes the SAC part of the loss."""

            if max_target_backup:
                next_dist_params = networks.policy_network.apply(
                    policy_params, transitions.next_observation
                )
                next_actions, _ = apply_and_sample_n(
                    key,
                    networks,
                    policy_params,
                    transitions.next_observation,
                    cql_num_samples,
                )
                next_q = batched_critic(
                    next_actions,
                    target_critic_params,
                    transitions.next_observation,
                )
                next_v = jnp.max(jnp.min(next_q, axis=-1), axis=0)
                target_q = jax.lax.stop_gradient(
                    transitions.reward * reward_scale
                    + transitions.discount * discount * next_v
                )
                chex.assert_shape(next_v, (batch_size,))
                chex.assert_shape(target_q, (batch_size,))
                return jnp.mean(
                    jnp.square(q_old_action - jnp.expand_dims(target_q, -1))
                )
            else:
                next_dist_params = networks.policy_network.apply(
                    policy_params, transitions.next_observation
                )
                next_action = networks.sample(next_dist_params, key)
                next_q = networks.critic_network.apply(
                    target_critic_params, transitions.next_observation, next_action
                )
                next_v = jnp.min(next_q, axis=-1)
                target_q = jax.lax.stop_gradient(
                    transitions.reward * reward_scale
                    + transitions.discount * discount * next_v
                )
                return jnp.mean(
                    jnp.square(q_old_action - jnp.expand_dims(target_q, -1))
                )

        def batched_critic(
            actions: jnp.ndarray,
            critic_params: networks_lib.Params,
            observation: jnp.ndarray,
        ) -> jnp.ndarray:
            """Applies the critic network to a batch of sampled actions."""
            actions = jax.lax.stop_gradient(actions)
            tiled_actions = jnp.reshape(actions, (batch_size * cql_num_samples, -1))
            tiled_states = jnp.tile(observation, [cql_num_samples, 1])
            tiled_q = networks.critic_network.apply(
                critic_params, tiled_states, tiled_actions
            )
            return jnp.reshape(tiled_q, (cql_num_samples, batch_size, -1))

        def cql_critic_loss(
            q_old_action: jnp.ndarray,
            critic_params: networks_lib.Params,
            policy_params: networks_lib.Params,
            transitions: types.Transition,
            key: networks_lib.PRNGKey,
        ) -> jnp.ndarray:
            """Computes the CQL part of the loss."""
            # The CQL part of the loss is
            #     logsumexp(Q(s,·)) - Q(s,a),
            # where s is the currrent state, and a the action in the dataset (so
            # Q(s,a) is simply q_old_action.
            # We need to estimate logsumexp(Q). This is done with importance sampling
            # (IS). This function implements the unlabeled equation page 29, Appx. F,
            # in https://arxiv.org/abs/2006.04779.
            # Here, IS is done with the uniform distribution and the policy in the
            # current state s. In their implementation, the authors also add the
            # policy in the transiting state s':
            # https://github.com/aviralkumar2907/CQL/blob/master/d4rl/rlkit/torch/sac/cql.py,
            # (l. 233-236).

            key_policy, key_policy_next, key_uniform = jax.random.split(key, 3)

            def sampled_q(obs, key):
                actions, log_probs = apply_and_sample_n(
                    key, networks, policy_params, obs, cql_num_samples
                )

                # (N, B, Q)
                q = batched_critic(actions, critic_params, transitions.observation)
                # Enable Cal-QL
                if use_calql:
                    mc_returns = transitions.extras["mc_return"]
                    lower_bounds = jnp.reshape(mc_returns, (1, batch_size, 1))
                    q = jnp.maximum(q, lower_bounds)

                return q - jax.lax.stop_gradient(jnp.expand_dims(log_probs, -1))

            # Sample wrt policy in s
            sampled_q_from_policy = sampled_q(transitions.observation, key_policy)

            # Sample wrt policy in s'
            sampled_q_from_policy_next = sampled_q(
                transitions.next_observation, key_policy_next
            )

            # Sample wrt uniform
            actions_uniform = jax.random.uniform(
                key_uniform,
                (cql_num_samples, batch_size) + action_spec.shape,
                minval=action_spec.minimum,
                maxval=action_spec.maximum,
            )
            log_prob_uniform = -jnp.sum(
                jnp.log(action_spec.maximum - action_spec.minimum)
            )
            sampled_q_from_uniform = (
                batched_critic(actions_uniform, critic_params, transitions.observation)
                - log_prob_uniform
            )

            # Combine the samplings
            combined = jnp.concatenate(
                (
                    sampled_q_from_uniform,
                    sampled_q_from_policy,
                    sampled_q_from_policy_next,
                ),
                axis=0,
            )
            # lse_q = jax.nn.logsumexp(combined, axis=0, b=1.0 / (3 * cql_num_samples))
            lse_q = jax.nn.logsumexp(combined, axis=0)

            return jnp.mean(lse_q - q_old_action)

        def critic_loss(
            critic_params: networks_lib.Params,
            policy_params: networks_lib.Params,
            target_critic_params: networks_lib.Params,
            cql_alpha: jnp.ndarray,
            transitions: types.Transition,
            key: networks_lib.PRNGKey,
        ) -> jnp.ndarray:
            """Computes the full critic loss."""
            key_cql, key_sac = jax.random.split(key, 2)
            q_old_action = networks.critic_network.apply(
                critic_params, transitions.observation, transitions.action
            )
            cql_loss = cql_critic_loss(
                q_old_action, critic_params, policy_params, transitions, key_cql
            )
            sac_loss = sac_critic_loss(
                q_old_action, policy_params, target_critic_params, transitions, key_sac
            )
            return cql_alpha * cql_loss + sac_loss

        def cql_lagrange_loss(
            log_cql_alpha: jnp.ndarray,
            critic_params: networks_lib.Params,
            policy_params: networks_lib.Params,
            transitions: types.Transition,
            key: jnp.ndarray,
        ) -> jnp.ndarray:
            """Computes the loss that optimizes the cql coefficient."""
            cql_alpha = jnp.exp(log_cql_alpha)
            q_old_action = networks.critic_network.apply(
                critic_params, transitions.observation, transitions.action
            )
            return -cql_alpha * (
                cql_critic_loss(
                    q_old_action, critic_params, policy_params, transitions, key
                )
                - cql_lagrange_threshold
            )

        def actor_loss(
            policy_params: networks_lib.Params,
            critic_params: networks_lib.Params,
            sac_alpha: jnp.ndarray,
            transitions: types.Transition,
            key: jnp.ndarray,
            in_initial_bc_iters: bool,
        ) -> jnp.ndarray:
            """Computes the loss for the policy."""
            dist_params = networks.policy_network.apply(
                policy_params, transitions.observation
            )
            if in_initial_bc_iters:
                log_prob = networks.log_prob(dist_params, transitions.action)
                actor_loss = -jnp.mean(log_prob)
            else:
                action = networks.sample(dist_params, key)
                log_prob = networks.log_prob(dist_params, action)
                q_action = networks.critic_network.apply(
                    critic_params, transitions.observation, action
                )
                min_q = jnp.min(q_action, axis=-1)
                actor_loss = jnp.mean(sac_alpha * log_prob - min_q)
            return actor_loss

        alpha_grad = jax.value_and_grad(alpha_loss)
        cql_lagrange_grad = jax.value_and_grad(cql_lagrange_loss)
        critic_grad = jax.value_and_grad(critic_loss)
        actor_grad = jax.value_and_grad(actor_loss)

        def update_step(
            state: TrainingState,
            rb_transitions: types.Transition,
            in_initial_bc_iters: bool,
        ) -> Tuple[TrainingState, Dict[str, jnp.ndarray]]:
            key, key_alpha, key_critic, key_actor = jax.random.split(state.key, 4)

            if adaptive_entropy_coefficient:
                alpha_loss, alpha_grads = alpha_grad(
                    state.log_sac_alpha, state.policy_params, rb_transitions, key_alpha
                )
                sac_alpha = jnp.exp(state.log_sac_alpha)
            else:
                sac_alpha = fixed_entropy_coefficient

            if adaptive_cql_coefficient:
                cql_lagrange_loss, cql_lagrange_grads = cql_lagrange_grad(
                    state.log_cql_alpha,
                    state.critic_params,
                    state.policy_params,
                    rb_transitions,
                    key_critic,
                )
                cql_lagrange_grads = jnp.clip(
                    cql_lagrange_grads,
                    -_CQL_GRAD_CLIPPING_VALUE,
                    _CQL_GRAD_CLIPPING_VALUE,
                )
                cql_alpha = jnp.exp(state.log_cql_alpha)
                cql_alpha = jnp.clip(
                    cql_alpha, a_min=0.0, a_max=_CQL_COEFFICIENT_MAX_VALUE
                )
            else:
                cql_alpha = fixed_cql_coefficient

            critic_loss, critic_grads = critic_grad(
                state.critic_params,
                state.policy_params,
                state.target_critic_params,
                cql_alpha,
                rb_transitions,
                key_critic,
            )
            actor_loss, actor_grads = actor_grad(
                state.policy_params,
                state.critic_params,
                sac_alpha,
                rb_transitions,
                key_actor,
                in_initial_bc_iters,
            )

            # Apply policy gradients
            actor_update, policy_optimizer_state = policy_optimizer.update(
                actor_grads, state.policy_optimizer_state
            )
            policy_params = optax.apply_updates(state.policy_params, actor_update)

            # Apply critic gradients
            critic_update, critic_optimizer_state = critic_optimizer.update(
                critic_grads, state.critic_optimizer_state
            )
            critic_params = optax.apply_updates(state.critic_params, critic_update)

            new_target_critic_params = jax.tree_map(
                lambda x, y: x * (1 - tau) + y * tau,
                state.target_critic_params,
                critic_params,
            )

            metrics = {
                "critic_loss": critic_loss,
                "actor_loss": actor_loss,
            }

            new_state = TrainingState(
                policy_optimizer_state=policy_optimizer_state,
                critic_optimizer_state=critic_optimizer_state,
                policy_params=policy_params,
                critic_params=critic_params,
                target_critic_params=new_target_critic_params,
                key=key,
                alpha_optimizer_state=state.alpha_optimizer_state,
                log_sac_alpha=state.log_sac_alpha,
                steps=state.steps + 1,
            )
            if adaptive_entropy_coefficient and (not in_initial_bc_iters):
                # Apply sac_alpha gradients
                alpha_update, alpha_optimizer_state = alpha_optimizer.update(
                    alpha_grads, state.alpha_optimizer_state
                )
                log_sac_alpha = optax.apply_updates(state.log_sac_alpha, alpha_update)
                metrics.update(
                    {
                        "alpha_loss": alpha_loss,
                        "sac_alpha": jnp.exp(log_sac_alpha),
                    }
                )
                new_state = new_state._replace(
                    alpha_optimizer_state=alpha_optimizer_state,
                    log_sac_alpha=log_sac_alpha,
                )
            else:
                metrics["alpha_loss"] = 0.0
                metrics["sac_alpha"] = fixed_cql_coefficient

            if adaptive_cql_coefficient:
                # Apply cql coeff gradients
                cql_update, cql_optimizer_state = cql_optimizer.update(
                    cql_lagrange_grads, state.cql_optimizer_state
                )
                log_cql_alpha = optax.apply_updates(state.log_cql_alpha, cql_update)
                metrics.update(
                    {
                        "cql_lagrange_loss": cql_lagrange_loss,
                        "cql_alpha": jnp.exp(log_cql_alpha),
                    }
                )
                new_state = new_state._replace(
                    cql_optimizer_state=cql_optimizer_state, log_cql_alpha=log_cql_alpha
                )

            return new_state, metrics

        # General learner book-keeping and loggers.
        self._counter = counter or counting.Counter()
        self._logger = logger or loggers.make_default_logger(
            "learner",
            asynchronous=True,
            serialize_fn=utils.fetch_devicearray,
            steps_key=self._counter.get_steps_key(),
        )

        # Iterator on demonstration transitions.
        self._demonstrations = demonstrations

        # Use the JIT compiler.
        update_step_in_initial_bc_iters = utils.process_multiple_batches(
            lambda x, y: update_step(x, y, True), num_sgd_steps_per_step
        )
        update_step_rest = utils.process_multiple_batches(
            lambda x, y: update_step(x, y, False), num_sgd_steps_per_step
        )

        self._update_step_in_initial_bc_iters = jax.jit(update_step_in_initial_bc_iters)
        self._update_step_rest = jax.jit(update_step_rest)
        self._num_sgd_steps_per_step = num_sgd_steps_per_step

        # Create initial state.
        key_policy, key_q, training_state_key = jax.random.split(random_key, 3)
        del random_key
        policy_params = networks.policy_network.init(key_policy)
        policy_optimizer_state = policy_optimizer.init(policy_params)
        critic_params = networks.critic_network.init(key_q)
        critic_optimizer_state = critic_optimizer.init(critic_params)

        self._state = TrainingState(
            policy_optimizer_state=policy_optimizer_state,
            critic_optimizer_state=critic_optimizer_state,
            policy_params=policy_params,
            critic_params=critic_params,
            target_critic_params=critic_params,
            key=training_state_key,
            steps=0,
        )

        if adaptive_entropy_coefficient:
            self._state = self._state._replace(
                alpha_optimizer_state=alpha_optimizer_state, log_sac_alpha=log_sac_alpha
            )
        if adaptive_cql_coefficient:
            self._state = self._state._replace(
                cql_optimizer_state=cql_optimizer_state, log_cql_alpha=log_cql_alpha
            )

        # 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

    def step(self):
        # Get data from replay (dropping extras if any). Note there is no
        # extra data here because we do not insert any into Reverb.
        transitions = next(self._demonstrations)

        counts = self._counter.get_counts()
        if "learner_steps" not in counts:
            cur_step = 0
        else:
            cur_step = counts["learner_steps"]
        in_initial_bc_iters = cur_step < self._num_bc_iters

        if in_initial_bc_iters:
            self._state, metrics = self._update_step_in_initial_bc_iters(
                self._state, transitions
            )
        else:
            self._state, metrics = self._update_step_rest(self._state, transitions)

        # Compute elapsed time.
        timestamp = time.time()
        elapsed_time = timestamp - self._timestamp if self._timestamp else 0
        self._timestamp = timestamp

        # Increment counts and record the current time
        counts = self._counter.increment(steps=1, walltime=elapsed_time)

        if elapsed_time > 0:
            metrics["steps_per_second"] = self._num_sgd_steps_per_step / elapsed_time

        # Attempts to write the logs.
        self._logger.write({**metrics, **counts})

    def get_variables(self, names: List[str]) -> List[Any]:
        variables = {
            "policy": self._state.policy_params,
        }
        return [variables[name] for name in names]

    def save(self) -> TrainingState:
        return self._state

    def restore(self, state: TrainingState):
        self._state = state

corax/agents/jax/calql/learning.py
ethanluoyc/corax
# type: ignore
"""Networks definitions for the CQL agent."""
import dataclasses
from typing import Tuple

import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np

from corax import specs
from corax.agents.jax import actor_core as actor_core_lib
from corax.jax import networks as networks_lib
from corax.jax import utils


@dataclasses.dataclass
class CQLNetworks:
    """Network and pure functions for the CQL agent."""

    policy_network: networks_lib.FeedForwardNetwork
    critic_network: networks_lib.FeedForwardNetwork
    log_prob: networks_lib.LogProbFn
    sample: networks_lib.SampleFn
    sample_eval: networks_lib.SampleFn
    environment_specs: specs.EnvironmentSpec


def apply_and_sample_n(
    key: networks_lib.PRNGKey,
    networks: CQLNetworks,
    params: networks_lib.Params,
    obs: jnp.ndarray,
    num_samples: int,
) -> Tuple[jnp.ndarray, jnp.ndarray]:
    """Applies the policy and samples num_samples actions."""
    dist_params = networks.policy_network.apply(params, obs)
    sampled_actions = jnp.array(
        [
            networks.sample(dist_params, key_n)
            for key_n in jax.random.split(key, num_samples)
        ]
    )
    sampled_log_probs = networks.log_prob(dist_params, sampled_actions)
    return sampled_actions, sampled_log_probs


def make_networks(
    spec: specs.EnvironmentSpec,
    policy_hidden_sizes=(256, 256),
    critic_hidden_sizes=(256, 256, 256),
) -> CQLNetworks:
    num_dimensions = np.prod(spec.actions.shape, dtype=int)

    def _actor_fn(obs):
        network = hk.Sequential(
            [
                hk.nets.MLP(
                    list(policy_hidden_sizes),
                    w_init=hk.initializers.VarianceScaling(1.0, "fan_in", "uniform"),
                    activation=jax.nn.relu,
                    activate_final=True,
                ),
                networks_lib.NormalTanhDistribution(num_dimensions),
            ]
        )
        return network(obs)

    def _critic_fn(obs, action):
        network1 = hk.Sequential(
            [
                hk.nets.MLP(
                    list(critic_hidden_sizes) + [1],
                    w_init=hk.initializers.VarianceScaling(1.0, "fan_in", "uniform"),
                    activation=jax.nn.relu,
                ),
            ]
        )
        network2 = hk.Sequential(
            [
                hk.nets.MLP(
                    list(critic_hidden_sizes) + [1],
                    w_init=hk.initializers.VarianceScaling(1.0, "fan_in", "uniform"),
                    activation=jax.nn.relu,
                ),
            ]
        )
        input_ = jnp.concatenate([obs, action], axis=-1)
        value1 = network1(input_)
        value2 = network2(input_)
        return jnp.concatenate([value1, value2], axis=-1)

    policy = hk.without_apply_rng(hk.transform(_actor_fn))
    critic = hk.without_apply_rng(hk.transform(_critic_fn))

    # Create dummy observations and actions to create network parameters.
    dummy_action = utils.zeros_like(spec.actions)
    dummy_obs = utils.zeros_like(spec.observations)
    dummy_action = utils.add_batch_dim(dummy_action)
    dummy_obs = utils.add_batch_dim(dummy_obs)

    return CQLNetworks(
        policy_network=networks_lib.FeedForwardNetwork(
            lambda key: policy.init(key, dummy_obs), policy.apply
        ),
        critic_network=networks_lib.FeedForwardNetwork(
            lambda key: critic.init(key, dummy_obs, dummy_action), critic.apply
        ),
        log_prob=lambda params, actions: params.log_prob(actions),
        sample=lambda params, key: params.sample(seed=key),
        sample_eval=lambda params, key: params.mode(),
        environment_specs=spec,
    )


def apply_policy_and_sample(
    networks: CQLNetworks, eval_mode: bool = False
) -> actor_core_lib.FeedForwardPolicy:
    """Returns a function that computes actions."""
    sample_fn = networks.sample if not eval_mode else networks.sample_eval
    if not sample_fn:
        raise ValueError("sample function is not provided")

    def apply_and_sample(params, key, obs):
        return sample_fn(networks.policy_network.apply(params, obs), key)

    return apply_and_sample

corax/agents/jax/calql/networks.py
ethanluoyc/corax
"""Decision transformer implementation in JAX."""

from corax.agents.jax.decision_transformer.acting import DecisionTransformerActor
from corax.agents.jax.decision_transformer.builder import DecisionTransformerBuilder
from corax.agents.jax.decision_transformer.config import DecisionTransformerConfig
from corax.agents.jax.decision_transformer.learning import DecisionTransformerLearner
from corax.agents.jax.decision_transformer.networks import DecisionTransformer
from corax.agents.jax.decision_transformer.networks import make_gym_networks

corax/agents/jax/decision_transformer/__init__.py
ethanluoyc/corax
# type: ignore
from typing import Any, Callable, NamedTuple, Optional

import dm_env
import jax
import jax.numpy as jnp
import tree

from corax import specs
from corax.jax import networks as networks_lib
from corax.jax import running_statistics
from corax.jax import utils
from corax.jax import variable_utils


def _append(context, update):
    # Alternatively, this works too.
    # return jnp.concatenate([context[1:], jnp.expand_dims(update, 0)], axis=0)
    new_context = jnp.roll(context, shift=-1, axis=0)
    new_context = new_context.at[-1].set(update)
    return new_context


class ActorState(NamedTuple):
    observations: jnp.ndarray
    actions: jnp.ndarray
    rewards: jnp.ndarray
    returns_to_go: jnp.ndarray
    timesteps: jnp.ndarray
    masks: jnp.ndarray

    running_rtg: jnp.ndarray
    running_time: jnp.ndarray


def initial_state(
    environment_spec: specs.EnvironmentSpec,
    target_return: float,
    context_size: int,
) -> ActorState:
    zero_values = utils.tile_nested(utils.zeros_like(environment_spec), context_size)
    return ActorState(
        observations=zero_values.observations,
        actions=zero_values.actions,
        rewards=zero_values.rewards,
        returns_to_go=zero_values.rewards,
        timesteps=jnp.zeros((context_size,), dtype=jnp.int32),
        masks=jnp.zeros((context_size,), dtype=jnp.bool_),
        running_rtg=target_return,
        running_time=jnp.zeros((), dtype=jnp.int32),
    )


class DecisionTransformerActor:
    def __init__(
        self,
        spec: specs.EnvironmentSpec,
        random_key: networks_lib.PRNGKey,
        forward_fn: Callable[[Any], Any],
        target_return: float,
        context_length: int,
        return_scale: float,
        variable_client: variable_utils.VariableClient,
        observation_mean_std: Optional[running_statistics.NestedMeanStd] = None,
        mode: str = "normal",
    ):
        self._key = random_key
        self._max_length = context_length
        self._forward_fn = forward_fn
        self._variable_client = variable_client
        self._target_return = target_return
        self._state = None

        def init(target_return) -> ActorState:
            return initial_state(spec, target_return / return_scale, context_length)

        def select_action(
            params: networks_lib.Params,
            state: ActorState,
            key: jnp.ndarray,
            observation: networks_lib.Observation,
        ):
            if observation_mean_std is not None:
                observation = running_statistics.normalize(
                    observation, observation_mean_std
                )
            new_state = state._replace(
                observations=_append(state.observations, observation),
                timesteps=_append(state.timesteps, state.running_time),
                returns_to_go=_append(state.returns_to_go, state.running_rtg),
                rewards=_append(state.actions, utils.zeros_like(spec.rewards)),
                actions=_append(state.actions, utils.zeros_like(spec.actions)),
                masks=_append(state.masks, True),
            )
            # Run dt_policy
            inputs = {
                "states": new_state.observations,
                "actions": new_state.actions,
                "timesteps": new_state.timesteps,
                "returns_to_go": jnp.expand_dims(new_state.returns_to_go, axis=-1),
                "attention_mask": new_state.masks,
            }
            inputs = tree.map_structure(lambda x: jnp.expand_dims(x, axis=0), inputs)
            key_forward, key = jax.random.split(key)
            _, action_preds, _ = forward_fn(
                params,
                key_forward,
                inputs["states"],
                inputs["actions"],
                None,
                inputs["returns_to_go"],
                inputs["timesteps"],
                attention_mask=inputs["attention_mask"],
                is_training=False,
            )
            return action_preds[0, -1], (new_state, key)

        def observe(state: ActorState, action: networks_lib.Action, reward: float):
            if mode != "delayed":
                running_rtg = state.running_rtg - reward / return_scale
            else:
                running_rtg = state.running_rtg
            return state._replace(
                actions=state.actions.at[-1].set(action),
                rewards=state.rewards.at[-1].set(reward),
                running_time=state.running_time + 1,
                running_rtg=running_rtg,
            )

        self._init = jax.jit(init)
        self._select_action = jax.jit(select_action)
        self._observe = jax.jit(observe)

    @property
    def _params(self):
        return self._variable_client.params

    def select_action(self, observation: networks_lib.Observation):
        action, (self._state, self._key) = self._select_action(
            self._params, self._state, self._key, observation
        )
        action = utils.to_numpy(action)
        return action

    def observe_first(self, timestep: dm_env.TimeStep):
        del timestep
        self._state = self._init(self._target_return)

    def observe(self, action: networks_lib.Action, next_timestep: dm_env.TimeStep):
        self._state = self._observe(self._state, action, next_timestep.reward)

    def update(self, wait: bool = False):
        self._variable_client.update(wait)

corax/agents/jax/decision_transformer/acting.py
ethanluoyc/corax
import jax
import optax
from absl.testing import absltest

from corax import specs
from corax.agents.jax import decision_transformer
from corax.agents.jax.decision_transformer import dataset as dataset_lib
from corax.testing import fakes


class RLDSDataset(absltest.TestCase):
    def test_dataset_generator(self):
        scale = 1.0
        env = fakes.ContinuousEnvironment(action_dim=2, observation_dim=2, bounded=True)
        env_specs = specs.make_environment_spec(env)

        episode_dataset = fakes.rlds_dataset_from_env_spec(env_specs)
        observation_mean_std = dataset_lib.get_observation_mean_std(episode_dataset)
        dataset = (
            dataset_lib.transform_decision_transformer_input(
                episode_dataset,
                sequence_length=10,
                scale=scale,
                observation_mean_std=observation_mean_std,
            )
            .repeat()
            .batch(10)
        )

        networks = decision_transformer.make_gym_networks(
            env_specs, 32, 2, 2, 0.0, 1000
        )

        optimizer = optax.adam(1e-3)
        key = jax.random.PRNGKey(0)
        learner = decision_transformer.DecisionTransformerLearner(
            model=networks,
            key=key,
            dataset=dataset.as_numpy_iterator(),
            optimizer=optimizer,
        )
        learner.step()


if __name__ == "__main__":
    absltest.main()

corax/agents/jax/decision_transformer/agent_test.py
ethanluoyc/corax
from typing import Any, Iterator, Optional

import optax

from corax import core
from corax import specs
from corax.agents.jax import builders
from corax.agents.jax.decision_transformer import acting
from corax.agents.jax.decision_transformer import config as dt_config
from corax.agents.jax.decision_transformer import learning
from corax.jax import networks as networks_lib
from corax.jax import running_statistics
from corax.jax import variable_utils
from corax.utils import counting
from corax.utils import loggers

DecisionTransformerNetwork = networks_lib.FeedForwardNetwork
DecisionTransformerPolicy = Any
Sample = Any


class DecisionTransformerBuilder(builders.OfflineBuilder):
    def __init__(
        self,
        config: dt_config.DecisionTransformerConfig,
        observation_mean_std: Optional[running_statistics.NestedMeanStd] = None,
        max_num_learner_steps: Optional[int] = None,
        actor_device=None,
    ):
        self._config = config
        self._max_num_learner_steps = max_num_learner_steps
        self._actor_device = actor_device
        self._observation_mean_std = observation_mean_std

    def make_learner(
        self,
        random_key: networks_lib.PRNGKey,
        networks: DecisionTransformerNetwork,
        dataset: Iterator[Sample],
        logger_fn: loggers.LoggerFactory,
        environment_spec: specs.EnvironmentSpec,
        *,
        counter: Optional[counting.Counter] = None,
    ) -> learning.DecisionTransformerLearner:
        del environment_spec
        learning_rate = optax.warmup_cosine_decay_schedule(
            0,
            self._config.learning_rate,
            warmup_steps=self._config.warmup_steps,
            decay_steps=self._max_num_learner_steps,  # type: ignore
        )
        optimizer = optax.chain(
            optax.clip_by_global_norm(self._config.grad_norm_clipping),
            optax.adamw(
                learning_rate=learning_rate,
                weight_decay=self._config.weight_decay,
                mask=self._config.weight_decay_mask,
            ),
        )
        return learning.DecisionTransformerLearner(
            model=networks,
            key=random_key,
            dataset=dataset,
            optimizer=optimizer,
            logger=logger_fn("learner"),
            counter=counter,
        )

    def make_actor(
        self,
        random_key: networks_lib.PRNGKey,
        policy: DecisionTransformerPolicy,
        environment_spec: specs.EnvironmentSpec,
        variable_source: core.VariableSource,
    ) -> acting.DecisionTransformerActor:
        variable_client = variable_utils.VariableClient(
            variable_source, "model", device=self._actor_device
        )
        return acting.DecisionTransformerActor(
            environment_spec,
            random_key=random_key,
            forward_fn=policy,
            context_length=self._config.context_length,
            target_return=self._config.target_return,
            return_scale=self._config.return_scale,
            variable_client=variable_client,
            observation_mean_std=self._observation_mean_std,
            mode=self._config.mode,
        )

    def make_policy(
        self,
        networks: DecisionTransformerNetwork,
        environment_spec: specs.EnvironmentSpec,
        evaluation: bool,
    ) -> DecisionTransformerPolicy:
        del environment_spec, evaluation
        return networks.apply

corax/agents/jax/decision_transformer/builder.py
ethanluoyc/corax
import dataclasses
from typing import Any, Callable

import jax.numpy as jnp
import tree


def exclude_bias_and_normalizers(params):
    def predicate(path, value: jnp.ndarray) -> jnp.ndarray:
        del value
        return path[-1] == "b" or "norm" in path[-2] or path[-1] == "embeddings"  # type: ignore

    return tree.map_structure_with_path(predicate, params)


@dataclasses.dataclass
class DecisionTransformerConfig:
    # DT parameters
    context_length: int
    target_return: float
    return_scale: float
    mode: str = "normal"

    # Optimizer parameters
    learning_rate: float = 1e-4
    warmup_steps: int = 10000
    grad_norm_clipping: float = 0.25
    weight_decay: float = 1e-4
    weight_decay_mask: Callable[[Any], Any] = exclude_bias_and_normalizers

    def __post_init__(self):
        assert self.mode in [
            "normal",
            "delayed",
        ], "mode must be either 'normal' or 'delayed'"

corax/agents/jax/decision_transformer/config.py
ethanluoyc/corax
import functools

import numpy as np
import rlds
import tensorflow as tf

from corax.jax import running_statistics


def get_observation_mean_std(episode_dataset):
    mean, std = rlds.transformations.mean_and_std(
        episode_dataset, get_step_fields=rlds.transformations.sar_fields_mask
    )

    observation_mean = np.asarray(mean["observation"], dtype=np.float32)
    observation_std = np.asarray(std["observation"], dtype=np.float32)

    return running_statistics.NestedMeanStd(mean=observation_mean, std=observation_std)  # type: ignore


def discounted_return(rewards, discounts):
    def discounted_return_fn(acc, reward_discount):
        reward, discount = reward_discount
        return acc * discount + reward

    return tf.scan(
        fn=discounted_return_fn,
        elems=(rewards, discounts),
        reverse=True,
        initializer=0.0,
    )


def transform_decision_transformer_input(
    episode_dataset,
    *,
    sequence_length: int,
    scale: float,
    observation_mean_std,
):
    def add_return_to_go(episode):
        episode_length = episode["steps"].cardinality()
        steps = episode["steps"].batch(episode_length).get_single_element()
        returns = discounted_return(steps["reward"], steps["discount"])
        timesteps = tf.range(episode_length)
        # Scale RTG
        steps["return_to_go"] = returns / scale  # type: ignore
        steps["timestep"] = timesteps
        # TODO(yl): Remove this
        if "info" in steps:
            steps.pop("infos")
        steps = tf.nest.map_structure(lambda x: x[:-1], steps)
        episode["steps"] = tf.data.Dataset.from_tensor_slices(steps)
        return episode

    def _pad_along_axis(x, padded_size, axis=0, value=0):
        pad_width = padded_size - tf.shape(x)[axis]
        if pad_width <= 0:
            return x
        padding = [(0, 0)] * len(x.shape.as_list())
        padding[axis] = (pad_width, 0)
        padded = tf.pad(x, padding, mode="CONSTANT", constant_values=value)
        return padded

    def pad_steps(steps, max_len):
        normalized_observation = (
            steps["observation"] - observation_mean_std.mean
        ) / observation_mean_std.std
        padded_obs = _pad_along_axis(normalized_observation, max_len, 0, 0)
        padded_act = _pad_along_axis(steps["action"], max_len, 0, -10)
        padded_rtg = _pad_along_axis(steps["return_to_go"], max_len, 0, 0)
        padded_discounts = _pad_along_axis(steps["discount"], max_len, 0, 2)
        padded_timesteps = _pad_along_axis(steps["timestep"], max_len, 0, 0)
        mask = _pad_along_axis(
            tf.ones(tf.shape(steps["reward"])[0], dtype=bool), max_len, 0, False  # type: ignore
        )
        return {
            "observation": padded_obs,
            "action": padded_act,
            "return_to_go": tf.expand_dims(padded_rtg, axis=-1),
            "discount": padded_discounts,
            "timestep": padded_timesteps,
            "mask": mask,
        }

    dataset = episode_dataset.map(add_return_to_go).cache()
    dataset = dataset.interleave(
        lambda episode: rlds.transformations.batch(
            episode["steps"], size=sequence_length, shift=1, drop_remainder=False
        ),
        num_parallel_calls=tf.data.AUTOTUNE,
        deterministic=False,
    )

    dataset = dataset.map(functools.partial(pad_steps, max_len=sequence_length))
    return dataset

corax/agents/jax/decision_transformer/dataset.py
ethanluoyc/corax
import time
from typing import NamedTuple, Sequence

import jax
import jax.numpy as jnp
import optax

import corax
from corax.jax import networks as networks_lib
from corax.utils import counting
from corax.utils import loggers


class TrainState(NamedTuple):
    params: networks_lib.Params
    opt_state: optax.OptState
    key: networks_lib.PRNGKey


class DecisionTransformerLearner(corax.Learner):
    def __init__(self, model, key, dataset, optimizer, logger=None, counter=None):
        def loss_fn(params, key, inputs):
            states = inputs["observation"]
            actions = inputs["action"]
            rewards = None
            # discounts = inputs["discount"]
            rtg = inputs["return_to_go"]
            timesteps = inputs["timestep"]
            attention_mask = inputs["mask"]

            action_target = actions
            _, action_preds, _ = model.apply(
                params,
                key,
                states,
                actions,
                rewards,
                rtg,
                timesteps,
                attention_mask=attention_mask,
            )

            act_dim = action_preds.shape[2]
            action_preds = action_preds.reshape(-1, act_dim)
            target = action_target.reshape(-1, act_dim)
            return jnp.mean(
                ((action_preds - target) ** 2)
                * jnp.expand_dims(attention_mask.reshape(-1) > 0, -1).astype(
                    jnp.float32
                )
            )

        def sgd_step(state: TrainState, inputs):
            key_step, key = jax.random.split(state.key)
            loss_value, grads = jax.value_and_grad(loss_fn)(
                state.params, key_step, inputs
            )
            updates, opt_state = optimizer.update(grads, state.opt_state, state.params)
            params = optax.apply_updates(state.params, updates)
            return TrainState(params, opt_state, key), {"loss": loss_value}

        self._sgd_step = jax.jit(sgd_step)
        self._dataset = dataset

        init_key, key = jax.random.split(key)
        init_params = model.init(init_key)
        init_opt_state = optimizer.init(init_params)

        self._state = TrainState(init_params, init_opt_state, key)
        self._logger = logger or loggers.make_default_logger(
            "learner",
            asynchronous=True,
        )
        self._counter = counter or counting.Counter()
        self._timestamp = None

    def step(self):
        inputs = next(self._dataset)

        self._state, metrics = self._sgd_step(self._state, inputs)

        # Compute elapsed time.
        timestamp = time.time()
        elapsed_time = timestamp - self._timestamp if self._timestamp else 0
        self._timestamp = timestamp

        # Increment counts
        counts = self._counter.increment(steps=1, walltime=elapsed_time)

        self._logger.write({**metrics, **counts})

    def get_variables(self, names: Sequence[str]):
        variables = {"model": self._state.params}
        return [variables[name] for name in names]

corax/agents/jax/decision_transformer/learning.py
ethanluoyc/corax
# type: ignore
"""Decision Transformer model definition."""
from typing import Optional

import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np

from corax import specs
from corax.jax import networks as networks_lib
from corax.jax import utils


class CausalSelfAttention(hk.MultiHeadAttention):
    """Self attention with a causal mask applied."""

    def __call__(
        self,
        query: jnp.ndarray,
        key: Optional[jnp.ndarray] = None,
        value: Optional[jnp.ndarray] = None,
        mask: Optional[jnp.ndarray] = None,
    ) -> jnp.ndarray:
        key = key if key is not None else query
        value = value if value is not None else query

        if query.ndim != 3:
            raise ValueError("Expect queries of shape [B, T, D].")

        seq_len = query.shape[1]
        causal_mask = np.tril(np.ones((1, 1, seq_len, seq_len)))
        mask = mask * causal_mask if mask is not None else causal_mask

        return super().__call__(query, key, value, mask)


class DenseBlock(hk.Module):
    """A 2-layer MLP which widens then narrows the input."""

    def __init__(
        self, init_scale: float, widening_factor: int = 4, name: Optional[str] = None
    ):
        super().__init__(name=name)
        self._init_scale = init_scale
        self._widening_factor = widening_factor

    def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
        hiddens = x.shape[-1]
        initializer = hk.initializers.VarianceScaling(self._init_scale)
        x = hk.Linear(self._widening_factor * hiddens, w_init=initializer)(x)
        x = jax.nn.gelu(x)
        return hk.Linear(hiddens, w_init=initializer)(x)


class Transformer(hk.Module):
    """A transformer stack."""

    def __init__(
        self,
        num_heads: int,
        num_layers: int,
        dropout_rate: float,
        name: Optional[str] = None,
    ):
        super().__init__(name=name)
        self._num_layers = num_layers
        self._num_heads = num_heads
        self._dropout_rate = dropout_rate

    def __call__(
        self, h: jnp.ndarray, mask: Optional[jnp.ndarray], is_training: bool
    ) -> jnp.ndarray:
        """Connects the transformer.

        Args:
          h: Inputs, [B, T, D].
          mask: Padding mask, [B, T].
          is_training: Whether we're training or not.

        Returns:
          Array of shape [B, T, D].
        """

        init_scale = 2.0 / self._num_layers
        dropout_rate = self._dropout_rate if is_training else 0.0
        if mask is not None:
            mask = mask[:, None, None, :]

        # Note: names chosen to approximately match those used in the GPT-2 code;
        # see https://github.com/openai/gpt-2/blob/master/src/model.py.
        for i in range(self._num_layers):
            h_norm = layer_norm(h, name=f"h{i}_ln_1")
            h_attn = CausalSelfAttention(
                num_heads=self._num_heads,
                key_size=32,
                model_size=h.shape[-1],
                w_init=hk.initializers.VarianceScaling(init_scale),
                name=f"h{i}_attn",
            )(h_norm, mask=mask)
            h_attn = hk.dropout(hk.next_rng_key(), dropout_rate, h_attn)
            h = h + h_attn
            h_norm = layer_norm(h, name=f"h{i}_ln_2")
            h_dense = DenseBlock(init_scale, name=f"h{i}_mlp")(h_norm)
            h_dense = hk.dropout(hk.next_rng_key(), dropout_rate, h_dense)
            h = h + h_dense
        h = layer_norm(h, name="ln_f")

        return h


class DecisionTransformer(hk.Module):
    def __init__(
        self,
        state_size: int,
        action_size: int,
        hidden_size: int,
        num_heads: int,
        num_layers: int,
        dropout_rate: int,
        action_tanh=True,
        max_ep_len=4096,
        name: str = "decision_transformer",
    ):
        super().__init__(name=name)
        self._state_size = state_size
        self._action_size = action_size
        self._hidden_size = hidden_size

        self.transformer = Transformer(
            num_heads=num_heads,
            num_layers=num_layers,
            dropout_rate=dropout_rate,
        )
        embed_init = hk.initializers.TruncatedNormal(stddev=0.02)
        w_init = embed_init

        self.embed_timestep = hk.Embed(max_ep_len, hidden_size, w_init=embed_init)
        self.embed_return = hk.Linear(hidden_size, w_init=w_init)
        self.embed_state = hk.Linear(hidden_size, w_init=w_init)
        self.embed_action = hk.Linear(hidden_size, w_init=w_init)
        self.embed_ln = hk.LayerNorm(axis=-1, create_scale=True, create_offset=True)

        # note: we don't predict states or returns for the paper
        self.predict_state = hk.Linear(self._state_size, w_init=w_init)
        self.predict_action = hk.Sequential(
            [hk.Linear(self._action_size, w_init=w_init)]
            + ([jnp.tanh] if action_tanh else [])
        )
        self.predict_return = hk.Linear(1, w_init=w_init)

    def __call__(
        self,
        states,
        actions,
        rewards,
        returns_to_go,
        timesteps,
        attention_mask=None,
        is_training=True,
    ):
        del rewards
        batch_size, seq_length = states.shape[0], states.shape[1]

        if attention_mask is None:
            # attention mask for GPT: 1 if can be attended to, 0 if not
            attention_mask = jnp.ones((batch_size, seq_length), dtype=jnp.int32)

        # embed each modality with a different head
        state_embeddings = self.embed_state(states)
        action_embeddings = self.embed_action(actions)
        returns_embeddings = self.embed_return(returns_to_go)
        time_embeddings = self.embed_timestep(timesteps)

        # time embeddings are treated similar to positional embeddings
        state_embeddings = state_embeddings + time_embeddings
        action_embeddings = action_embeddings + time_embeddings
        returns_embeddings = returns_embeddings + time_embeddings

        # this makes the sequence look like (R_1, s_1, a_1, R_2, s_2, a_2, ...)
        # which works nice in an autoregressive sense since states predict actions
        stacked_inputs = jnp.transpose(
            jnp.stack(
                (returns_embeddings, state_embeddings, action_embeddings), axis=1
            ),
            (0, 2, 1, 3),
        ).reshape((batch_size, 3 * seq_length, self._hidden_size))
        stacked_inputs = self.embed_ln(stacked_inputs)

        # to make the attention mask fit the stacked inputs, have to stack it as well
        stacked_attention_mask = jnp.transpose(
            jnp.stack((attention_mask, attention_mask, attention_mask), axis=1),
            (0, 2, 1),
        ).reshape((batch_size, 3 * seq_length))

        # we feed in the input embeddings (not word indices as in NLP) to the model
        transformer_outputs = self.transformer(
            stacked_inputs,
            mask=stacked_attention_mask,
            is_training=is_training,
        )
        x = transformer_outputs

        # reshape x so that the second dimension corresponds to the original
        # returns (0), states (1), or actions (2); i.e. x[:,1,t] is the token for s_t
        x = jnp.transpose(
            x.reshape(batch_size, seq_length, 3, self._hidden_size), (0, 2, 1, 3)
        )

        # get predictions
        return_preds = self.predict_return(
            x[:, 2]
        )  # predict next return given state and action
        state_preds = self.predict_state(
            x[:, 2]
        )  # predict next state given state and action
        action_preds = self.predict_action(x[:, 1])  # predict next action given state

        return state_preds, action_preds, return_preds


def layer_norm(x: jnp.ndarray, name: Optional[str] = None) -> jnp.ndarray:
    """Apply a unique LayerNorm to x with default settings."""
    return hk.LayerNorm(axis=-1, create_scale=True, create_offset=True, name=name)(x)


def make_gym_networks(
    spec: specs.EnvironmentSpec,
    hidden_size: int,
    num_layers: int,
    num_heads: int,
    dropout_rate: float,
    episode_length: int,
):
    state_dim = spec.observations.shape[-1]
    act_dim = spec.actions.shape[-1]

    def _model_fn(
        states,
        actions,
        rewards,
        returns_to_go,
        timesteps,
        attention_mask,
        is_training=True,
    ):
        model = DecisionTransformer(
            state_size=state_dim,
            action_size=act_dim,
            max_ep_len=episode_length,
            hidden_size=hidden_size,
            num_layers=num_layers,
            num_heads=num_heads,
            dropout_rate=dropout_rate,
        )
        return model(
            states,
            actions,
            rewards,
            returns_to_go,
            timesteps,
            attention_mask=attention_mask,
            is_training=is_training,
        )

    dummy_obs = utils.zeros_like(spec.observations)
    dummy_act = utils.zeros_like(spec.actions)

    dummy_obs = utils.add_batch_dim(jnp.stack([dummy_obs]))
    dummy_act = utils.add_batch_dim(jnp.stack([dummy_act]))
    dummy_reward = jnp.zeros((1, 1, 1))
    dummy_rtg = jnp.zeros((1, 1, 1))
    dummy_timestep = jnp.zeros(
        (
            1,
            1,
        ),
        dtype=jnp.int32,
    )
    dummy_mask = jnp.zeros((1, 1), dtype=jnp.bool_)

    model = hk.transform(_model_fn)
    init_fn = lambda key: model.init(
        key,
        dummy_obs,
        dummy_act,
        dummy_reward,
        dummy_rtg,
        dummy_timestep,
        dummy_mask,
        is_training=True,
    )
    return networks_lib.FeedForwardNetwork(init=init_fn, apply=model.apply)

corax/agents/jax/decision_transformer/networks.py
ethanluoyc/corax
# type: ignore
import chex
import haiku as hk
import jax
import jax.numpy as jnp
from absl.testing import absltest

from corax.agents.jax.decision_transformer import networks


class NetworksTest(absltest.TestCase):
    def test_dt_run(self):
        batch_size = 10
        sequence_size = 5
        num_heads = 2
        num_layers = 2
        dropout_rate = 0.0
        state_size = 3
        action_size = 5
        hidden_size = 8

        def forward_fn(data, is_training: bool = True):
            dt = networks.DecisionTransformer(
                state_size=state_size,
                action_size=action_size,
                hidden_size=hidden_size,
                num_heads=num_heads,
                num_layers=num_layers,
                dropout_rate=dropout_rate,
                action_tanh=True,
            )
            output_embeddings = dt(
                data["states"],
                data["actions"],
                data["rewards"],
                data["returns_to_go"],
                data["timesteps"],
                attention_mask=data["mask"],
                is_training=is_training,
            )
            return output_embeddings

        forward = hk.transform(forward_fn)
        examples_input = {
            "states": jnp.zeros((batch_size, sequence_size, state_size)),
            "actions": jnp.zeros((batch_size, sequence_size, action_size)),
            "rewards": jnp.zeros((batch_size, sequence_size, 1)),
            "returns_to_go": jnp.zeros((batch_size, sequence_size, 1)),
            "timesteps": jnp.zeros((batch_size, sequence_size), dtype=jnp.int32),
            "mask": jnp.ones((batch_size, sequence_size), dtype=jnp.bool_),
        }
        key = jax.random.PRNGKey(0)
        key_init, key_apply = jax.random.split(key)
        del key
        params = forward.init(key_init, examples_input, True)
        state_preds, action_preds, return_preds = forward.apply(
            params, key_apply, examples_input, True
        )
        chex.assert_equal_shape((state_preds, examples_input["states"]))
        chex.assert_equal_shape((action_preds, examples_input["actions"]))
        chex.assert_equal_shape((return_preds, examples_input["rewards"]))

    def test_transformer_forward(self):
        batch_size = 10
        sequence_size = 5
        input_size = 4
        num_heads = 2
        num_layers = 2
        dropout_rate = 0.0

        def forward_fn(data, is_training=True):
            tokens = data["inputs"]
            input_embeddings = tokens
            # input_mask = jnp.greater(tokens, 0)
            input_mask = None

            # Run the transformer over the inputs.
            transformer = networks.Transformer(
                num_heads=num_heads, num_layers=num_layers, dropout_rate=dropout_rate
            )
            output_embeddings = transformer(input_embeddings, input_mask, is_training)
            return output_embeddings

        forward = hk.transform(forward_fn)
        examples_input = {"inputs": jnp.zeros((batch_size, sequence_size, input_size))}
        key = jax.random.PRNGKey(0)
        key_init, key_apply = jax.random.split(key)
        del key
        params = forward.init(key_init, examples_input, True)
        output = forward.apply(params, key_apply, examples_input, True)
        self.assertEqual(output.shape, (batch_size, sequence_size, input_size))


if __name__ == "__main__":
    absltest.main()

corax/agents/jax/decision_transformer/networks_test.py
ethanluoyc/corax
"""DrQ-v2 agent implementation."""

from corax.agents.jax.drq_v2.acting import DrQV2Actor
from corax.agents.jax.drq_v2.learning import DrQV2Learner
from corax.agents.jax.drq_v2.networks import DrQV2Networks
from corax.agents.jax.drq_v2.networks import apply_policy_and_sample
from corax.agents.jax.drq_v2.networks import make_networks

corax/agents/jax/drq_v2/__init__.py
ethanluoyc/corax
"""Acting component for DrQ-V2."""
from typing import Optional

import jax

from corax import adders
from corax import core
from corax import types
from corax.agents.jax.drq_v2 import networks as drq_v2_networks
from corax.jax import networks as networks_lib
from corax.jax import utils
from corax.jax import variable_utils


class DrQV2Actor(core.Actor):
    """Actor component for DrQ-V2 agent."""

    def __init__(
        self,
        policy: drq_v2_networks.DrQV2PolicyNetwork,
        random_key: networks_lib.PRNGKey,
        variable_client: Optional[variable_utils.VariableClient],
        adder: Optional[adders.Adder] = None,
        jit: bool = True,
        backend: Optional[str] = "cpu",
        per_episode_update: bool = False,
    ):
        """Initializes a DrQV2Actor.
        Args:
            policy: DrQV2 policy network.
            random_key: Random key.
            variable_client: The variable client to get policy parameters from.
            adder: An adder to add experiences to.
            jit: Whether or not to jit the passed ActorCore's pure functions.
            backend: Which backend to use when jitting the policy.
            per_episode_update: if True, updates variable client params once at the
                beginning of each episode
        """
        self._random_key = random_key
        self._variable_client = variable_client
        self._adder = adder
        self._state = None
        self._num_steps = 0
        self._per_episode_update = per_episode_update

        def init(rng):
            return rng

        def select_action(params, observation, state, step):
            rng = state
            rng1, rng2 = jax.random.split(rng)
            observation = utils.add_batch_dim(observation)
            action = utils.squeeze_batch_dim(policy(params, rng1, observation, step))
            return action, rng2

        def get_extras(unused):
            del unused

        # Unpack ActorCore, jitting if requested.
        if jit:
            self._init = jax.jit(init, backend=backend)
            self._policy = jax.jit(select_action, backend=backend)
        else:
            self._init = init
            self._policy = select_action
            self._get_extras = get_extras
            self._per_episode_update = per_episode_update
        self._get_extras = get_extras

    @property
    def _params(self):
        return self._variable_client.params if self._variable_client else []

    def select_action(self, observation) -> types.NestedArray:
        action, self._state = self._policy(
            self._params, observation, self._state, self._num_steps
        )
        self._num_steps += 1
        return utils.to_numpy(action)

    def observe_first(self, timestep):
        self._random_key, key = jax.random.split(self._random_key)
        self._state = self._init(key)
        if self._adder:
            self._adder.add_first(timestep)
        if self._variable_client and self._per_episode_update:
            self._variable_client.update_and_wait()

    def observe(self, action, next_timestep):
        if self._adder:
            self._adder.add(action, next_timestep, extras=self._get_extras(self._state))

    def update(self, wait: bool = False):
        if self._variable_client and not self._per_episode_update:
            self._variable_client.update(wait)

corax/agents/jax/drq_v2/acting.py
ethanluoyc/corax
"""Data augmentations used in DrQ."""
from typing import Callable

import jax
import jax.numpy as jnp

from corax import types
from corax.jax import types as jax_types

DataAugmentation = Callable[[jax_types.PRNGKey, types.NestedArray], types.NestedArray]


# From https://github.com/ikostrikov/jax-rl/blob/main/jax_rl/agents/drq/augmentations.py
def random_crop(key: jax_types.PRNGKey, img, padding):
    crop_from = jax.random.randint(key, (2,), 0, 2 * padding + 1)
    crop_from = jnp.concatenate([crop_from, jnp.zeros((1,), dtype=jnp.int32)])
    padded_img = jnp.pad(
        img, ((padding, padding), (padding, padding), (0, 0)), mode="edge"
    )
    return jax.lax.dynamic_slice(padded_img, crop_from, img.shape)


def batched_random_crop(key, imgs, padding=4):
    keys = jax.random.split(key, imgs.shape[0])
    return jax.vmap(random_crop, (0, 0, None))(keys, imgs, padding)


def random_shift_aug(key, img, padding):
    """Perform the random shift augmentation used in DrQ-V2"""
    # Pad the image
    padded_img = jnp.pad(
        img, ((padding, padding), (padding, padding), (0, 0)), mode="edge"
    )
    # Shift and Resample
    shifts = jax.random.randint(key, (2,), 0, 2 * padding + 1)
    return jax.image.scale_and_translate(
        padded_img,
        img.shape,
        spatial_dims=(0, 1),
        scale=jnp.ones(2, dtype=jnp.float32),
        translation=-shifts,
        method="bilinear",
    ).astype(img.dtype)


def batched_random_shift_aug(key, img, padding=4):
    keys = jax.random.split(key, img.shape[0])
    return jax.vmap(random_shift_aug, (0, 0, None))(keys, img, padding)

corax/agents/jax/drq_v2/augmentations.py
ethanluoyc/corax
"""Learner component for DrQV2."""
import time
from typing import Iterator, NamedTuple, Optional

import jax
import jax.numpy as jnp
import optax
import reverb
import rlax

from corax import core
from corax import types as acme_types
from corax.agents.jax.drq_v2 import augmentations
from corax.agents.jax.drq_v2 import networks as drq_v2_networks
from corax.jax import networks as networks_lib
from corax.jax import types as jax_types
from corax.jax import utils
from corax.utils import counting
from corax.utils import loggers


class TrainingState(NamedTuple):
    """Holds training state for the DrQ learner."""

    policy_params: networks_lib.Params
    policy_opt_state: optax.OptState

    encoder_params: networks_lib.Params
    # There is not target encoder parameters in v2.
    encoder_opt_state: optax.OptState

    critic_params: networks_lib.Params
    critic_target_params: networks_lib.Params
    critic_opt_state: optax.OptState

    key: jax_types.PRNGKey
    steps: int


class DrQV2Learner(core.Learner):
    """Learner for DrQ-v2"""

    def __init__(
        self,
        random_key: jax_types.PRNGKey,
        dataset: Iterator[reverb.ReplaySample],
        networks: drq_v2_networks.DrQV2Networks,
        sigma_schedule: optax.Schedule,
        augmentation: augmentations.DataAugmentation,
        policy_optimizer: optax.GradientTransformation,
        critic_optimizer: optax.GradientTransformation,
        encoder_optimizer: optax.GradientTransformation,
        noise_clip: float = 0.3,
        critic_soft_update_rate: float = 0.005,
        discount: float = 0.99,
        num_sgd_steps_per_step: int = 1,
        clipping: bool = False,
        bc_alpha: Optional[float] = None,
        counter: Optional[counting.Counter] = None,
        logger: Optional[loggers.Logger] = None,
    ):
        def critic_loss_fn(
            critic_params: networks_lib.Params,
            encoder_params: networks_lib.Params,
            critic_target_params: networks_lib.Params,
            policy_params: networks_lib.Params,
            transitions: acme_types.Transition,
            key: jax_types.PRNGKey,
            sigma: jnp.ndarray,
        ):
            next_encoded = networks.encoder_network.apply(
                encoder_params, transitions.next_observation
            )
            next_action = networks.policy_network.apply(policy_params, next_encoded)
            next_action = networks.add_policy_noise(next_action, key, sigma, noise_clip)  # type: ignore
            next_q1, next_q2 = networks.critic_network.apply(
                critic_target_params, next_encoded, next_action
            )
            # Calculate q target values
            next_q = jnp.minimum(next_q1, next_q2)
            target_q = transitions.reward + transitions.discount * discount * next_q
            target_q = jax.lax.stop_gradient(target_q)
            # Calculate predicted Q
            encoded = networks.encoder_network.apply(
                encoder_params, transitions.observation
            )
            q1, q2 = networks.critic_network.apply(
                critic_params, encoded, transitions.action
            )
            loss_critic = (jnp.square(target_q - q1) + jnp.square(target_q - q2)).mean(
                axis=0
            )
            return loss_critic, {"q1": q1.mean(), "q2": q2.mean()}

        def critic_fn(params, h, a):
            q1, q2 = networks.critic_network.apply(params, h, a)
            q = jnp.minimum(q1, q2)
            return q

        def policy_loss_fn(
            policy_params: networks_lib.Params,
            critic_params: networks_lib.Params,
            encoder_params: networks_lib.Params,
            transition: acme_types.Transition,
            sigma: jnp.ndarray,
            key,
        ):
            encoded = networks.encoder_network.apply(
                encoder_params, transition.observation
            )
            action = networks.policy_network.apply(policy_params, encoded)
            action = networks.add_policy_noise(action, key, sigma, noise_clip)  # type: ignore
            if bc_alpha is None:
                grad_critic = jax.vmap(
                    jax.grad(critic_fn, argnums=2), in_axes=(None, 0, 0)
                )
                dq_da = grad_critic(critic_params, encoded, action)
                dqda_clipping = 1.0 if clipping else None
                batch_dpg_learning = jax.vmap(rlax.dpg_loss, in_axes=(0, 0, None))
                loss = batch_dpg_learning(action, dq_da, dqda_clipping)
                return jnp.mean(loss), {}
            else:
                q = critic_fn(critic_params, encoded, action)
                lmbda = jax.lax.stop_gradient(bc_alpha / jnp.mean(jnp.abs(q)))
                loss = -lmbda * jnp.mean(q) + jnp.mean(
                    jnp.square(action - transition.action)
                )
                return loss, {}

        policy_grad_fn = jax.value_and_grad(policy_loss_fn, has_aux=True)
        critic_grad_fn = jax.value_and_grad(
            critic_loss_fn, argnums=(0, 1), has_aux=True
        )

        def update_step(
            state: TrainingState,
            transitions: acme_types.Transition,
        ):
            key_aug1, key_aug2, key_policy, key_critic, key = jax.random.split(
                state.key, 5
            )
            sigma = sigma_schedule(state.steps)
            # Perform data augmentation on o_tm1 and o_t
            observation_aug = augmentation(key_aug1, transitions.observation)
            next_observation_aug = augmentation(key_aug2, transitions.next_observation)
            transitions = transitions._replace(
                observation=observation_aug,
                next_observation=next_observation_aug,
            )
            # Update critic
            (critic_loss, critic_aux), (critic_grad, encoder_grad) = critic_grad_fn(
                state.critic_params,
                state.encoder_params,
                state.critic_target_params,
                state.policy_params,
                transitions,
                key_critic,
                sigma,
            )
            (policy_loss, policy_aux), actor_grad = policy_grad_fn(
                state.policy_params,
                state.critic_params,
                state.encoder_params,
                transitions,
                sigma,
                key_policy,
            )
            encoder_update, encoder_opt_state = encoder_optimizer.update(
                encoder_grad, state.encoder_opt_state
            )
            critic_update, critic_opt_state = critic_optimizer.update(
                critic_grad, state.critic_opt_state
            )
            encoder_params = optax.apply_updates(state.encoder_params, encoder_update)
            critic_params = optax.apply_updates(state.critic_params, critic_update)
            # Update policy
            policy_update, policy_opt_state = policy_optimizer.update(
                actor_grad, state.policy_opt_state
            )
            policy_params = optax.apply_updates(state.policy_params, policy_update)

            critic_target_params = optax.incremental_update(
                critic_params, state.critic_target_params, critic_soft_update_rate
            )
            metrics = {
                "policy_loss": policy_loss,
                "critic_loss": critic_loss,
                "sigma": sigma,
                **critic_aux,
                **policy_aux,
            }
            new_state = TrainingState(
                policy_params=policy_params,
                policy_opt_state=policy_opt_state,
                encoder_params=encoder_params,
                encoder_opt_state=encoder_opt_state,
                critic_params=critic_params,
                critic_target_params=critic_target_params,
                critic_opt_state=critic_opt_state,
                key=key,
                steps=state.steps + 1,
            )
            return new_state, metrics

        self._iterator = dataset
        self._counter = counter or counting.Counter()
        self._logger = logger or loggers.make_default_logger(
            label="learner",
            asynchronous=True,
            serialize_fn=utils.fetch_devicearray,
        )
        self._update_step = utils.process_multiple_batches(
            update_step, num_sgd_steps_per_step
        )
        self._update_step = jax.jit(self._update_step)

        # Initialize training state
        def make_initial_state(key: jax_types.PRNGKey):
            key_encoder, key_critic, key_policy, key = jax.random.split(key, 4)
            encoder_init_params = networks.encoder_network.init(key_encoder)
            encoder_init_opt_state = encoder_optimizer.init(encoder_init_params)

            critic_init_params = networks.critic_network.init(key_critic)
            critic_init_opt_state = critic_optimizer.init(critic_init_params)

            policy_init_params = networks.policy_network.init(key_policy)
            policy_init_opt_state = policy_optimizer.init(policy_init_params)

            return TrainingState(
                policy_params=policy_init_params,
                policy_opt_state=policy_init_opt_state,
                encoder_params=encoder_init_params,
                critic_params=critic_init_params,
                critic_target_params=critic_init_params,
                encoder_opt_state=encoder_init_opt_state,
                critic_opt_state=critic_init_opt_state,
                key=key,
                steps=0,
            )

        # Create initial state.
        self._state = make_initial_state(random_key)

        # 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

    def step(self):
        # Get the next batch from the replay iterator
        sample = next(self._iterator)
        transitions = acme_types.Transition(*sample.data)

        # Perform a single learner step
        self._state, metrics = self._update_step(self._state, transitions)

        # Compute elapsed time
        timestamp = time.time()
        elapsed_time = timestamp - self._timestamp if self._timestamp else 0
        self._timestamp = timestamp

        # Increment counts and record the current time
        counts = self._counter.increment(steps=1, walltime=elapsed_time)

        # Attempts to write the logs.
        self._logger.write({**metrics, **counts})

    def get_variables(self, names):
        variables = {
            "policy": {
                "encoder": self._state.encoder_params,
                "policy": self._state.policy_params,
            },
        }
        return [variables[name] for name in names]

    def save(self) -> TrainingState:
        return self._state

    def restore(self, state: TrainingState) -> None:
        self._state = state

corax/agents/jax/drq_v2/learning.py
ethanluoyc/corax
# type: ignore
"""Network definitions for DrQ-v2."""
import dataclasses
from typing import Callable

import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np

from corax import specs
from corax import types
from corax.agents.jax import actor_core
from corax.jax import networks as networks_lib
from corax.jax import utils

# Unlike standard FF-policy, in our DrQ-V2 implementation we use
# scheduled stddev parameters, the pure function for the policy
# thus needs to know the current time step of the actor to calculate
# the current stddev.
Step = int
DrQV2PolicyNetwork = Callable[
    [networks_lib.Params, networks_lib.PRNGKey, types.NestedArray, Step],
    types.NestedArray,
]


class DrQTorso(hk.Module):
    """DrQ Torso used in DrQ-v2."""

    def __init__(self, name: str = "drq_torso"):
        super().__init__(name)
        # pylint: disable=use-dict-literal
        conv_kwargs = {
            "kernel_shape": (3, 3),
            "output_channels": 32,
            "padding": "VALID",
            # This follows from the reference implementation, the scale accounts for
            # using the ReLU activation.
            "w_init": hk.initializers.Orthogonal(2**0.5),
        }
        self._network = hk.Sequential(
            [
                hk.Conv2D(stride=2, **conv_kwargs),
                jax.nn.relu,
                hk.Conv2D(stride=1, **conv_kwargs),
                jax.nn.relu,
                hk.Conv2D(stride=1, **conv_kwargs),
                jax.nn.relu,
                hk.Conv2D(stride=1, **conv_kwargs),
                jax.nn.relu,
                hk.Flatten(),
            ]
        )

    def __call__(self, inputs: jnp.ndarray):
        if not jnp.issubdtype(inputs.dtype, jnp.uint8):
            msg = "Expect inputs to be uint8 pixel values between 0 to 255."
            raise ValueError(msg)
        if inputs.ndim != 4:
            msg = (
                "Input array should have 4 dimensions "
                "(for batch size, height, width, and channels),"
                f" but it has {inputs.ndim}"
            )
            raise ValueError(msg)
        # Floatify the image.
        preprocessed_inputs = inputs.astype(jnp.float32) / 255.0 - 0.5
        return self._network(preprocessed_inputs)


@dataclasses.dataclass
class DrQV2Networks:
    encoder_network: networks_lib.FeedForwardNetwork
    policy_network: networks_lib.FeedForwardNetwork
    critic_network: networks_lib.FeedForwardNetwork
    add_policy_noise: Callable[
        [types.NestedArray, networks_lib.PRNGKey, float, float], types.NestedArray
    ]
    get_policy_feature: Callable


def apply_policy_and_sample(
    networks: DrQV2Networks, action_specs: specs.BoundedArray, sigma: float
) -> actor_core.FeedForwardPolicy:
    """Create a pure policy function for actor_core."""

    def policy(params, key, obs):
        feature_map = networks.encoder_network.apply(params["encoder"], obs)
        action = networks.policy_network.apply(params["policy"], feature_map)
        noise = jax.random.normal(key, shape=action.shape) * sigma
        return jnp.clip(action + noise, action_specs.minimum, action_specs.maximum)

    return policy


def make_networks(
    spec: specs.EnvironmentSpec,
    hidden_size: int = 1024,
    latent_size: int = 50,
) -> DrQV2Networks:
    """Create networks for the DrQ-v2 agent."""
    action_size = np.prod(spec.actions.shape, dtype=int)

    def _encoder_fn(obs):
        return DrQTorso()(obs)

    def _trunk():
        w_init = hk.initializers.Orthogonal(1.0)
        return hk.Sequential(
            [
                hk.Linear(latent_size, w_init=w_init),
                hk.LayerNorm(axis=-1, create_scale=True, create_offset=True),
                jnp.tanh,
            ]
        )

    def _critic_fn(obs, action):
        w_init = hk.initializers.Orthogonal(1.0)
        embed = _trunk()(obs)
        # NOTE: Actions are concatenated after the trunk
        embed = jnp.concatenate([embed, action], axis=-1)
        critic_network1 = hk.nets.MLP(
            output_sizes=(hidden_size, hidden_size, 1),
            w_init=w_init,
            activate_final=False,
        )
        critic_network2 = hk.nets.MLP(
            output_sizes=(hidden_size, hidden_size, 1),
            w_init=w_init,
            activate_final=False,
        )
        q1 = critic_network1(embed)
        q2 = critic_network2(embed)
        q1 = jnp.squeeze(q1, axis=-1)
        q2 = jnp.squeeze(q2, axis=-1)
        return q1, q2

    def _policy_fn(obs):
        embed = _trunk()(obs)
        head = hk.Sequential(
            [
                hk.nets.MLP(
                    output_sizes=(hidden_size, hidden_size),
                    w_init=hk.initializers.Orthogonal(1.0),
                    activate_final=True,
                ),
                hk.Linear(action_size, w_init=hk.initializers.Orthogonal(1.0)),
                jnp.tanh,
            ]
        )
        return head(embed)

    def _policy_features_fn(obs):
        return _trunk()(obs)

    def add_policy_noise(
        action: types.NestedArray,
        key: networks_lib.PRNGKey,
        sigma: float,
        noise_clip: float,
    ) -> types.NestedArray:
        """Adds action noise to bootstrapped Q-value estimate in critic loss."""
        noise = jax.random.normal(key=key, shape=action.shape) * sigma
        noise = jnp.clip(noise, -noise_clip, noise_clip)
        action = action + noise
        clipped_action = jnp.clip(action, spec.actions.minimum, spec.actions.maximum)
        return (
            action
            - jax.lax.stop_gradient(action)
            + jax.lax.stop_gradient(clipped_action)
        )

    def transform_without_rng(f):
        return hk.without_apply_rng(hk.transform(f))

    policy = transform_without_rng(_policy_fn)
    critic = transform_without_rng(_critic_fn)
    encoder = transform_without_rng(_encoder_fn)
    policy_feature = transform_without_rng(_policy_features_fn)

    dummy_action = utils.add_batch_dim(utils.zeros_like(spec.actions))
    dummy_obs = utils.add_batch_dim(utils.zeros_like(spec.observations))
    dummy_encoded = jax.eval_shape(encoder.init, jax.random.PRNGKey(0), dummy_obs)
    dummy_encoded = jax.eval_shape(encoder.apply, dummy_encoded, dummy_obs)
    dummy_encoded = jnp.zeros(shape=dummy_encoded.shape, dtype=dummy_encoded.dtype)

    return DrQV2Networks(
        encoder_network=networks_lib.FeedForwardNetwork(
            lambda key: encoder.init(key, dummy_obs), encoder.apply
        ),
        policy_network=networks_lib.FeedForwardNetwork(
            lambda key: policy.init(key, dummy_encoded), policy.apply
        ),
        critic_network=networks_lib.FeedForwardNetwork(
            lambda key: critic.init(key, dummy_encoded, dummy_action), critic.apply
        ),
        add_policy_noise=add_policy_noise,
        get_policy_feature=policy_feature.apply,
    )

corax/agents/jax/drq_v2/networks.py
ethanluoyc/corax
"""Implicit Q Learning agent."""
from corax.agents.jax.iql.builder import IQLBuilder
from corax.agents.jax.iql.config import IQLConfig
from corax.agents.jax.iql.learning import IQLLearnerCore
from corax.agents.jax.iql.networks import IQLNetworks
from corax.agents.jax.iql.networks import apply_policy_and_sample
from corax.agents.jax.iql.networks import make_networks

corax/agents/jax/iql/__init__.py
ethanluoyc/corax
"""IQL Builder."""
from typing import Iterator, Optional

import optax
import reverb

from corax import adders
from corax import core
from corax import specs
from corax import types
from corax.agents.jax import actor_core as actor_core_lib
from corax.agents.jax import actors
from corax.agents.jax import builders
from corax.agents.jax.iql import config as iql_config
from corax.agents.jax.iql import learning as iql_learning
from corax.agents.jax.iql import networks as iql_networks
from corax.jax import networks as networks_lib
from corax.jax import variable_utils
from corax.utils import counting
from corax.utils import loggers


class IQLBuilder(
    builders.OfflineBuilder[
        iql_networks.IQLNetworks, actor_core_lib.FeedForwardPolicy, types.Transition
    ]
):
    """IQL Builder."""

    def __init__(self, config: iql_config.IQLConfig):
        """Creates a IQL learner, a behavior policy and an eval actor.

        Args:
          config: a config with IQL hps
        """
        self._config = config

    def make_learner(
        self,
        random_key: networks_lib.PRNGKey,
        networks: iql_networks.IQLNetworks,
        dataset: Iterator[types.Transition],
        logger_fn: loggers.LoggerFactory,
        environment_spec: specs.EnvironmentSpec,
        replay_client: Optional[reverb.Client] = None,
        counter: Optional[counting.Counter] = None,
    ) -> core.Learner:
        del environment_spec, replay_client

        policy_optimizer = optax.adam(self._config.learning_rate)
        critic_optimizer = optax.adam(self._config.learning_rate)
        value_optimizer = optax.adam(self._config.learning_rate)

        return iql_learning.IQLLearner(
            networks=networks,
            random_key=random_key,
            dataset=dataset,
            policy_optimizer=policy_optimizer,
            critic_optimizer=critic_optimizer,
            value_optimizer=value_optimizer,
            discount=self._config.discount,
            tau=self._config.tau,
            expectile=self._config.expectile,
            temperature=self._config.temperature,
            logger=logger_fn("learner"),
            counter=counter,
        )

    def make_actor(
        self,
        random_key: networks_lib.PRNGKey,
        policy: actor_core_lib.FeedForwardPolicy,
        environment_spec: specs.EnvironmentSpec,
        variable_source: Optional[core.VariableSource] = None,
        adder: Optional[adders.Adder] = None,
    ) -> core.Actor:
        del environment_spec
        assert variable_source is not None
        actor_core = actor_core_lib.batched_feed_forward_to_actor_core(policy)
        # Inference happens on CPU, so it's better to move variables there too.
        variable_client = variable_utils.VariableClient(
            variable_source, key="policy", device="cpu"
        )
        return actors.GenericActor(
            actor_core, random_key, variable_client, adder, backend="cpu"
        )

    def make_policy(
        self,
        networks: iql_networks.IQLNetworks,
        environment_spec: specs.EnvironmentSpec,
        evaluation: bool = False,
    ) -> actor_core_lib.FeedForwardPolicy:
        """Creates a policy."""
        return iql_networks.apply_policy_and_sample(
            networks,
            environment_spec.actions,
            eval_mode=evaluation,
        )

corax/agents/jax/iql/builder.py
ethanluoyc/corax
import dataclasses


@dataclasses.dataclass
class IQLConfig:
    learning_rate: float = 3e-4
    discount: float = 0.99
    tau: float = 0.005
    expectile: float = 0.8
    temperature: float = 0.1

corax/agents/jax/iql/config.py
ethanluoyc/corax
"""Implementations of Implicit Q Learning (IQL) learner component."""
from typing import Any, Dict, Iterator, NamedTuple, Optional, Tuple

import jax
import jax.numpy as jnp
import optax

from corax import types as core_types
from corax.agents.jax import learner_core
from corax.agents.jax import learners
from corax.agents.jax.iql import networks as iql_networks
from corax.jax import networks as networks_lib
from corax.jax import types as jax_types
from corax.utils import counting
from corax.utils import loggers

_Metrics = Dict[str, jnp.ndarray]


class TrainingState(NamedTuple):
    policy_params: networks_lib.Params
    policy_opt_state: optax.OptState
    value_params: networks_lib.Params
    value_opt_state: optax.OptState
    critic_params: networks_lib.Params
    critic_opt_state: optax.OptState
    target_critic_params: networks_lib.Params
    key: jax_types.PRNGKey
    steps: int


def expectile_loss(diff, expectile=0.8):
    weight = jnp.where(diff > 0, expectile, (1 - expectile))
    return weight * (diff**2)


class IQLLearnerCore(learner_core.LearnerCore[TrainingState, core_types.Transition]):
    """IQL Learner core."""

    def __init__(
        self,
        networks: iql_networks.IQLNetworks,
        policy_optimizer: optax.GradientTransformation,
        critic_optimizer: optax.GradientTransformation,
        value_optimizer: optax.GradientTransformation,
        discount: float = 0.99,
        tau: float = 0.005,
        expectile: float = 0.8,
        temperature: float = 0.1,
    ):
        """Create an instance of the IQLLearner.

        Args:
            random_key (types.PRNGKey): random seed used by the learner.
            networks (iql_networks.IQLNetworks): networks used by the learner.
            policy_optimizer (optax.GradientTransformation): optimizer for policy.
            critic_optimizer (optax.GradientTransformation): optimizer for critic.
            value_optimizer (optax.GradientTransformation): optimizer for value critic.
            discount (float, optional): additional discount. Defaults to 0.99.
            tau (float, optional): target soft update rate. Defaults to 0.005.
            expectile (float, optional): expectile for training critic. Defaults to 0.8.
            temperature (float, optional): temperature for the AWR. Defaults to 0.1.

        Returns:
            An instance of IQLLearner
        """

        policy_network = networks.policy_network
        value_network = networks.value_network
        critic_network = networks.critic_network

        def awr_actor_loss_fn(
            policy_params: networks_lib.Params,
            key: jax_types.PRNGKey,
            target_critic_params: networks_lib.Params,
            value_params: networks_lib.Params,
            batch: core_types.Transition,
        ) -> Tuple[jnp.ndarray, Any]:
            v = value_network.apply(value_params, batch.observation)
            q1, q2 = critic_network.apply(
                target_critic_params, batch.observation, batch.action
            )
            q = jnp.minimum(q1, q2)
            exp_a = jnp.exp((q - v) * temperature)
            exp_a = jnp.minimum(exp_a, 100.0)
            dist = policy_network.apply(
                policy_params, batch.observation, is_training=True, key=key
            )
            log_probs = dist.log_prob(batch.action)
            actor_loss = -(exp_a * log_probs).mean()

            return actor_loss, {"actor_loss": actor_loss, "advantage": jnp.mean(q - v)}

        def value_loss_fn(
            value_params: networks_lib.Params,
            target_critic_params: networks_lib.Params,
            batch: core_types.Transition,
        ) -> Tuple[jnp.ndarray, Any]:
            q1, q2 = critic_network.apply(
                target_critic_params, batch.observation, batch.action
            )
            q = jnp.minimum(q1, q2)
            v = value_network.apply(value_params, batch.observation)
            value_loss = expectile_loss(q - v, expectile).mean()
            return value_loss, {"value_loss": value_loss, "value": v.mean()}

        def critic_loss_fn(
            critic_params: networks_lib.Params,
            target_value_params: networks_lib.Params,
            batch: core_types.Transition,
        ) -> Tuple[jnp.ndarray, Any]:
            next_v = value_network.apply(target_value_params, batch.next_observation)
            target_q = batch.reward + discount * batch.discount * next_v
            q1, q2 = critic_network.apply(
                critic_params, batch.observation, batch.action
            )
            critic_loss = ((q1 - target_q) ** 2 + (q2 - target_q) ** 2).mean()
            return critic_loss, {
                "critic_loss": critic_loss,
                "q1": q1.mean(),
                "q2": q2.mean(),
            }

        actor_grad_fn = jax.grad(awr_actor_loss_fn, has_aux=True)
        value_grad_fn = jax.grad(value_loss_fn, has_aux=True)
        critic_grad_fn = jax.grad(critic_loss_fn, has_aux=True)

        def update_step(
            state: TrainingState, batch: core_types.Transition
        ) -> Tuple[TrainingState, _Metrics]:
            # Update value network first
            policy_key, key = jax.random.split(state.key)
            value_grads, value_metrics = value_grad_fn(
                state.value_params, state.target_critic_params, batch
            )
            value_updates, value_opt_state = value_optimizer.update(
                value_grads, state.value_opt_state
            )
            value_params = optax.apply_updates(state.value_params, value_updates)
            # Update policy network
            policy_grads, policy_metrics = actor_grad_fn(
                state.policy_params,
                policy_key,
                state.target_critic_params,
                value_params,
                batch,
            )
            policy_updates, policy_opt_state = policy_optimizer.update(
                policy_grads, state.policy_opt_state
            )
            policy_params = optax.apply_updates(state.policy_params, policy_updates)
            # Update critic network
            critic_grads, critic_metrics = critic_grad_fn(
                state.critic_params, value_params, batch
            )
            critic_updates, critic_opt_state = critic_optimizer.update(
                critic_grads, state.critic_opt_state
            )
            critic_params = optax.apply_updates(state.critic_params, critic_updates)

            target_critic_params = optax.incremental_update(
                critic_params, state.target_critic_params, tau
            )
            state = TrainingState(
                policy_params=policy_params,
                policy_opt_state=policy_opt_state,
                critic_params=critic_params,
                critic_opt_state=critic_opt_state,
                value_params=value_params,
                value_opt_state=value_opt_state,
                target_critic_params=target_critic_params,
                key=key,
                steps=state.steps + 1,
            )
            return state, {**critic_metrics, **value_metrics, **policy_metrics}

        self._update_step = jax.jit(update_step)

        def make_initial_state(key):
            policy_key, critic_key, value_key, key = jax.random.split(key, 4)
            policy_params = policy_network.init(policy_key)
            policy_opt_state = policy_optimizer.init(policy_params)
            critic_params = critic_network.init(critic_key)
            critic_opt_state = critic_optimizer.init(critic_params)
            value_params = value_network.init(value_key)
            value_opt_state = value_optimizer.init(value_params)
            state = TrainingState(
                policy_params=policy_params,
                policy_opt_state=policy_opt_state,
                critic_params=critic_params,
                critic_opt_state=critic_opt_state,
                target_critic_params=critic_params,
                value_params=value_params,
                value_opt_state=value_opt_state,
                key=key,
                steps=0,
            )
            return state

        self._init_fn = make_initial_state

    def init(self, key: jax_types.PRNGKey) -> TrainingState:
        return self._init_fn(key)

    def step(self, state: TrainingState, transitions: core_types.Transition):
        # Perform a single learner step
        state, metrics = self._update_step(state, transitions)
        return state, metrics

    def get_variables(self, state, names):
        variables = {
            "policy": state.policy_params,
            "critic": state.critic_params,
        }
        return [variables[name] for name in names]


class IQLLearner(learners.DefaultJaxLearner[TrainingState, core_types.Transition]):
    def __init__(
        self,
        networks: iql_networks.IQLNetworks,
        dataset: Iterator[core_types.Transition],
        random_key: jax_types.PRNGKey,
        policy_optimizer: optax.GradientTransformation,
        critic_optimizer: optax.GradientTransformation,
        value_optimizer: optax.GradientTransformation,
        discount: float = 0.99,
        tau: float = 0.005,
        expectile: float = 0.8,
        temperature: float = 0.1,
        counter: Optional[counting.Counter] = None,
        logger: Optional[loggers.Logger] = None,
    ):
        core = IQLLearnerCore(
            networks=networks,
            policy_optimizer=policy_optimizer,
            critic_optimizer=critic_optimizer,
            value_optimizer=value_optimizer,
            discount=discount,
            tau=tau,
            expectile=expectile,
            temperature=temperature,
        )
        super().__init__(
            core, random_key, iterator=dataset, logger=logger, counter=counter
        )

corax/agents/jax/iql/learning.py
ethanluoyc/corax
# type: ignore
"""Default networks used by the IQL agent."""
import dataclasses
from typing import Callable, Optional, Sequence, Tuple

import haiku as hk
import jax.numpy as jnp
import numpy as onp
from jax import nn
from tensorflow_probability.substrates import jax as tfp  # type: ignore

from corax import specs
from corax.jax import networks as networks_lib

tfd = tfp.distributions
tfb = tfp.bijectors

LOG_STD_MIN = -10.0
LOG_STD_MAX = 2.0


# The default uses onp as we do not want to trigger JAX init
def _default_init(scale: float = onp.sqrt(2)):
    return hk.initializers.Orthogonal(scale)


class Policy(hk.Module):
    """Policy network for IQL."""

    def __init__(
        self,
        hidden_dims: Sequence[int],
        action_dim: int,
        state_dependent_std: bool = True,
        log_std_scale: float = 1.0,
        log_std_min: Optional[float] = None,
        log_std_max: Optional[float] = None,
        tanh_squash_distribution: bool = True,
        dropout_rate=None,
        name=None,
    ):
        super().__init__(name=name)
        self.hidden_dims = hidden_dims
        self.action_dim = action_dim
        self.state_dependent_std = state_dependent_std
        self.log_std_scale = log_std_scale
        self.log_std_min = log_std_min
        self.log_std_max = log_std_max
        self.tanh_squash_distribution = tanh_squash_distribution
        self.dropout_rate = dropout_rate

    def __call__(
        self, observations: jnp.ndarray, *, is_training=True, key=None
    ) -> tfd.Distribution:  # type: ignore
        torso = hk.nets.MLP(
            output_sizes=self.hidden_dims,
            activate_final=True,
            w_init=_default_init(),
        )  # type: ignore
        outputs = torso(
            observations,
            dropout_rate=self.dropout_rate if is_training else None,
            rng=key if is_training and self.dropout_rate is not None else None,
        )

        means = hk.Linear(self.action_dim, w_init=_default_init())(outputs)  # type: ignore

        if self.state_dependent_std:
            log_stds = hk.Linear(
                self.action_dim, w_init=_default_init(self.log_std_scale)  # type: ignore
            )(outputs)
        else:
            log_stds = hk.get_parameter("log_stds", (self.action_dim,), init=jnp.zeros)

        log_std_min = self.log_std_min or LOG_STD_MIN
        log_std_max = self.log_std_max or LOG_STD_MAX
        log_stds = jnp.clip(log_stds, log_std_min, log_std_max)

        if not self.tanh_squash_distribution:
            means = nn.tanh(means)

        base_dist = tfd.MultivariateNormalDiag(loc=means, scale_diag=jnp.exp(log_stds))
        if self.tanh_squash_distribution:
            return tfd.TransformedDistribution(
                distribution=base_dist, bijector=tfb.Tanh()
            )
        else:
            return base_dist


class ValueCritic(hk.Module):
    """Value network."""

    def __init__(self, hidden_dims: Sequence[int]):
        super().__init__()
        self.hidden_dims = hidden_dims

    def __call__(self, observations: jnp.ndarray) -> jnp.ndarray:
        critic_value = hk.nets.MLP(
            output_sizes=(*self.hidden_dims, 1),  # type: ignore
            w_init=_default_init(),
        )(observations)
        return jnp.squeeze(critic_value, -1)


class Critic(hk.Module):
    """A single critic network."""

    def __init__(
        self,
        hidden_dims: Sequence[int],
        activations: Callable[[jnp.ndarray], jnp.ndarray] = nn.relu,
    ):
        super().__init__()
        self.hidden_dims = hidden_dims
        self.activations = activations

    def __call__(self, observations: jnp.ndarray, actions: jnp.ndarray) -> jnp.ndarray:
        inputs = jnp.concatenate([observations, actions], -1)
        critic = hk.nets.MLP(
            output_sizes=(*self.hidden_dims, 1),  # type: ignore
            w_init=_default_init(),
        )(inputs)
        return jnp.squeeze(critic, -1)


class DoubleCritic(hk.Module):
    """Double critic network."""

    def __init__(
        self,
        hidden_dims: Sequence[int],
        activations: Callable[[jnp.ndarray], jnp.ndarray] = nn.relu,
    ):
        super().__init__()
        self.hidden_dims = hidden_dims
        self.activations = activations

    def __call__(
        self, observations: jnp.ndarray, actions: jnp.ndarray
    ) -> Tuple[jnp.ndarray, jnp.ndarray]:
        critic1 = Critic(self.hidden_dims, activations=self.activations)(  # type: ignore
            observations, actions
        )
        critic2 = Critic(self.hidden_dims, activations=self.activations)(  # type: ignore
            observations, actions
        )
        return critic1, critic2


@dataclasses.dataclass
class IQLNetworks:
    policy_network: networks_lib.FeedForwardNetwork
    critic_network: networks_lib.FeedForwardNetwork
    value_network: networks_lib.FeedForwardNetwork


def apply_policy_and_sample(
    networks: IQLNetworks,
    action_spec: specs.BoundedArray,
    eval_mode: bool = False,
):
    def policy_network(params, key, obs):
        # Note: the policy network is only used for inference, regardless
        # of whether we are sampling actions or using the mode. Thus
        # is_training is False
        action_dist = networks.policy_network.apply(params, obs, is_training=False)
        action = action_dist.mode() if eval_mode else action_dist.sample(seed=key)
        return jnp.clip(action, action_spec.minimum, action_spec.maximum)

    return policy_network


def make_networks(
    spec: specs.EnvironmentSpec,
    hidden_dims=(256, 256),
    dropout_rate=None,
):
    action_dim = spec.actions.shape[-1]

    def _actor_fn(observations, *, is_training=True, key=None):
        # To handle potential use of dropout, the policy used by IQL
        # requires two additional args: `is_training` and `rng`
        # which are used by the IQL learner to optionally enable dropout
        # in learning the policy.
        return Policy(
            hidden_dims,
            action_dim,
            log_std_scale=1e-3,
            log_std_min=-5.0,
            state_dependent_std=False,
            tanh_squash_distribution=False,
            dropout_rate=dropout_rate,
        )(observations, is_training=is_training, key=key)

    def _critic_fn(o, a):
        a = jnp.clip(a, spec.actions.minimum, spec.actions.maximum)
        return DoubleCritic(hidden_dims)(o, a)

    def _value_fn(o):
        return ValueCritic(hidden_dims)(o)

    policy = hk.without_apply_rng(hk.transform(_actor_fn))
    critic = hk.without_apply_rng(hk.transform(_critic_fn))
    value = hk.without_apply_rng(hk.transform(_value_fn))
    dummy_obs = jnp.zeros((1, *spec.observations.shape), dtype=spec.observations.dtype)
    dummy_action = jnp.zeros((1, *spec.actions.shape), dtype=spec.actions.dtype)

    return IQLNetworks(
        policy_network=networks_lib.FeedForwardNetwork(
            init=lambda key: policy.init(key, dummy_obs, is_training=False),
            apply=policy.apply,
        ),
        critic_network=networks_lib.FeedForwardNetwork(
            init=lambda key: critic.init(key, dummy_obs, dummy_action),
            apply=critic.apply,
        ),
        value_network=networks_lib.FeedForwardNetwork(
            init=lambda key: value.init(key, dummy_obs),
            apply=value.apply,
        ),
    )

corax/agents/jax/iql/networks.py
ethanluoyc/corax
import abc
from typing import Generic, List, Sequence, Tuple

from corax import types
from corax.jax import types as jax_types


class LearnerCore(abc.ABC, Generic[jax_types.TrainingState, jax_types.Sample]):
    @abc.abstractmethod
    def init(self, key: jax_types.PRNGKey) -> jax_types.TrainingState:
        ...

    @abc.abstractmethod
    def step(
        self, state: jax_types.TrainingState, batch: jax_types.Sample
    ) -> Tuple[jax_types.TrainingState, jax_types.TrainingMetrics]:
        ...

    @abc.abstractmethod
    def get_variables(
        self, state: jax_types.TrainingState, names: Sequence[str]
    ) -> List[types.NestedArray]:
        ...

corax/agents/jax/learner_core.py
ethanluoyc/corax
import time
from typing import Any, Callable, Generic, Iterator, List, Optional, Sequence

import jax

from corax import core
from corax import types
from corax.agents.jax import learner_core as learner_core_lib
from corax.jax import types as jax_types
from corax.jax import utils as jax_utils
from corax.utils import counting
from corax.utils import loggers


class DefaultJaxLearner(
    core.Learner, Generic[jax_types.TrainingState, jax_types.Sample]
):
    """A generic JAX learner that wraps a pure LearnerCore."""

    _state: jax_types.TrainingState

    def __init__(
        self,
        learner_core: learner_core_lib.LearnerCore[
            jax_types.TrainingState, jax_types.Sample
        ],
        key: jax_types.PRNGKey,
        iterator: Iterator[jax_types.Sample],
        logger: Optional[loggers.Logger] = None,
        counter: Optional[counting.Counter] = None,
        num_sgd_steps_per_step: int = 1,
        postprocess_aux: Optional[Callable[[Any], Any]] = None,
    ) -> None:
        self._learner_core = learner_core
        self._iterator = iterator
        self._state = self._learner_core.init(key)

        update_step = jax_utils.process_multiple_batches(
            self._learner_core.step,
            num_sgd_steps_per_step,
            postprocess_aux=postprocess_aux,
        )
        self._update_step = jax.jit(update_step)

        self._logger = logger or loggers.make_default_logger(
            "learner", asynchronous=True, serialize_fn=jax_utils.fetch_devicearray
        )
        self._counter = counter or counting.Counter()

        self._timestamp = None

    def step(self):
        sample = next(self._iterator)
        self._state, metrics = self._update_step(self._state, sample)

        # Compute elapsed time.
        timestamp = time.time()
        elapsed_time = timestamp - self._timestamp if self._timestamp else 0
        self._timestamp = timestamp

        counts = self._counter.increment(steps=1, walltime=elapsed_time)

        if elapsed_time > 0:
            metrics["steps_per_second"] = 1 / elapsed_time
        else:
            metrics["steps_per_second"] = 0.0

        if self._logger:
            self._logger.write({**counts, **metrics})

    def get_variables(self, names: Sequence[str]) -> List[types.NestedArray]:
        return self._learner_core.get_variables(self._state, names)

    def save(self) -> Any:
        return self._state

    def restore(self, state: Any) -> None:
        self._state = state

corax/agents/jax/learners.py
ethanluoyc/corax
"""Runners used for executing local agents."""

import time
from typing import Callable

import dm_env

from corax import core
from corax import types


class LocalLayout(core.Actor):
    """Actor which learns (updates its parameters) when `update` is called.

    This combines a base actor and a learner. Whenever `update` is called
    on the wrapping actor the learner will take a step (e.g. one step of gradient
    descent) as long as there is data available for training
    (provided iterator and replay_tables are used to check for that).
    Selecting actions and making observations are handled by the base actor.
    Intended to be used by the `run_experiment` only.
    """

    def __init__(
        self,
        actor: core.Actor,
        learner: core.Learner,
        iterator: core.PrefetchingIterator,
        can_sample: Callable[[], bool],
    ):
        """Initializes _LearningActor.

        Args:
          actor: Actor to be wrapped.
          learner: Learner on which step() is to be called when there is data.
          iterator: Iterator used by the Learner to fetch training data.
          replay_tables: Collection of tables from which Learner fetches data
            through the iterator.
          sample_sizes: For each table from `replay_tables`, how many elements the
            table should have available for sampling to wait for the `iterator` to
            prefetch a batch of data. Otherwise more experience needs to be
            collected by the actor.
          checkpointer: Checkpointer to save the state on update.
        """
        self._actor = actor
        self._learner = learner
        self._iterator = iterator
        self._learner_steps = 0
        self._can_sample = can_sample

    def select_action(self, observation: types.NestedArray) -> types.NestedArray:
        return self._actor.select_action(observation)

    def observe_first(self, timestep: dm_env.TimeStep):
        self._actor.observe_first(timestep)

    def observe(self, action: types.NestedArray, next_timestep: dm_env.TimeStep):
        self._actor.observe(action, next_timestep)

    def _maybe_train(self):
        trained = False
        while True:
            if self._iterator.ready():
                self._learner.step()
                batches = self._iterator.retrieved_elements() - self._learner_steps
                self._learner_steps += 1
                assert batches == 1, (
                    "Learner step must retrieve exactly one element from the iterator"
                    f" (retrieved {batches}). Otherwise agent can deadlock. Example "
                    "cause is that your chosen agent"
                    "s Builder has a `make_learner` "
                    "factory that prefetches the data but it shouldn"
                    "t."
                )
                trained = True
            else:
                # Wait for the iterator to fetch more data from the table(s) only
                # if there plenty of data to sample from each table.
                if not self._can_sample():
                    return trained
                # Let iterator's prefetching thread get data from the table(s).
                time.sleep(0.001)

    def update(
        self,
    ):  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
        if self._maybe_train():
            # Update the actor weights only when learner was updated.
            self._actor.update()

corax/agents/jax/local_layout.py
ethanluoyc/corax
from corax.agents.jax.oril.learning import ORILLearner
from corax.agents.jax.oril.learning import ORILSample
from corax.agents.jax.oril.losses import oril_loss
from corax.agents.jax.oril.losses import oril_pu_loss
from corax.agents.jax.oril.networks import make_reward_network

corax/agents/jax/oril/__init__.py
ethanluoyc/corax
"""Tests for ORIL agent."""
import jax
import optax
from absl.testing import absltest

from corax import specs
from corax.agents.jax import iql
from corax.agents.jax import learners
from corax.agents.jax import oril
from corax.testing import fakes


class ORILLearnerTest(absltest.TestCase):
    """Integration tests for ORIL learner"""

    def test_oril_learner(self):
        # Create a fake environment to test with.
        environment = fakes.ContinuousEnvironment(
            action_dim=2, observation_dim=3, episode_length=10, bounded=True
        )
        environment_spec = specs.make_environment_spec(environment)
        expert_iterator = fakes.transition_iterator(environment)(32)
        unlabeled_iterator = fakes.transition_iterator(environment)(32)
        reward_network = oril.make_reward_network(environment_spec)
        key = jax.random.PRNGKey(0)
        oril_key, learner_key = jax.random.split(key)

        def make_oril_iterator(expert_iterator, unlabeled_iterator, offline_iterator):
            while True:
                # pylint: disable=stop-iteration-return
                expert_sample = next(expert_iterator)
                unlabeled_sample = next(unlabeled_iterator)
                offline_rl_sample = next(offline_iterator)
                yield oril.ORILSample(
                    expert_sample, unlabeled_sample, offline_rl_sample
                )

        def make_offline_learner(iterator):
            iql_networks = iql.make_networks(environment_spec)
            learner_core = iql.IQLLearnerCore(
                iql_networks,
                policy_optimizer=optax.adam(1e-4),
                critic_optimizer=optax.adam(1e-4),
                value_optimizer=optax.adam(1e-4),
            )
            return learners.DefaultJaxLearner(learner_core, learner_key, iterator)

        learner = oril.ORILLearner(
            make_oril_iterator(expert_iterator, unlabeled_iterator, unlabeled_iterator),
            make_offline_learner,
            reward_network,
            loss_fn=oril.oril_loss,
            random_key=oril_key,
            optimizer=optax.adam(1e-4),
        )
        learner.step()


if __name__ == "__main__":
    absltest.main()

corax/agents/jax/oril/agent_test.py
ethanluoyc/corax
"""Learner for ORIL rewarder."""
import functools
import itertools
import time
from typing import Any, Callable, Iterator, List, NamedTuple, Optional, Sequence, Tuple

import jax
import optax

import corax
from corax import types
from corax.jax import networks as networks_lib
from corax.jax import types as jax_types
from corax.jax import utils
from corax.utils import counting
from corax.utils import loggers


class RewarderTrainingState(NamedTuple):
    params: networks_lib.Params
    optimizer_state: optax.OptState
    key: jax_types.PRNGKey
    steps: int


class TrainingState(NamedTuple):
    rewarder_state: RewarderTrainingState
    learner_state: Any


class ORILSample(NamedTuple):
    expert_sample: types.Transition
    unlabeled_sample: types.Transition
    learner_sample: types.Transition


class ORILLearner(corax.Learner):
    def __init__(
        self,
        iterator: Iterator[ORILSample],
        offline_learner_factory: Callable[[Iterator[types.Transition]], corax.Learner],
        network: networks_lib.FeedForwardNetwork,
        loss_fn,
        random_key: jax_types.PRNGKey,
        optimizer: optax.GradientTransformation,
        counter: Optional[counting.Counter] = None,
        logger: Optional[loggers.Logger] = None,
    ):
        # self._iterator = iterator
        iterator, offline_rl_iterator = itertools.tee(iterator)
        self._iterator = iterator
        offline_rl_iterator = (
            self._process_sample(sample.learner_sample)
            for sample in offline_rl_iterator
        )
        self._offline_learner = offline_learner_factory(offline_rl_iterator)

        def update_step(
            state: RewarderTrainingState,
            data: Tuple[types.Transition, types.Transition],
        ):
            key, _ = jax.random.split(state.key)
            compute_loss = functools.partial(loss_fn, network.apply)
            expert_transitions, unlabaled_transitions = data
            grads, metrics = jax.grad(compute_loss, has_aux=True)(
                state.params, expert_transitions, unlabaled_transitions
            )

            (updates, opt_state) = optimizer.update(grads, state.optimizer_state)

            # Apply optimizer updates to parameters.
            params = optax.apply_updates(state.params, updates)
            new_state = RewarderTrainingState(
                params=params, optimizer_state=opt_state, steps=state.steps + 1, key=key
            )
            return new_state, metrics

        self._update_step = jax.jit(update_step)
        init_key, state_key = jax.random.split(random_key)
        params = network.init(init_key)

        self._state = RewarderTrainingState(
            params=params,
            optimizer_state=optimizer.init(params),
            key=state_key,
            steps=0,
        )

        self._counter = counter or counting.Counter()
        self._logger = logger or loggers.make_default_logger(
            "learner",
            asynchronous=True,
            serialize_fn=utils.fetch_devicearray,
            steps_key=self._counter.get_steps_key(),
        )

        # 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

        def compute_reward(params, transition):
            logits = network.apply(params, transition.observation)
            return jax.nn.sigmoid(logits)

        self._get_reward = jax.jit(compute_reward)

    def _process_sample(self, transitions):
        rewards = self._get_reward(self._state.params, transitions)
        return transitions._replace(reward=rewards)

    def step(self):
        sample = next(self._iterator)
        expert_sample = sample.expert_sample
        unlabeled_sample = sample.unlabeled_sample
        self._state, metrics = self._update_step(
            self._state, (expert_sample, unlabeled_sample)
        )
        # Update the offline learner
        self._offline_learner.step()

        # Compute elapsed time.
        timestamp = time.time()
        elapsed_time = timestamp - self._timestamp if self._timestamp else 0
        self._timestamp = timestamp

        # Increment counts and record the current time.
        counts = self._counter.increment(steps=1, walltime=elapsed_time)

        # Attempts to write the logs.
        self._logger.write({**metrics, **counts})

    def get_variables(self, names: Sequence[str]) -> List[types.NestedArray]:
        rewarder_dict = {"rewarder": self._state.params}

        learner_names = [name for name in names if name not in rewarder_dict]
        learner_dict = {}
        if learner_names:
            learner_dict = dict(
                zip(learner_names, self._offline_learner.get_variables(learner_names))
            )

        variables = [
            rewarder_dict.get(name, learner_dict.get(name, None)) for name in names
        ]
        return variables

    def save(self) -> TrainingState:
        return TrainingState(
            rewarder_state=self._state, learner_state=self._offline_learner.save()
        )

    def restore(self, state: TrainingState):
        self._state = state.rewarder_state
        self._offline_learner.restore(state.learner_state)

corax/agents/jax/oril/learning.py
ethanluoyc/corax
from typing import Callable

import jax
import jax.numpy as jnp

from corax import types
from corax.jax import networks as networks_lib

RewarderFn = Callable[
    [networks_lib.Params, networks_lib.Observation], networks_lib.Logits
]


def oril_loss(
    rewarder_fn: RewarderFn,
    params: networks_lib.Params,
    expert_transitions: types.Transition,
    unlabeled_transitions: types.Transition,
):
    """Loss function for learning reward function.

    See https://arxiv.org/pdf/2011.13885.pdf Eqn 1
    """
    expert_logits = rewarder_fn(params, expert_transitions.observation)
    unlabeled_logits = rewarder_fn(params, unlabeled_transitions.observation)
    # R(s) = signmoid(NN(s))
    # -log R(expert) + -log(1-R(unlabeled))
    # = -log(sigmoid(NN(expert))) + -log(1-sigmoid(NN(unlabeled)))
    # = -log_sigmoid(NN(expert)) + -log(-sigmoid(NN(unlabeled)))
    expert_loss = -jax.nn.log_sigmoid(expert_logits)
    unlabeled_loss = jax.nn.softplus(unlabeled_logits)
    total_loss = jnp.mean(expert_loss + unlabeled_loss)
    metrics = {
        "total_loss": total_loss,
        "expert_loss": jnp.mean(expert_loss),
        "unlabeled_loss": jnp.mean(unlabeled_loss),
    }
    return total_loss, metrics


# Positive-unlabeled learning (PU-learning)
def oril_pu_loss(
    rewarder_fn: RewarderFn,
    params: networks_lib.Params,
    expert_transitions: types.Transition,
    unlabeled_transitions: types.Transition,
    gamma: float = 0.5,
):
    """ORIL Positive-Unlabeled (PU) Learning loss.

    See https://arxiv.org/pdf/2011.13885.pdf Eqn 5
    """
    expert_logits = rewarder_fn(params, expert_transitions.observation)
    unlabeled_logits = rewarder_fn(params, unlabeled_transitions.observation)
    # https://arxiv.org/pdf/2011.13885.pdf Eqn (5)
    #   gamma * -log(R(expert)) + (1 - gamma) * -log(1-R(false))
    # = gamma * -log(R(expert)) + -log(1-R(unlabaled)) - gamma * -log(1-R(expert))

    # gamma * -log(R(expert))
    positive_loss = gamma * -jax.nn.log_sigmoid(expert_logits)
    # -log(1-R(unlabaled)) - gamma * -log(1-R(expert))
    negative_loss = jax.nn.softplus(unlabeled_logits) - gamma * jax.nn.softplus(
        expert_logits
    )
    total_loss = jnp.mean(positive_loss + negative_loss)
    metrics = {
        "total_loss": total_loss,
        "positive_loss": jnp.mean(positive_loss),
        "negative_loss": jnp.mean(negative_loss),
        "expert_logits": jnp.mean(expert_logits),
        "unlabeled_logits": jnp.mean(unlabeled_logits),
    }
    return total_loss, metrics

corax/agents/jax/oril/losses.py
ethanluoyc/corax
# type: ignore
"""Tests for ORIL loss functions."""
import jax
import jax.numpy as jnp
import tree
from absl.testing import absltest

from corax import types
from corax.agents.jax.oril import losses


class ORILLossTest(absltest.TestCase):
    def setUp(self):
        def dummy_rewarder_fn(params, obs):
            del params
            return jnp.square(obs)

        self.rewarder_fn = dummy_rewarder_fn

        zero_transition = types.Transition(0.1, 0.0, 0.0, 0.0, 0.0)
        self.unlabeled_transitions = tree.map_structure(
            lambda x: jnp.expand_dims(x, axis=0), zero_transition
        )

        one_transition = types.Transition(1.0, 0.0, 0.0, 0.0, 0.0)
        self.expert_transitions = tree.map_structure(
            lambda x: jnp.expand_dims(x, axis=0), one_transition
        )

    def test_oril_loss(self):
        loss, _ = losses.oril_loss(
            self.rewarder_fn, {}, self.expert_transitions, self.unlabeled_transitions  # type: ignore
        )

        r_e = jax.nn.sigmoid(
            self.rewarder_fn(None, self.expert_transitions.observation)  # type: ignore
        )
        r_u = jax.nn.sigmoid(
            self.rewarder_fn(None, self.unlabeled_transitions.observation)  # type: ignore
        )

        expected_loss = -jnp.log(r_e) - jnp.log(1 - r_u)

        self.assertAlmostEqual(loss, expected_loss, places=6)

    def test_oril_pu_loss(self):
        gamma = 0.7
        loss, _ = losses.oril_pu_loss(
            self.rewarder_fn,
            {},
            self.expert_transitions,
            self.unlabeled_transitions,
            gamma=gamma,
        )

        r_e = jax.nn.sigmoid(
            self.rewarder_fn(None, self.expert_transitions.observation)  # type: ignore
        )
        r_u = jax.nn.sigmoid(
            self.rewarder_fn(None, self.unlabeled_transitions.observation)  # type: ignore
        )

        expected_loss = (
            gamma * (-jnp.log(r_e)) + -jnp.log(1 - r_u) - gamma * -jnp.log(1 - r_e)
        )

        self.assertAlmostEqual(loss, expected_loss, places=6)


if __name__ == "__main__":
    absltest.main()

corax/agents/jax/oril/losses_test.py
ethanluoyc/corax
from typing import Sequence

import haiku as hk
import jax.numpy as jnp

from corax import specs
from corax.jax import networks as networks_lib
from corax.jax import utils


def make_reward_network(
    spec: specs.EnvironmentSpec, hidden_layer_sizes: Sequence[int] = (256, 256)
):
    def _reward_fn(obs):
        network = hk.nets.MLP(tuple(hidden_layer_sizes) + (1,))  # type: ignore
        reward = network(obs)
        return jnp.squeeze(reward)

    reward_network = hk.without_apply_rng(hk.transform(_reward_fn))
    dummy_obs = utils.add_batch_dim(utils.zeros_like(spec.observations))
    return networks_lib.FeedForwardNetwork(
        init=lambda key: reward_network.init(key, dummy_obs),
        apply=reward_network.apply,
    )

corax/agents/jax/oril/networks.py
ethanluoyc/corax
import dm_env

from corax import adders
from corax import types
from corax.agents.jax.otr import rewarder


class OTAdder(adders.Adder):
    """Adder wrapper substituting OT rewards."""

    def __init__(
        self,
        direct_rl_adder: adders.Adder,
        otil_rewarder: rewarder.OTRewarder,
    ):
        self._adder = direct_rl_adder
        self._rewarder = otil_rewarder
        self._steps = []
        self._timesteps = []

    def add_first(self, timestep: dm_env.TimeStep):
        self._steps = []
        self._timesteps = []
        self._timesteps.append(timestep)

    def add(
        self,
        action: types.NestedArray,
        next_timestep: dm_env.TimeStep,
        extras: types.NestedArray = (),
    ):
        # TODO(yl): Handle extras
        del extras
        self._steps.append(
            types.Transition(self._timesteps[-1].observation, action, (), (), ())
        )
        self._timesteps.append(next_timestep)
        if next_timestep.last():
            self._add_episode()

    def _add_episode(self):
        # Compute pseudo-rewards and add
        psuedo_rewards = self._rewarder.compute_offline_rewards(self._steps[:])
        first_timestep = self._timesteps[0]
        self._adder.add_first(first_timestep._replace(reward=psuedo_rewards[0]))
        actions = [step.action for step in self._steps]
        assert len(actions) == len(self._timesteps) - 1 == len(psuedo_rewards)
        for action, next_ts, pr in zip(actions, self._timesteps[1:], psuedo_rewards):
            self._adder.add(action, next_ts._replace(reward=pr))
        self._adder.reset()

    def reset(self):
        self._adder.reset()
        self._steps = []
        self._timesteps = []

corax/agents/jax/otr/adder.py
ethanluoyc/corax
# type: ignore
"""Rewarder for computing rewards using Optimal Transport."""
from typing import Any, Callable, Optional, Protocol, Sequence

import chex
import jax
import jax.numpy as jnp
import numpy as onp
import ott
import ott.geometry.costs
from ott.geometry import pointcloud
from ott.problems.linear import linear_problem
from ott.solvers.linear import sinkhorn

from corax import types
from corax.jax import networks as networks_lib
from corax.jax import running_statistics
from corax.jax import variable_utils

EncoderFn = Callable[[networks_lib.Params, networks_lib.Observation], types.NestedArray]
PreprocessorState = Any


class Preprocessor(Protocol):
    """Interface for preprocessor.

    The preprocessor is used for extracting informative state
    representation from the observations.

    Implementation of the Preprocessor protocol is expected
    to be stateless. This allows the preprocessor to be used
    in a JAX jitted function.
    """

    def init(self):
        ...

    def update(self, state: PreprocessorState, atoms) -> PreprocessorState:
        ...

    def preprocess(self, params: networks_lib.Params, state: PreprocessorState, atoms):
        ...


class EncoderPreprocessor(Preprocessor):
    """Parametric encoder preprocessor."""

    def __init__(self, encoder_fn: EncoderFn):
        self._encoder_fn = encoder_fn

    def init(self):
        return ()

    def update(self, state, atoms):
        del state, atoms
        return ()

    def preprocess(self, params, state, atoms):
        del state
        return self._encoder_fn(params, atoms)  # pylint: disable=not-callable


class MeanStdPreprocessor(Preprocessor):
    """Running mean/std preprocessor."""

    def __init__(self, spec: types.NestedArray, partial_update: bool = False):
        self._observation_spec = spec
        self._partial_update = partial_update

    def init(self):
        return running_statistics.init_state(self._observation_spec)

    def update(self, state, atoms):
        assert atoms.ndim == 2
        if self._partial_update:
            state = running_statistics.init_state(self._observation_spec)
            state = running_statistics.update(state, atoms)
            return state
        else:
            state = running_statistics.update(state, atoms)
            return state

    def preprocess(self, params, state, atoms):
        del params
        return running_statistics.normalize(atoms, state)


class NoOpPreprocessor(Preprocessor):
    """Identity preprocessor."""

    def init(self):
        return ()

    def update(self, state, atoms):
        del state, atoms
        return ()

    def preprocess(self, params, state, atoms):
        del params, state
        return atoms


class AggregateFn(Protocol):
    """Function for combining pseudo-rewards from multiple expert demonstrations."""

    def __call__(self, rewards: chex.Array, **kwargs) -> chex.Array:
        ...


def aggregate_top_k(rewards, k=1):
    """Aggregate rewards from multiple expert demonstrations by mean of top-K demos."""
    scores = jnp.sum(rewards, axis=-1)
    _, indices = jax.lax.top_k(scores, k=k)
    return jnp.mean(rewards[indices], axis=0)


def aggregate_mean(rewards):
    """Aggregate rewards from multiple expert demonstrations by taking the mean"""
    return jnp.mean(rewards, axis=0)


class SquashingFn(Protocol):
    """Function for squashing pseudo-rewards computed from the OT computation."""

    def __call__(self, rewards: chex.Array, **kwargs) -> chex.Array:
        ...


def squashing_linear(rewards, alpha: float = 10.0):
    """Compute squashed rewards with alpha * rewards."""
    return alpha * rewards


def squashing_exponential(rewards, alpha: float = 5.0, beta: float = 5.0):
    """Compute squashed rewards with alpha * exp(beta * rewards)."""
    return alpha * jnp.exp(beta * rewards)


class OTRewarder:
    """OT rewarder.

    The rewarder measures similarities with the expert demonstration.
    """

    def __init__(
        self,
        demonstrations: Sequence[Sequence[types.Transition]],
        episode_length: int,
        preprocessor: Optional[Preprocessor] = None,
        aggregate_fn: AggregateFn = aggregate_top_k,
        squashing_fn: SquashingFn = squashing_linear,
        max_iterations: float = 100,
        threshold: float = 1e-9,
        epsilon: float = 1e-2,
        preprocessor_update_period: int = 1,
        use_actions_for_distance: bool = False,
        variable_client: Optional[variable_utils.VariableClient] = None,
    ):
        if use_actions_for_distance and preprocessor is not None:
            raise NotImplementedError("Use actions with preprocessor not yet supported")

        self._episode_length = episode_length
        self._aggregate_fn = aggregate_fn
        self._squashing_fn = squashing_fn
        self._max_iterations = max_iterations
        self._threshold = threshold
        self._epsilon = epsilon
        self._preprocessor_update_period = preprocessor_update_period
        self._variable_client = variable_client
        self._params = variable_client.params if variable_client is not None else None
        self._num_episodes_seen = 0
        self._use_actions_for_distance = use_actions_for_distance

        # Prepare expert atoms
        self._expert_atoms = []
        self._expert_weights = []
        # Vectorize atoms, pad the atoms and compute weights
        for demo in demonstrations:
            atoms, weights, _, _ = _pack_trajectory(
                demo, self._episode_length, self._use_actions_for_distance
            )
            self._expert_atoms.append(atoms)
            self._expert_weights.append(weights)

        self._expert_atoms = onp.stack(self._expert_atoms)
        self._expert_weights = onp.stack(self._expert_weights)

        self._preprocessor = preprocessor or NoOpPreprocessor()
        self._preprocessor_state = self._preprocessor.init()

        self._batched_ot_solve = jax.jit(
            jax.vmap(self._solve_ot, in_axes=(None, None, 0, 0, None, None))
        )
        self._compute_rewards = jax.jit(self._compute_otil_rewards)
        self._update_preprocessor = jax.jit(self._preprocessor.update)

    def _solve_ot(
        self, params, state, expert_atoms, expert_weights, agent_atoms, agent_weights
    ):
        agent_atoms = self._preprocessor.preprocess(params, state, agent_atoms)
        expert_atoms = self._preprocessor.preprocess(params, state, expert_atoms)
        cost_fn = ott.geometry.costs.Cosine()
        geom = pointcloud.PointCloud(
            agent_atoms, expert_atoms, cost_fn=cost_fn, epsilon=self._epsilon
        )
        solver = sinkhorn.Sinkhorn(
            threshold=self._threshold, max_iterations=self._max_iterations
        )
        problem = linear_problem.LinearProblem(geom, a=agent_weights, b=expert_weights)
        sinkhorn_output = solver(problem)
        coupling_matrix = geom.transport_from_potentials(
            sinkhorn_output.f, sinkhorn_output.g
        )
        cost_matrix = cost_fn.all_pairs(agent_atoms, expert_atoms)
        ot_costs = jnp.einsum("ij,ij->i", coupling_matrix, cost_matrix)
        return ot_costs

    def _compute_otil_rewards(
        self,
        params,
        preprocessor_state,
        all_expert_atoms,
        all_expert_weights,
        agent_atoms,
        agent_weights,
        agent_mask,
    ):
        ot_costs = self._batched_ot_solve(
            params,
            preprocessor_state,
            all_expert_atoms,
            all_expert_weights,
            agent_atoms,
            agent_weights,
        )
        # Compute pseudo rewards based on alignments
        pseudo_rewards = -ot_costs
        rewards = self._squashing_fn(pseudo_rewards)
        # Maskout rewards for padded atoms with zero
        rewards = jnp.where(agent_mask, rewards, 0.0)
        rewards = self._aggregate_fn(rewards)
        return rewards

    def compute_offline_rewards(self, agent_steps):
        """Compute rewards based on optimal transport."""
        # Vectorize atoms, pad the atoms and compute weights
        agent_atoms, agent_weights, num_agent_atoms, agent_mask = _pack_trajectory(
            agent_steps, self._episode_length, self._use_actions_for_distance
        )

        # Update preprocessor state if necessary
        if self._num_episodes_seen % self._preprocessor_update_period == 0:
            # Maybe update parameters
            if self._variable_client is not None:
                self._variable_client.update_and_wait()
                self._params = self._variable_client.params
            # Take only the first num_agent_atoms for updating the preprocessor
            # due to padding.
            self._preprocessor_state = self._update_preprocessor(
                self._preprocessor_state, agent_atoms[:num_agent_atoms]
            )

        rewards = self._compute_otil_rewards(
            self._params,
            self._preprocessor_state,
            self._expert_atoms,
            self._expert_weights,
            agent_atoms,
            agent_weights,
            agent_mask,
        )
        # "Unpad" the computed rewards
        rewards = rewards[:num_agent_atoms]

        self._num_episodes_seen += 1

        return jax.device_get(rewards)


def _pack_trajectory(
    trajectory: Sequence[types.Transition],
    max_sequence_length: int,
    use_actions: bool = False,
):
    """Pack a list of observations in a trajectory into array."""
    num_atoms = len(trajectory)
    if use_actions:
        observations = [
            onp.concatenate([atom.observation, atom.action], axis=-1)
            for atom in trajectory
        ]
    else:
        observations = [atom.observation for atom in trajectory]

    atoms = onp.stack(observations, axis=0)
    atoms = _pad(atoms, max_sequence_length)
    # Compute weights for used in OT
    # NOTE: The weights are padded with zeros so that the
    # padded atoms has zero probability.
    # This ensures that the computed OT solution is identical
    # with non-padded atoms.
    weights = onp.ones((num_atoms,)) / max_sequence_length
    weights = _pad(weights, max_sequence_length)
    # Compute a mask for the valid steps in the sequence
    mask = _pad(onp.ones(num_atoms, dtype=bool), max_sequence_length)
    return atoms, weights, num_atoms, mask


def _pad(x, max_sequence_length: int):
    paddings = [(0, max_sequence_length - x.shape[0])]
    paddings.extend([(0, 0) for _ in range(x.ndim - 1)])
    return onp.pad(x, paddings, mode="constant", constant_values=0.0)

corax/agents/jax/otr/rewarder.py
ethanluoyc/corax
import functools

import jax
import numpy as onp
import numpy.testing as np_test
from absl.testing import absltest

from corax import types
from corax.agents.jax.otr import rewarder as ot_rewarder


class OTRewarderTest(absltest.TestCase):
    def test_rewarder(self):
        num_expert_demos = 2
        episode_length = 10
        obs_size = 3
        expert_observations = jax.random.normal(
            jax.random.PRNGKey(0), (num_expert_demos, episode_length, obs_size)
        )
        expert_demos = []
        # Create expert demonstrations
        for n in range(num_expert_demos):
            traj = []
            for t in range(episode_length):
                traj.append(types.Transition(expert_observations[n, t], (), (), (), ()))
            expert_demos.append(traj)
        # Create agent demonstrations
        agent_demo = []
        agent_observations = jax.random.normal(
            jax.random.PRNGKey(0), (episode_length, obs_size)
        )
        for t in range(episode_length):
            agent_demo.append(types.Transition(agent_observations[t], (), (), (), ()))
        # Compute rewards with the rewarder
        rewarder = ot_rewarder.OTRewarder(expert_demos, episode_length)
        rewards = (rewarder.compute_offline_rewards(agent_demo)).tolist()
        self.assertLen(rewards, episode_length)

    def test_zero_rewards_for_expert_demo(self):
        num_expert_demos = 1
        episode_length = 10
        obs_size = 3
        expert_observations = jax.random.normal(
            jax.random.PRNGKey(0), (num_expert_demos, episode_length, obs_size)
        )
        expert_demos = []
        # Create expert demonstrations
        for n in range(num_expert_demos):
            traj = []
            for t in range(episode_length):
                traj.append(types.Transition(expert_observations[n, t], (), (), (), ()))
            expert_demos.append(traj)
        # Create agent demonstrations
        agent_demos = expert_demos[0]
        # Compute rewards with the rewarder
        rewarder = ot_rewarder.OTRewarder(
            expert_demos,
            episode_length,
            squashing_fn=functools.partial(ot_rewarder.squashing_linear, alpha=1.0),  # type: ignore
        )
        rewards = rewarder.compute_offline_rewards(agent_demos)
        np_test.assert_allclose(rewards, onp.zeros_like(rewards), atol=1e-3)


if __name__ == "__main__":
    absltest.main()

corax/agents/jax/otr/rewarder_test.py
ethanluoyc/corax
from corax.agents.jax.redq.learning import REDQLearnerCore
from corax.agents.jax.redq.networks import apply_policy_and_sample
from corax.agents.jax.redq.networks import make_networks
from corax.agents.jax.redq.networks import target_entropy_from_spec
from corax.utils import lazy_loader

with lazy_loader.LazyImports(__name__, False):
    from corax.agents.jax.redq.builder import REDQBuilder
    from corax.agents.jax.redq.config import REDQConfig

del lazy_loader

corax/agents/jax/redq/__init__.py
ethanluoyc/corax
"""Tests for the REDQ agent."""

import jax
import optax
from absl.testing import absltest

import corax
from corax.agents.jax import redq
from corax.testing import fakes


class AgentTest(absltest.TestCase):
    def test_train(self):
        env = fakes.ContinuousEnvironment(bounded=True)
        env_spec = corax.make_environment_spec(env)
        networks = redq.make_networks(
            env_spec, hidden_sizes=(8, 8), num_qs=5, num_min_qs=2
        )
        utd_ratio = 2
        batch_size = 2
        learner_core = redq.REDQLearnerCore(
            networks,
            policy_optimizer=optax.adam(1e-4),
            critic_optimizer=optax.adam(1e-4),
            temperature_optimizer=optax.adam(1e-4),
            target_entropy=redq.target_entropy_from_spec(env_spec.actions),
        )
        iterator = fakes.transition_iterator(env)(batch_size * utd_ratio)
        key = jax.random.PRNGKey(0)
        state = learner_core.init(key)
        state, _ = learner_core.step(state, next(iterator))


if __name__ == "__main__":
    absltest.main()

corax/agents/jax/redq/agent_test.py
ethanluoyc/corax
"""REDQ Builder."""
from typing import Callable, Iterator, List, Optional

import jax
import numpy as np
import optax
import reverb

from corax import adders
from corax import core
from corax import specs
from corax import types
from corax.adders import reverb as adders_reverb
from corax.agents.jax import actor_core as actor_core_lib
from corax.agents.jax import actors
from corax.agents.jax import builders
from corax.agents.jax import learners
from corax.agents.jax.redq import config as red_config
from corax.agents.jax.redq import learning
from corax.agents.jax.redq import networks as redq_networks
from corax.datasets import reverb as datasets
from corax.jax import networks as networks_lib
from corax.jax import utils
from corax.jax import variable_utils
from corax.utils import counting
from corax.utils import loggers

REDQNetworks = redq_networks.REDQNetworks
DemonstrationFactory = Callable[[int], Iterator[types.Transition]]


def _generate_rlpd_samples(
    demonstration_iterator: Iterator[types.Transition],
    replay_iterator: Iterator[types.Transition],
) -> Iterator[reverb.ReplaySample]:
    """Generator which creates the sample iterator for RLPD.

    Args:
      demonstration_iterator: Iterator of demonstrations.
      replay_iterator: Replay buffer sample iterator.

    Yields:
      Samples having a mix of offline and online data.
    """

    def combine(one, other):
        def combine_leaf(a, b):
            # No interleave is performed here but it is important to shuffle
            # the batch. See comments in `learning.py`.
            return np.concatenate([a, b], axis=0)

        return jax.tree_map(combine_leaf, one, other)

    while True:
        replay_transitions = types.Transition(*next(replay_iterator))
        offline_transitions = next(demonstration_iterator)
        combined = combine(offline_transitions, replay_transitions)
        yield combined


class REDQBuilder(
    builders.ActorLearnerBuilder[
        REDQNetworks,
        actor_core_lib.FeedForwardPolicy,
        reverb.ReplaySample,
    ]
):
    """REDQ Builder."""

    def __init__(
        self,
        config: red_config.REDQConfig,
        make_demonstrations: Optional[DemonstrationFactory] = None,
    ):
        self._config = config
        self._make_demonstrations = make_demonstrations

    def make_learner(
        self,
        random_key: networks_lib.PRNGKey,
        networks: REDQNetworks,
        dataset: Iterator[reverb.ReplaySample],
        logger_fn: loggers.LoggerFactory,
        environment_spec: specs.EnvironmentSpec,
        replay_client: Optional[reverb.Client] = None,
        counter: Optional[counting.Counter] = None,
    ) -> learners.DefaultJaxLearner:
        del replay_client
        policy_optimizer = optax.adam(self._config.actor_learning_rate)
        critic_optimizer = optax.adam(self._config.critic_learning_rate)
        temperature_optimizer = optax.adam(self._config.temperature_learning_rate)

        target_entropy = self._config.target_entropy
        if target_entropy is None:
            target_entropy = redq_networks.target_entropy_from_spec(
                environment_spec.actions
            )

        learner_core = learning.REDQLearnerCore(
            networks,
            policy_optimizer=policy_optimizer,
            critic_optimizer=critic_optimizer,
            temperature_optimizer=temperature_optimizer,
            target_entropy=target_entropy,
            discount=self._config.discount,
            tau=self._config.tau,
            init_temperature=self._config.init_temperature,
            backup_entropy=self._config.backup_entropy,
            utd_ratio=self._config.num_sgd_steps_per_step,
            reward_scale=self._config.reward_scale,
            reward_bias=self._config.reward_bias,
        )

        return learners.DefaultJaxLearner(
            learner_core, random_key, dataset, logger_fn("learner"), counter
        )

    def make_actor(
        self,
        random_key: networks_lib.PRNGKey,
        policy: actor_core_lib.FeedForwardPolicy,
        environment_spec: specs.EnvironmentSpec,
        variable_source: Optional[core.VariableSource] = None,
        adder: Optional[adders.Adder] = None,
    ) -> actors.GenericActor:
        del environment_spec
        assert variable_source is not None
        actor_core = actor_core_lib.batched_feed_forward_to_actor_core(policy)
        variable_client = variable_utils.VariableClient(
            variable_source, "policy", device="cpu"
        )
        return actors.GenericActor(
            actor_core, random_key, variable_client, adder, backend="cpu"
        )

    def make_replay_tables(
        self,
        environment_spec: specs.EnvironmentSpec,
        policy: actor_core_lib.FeedForwardPolicy,
    ) -> List[reverb.Table]:
        """Create tables to insert data into."""
        del policy
        samples_per_insert_tolerance = (
            self._config.samples_per_insert_tolerance_rate
            * self._config.samples_per_insert
        )
        error_buffer = self._config.min_replay_size * samples_per_insert_tolerance
        limiter = reverb.rate_limiters.SampleToInsertRatio(
            min_size_to_sample=self._config.min_replay_size,
            samples_per_insert=self._config.samples_per_insert,
            error_buffer=max(error_buffer, 2 * self._config.samples_per_insert),
        )
        return [
            reverb.Table(
                name=self._config.replay_table_name,
                sampler=reverb.selectors.Uniform(),
                remover=reverb.selectors.Fifo(),
                max_size=self._config.max_replay_size,
                rate_limiter=limiter,
                signature=adders_reverb.NStepTransitionAdder.signature(
                    environment_spec
                ),
            )
        ]

    def make_dataset_iterator(
        self, replay_client: reverb.Client
    ) -> Iterator[reverb.ReplaySample]:
        """Create a dataset iterator to use for learning/updating the agent."""
        sgd_steps_per_step = self._config.num_sgd_steps_per_step
        batch_size = self._config.batch_size

        if self._make_demonstrations is not None:
            offline_batch_size = int(batch_size * self._config.offline_fraction)
            online_batch_size = batch_size - offline_batch_size

            replay_iterator = datasets.make_reverb_dataset(
                table=self._config.replay_table_name,
                server_address=replay_client.server_address,
                batch_size=online_batch_size * sgd_steps_per_step,
                num_parallel_calls=max(16, 4 * jax.local_device_count()),
            ).as_numpy_iterator()
            replay_iterator = (sample.data for sample in replay_iterator)  # type: ignore

            if offline_batch_size > 0:
                offline_iterator = self._make_demonstrations(
                    offline_batch_size * sgd_steps_per_step
                )
                iterator = _generate_rlpd_samples(offline_iterator, replay_iterator)  # type: ignore
            else:
                iterator = replay_iterator

        else:
            online_batch_size = batch_size
            iterator = datasets.make_reverb_dataset(
                table=self._config.replay_table_name,
                server_address=replay_client.server_address,
                batch_size=online_batch_size * sgd_steps_per_step,
                num_parallel_calls=max(16, 4 * jax.local_device_count()),
            ).as_numpy_iterator()
            iterator = (sample.data for sample in iterator)  # type: ignore

        return utils.device_put(iterator, jax.local_devices()[0])

    def make_adder(
        self,
        replay_client: reverb.Client,
        environment_spec: Optional[specs.EnvironmentSpec],
        policy: Optional[actor_core_lib.FeedForwardPolicy],
    ) -> adders_reverb.NStepTransitionAdder:
        """Create an adder which records data generated by the actor/environment."""
        del environment_spec, policy
        return adders_reverb.NStepTransitionAdder(
            priority_fns={self._config.replay_table_name: None},
            client=replay_client,
            n_step=self._config.n_step,
            discount=self._config.discount,
        )

    def make_policy(
        self,
        networks: REDQNetworks,
        environment_spec: specs.EnvironmentSpec,
        evaluation: bool = False,
    ) -> actor_core_lib.FeedForwardPolicy:
        """Construct the policy."""
        del environment_spec

        return redq_networks.apply_policy_and_sample(networks, evaluation)

corax/agents/jax/redq/builder.py
ethanluoyc/corax
import dataclasses
from typing import Optional

from corax.adders import reverb as adders_reverb


@dataclasses.dataclass
class REDQConfig:
    """Configuration options for REDQ."""

    # Optimizer options
    batch_size: int = 256
    actor_learning_rate: float = 3e-4
    critic_learning_rate: float = 3e-4
    temperature_learning_rate: float = 3e-4

    init_temperature: float = 1.0
    backup_entropy: bool = True
    discount: float = 0.99
    n_step: int = 1
    target_entropy: Optional[float] = None
    # Target smoothing coefficient.
    tau: float = 0.005

    # Replay options
    min_replay_size: int = 10000
    max_replay_size: int = 1000000
    replay_table_name: str = adders_reverb.DEFAULT_PRIORITY_TABLE

    samples_per_insert: float = 256
    samples_per_insert_tolerance_rate: float = 0.1

    offline_fraction: float = 0.5

    # How many gradient updates to perform per step.
    num_sgd_steps_per_step: int = 1

    # Reward scaling options
    # reward may be scaled as reward_scale * reward + reward_bias
    reward_scale: float = 1.0
    reward_bias: float = 0.0

corax/agents/jax/redq/config.py
ethanluoyc/corax
from typing import Dict, List, NamedTuple, Sequence, Tuple

import jax
import jax.numpy as jnp
import optax

from corax import types
from corax.agents.jax import learner_core
from corax.agents.jax.redq import networks as red_networks
from corax.jax import networks as networks_lib
from corax.jax import types as jax_types

Metrics = Dict[str, jax.Array]


class TrainingState(NamedTuple):
    policy_params: networks_lib.Params
    critic_params: networks_lib.Params
    target_critic_params: networks_lib.Params
    log_temperature_params: networks_lib.Params

    policy_opt_state: optax.OptState
    critic_opt_state: optax.OptState
    temperature_opt_state: optax.OptState

    key: jax_types.PRNGKey


class REDQLearnerCore(learner_core.LearnerCore):
    def __init__(
        self,
        networks: red_networks.REDQNetworks,
        *,
        policy_optimizer: optax.GradientTransformation,
        critic_optimizer: optax.GradientTransformation,
        temperature_optimizer: optax.GradientTransformation,
        target_entropy: float,
        discount: float = 0.99,
        tau: float = 0.005,
        init_temperature: float = 1.0,
        backup_entropy: bool = True,
        utd_ratio: int = 1,
        reward_scale: float = 1.0,
        reward_bias: float = 0.0,
    ) -> None:
        self._networks = networks
        self._policy_optimizer = policy_optimizer
        self._critic_optimizer = critic_optimizer
        self._temperature_optimizer = temperature_optimizer
        self._tau = tau
        self._target_entropy = target_entropy
        self._discount = discount
        self._init_temperature = init_temperature
        self._backup_entropy = backup_entropy
        self._utd_ratio = utd_ratio
        self._reward_scale = reward_scale
        self._reward_bias = reward_bias

        def actor_loss_fn(
            policy_params: networks_lib.Params,
            critic_params: networks_lib.Params,
            log_temperature: networks_lib.Params,
            transitions: types.Transition,
            key: jax_types.PRNGKey,
        ) -> Tuple[jax.Array, Metrics]:
            dist = self._networks.policy_network.apply(
                policy_params, transitions.observation
            )
            actions = dist.sample(seed=key)
            log_probs = dist.log_prob(actions)
            qs = self._networks.critic_network.apply(
                critic_params, transitions.observation, actions
            )
            q = jnp.mean(qs, axis=0)
            temperature = jnp.exp(log_temperature)
            actor_loss = log_probs * temperature - q
            actor_loss = jnp.mean(actor_loss)
            return actor_loss, {"actor_loss": actor_loss, "entropy": -log_probs.mean()}

        def critic_loss_fn(
            critic_params: networks_lib.Params,
            state: TrainingState,
            transitions: types.Transition,
            key: jax_types.PRNGKey,
        ) -> Tuple[jax.Array, Metrics]:
            dist = self._networks.policy_network.apply(
                state.policy_params, transitions.next_observation
            )

            policy_key, critic_key = jax.random.split(key)
            next_actions = dist.sample(seed=policy_key)
            # [M, ...]
            subsampled_target_params = red_networks.subsample_ensemble_params(
                state.target_critic_params, critic_key, self._networks.num_min_qs
            )
            # [M, B]
            next_qs = self._networks.critic_network.apply(
                subsampled_target_params, transitions.next_observation, next_actions
            )
            # [B]
            next_q = next_qs.min(axis=0)
            target_q = (
                transitions.reward + self._discount * transitions.discount * next_q
            )

            if self._backup_entropy:
                next_log_probs = dist.log_prob(next_actions)
                temperature = jnp.exp(state.log_temperature_params)
                target_q = target_q - (
                    self._discount * transitions.discount * temperature * next_log_probs
                )

            qs = self._networks.critic_network.apply(
                critic_params, transitions.observation, transitions.action
            )
            critic_loss = jnp.mean(jnp.square(qs - target_q))
            return critic_loss, {"critic_loss": critic_loss, "q": qs.mean()}

        def temperature_loss_fn(
            log_temperature: jax.Array, entropy: jax.Array
        ) -> Tuple[jax.Array, Metrics]:
            temperature = jnp.exp(log_temperature)
            temp_loss = temperature * jnp.mean(entropy - self._target_entropy)
            return temp_loss, {
                "temperature": temperature,
                "temperature_loss": temp_loss,
            }

        self._actor_loss = actor_loss_fn
        self._critic_loss = critic_loss_fn
        self._temperature_loss = temperature_loss_fn

    def init(self, key: jax_types.PRNGKey) -> TrainingState:
        policy_key, critic_key, train_key = jax.random.split(key, 3)

        policy_params = self._networks.policy_network.init(policy_key)
        critic_params = self._networks.critic_network.init(critic_key)
        temperature_params = jnp.log(self._init_temperature)

        policy_opt_state = self._policy_optimizer.init(policy_params)
        critic_opt_state = self._critic_optimizer.init(critic_params)
        temperature_opt_state = self._temperature_optimizer.init(temperature_params)

        return TrainingState(
            policy_params=policy_params,
            critic_params=critic_params,
            target_critic_params=critic_params,
            log_temperature_params=temperature_params,
            policy_opt_state=policy_opt_state,
            critic_opt_state=critic_opt_state,
            temperature_opt_state=temperature_opt_state,
            key=train_key,
        )

    def _update_critic(
        self, state: TrainingState, transitions: types.Transition
    ) -> Tuple[TrainingState, Metrics]:
        grad_fn = jax.grad(self._critic_loss, has_aux=True)

        key, update_key = jax.random.split(state.key)

        critic_grads, critic_metrics = grad_fn(
            state.critic_params,
            state,
            transitions,
            update_key,
        )

        critic_updates, critic_opt_state = self._critic_optimizer.update(
            critic_grads, state.critic_opt_state
        )

        critic_params = optax.apply_updates(state.critic_params, critic_updates)

        target_critic_params = optax.incremental_update(
            critic_params,
            state.target_critic_params,
            self._tau,
        )

        new_state = state._replace(
            critic_params=critic_params,
            critic_opt_state=critic_opt_state,
            target_critic_params=target_critic_params,
            key=key,
        )

        return (new_state, critic_metrics)

    def _update_actor(
        self, state: TrainingState, transitions: types.Transition
    ) -> Tuple[TrainingState, Metrics]:
        grad_fn = jax.grad(self._actor_loss, has_aux=True)

        key, update_key = jax.random.split(state.key)

        policy_grads, actor_metrics = grad_fn(
            state.policy_params,
            state.critic_params,
            state.log_temperature_params,
            transitions,
            update_key,
        )

        policy_updates, policy_opt_state = self._policy_optimizer.update(
            policy_grads, state.policy_opt_state
        )

        policy_params = optax.apply_updates(state.policy_params, policy_updates)

        new_state = state._replace(
            policy_params=policy_params, policy_opt_state=policy_opt_state, key=key
        )

        return (new_state, actor_metrics)

    def _update_temperature(
        self, state: TrainingState, entropy: jax.Array
    ) -> Tuple[TrainingState, Metrics]:
        grad_fn = jax.grad(self._temperature_loss, has_aux=True)
        grads, metrics = grad_fn(state.log_temperature_params, entropy)

        temperature_updates, opt_state = self._temperature_optimizer.update(
            grads, state.temperature_opt_state
        )

        log_temperature_params = optax.apply_updates(
            state.log_temperature_params, temperature_updates
        )

        new_state = state._replace(
            log_temperature_params=log_temperature_params,
            temperature_opt_state=opt_state,
        )
        return new_state, metrics

    def step(
        self, state: TrainingState, transitions: types.Transition
    ) -> Tuple[TrainingState, Dict[str, jax.Array]]:
        # Maybe rescale reward
        transitions = transitions._replace(
            reward=transitions.reward * self._reward_scale + self._reward_bias
        )

        # We simply concatenate the replay + offline when making the iterator
        # So we need to shuffle the transitions here
        # This could alternatively be done in the iterator
        # Note that shuffling is crucial since otherwise
        # some minibatches below would only contain either replay or offline data.
        key, subkey = jax.random.split(state.key)
        # Reuse the same key to ensure all leaves are shuffled in the same way.
        transitions = jax.tree_util.tree_map(
            lambda a: jax.random.permutation(subkey, a, axis=0, independent=False),
            transitions,
        )
        state = state._replace(key=key)

        batch_size = transitions.reward.shape[0]
        if batch_size % self._utd_ratio != 0:
            raise ValueError("Total batch size to be divisible by utd_ratio")

        num_batches = self._utd_ratio
        minibatches = jax.tree_util.tree_map(
            lambda a: jnp.reshape(a, (num_batches, -1, *a.shape[1:])), transitions
        )

        state, critic_metrics = jax.lax.scan(
            self._update_critic, state, minibatches, length=num_batches
        )
        critic_metrics = jax.tree_map(jnp.mean, critic_metrics)

        # Update policy and temperature with last minibatch
        batch = jax.tree_util.tree_map(lambda a: a[-1], minibatches)

        state, actor_metrics = self._update_actor(state, batch)
        state, temperature_metrics = self._update_temperature(
            state, actor_metrics["entropy"]
        )
        return state, {**critic_metrics, **actor_metrics, **temperature_metrics}

    def get_variables(
        self, state: TrainingState, names: Sequence[str]
    ) -> List[types.NestedArray]:
        variables = {
            "policy": state.policy_params,
            "critic": state.critic_params,
            "target_critic": state.target_critic_params,
        }
        return [variables[name] for name in names]

corax/agents/jax/redq/learning.py
ethanluoyc/corax
# type: ignore
import dataclasses
import functools
from typing import Callable, Optional, Sequence

import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np

from corax import specs
from corax.jax import networks as networks_lib
from corax.jax import types as jax_types
from corax.jax import utils as jax_utils


def default_init(scale=1.0):
    return hk.initializers.VarianceScaling(scale, "fan_avg", "uniform")


@dataclasses.dataclass
class MLP(hk.Module):
    hidden_dims: Sequence[int]
    activations: Callable[[jnp.ndarray], jnp.ndarray] = jax.nn.relu
    activate_final: bool = False
    use_layer_norm: bool = False
    scale_final: Optional[float] = None
    use_pnorm: bool = False
    name: Optional[str] = None

    def __call__(self, x: jax.Array) -> jax.Array:
        for i, size in enumerate(self.hidden_dims):
            if i + 1 == len(self.hidden_dims) and self.scale_final is not None:
                x = hk.Linear(size, w_init=default_init(self.scale_final))(x)
            else:
                x = hk.Linear(size, w_init=default_init())(x)

            if i + 1 < len(self.hidden_dims) or self.activate_final:
                if self.use_layer_norm:
                    x = hk.LayerNorm(axis=-1, create_scale=True, create_offset=True)(x)
                x = self.activations(x)

        if self.use_pnorm:
            x /= jnp.linalg.norm(x, axis=-1, keepdims=True).clip(1e-10)

        return x


def ensemble_init(
    base_init: Callable[[jax_types.PRNGKey], hk.Params],
    key: jax_types.PRNGKey,
    ensemble_size: int,
):
    keys = jax.random.split(key, ensemble_size)
    return jax.vmap(base_init)(keys)


def ensemble_all_apply(base_apply, params, *args):
    return jax.vmap(base_apply, in_axes=(0,) + (None,) * len(args))(params, *args)


def subsample_ensemble_params(params, key, subset_size):
    ensemble_size = jax.tree_util.tree_leaves(params)[0].shape[0]
    ensemble_indices = jnp.arange(ensemble_size, dtype=jnp.int32)
    subset_indices = jax.random.choice(
        key, a=ensemble_indices, shape=(subset_size,), replace=False
    )
    return jax.tree_util.tree_map(lambda x: x[subset_indices], params)


def ensemble_subsample_apply(base_apply, subset_size: int, params, key, *args):
    subset_params = subsample_ensemble_params(params, key, subset_size)
    return ensemble_all_apply(base_apply, subset_params, *args)


@dataclasses.dataclass
class REDQNetworks:
    policy_network: networks_lib.FeedForwardNetwork
    critic_network: networks_lib.FeedForwardNetwork
    num_min_qs: int


def make_networks(
    spec: specs.EnvironmentSpec,
    hidden_sizes: Sequence[int] = (256, 256),
    num_qs: int = 2,
    num_min_qs: Optional[int] = None,
    critic_layer_norm: bool = True,
):
    num_dimensions = np.prod(spec.actions.shape, dtype=int)
    num_min_qs = num_min_qs or num_qs

    # TODO(yl): Support stochastic networks (e.g., dropout)
    def _actor_fn(obs):
        torso = MLP(
            hidden_sizes,
            use_layer_norm=False,
            activate_final=True,
            use_pnorm=False,
        )
        network = hk.Sequential(
            [torso, networks_lib.NormalTanhDistribution(num_dimensions)]
        )
        return network(obs)

    def _critic_fn(obs, action):
        inputs = jnp.concatenate([obs, action], axis=-1)
        network = MLP(
            list(hidden_sizes) + [1],
            use_layer_norm=critic_layer_norm,
            activate_final=False,
            use_pnorm=False,
        )
        return jnp.squeeze(network(inputs), axis=-1)

    policy = hk.without_apply_rng(hk.transform(_actor_fn))
    critic = hk.without_apply_rng(hk.transform(_critic_fn))
    dummy_obs = jax_utils.add_batch_dim(jax_utils.zeros_like(spec.observations))
    dummy_act = jax_utils.add_batch_dim(jax_utils.zeros_like(spec.actions))

    ensemble_critic_init = lambda key: ensemble_init(
        lambda rng: critic.init(rng, dummy_obs, dummy_act), key, num_qs
    )
    ensemble_critic_apply = functools.partial(ensemble_all_apply, critic.apply)

    return REDQNetworks(
        policy_network=networks_lib.FeedForwardNetwork(
            init=lambda key: policy.init(key, dummy_obs), apply=policy.apply
        ),
        critic_network=networks_lib.FeedForwardNetwork(
            init=ensemble_critic_init,
            apply=ensemble_critic_apply,
        ),
        num_min_qs=num_min_qs,
    )


def target_entropy_from_spec(action_spec: specs.BoundedArray) -> float:
    num_dimensions = action_spec.shape[-1]
    return -num_dimensions / 2


def apply_policy_and_sample(networks: REDQNetworks, evaluation: bool):
    def policy_network(params, key, observations):
        dist = networks.policy_network.apply(params, observations)
        if evaluation:
            return dist.mode()
        return dist.sample(seed=key)

    return policy_network

corax/agents/jax/redq/networks.py
ethanluoyc/corax
"""Tests for the REDQ agent."""

import chex
import jax
from absl.testing import absltest

import corax
from corax.agents.jax.redq import networks as red_networks
from corax.jax import utils as jax_utils
from corax.testing import fakes


class NetworkTest(absltest.TestCase):
    def test_network(self):
        env = fakes.ContinuousEnvironment(bounded=True)
        env_spec = corax.make_environment_spec(env)
        networks = red_networks.make_networks(
            env_spec, hidden_sizes=(8, 8), num_qs=5, num_min_qs=2
        )

        key = jax.random.PRNGKey(0)
        dummy_obs = jax_utils.tile_nested(
            jax_utils.zeros_like(env_spec.observations), 2
        )
        dummy_act = jax_utils.tile_nested(jax_utils.zeros_like(env_spec.actions), 2)

        critic_params = networks.critic_network.init(key)

        qs = networks.critic_network.apply(critic_params, dummy_obs, dummy_act)
        chex.assert_shape(qs, (5, 2))

        subset_params = red_networks.subsample_ensemble_params(critic_params, key, 2)
        subset_qs = networks.critic_network.apply(subset_params, dummy_obs, dummy_act)

        chex.assert_shape(subset_qs, (2, 2))


if __name__ == "__main__":
    absltest.main()

corax/agents/jax/redq/networks_test.py
ethanluoyc/corax
"""Twin-Delayed DDPG agent."""

from corax.agents.jax.td3.builder import TD3Builder
from corax.agents.jax.td3.config import TD3Config
from corax.agents.jax.td3.learning import TD3LearnerCore
from corax.agents.jax.td3.networks import TD3Networks
from corax.agents.jax.td3.networks import get_default_behavior_policy
from corax.agents.jax.td3.networks import make_networks

corax/agents/jax/td3/__init__.py
ethanluoyc/corax
"""Tests for the REDQ agent."""

import jax
import optax
from absl.testing import absltest

import corax
from corax.agents.jax import td3
from corax.testing import fakes


class TD3Test(absltest.TestCase):
    def test_learner_core_step(self):
        env = fakes.ContinuousEnvironment(bounded=True)
        env_spec = corax.make_environment_spec(env)
        networks = td3.make_networks(env_spec, hidden_layer_sizes=(8, 8))
        batch_size = 2
        learner_core = td3.TD3LearnerCore(
            networks,
            discount=0.99,
            policy_optimizer=optax.adam(1e-4),
            critic_optimizer=optax.adam(1e-4),
            twin_critic_optimizer=optax.adam(1e-4),
        )
        iterator = fakes.transition_iterator_from_spec(env_spec)(batch_size)
        key = jax.random.PRNGKey(0)
        state = learner_core.init(key)
        state, _ = learner_core.step(state, next(iterator))


if __name__ == "__main__":
    absltest.main()

corax/agents/jax/td3/agent_test.py
ethanluoyc/corax
# 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.

"""TD3 Builder."""
from typing import Iterator, List, Optional

import jax
import optax
import reverb
from reverb import rate_limiters

from corax import adders
from corax import core
from corax import specs
from corax import types
from corax.adders import reverb as adders_reverb
from corax.agents.jax import actor_core as actor_core_lib
from corax.agents.jax import actors
from corax.agents.jax import builders
from corax.agents.jax.td3 import config as td3_config
from corax.agents.jax.td3 import learning
from corax.agents.jax.td3 import networks as td3_networks
from corax.datasets import reverb as datasets
from corax.jax import networks as networks_lib
from corax.jax import utils
from corax.jax import variable_utils
from corax.utils import counting
from corax.utils import loggers


class TD3Builder(
    builders.ActorLearnerBuilder[
        td3_networks.TD3Networks,
        actor_core_lib.FeedForwardPolicy,
        reverb.ReplaySample,
    ]
):
    """TD3 Builder."""

    def __init__(
        self,
        config: td3_config.TD3Config,
    ):
        """Creates a TD3 learner, a behavior policy and an eval actor.

        Args:
          config: a config with TD3 hps
        """
        self._config = config

    def make_learner(
        self,
        random_key: networks_lib.PRNGKey,
        networks: td3_networks.TD3Networks,
        dataset: Iterator[reverb.ReplaySample],
        logger_fn: loggers.LoggerFactory,
        environment_spec: specs.EnvironmentSpec,
        replay_client: Optional[reverb.Client] = None,
        counter: Optional[counting.Counter] = None,
    ) -> core.Learner:
        del environment_spec, replay_client

        critic_optimizer = optax.adam(self._config.critic_learning_rate)
        twin_critic_optimizer = optax.adam(self._config.critic_learning_rate)
        policy_optimizer = optax.adam(self._config.policy_learning_rate)

        if self._config.policy_gradient_clipping is not None:
            policy_optimizer = optax.chain(
                optax.clip_by_global_norm(self._config.policy_gradient_clipping),
                policy_optimizer,
            )

        iterator = (types.Transition(*sample.data) for sample in dataset)
        return learning.TD3Learner(
            networks=networks,
            random_key=random_key,
            discount=self._config.discount,
            target_sigma=self._config.target_sigma,
            noise_clip=self._config.noise_clip,
            policy_optimizer=policy_optimizer,
            critic_optimizer=critic_optimizer,
            twin_critic_optimizer=twin_critic_optimizer,
            bc_alpha=self._config.bc_alpha,
            iterator=iterator,
            logger=logger_fn("learner"),
            counter=counter,
        )

    def make_actor(
        self,
        random_key: networks_lib.PRNGKey,
        policy: actor_core_lib.FeedForwardPolicy,
        environment_spec: specs.EnvironmentSpec,
        variable_source: Optional[core.VariableSource] = None,
        adder: Optional[adders.Adder] = None,
    ) -> core.Actor:
        del environment_spec
        assert variable_source is not None
        actor_core = actor_core_lib.batched_feed_forward_to_actor_core(policy)
        # Inference happens on CPU, so it's better to move variables there too.
        variable_client = variable_utils.VariableClient(
            variable_source, "policy", device="cpu"
        )
        return actors.GenericActor(
            actor_core, random_key, variable_client, adder, backend="cpu"
        )

    def make_replay_tables(
        self,
        environment_spec: specs.EnvironmentSpec,
        policy: actor_core_lib.FeedForwardPolicy,
    ) -> List[reverb.Table]:
        """Creates reverb tables for the algorithm."""
        del policy
        samples_per_insert_tolerance = (
            self._config.samples_per_insert_tolerance_rate
            * self._config.samples_per_insert
        )
        error_buffer = self._config.min_replay_size * samples_per_insert_tolerance
        limiter = rate_limiters.SampleToInsertRatio(
            min_size_to_sample=self._config.min_replay_size,
            samples_per_insert=self._config.samples_per_insert,
            error_buffer=error_buffer,
        )
        return [
            reverb.Table(
                name=self._config.replay_table_name,
                sampler=reverb.selectors.Uniform(),
                remover=reverb.selectors.Fifo(),
                max_size=self._config.max_replay_size,
                rate_limiter=limiter,
                signature=adders_reverb.NStepTransitionAdder.signature(
                    environment_spec
                ),
            )
        ]

    def make_dataset_iterator(
        self, replay_client: reverb.Client
    ) -> Iterator[reverb.ReplaySample]:
        """Creates a dataset iterator to use for learning."""
        dataset = datasets.make_reverb_dataset(
            table=self._config.replay_table_name,
            server_address=replay_client.server_address,
            batch_size=(self._config.batch_size),
            prefetch_size=self._config.prefetch_size,
        )
        return utils.device_put(dataset.as_numpy_iterator(), jax.devices()[0])

    def make_adder(
        self,
        replay_client: reverb.Client,
        environment_spec: Optional[specs.EnvironmentSpec],
        policy: Optional[actor_core_lib.FeedForwardPolicy],
    ) -> Optional[adders.Adder]:
        """Creates an adder which handles observations."""
        del environment_spec, policy
        return adders_reverb.NStepTransitionAdder(
            priority_fns={self._config.replay_table_name: None},
            client=replay_client,
            n_step=self._config.n_step,
            discount=self._config.discount,
        )

    def make_policy(
        self,
        networks: td3_networks.TD3Networks,
        environment_spec: specs.EnvironmentSpec,
        evaluation: bool = False,
    ) -> actor_core_lib.FeedForwardPolicy:
        """Creates a policy."""
        sigma = 0 if evaluation else self._config.sigma
        return td3_networks.get_default_behavior_policy(
            networks=networks, action_specs=environment_spec.actions, sigma=sigma
        )

corax/agents/jax/td3/builder.py
ethanluoyc/corax
# 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.

"""TD3 config."""
import dataclasses
from typing import Optional, Union

import optax

from corax.adders import reverb as adders_reverb


@dataclasses.dataclass
class TD3Config:
    """Configuration options for TD3."""

    # Loss options
    batch_size: int = 256
    policy_learning_rate: Union[optax.Schedule, float] = 3e-4
    critic_learning_rate: Union[optax.Schedule, float] = 3e-4
    # Policy gradient clipping is not part of the original TD3 implementation,
    # used e.g. in DAC https://arxiv.org/pdf/1809.02925.pdf
    policy_gradient_clipping: Optional[float] = None
    discount: float = 0.99
    n_step: int = 1

    # TD3 specific options (https://arxiv.org/pdf/1802.09477.pdf)
    sigma: float = 0.1
    delay: int = 2
    target_sigma: float = 0.2
    noise_clip: float = 0.5
    tau: float = 0.005

    # Replay options
    min_replay_size: int = 1000
    max_replay_size: int = 1000000
    replay_table_name: str = adders_reverb.DEFAULT_PRIORITY_TABLE
    prefetch_size: int = 4
    samples_per_insert: float = 256
    # Rate to be used for the SampleToInsertRatio rate limiter tolerance.
    # See a formula in make_replay_tables for more details.
    samples_per_insert_tolerance_rate: float = 0.1

    # Offline RL options
    # if bc_alpha: if given, will add a bc regularization term to the policy loss,
    # (https://arxiv.org/pdf/2106.06860.pdf), useful for offline training.
    bc_alpha: Optional[float] = None

corax/agents/jax/td3/config.py
ethanluoyc/corax
"""TD3 learner implementation."""

from typing import Dict, Iterator, List, NamedTuple, Optional, Tuple

import jax
import jax.numpy as jnp
import optax
import rlax

from corax import types
from corax.agents.jax import learner_core as learner_core_lib
from corax.agents.jax import learners
from corax.agents.jax.td3 import networks as td3_networks
from corax.jax import networks as networks_lib
from corax.jax import types as jax_types
from corax.utils import counting
from corax.utils import loggers


class TrainingState(NamedTuple):
    """Contains training state for the learner."""

    policy_params: networks_lib.Params
    target_policy_params: networks_lib.Params
    critic_params: networks_lib.Params
    target_critic_params: networks_lib.Params
    twin_critic_params: networks_lib.Params
    target_twin_critic_params: networks_lib.Params
    policy_opt_state: optax.OptState
    critic_opt_state: optax.OptState
    twin_critic_opt_state: optax.OptState
    steps: int
    random_key: networks_lib.PRNGKey


class TD3LearnerCore(learner_core_lib.LearnerCore):
    """TD3 learner."""

    def __init__(
        self,
        networks: td3_networks.TD3Networks,
        discount: float,
        policy_optimizer: optax.GradientTransformation,
        critic_optimizer: optax.GradientTransformation,
        twin_critic_optimizer: optax.GradientTransformation,
        delay: int = 2,
        target_sigma: float = 0.2,
        noise_clip: float = 0.5,
        tau: float = 0.005,
        use_sarsa_target: bool = False,
        bc_alpha: Optional[float] = None,
    ):
        """Initializes the TD3 learner.

        Args:
          networks: TD3 networks.
          discount: discount to use for TD updates
          policy_optimizer: the policy optimizer.
          critic_optimizer: the Q-function optimizer.
          twin_critic_optimizer: the twin Q-function optimizer.
          delay: ratio of policy updates for critic updates (see TD3),
            delay=2 means 2 updates of the critic for 1 policy update.
          target_sigma: std of zero mean Gaussian added to the action of
            the next_state, for critic evaluation (reducing overestimation bias).
          noise_clip: hard constraint on target noise.
          tau: target parameters smoothing coefficient.
          use_sarsa_target: compute on-policy target using iterator's actions rather
            than sampled actions.
            Useful for 1-step offline RL (https://arxiv.org/pdf/2106.08909.pdf).
            When set to `True`, `target_policy_params` are unused.
            This is only working when the learner is used as an offline algorithm.
            I.e. TD3Builder does not support adding the SARSA target to the replay
            buffer.
          bc_alpha: bc_alpha: Implements TD3+BC.
            See comments in TD3Config.bc_alpha for details.
        """

        def policy_loss(
            policy_params: networks_lib.Params,
            critic_params: networks_lib.Params,
            transition: types.NestedArray,
        ) -> jnp.ndarray:
            # Computes the discrete policy gradient loss.
            action = networks.policy_network.apply(
                policy_params, transition.observation
            )
            grad_critic = jax.vmap(
                jax.grad(networks.critic_network.apply, argnums=2), in_axes=(None, 0, 0)
            )
            dq_da = grad_critic(critic_params, transition.observation, action)
            batch_dpg_learning = jax.vmap(rlax.dpg_loss, in_axes=(0, 0))
            loss = jnp.mean(jnp.sum(batch_dpg_learning(action, dq_da), axis=-1))
            if bc_alpha is not None:
                # BC regularization for offline RL
                q_sa = networks.critic_network.apply(
                    critic_params, transition.observation, action
                )
                bc_factor = jax.lax.stop_gradient(bc_alpha / jnp.mean(jnp.abs(q_sa)))
                loss += jnp.mean(jnp.square(action - transition.action)) / bc_factor
            return loss

        def critic_loss(
            critic_params: networks_lib.Params,
            state: TrainingState,
            transition: types.Transition,
            random_key: jnp.ndarray,
        ):
            # Computes the critic loss.
            q_tm1 = networks.critic_network.apply(
                critic_params, transition.observation, transition.action
            )

            if use_sarsa_target:
                # TODO(b/222674779): use N-steps Trajectories to get the next actions.
                assert (
                    "next_action" in transition.extras
                ), "next actions should be given as extras for one step RL."
                action = transition.extras["next_action"]
            else:
                action = networks.policy_network.apply(
                    state.target_policy_params, transition.next_observation
                )
                action = networks.add_policy_noise(
                    action, random_key, target_sigma, noise_clip
                )

            q_t = networks.critic_network.apply(
                state.target_critic_params, transition.next_observation, action
            )
            twin_q_t = networks.twin_critic_network.apply(
                state.target_twin_critic_params, transition.next_observation, action
            )

            q_t = jnp.minimum(q_t, twin_q_t)

            target_q_tm1 = transition.reward + discount * transition.discount * q_t
            td_error = jax.lax.stop_gradient(target_q_tm1) - q_tm1

            return jnp.mean(jnp.square(td_error))

        def update_step(
            state: TrainingState,
            transitions: types.Transition,
        ) -> Tuple[TrainingState, Dict[str, jnp.ndarray]]:
            random_key, key_critic, key_twin = jax.random.split(state.random_key, 3)

            # Updates on the critic: compute the gradients, and update using
            # Polyak averaging.
            critic_loss_and_grad = jax.value_and_grad(critic_loss)
            critic_loss_value, critic_gradients = critic_loss_and_grad(
                state.critic_params, state, transitions, key_critic
            )
            critic_updates, critic_opt_state = critic_optimizer.update(
                critic_gradients, state.critic_opt_state
            )
            critic_params = optax.apply_updates(state.critic_params, critic_updates)
            # In the original authors' implementation the critic target update is
            # delayed similarly to the policy update which we found empirically to
            # perform slightly worse.
            target_critic_params = optax.incremental_update(
                new_tensors=critic_params,
                old_tensors=state.target_critic_params,
                step_size=tau,
            )

            # Updates on the twin critic: compute the gradients, and update using
            # Polyak averaging.
            twin_critic_loss_value, twin_critic_gradients = critic_loss_and_grad(
                state.twin_critic_params, state, transitions, key_twin
            )
            twin_critic_updates, twin_critic_opt_state = twin_critic_optimizer.update(
                twin_critic_gradients, state.twin_critic_opt_state
            )
            twin_critic_params = optax.apply_updates(
                state.twin_critic_params, twin_critic_updates
            )
            # In the original authors' implementation the twin critic target update is
            # delayed similarly to the policy update which we found empirically to
            # perform slightly worse.
            target_twin_critic_params = optax.incremental_update(
                new_tensors=twin_critic_params,
                old_tensors=state.target_twin_critic_params,
                step_size=tau,
            )

            # Updates on the policy: compute the gradients, and update using
            # Polyak averaging (if delay enabled, the update might not be applied).
            policy_loss_and_grad = jax.value_and_grad(policy_loss)
            policy_loss_value, policy_gradients = policy_loss_and_grad(
                state.policy_params, state.critic_params, transitions
            )

            def update_policy_step():
                policy_updates, policy_opt_state = policy_optimizer.update(
                    policy_gradients, state.policy_opt_state
                )
                policy_params = optax.apply_updates(state.policy_params, policy_updates)
                target_policy_params = optax.incremental_update(
                    new_tensors=policy_params,
                    old_tensors=state.target_policy_params,
                    step_size=tau,
                )
                return policy_params, target_policy_params, policy_opt_state

            # The update on the policy is applied every `delay` steps.
            current_policy_state = (
                state.policy_params,
                state.target_policy_params,
                state.policy_opt_state,
            )
            policy_params, target_policy_params, policy_opt_state = jax.lax.cond(
                state.steps % delay == 0,
                lambda _: update_policy_step(),
                lambda _: current_policy_state,
                operand=None,
            )

            steps = state.steps + 1

            new_state = TrainingState(
                policy_params=policy_params,
                critic_params=critic_params,
                twin_critic_params=twin_critic_params,
                target_policy_params=target_policy_params,
                target_critic_params=target_critic_params,
                target_twin_critic_params=target_twin_critic_params,
                policy_opt_state=policy_opt_state,
                critic_opt_state=critic_opt_state,
                twin_critic_opt_state=twin_critic_opt_state,
                steps=steps,
                random_key=random_key,
            )

            metrics = {
                "policy_loss": policy_loss_value,
                "critic_loss": critic_loss_value,
                "twin_critic_loss": twin_critic_loss_value,
            }

            return new_state, metrics

        self._update_step = update_step

        # Create initial state.
        def make_initial_state(random_key):
            (
                key_init_policy,
                key_init_twin,
                key_init_target,
                key_state,
            ) = jax.random.split(random_key, 4)
            # Create the network parameters and copy into the target network parameters.
            initial_policy_params = networks.policy_network.init(key_init_policy)
            initial_critic_params = networks.critic_network.init(key_init_twin)
            initial_twin_critic_params = networks.twin_critic_network.init(
                key_init_target
            )

            initial_target_policy_params = initial_policy_params
            initial_target_critic_params = initial_critic_params
            initial_target_twin_critic_params = initial_twin_critic_params

            # Initialize optimizers.
            initial_policy_opt_state = policy_optimizer.init(initial_policy_params)
            initial_critic_opt_state = critic_optimizer.init(initial_critic_params)
            initial_twin_critic_opt_state = twin_critic_optimizer.init(
                initial_twin_critic_params
            )

            return TrainingState(
                policy_params=initial_policy_params,
                target_policy_params=initial_target_policy_params,
                critic_params=initial_critic_params,
                twin_critic_params=initial_twin_critic_params,
                target_critic_params=initial_target_critic_params,
                target_twin_critic_params=initial_target_twin_critic_params,
                policy_opt_state=initial_policy_opt_state,
                critic_opt_state=initial_critic_opt_state,
                twin_critic_opt_state=initial_twin_critic_opt_state,
                steps=0,
                random_key=key_state,
            )

        self._init_state = make_initial_state

    def init(self, random_key: networks_lib.PRNGKey):
        return self._init_state(random_key)

    def step(self, state: TrainingState, transitions: types.Transition):
        # Get data from replay (dropping extras if any). Note there is no
        # extra data here because we do not insert any into Reverb.

        return self._update_step(state, transitions)

    def get_variables(
        self, state: TrainingState, names: List[str]
    ) -> List[networks_lib.Params]:
        variables = {
            "policy": state.policy_params,
            "critic": state.critic_params,
            "twin_critic": state.twin_critic_params,
        }
        return [variables[name] for name in names]


# TODO(yl): Consider removing this
class TD3Learner(learners.DefaultJaxLearner):
    """TD3 learner."""

    def __init__(
        self,
        networks: td3_networks.TD3Networks,
        random_key: jax_types.PRNGKey,
        iterator: Iterator[types.Transition],
        discount: float,
        policy_optimizer: optax.GradientTransformation,
        critic_optimizer: optax.GradientTransformation,
        twin_critic_optimizer: optax.GradientTransformation,
        delay: int = 2,
        target_sigma: float = 0.2,
        noise_clip: float = 0.5,
        tau: float = 0.005,
        use_sarsa_target: bool = False,
        bc_alpha: Optional[float] = None,
        logger: Optional[loggers.Logger] = None,
        counter: Optional[counting.Counter] = None,
    ):
        core = TD3LearnerCore(
            networks=networks,
            discount=discount,
            target_sigma=target_sigma,
            noise_clip=noise_clip,
            policy_optimizer=policy_optimizer,
            critic_optimizer=critic_optimizer,
            twin_critic_optimizer=twin_critic_optimizer,
            delay=delay,
            tau=tau,
            use_sarsa_target=use_sarsa_target,
            bc_alpha=bc_alpha,
        )
        super().__init__(
            core, random_key, iterator=iterator, logger=logger, counter=counter
        )

corax/agents/jax/td3/learning.py
ethanluoyc/corax
# type: ignore
# 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.

"""TD3 networks definition."""
import dataclasses
from typing import Callable, Sequence

import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np

from corax import specs
from corax import types
from corax.agents.jax import actor_core as actor_core_lib
from corax.jax import networks as networks_lib
from corax.jax import utils


@dataclasses.dataclass
class TD3Networks:
    """Network and pure functions for the TD3 agent."""

    policy_network: networks_lib.FeedForwardNetwork
    critic_network: networks_lib.FeedForwardNetwork
    twin_critic_network: networks_lib.FeedForwardNetwork
    add_policy_noise: Callable[
        [types.NestedArray, networks_lib.PRNGKey, float, float], types.NestedArray
    ]


def get_default_behavior_policy(
    networks: TD3Networks, action_specs: specs.BoundedArray, sigma: float
) -> actor_core_lib.FeedForwardPolicy:
    """Selects action according to the policy."""

    def behavior_policy(
        params: networks_lib.Params,
        key: networks_lib.PRNGKey,
        observation: types.NestedArray,
    ):
        action = networks.policy_network.apply(params, observation)
        noise = jax.random.normal(key, shape=action.shape) * sigma
        noisy_action = jnp.clip(
            action + noise, action_specs.minimum, action_specs.maximum
        )
        return noisy_action

    return behavior_policy


def make_networks(
    spec: specs.EnvironmentSpec, hidden_layer_sizes: Sequence[int] = (256, 256)
) -> TD3Networks:
    """Creates networks used by the agent.

    The networks used are based on LayerNormMLP, which is different than the
    MLP with relu activation described in TD3 (which empirically performs worse).

    Args:
      spec: Environment specs
      hidden_layer_sizes: list of sizes of hidden layers in actor/critic networks

    Returns:
      network: TD3Networks
    """

    action_specs = spec.actions
    num_dimensions = np.prod(action_specs.shape, dtype=int)

    def add_policy_noise(
        action: types.NestedArray,
        key: networks_lib.PRNGKey,
        target_sigma: float,
        noise_clip: float,
    ) -> types.NestedArray:
        """Adds action noise to bootstrapped Q-value estimate in critic loss."""
        noise = jax.random.normal(key=key, shape=action_specs.shape) * target_sigma
        noise = jnp.clip(noise, -noise_clip, noise_clip)
        return jnp.clip(action + noise, action_specs.minimum, action_specs.maximum)

    def _actor_fn(obs: types.NestedArray) -> types.NestedArray:
        network = hk.Sequential(
            [
                networks_lib.LayerNormMLP(hidden_layer_sizes, activate_final=True),
                networks_lib.NearZeroInitializedLinear(num_dimensions),
                networks_lib.TanhToSpec(spec.actions),
            ]
        )
        return network(obs)

    def _critic_fn(
        obs: types.NestedArray, action: types.NestedArray
    ) -> types.NestedArray:
        network1 = hk.Sequential(
            [
                networks_lib.LayerNormMLP(list(hidden_layer_sizes) + [1]),
            ]
        )
        input_ = jnp.concatenate([obs, action], axis=-1)
        value = network1(input_)
        return jnp.squeeze(value)

    policy = hk.without_apply_rng(hk.transform(_actor_fn))
    critic = hk.without_apply_rng(hk.transform(_critic_fn))

    # Create dummy observations and actions to create network parameters.
    dummy_action = utils.zeros_like(spec.actions)
    dummy_obs = utils.zeros_like(spec.observations)
    dummy_action = utils.add_batch_dim(dummy_action)
    dummy_obs = utils.add_batch_dim(dummy_obs)

    network = TD3Networks(
        policy_network=networks_lib.FeedForwardNetwork(
            lambda key: policy.init(key, dummy_obs), policy.apply
        ),
        critic_network=networks_lib.FeedForwardNetwork(
            lambda key: critic.init(key, dummy_obs, dummy_action), critic.apply
        ),
        twin_critic_network=networks_lib.FeedForwardNetwork(
            lambda key: critic.init(key, dummy_obs, dummy_action), critic.apply
        ),
        add_policy_noise=add_policy_noise,
    )

    return network

corax/agents/jax/td3/networks.py
ethanluoyc/corax
"""TDMPC agent module."""
from corax.agents.jax.tdmpc.acting import TDMPCActor
from corax.agents.jax.tdmpc.builder import TDMPCBuilder
from corax.agents.jax.tdmpc.config import TDMPCConfig
from corax.agents.jax.tdmpc.learning import TDMPCLearner
from corax.agents.jax.tdmpc.networks import TDMPCNetworks
from corax.agents.jax.tdmpc.networks import TDMPCParams
from corax.agents.jax.tdmpc.networks import make_networks

corax/agents/jax/tdmpc/__init__.py
ethanluoyc/corax
# type: ignore
from typing import Optional

import dm_env
import jax
import jax.numpy as jnp
import numpy as np
import optax
import rlax

from corax import adders
from corax import core
from corax import specs
from corax.agents.jax.tdmpc import networks as tdmpc_networks
from corax.agents.jax.tdmpc import planning
from corax.jax import networks as networks_lib
from corax.jax import variable_utils


class TDMPCActor(core.Actor):
    def __init__(
        self,
        variable_client: variable_utils.VariableClient,
        spec: specs.EnvironmentSpec,
        networks: tdmpc_networks.TDMPCNetworks,
        random_key: jax.random.PRNGKeyArray,
        *,
        std_schedule: optax.Schedule,
        horizon_schedule: optax.Schedule,
        discount: float = 0.99,
        num_samples: int = 512,
        min_std: float = 0.05,
        temperature: float = 0.5,
        momentum: float = 0.1,
        num_elites: int = 64,
        iterations: int = 6,
        seed_steps: int = 5000,
        mixture_coef: float = 0.05,
        horizon: int = 5,
        adder: Optional[adders.Adder] = None,
        value_tx_pair: rlax.TxPair = rlax.IDENTITY_PAIR,
        evaluation: bool = False,
    ) -> None:
        action_dim = spec.actions.shape[0]
        num_pi_trajs = int(mixture_coef * num_samples)

        def random_policy(params, obs, state):
            del params, obs
            (key, prev_mean, epsilon) = state
            sample_key, key = jax.random.split(key)
            action = jax.random.uniform(
                key=sample_key,
                minval=-1.0,
                maxval=1.0,
                dtype=jnp.float32,
                shape=spec.actions.shape,
            )
            return action, (key, prev_mean, epsilon)

        self._random_policy = jax.jit(random_policy)

        def plan_init(key: jax.random.PRNGKeyArray, step: int):
            epsilon = std_schedule(step) if not evaluation else min_std
            planning_horizon = int(min(horizon, horizon_schedule(step)))
            previous_trajectory = planning.get_initial_trajectory(
                action_dim,
                planning_horizon,
            )
            return (key, previous_trajectory, epsilon)

        def plan_step(params, obs, state):
            encoder_fn = networks.h

            def policy_fn(params, z, key):
                return networks.pi(params, z, min_std, key)

            def critic_fn(params, z, action):
                value = jnp.squeeze(
                    jnp.minimum(*networks.q(params, z, action)), axis=-1
                )
                return value_tx_pair.apply_inv(value)

            def model_fn(params, z, action):
                next_z, reward = networks.next(params, z, action)
                reward = value_tx_pair.apply_inv(reward)
                return next_z, jnp.squeeze(reward, axis=-1)

            (key, previous_trajectory, epsilon) = state

            step_key, sample_key, key = jax.random.split(key, 3)

            trajectory_mean, trajectory_std, _, actions = planning.td_mpc_planner(
                policy_fn=policy_fn,
                encoder_fn=encoder_fn,
                model_fn=model_fn,
                critic_fn=critic_fn,
                params=params,
                observation=obs,
                previous_trajectory=previous_trajectory,
                key=step_key,
                n_policy_trajectories=num_pi_trajs,
                n_sample_trajectories=num_samples,
                n_iterations=iterations,
                epsilon=epsilon,
                k=num_elites,
                discount=discount,
                temperature=temperature,
                momentum=momentum,
            )

            previous_trajectory = trajectory_mean
            mean, std = actions[0], trajectory_std[0]
            action = mean
            if not evaluation:
                action = action + std * jax.random.normal(sample_key, (action_dim,))
            new_state = (key, previous_trajectory, epsilon)
            return action, new_state

        self._plan_init = plan_init
        self._plan_step = jax.jit(plan_step)

        # TODO(yl): What to do with exploration in distribution experiments?
        self._seed_steps = seed_steps
        self._actor_steps = 0
        self._key = random_key

        self._state = None
        self._adder = adder
        self._variable_client = variable_client
        self._evaluation = evaluation

    @property
    def _params(self):
        return self._variable_client.params

    def observe_first(self, timestep: dm_env.TimeStep) -> None:
        init_key, self._key = jax.random.split(self._key)
        self._state = self._plan_init(init_key, self._actor_steps)

        if self._adder is not None:
            self._adder.add_first(timestep)

    def observe(
        self, action: networks_lib.Action, next_timestep: dm_env.TimeStep
    ) -> None:
        if self._adder is not None:
            self._adder.add(action, next_timestep)

    def select_action(
        self, observation: networks_lib.Observation
    ) -> networks_lib.Action:
        if self._actor_steps < self._seed_steps and not self._evaluation:
            action, self._state = self._random_policy(
                self._params, observation, self._state
            )
        else:
            action, self._state = self._plan_step(
                self._params, observation, self._state
            )

        self._actor_steps += 1
        return np.asarray(action)

    def update(self, wait: bool = False) -> None:
        self._variable_client.update(wait)

corax/agents/jax/tdmpc/acting.py
ethanluoyc/corax
from typing import Iterator, List, Optional, Tuple

import jax
import reverb

from corax import adders
from corax import specs
from corax.adders import reverb as adders_reverb
from corax.agents.jax import builders
from corax.agents.jax.tdmpc import acting
from corax.agents.jax.tdmpc import config as tdmpc_config
from corax.agents.jax.tdmpc import learning
from corax.agents.jax.tdmpc import networks as tdmpc_networks
from corax.agents.jax.tdmpc import types as tdmpc_types
from corax.datasets import reverb as reverb_datasets
from corax.jax import networks as networks_lib
from corax.jax import utils as jax_utils
from corax.jax import variable_utils
from corax.utils import counting
from corax.utils import loggers

TDMPCNetworks = tdmpc_networks.TDMPCNetworks
TDMPCPolicy = Tuple[tdmpc_networks.TDMPCNetworks, bool]


class TDMPCBuilder(builders.ActorLearnerBuilder):
    def __init__(self, config: tdmpc_config.TDMPCConfig):
        self._config = config

    def make_replay_tables(
        self,
        environment_spec: specs.EnvironmentSpec,
        policy: Optional[TDMPCPolicy] = None,
    ) -> List[reverb.Table]:
        del policy
        min_replay_size = self._config.min_replay_size
        samples_per_insert_tolerance = (
            self._config.samples_per_insert_tolerance_rate
            * self._config.samples_per_insert
        )
        error_buffer = min_replay_size * samples_per_insert_tolerance
        limiter = reverb.rate_limiters.SampleToInsertRatio(
            min_size_to_sample=min_replay_size,
            samples_per_insert=self._config.samples_per_insert,
            error_buffer=error_buffer,
        )

        return [
            reverb.Table(
                name=adders_reverb.DEFAULT_PRIORITY_TABLE,
                sampler=reverb.selectors.Prioritized(
                    self._config.importance_sampling_exponent
                ),
                remover=reverb.selectors.Fifo(),
                rate_limiter=limiter,
                max_size=self._config.max_replay_size,
                # For horizon of H, we insert sequences of size H+1 for bootstrapping
                signature=adders_reverb.SequenceAdder.signature(
                    environment_spec,
                    sequence_length=self._config.horizon + 1,
                ),
            )
        ]

    def make_adder(
        self,
        replay_client: reverb.Client,
        environment_spec: Optional[specs.EnvironmentSpec] = None,
        policy: Optional[TDMPCPolicy] = None,
    ) -> Optional[adders.Adder]:
        del policy, environment_spec
        return adders_reverb.SequenceAdder(
            replay_client,
            sequence_length=self._config.horizon + 1,
            period=1,
            end_of_episode_behavior=adders_reverb.EndBehavior.WRITE,
            max_in_flight_items=1,
        )

    def make_actor(
        self,
        random_key: networks_lib.PRNGKey,
        policy: TDMPCPolicy,
        environment_spec: specs.EnvironmentSpec,
        variable_source: learning.TDMPCLearner,
        adder: Optional[adders.Adder] = None,
    ) -> acting.TDMPCActor:
        networks, evaluation = policy
        variable_client = variable_utils.VariableClient(
            variable_source,
            "policy",
            update_period=self._config.variable_update_period,
        )
        return acting.TDMPCActor(
            variable_client,
            environment_spec,
            networks,
            random_key,
            std_schedule=self._config.std_schedule,
            horizon_schedule=self._config.horizon_schedule,
            discount=self._config.discount,
            num_samples=self._config.num_trajectories,
            min_std=self._config.min_std,
            temperature=self._config.temperature,
            momentum=self._config.momentum,
            num_elites=self._config.num_elites,
            iterations=self._config.num_iterations,
            seed_steps=self._config.min_replay_size,
            mixture_coef=self._config.policy_trajectory_fraction,
            value_tx_pair=self._config.value_tx_pair,
            horizon=self._config.horizon,
            adder=adder,
            evaluation=evaluation,
        )

    def make_learner(
        self,
        random_key: networks_lib.PRNGKey,
        networks: TDMPCNetworks,
        dataset: Iterator[learning.TDMPCReplaySample],
        logger_fn: loggers.LoggerFactory,
        environment_spec: specs.EnvironmentSpec,
        replay_client: reverb.Client,
        counter: Optional[counting.Counter] = None,
    ) -> learning.TDMPCLearner:
        loss_scale = tdmpc_types.LossScalesConfig(
            consistency=self._config.consistency_loss_scale,
            reward=self._config.reward_loss_scale,
            value=self._config.value_loss_scale,
        )

        return learning.TDMPCLearner(
            spec=environment_spec,
            networks=networks,
            random_key=random_key,
            iterator=dataset,
            replay_client=replay_client,
            optimizer=self._config.optimizer,
            importance_sampling_exponent=self._config.priority_exponent,
            discount=self._config.discount,
            min_std=self._config.min_std,
            target_update_rate=self._config.critic_update_rate,
            value_tx_pair=self._config.value_tx_pair,
            loss_scale=loss_scale,
            rho=self._config.rho,
            counter=counter,
            logger=logger_fn("learner"),
        )

    def make_dataset_iterator(
        self, replay_client: reverb.Client
    ) -> Iterator[learning.TDMPCReplaySample]:
        dataset = reverb_datasets.make_reverb_dataset(
            server_address=replay_client.server_address,
            batch_size=self._config.batch_size,
            num_parallel_calls=4,
        )
        return jax_utils.device_put(
            dataset.as_numpy_iterator(),
            split_fn=jax_utils.keep_key_on_host,
            device=jax.devices()[0],
        )

    def make_policy(
        self,
        networks: TDMPCNetworks,
        environment_spec: specs.EnvironmentSpec,
        evaluation: bool = False,
    ) -> TDMPCPolicy:
        del environment_spec
        return networks, evaluation

corax/agents/jax/tdmpc/builder.py
ethanluoyc/corax
"""Configuration for TD-MPC agents"""
import dataclasses

import optax
import rlax


@dataclasses.dataclass
class TDMPCConfig:
    std_schedule: optax.Schedule
    horizon_schedule: optax.Schedule

    optimizer: optax.GradientTransformation
    discount: float = 0.99

    # Training configuration
    critic_update_rate: float = 0.01  # tau
    value_tx_pair: rlax.TxPair = rlax.IDENTITY_PAIR
    variable_update_period: int = 1

    # Loss scales for model learning
    consistency_loss_scale: float = 2.0
    reward_loss_scale: float = 0.5
    value_loss_scale: float = 0.1

    rho: float = 0.5  # model prediction loss discount

    # Planning options
    horizon: int = 5  # Planning horizon
    num_trajectories: int = 512  # Number of trajectories to sample per planning step
    # Fraction of trajectories to sample from policy
    policy_trajectory_fraction: float = 0.05
    num_elites: int = 64  # Number of elite trajectories for CEM
    num_iterations: int = 6  # Number of CEM iterations
    min_std: float = 0.05  # Minimum noise to add to sample trajectories

    # parameter used for computing the normalized empirical estimate of
    # the value of the trajectories.
    temperature: float = 0.5
    momentum: float = 0.1  # Momentum for updating trajectories distribution parameters

    # Replay configuration
    samples_per_insert: float = 512.0
    samples_per_insert_tolerance_rate: float = 0.1
    # min_replay_size currently has dual meaning:
    # 1. The number of samples to insert into Reverb prior to learning begins
    # 2. The number of random exploration steps
    min_replay_size: int = 5000
    max_replay_size: int = int(1e6)
    batch_size: int = 512

    # Prioritized Experience Replay configuration
    # See https://arxiv.org/pdf/1511.05952.pdf
    importance_sampling_exponent: float = 0.6  # alpha
    priority_exponent: float = 0.4  # beta

corax/agents/jax/tdmpc/config.py
ethanluoyc/corax
# type: ignore
import time
from typing import Iterator, List, NamedTuple, Optional, Tuple

import chex
import haiku as hk
import jax
import jax.numpy as jnp
import optax
import reverb
import rlax
import tree

import corax
from corax import specs
from corax import types
from corax.adders import reverb as adders_reverb
from corax.agents.jax.tdmpc import networks as tdmpc_networks
from corax.agents.jax.tdmpc import types as tdmpc_types
from corax.jax import networks as networks_lib
from corax.jax import utils as jax_utils
from corax.utils import async_utils
from corax.utils import counting
from corax.utils import loggers

TDMPCReplaySample = jax_utils.PrefetchingSplit


def _squared_error(predictions: chex.Array, targets: chex.Array) -> chex.Array:
    chex.assert_type([predictions], float)
    chex.assert_equal_shape((predictions, targets))
    errors = predictions - targets
    return jnp.square(errors)


def _l1_loss(predictions: chex.Array, targets: chex.Array) -> chex.Array:
    chex.assert_type([predictions], float)
    chex.assert_equal_shape((predictions, targets))
    errors = predictions - targets
    return jnp.abs(errors)


class TrainingState(NamedTuple):
    """Contains training state for the learner."""

    params: tdmpc_networks.TDMPCParams
    target_params: tdmpc_networks.TDMPCParams
    opt_state: optax.OptState
    steps: int
    key: jax.random.PRNGKeyArray


class ModelOutputs(NamedTuple):
    """Contains model outputs."""

    z_prediction: types.NestedArray
    z_prior: types.NestedArray
    reward_pred: types.NestedArray
    q_value: types.NestedArray


class LossTargets(NamedTuple):
    """Contains loss targets."""

    td_target: types.NestedArray
    z_target: types.NestedArray
    reward: types.NestedArray


class TDMPCLearner(corax.Learner):
    """TDMPC Learner."""

    _state: TrainingState

    def __init__(
        self,
        spec: specs.EnvironmentSpec,
        networks: tdmpc_networks.TDMPCNetworks,
        random_key: jax.random.PRNGKeyArray,
        replay_client: reverb.Client,
        iterator: Iterator[TDMPCReplaySample],
        *,
        optimizer: optax.GradientTransformation,
        discount: float = 0.99,
        min_std: float = 0.05,
        loss_scale: Optional[tdmpc_types.LossScalesConfig] = None,
        value_tx_pair: rlax.TxPair = rlax.IDENTITY_PAIR,
        importance_sampling_exponent: float = 0.4,
        target_update_rate: float = 0.01,
        rho: float = 0.5,
        logger: Optional[loggers.Logger] = None,
        counter: Optional[counting.Counter] = None,
    ):
        """Initializes the learner.

        Args:
            spec: A description of the actions, observations, etc.
            networks: networks used by the learner.
            random_key: A random key for initializing the learner.
            replay_client: A Reverb client for updating priorities.
            iterator: A dataset iterator.
            optimizer: The optimizer to use for training.
            discount: Discount to use for TD updates.
            min_std: Minimum standard deviation to use for the policy.
            loss_scale: Loss scales to use for the TDMPC losses.
            importance_sampling_exponent: Exponent to use for importance sampling.
            target_update_rate: Rate at which to update the target networks.
            rho: The rho parameter for the TDMPC loss.
            logger: Logger object for writing training statistics to.
            counter: Counter object for keeping track of iterations.
        """
        if loss_scale is None:
            loss_scale = tdmpc_types.LossScalesConfig()
        self._discount = discount
        self._min_std = min_std
        self._target_update_rate = target_update_rate
        self._loss_scale = loss_scale
        self._rho = rho
        self._networks = networks
        self._optimizer = optimizer
        self._importance_sampling_exponent = importance_sampling_exponent
        self._replay_client = replay_client
        self._tx_pair = value_tx_pair

        # Initialize parameters.
        param_key, key = jax.random.split(random_key)
        params = tdmpc_networks.init_params(self._networks, spec, param_key)

        # Initialize optimizer state.
        opt_state = self._optimizer.init(params)

        # Initialize learner state.
        self._state = TrainingState(
            params=params,
            target_params=params,
            opt_state=opt_state,
            key=key,
            steps=0,
        )

        self._update_step = jax.jit(self._sgd_step)

        self._counter = counter or counting.Counter()
        self._logger = logger

        def update_priorities(keys_and_priorities: Tuple[jnp.ndarray, jnp.ndarray]):
            keys, priorities = keys_and_priorities
            keys, priorities = tree.map_structure(
                # Fetch array and combine device and batch dimensions.
                lambda x: jax_utils.fetch_devicearray(x).reshape((-1,) + x.shape[2:]),
                (keys, priorities),
            )
            replay_client.mutate_priorities(
                table=adders_reverb.DEFAULT_PRIORITY_TABLE,
                updates=dict(zip(keys, priorities)),
            )

        self._async_priority_updater = async_utils.AsyncExecutor(update_priorities)

        self._iterator = iterator

        # 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

    def _compute_targets(
        self,
        target_params: tdmpc_networks.TDMPCParams,
        sequences: adders_reverb.Step,
        online_params: adders_reverb.Step,
        z_prior: types.NestedArray,
        key: networks_lib.PRNGKey,
    ):
        """Computes the targets for the TDMPC loss."""

        def policy(params, obs, key):
            return self._networks.pi(params, obs, self._min_std, key)

        batched_policy = jax.vmap(policy, in_axes=(None, 0, 0))
        batched_critic = jax.vmap(self._networks.q, (None, 0, 0))
        batched_encoder = jax.vmap(self._networks.h, (None, 0))

        z_target = batched_encoder(target_params, sequences.observation)

        # NOTE: Actions for policy evaluation is conditioned on
        # actions selected by the online policy based on prior encoding, See
        # https://github.com/nicklashansen/tdmpc/blob/main/src/algorithm/tdmpc.py#L172-L177
        policy_action = batched_policy(
            online_params, z_prior, jax.random.split(key, z_prior.shape[0])
        )
        policy_action = jax.lax.stop_gradient(policy_action)

        v1_target, v2_target = batched_critic(
            target_params, z_prior[1:], policy_action[1:]
        )

        # Transform to raw reward space.
        v1_target, v2_target = jax.tree_map(
            self._tx_pair.apply_inv, (v1_target, v2_target)
        )
        target_values = jnp.squeeze(jnp.minimum(v1_target, v2_target), axis=-1)

        discounts = self._discount * sequences.discount[:-1]
        td_target = sequences.reward[:-1] + discounts * target_values
        reward_target = sequences.reward[:-1]

        # Transform back to canonical space.
        td_target = jax.lax.stop_gradient(self._tx_pair.apply(td_target))
        reward_target = jax.lax.stop_gradient(self._tx_pair.apply(reward_target))

        return LossTargets(
            td_target=td_target,
            z_target=z_target,
            reward=reward_target,
        )

    def _compute_predictions(
        self,
        params: tdmpc_networks.TDMPCNetworks,
        sequences: adders_reverb.Step,
    ):
        """Computes the predictions from the model."""

        def unroll_dynamics(params, actions, z):
            def _next_core(action, z):
                next_z, reward = self._networks.next(params, z, action)
                return (next_z, reward), next_z

            (online_z_posterior, reward_pred), _ = hk.static_unroll(
                _next_core, actions, z
            )

            z_predictions = jnp.concatenate(
                [jnp.expand_dims(z, axis=0), online_z_posterior], axis=0
            )

            return z_predictions, jnp.squeeze(reward_pred, axis=-1)

        batched_critic = jax.vmap(self._networks.q, (None, 0, 0))
        batched_encoder = jax.vmap(self._networks.h, (None, 0))

        z_prior = batched_encoder(params, sequences.observation)
        # [H+1, B, Z], [H, B]
        z_predictions, reward_pred = unroll_dynamics(
            params, sequences.action[:-1], z_prior[0]
        )

        # ([H, B], [H, B])
        q_values = batched_critic(params, z_predictions[:-1], sequences.action[:-1])
        q_values = jax.tree_map(lambda x: jnp.squeeze(x, axis=-1), q_values)

        return ModelOutputs(
            z_prediction=z_predictions,
            z_prior=z_prior,
            reward_pred=reward_pred,
            q_value=q_values,
        )

    def _policy_loss(
        self,
        params: tdmpc_networks.TDMPCNetworks,
        z_predictions: types.NestedArray,
        key: networks_lib.PRNGKey,
    ):
        """Computes the policy loss."""

        def policy(params, obs, key):
            return self._networks.pi(params, obs, self._min_std, key)

        batched_policy = jax.vmap(policy, in_axes=(None, 0, 0))

        # stop gradient from policy_loss to other model components
        frozen_params = jax.lax.stop_gradient(params)
        z_policy = jax.lax.stop_gradient(z_predictions)

        actions = batched_policy(
            params, z_policy, jax.random.split(key, z_predictions.shape[0])
        )
        # Compute policy loss
        q1, q2 = jax.vmap(self._networks.q, (None, 0, 0))(
            frozen_params, z_policy, actions
        )

        q = jnp.squeeze(jnp.minimum(q1, q2), axis=-1)

        return -q

    def _loss_fn(
        self,
        params: tdmpc_networks.TDMPCParams,
        target_params: tdmpc_networks.TDMPCParams,
        batch: reverb.ReplaySample,
        key: jax.random.PRNGKeyArray,
    ):
        """Computes the loss for the model and policy."""
        sequences: adders_reverb.Step = jax_utils.batch_to_sequence(batch.data)
        # [H+1, B, ...]
        # For horizon of H, we insert sequences of size H+1 for bootstrapping
        horizon = sequences.reward.shape[0] - 1

        predictions = self._compute_predictions(params, sequences)
        q1_t, q2_t = predictions.q_value

        key, target_key = jax.random.split(key)
        targets = self._compute_targets(
            target_params, sequences, params, predictions.z_prior, target_key
        )

        # Compute model loss
        consistency_loss = jnp.mean(
            _squared_error(predictions.z_prediction[1:], targets.z_target[1:]), axis=-1
        )
        reward_loss = _squared_error(predictions.reward_pred, targets.reward)
        value_loss = _squared_error(q1_t, targets.td_target) + _squared_error(
            q2_t, targets.td_target
        )
        priorities = _l1_loss(q1_t, targets.td_target) + _l1_loss(
            q2_t, targets.td_target
        )

        chex.assert_equal_shape([reward_loss, value_loss, priorities, consistency_loss])

        rhos = jnp.expand_dims(jnp.power(self._rho, jnp.arange(horizon + 1)), axis=-1)

        # Discount the model losses
        consistency_loss = jnp.sum(rhos[:-1] * consistency_loss, axis=0)
        reward_loss = jnp.sum(rhos[:-1] * reward_loss, axis=0)
        value_loss = jnp.sum(rhos[:-1] * value_loss, axis=0)
        priorities = jnp.sum(rhos[:-1] * priorities, axis=0)

        model_loss = (
            self._loss_scale.consistency * jnp.clip(consistency_loss, 0, 1e4)
            + self._loss_scale.reward * jnp.clip(reward_loss, 0, 1e4)
            + self._loss_scale.value * jnp.clip(value_loss, 0, 1e4)
        )
        model_loss = optax.scale_gradient(model_loss, 1.0 / horizon)

        # Compute policy loss
        policy_loss = self._policy_loss(params, predictions.z_prediction, key)
        chex.assert_rank([rhos, policy_loss], 2)
        chex.assert_equal_rank([rhos, policy_loss])
        policy_loss = jnp.mean(jnp.sum(rhos * policy_loss, axis=0), axis=0)

        # Compute weights for the model loss
        # NOTE: In the original paper, weights are only used for the model loss.
        # The policy loss is not weighted.
        probabilities = batch.info.probability
        importance_sampling_weights = (1 / probabilities).astype(jnp.float32)
        importance_sampling_weights **= self._importance_sampling_exponent
        importance_sampling_weights /= jnp.max(importance_sampling_weights)

        weighted_model_loss = importance_sampling_weights * model_loss
        weighted_model_loss = jnp.mean(weighted_model_loss, axis=0)

        total_loss = weighted_model_loss + policy_loss

        metrics = {
            "total_loss": jnp.mean(total_loss),
            "policy_loss": jnp.mean(policy_loss),
            "model/model_loss": jnp.mean(model_loss),
            "model/weighted_model_loss": jnp.mean(weighted_model_loss),
            "model/consistentcy_loss": jnp.mean(consistency_loss),
            "model/reward_loss": jnp.mean(reward_loss),
            "model/critic_loss": jnp.mean(value_loss),
        }

        return total_loss, (priorities, metrics)

    def _sgd_step(self, state: TrainingState, batch: TDMPCReplaySample):
        key, random_key = jax.random.split(state.key)
        grad_fn = jax.value_and_grad(self._loss_fn, has_aux=True)
        (_, (priorities, metrics)), gradients = grad_fn(
            state.params, state.target_params, batch, key
        )

        metrics["grad_norm"] = optax.global_norm(gradients)

        updates, opt_state = self._optimizer.update(
            gradients, state.opt_state, state.params
        )

        params = optax.apply_updates(state.params, updates)
        target_params = optax.incremental_update(
            params, state.target_params, self._target_update_rate
        )

        steps = state.steps + 1
        new_state = TrainingState(
            params=params,
            target_params=target_params,
            opt_state=opt_state,
            key=random_key,
            steps=steps,
        )

        return (new_state, metrics, priorities)

    def step(self):
        prefetching_splits = next(self._iterator)

        keys = prefetching_splits.host
        samples = prefetching_splits.device

        (self._state, metrics, priorities) = self._update_step(self._state, samples)

        if self._replay_client:
            self._async_priority_updater.put((keys, priorities))

        # Compute elapsed time.
        timestamp = time.time()
        elapsed_time = timestamp - self._timestamp if self._timestamp else 0
        self._timestamp = timestamp

        counts = self._counter.increment(steps=1, walltime=elapsed_time)

        if self._logger is not None:
            self._logger.write({**counts, **metrics})

    def get_variables(self, names: List[str]) -> List[networks_lib.Params]:
        variables = {"policy": self._state.params}
        return [variables[name] for name in names]

    def save(self) -> TrainingState:
        return self._state

    def restore(self, state: TrainingState):
        self._state = state

corax/agents/jax/tdmpc/learning.py
ethanluoyc/corax
import jax
import jax.numpy as jnp
import optax
import reverb
from absl.testing import absltest

from corax import specs
from corax.adders import reverb as adders_reverb
from corax.agents.jax import tdmpc
from corax.agents.jax.tdmpc import learning
from corax.jax import utils
from corax.testing import fakes


class LossTest(absltest.TestCase):
    def test_loss_gradients(self):
        """Test that single pass loss computes gradients correctly."""
        environment = fakes.ContinuousEnvironment(bounded=True)

        environment_spec = specs.make_environment_spec(environment)

        networks = tdmpc.make_networks(
            environment_spec,
            mlp_hidden_size=10,
            latent_size=10,
            encoder_hidden_size=10,
            zero_init=False,
        )

        learner = learning.TDMPCLearner(
            environment_spec,
            networks,
            jax.random.PRNGKey(0),
            iterator=None,  # type: ignore
            replay_client=None,  # type: ignore
            optimizer=optax.adam(1e-3),
        )

        dummy_values = utils.zeros_like(environment_spec)
        horizon = 3

        sequences = adders_reverb.Step(
            observation=utils.add_batch_dim(
                utils.tile_nested(dummy_values.observations, horizon)
            ),
            reward=utils.add_batch_dim(
                utils.tile_nested(dummy_values.rewards, horizon)
            ),
            action=utils.add_batch_dim(
                utils.tile_nested(dummy_values.actions, horizon)
            ),
            discount=jnp.ones((1, horizon)),
            start_of_episode=jnp.ones((1, horizon)),
        )
        sequences = jax.tree_map(lambda x: x + 1, sequences)

        batch = reverb.ReplaySample(
            info=reverb.SampleInfo(
                probability=jnp.ones((1,)),
                key=None,
                table_size=1,
                priority=1,
                times_sampled=1.0,
            ),
            data=sequences,
        )

        params = learner._state.params

        def get_grad(name):
            def loss_fn(params):
                _, (_, logging_dict) = learner._loss_fn(
                    params, params, batch, jax.random.PRNGKey(0)
                )
                return logging_dict[name]

            return jax.grad(loss_fn)(params)

        policy_grads = get_grad("policy_loss")
        self.assertEqual(
            optax.global_norm((policy_grads._replace(policy_params={}))), 0
        )
        model_grads = get_grad("model/model_loss")
        self.assertEqual(optax.global_norm((model_grads.policy_params)), 0.0)


if __name__ == "__main__":
    absltest.main()

corax/agents/jax/tdmpc/loss_test.py
ethanluoyc/corax
# type: ignore
"""Networks used in TD-MPC."""
import dataclasses
from typing import NamedTuple

import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np

from corax import specs
from corax.jax import utils


class TDMPCParams(NamedTuple):
    encoder_params: hk.Params
    reward_params: hk.Params
    dynamics_params: hk.Params
    critic_params: hk.Params
    twin_critic_params: hk.Params
    policy_params: hk.Params


def _truncated_normal_sample(key, mean, std, noise_clip=0.3):
    noise = jax.random.normal(key=key, shape=mean.shape) * std
    noise = jnp.clip(noise, -noise_clip, noise_clip)
    sample = mean + noise
    clipped_sample = jnp.clip(sample, -1, 1)
    return (
        sample - jax.lax.stop_gradient(sample) + jax.lax.stop_gradient(clipped_sample)
    )


@dataclasses.dataclass
class TDMPCNetworks:
    encoder_network: hk.Transformed
    reward_network: hk.Transformed
    dynamics_network: hk.Transformed
    critic_network: hk.Transformed
    policy_network: hk.Transformed

    def h(self, params, obs):
        return self.encoder_network.apply(params.encoder_params, obs)

    def next(self, params, z, act):
        next_z = self.dynamics_network.apply(params.dynamics_params, z, act)
        reward = self.reward_network.apply(params.reward_params, z, act)
        return next_z, reward

    def pi(self, params, z, std, key):
        a = jnp.tanh(self.policy_network.apply(params.policy_params, z))
        return _truncated_normal_sample(key, a, std, noise_clip=0.3)

    def q(self, params, z, action):
        q1 = self.critic_network.apply(params.critic_params, z, action)
        q2 = self.critic_network.apply(params.twin_critic_params, z, action)
        return q1, q2


def init_params(
    networks: TDMPCNetworks,
    spec: specs.EnvironmentSpec,
    key: jax.random.PRNGKeyArray,
):
    keys = jax.random.split(key, 6)
    dummy_obs = utils.add_batch_dim(utils.zeros_like(spec.observations))
    dummy_act = utils.add_batch_dim(utils.zeros_like(spec.actions))

    encoder_params = networks.encoder_network.init(keys[0], dummy_obs)
    embedding = networks.encoder_network.apply(encoder_params, dummy_obs)

    critic_params = networks.critic_network.init(keys[1], embedding, dummy_act)
    twin_critic_params = networks.critic_network.init(keys[2], embedding, dummy_act)
    dynamics_params = networks.dynamics_network.init(keys[3], embedding, dummy_act)
    reward_params = networks.reward_network.init(keys[4], embedding, dummy_act)
    policy_params = networks.policy_network.init(keys[5], embedding)

    return TDMPCParams(
        encoder_params=encoder_params,
        critic_params=critic_params,
        dynamics_params=dynamics_params,
        twin_critic_params=twin_critic_params,
        reward_params=reward_params,
        policy_params=policy_params,
    )


def make_networks(
    spec: specs.EnvironmentSpec,
    latent_size,
    encoder_hidden_size,
    mlp_hidden_size,
    zero_init: bool = True,
):
    num_dimensions = np.prod(spec.actions.shape, dtype=int)

    w_init = hk.initializers.Orthogonal(1.0)

    def _encoder_fn(obs):
        obs = utils.batch_concat(obs)
        return hk.nets.MLP(
            [encoder_hidden_size, latent_size],
            activation=jax.nn.elu,
            w_init=w_init,
            activate_final=False,
        )(obs)

    def _dynamics_fn(obs_embedding, act):
        inputs = jnp.concatenate([obs_embedding, act], axis=-1)
        return hk.Sequential(
            [
                hk.nets.MLP(
                    [mlp_hidden_size, mlp_hidden_size, obs_embedding.shape[-1]],
                    activation=jax.nn.elu,
                    w_init=w_init,
                    activate_final=False,
                ),
            ]
        )(inputs)

    def _actor_fn(obs_embedding):
        return hk.Sequential(
            [
                hk.nets.MLP(
                    [mlp_hidden_size, mlp_hidden_size, num_dimensions],
                    activation=jax.nn.elu,
                    w_init=w_init,
                    activate_final=False,
                ),
            ]
        )(obs_embedding)

    def _critic_fn(obs_embedding, act):
        inputs = jnp.concatenate([obs_embedding, act], axis=-1)
        return hk.Sequential(
            [
                hk.Linear(mlp_hidden_size, w_init=w_init),
                hk.LayerNorm(axis=-1, create_scale=True, create_offset=True),
                jax.lax.tanh,
                hk.Linear(mlp_hidden_size, w_init=w_init),
                jax.nn.elu,
                hk.Linear(
                    1,
                    w_init=w_init if not zero_init else jnp.zeros,
                ),
            ]
        )(inputs)

    def _reward_fn(obs_embedding, act):
        inputs = jnp.concatenate([obs_embedding, act], axis=-1)
        return hk.Sequential(
            [
                hk.nets.MLP(
                    [mlp_hidden_size, mlp_hidden_size],
                    activation=jax.nn.elu,
                    w_init=w_init,
                    activate_final=True,
                ),
                hk.Linear(1, w_init=w_init if not zero_init else jnp.zeros),
            ]
        )(inputs)

    encoder_network = hk.without_apply_rng(hk.transform(_encoder_fn))
    reward_network = hk.without_apply_rng(hk.transform(_reward_fn))
    dynamics_network = hk.without_apply_rng(hk.transform(_dynamics_fn))
    critic_network = hk.without_apply_rng(hk.transform(_critic_fn))
    policy_network = hk.without_apply_rng(hk.transform(_actor_fn))
    return TDMPCNetworks(
        encoder_network=encoder_network,
        reward_network=reward_network,
        critic_network=critic_network,
        dynamics_network=dynamics_network,
        policy_network=policy_network,
    )

corax/agents/jax/tdmpc/networks.py
ethanluoyc/corax
"""TD-MPC inference algorithm."""
from typing import Any, Callable, Tuple

import chex
import jax
import jax.numpy as jnp

Params = Any
Observation = jax.Array
Embedding = jax.Array
Action = jax.Array
Reward = jax.Array
Trajectory = jax.Array

EncoderFn = Callable[[Params, Observation], Embedding]
ModelFn = Callable[[Params, Embedding, Action], Tuple[Embedding, Reward]]
CriticFn = Callable[[Params, Embedding, Action], jax.Array]
PolicyFn = Callable[[Params, Embedding, jax.random.PRNGKeyArray], Action]


def get_initial_trajectory(action_dims: int, horizon: int) -> Trajectory:
    """Create empty planning trajectory."""
    return jnp.zeros((horizon, action_dims))


def td_mpc_planner(
    policy_fn: PolicyFn,
    encoder_fn: EncoderFn,
    model_fn: ModelFn,
    critic_fn: CriticFn,
    params: Params,
    observation: Observation,
    previous_trajectory: Trajectory,
    key: jax.random.PRNGKeyArray,
    *,
    n_policy_trajectories: int,
    n_sample_trajectories: int,
    n_iterations: int,
    k: int,
    discount: float,
    epsilon: float,
    temperature: float,
    momentum: float,
) -> Tuple[Trajectory, Trajectory, jax.Array, Action]:
    """Run the TD-MPC planner.

    This will perform the planning step as described in Alg. 1 of the TD-MPC paper:
    [https://arxiv.org/abs/2203.04955].

    Args:
        policy_fn: policy for selection actions.
        encoder_fn: encoder for encoding observations.
        model_fn: latent dynamics model for computing next latent state and reward.
        critic_fn: Q-function for computing bootstrap return estimate. Expects
            the output to be 1D array of the value estimate. Therefore, twin critic
            should be transformed to adapt to using this function.
        params: network parameters for the learned models. Will be forwarded to
            the policy, encoder, model and critic.
        observation: current observation of shape [O].
        previous_trajectory: array of shape [H, A]. Previous trajectory computed by
            the planner.
        key: JAX random key.
        n_policy_trajectories: number of policy trajectories (N_pi).
        n_sample_trajectories: number of sample trajectories (N).
        n_iterations: number of iterations to run the TD-MPC planner.
        k: top-K number of trajectories used for updating the parameters.
        discount: discount used for computing the value estimate.
        epsilon: minimum stddev for updating the std parameters to prevent pre-mature
            convergence.
        temperature: temperature parameter used for computing the normalized empirical
            estimate of the value of the trajectories.
        momentum: momentum parameter used for updating the mean parameter.

    Returns:
        mean: [H, A] array. Final mean parameter for adapted trajectory distribution.
        std: [H, A] array. Final std parameter for adapted trajectory distribution.
        scores: [k] array. Final normalized empirical value estimate.
        actions: [A] array. The actions selected by the planner.

    """
    chex.assert_rank(previous_trajectory, 2)
    horizon = previous_trajectory.shape[0]
    action_dims = previous_trajectory.shape[1]

    # Encode state z_t = h(s_t)
    observation = jax.tree_map(
        lambda x: jnp.expand_dims(x, axis=0), observation
    )  # [1, O]
    z = encoder_fn(params, observation)  # [1, Z]

    policy_actions = None
    if n_policy_trajectories > 0:
        rollout_key, key = jax.random.split(key)
        policy_actions = _rollout_policy(
            params, model_fn, policy_fn, z, rollout_key, horizon, n_policy_trajectories
        )

    # Repeat for sample trajectories
    z = jnp.tile(z, (n_sample_trajectories + n_policy_trajectories, 1))  # [N + N_pi, Z]

    # Reuse the trajectory optimization using previous shifted mean
    initial_mean = jnp.roll(previous_trajectory, shift=-1)  # [H, A]
    initial_mean = initial_mean.at[-1].set(0)

    # Always use large initial variance to avoid local minima
    initial_std = 2 * jnp.ones((horizon, action_dims))  # [H, A]
    initial_state = (
        initial_mean,
        initial_std,
        # Score
        jnp.zeros((k,)),
        # Elite actions
        jnp.zeros((k, horizon, action_dims)),
        # Random key
        key,
    )

    def loop_body_fn(_, state):
        """A single iteration of the TD-MPC planning loop"""
        mean, std, _, _, key = state
        action_key, value_key, key = jax.random.split(key, 3)

        # Sample N trajectories of length H from N(mu^{j-1}, (sigma^{j-1})^2 I)
        action_noise = jax.random.normal(
            action_key, (n_sample_trajectories, horizon, action_dims)
        )  # [N, H, A]
        actions = mean + std * action_noise  # [N, H, A]
        # clip actions to be in the canonical space.
        actions = jnp.clip(actions, -1.0, 1.0)

        # Sample N_pi trajectories of length H using learned model.
        # Reuse samples sampled at the beginning of the iteration.
        if policy_actions is not None:
            actions = jnp.concatenate(
                [actions, policy_actions], axis=0
            )  # [N + N_pi, H, A]

        chex.assert_shape(
            [actions],
            (n_sample_trajectories + n_policy_trajectories, horizon, action_dims),
        )

        # Compute the top-k action trajectories.
        # Estimate trajectory returns (\phi_Gamma) using (d_theta, R_theta, Q_theta)
        # [N + N_pi]
        values = _compute_n_step_return(
            model_fn,
            critic_fn,
            policy_fn,
            params,
            z,
            jnp.swapaxes(actions, 0, 1),
            value_key,
            discount=discount,
        )

        # Avoid NaNs
        # TODO(yl): Replace 0 with something else in case we use negative rewards.
        values = jnp.nan_to_num(values, 0)
        chex.assert_shape([values], (n_sample_trajectories + n_policy_trajectories,))

        ## Update parameters mu, signa for next iteration (Eqn. 5)
        # Select top-k returns trajectories
        # top_k_idxs [K]
        top_k_indices = jax.lax.top_k(values, k)[1]  # [K]

        # Select \phi^* and the corresponding action
        top_k_values = values[top_k_indices]  # [K]
        top_k_actions = actions[top_k_indices]  # [K, H, A]
        chex.assert_shape([top_k_values], (k,))
        chex.assert_shape([top_k_actions], (k, horizon, action_dims))

        ## Update parameters according to (Eqn. 4)
        # Compute the normalized empirical estimate \Omega_i
        max_value = jnp.max(top_k_values, axis=0)
        scores = jnp.exp(temperature * (top_k_values - max_value))  # [K]
        chex.assert_shape([scores], (k,))

        # Omega_i / sum_i(Omega_i)
        scores = scores / jnp.sum(scores, axis=0)  # [K]
        chex.assert_shape([scores], (k,))
        # Expand right-most dimensions of weights to allow broadcasting
        weights = jnp.reshape(
            scores, scores.shape + (1,) * (top_k_actions.ndim - scores.ndim)
        )

        new_mean = jnp.sum(weights * top_k_actions, axis=0) / (
            jnp.sum(scores, axis=0) + 1e-9
        )  # [H, A]
        new_std = jnp.sqrt(
            jnp.sum(weights * jnp.square(top_k_actions - new_mean), axis=0)
            / (jnp.sum(scores, axis=0) + 1e-9)
        )  # [H, A]
        chex.assert_shape([new_mean, new_std], (horizon, action_dims))

        ## Constrain the stddev using Eqn 5.
        #  Eqn. 5 suggests lower-bounding the std to avoid exploration collapse
        #  In the original implementation, the std is also clipped to be below 2 stddev.
        new_std = jnp.clip(new_std, epsilon, 2.0)
        new_mean = momentum * mean + (1 - momentum) * new_mean
        return (new_mean, new_std, scores, top_k_actions, key)

    # For each iteraction j = 1...J do
    mean, std, scores, top_k_actions, key = jax.lax.fori_loop(
        0, n_iterations, loop_body_fn, initial_state
    )

    chex.assert_shape(scores, (k,))  # [K]
    chex.assert_axis_dimension(scores, 0, k)
    actions = top_k_actions[jax.random.choice(key, k, replace=True, p=scores)]  # [H, A]
    return mean, std, scores, actions


def _compute_n_step_return(
    model_fn: ModelFn,
    critic_fn: CriticFn,
    policy_fn: PolicyFn,
    params: Params,
    z: jax.Array,
    actions: jax.Array,
    key: jax.random.PRNGKeyArray,
    *,
    discount: float,
) -> jax.Array:
    horizon = actions.shape[0]
    return_, discount_ = 0, 1

    for t in range(horizon):
        z, reward = model_fn(params, z, actions[t])
        return_ += discount_ * reward
        discount_ *= discount

    return_ += discount_ * critic_fn(params, z, policy_fn(params, z, key))

    return return_


def _rollout_policy(
    params: Params,
    model_fn: ModelFn,
    policy_fn: PolicyFn,
    z: Embedding,
    key: jax.random.PRNGKeyArray,
    horizon: int,
    num_trajectories: int,
) -> Trajectory:
    policy_actions = []
    z = jnp.tile(z, [num_trajectories, 1])
    for _ in range(horizon):
        key_, key = jax.random.split(key)
        actions = policy_fn(params, z, key_)
        policy_actions.append(actions)
        z, _ = model_fn(params, z, actions)
    return jnp.stack(policy_actions, axis=1)

corax/agents/jax/tdmpc/planning.py
ethanluoyc/corax
"""TDMPC planner"""

import chex
import jax
import jax.numpy as jnp
import numpy as np
from absl.testing import absltest
from absl.testing import parameterized

from corax.agents.jax.tdmpc import planning


class PlannerTest(parameterized.TestCase):
    @parameterized.parameters(
        {"offset": 0.1},
        {"offset": 0.5},
    )
    def test_planner_solve(self, offset):
        horizon, obs_size, act_size = 1, 1, 3

        def policy_fn(params, z, key):
            del params, key
            return jnp.ones(z.shape[:-1] + (act_size,))

        def next_fn(params, z, action):
            del params
            return z, -jnp.sum(jnp.square(action - offset), axis=-1)

        def critic_fn(params, z, action):
            del params, action
            return jnp.zeros(z.shape[:-1])

        def encoder_fn(params, obs):
            del params
            return obs

        params = ()
        key = jax.random.PRNGKey(0)
        observation = jnp.zeros((obs_size,))
        previous_trajectory = planning.get_initial_trajectory(act_size, horizon)
        trajectory_mean, trajectory_std, score, actions = planning.td_mpc_planner(
            policy_fn,
            encoder_fn,
            next_fn,
            critic_fn,
            params,
            observation,
            previous_trajectory,
            key,
            n_policy_trajectories=32,
            n_sample_trajectories=512,
            epsilon=0.0,
            k=64,
            discount=1.0,
            temperature=1.0,
            momentum=0.0,
            n_iterations=20,
        )
        chex.assert_shape(trajectory_mean, (horizon, act_size))
        chex.assert_shape(trajectory_std, (horizon, act_size))
        chex.assert_shape(score, (64,))
        chex.assert_trees_all_close(actions, jnp.full((horizon, act_size), offset))

    def test_compute_n_step_return(self):
        params = ()
        act_size = 1

        def next_fn(params, z, action):
            del params
            return z, jnp.sum(jnp.square(action), axis=-1)

        def policy_fn(params, z, key):
            del params, key
            return jnp.ones(z.shape[:-1] + (act_size,))

        actions = jnp.array([1.0, 2.0, 3.0])
        actions = jnp.reshape(actions, actions.shape + (1, 1))
        horizon = actions.shape[0]
        z = jnp.ones((horizon, 1, 1))

        bootstrap_value = 8.0

        def critic_fn(params, z, action):
            del params, action, z
            return bootstrap_value

        discount = 0.99
        key = jax.random.PRNGKey(0)

        return_prediction = (
            planning._compute_n_step_return(  # pylint: disable=protected-access
                next_fn,
                critic_fn,  # type: ignore
                policy_fn,
                params,
                z,
                actions,
                key,
                discount=discount,
            )
        )

        return_expected = (
            jnp.sum(jnp.square(actions[0]))
            + jnp.sum(jnp.square(actions[1])) * (discount**1)
            + jnp.sum(jnp.square(actions[2])) * (discount**2)
            + bootstrap_value * (discount**3)
        )
        np.testing.assert_allclose(return_expected, return_prediction)


if __name__ == "__main__":
    absltest.main()

corax/agents/jax/tdmpc/planning_test.py
ethanluoyc/corax
import dataclasses


@dataclasses.dataclass(frozen=True)
class LossScalesConfig:
    consistency: float = 2.0
    reward: float = 0.5
    value: float = 0.1

corax/agents/jax/tdmpc/types.py
ethanluoyc/corax
"""Core Corax interfaces.

This file specifies and documents the notions of `Actor` and `Learner`.
"""

import abc
import itertools
from typing import Generic, Iterator, List, Optional, Sequence, TypeVar

import dm_env

from corax import types

T = TypeVar("T")


class Actor(abc.ABC):
    """Interface for an agent that can act.

    This interface defines an API for an Actor to interact with an EnvironmentLoop
    (see acme.environment_loop), e.g. a simple RL loop where each step is of the
    form:

      # Make the first observation.
      timestep = env.reset()
      actor.observe_first(timestep)

      # Take a step and observe.
      action = actor.select_action(timestep.observation)
      next_timestep = env.step(action)
      actor.observe(action, next_timestep)

      # Update the actor policy/parameters.
      actor.update()
    """

    @abc.abstractmethod
    def select_action(self, observation: types.NestedArray) -> types.NestedArray:
        """Samples from the policy and returns an action."""

    @abc.abstractmethod
    def observe_first(self, timestep: dm_env.TimeStep):
        """Make a first observation from the environment.

        Note that this need not be an initial state, it is merely beginning the
        recording of a trajectory.

        Args:
          timestep: first timestep.
        """

    @abc.abstractmethod
    def observe(
        self,
        action: types.NestedArray,
        next_timestep: dm_env.TimeStep,
    ):
        """Make an observation of timestep data from the environment.

        Args:
          action: action taken in the environment.
          next_timestep: timestep produced by the environment given the action.
        """

    @abc.abstractmethod
    def update(self, wait: bool = False):
        """Perform an update of the actor parameters from past observations.

        Args:
          wait: if True, the update will be blocking.
        """


class VariableSource(abc.ABC):
    """Abstract source of variables.

    Objects which implement this interface provide a source of variables, returned
    as a collection of (nested) numpy arrays. Generally this will be used to
    provide variables to some learned policy/etc.
    """

    @abc.abstractmethod
    def get_variables(self, names: Sequence[str]) -> List[types.NestedArray]:
        """Return the named variables as a collection of (nested) numpy arrays.

        Args:
          names: args where each name is a string identifying a predefined subset of
            the variables.

        Returns:
          A list of (nested) numpy arrays `variables` such that `variables[i]`
          corresponds to the collection named by `names[i]`.
        """


class Worker(abc.ABC):
    """An interface for (potentially) distributed workers."""

    @abc.abstractmethod
    def run(self):
        """Runs the worker."""


class Saveable(abc.ABC, Generic[T]):
    """An interface for saveable objects."""

    @abc.abstractmethod
    def save(self) -> T:
        """Returns the state from the object to be saved."""

    @abc.abstractmethod
    def restore(self, state: T):
        """Given the state, restores the object."""


class Learner(VariableSource, Worker, Saveable):
    """Abstract learner object.

    This corresponds to an object which implements a learning loop. A single step
    of learning should be implemented via the `step` method and this step
    is generally interacted with via the `run` method which runs update
    continuously.

    All objects implementing this interface should also be able to take in an
    external dataset (see acme.datasets) and run updates using data from this
    dataset. This can be accomplished by explicitly running `learner.step()`
    inside a for/while loop or by using the `learner.run()` convenience function.
    Data will be read from this dataset asynchronously and this is primarily
    useful when the dataset is filled by an external process.
    """

    @abc.abstractmethod
    def step(self):
        """Perform an update step of the learner's parameters."""

    def run(self, num_steps: Optional[int] = None) -> None:
        """Run the update loop; typically an infinite loop which calls step."""

        iterator = range(num_steps) if num_steps is not None else itertools.count()

        for _ in iterator:
            self.step()

    def save(self):
        raise NotImplementedError('Method "save" is not implemented.')

    def restore(self, state):
        raise NotImplementedError('Method "restore" is not implemented.')


class PrefetchingIterator(Iterator[T], abc.ABC):
    """Abstract iterator object which supports `ready` method."""

    @abc.abstractmethod
    def ready(self) -> bool:
        """Is there any data waiting for processing."""

    @abc.abstractmethod
    def retrieved_elements(self) -> int:
        """How many elements were retrieved from the iterator."""

corax/core.py
ethanluoyc/corax
# 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.

"""Tests for core.py."""

from typing import List

from absl.testing import absltest

from corax import core
from corax import types


class StepCountingLearner(core.Learner):
    """A learner which counts `num_steps` and then raises `StopIteration`."""

    def __init__(self, num_steps: int):
        self.step_count = 0
        self.num_steps = num_steps

    def step(self):
        self.step_count += 1
        if self.step_count >= self.num_steps:
            raise StopIteration()

    def get_variables(self, unused: List[str]) -> List[types.NestedArray]:
        del unused
        return []


class CoreTest(absltest.TestCase):
    def test_learner_run_with_limit(self):
        learner = StepCountingLearner(100)
        learner.run(7)
        self.assertEqual(learner.step_count, 7)

    def test_learner_run_no_limit(self):
        learner = StepCountingLearner(100)
        with self.assertRaises(StopIteration):
            learner.run()
        self.assertEqual(learner.step_count, 100)


if __name__ == "__main__":
    absltest.main()

corax/core_test.py
ethanluoyc/corax
# 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.

"""Functions for making TensorFlow datasets for sampling from Reverb replay."""

import collections
import os
from typing import Callable, Mapping, Optional, Union

import reverb
import tensorflow as tf

from corax.adders import reverb as adders

Transform = Callable[[reverb.ReplaySample], reverb.ReplaySample]


def make_reverb_dataset(
    server_address: str,
    batch_size: Optional[int] = None,
    prefetch_size: Optional[int] = None,
    table: Union[str, Mapping[str, float]] = adders.DEFAULT_PRIORITY_TABLE,
    num_parallel_calls: Optional[int] = 12,
    max_in_flight_samples_per_worker: Optional[int] = None,
    postprocess: Optional[Transform] = None,
) -> tf.data.Dataset:
    """Make a TensorFlow dataset backed by a Reverb trajectory replay service.

    Arguments:
      server_address: Address of the Reverb server.
      batch_size: Batch size of the returned dataset.
      prefetch_size: The number of elements to prefetch from the original dataset.
        Note that Reverb may do some internal prefetching in addition to this.
      table: The name of the Reverb table to use, or a mapping of (table_name,
        float_weight) for mixing multiple tables in the input (e.g. mixing online
        and offline experiences).
      num_parallel_calls: The parralelism to use. Setting it to `tf.data.AUTOTUNE`
        will allow `tf.data` to automatically find a reasonable value.
      max_in_flight_samples_per_worker: see reverb.TrajectoryDataset for details.
      postprocess: User-specified transformation to be applied to the dataset (as
        `ds.map(postprocess)`).

    Returns:
      A `tf.data.Dataset` iterating over the contents of the Reverb table.

    Raises:
      ValueError if `environment_spec` or `extra_spec` are set, or `table` is a
      mapping with no positive weight values.
    """

    # This is the default that used to be set by reverb.TFClient.dataset().
    if max_in_flight_samples_per_worker is None and batch_size is None:
        max_in_flight_samples_per_worker = 100
    elif max_in_flight_samples_per_worker is None:
        assert batch_size is not None
        max_in_flight_samples_per_worker = 2 * batch_size

    # Create mapping from tables to non-zero weights.
    if isinstance(table, str):
        tables = collections.OrderedDict([(table, 1.0)])
    else:
        tables = collections.OrderedDict(
            [(name, weight) for name, weight in table.items() if weight > 0.0]
        )
        if len(tables) <= 0:
            raise ValueError(f"No positive weights in input tables {tables}")

    # Normalize weights.
    total_weight = sum(tables.values())
    tables = collections.OrderedDict(
        [(name, weight / total_weight) for name, weight in tables.items()]
    )

    def _make_dataset(unused_idx: tf.Tensor) -> tf.data.Dataset:
        datasets = ()
        for table_name, weight in tables.items():
            max_in_flight_samples = max(
                1, int(max_in_flight_samples_per_worker * weight)
            )
            dataset = reverb.TrajectoryDataset.from_table_signature(
                server_address=server_address,
                table=table_name,
                max_in_flight_samples_per_worker=max_in_flight_samples,
            )
            datasets += (dataset,)
        if len(datasets) > 1:
            dataset = tf.data.Dataset.sample_from_datasets(
                datasets, weights=tables.values()
            )
        else:
            dataset = datasets[0]

        # Post-process each element if a post-processing function is passed, e.g.
        # observation-stacking or data augmenting transformations.
        if postprocess:
            dataset = dataset.map(postprocess)

        if batch_size:
            dataset = dataset.batch(batch_size, drop_remainder=True)

        return dataset

    if num_parallel_calls is not None:
        # Create a datasets and interleaves it to create `num_parallel_calls`
        # `TrajectoryDataset`s.
        num_datasets_to_interleave = (
            os.cpu_count()
            if num_parallel_calls == tf.data.AUTOTUNE
            else num_parallel_calls
        )
        dataset = tf.data.Dataset.range(num_datasets_to_interleave).interleave(
            map_func=_make_dataset,
            cycle_length=num_parallel_calls,
            num_parallel_calls=num_parallel_calls,
            deterministic=False,
        )
    else:
        dataset = _make_dataset(tf.constant(0))  # type: ignore

    if prefetch_size:
        dataset = dataset.prefetch(prefetch_size)

    return dataset

corax/datasets/reverb.py
ethanluoyc/corax
# Import all TFDS datasets here for dataset registration.
import corax.datasets.tfds.adroit_binary
import corax.datasets.tfds.exorl
import corax.datasets.tfds.gwil
import corax.datasets.tfds.rot
import corax.datasets.tfds.vd4rl

# Import utilities
from corax.datasets.tfds.utils import JaxInMemoryRandomSampleIterator
from corax.datasets.tfds.utils import get_tfds_dataset
from corax.datasets.tfds.utils import load_tfds_dataset

corax/datasets/tfds/__init__.py
ethanluoyc/corax
from corax.datasets.tfds.adroit_binary.adroit_binary import AdroitBinary

corax/datasets/tfds/adroit_binary/__init__.py
ethanluoyc/corax
"""adroit_binary dataset."""

import dataclasses
import os

import numpy as np
import tensorflow_datasets as tfds
import tree


@dataclasses.dataclass
class BuilderConfig(tfds.core.BuilderConfig):
    filename: str = ""


@dataclasses.dataclass
class DatasetConfig:
    task: str
    obs_dim: int
    action_dim: int


_DATASET_CONFIGS = {
    "pen": DatasetConfig(task="pen", obs_dim=45, action_dim=24),
    "door": DatasetConfig(task="door", obs_dim=39, action_dim=28),
    "relocate": DatasetConfig(task="relocate", obs_dim=39, action_dim=30),
}


class AdroitBinary(tfds.core.GeneratorBasedBuilder):
    """DatasetBuilder for adroit_binary dataset."""

    VERSION = tfds.core.Version("1.0.0")
    RELEASE_NOTES = {
        "1.0.0": "Initial release.",
    }

    # pylint: disable=unexpected-keyword-arg
    BUILDER_CONFIGS = [
        BuilderConfig(name="pen-binary-expert-v0", filename="pen2_sparse.npy"),
        BuilderConfig(name="door-binary-expert-v0", filename="door2_sparse.npy"),
        BuilderConfig(
            name="relocate-binary-expert-v0", filename="relocate2_sparse.npy"
        ),
        BuilderConfig(name="pen-binary-bc-v0", filename="pen_bc_sparse4.npy"),
        BuilderConfig(name="door-binary-bc-v0", filename="door_bc_sparse4.npy"),
        BuilderConfig(name="relocate-binary-bc-v0", filename="relocate_bc_sparse4.npy"),
    ]

    # pylint: enable=unexpected-keyword-arg

    def _info(self) -> tfds.core.DatasetInfo:
        """Returns the dataset metadata."""
        task = self.builder_config.name.split("-")[0]
        dataset_config = _DATASET_CONFIGS[task]

        action_dtype = np.float64
        if "bc-v0" in self.builder_config.name:
            action_dtype = np.float32

        steps_dict = {
            "observation": tfds.features.Tensor(
                shape=(dataset_config.obs_dim,),
                encoding=tfds.features.Encoding.ZLIB,
                dtype=np.float64,
            ),
            "action": tfds.features.Tensor(
                shape=(dataset_config.action_dim,),
                encoding=tfds.features.Encoding.ZLIB,
                dtype=action_dtype,
            ),
            "reward": np.float64,
            # "discount": np.float32,
            "is_terminal": np.bool_,
            "is_first": np.bool_,
            "is_last": np.bool_,
        }
        if "bc-v0" in self.builder_config.name:
            steps_dict["infos"] = {
                "goal_achieved": tfds.features.Tensor(shape=(), dtype=np.bool_),
                "env_reward": tfds.features.Tensor(shape=(), dtype=np.float64),
            }
        return tfds.core.DatasetInfo(
            builder=self,
            description="TODO",
            features=tfds.features.FeaturesDict(
                {
                    # These are the features of your dataset like images, labels ...
                    # "episode_id": np.int64,
                    "steps": tfds.features.Dataset(steps_dict),
                }
            ),
            # If there's a common (input, target) tuple from the
            # features, specify them here. They'll be used if
            # `as_supervised=True` in `builder.as_dataset`.
            supervised_keys=None,  # disabled
            homepage="""TODO""",
            citation="""TODO""",
        )

    def _split_generators(self, dl_manager: tfds.download.DownloadManager):
        """Returns SplitGenerators."""
        dataset_url = "https://drive.google.com/uc?export=download&id=1SsVaQKZnY5UkuR78WrInp9XxTdKHbF0x"
        extracted_path = dl_manager.download_and_extract({"file_path": dataset_url})
        file_path = extracted_path["file_path"] / self.builder_config.filename

        # TODO(exorl): Returns the Dict[split names, Iterator[Key, Example]]
        return {
            "train": self._generate_examples(file_path),
        }

    def _generate_examples(self, path):
        """Yields examples."""
        # TODO(exorl): Yields (key, example) tuples from the dataset
        if "bc_sparse" in os.fspath(path):
            for i, episode in enumerate(load_bc_data(path)):
                yield i, {"steps": episode}
        else:
            for i, episode in enumerate(load_expert_data(path)):
                yield i, {"steps": episode}


def load_expert_data(filename):
    dataset = np.load(filename, allow_pickle=True)
    episodes = []
    for episode in dataset:
        observations = tree.map_structure(
            lambda *x: np.stack(x), *episode["observations"]
        )
        actions = tree.map_structure(lambda *x: np.stack(x), *episode["actions"])
        next_observations = tree.map_structure(
            lambda *x: np.stack(x), *episode["next_observations"]
        )
        rewards = tree.map_structure(lambda *x: np.stack(x), *episode["rewards"])
        terminals = tree.map_structure(lambda *x: np.stack(x), *episode["terminals"])
        agent_infos = tree.map_structure(
            lambda *x: np.stack(x), *episode["agent_infos"]
        )
        env_infos = tree.map_structure(lambda *x: np.stack(x), *episode["env_infos"])
        assert len(agent_infos) == 0
        assert len(env_infos) == 0
        # Expert data has one additional reward
        assert terminals.shape[0] == rewards.shape[0] - 1
        # Expert data observation is redundant
        np.testing.assert_allclose(
            observations["observation"], observations["state_observation"]
        )
        np.testing.assert_allclose(
            next_observations["observation"], next_observations["state_observation"]
        )

        # Assert next observation is the same as the previous observation offset by 1
        np.testing.assert_allclose(
            observations["observation"][1:], next_observations["observation"][:-1]
        )
        np.testing.assert_allclose(
            observations["state_observation"][1:],
            next_observations["state_observation"][:-1],
        )
        episodes.append(
            {
                "is_first": np.asarray(
                    [True] + observations["observation"].shape[0] * [False]
                ),
                "is_last": np.asarray(
                    ([False] * observations["observation"].shape[0] + [True])
                ),
                "is_terminal": np.concatenate([[False], terminals], axis=0),
                "observation": np.concatenate(
                    [
                        observations["observation"],
                        next_observations["observation"][-1:],
                    ],
                    axis=0,
                ),
                "action": np.concatenate(
                    [actions, np.zeros_like(actions[-1:])], axis=0
                ),
                "reward": rewards,
            }
        )
    return episodes


def load_bc_data(filename):
    dataset = np.load(filename, allow_pickle=True)
    episodes = []
    for episode in dataset:
        observations = tree.map_structure(
            lambda *x: np.stack(x), *episode["observations"]
        )
        actions = tree.map_structure(lambda *x: np.stack(x), *episode["actions"])
        next_observations = tree.map_structure(
            lambda *x: np.stack(x), *episode["next_observations"]
        )
        rewards = tree.map_structure(lambda *x: np.concatenate(x), *episode["rewards"])
        # terminals in BC data is rank 1
        terminals = tree.map_structure(
            lambda *x: np.concatenate(x), *episode["terminals"]
        )
        assert terminals.shape[0] == rewards.shape[0]
        agent_infos = tree.map_structure(
            lambda *x: np.stack(x), *episode["agent_infos"]
        )
        # env_infos = tree.map_structure(lambda *x: np.stack(x), *episode["env_infos"])
        np.testing.assert_allclose(observations[1:], next_observations[:-1])
        goal_achieved = np.stack(
            [info["goal_achieved"] for info in episode["env_infos"]]
        )
        env_reward = np.stack([info["env_reward"] for info in episode["env_infos"]])
        is_truncated = episode["env_infos"][-1]["TimeLimit.truncated"]
        # Check that terminals is actually referring to done in Gym 21 API.
        assert terminals[-1] == is_truncated
        # assertion passed, update terminal to use truncation
        terminals[-1] = not is_truncated
        assert len(agent_infos) == 0
        episodes.append(
            {
                "is_first": np.asarray([True] + observations.shape[0] * [False]),
                "is_last": np.asarray(([False] * observations.shape[0] + [True])),
                "is_terminal": np.concatenate([[False], terminals], axis=0),
                "reward": np.concatenate(
                    [rewards, np.zeros_like(rewards[-1:])], axis=0
                ),
                "observation": np.concatenate(
                    [
                        observations,
                        next_observations[-1:],
                    ],
                    axis=0,
                ),
                "action": np.concatenate(
                    [actions, np.zeros_like(actions[-1:])], axis=0
                ),
                "infos": {
                    "goal_achieved": np.concatenate(
                        [goal_achieved, np.zeros_like(goal_achieved[-1:])], axis=0
                    ),
                    "env_reward": np.concatenate(
                        [env_reward, np.zeros_like(env_reward[-1:])], axis=0
                    ),
                },
            }
        )
    return episodes

corax/datasets/tfds/adroit_binary/adroit_binary.py
ethanluoyc/corax
"""exorl dataset."""

from corax.datasets.tfds.exorl.exorl import Exorl

corax/datasets/tfds/exorl/__init__.py
ethanluoyc/corax
"""exorl dataset."""

import dataclasses

import numpy as np
import tensorflow_datasets as tfds

_DESCRIPTION = """
Exploratory dataset from the

"Don't Change the Algorithm, Change the Data: Exploratory Data for Offline Reinforcement Learning"

paper by Denis Yarats, David Brandfonbrener, Hao Liu, Misha Laskin, Pieter Abbeel, Alessandro Lazaric, and Lerrel Pinto.

The datasets follow the RLDS format to represent steps and episodes.
"""

# TODO(exorl): BibTeX citation
_CITATION = """
@article{yarats2022exorl,
  title={Don't Change the Algorithm, Change the Data: Exploratory Data for Offline Reinforcement Learning},
  author={Denis Yarats, David Brandfonbrener, Hao Liu, Michael Laskin, Pieter Abbeel, Alessandro Lazaric, Lerrel Pinto},
  journal={arXiv preprint arXiv:2201.13425},
  year={2022}
}
"""

_HOMEPAGE = "https://github.com/denisyarats/exorl"
_BASE_URL = "https://dl.fbaipublicfiles.com/exorl/{domain}/{algorithm}.zip"

_DATASET_DESCRIPTION = """
See more details about the tasks and algorithms in https://github.com/denisyarats/exorl
"""


@dataclasses.dataclass
class ExorlBuilderConfig(tfds.core.BuilderConfig):
    algorithm: str = "proto"
    domain: str = "walker"


@dataclasses.dataclass
class DatasetConfig:
    domain: str
    obs_dim: int = 24
    action_dim: int = 6
    physics_dim: int = 18


_DATASET_CONFIGS = {
    "cartpole": DatasetConfig(
        domain="cartpole", obs_dim=5, action_dim=1, physics_dim=4
    ),
    "cheetah": DatasetConfig(
        domain="cheetah", obs_dim=17, action_dim=6, physics_dim=18
    ),
    "jaco": DatasetConfig(domain="jaco", obs_dim=42, action_dim=9, physics_dim=18),
    "quadruped": DatasetConfig(
        domain="quadruped", obs_dim=78, action_dim=12, physics_dim=57
    ),
    "point_mass_maze": DatasetConfig(
        domain="point_mass_maze", obs_dim=4, action_dim=2, physics_dim=4
    ),
    "walker": DatasetConfig(domain="walker", obs_dim=24, action_dim=6, physics_dim=18),
}


class Exorl(tfds.core.GeneratorBasedBuilder):
    """DatasetBuilder for exorl dataset."""

    VERSION = tfds.core.Version("1.0.0")
    RELEASE_NOTES = {
        "1.0.0": "Initial release.",
    }

    # pylint: disable=unexpected-keyword-arg
    BUILDER_CONFIGS = [
        # proto
        ExorlBuilderConfig(
            name="cartpole_proto",
            algorithm="proto",
            domain="cartpole",
            description=_DATASET_DESCRIPTION,
        ),
        ExorlBuilderConfig(
            name="cheetah_proto",
            algorithm="proto",
            domain="cheetah",
            description=_DATASET_DESCRIPTION,
        ),
        ExorlBuilderConfig(
            name="jaco_proto",
            algorithm="proto",
            domain="jaco",
            description=_DATASET_DESCRIPTION,
        ),
        ExorlBuilderConfig(
            name="quadruped_proto",
            algorithm="proto",
            domain="quadruped",
            description=_DATASET_DESCRIPTION,
        ),
        ExorlBuilderConfig(
            name="point_mass_maze_proto",
            algorithm="proto",
            domain="point_mass_maze",
            description=_DATASET_DESCRIPTION,
        ),
        ExorlBuilderConfig(
            name="walker_proto",
            algorithm="proto",
            domain="walker",
            description=_DATASET_DESCRIPTION,
        ),
    ]

    # pylint: enable=unexpected-keyword-arg

    def _info(self) -> tfds.core.DatasetInfo:
        """Returns the dataset metadata."""
        dataset_config = _DATASET_CONFIGS[self.builder_config.domain]
        steps_dict = {
            "observation": tfds.features.Tensor(
                shape=(dataset_config.obs_dim,),
                encoding=tfds.features.Encoding.ZLIB,
                dtype=np.float32,
            ),
            "action": tfds.features.Tensor(
                shape=(dataset_config.action_dim,),
                encoding=tfds.features.Encoding.ZLIB,
                dtype=np.float32,
            ),
            "reward": np.float32,
            "discount": np.float32,
            "physics": tfds.features.Tensor(
                shape=(dataset_config.physics_dim,),
                encoding=tfds.features.Encoding.ZLIB,
                dtype=np.float64,
            ),
            "is_terminal": np.bool_,
            "is_first": np.bool_,
            "is_last": np.bool_,
        }
        return tfds.core.DatasetInfo(
            builder=self,
            description=_DESCRIPTION,
            features=tfds.features.FeaturesDict(
                {
                    # These are the features of your dataset like images, labels ...
                    "episode_id": np.int64,
                    "steps": tfds.features.Dataset(steps_dict),
                }
            ),
            # If there's a common (input, target) tuple from the
            # features, specify them here. They'll be used if
            # `as_supervised=True` in `builder.as_dataset`.
            supervised_keys=None,  # disabled
            homepage=_HOMEPAGE,
            citation=_CITATION,
        )

    def _split_generators(self, dl_manager: tfds.download.DownloadManager):
        """Returns SplitGenerators."""
        # TODO(exorl): Downloads the data and defines the splits
        dataset_url = _BASE_URL.format(
            domain=self.builder_config.domain, algorithm=self.builder_config.algorithm
        )
        extracted_path = dl_manager.download_and_extract({"file_path": dataset_url})

        # TODO(exorl): Returns the Dict[split names, Iterator[Key, Example]]
        return {
            "train": self._generate_examples(extracted_path),
        }

    def _generate_examples(self, path):
        """Yields examples."""
        # TODO(exorl): Yields (key, example) tuples from the dataset
        buffer_path = path["file_path"] / "buffer"
        for i, f in enumerate(sorted(buffer_path.glob("*.npz"))):
            episode = _decode_episode(f)
            yield i, episode


def _decode_episode(filename):
    with filename.open("rb") as f:
        data = np.load(f)
        data = {k: data[k] for k in data.keys()}
        episode = {}
        # filenames has the format
        # episode_{episode_number}_{episode_length}.npz
        episode_id = int(filename.name.split("_")[1])
        episode["observation"] = data["observation"]
        episode["physics"] = data["physics"]
        # Change to SAR alignment following the RLDS convention
        episode["action"] = np.concatenate(
            (data["action"][1:], [np.zeros_like(data["action"][0])])
        )
        episode["reward"] = np.concatenate(
            (data["reward"][1:], [np.zeros_like(data["reward"][0])])
        )
        episode["reward"] = np.squeeze(episode["reward"], axis=-1)
        episode["discount"] = np.concatenate(
            (data["discount"][1:], [np.zeros_like(data["discount"][0])])
        )
        episode["discount"] = np.squeeze(episode["discount"], axis=-1)
        episode["is_first"] = [True] + [False] * (len(data["observation"]) - 1)
        episode["is_last"] = [False] * (len(data["observation"]) - 1) + [True]
        episode["is_terminal"] = [False] * (len(data["observation"]))

        return {"episode_id": episode_id, "steps": episode}

corax/datasets/tfds/exorl/exorl.py
ethanluoyc/corax
"""exorl dataset."""

import tensorflow_datasets as tfds

from corax.datasets.tfds.exorl import exorl


class ExorlTest(tfds.testing.DatasetBuilderTestCase):
    """Tests for exorl dataset."""

    DATASET_CLASS = exorl.Exorl
    SPLITS = {
        "train": 1,  # Number of fake train example
    }

    SKIP_CHECKSUMS = True
    SKIP_TF1_GRAPH_MODE = True

    BUILDER_CONFIG_NAMES_TO_TEST = ["walker_proto"]

    # If you are calling `download/download_and_extract` with a dict, like:
    #   dl_manager.download({'some_key': 'http://a.org/out.txt', ...})
    # then the tests needs to provide the fake output paths relative to the
    # fake data directory
    DL_EXTRACT_RESULT = {"file_path": "walker_proto"}
    DL_DOWNLOAD_RESULT = {"file_path": "walker_proto"}


if __name__ == "__main__":
    tfds.testing.test_main()

corax/datasets/tfds/exorl/exorl_test.py
ethanluoyc/corax
"""gwil dataset."""

from corax.datasets.tfds.gwil.gwil import Gwil

corax/datasets/tfds/gwil/__init__.py
ethanluoyc/corax
"""gwil dataset."""

import dataclasses
import os
import pickle

import numpy as np
import tensorflow as tf
import tensorflow_datasets as tfds
from tensorflow_datasets.rlds import rlds_base

_DESCRIPTION = """
Expert demonstrations used in [Fickinger et al., 2022](https://arxiv.org/pdf/2110.03684.pdf).
"""

_CITATION = """
@inproceedings{fickinger2022gromov,
  title={Cross-Domain Imitation Learning via Optimal Transport},
  author={Fickinger, Arnaud and Cohen, Samuel and Russell, Stuart and Amos, Brandon},
  booktitle={10th International Conference on Learning Representations, ICLR},
  year={2022}
}
"""

_HOMEPAGE = "https://arnaudfickinger.github.io/gwil/"

_DRIVE_PATH = (
    "https://drive.google.com/uc?export=download&id=1xE882IuQkXUuaeXHInYaP9eqvhQm48Et"
)
_DATA_PATHS = {
    "cartpole_swingup": "expert_trajectory_cartpole_swingup.pickle",
    "walker_walk": "expert_trajectory_walker_walk.pickle",
    "cheetah_run": "expert_trajectory_cheetah_run.pickle",
    "pendulum_swingup": "expert_trajectory_pendulum_swingup.pickle",
    "maze_end_pointmass": "expert_trajectory_MazeEnd_PointMass_0.pickle",
}


def build_episode(demo_file: os.PathLike):
    # GWIL dataset contains only a single episode
    time_limit = 1000
    with open(demo_file, "rb") as infile:
        traj = pickle.load(infile)
    # if "Maze" in os.path.normpath(demo_file):
    #   dmc_observations = traj['dmc_obs']
    #   observation_keys = traj['dmc_obs'].keys()
    #   flat_observation = np.concatenate([traj['dmc_obs'][key] for key in observation_keys], axis=-1)
    #   print(traj['dmc_obs'].keys())
    # Data validation
    # Check next_observation and observation is consistent.
    assert (traj["obs"][1:] == traj["nobs"][:-1]).all()
    # NOTE: Pixel observation in the GWIL episode is inconsistent
    # In addition, they are not the same as the pixel observations used in typical pixel agents.
    # Therefore, don't use the pixel observation for training.
    episode_length = len(traj["reward"])
    ends_in_terminal = episode_length < time_limit
    steps_dict = {
        "observation": {
            "states": np.concatenate([traj["obs"], [traj["nobs"][-1]]], axis=0),
            "pixels": np.transpose(
                np.concatenate([traj["pixel_obs"], [traj["pixel_nobs"][-1]]], axis=0),
                [0, 2, 3, 1],
            ),
        },
        "action": np.concatenate(
            [traj["action"], [np.zeros_like(traj["action"][0])]], axis=0
        ),
        "reward": np.concatenate(
            [np.asarray(traj["reward"]), [np.zeros_like(traj["reward"][0])]], axis=0
        ),
        "is_first": np.asarray([True] + [False] * episode_length),
        "is_last": np.asarray([False] * episode_length + [True]),
        "is_terminal": np.asarray([False] * episode_length + [ends_in_terminal]),
    }
    episode = {"episode_return": traj["cumulative_reward"], "steps": steps_dict}
    return episode


@dataclasses.dataclass
class GwilBuilderConfig(tfds.core.BuilderConfig):
    state_size: int = 0
    action_size: int = 0


class Gwil(tfds.core.GeneratorBasedBuilder):
    """DatasetBuilder for gwil dataset."""

    VERSION = tfds.core.Version("1.0.0")
    RELEASE_NOTES = {
        "1.0.0": "Initial release.",
    }

    # pylint: disable=unexpected-keyword-arg
    BUILDER_CONFIGS = [
        GwilBuilderConfig(
            name="cartpole_swingup",
            state_size=5,
            action_size=1,
        ),
        GwilBuilderConfig(
            name="pendulum_swingup",
            state_size=3,
            action_size=1,
        ),
        GwilBuilderConfig(
            name="walker_walk",
            state_size=24,
            action_size=6,
        ),
        GwilBuilderConfig(
            name="cheetah_run",
            state_size=17,
            action_size=6,
        ),
        GwilBuilderConfig(
            name="maze_end_pointmass",
            state_size=6,
            action_size=2,
        ),
    ]

    # pylint: enable=unexpected-keyword-arg

    def _info(self) -> tfds.core.DatasetInfo:
        """Returns the dataset metadata."""
        # pylint: disable=unexpected-keyword-arg
        builder_config = self.builder_config
        observation_info = {
            "states": tfds.features.Tensor(
                shape=(builder_config.state_size,),
                dtype=tf.float64,
                encoding=tfds.features.Encoding.ZLIB,
            ),
            "pixels": tfds.features.Image(
                shape=(84, 84, 3), dtype=tf.uint8, encoding_format="png"
            ),
        }
        action_info = tfds.features.Tensor(
            shape=(builder_config.action_size,), dtype=tf.float32
        )
        ds_config = rlds_base.DatasetConfig(
            name=builder_config.name,
            observation_info=observation_info,
            action_info=action_info,
            reward_info=tfds.features.Tensor(shape=(), dtype=tf.float64),
            episode_metadata_info={
                "episode_return": tfds.features.Tensor(shape=(), dtype=tf.float64)
            },
            citation=_CITATION,
            homepage=_HOMEPAGE,
            overall_description=_DESCRIPTION,
            supervised_keys=None,  # pytype: disable=wrong-arg-types  # gen-stub-imports
        )
        return rlds_base.build_info(ds_config, self)

    def _split_generators(self, dl_manager: tfds.download.DownloadManager):
        """Returns SplitGenerators."""
        # TODO(gwil): Downloads the data and defines the splits
        path = dl_manager.download_and_extract({"file_path": _DRIVE_PATH})["file_path"]

        filename = _DATA_PATHS[self.builder_config.name]

        # TODO(gwil): Returns the Dict[split names, Iterator[Key, Example]]
        return {
            "train": self._generate_examples(path / "exp" / filename),
        }

    def _generate_examples(self, path):
        """Yields examples."""
        # TODO(gwil): Yields (key, example) tuples from the dataset
        key = os.path.basename(path)
        yield key, build_episode(path)

corax/datasets/tfds/gwil/gwil.py
ethanluoyc/corax
"""rot dataset."""

from corax.datasets.tfds.rot.rot_dmc import RotDmc

corax/datasets/tfds/rot/__init__.py
ethanluoyc/corax
"""rot dataset."""

import dataclasses
import os
import pickle

import numpy as np
import tensorflow as tf
import tensorflow_datasets as tfds
from tensorflow_datasets.rlds import rlds_base

_DESCRIPTION = """
Datasets for the [ROT paper](https://arxiv.org/abs/2206.15469).

NOTE: last observations are not recorded in the dataset.
"""

_CITATION = """
@inproceedings{
haldar2022watch,
title={Watch and Match: Supercharging Imitation with Regularized Optimal Transport},
author={Siddhant Haldar and Vaibhav Mathur and Denis Yarats and Lerrel Pinto},
booktitle={6th Annual Conference on Robot Learning},
year={2022},
url={https://openreview.net/forum?id=ZUtgUA0Fuwd}
}
"""
_HOMEPAGE = """
https://rot-robot.github.io/
"""

_FILE_URL = "https://osf.io/zgh2f/download"


def _build_rot_dataset(file_path: os.PathLike):
    with tf.io.gfile.GFile(file_path, "rb") as f:
        demo = pickle.load(f)
        expert_pixels, expert_states, expert_actions, expert_reward = demo
        # [N, T, C, H, W] -> [N, T, H, W, C]
        expert_pixels = np.transpose(expert_pixels, [0, 1, 3, 4, 2])

    num_demonstrations = expert_pixels.shape[0]
    episode_length = expert_pixels.shape[1]

    episodes = []
    for i in range(num_demonstrations):
        episodes.append(
            {
                "steps": {
                    "observation": {
                        "pixels": expert_pixels[i],
                        "states": expert_states[i],
                    },
                    "action": expert_actions[i],
                    "reward": expert_reward[i],
                    "is_first": [True] + [False] * (episode_length - 1),
                    "is_last": [False] * (episode_length - 1) + [True],
                    "is_terminal": [False] * episode_length,
                }
            }
        )
    return episodes


@dataclasses.dataclass
class RotDmcConfig(tfds.core.BuilderConfig):
    suite: str = "dmc"
    observation_size: int = 0
    action_size: int = 0


class RotDmc(tfds.core.GeneratorBasedBuilder):
    """DatasetBuilder for rot dataset."""

    VERSION = tfds.core.Version("1.0.0")
    RELEASE_NOTES = {
        "1.0.0": "Initial release.",
    }

    # pylint: disable=unexpected-keyword-arg
    BUILDER_CONFIGS = [
        RotDmcConfig(
            name="cheetah_run",
            suite="dmc",
            observation_size=17,
            action_size=6,
        ),
        RotDmcConfig(
            name="cartpole_swingup",
            suite="dmc",
            observation_size=5,
            action_size=1,
        ),
        RotDmcConfig(
            name="acrobot_swingup",
            suite="dmc",
            observation_size=6,
            action_size=1,
        ),
        RotDmcConfig(
            name="finger_spin",
            suite="dmc",
            observation_size=9,
            action_size=2,
        ),
        RotDmcConfig(
            name="hopper_hop",
            suite="dmc",
            observation_size=15,
            action_size=4,
        ),
        RotDmcConfig(
            name="hopper_stand",
            suite="dmc",
            observation_size=15,
            action_size=4,
        ),
        RotDmcConfig(
            name="quadruped_run",
            suite="dmc",
            observation_size=78,
            action_size=12,
        ),
        RotDmcConfig(
            name="walker_run",
            suite="dmc",
            observation_size=24,
            action_size=6,
        ),
        RotDmcConfig(
            name="walker_stand",
            suite="dmc",
            observation_size=24,
            action_size=6,
        ),
        RotDmcConfig(
            name="walker_walk",
            suite="dmc",
            observation_size=24,
            action_size=6,
        ),
    ]

    # pylint: disable=unexpected-keyword-arg

    def _info(self) -> tfds.core.DatasetInfo:
        """Returns the dataset metadata."""
        builder_config = self.builder_config
        observation_info = {
            "states": tfds.features.Tensor(
                shape=(builder_config.observation_size,),
                dtype=tf.float32,
                encoding=tfds.features.Encoding.ZLIB,
            ),
            "pixels": tfds.features.Tensor(
                shape=(84, 84, 9),
                dtype=tf.uint8,
                encoding=tfds.features.Encoding.ZLIB,
            ),
        }
        action_info = tfds.features.Tensor(
            shape=(builder_config.action_size,),
            dtype=tf.float32,
            encoding=tfds.features.Encoding.ZLIB,
        )
        ds_config = rlds_base.DatasetConfig(
            name=builder_config.name,
            observation_info=observation_info,
            action_info=action_info,
            reward_info=tfds.features.Tensor(shape=(), dtype=tf.float64),
            citation=_CITATION,
            homepage=_HOMEPAGE,
            overall_description=_DESCRIPTION,
            supervised_keys=None,  # pytype: disable=wrong-arg-types  # gen-stub-imports
        )
        domain, _ = self.builder_config.name.split("_")
        camera_id = dict(quadruped=2).get(domain, 0)
        metadata = {
            "camera_id": camera_id,
            "frame_stack": 3,
            "action_repeat": 2,
        }
        return rlds_base.build_info(ds_config, self, metadata)

    def _split_generators(self, dl_manager: tfds.download.DownloadManager):
        """Returns SplitGenerators."""
        # TODO(rot): Downloads the data and defines the splits
        demo_path = dl_manager.download_and_extract({"file_path": _FILE_URL})[
            "file_path"
        ]
        name = self.builder_config.name

        return {
            "train": self._generate_examples(
                demo_path / "expert_demos" / self.builder_config.suite / name
            ),
        }

    def _generate_examples(self, path):
        """Yields examples."""
        # TODO(rot): Yields (key, example) tuples from the dataset
        for f in path.glob("*.pkl"):
            for episode_id, episode in enumerate(_build_rot_dataset(f)):
                yield episode_id, episode

corax/datasets/tfds/rot/rot_dmc.py
ethanluoyc/corax
"""rot dataset."""

import tensorflow_datasets as tfds

from corax.datasets.tfds.rot import rot_dmc


class RotDmcTest(tfds.testing.DatasetBuilderTestCase):
    """Tests for rot dataset."""

    DATASET_CLASS = rot_dmc.RotDmc

    BUILDER_CONFIG_NAMES_TO_TEST = ["cheetah_run"]

    SPLITS = {
        "train": 1,  # Number of fake train example
    }

    SKIP_CHECKSUMS = True
    SKIP_TF1_GRAPH_MODE = True
    DL_EXTRACT_RESULT = {"file_path": "."}


if __name__ == "__main__":
    tfds.testing.test_main()

corax/datasets/tfds/rot/rot_test.py
ethanluoyc/corax
# type: ignore
# 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.

"""Utilities related to loading TFDS datasets.
CHANGE: Fixes issues with inlining large dataset.
"""

import logging
from typing import Any, Iterator, Optional, Sequence, Tuple

import jax
import jax.numpy as jnp
import numpy as np
import rlds
import tensorflow as tf
import tensorflow_datasets as tfds
from flax import jax_utils

from corax import specs
from corax import types


def _batched_step_to_transition(step: rlds.BatchedStep) -> types.Transition:
    return types.Transition(
        observation=tf.nest.map_structure(lambda x: x[0], step[rlds.OBSERVATION]),
        action=tf.nest.map_structure(lambda x: x[0], step[rlds.ACTION]),
        reward=tf.nest.map_structure(lambda x: x[0], step[rlds.REWARD]),
        discount=1.0 - tf.cast(step[rlds.IS_TERMINAL][1], dtype=tf.float32),
        # If next step is terminal, then the observation may be arbitrary.
        next_observation=tf.nest.map_structure(lambda x: x[1], step[rlds.OBSERVATION]),
    )


def _batch_steps(episode: rlds.Episode) -> tf.data.Dataset:
    return rlds.transformations.batch(
        episode[rlds.STEPS], size=2, shift=1, drop_remainder=True
    )


def _dataset_size_upperbound(dataset: tf.data.Dataset) -> int:
    if dataset.cardinality() != tf.data.experimental.UNKNOWN_CARDINALITY:
        return dataset.cardinality()
    return tf.cast(dataset.batch(1000).reduce(0, lambda x, step: x + 1000), tf.int64)


def load_tfds_dataset(
    dataset_name: str,
    num_episodes: Optional[int] = None,
    env_spec: Optional[specs.EnvironmentSpec] = None,
    download: bool = True,
) -> tf.data.Dataset:
    """Returns a TFDS dataset with the given name."""
    # Used only in tests.
    del env_spec

    dataset = tfds.load(dataset_name, download=download)["train"]
    if num_episodes:
        dataset = dataset.take(num_episodes)
    return dataset


# TODO(sinopalnikov): replace get_ftds_dataset with a pair of load/transform.
def get_tfds_dataset(
    dataset_name: str,
    num_episodes: Optional[int] = None,
    env_spec: Optional[specs.EnvironmentSpec] = None,
) -> tf.data.Dataset:
    """Returns a TFDS dataset transformed to a dataset of transitions."""
    dataset = load_tfds_dataset(dataset_name, num_episodes, env_spec)
    batched_steps = dataset.flat_map(_batch_steps)
    return rlds.transformations.map_steps(batched_steps, _batched_step_to_transition)


# In order to avoid excessive copying on TPU one needs to make the last
# dimension a multiple of this number.
_BEST_DIVISOR = 128


def _pad(x: jnp.ndarray) -> jnp.ndarray:
    if len(x.shape) != 2:
        return x
    # Find a more scientific way to find this threshold (30). Depending on various
    # conditions for low enough sizes the excessive copying is not triggered.
    if x.shape[-1] % _BEST_DIVISOR != 0 and x.shape[-1] > 30:
        n = _BEST_DIVISOR - (x.shape[-1] % _BEST_DIVISOR)
        x = np.pad(x, [(0, 0)] * (x.ndim - 1) + [(0, n)], "constant")
    return x


# Undo the padding.
def _unpad(x: jnp.ndarray, shape: Sequence[int]) -> jnp.ndarray:
    if len(shape) == 2 and x.shape[-1] != shape[-1]:
        return x[..., : shape[-1]]
    return x


_PMAP_AXIS_NAME = "data"


class JaxInMemoryRandomSampleIterator(Iterator[Any]):
    """In memory random sample iterator implemented in JAX.

    Loads the whole dataset in memory and performs random sampling with
    replacement of batches of `batch_size`.
    This class provides much faster sampling functionality compared to using
    an iterator on tf.data.Dataset.
    """

    def __init__(
        self,
        dataset: tf.data.Dataset,
        key: jnp.ndarray,
        batch_size: int,
        shard_dataset_across_devices: bool = False,
    ):
        """Creates an iterator.

        Args:
          dataset: underlying tf Dataset
          key: a key to be used for random number generation
          batch_size: batch size
          shard_dataset_across_devices: whether to use all available devices
            for storing the underlying dataset. The upside is a larger
            dataset capacity that fits into memory. Downsides are:
              - execution of pmapped functions is usually slower than jitted
              - few last elements in the dataset might be dropped (if not multiple)
              - sampling is not 100% uniform, since each core will be doing sampling
                only within its data chunk
            The number of available devices must divide the batch_size evenly.
        """
        # Read the whole dataset. We use artificially large batch_size to make sure
        # we capture the whole dataset.
        size = _dataset_size_upperbound(dataset)
        data = next(dataset.batch(size).as_numpy_iterator())
        self._dataset_size = jax.tree_util.tree_flatten(
            jax.tree_util.tree_map(lambda x: x.shape[0], data)
        )[0][0]
        device = jax_utils._pmap_device_order()
        if not shard_dataset_across_devices:
            device = device[:1]
        should_pmap = len(device) > 1
        assert batch_size % len(device) == 0
        self._dataset_size = self._dataset_size - self._dataset_size % len(device)
        # len(device) needs to divide self._dataset_size evenly.
        assert self._dataset_size % len(device) == 0
        logging.info("Trying to load %s elements to %s", self._dataset_size, device)
        logging.info(
            "Dataset %s %s",
            ("before padding" if should_pmap else ""),
            jax.tree_util.tree_map(lambda x: x.shape, data),
        )
        if should_pmap:
            shapes = jax.tree_util.tree_map(lambda x: x.shape, data)
            # Padding to a multiple of 128 is needed to avoid excessive copying on TPU
            data = jax.tree_util.tree_map(_pad, data)
            logging.info(
                "Dataset after padding %s",
                jax.tree_util.tree_map(lambda x: x.shape, data),
            )

            def split_and_put(x: jnp.ndarray) -> jnp.ndarray:
                return jax.device_put_sharded(
                    np.split(x[: self._dataset_size], len(device)), devices=device
                )

            self._jax_dataset = jax.tree_util.tree_map(split_and_put, data)
        else:
            self._jax_dataset = jax.tree_util.tree_map(jax.device_put, data)

        self._key = (
            jnp.stack(jax.random.split(key, len(device))) if should_pmap else key
        )

        def sample_per_shard(data: Any, key: jnp.ndarray) -> Tuple[Any, jnp.ndarray]:
            key1, key2 = jax.random.split(key)
            indices = jax.random.randint(
                key1,
                (batch_size // len(device),),
                minval=0,
                maxval=self._dataset_size // len(device),
            )
            data_sample = jax.tree_util.tree_map(
                lambda d: jnp.take(d, indices, axis=0), data
            )
            return data_sample, key2

        if should_pmap:

            def sample(data, key):
                data_sample, key = sample_per_shard(data, key)
                # Gathering data on TPUs is much more efficient that doing so on a host
                # since it avoids Host - Device communications.
                data_sample = jax.lax.all_gather(
                    data_sample, axis_name=_PMAP_AXIS_NAME, axis=0, tiled=True
                )
                data_sample = jax.tree_util.tree_map(_unpad, data_sample, shapes)
                return data_sample, key

            pmapped_sample = jax.pmap(sample, axis_name=_PMAP_AXIS_NAME)

            def sample_and_postprocess(
                jax_dataset, key: jnp.ndarray
            ) -> Tuple[Any, jnp.ndarray]:
                data, key = pmapped_sample(jax_dataset, key)
                # All pmapped devices return the same data, so we just take the one from
                # the first device.
                return jax.tree_util.tree_map(lambda x: x[0], data), key

            self._sample = sample_and_postprocess
        else:
            self._sample = jax.jit(sample_per_shard)

    def __next__(self) -> Any:
        data, self._key = self._sample(self._jax_dataset, self._key)
        return data

    @property
    def dataset_size(self) -> int:
        """An integer of the dataset cardinality."""
        return self._dataset_size

corax/datasets/tfds/utils.py
ethanluoyc/corax
from corax.datasets.tfds.vd4rl.vd4rl import Vd4rl

corax/datasets/tfds/vd4rl/__init__.py
ethanluoyc/corax
import collections
import json
import pathlib

from absl import app
from absl import flags

_DATA_ROOT = flags.DEFINE_string("data_root", None, "Data root", required=True)
_MANIFEST_PATH = flags.DEFINE_string("manifest_path", "manifest.json", "Data root")


def nest_dict(dict1):
    result = {}
    for k, v in dict1.items():
        # for each key call method split_rec which
        # will split keys to form recursively
        # nested dictionary
        split_rec(k, v, result)
    return result


def split_rec(k, v, out):
    # splitting keys in dict
    # calling_recursively to break items on '_'
    k, *rest = k.split("/", 1)
    if rest:
        split_rec(rest[0], v, out.setdefault(k, {}))
    else:
        out[k] = v


def unflatten(dictionary):
    resultDict = {}
    for key, value in dictionary.items():
        parts = key.split("/")
        d = resultDict
        for part in parts[:-1]:
            if part not in d:
                d[part] = {}
            d = d[part]
        d[parts[-1]] = value
    return resultDict


def main(_):
    data_root = pathlib.Path(_DATA_ROOT.value)

    paths = [str(p.relative_to(data_root)) for p in data_root.rglob("*") if p.is_file()]

    manifest = collections.defaultdict(list)
    for path in paths:
        parts = path.split("/")
        manifest["/".join(parts[:-1])].append(parts[-1])

    for k in manifest:
        manifest[k] = sorted(manifest[k])

    with open(_MANIFEST_PATH.value, "w") as f:
        json.dump(manifest, f, indent=2)


if __name__ == "__main__":
    app.run(main)

corax/datasets/tfds/vd4rl/create_manifest.py
ethanluoyc/corax
import json

from gdown.download_folder import download_and_parse_google_drive_link  # noqa
from gdown.download_folder import get_directory_structure  # noqa


def nest_dict(dict1):
    result = {}
    for k, v in dict1.items():
        # for each key call method split_rec which
        # will split keys to form recursively
        # nested dictionary
        split_rec(k, v, result)
    return result


def split_rec(k, v, out):
    # splitting keys in dict
    # calling_recursively to break items on '_'
    k, *rest = k.split("/", 1)
    if rest:
        split_rec(rest[0], v, out.setdefault(k, {}))
    else:
        out[k] = v


def main():
    url = "https://drive.google.com/drive/folders/16CmKNZWIK907Cj5J-_zaEy4YuKOTjupK"
    return_code, gdrive_file = download_and_parse_google_drive_link(
        url, remaining_ok=True
    )
    assert return_code
    d = {
        folder: file_id for file_id, folder in get_directory_structure(gdrive_file, "")
    }
    d = {k: v for k, v in d.items() if v is not None and "64px" not in k}
    print(json.dumps(nest_dict(d), indent=4))


if __name__ == "__main__":
    main()

corax/datasets/tfds/vd4rl/generate_vd4rl_file_list.py
ethanluoyc/corax
"""vd4rl dataset."""
import dataclasses
import json
import os
from typing import Any, NamedTuple

import dm_env
import h5py
import numpy as np
import tensorflow_datasets as tfds
import tree

_DESCRIPTION = """
TODO
"""

_CITATION = """
TODO
"""

_HOMEPAGE = "TODO"

_DATASET_DESCRIPTION = """
TODO
"""


@dataclasses.dataclass
class Vd4rlBuilderConfig(tfds.core.BuilderConfig):
    prefix: str = "main"
    env: str = "walker_walk"
    quality: str = "expert"
    resolution: int = 84


QUALITIES = ["expert", "medium", "medium_replay", "medium_expert", "random"]
ENVS = ["walker_walk", "cheetah_run", "humanoid_walk"]


def _get_all_configs():
    configs = []
    for env_name in ENVS:
        for quality in QUALITIES:
            # pylint: disable=unexpected-keyword-arg
            configs.append(
                Vd4rlBuilderConfig(
                    name=f"main_{env_name}_{quality}_84px",
                    prefix="main",
                    env=env_name,
                    quality=quality,
                    resolution=84,
                )
            )
            # pylint: enable=unexpected-keyword-arg
    return configs


_BUILDER_CONFIGS = _get_all_configs()


class Vd4rl(tfds.core.GeneratorBasedBuilder):
    """DatasetBuilder for exorl dataset."""

    VERSION = tfds.core.Version("1.0.0")
    RELEASE_NOTES = {
        "1.0.0": "Initial release.",
    }

    # pylint: disable=unexpected-keyword-arg
    BUILDER_CONFIGS = _BUILDER_CONFIGS

    # pylint: enable=unexpected-keyword-arg

    def _info(self) -> tfds.core.DatasetInfo:
        """Returns the dataset metadata."""
        resolution = self.builder_config.resolution
        action_dim = {"walker_walk": 6, "humanoid_walk": 21, "cheetah_run": 6}[
            self.builder_config.env
        ]
        steps_dict = {
            "observation": tfds.features.Image(
                shape=(resolution, resolution, 3), encoding_format="png", dtype=np.uint8
            ),
            "action": tfds.features.Tensor(
                shape=(action_dim,),
                encoding=tfds.features.Encoding.ZLIB,
                dtype=np.float32,
            ),
            "reward": np.float32,
            "discount": np.float32,
            "is_terminal": np.bool_,
            "is_first": np.bool_,
            "is_last": np.bool_,
        }
        return tfds.core.DatasetInfo(
            builder=self,
            description=_DESCRIPTION,
            features=tfds.features.FeaturesDict(
                {
                    # These are the features of your dataset like images, labels ...
                    # 'episode_id': tf.int64,
                    "steps": tfds.features.Dataset(steps_dict),
                }
            ),
            # If there's a common (input, target) tuple from the
            # features, specify them here. They'll be used if
            # `as_supervised=True` in `builder.as_dataset`.
            supervised_keys=None,  # disabled
            homepage=_HOMEPAGE,
            citation=_CITATION,
        )

    def _split_generators(self, dl_manager: tfds.download.DownloadManager):
        """Returns SplitGenerators."""
        manifest_path = os.path.join(
            os.path.abspath(os.path.dirname(__file__)), "manifest.json"
        )
        with open(manifest_path, encoding="utf-8") as f:
            manifest = json.load(f)
        prefix = self.builder_config.prefix
        env_name = self.builder_config.env
        quality = self._builder_config.quality
        resolution = self._builder_config.resolution
        key = f"{prefix}/{env_name}/{quality}/{resolution}px"
        filenames = manifest[key]
        files = {
            fname: f"https://huggingface.co/datasets/conglu/vd4rl/resolve/main/vd4rl/{key}/{fname}"
            for fname in filenames
        }
        downloaded_files = dl_manager.download(files)
        filenames = sorted(downloaded_files.keys())
        files = [downloaded_files[k] for k in filenames]
        return {
            "train": self._generate_examples(files),
        }

    def _generate_examples(self, paths):
        """Yields examples."""
        current_episode = None

        counter = 0
        for step_data in _load_trajectories(paths):
            # action is the action that generated the timestep
            if step_data.timestep.first():
                if current_episode is not None:
                    # episodes.append(current_episode)
                    rlds_steps = convert_to_rlds_steps(current_episode)
                    rlds_steps = tree.map_structure(
                        lambda *xs: np.stack(xs), *rlds_steps
                    )
                    yield counter, {"steps": rlds_steps}
                    counter += 1
                    current_episode = None

            if current_episode is None:
                current_episode = [step_data]
            else:
                current_episode.append(step_data)

        if current_episode is not None:
            rlds_steps = convert_to_rlds_steps(current_episode)
            rlds_steps = tree.map_structure(lambda *xs: np.stack(xs), *rlds_steps)
            yield counter, {"steps": rlds_steps}
            current_episode = None


step_type_lookup = {
    0: dm_env.StepType.FIRST,
    1: dm_env.StepType.MID,
    2: dm_env.StepType.LAST,
}


def _get_timestep_from_idx(offline_data: dict, idx: int):
    # Convert to channel last
    observation = np.transpose(offline_data["observation"][idx], [1, 2, 0])
    timestep = dm_env.TimeStep(
        step_type=step_type_lookup[offline_data["step_type"][idx]],
        reward=offline_data["reward"][idx].astype(np.float32),
        observation=observation,
        discount=offline_data["discount"][idx].astype(np.float32),
    )
    action = offline_data["action"][idx]
    return timestep, action


class _StepData(NamedTuple):
    timestep: dm_env.TimeStep
    action: Any = None


def _load_trajectories(filenames):
    for filename in filenames:
        episodes = h5py.File(filename, "r")
        episodes = {k: v[:] for k, v in episodes.items()}
        offline_data = episodes
        offline_data_length = offline_data["reward"].shape[0]
        for v in offline_data.values():
            assert v.shape[0] == offline_data_length
        for idx in range(offline_data_length):
            time_step, action = _get_timestep_from_idx(offline_data, idx)
            yield _StepData(time_step, action)


## https://github.com/deepmind/envlogger/blob/main/envlogger/backends/rlds_utils.py
def to_rlds_step(prev_step, step: None):
    """Builds an RLDS step from two Envlogger steps.
    Steps follow the RLDS convention from https://github.com/google-research/rlds.
    Args:
      prev_step: previous step.
      step: current step. If None, it builds the last step (where the observation
        is the last one, and the action, reward and discount are undefined).
    Returns:
       RLDS Step.
    """
    metadata = {}
    return {
        "action": step.action
        if step
        else tree.map_structure(np.zeros_like, prev_step.action),
        "discount": step.timestep.discount
        if step
        else tree.map_structure(np.zeros_like, prev_step.timestep.discount),
        "is_first": prev_step.timestep.first(),
        "is_last": prev_step.timestep.last(),
        "is_terminal": (
            prev_step.timestep.last() and prev_step.timestep.discount == 0.0
        ),
        "observation": prev_step.timestep.observation,
        "reward": step.timestep.reward
        if step
        else tree.map_structure(np.zeros_like, prev_step.timestep.reward),
        **metadata,
    }


def convert_to_rlds_steps(episode):
    prev_steps = episode[:-1]
    steps = episode[1:]
    rlds_steps = []
    for prev_step, step in zip(prev_steps, steps):
        rlds_steps.append(to_rlds_step(prev_step, step))
    rlds_steps.append(to_rlds_step(episode[-1], None))
    return rlds_steps

corax/datasets/tfds/vd4rl/vd4rl.py
ethanluoyc/corax
import numpy as np
import rlds
import tensorflow as tf
import tree


def stack_observations_np(observations, stack_size):
    def _repeat_first(sequence):
        reps = [1] * (sequence[0].ndim + 1)
        reps[0] = stack_size - 1
        repeated_first_step = np.tile(sequence[0], reps)
        return np.concatenate(
            [repeated_first_step, sequence], 0
        )  # [:-(stack_size - 1)]

    def _stack_observation(observation):
        stack = [np.roll(observation, i, axis=0) for i in range(stack_size)]
        stack.reverse()  # chronological order
        return np.concatenate(stack, axis=-1)

    observations = _repeat_first(observations)
    return _stack_observation(observations)[(stack_size - 1) :]


def get_n_step_transitions(episode_steps, nsteps, discount):
    """Get n_step transitions from an episode of steps"""
    episode_length = episode_steps["reward"].shape[0] - 1
    transitions = []

    for t in range(0, episode_length - nsteps + 1):
        observations = episode_steps["observation"][t : t + nsteps + 1]
        actions = episode_steps["action"][t : t + nsteps + 1]
        rewards = episode_steps["reward"][t : t + nsteps + 1]
        discounts = episode_steps["discount"][t : t + nsteps + 1]

        total_reward = rewards[0]
        total_discount = discounts[0]

        for n in range(1, len(discounts) - 1):
            total_discount *= discount
            total_reward += rewards[n] * total_discount
            total_discount *= discounts[n]

        transitions.append(
            {
                "observation": observations[0],
                "next_observation": observations[-1],
                "reward": total_reward,
                "discount": total_discount,
                "action": actions[0],
            }
        )
    # This can be faster
    return tree.map_structure(lambda *xs: np.stack(xs), *transitions)


def stack_observation_tfds(observations: tf.data.Dataset, stack_size: int):
    first_obs = observations.take(1)
    observations = first_obs.repeat(stack_size - 1).concatenate(observations)
    data = observations.window(stack_size, shift=1, drop_remainder=True).flat_map(
        lambda x: x.batch(stack_size, drop_remainder=True)
    )
    return data.map(
        lambda x: tf.concat(tf.unstack(x, axis=0), axis=-1),
        num_parallel_calls=tf.data.AUTOTUNE,
    )


def process_data(steps_dataset, stack_size):
    obs = steps_dataset.map(
        lambda step: step["observation"], num_parallel_calls=tf.data.AUTOTUNE
    )
    stacked_obs = stack_observation_tfds(obs, stack_size)
    zipped_ds = tf.data.Dataset.zip((stacked_obs, steps_dataset))

    def combine(stacked_obs, steps):
        steps["observation"] = stacked_obs
        return steps

    return zipped_ds.map(combine, num_parallel_calls=tf.data.AUTOTUNE)


@tf.function
def _sequence_to_transitions(sequence, discount):
    total_reward = sequence["reward"][0]
    total_discount = sequence["discount"][0]
    n_steps = len(sequence["reward"]) - 1

    for t in range(1, n_steps):
        total_discount *= discount
        total_reward += total_discount * sequence["reward"][t]
        total_discount *= sequence["discount"][t]

    return {
        "observation": sequence["observation"][0],
        "next_observation": sequence["observation"][-1],
        "action": sequence["action"][0],
        "reward": total_reward,
        "discount": total_discount,
    }


def tfds_get_n_step_transitions(steps_dataset, nsteps, discount):
    """Get n_step transitions from an episode of steps"""
    batched_steps = rlds.transformations.batch(
        steps_dataset, size=nsteps + 1, shift=1, drop_remainder=True
    )
    return batched_steps.map(
        lambda seq: _sequence_to_transitions(seq, discount),
        num_parallel_calls=tf.data.experimental.AUTOTUNE,
    )

corax/datasets/tfds/vd4rl/vd4rl_preprocessor.py
ethanluoyc/corax
import numpy as np
import tensorflow as tf
from absl.testing import absltest
from absl.testing import parameterized

from corax.datasets.tfds.vd4rl import vd4rl_preprocessor

_TRANSITION_TEST_CASES = [
    {
        "testcase_name": "test_normal",
        "steps": {
            "observation": np.arange(4, dtype=np.float32),
            "reward": np.array([1, 2, 3, 4], dtype=np.float32),
            "discount": np.ones(4, dtype=np.float32),
            "action": np.arange(4, dtype=np.float32),
        },
        "expected_transitions": {
            "observation": np.array([0, 1]),
            "next_observation": np.array([2, 3]),
            "discount": np.array([0.99, 0.99]),
            "action": np.array([0, 1]),
            "reward": np.array([1 + 2 * 0.99, 2 + 3 * 0.99]),
        },
        "n_steps": 2,
        "discount": 0.99,
    },
    {
        "testcase_name": "test_normal_2",
        "steps": {
            "observation": np.array([0, 1, 2, 3, 4, 5], dtype=np.float32),
            "reward": np.array([1, 2, 3, 4, 5, 6], dtype=np.float32),
            "discount": np.array([1, 1, 1, 1, 1, 1], dtype=np.float32),
            "action": np.array([0, 1, 2, 3, 4, 5], dtype=np.float32),
        },
        "expected_transitions": {
            "observation": np.array([0, 1, 2]),
            "next_observation": np.array([3, 4, 5]),
            "discount": np.array([0.99 * 0.99, 0.99 * 0.99, 0.99 * 0.99]),
            "action": np.array([0, 1, 2]),
            "reward": np.array(
                [
                    1 + 2 * 0.99 + 3 * 0.99 * 0.99,
                    2 + 3 * 0.99 + 4 * 0.99 * 0.99,
                    3 + 4 * 0.99 + 5 * 0.99 * 0.99,
                ]
            ),
        },
        "n_steps": 3,
        "discount": 0.99,
    },
    {
        "testcase_name": "test_terminal",
        "steps": {
            "observation": np.array([0, 1, 2, 3, 4, 5], dtype=np.float32),
            "reward": np.array([1, 2, 3, 4, 5, 6], dtype=np.float32),
            "discount": np.array([1, 1, 1, 1, 0, 0], dtype=np.float32),
            "action": np.array([0, 1, 2, 3, 4, 5], dtype=np.float32),
        },
        "expected_transitions": {
            "observation": np.array([0, 1, 2]),
            "next_observation": np.array([3, 4, 5]),
            "discount": np.array([0.99 * 0.99, 0.99 * 0.99, 0]),
            "action": np.array([0, 1, 2]),
            "reward": np.array(
                [
                    1 + 2 * 0.99 + 3 * 0.99 * 0.99,
                    2 + 3 * 0.99 + 4 * 0.99 * 0.99,
                    # 3 + 4 * 0.99 + 5 * 0.99 * 0.99,
                    3 + 4 * 0.99 + 5 * 0.99 * 0.99,
                ]
            ),
        },
        "n_steps": 3,
        "discount": 0.99,
    },
    {
        "testcase_name": "test_terminal_2",
        "steps": {
            "observation": np.array([1, 2, 3, 4, 5], dtype=np.float32),
            "reward": np.array([2, 3, 5, 7, 0], dtype=np.float32),
            "discount": np.array([1, 1, 1, 0, 0], dtype=np.float32),
            "action": np.array([0, 0, 0, 0, 0], dtype=np.float32),
        },
        "expected_transitions": {
            "observation": np.array([1, 2]),
            "next_observation": np.array([4, 5]),
            "discount": np.array([0.5 * 0.5, 0]),
            "action": np.array([0, 0]),
            "reward": np.array(
                [
                    2 + 0.5 * 3 + 0.25 * 5,
                    3 + 0.5 * 5 + 0.25 * 7,
                ]
            ),
        },
        "n_steps": 3,
        "discount": 0.5,
    },
    {
        "testcase_name": "test_one_step",
        "steps": {
            "observation": np.array([1, 2, 3], dtype=np.float32),
            "reward": np.array([2, 3, -1], dtype=np.float32),
            "discount": np.array([1, 0, 0], dtype=np.float32),
            "action": np.array([0, 1, 0], dtype=np.float32),
        },
        "expected_transitions": {
            "observation": np.array([1, 2]),
            "next_observation": np.array([2, 3]),
            "discount": np.array([1, 0]),
            "action": np.array([0, 1]),
            "reward": np.array([2, 3]),
        },
        "n_steps": 1,
        "discount": 0.5,
    },
]


def _assert_transitions_equal(np_transitions, expected_transitions):
    np.testing.assert_allclose(
        np_transitions["observation"], expected_transitions["observation"]
    )
    np.testing.assert_allclose(
        np_transitions["next_observation"], expected_transitions["next_observation"]
    )
    np.testing.assert_allclose(np_transitions["reward"], expected_transitions["reward"])
    np.testing.assert_allclose(
        np_transitions["discount"], expected_transitions["discount"]
    )
    np.testing.assert_allclose(np_transitions["action"], expected_transitions["action"])


class VD4RLDatasetTest(parameterized.TestCase):
    def test_observation_stacking_np(self):
        test_obs = np.arange(4).reshape((4, 1))
        stacked_obs = vd4rl_preprocessor.stack_observations_np(test_obs, 2)
        expected_stacked_obs = np.array(
            [
                [0, 0],
                [0, 1],
                [1, 2],
                [2, 3],
            ]
        )
        np.testing.assert_equal(stacked_obs, expected_stacked_obs)

    def test_observation_stacking_tfds(self):
        test_obs = np.arange(4).reshape((4, 1))
        stacked_obs = next(
            vd4rl_preprocessor.stack_observation_tfds(
                tf.data.Dataset.from_tensor_slices(test_obs), 2
            )
            .batch(100)
            .as_numpy_iterator()
        )
        expected_stacked_obs = np.array(
            [
                [0, 0],
                [0, 1],
                [1, 2],
                [2, 3],
            ]
        )
        np.testing.assert_equal(stacked_obs, expected_stacked_obs)

    @parameterized.named_parameters(*_TRANSITION_TEST_CASES)
    def test_create_transitions_np(
        self, n_steps, steps, discount, expected_transitions
    ):
        np_transitions = vd4rl_preprocessor.get_n_step_transitions(
            steps, n_steps, discount
        )
        _assert_transitions_equal(np_transitions, expected_transitions)

    @parameterized.named_parameters(*_TRANSITION_TEST_CASES)
    def test_create_transitions_tfds(
        self, n_steps, steps, discount, expected_transitions
    ):
        steps_dataset = tf.data.Dataset.from_tensor_slices(steps)
        transition_ds = vd4rl_preprocessor.tfds_get_n_step_transitions(
            steps_dataset, n_steps, discount
        )
        tf_transitions = next(transition_ds.batch(100).as_numpy_iterator())
        _assert_transitions_equal(tf_transitions, expected_transitions)


if __name__ == "__main__":
    tf.config.set_visible_devices([], "GPU")
    absltest.main()

corax/datasets/tfds/vd4rl/vd4rl_preprocessor_test.py
ethanluoyc/corax
# 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.

"""A simple agent-environment training loop."""

import operator
import time
from typing import List, Optional, Sequence

import dm_env
import numpy as np
import tree
from dm_env import specs

from corax import core
from corax.utils import counting
from corax.utils import loggers
from corax.utils import observers as observers_lib


class EnvironmentLoop(core.Worker):
    """A simple RL environment loop.

    This takes `Environment` and `Actor` instances and coordinates their
    interaction. Agent is updated if `should_update=True`. This can be used as:

      loop = EnvironmentLoop(environment, actor)
      loop.run(num_episodes)

    A `Counter` instance can optionally be given in order to maintain counts
    between different Acme components. If not given a local Counter will be
    created to maintain counts between calls to the `run` method.

    A `Logger` instance can also be passed in order to control the output of the
    loop. If not given a platform-specific default logger will be used as defined
    by utils.loggers.make_default_logger. A string `label` can be passed to easily
    change the label associated with the default logger; this is ignored if a
    `Logger` instance is given.

    A list of 'Observer' instances can be specified to generate additional metrics
    to be logged by the logger. They have access to the 'Environment' instance,
    the current timestep datastruct and the current action.
    """

    def __init__(
        self,
        environment: dm_env.Environment,
        actor: core.Actor,
        counter: Optional[counting.Counter] = None,
        logger: Optional[loggers.Logger] = None,
        should_update: bool = True,
        label: str = "environment_loop",
        observers: Sequence[observers_lib.EnvLoopObserver] = (),
    ):
        # Internalize agent and environment.
        self._environment = environment
        self._actor = actor
        self._counter = counter or counting.Counter()
        self._logger = logger or loggers.make_default_logger(
            label, steps_key=self._counter.get_steps_key()
        )
        self._should_update = should_update
        self._observers = observers

    def run_episode(self) -> loggers.LoggingData:
        """Run one episode.

        Each episode is a loop which interacts first with the environment to get an
        observation and then give that observation to the agent in order to retrieve
        an action.

        Returns:
          An instance of `loggers.LoggingData`.
        """
        # Reset any counts and start the environment.
        episode_start_time = time.time()
        select_action_durations: List[float] = []
        env_step_durations: List[float] = []
        episode_steps: int = 0

        # For evaluation, this keeps track of the total undiscounted reward
        # accumulated during the episode.
        episode_return = tree.map_structure(
            _generate_zeros_from_spec, self._environment.reward_spec()
        )
        env_reset_start = time.time()
        timestep = self._environment.reset()
        env_reset_duration = time.time() - env_reset_start
        # Make the first observation.
        self._actor.observe_first(timestep)
        for observer in self._observers:
            # Initialize the observer with the current state of the env after reset
            # and the initial timestep.
            observer.observe_first(self._environment, timestep)

        # Run an episode.
        while not timestep.last():
            # Book-keeping.
            episode_steps += 1

            # Generate an action from the agent's policy.
            select_action_start = time.time()
            action = self._actor.select_action(timestep.observation)
            select_action_durations.append(time.time() - select_action_start)

            # Step the environment with the agent's selected action.
            env_step_start = time.time()
            timestep = self._environment.step(action)
            env_step_durations.append(time.time() - env_step_start)

            # Have the agent and observers observe the timestep.
            self._actor.observe(action, next_timestep=timestep)
            for observer in self._observers:
                # One environment step was completed. Observe the current state of the
                # environment, the current timestep and the action.
                observer.observe(self._environment, timestep, action)

            # Give the actor the opportunity to update itself.
            if self._should_update:
                self._actor.update()

            # Equivalent to: episode_return += timestep.reward
            # We capture the return value because if timestep.reward is a JAX
            # DeviceArray, episode_return will not be mutated in-place. (In all other
            # cases, the returned episode_return will be the same object as the
            # argument episode_return.)
            episode_return = tree.map_structure(
                operator.iadd, episode_return, timestep.reward
            )

        # Record counts.
        counts = self._counter.increment(episodes=1, steps=episode_steps)

        # Collect the results and combine with counts.
        steps_per_second = episode_steps / (time.time() - episode_start_time)
        result = {
            "episode_length": episode_steps,
            "episode_return": episode_return,
            "steps_per_second": steps_per_second,
            "env_reset_duration_sec": env_reset_duration,
            "select_action_duration_sec": np.mean(select_action_durations),
            "env_step_duration_sec": np.mean(env_step_durations),
        }
        result.update(counts)
        for observer in self._observers:
            result.update(observer.get_metrics())
        return result

    def run(
        self,
        num_episodes: Optional[int] = None,
        num_steps: Optional[int] = None,
    ) -> int:
        """Perform the run loop.

        Run the environment loop either for `num_episodes` episodes or for at
        least `num_steps` steps (the last episode is always run until completion,
        so the total number of steps may be slightly more than `num_steps`).
        At least one of these two arguments has to be None.

        Upon termination of an episode a new episode will be started. If the number
        of episodes and the number of steps are not given then this will interact
        with the environment infinitely.

        Args:
          num_episodes: number of episodes to run the loop for.
          num_steps: minimal number of steps to run the loop for.

        Returns:
          Actual number of steps the loop executed.

        Raises:
          ValueError: If both 'num_episodes' and 'num_steps' are not None.
        """

        if not (num_episodes is None or num_steps is None):
            raise ValueError('Either "num_episodes" or "num_steps" should be None.')

        def should_terminate(episode_count: int, step_count: int) -> bool:
            return (num_episodes is not None and episode_count >= num_episodes) or (
                num_steps is not None and step_count >= num_steps
            )

        episode_count: int = 0
        step_count: int = 0
        # TODO(yl): Add signal handler when running with launchpad
        while not should_terminate(episode_count, step_count):
            episode_start = time.time()
            result = self.run_episode()
            result = {**result, **{"episode_duration": time.time() - episode_start}}
            episode_count += 1
            step_count += int(result["episode_length"])
            # Log the given episode results.
            self._logger.write(result)

        return step_count


def _generate_zeros_from_spec(spec: specs.Array) -> np.ndarray:
    return np.zeros(spec.shape, spec.dtype)

corax/environment_loop.py
ethanluoyc/corax
# type: ignore
# 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.

"""Tests for the environment loop."""

from typing import Optional

import numpy as np
from absl.testing import absltest
from absl.testing import parameterized

from corax import environment_loop
from corax import specs
from corax import types
from corax.testing import fakes

EPISODE_LENGTH = 10

# Discount specs
F32_2_MIN_0_MAX_1 = specs.BoundedArray(
    dtype=np.float32, shape=(2,), minimum=0.0, maximum=1.0
)
F32_2x1_MIN_0_MAX_1 = specs.BoundedArray(
    dtype=np.float32, shape=(2, 1), minimum=0.0, maximum=1.0
)
TREE_MIN_0_MAX_1 = {"a": F32_2_MIN_0_MAX_1, "b": F32_2x1_MIN_0_MAX_1}

# Reward specs
F32 = specs.Array(dtype=np.float32, shape=())
F32_1x3 = specs.Array(dtype=np.float32, shape=(1, 3))
TREE = {"a": F32, "b": F32_1x3}

TEST_CASES = (
    ("scalar_discount_scalar_reward", None, None),
    ("vector_discount_scalar_reward", F32_2_MIN_0_MAX_1, F32),
    ("matrix_discount_matrix_reward", F32_2x1_MIN_0_MAX_1, F32_1x3),
    ("tree_discount_tree_reward", TREE_MIN_0_MAX_1, TREE),
)


class EnvironmentLoopTest(parameterized.TestCase):
    @parameterized.named_parameters(*TEST_CASES)
    def test_one_episode(self, discount_spec, reward_spec):
        _, loop = _parameterized_setup(discount_spec, reward_spec)
        result = loop.run_episode()
        self.assertIn("episode_length", result)
        self.assertEqual(EPISODE_LENGTH, result["episode_length"])
        self.assertIn("episode_return", result)
        self.assertIn("steps_per_second", result)

    @parameterized.named_parameters(*TEST_CASES)
    def test_run_episodes(self, discount_spec, reward_spec):
        actor, loop = _parameterized_setup(discount_spec, reward_spec)

        # Run the loop. There should be EPISODE_LENGTH update calls per episode.
        loop.run(num_episodes=10)
        self.assertEqual(actor.num_updates, 10 * EPISODE_LENGTH)

    @parameterized.named_parameters(*TEST_CASES)
    def test_run_steps(self, discount_spec, reward_spec):
        actor, loop = _parameterized_setup(discount_spec, reward_spec)

        # Run the loop. This will run 2 episodes so that total number of steps is
        # at least 15.
        loop.run(num_steps=EPISODE_LENGTH + 5)
        self.assertEqual(actor.num_updates, 2 * EPISODE_LENGTH)


def _parameterized_setup(
    discount_spec: Optional[types.NestedSpec] = None,
    reward_spec: Optional[types.NestedSpec] = None,
):
    """Common setup code that, unlike self.setUp, takes arguments.

    Args:
      discount_spec: None, or a (nested) specs.BoundedArray.
      reward_spec: None, or a (nested) specs.Array.
    Returns:
      environment, actor, loop
    """
    env_kwargs = {"episode_length": EPISODE_LENGTH}
    if discount_spec:
        env_kwargs["discount_spec"] = discount_spec
    if reward_spec:
        env_kwargs["reward_spec"] = reward_spec

    environment = fakes.DiscreteEnvironment(**env_kwargs)
    actor = fakes.Actor(specs.make_environment_spec(environment))
    loop = environment_loop.EnvironmentLoop(environment, actor)
    return actor, loop


if __name__ == "__main__":
    absltest.main()

corax/environment_loop_test.py
ethanluoyc/corax
"""JAX experiment utils."""

from corax.jax.experiments.config import CheckpointingConfig
from corax.jax.experiments.config import DeprecatedPolicyFactory
from corax.jax.experiments.config import EvaluatorFactory
from corax.jax.experiments.config import ExperimentConfig
from corax.jax.experiments.config import MakeActorFn
from corax.jax.experiments.config import NetworkFactory
from corax.jax.experiments.config import OfflineExperimentConfig
from corax.jax.experiments.config import PolicyFactory
from corax.jax.experiments.config import default_evaluator_factory
from corax.jax.experiments.config import make_policy
from corax.jax.experiments.imitation_experiment import ImitationExperimentConfig
from corax.jax.experiments.imitation_experiment import run_imitation_experiment
from corax.jax.experiments.imitation_loop import EpisodeRewarder
from corax.jax.experiments.imitation_loop import ImitationEnvironmentLoop
from corax.jax.experiments.run_experiment import run_experiment
from corax.jax.experiments.run_offline_experiment import run_offline_experiment

corax/jax/experiments/__init__.py
ethanluoyc/corax
# type: ignore
# 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.

"""JAX experiment config."""

import dataclasses
import datetime
from typing import Any, Callable, Generic, Iterator, Optional, Sequence

import jax
from typing_extensions import Protocol

from corax import core
from corax import environment_loop
from corax import specs
from corax.agents.jax import builders
from corax.jax import types
from corax.jax import utils
from corax.utils import counting
from corax.utils import loggers
from corax.utils import observers as observers_lib


class MakeActorFn(Protocol, Generic[builders.Policy]):
    def __call__(
        self,
        random_key: types.PRNGKey,
        policy: builders.Policy,
        environment_spec: specs.EnvironmentSpec,
        variable_source: core.VariableSource,
    ) -> core.Actor:
        ...


class NetworkFactory(Protocol, Generic[builders.Networks]):
    def __call__(self, environment_spec: specs.EnvironmentSpec) -> builders.Networks:
        ...


class DeprecatedPolicyFactory(Protocol, Generic[builders.Networks, builders.Policy]):
    def __call__(self, networks: builders.Networks) -> builders.Policy:
        ...


class PolicyFactory(Protocol, Generic[builders.Networks, builders.Policy]):
    def __call__(
        self,
        networks: builders.Networks,
        environment_spec: specs.EnvironmentSpec,
        evaluation: bool,
    ) -> builders.Policy:
        ...


class EvaluatorFactory(Protocol, Generic[builders.Policy]):
    def __call__(
        self,
        random_key: types.PRNGKey,
        variable_source: core.VariableSource,
        counter: counting.Counter,
        make_actor_fn: MakeActorFn[builders.Policy],
    ) -> core.Worker:
        ...


@dataclasses.dataclass(frozen=True)
class CheckpointingConfig:
    """Configuration options for checkpointing.

    Attributes:
      max_to_keep: Maximum number of checkpoints to keep. Unless preserved by
        keep_checkpoint_every_n_hours, checkpoints will be deleted from the active
        set, oldest first, until only max_to_keep checkpoints remain. Does not
        apply to replay checkpointing.
      directory: Where to store the checkpoints.
      add_uid: Whether or not to add a unique identifier, see
        `paths.get_unique_id()` for how it is generated.
      time_delta_minutes: How often to save the checkpoint, in minutes.
      keep_checkpoint_every_n_hours: Upon removal from the active set, a
        checkpoint will be preserved if it has been at least
        keep_checkpoint_every_n_hours since the last preserved checkpoint. The
        default setting of None does not preserve any checkpoints in this way.
      replay_checkpointing_time_delta_minutes: How frequently to write replay
        checkpoints; defaults to None, which disables periodic checkpointing.
        Warning! These are written asynchronously so as not to interrupt other
        replay duties, however this does pose a risk of OOM since items that would
        otherwise be removed are temporarily kept alive for checkpointing
        purposes.
        Note: Since replay buffers tend to be quite large O(100GiB), writing can
          take up to 10 minutes so keep that in mind when setting this frequency.
      checkpoint_ttl_seconds: TTL (time to leave) in seconds for checkpoints.
        Indefinite if set to None.
    """

    max_to_keep: int = 1
    directory: str = "~/acme"
    add_uid: bool = True
    time_delta_minutes: int = 5
    keep_checkpoint_every_n_hours: Optional[int] = None
    replay_checkpointing_time_delta_minutes: Optional[int] = None
    checkpoint_ttl_seconds: Optional[int] = int(
        datetime.timedelta(days=5).total_seconds()
    )


@dataclasses.dataclass(frozen=True)
class ExperimentConfig(Generic[builders.Networks, builders.Policy, builders.Sample]):
    """Config which defines aspects of constructing an experiment.

    Attributes:
      builder: Builds components of an RL agent (Learner, Actor...).
      network_factory: Builds networks used by the agent.
      environment_factory: Returns an instance of an environment.
      max_num_actor_steps: How many environment steps to perform.
      seed: Seed used for agent initialization.
      policy_network_factory: Policy network factory which is used actors to
        perform inference.
      evaluator_factories: Factories of policy evaluators. When not specified the
        default evaluators are constructed using eval_policy_network_factory. Set
        to an empty list to disable evaluators.
      eval_policy_network_factory: Policy network factory used by evaluators.
        Should be specified to use the default evaluators (when
        evaluator_factories is not provided).
      environment_spec: Specification of the environment. Can be specified to
        reduce the number of times environment_factory is invoked (for performance
        or resource usage reasons).
      observers: Observers used for extending logs with custom information.
      logger_factory: Loggers factory used to construct loggers for learner,
        actors and evaluators.
      checkpointing: Configuration options for checkpointing. If None,
        checkpointing and snapshotting is disabled.
    """

    # Below fields must be explicitly specified for any Agent.
    builder: builders.ActorLearnerBuilder[
        builders.Networks, builders.Policy, builders.Sample
    ]
    network_factory: NetworkFactory[builders.Networks]
    environment_factory: types.EnvironmentFactory
    max_num_actor_steps: int
    seed: int
    logger_factory: loggers.LoggerFactory
    # policy_network_factory is deprecated. Use builder.make_policy to
    # create the policy.
    policy_network_factory: Optional[
        DeprecatedPolicyFactory[builders.Networks, builders.Policy]
    ] = None
    # Fields below are optional. If you just started with Acme do not worry about
    # them. You might need them later when you want to customize your RL agent.
    # TODO(stanczyk): Introduce a marker for the default value (instead of None).
    evaluator_factories: Optional[Sequence[EvaluatorFactory[builders.Policy]]] = None
    # eval_policy_network_factory is deprecated. Use builder.make_policy to
    # create the policy.
    eval_policy_network_factory: Optional[
        DeprecatedPolicyFactory[builders.Networks, builders.Policy]
    ] = None
    environment_spec: Optional[specs.EnvironmentSpec] = None
    observers: Sequence[observers_lib.EnvLoopObserver] = ()
    checkpointing: Optional[CheckpointingConfig] = CheckpointingConfig()

    # TODO(stanczyk): Make get_evaluator_factories a standalone function.
    def get_evaluator_factories(self):
        """Constructs the evaluator factories."""
        if self.evaluator_factories is not None:
            return self.evaluator_factories

        def eval_policy_factory(
            networks: builders.Networks,
            environment_spec: specs.EnvironmentSpec,
            evaluation: bool,
        ) -> builders.Policy:
            del evaluation
            # The config factory has precedence until all agents are migrated to use
            # builder.make_policy
            if self.eval_policy_network_factory is not None:
                return self.eval_policy_network_factory(networks)
            else:
                return self.builder.make_policy(
                    networks=networks,
                    environment_spec=environment_spec,
                    evaluation=True,
                )

        return [
            default_evaluator_factory(
                environment_factory=self.environment_factory,
                network_factory=self.network_factory,
                policy_factory=eval_policy_factory,
                logger_factory=self.logger_factory,
                observers=self.observers,
            )
        ]


@dataclasses.dataclass
class OfflineExperimentConfig(
    Generic[builders.Networks, builders.Policy, builders.Sample]
):
    """Config which defines aspects of constructing an offline RL experiment.

    This class is similar to the ExperimentConfig, but is tailored to offline RL
    setting, so it excludes attributes related to training via interaction with
    the environment (max_num_actor_steps, policy_network_factory) and instead
    includes attributes specific to learning from demonstration.

    Attributes:
      builder: Builds components of an offline RL agent (Learner and Evaluator).
      network_factory: Builds networks used by the agent.
      demonstration_dataset_factory: Function that returns an iterator over
        demonstrations.
      environment_spec: Specification of the environment.
      max_num_learner_steps: How many learner steps to perform.
      seed: Seed used for agent initialization.
      evaluator_factories: Factories of policy evaluators. When not specified the
        default evaluators are constructed using eval_policy_network_factory. Set
        to an empty list to disable evaluators.
      eval_policy_factory: Policy factory used by evaluators. Should be specified
        to use the default evaluators (when evaluator_factories is not provided).
      environment_factory: Returns an instance of an environment to be used for
        evaluation. Should be specified to use the default evaluators (when
        evaluator_factories is not provided).
      observers: Observers used for extending logs with custom information.
      logger_factory: Loggers factory used to construct loggers for learner,
        actors and evaluators.
      checkpointing: Configuration options for checkpointing. If None,
        checkpointing and snapshotting is disabled.
    """

    # Below fields must be explicitly specified for any Agent.
    builder: builders.OfflineBuilder[
        builders.Networks, builders.Policy, builders.Sample
    ]
    network_factory: Callable[[specs.EnvironmentSpec], builders.Networks]
    demonstration_dataset_factory: Callable[[types.PRNGKey], Iterator[builders.Sample]]
    environment_factory: types.EnvironmentFactory
    logger_factory: loggers.LoggerFactory
    max_num_learner_steps: int
    seed: int
    # Fields below are optional. If you just started with Acme do not worry about
    # them. You might need them later when you want to customize your RL agent.
    # TODO(stanczyk): Introduce a marker for the default value (instead of None).
    evaluator_factories: Optional[Sequence[EvaluatorFactory]] = None
    environment_spec: Optional[specs.EnvironmentSpec] = None
    observers: Sequence[observers_lib.EnvLoopObserver] = ()
    checkpointing: Optional[CheckpointingConfig] = CheckpointingConfig()

    # TODO(stanczyk): Make get_evaluator_factories a standalone function.
    def get_evaluator_factories(self):
        """Constructs the evaluator factories."""
        if self.evaluator_factories is not None:
            return self.evaluator_factories
        if self.environment_factory is None:
            raise ValueError(
                "You need to set `environment_factory` in `OfflineExperimentConfig` "
                "when `evaluator_factories` are not specified. To disable evaluation "
                "altogether just set `evaluator_factories = []`"
            )

        return [
            default_evaluator_factory(
                environment_factory=self.environment_factory,
                network_factory=self.network_factory,
                policy_factory=self.builder.make_policy,
                logger_factory=self.logger_factory,
                observers=self.observers,
            )
        ]


def default_evaluator_factory(
    environment_factory: types.EnvironmentFactory,
    network_factory: NetworkFactory[builders.Networks],
    policy_factory: PolicyFactory[builders.Networks, builders.Policy],
    logger_factory: loggers.LoggerFactory,
    observers: Sequence[observers_lib.EnvLoopObserver] = (),
) -> EvaluatorFactory[builders.Policy]:
    """Returns a default evaluator process."""

    def evaluator(
        random_key: types.PRNGKey,
        variable_source: core.VariableSource,
        counter: counting.Counter,
        make_actor: MakeActorFn[builders.Policy],
    ):
        """The evaluation process."""

        # Create environment and evaluator networks
        environment_key, actor_key = jax.random.split(random_key)
        # Environments normally require uint32 as a seed.
        environment = environment_factory(utils.sample_uint32(environment_key))
        environment_spec = specs.make_environment_spec(environment)
        networks = network_factory(environment_spec)
        policy = policy_factory(networks, environment_spec, True)
        actor = make_actor(actor_key, policy, environment_spec, variable_source)

        # Create logger and counter.
        counter = counting.Counter(counter, "evaluator")
        logger = logger_factory("evaluator", "actor_steps", 0)

        # Create the run loop and return it.
        return environment_loop.EnvironmentLoop(
            environment, actor, counter, logger, observers=observers
        )

    return evaluator


def make_policy(
    experiment: ExperimentConfig[builders.Networks, builders.Policy, Any],
    networks: builders.Networks,
    environment_spec: specs.EnvironmentSpec,
    evaluation: bool,
) -> builders.Policy:
    """Constructs a policy. It is only meant to be used internally."""
    # TODO(sabela): remove and update callers once all agents use
    # builder.make_policy
    if not evaluation and experiment.policy_network_factory:
        return experiment.policy_network_factory(networks)
    if evaluation and experiment.eval_policy_network_factory:
        return experiment.eval_policy_network_factory(networks)
    return experiment.builder.make_policy(
        networks=networks, environment_spec=environment_spec, evaluation=evaluation
    )

corax/jax/experiments/config.py
ethanluoyc/corax
# type: ignore
"""A training loop for IL agent."""

import dataclasses
from typing import Callable

import jax
import reverb

import corax
from corax import specs as env_specs
from corax.jax import experiments
from corax.jax import utils
from corax.jax.experiments import config
from corax.jax.experiments import imitation_loop
from corax.jax.experiments import savers
from corax.jax.experiments.run_experiment import _disable_insert_blocking
from corax.jax.experiments.run_experiment import _LearningActor
from corax.utils import counting

Config = experiments.ExperimentConfig
RewarderFactory = Callable[
    [corax.VariableSource, env_specs.EnvironmentSpec], imitation_loop.EpisodeRewarder
]


@dataclasses.dataclass(frozen=True)
class ImitationExperimentConfig:
    config: experiments.ExperimentConfig
    rewarder_factory: RewarderFactory


def run_imitation_experiment(
    imitation_experiment: ImitationExperimentConfig,
    eval_every: int,
    num_eval_episodes: int,
):
    rewarder_factory = imitation_experiment.rewarder_factory
    experiment = imitation_experiment.config
    # Create an environment, grab the spec, and use it to create networks.
    key = jax.random.PRNGKey(experiment.seed)
    # key_train_env, key_eval, key_actor, key_eval_actor = jax.random.split(key, 4)
    environment_key, key = jax.random.split(key)
    environment = experiment.environment_factory(utils.sample_uint32(environment_key))
    environment_spec = experiment.environment_spec or corax.make_environment_spec(
        environment
    )
    networks = experiment.network_factory(environment_spec)
    policy = config.make_policy(
        experiment=experiment,
        networks=networks,
        environment_spec=environment_spec,
        evaluation=False,
    )

    replay_tables = experiment.builder.make_replay_tables(environment_spec, policy)

    # Disable blocking of inserts by tables' rate limiters, as this function
    # executes learning (sampling from the table) and data generation
    # (inserting into the table) sequentially from the same thread
    # which could result in blocked insert making the algorithm hang.
    replay_tables, rate_limiters_max_diff = _disable_insert_blocking(replay_tables)

    replay_server = reverb.Server(replay_tables, port=None)
    replay_client = reverb.Client(f"localhost:{replay_server.port}")

    # Parent counter allows to share step counts between train and eval loops and
    # the learner, so that it is possible to plot for example evaluator's return
    # value as a function of the number of training episodes.
    parent_counter = counting.Counter(time_delta=0.0)

    # Create actor, and learner for generating, storing, and consuming
    # data respectively.
    dataset = experiment.builder.make_dataset_iterator(replay_client)
    # We always use prefetch, as it provides an iterator with additional
    # 'ready' method.
    dataset = utils.prefetch(dataset, buffer_size=1)
    learner_key, key = jax.random.split(key)
    learner_counter = counting.Counter(parent_counter, prefix="learner", time_delta=0.0)

    learner = experiment.builder.make_learner(
        random_key=learner_key,
        networks=networks,
        dataset=dataset,
        logger_fn=experiment.logger_factory,
        environment_spec=environment_spec,
        replay_client=replay_client,
        counter=learner_counter,
    )

    rewarder = rewarder_factory(learner, environment_spec)
    actor_key, key = jax.random.split(key)
    # adder is not used in the actor but instead passed to the imitation loop.
    actor = experiment.builder.make_actor(
        actor_key, policy, environment_spec, variable_source=learner, adder=None
    )

    checkpointer = None
    if experiment.checkpointing is not None:
        checkpointing = experiment.checkpointing
        checkpointer = savers.Checkpointer(
            objects_to_save={"learner": learner, "counter": parent_counter},
            time_delta_minutes=checkpointing.time_delta_minutes,
            directory=checkpointing.directory,
            add_uid=checkpointing.add_uid,
            max_to_keep=checkpointing.max_to_keep,
            keep_checkpoint_every_n_hours=checkpointing.keep_checkpoint_every_n_hours,
            checkpoint_ttl_seconds=checkpointing.checkpoint_ttl_seconds,
        )

    # Replace the actor with a LearningActor. This makes sure that every time
    # that `update` is called on the actor it checks to see whether there is
    # any new data to learn from and if so it runs a learner step. The rate
    # at which new data is released is controlled by the replay table's
    # rate_limiter which is created by the builder.make_replay_tables call above.
    actor = _LearningActor(
        actor, learner, dataset, replay_tables, rate_limiters_max_diff, checkpointer
    )

    parent_counter = counting.Counter(time_delta=0.0)
    train_counter = counting.Counter(parent_counter, "actor", time_delta=0.0)
    train_logger = experiment.logger_factory("actor", train_counter.get_steps_key(), 0)

    adder = experiment.builder.make_adder(replay_client, environment_spec, policy)
    train_loop = imitation_loop.ImitationEnvironmentLoop(
        environment,
        actor,
        adder=adder,
        rewarder=rewarder,
        counter=train_counter,
        logger=train_logger,
        observers=experiment.observers,
    )

    eval_loop = None
    if num_eval_episodes > 0:
        # Create the evaluation actor and loop.
        eval_policy = config.make_policy(
            experiment=experiment,
            networks=networks,
            environment_spec=environment_spec,
            evaluation=True,
        )
        eval_actor_key, key = jax.random.split(key)
        eval_actor = experiment.builder.make_actor(
            random_key=eval_actor_key,
            policy=eval_policy,
            environment_spec=environment_spec,
            variable_source=learner,
        )
        eval_environment_key, key = jax.random.split(key)
        if experiment.eval_environment_factory is None:
            eval_env = experiment.environment_factory(
                utils.sample_uint32(eval_environment_key)
            )
        else:
            eval_env = experiment.eval_environment_factory(
                utils.sample_uint32(eval_environment_key)
            )

        eval_counter = counting.Counter(parent_counter, "evaluator", time_delta=0.0)
        eval_logger = experiment.logger_factory(
            "evaluator", eval_counter.get_steps_key(), 0
        )
        eval_loop = imitation_loop.ImitationEnvironmentLoop(
            eval_env,
            eval_actor,
            adder=None,
            rewarder=rewarder,
            counter=eval_counter,
            logger=eval_logger,
            should_update_rewarder=False,
            observers=experiment.observers,
        )

    steps = 0
    while steps < experiment.max_num_actor_steps:
        if eval_loop:
            eval_loop.run(num_episodes=num_eval_episodes)
        train_loop.run(num_steps=eval_every)
        steps += eval_every
    if eval_loop:
        eval_loop.run(num_episodes=num_eval_episodes)
        eval_env.close()  # type: ignore
    environment.close()

corax/jax/experiments/imitation_experiment.py
ethanluoyc/corax
# type: ignore
"""A training loop for IL agent."""
import operator
import time
from typing import List, Optional, Protocol, Sequence

import dm_env
import jax.numpy as jnp
import numpy as np
import tree
from dm_env import specs

from corax import adders
from corax import core
from corax import types
from corax.utils import counting
from corax.utils import loggers
from corax.utils import observers as observers_lib


class EpisodeRewarder(Protocol):
    def compute_offline_rewards(self, agent_steps, update: bool):
        """Compute pseudo-rewards for an episode."""


class ImitationEnvironmentLoop(core.Worker):
    def __init__(
        self,
        environment: dm_env.Environment,
        actor: core.Actor,
        adder: Optional[adders.Adder] = None,
        rewarder: Optional[EpisodeRewarder] = None,
        counter: Optional[counting.Counter] = None,
        logger: Optional[loggers.Logger] = None,
        should_update_actor: bool = True,
        should_update_rewarder: bool = True,
        label: str = "environment_loop",
        observers: Sequence[observers_lib.EnvLoopObserver] = (),
    ):
        # Internalize agent and environment.
        self._environment = environment
        self._actor = actor
        self._adder = adder
        self._rewarder = rewarder
        self._counter = counter or counting.Counter()
        self._logger = logger or loggers.make_default_logger(
            label, steps_key=self._counter.get_steps_key()
        )
        self._should_update_actor = should_update_actor
        self._should_update_rewarder = should_update_rewarder
        self._observers = observers

    def run_episode(self) -> loggers.LoggingData:
        """Run one episode.

        Each episode is a loop which interacts first with the environment to get an
        observation and then give that observation to the agent in order to retrieve
        an action.

        Returns:
          An instance of `loggers.LoggingData`.
        """
        # Reset any counts and start the environment.
        episode_start_time = time.time()
        select_action_durations: List[float] = []
        env_step_durations: List[float] = []
        episode_steps: int = 0

        # For evaluation, this keeps track of the total undiscounted reward
        # accumulated during the episode.
        episode_return = tree.map_structure(
            _generate_zeros_from_spec, self._environment.reward_spec()
        )
        env_reset_start = time.time()
        timestep = self._environment.reset()
        env_reset_duration = time.time() - env_reset_start
        # Make the first observation.
        timesteps = []
        steps = []
        self._actor.observe_first(timestep)
        timesteps.append(timestep)

        for observer in self._observers:
            # Initialize the observer with the current state of the env after reset
            # and the initial timestep.
            observer.observe_first(self._environment, timestep)

        # Run an episode.
        while not timestep.last():
            # Book-keeping.
            episode_steps += 1

            # Generate an action from the agent's policy.
            select_action_start = time.time()
            action = self._actor.select_action(timestep.observation)
            select_action_durations.append(time.time() - select_action_start)

            # Step the environment with the agent's selected action.
            env_step_start = time.time()
            timestep = self._environment.step(action)
            env_step_durations.append(time.time() - env_step_start)

            # Have the agent and observers observe the timestep.
            # TODO(https://github.com/ethanluoyc/cdrl/issues/26).
            # Investigate a different approach to align observation, action.
            self._actor.observe(action, next_timestep=timestep)
            steps.append(
                types.Transition(timesteps[-1].observation, action, (), (), ())
            )
            timesteps.append(timestep)

            for observer in self._observers:
                # One environment step was completed. Observe the current state of the
                # environment, the current timestep and the action.
                observer.observe(self._environment, timestep, action)

            # Give the actor the opportunity to update itself.
            if self._should_update_actor:
                self._actor.update()

            # Equivalent to: episode_return += timestep.reward
            # We capture the return value because if timestep.reward is a JAX
            # DeviceArray, episode_return will not be mutated in-place. (In all other
            # cases, the returned episode_return will be the same object as the
            # argument episode_return.)
            episode_return = tree.map_structure(
                operator.iadd, episode_return, timestep.reward
            )

        steps.append(
            types.Transition(
                observation=timesteps[-1].observation,
                action=np.zeros_like(steps[0].action),
                reward=(),
                discount=(),
                next_observation=(),
            )
        )

        # Compute episode imitation return
        rewarder_start_time = time.time()

        def _convert_timestep(timestep):
            # TODO(yl): LP has a bug when passing non-JAX TF arrays.
            # Convert to JAX arrays until that is fixed.
            return tree.map_structure(
                lambda x: jnp.asarray(x) if x is not None else None, timestep
            )

        psuedo_rewards = self._rewarder.compute_offline_rewards(
            [_convert_timestep(step) for step in steps],
            self._should_update_rewarder,
        )
        assert len(timesteps) - 1 == len(psuedo_rewards)

        episode_return_imitation = np.sum(psuedo_rewards)
        rewarder_duration_sec = time.time() - rewarder_start_time

        # Add the updated reward timesteps to adder
        if self._adder is not None:
            first_timestep = timesteps[0]
            self._adder.add_first(first_timestep._replace(reward=psuedo_rewards[0]))
            actions = [step.action for step in steps[:-1]]
            assert len(actions) == len(timesteps) - 1 == len(psuedo_rewards)
            for action, next_ts, relabeled_reward in zip(
                actions, timesteps[1:], psuedo_rewards
            ):
                self._adder.add(action, next_ts._replace(reward=relabeled_reward))
            self._adder.reset()

        # Record counts.
        counts = self._counter.increment(episodes=1, steps=episode_steps)

        # Collect the results and combine with counts.
        steps_per_second = episode_steps / (time.time() - episode_start_time)
        result = {
            "episode_length": episode_steps,
            "episode_return": episode_return,
            "episode_return_imitation": episode_return_imitation,
            "steps_per_second": steps_per_second,
            "env_reset_duration_sec": env_reset_duration,
            "select_action_duration_sec": np.mean(select_action_durations),
            "env_step_duration_sec": np.mean(env_step_durations),
            "rewarder_duration_sec": rewarder_duration_sec,
        }
        result.update(counts)
        for observer in self._observers:
            result.update(observer.get_metrics())
        return result

    def run(
        self,
        num_episodes: Optional[int] = None,
        num_steps: Optional[int] = None,
    ) -> int:
        """Perform the run loop.

        Run the environment loop either for `num_episodes` episodes or for at
        least `num_steps` steps (the last episode is always run until completion,
        so the total number of steps may be slightly more than `num_steps`).
        At least one of these two arguments has to be None.

        Upon termination of an episode a new episode will be started. If the number
        of episodes and the number of steps are not given then this will interact
        with the environment infinitely.

        Args:
          num_episodes: number of episodes to run the loop for.
          num_steps: minimal number of steps to run the loop for.

        Returns:
          Actual number of steps the loop executed.

        Raises:
          ValueError: If both 'num_episodes' and 'num_steps' are not None.
        """

        if not (num_episodes is None or num_steps is None):
            raise ValueError('Either "num_episodes" or "num_steps" should be None.')

        def should_terminate(episode_count: int, step_count: int) -> bool:
            return (num_episodes is not None and episode_count >= num_episodes) or (
                num_steps is not None and step_count >= num_steps
            )

        episode_count: int = 0
        step_count: int = 0
        while not should_terminate(episode_count, step_count):
            episode_start = time.time()
            result = self.run_episode()
            result = {**result, **{"episode_duration": time.time() - episode_start}}
            episode_count += 1
            step_count += int(result["episode_length"])
            # Log the given episode results.
            self._logger.write(result)

        return step_count


def _generate_zeros_from_spec(spec: specs.Array) -> np.ndarray:
    return np.zeros(spec.shape, spec.dtype)

corax/jax/experiments/imitation_loop.py
ethanluoyc/corax
import numpy as np
from absl.testing import absltest

from corax import specs
from corax.jax.experiments import imitation_loop
from corax.testing import fakes


class _DummyRewarder(imitation_loop.EpisodeRewarder):
    def compute_offline_rewards(self, agent_steps, update: bool):
        del update
        return np.zeros((len(agent_steps) - 1))


class ImitationLoopTest(absltest.TestCase):
    def test_imitation_loop(self):
        environment = fakes.DiscreteEnvironment()
        actor = fakes.Actor(specs.make_environment_spec(environment))
        rewarder = _DummyRewarder()
        loop = imitation_loop.ImitationEnvironmentLoop(
            environment, actor, adder=None, rewarder=rewarder
        )
        loop.run_episode()


if __name__ == "__main__":
    absltest.main()

corax/jax/experiments/imitation_loop_test.py
ethanluoyc/corax
# type: ignore
# 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.

"""Runners used for executing local agents."""

import sys
import time
from typing import Optional, Sequence, Tuple

import dm_env
import jax
import reverb

import corax
from corax import core
from corax import specs
from corax import types
from corax.jax import utils
from corax.jax.experiments import config
from corax.jax.experiments import savers
from corax.utils import counting


def run_experiment(
    experiment: config.ExperimentConfig,
    eval_every: int = 100,
    num_eval_episodes: int = 1,
):
    """Runs a simple, single-threaded training loop using the default evaluators.

    It targets simplicity of the code and so only the basic features of the
    ExperimentConfig are supported.

    Arguments:
      experiment: Definition and configuration of the agent to run.
      eval_every: After how many actor steps to perform evaluation.
      num_eval_episodes: How many evaluation episodes to execute at each
        evaluation step.
    """

    key = jax.random.PRNGKey(experiment.seed)

    # Create the environment and get its spec.
    environment = experiment.environment_factory(experiment.seed)
    environment_spec = experiment.environment_spec or specs.make_environment_spec(
        environment
    )

    # Create the networks and policy.
    networks = experiment.network_factory(environment_spec)
    policy = config.make_policy(
        experiment=experiment,
        networks=networks,
        environment_spec=environment_spec,
        evaluation=False,
    )

    # Create the replay server and grab its address.
    replay_tables = experiment.builder.make_replay_tables(environment_spec, policy)

    # Disable blocking of inserts by tables' rate limiters, as this function
    # executes learning (sampling from the table) and data generation
    # (inserting into the table) sequentially from the same thread
    # which could result in blocked insert making the algorithm hang.
    replay_tables, rate_limiters_max_diff = _disable_insert_blocking(replay_tables)

    replay_server = reverb.Server(replay_tables, port=None)
    replay_client = reverb.Client(f"localhost:{replay_server.port}")

    # Parent counter allows to share step counts between train and eval loops and
    # the learner, so that it is possible to plot for example evaluator's return
    # value as a function of the number of training episodes.
    parent_counter = counting.Counter(time_delta=0.0)

    dataset = experiment.builder.make_dataset_iterator(replay_client)
    # We always use prefetch as it provides an iterator with an additional
    # 'ready' method.
    dataset = utils.prefetch(dataset, buffer_size=1)

    # Create actor, adder, and learner for generating, storing, and consuming
    # data respectively.
    # NOTE: These are created in reverse order as the actor needs to be given the
    # adder and the learner (as a source of variables).
    learner_key, key = jax.random.split(key)
    learner = experiment.builder.make_learner(
        random_key=learner_key,
        networks=networks,
        dataset=dataset,
        logger_fn=experiment.logger_factory,
        environment_spec=environment_spec,
        replay_client=replay_client,
        counter=counting.Counter(parent_counter, prefix="learner", time_delta=0.0),
    )

    adder = experiment.builder.make_adder(replay_client, environment_spec, policy)

    actor_key, key = jax.random.split(key)
    actor = experiment.builder.make_actor(
        actor_key, policy, environment_spec, variable_source=learner, adder=adder
    )

    # Create the environment loop used for training.
    train_counter = counting.Counter(parent_counter, prefix="actor", time_delta=0.0)
    train_logger = experiment.logger_factory("actor", train_counter.get_steps_key(), 0)

    checkpointer = None
    if experiment.checkpointing is not None:
        checkpointing = experiment.checkpointing
        checkpointer = savers.Checkpointer(
            objects_to_save={"learner": learner, "counter": parent_counter},
            time_delta_minutes=checkpointing.time_delta_minutes,
            directory=checkpointing.directory,
            add_uid=checkpointing.add_uid,
            max_to_keep=checkpointing.max_to_keep,
            keep_checkpoint_every_n_hours=checkpointing.keep_checkpoint_every_n_hours,
            checkpoint_ttl_seconds=checkpointing.checkpoint_ttl_seconds,
        )

    # Replace the actor with a LearningActor. This makes sure that every time
    # that `update` is called on the actor it checks to see whether there is
    # any new data to learn from and if so it runs a learner step. The rate
    # at which new data is released is controlled by the replay table's
    # rate_limiter which is created by the builder.make_replay_tables call above.
    actor = _LearningActor(
        actor, learner, dataset, replay_tables, rate_limiters_max_diff, checkpointer
    )

    train_loop = corax.EnvironmentLoop(
        environment,
        actor,
        counter=train_counter,
        logger=train_logger,
        observers=experiment.observers,
    )

    max_num_actor_steps = (
        experiment.max_num_actor_steps
        - parent_counter.get_counts().get(train_counter.get_steps_key(), 0)
    )

    if num_eval_episodes == 0:
        # No evaluation. Just run the training loop.
        train_loop.run(num_steps=max_num_actor_steps)
        return

    # Create the evaluation actor and loop.
    eval_counter = counting.Counter(parent_counter, prefix="evaluator", time_delta=0.0)
    eval_logger = experiment.logger_factory(
        "evaluator", eval_counter.get_steps_key(), 0
    )
    eval_policy = config.make_policy(
        experiment=experiment,
        networks=networks,
        environment_spec=environment_spec,
        evaluation=True,
    )
    eval_actor = experiment.builder.make_actor(
        random_key=jax.random.PRNGKey(experiment.seed),
        policy=eval_policy,
        environment_spec=environment_spec,
        variable_source=learner,
    )
    eval_loop = corax.EnvironmentLoop(
        environment,
        eval_actor,
        counter=eval_counter,
        logger=eval_logger,
        observers=experiment.observers,
    )

    steps = 0
    while steps < max_num_actor_steps:
        eval_loop.run(num_episodes=num_eval_episodes)
        num_steps = min(eval_every, max_num_actor_steps - steps)
        steps += train_loop.run(num_steps=num_steps)
    eval_loop.run(num_episodes=num_eval_episodes)

    environment.close()


class _LearningActor(core.Actor):
    """Actor which learns (updates its parameters) when `update` is called.

    This combines a base actor and a learner. Whenever `update` is called
    on the wrapping actor the learner will take a step (e.g. one step of gradient
    descent) as long as there is data available for training
    (provided iterator and replay_tables are used to check for that).
    Selecting actions and making observations are handled by the base actor.
    Intended to be used by the `run_experiment` only.
    """

    def __init__(
        self,
        actor: core.Actor,
        learner: core.Learner,
        iterator: core.PrefetchingIterator,
        replay_tables: Sequence[reverb.Table],
        sample_sizes: Sequence[int],
        checkpointer: Optional[savers.Checkpointer],
    ):
        """Initializes _LearningActor.

        Args:
          actor: Actor to be wrapped.
          learner: Learner on which step() is to be called when there is data.
          iterator: Iterator used by the Learner to fetch training data.
          replay_tables: Collection of tables from which Learner fetches data
            through the iterator.
          sample_sizes: For each table from `replay_tables`, how many elements the
            table should have available for sampling to wait for the `iterator` to
            prefetch a batch of data. Otherwise more experience needs to be
            collected by the actor.
          checkpointer: Checkpointer to save the state on update.
        """
        self._actor = actor
        self._learner = learner
        self._iterator = iterator
        self._replay_tables = replay_tables
        self._sample_sizes = sample_sizes
        self._learner_steps = 0
        self._checkpointer = checkpointer

    def select_action(self, observation: types.NestedArray) -> types.NestedArray:
        return self._actor.select_action(observation)

    def observe_first(self, timestep: dm_env.TimeStep):
        self._actor.observe_first(timestep)

    def observe(self, action: types.NestedArray, next_timestep: dm_env.TimeStep):
        self._actor.observe(action, next_timestep)

    def _maybe_train(self):
        trained = False
        while True:
            if self._iterator.ready():
                self._learner.step()
                batches = self._iterator.retrieved_elements() - self._learner_steps
                self._learner_steps += 1
                assert batches == 1, (
                    "Learner step must retrieve exactly one element from the iterator"
                    f" (retrieved {batches}). Otherwise agent can deadlock. Example "
                    "cause is that your chosen agent"
                    "s Builder has a `make_learner` "
                    "factory that prefetches the data but it shouldn"
                    "t."
                )
                trained = True
            else:
                # Wait for the iterator to fetch more data from the table(s) only
                # if there plenty of data to sample from each table.
                for table, sample_size in zip(self._replay_tables, self._sample_sizes):
                    if not table.can_sample(sample_size):
                        return trained
                # Let iterator's prefetching thread get data from the table(s).
                time.sleep(0.001)

    def update(
        self,
    ):  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
        if self._maybe_train():
            # Update the actor weights only when learner was updated.
            self._actor.update()
        if self._checkpointer:
            self._checkpointer.save()


def _disable_insert_blocking(
    tables: Sequence[reverb.Table],
) -> Tuple[Sequence[reverb.Table], Sequence[int]]:
    """Disables blocking of insert operations for a given collection of tables."""
    modified_tables = []
    sample_sizes = []
    for table in tables:
        rate_limiter_info = table.info.rate_limiter_info
        rate_limiter = reverb.rate_limiters.RateLimiter(
            samples_per_insert=rate_limiter_info.samples_per_insert,
            min_size_to_sample=rate_limiter_info.min_size_to_sample,
            min_diff=rate_limiter_info.min_diff,
            max_diff=sys.float_info.max,
        )
        modified_tables.append(table.replace(rate_limiter=rate_limiter))
        # Target the middle of the rate limiter's insert-sample balance window.
        sample_sizes.append(
            max(1, int((rate_limiter_info.max_diff - rate_limiter_info.min_diff) / 2))
        )
    return modified_tables, sample_sizes

corax/jax/experiments/run_experiment.py
ethanluoyc/corax
# type: ignore
# 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.

"""Runner used for executing local offline RL agents."""

import jax

import corax
from corax import specs
from corax.jax.experiments import config
from corax.jax.experiments import savers
from corax.utils import counting


def run_offline_experiment(
    experiment: config.OfflineExperimentConfig,
    eval_every: int = 100,
    num_eval_episodes: int = 1,
):
    """Runs a simple, single-threaded training loop using the default evaluators.

    It targets simplicity of the code and so only the basic features of the
    OfflineExperimentConfig are supported.

    Arguments:
      experiment: Definition and configuration of the agent to run.
      eval_every: After how many learner steps to perform evaluation.
      num_eval_episodes: How many evaluation episodes to execute at each
        evaluation step.
    """

    key = jax.random.PRNGKey(experiment.seed)

    # Create the environment and get its spec.
    environment = experiment.environment_factory(experiment.seed)
    environment_spec = experiment.environment_spec or specs.make_environment_spec(
        environment
    )

    # Create the networks and policy.
    networks = experiment.network_factory(environment_spec)

    # Parent counter allows to share step counts between train and eval loops and
    # the learner, so that it is possible to plot for example evaluator's return
    # value as a function of the number of training episodes.
    parent_counter = counting.Counter(time_delta=0.0)

    # Create the demonstrations dataset.
    dataset_key, key = jax.random.split(key)
    dataset = experiment.demonstration_dataset_factory(dataset_key)

    # Create the learner.
    learner_key, key = jax.random.split(key)
    learner = experiment.builder.make_learner(
        random_key=learner_key,
        networks=networks,
        dataset=dataset,
        logger_fn=experiment.logger_factory,
        environment_spec=environment_spec,
        counter=counting.Counter(parent_counter, prefix="learner", time_delta=0.0),
    )

    # Define the evaluation loop.
    eval_loop = None
    if num_eval_episodes > 0:
        # Create the evaluation actor and loop.
        eval_counter = counting.Counter(
            parent_counter, prefix="evaluator", time_delta=0.0
        )
        eval_logger = experiment.logger_factory(
            "evaluator", eval_counter.get_steps_key(), 0
        )
        eval_key, key = jax.random.split(key)
        eval_actor = experiment.builder.make_actor(
            random_key=eval_key,
            policy=experiment.builder.make_policy(networks, environment_spec, True),
            environment_spec=environment_spec,
            variable_source=learner,
        )
        eval_loop = corax.EnvironmentLoop(
            environment,
            eval_actor,
            counter=eval_counter,
            logger=eval_logger,
            observers=experiment.observers,
        )

    checkpointer = None
    if experiment.checkpointing is not None:
        checkpointing = experiment.checkpointing
        checkpointer = savers.Checkpointer(
            objects_to_save={"learner": learner, "counter": parent_counter},
            time_delta_minutes=checkpointing.time_delta_minutes,
            directory=checkpointing.directory,
            add_uid=checkpointing.add_uid,
            max_to_keep=checkpointing.max_to_keep,
            keep_checkpoint_every_n_hours=checkpointing.keep_checkpoint_every_n_hours,
            checkpoint_ttl_seconds=checkpointing.checkpoint_ttl_seconds,
        )

    max_num_learner_steps = (
        experiment.max_num_learner_steps
        - parent_counter.get_counts().get("learner_steps", 0)
    )

    # Run the training loop.
    if eval_loop:
        eval_loop.run(num_eval_episodes)
    steps = 0
    while steps < max_num_learner_steps:
        learner_steps = min(eval_every, max_num_learner_steps - steps)
        for _ in range(learner_steps):
            learner.step()
            if checkpointer is not None:
                checkpointer.save()
        if eval_loop:
            eval_loop.run(num_eval_episodes)
        steps += learner_steps

corax/jax/experiments/run_offline_experiment.py
ethanluoyc/corax
# type: ignore
# 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.

"""Utility classes for saving model checkpoints and snapshots."""

import abc
import datetime
import pickle
import time
from typing import Mapping, Optional, Union

import tensorflow as tf
from absl import logging

from corax import core
from corax.utils import paths

# from corax.utils import signals

PythonState = tf.train.experimental.PythonState
Checkpointable = Union[tf.Module, tf.Variable, PythonState]

_DEFAULT_CHECKPOINT_TTL = int(datetime.timedelta(days=5).total_seconds())
_DEFAULT_SNAPSHOT_TTL = int(datetime.timedelta(days=90).total_seconds())


class TFSaveable(abc.ABC):
    """An interface for objects that expose their checkpointable TF state."""

    @property
    @abc.abstractmethod
    def state(self) -> Mapping[str, Checkpointable]:
        """Returns TensorFlow checkpointable state."""


class Checkpointer:
    """Convenience class for periodically checkpointing.

    This can be used to checkpoint any object with trackable state (e.g.
    tensorflow variables or modules); see tf.train.Checkpoint for
    details. Objects inheriting from tf.train.experimental.PythonState can also
    be checkpointed.

    Typically people use Checkpointer to make sure that they can correctly recover
    from a machine going down during learning. For more permanent storage of self-
    contained "networks" see the Snapshotter object.

    Usage example:

    ```python
    model = snt.Linear(10)
    checkpointer = Checkpointer(objects_to_save={'model': model})

    for _ in range(100):
      # ...
      checkpointer.save()
    ```
    """

    def __init__(
        self,
        objects_to_save: Mapping[str, Union[Checkpointable, core.Saveable]],
        *,
        directory: str = "~/acme/",
        subdirectory: str = "default",
        time_delta_minutes: float = 10.0,
        enable_checkpointing: bool = True,
        add_uid: bool = True,
        max_to_keep: int = 1,
        checkpoint_ttl_seconds: Optional[int] = _DEFAULT_CHECKPOINT_TTL,
        keep_checkpoint_every_n_hours: Optional[int] = None,
    ):
        """Builds the saver object.

        Args:
          objects_to_save: Mapping specifying what to checkpoint.
          directory: Which directory to put the checkpoint in.
          subdirectory: Sub-directory to use (e.g. if multiple checkpoints are being
            saved).
          time_delta_minutes: How often to save the checkpoint, in minutes.
          enable_checkpointing: whether to checkpoint or not.
          add_uid: If True adds a UID to the checkpoint path, see
            `paths.get_unique_id()` for how this UID is generated.
          max_to_keep: The maximum number of checkpoints to keep.
          checkpoint_ttl_seconds: TTL (time to leave) in seconds for checkpoints.
          keep_checkpoint_every_n_hours: keep_checkpoint_every_n_hours passed to
            tf.train.CheckpointManager.
        """

        # Convert `Saveable` objects to TF `Checkpointable` first, if necessary.
        def to_ckptable(x: Union[Checkpointable, core.Saveable]) -> Checkpointable:
            if isinstance(x, core.Saveable):
                return SaveableAdapter(x)
            return x

        objects_to_save = {k: to_ckptable(v) for k, v in objects_to_save.items()}

        self._time_delta_minutes = time_delta_minutes
        self._last_saved = 0.0
        self._enable_checkpointing = enable_checkpointing
        self._checkpoint_manager = None

        if enable_checkpointing:
            # Checkpoint object that handles saving/restoring.
            self._checkpoint = tf.train.Checkpoint(**objects_to_save)
            self._checkpoint_dir = paths.process_path(
                directory,
                "checkpoints",
                subdirectory,
                ttl_seconds=checkpoint_ttl_seconds,
                backups=False,
                add_uid=add_uid,
            )

            # Create a manager to maintain different checkpoints.
            self._checkpoint_manager = tf.train.CheckpointManager(
                self._checkpoint,
                directory=self._checkpoint_dir,
                max_to_keep=max_to_keep,
                keep_checkpoint_every_n_hours=keep_checkpoint_every_n_hours,
            )

            self.restore()

    def save(self, force: bool = False) -> bool:
        """Save the checkpoint if it's the appropriate time, otherwise no-ops.

        Args:
          force: Whether to force a save regardless of time elapsed since last save.

        Returns:
          A boolean indicating if a save event happened.
        """
        if not self._enable_checkpointing:
            return False

        if not force and time.time() - self._last_saved < 60 * self._time_delta_minutes:
            return False

        # Save any checkpoints.
        logging.info("Saving checkpoint: %s", self._checkpoint_manager.directory)
        self._checkpoint_manager.save()
        self._last_saved = time.time()

        return True

    def restore(self):
        # Restore from the most recent checkpoint (if it exists).
        checkpoint_to_restore = self._checkpoint_manager.latest_checkpoint
        logging.info("Attempting to restore checkpoint: %s", checkpoint_to_restore)
        self._checkpoint.restore(checkpoint_to_restore)

    @property
    def directory(self):
        return self._checkpoint_manager.directory


class CheckpointingRunner(core.Worker):
    """Wrap an object and expose a run method which checkpoints periodically.

    This internally creates a Checkpointer around `wrapped` object and exposes
    all of the methods of `wrapped`. Additionally, any `**kwargs` passed to the
    runner are forwarded to the internal Checkpointer.
    """

    def __init__(
        self,
        wrapped: Union[Checkpointable, core.Saveable, TFSaveable],
        key: str = "wrapped",
        *,
        time_delta_minutes: int = 30,
        **kwargs,
    ):
        if isinstance(wrapped, TFSaveable):
            # If the object to be wrapped exposes its TF State, checkpoint that.
            objects_to_save = wrapped.state
        else:
            # Otherwise checkpoint the wrapped object itself.
            objects_to_save = wrapped

        self._wrapped = wrapped
        self._time_delta_minutes = time_delta_minutes
        self._checkpointer = Checkpointer(
            objects_to_save={key: objects_to_save},
            time_delta_minutes=time_delta_minutes,
            **kwargs,
        )

    # Handle preemption signal. Note that this must happen in the main thread.
    def _signal_handler(self):
        logging.info("Caught SIGTERM: forcing a checkpoint save.")
        self._checkpointer.save(force=True)

    def step(self):
        if isinstance(self._wrapped, core.Learner):
            # Learners have a step() method, so alternate between that and ckpt call.
            self._wrapped.step()
            self._checkpointer.save()
        else:
            # Wrapped object doesn't have a run method; set our run method to ckpt.
            self.checkpoint()

    def run(self):
        # TODO(yl): Add back signal handling logic
        """Runs the checkpointer."""
        while True:
            self.step()

    def __dir__(self):
        return dir(self._wrapped) + ["get_directory"]

    # TODO(b/195915583) : Throw when wrapped object has get_directory() method.
    def __getattr__(self, name):
        if name == "get_directory":
            return self.get_directory
        return getattr(self._wrapped, name)

    def checkpoint(self):
        self._checkpointer.save()
        # Do not sleep for a long period of time to avoid LaunchPad program
        # termination hangs (time.sleep is not interruptible).
        for _ in range(self._time_delta_minutes * 60):
            time.sleep(1)

    def get_directory(self):
        return self._checkpointer.directory


class SaveableAdapter(tf.train.experimental.PythonState):
    """Adapter which allows `Saveable` object to be checkpointed by TensorFlow."""

    def __init__(self, object_to_save: core.Saveable):
        self._object_to_save = object_to_save

    def serialize(self):
        state = self._object_to_save.save()
        return pickle.dumps(state)

    def deserialize(self, pickled: bytes):
        state = pickle.loads(pickled)
        self._object_to_save.restore(state)

corax/jax/experiments/savers.py
ethanluoyc/corax
from corax.jax.networks.base import Action
from corax.jax.networks.base import BatchSize
from corax.jax.networks.base import FeedForwardNetwork
from corax.jax.networks.base import Logits
from corax.jax.networks.base import LogProbFn
from corax.jax.networks.base import Observation
from corax.jax.networks.base import Params
from corax.jax.networks.base import PRNGKey
from corax.jax.networks.base import QValues
from corax.jax.networks.base import SampleFn
from corax.jax.networks.continuous import LayerNormMLP
from corax.jax.networks.continuous import NearZeroInitializedLinear
from corax.jax.networks.distributional import GaussianMixture
from corax.jax.networks.distributional import MultivariateNormalDiagHead
from corax.jax.networks.distributional import NormalTanhDistribution
from corax.jax.networks.distributional import TanhTransformedDistribution
from corax.jax.networks.rescaling import ClipToSpec
from corax.jax.networks.rescaling import RescaleToSpec
from corax.jax.networks.rescaling import TanhToSpec

corax/jax/networks/__init__.py
ethanluoyc/corax
# 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.

"""Base interfaces for networks."""

import dataclasses
from typing import Callable, Optional, Tuple

import haiku as hk
import jax.numpy as jnp
from typing_extensions import Protocol

from corax import types
from corax.jax import types as jax_types

# This definition is deprecated. Use jax_types.PRNGKey directly instead.
# TODO(sinopalnikov): migrate all users and remove this definition.
PRNGKey = jax_types.PRNGKey

# Commonly-used types.
BatchSize = int
Observation = types.NestedArray
Action = types.NestedArray
Params = types.NestedArray
NetworkOutput = types.NestedArray
QValues = jnp.ndarray
Logits = jnp.ndarray
LogProb = jnp.ndarray
Value = jnp.ndarray
RecurrentState = types.NestedArray
Entropy = jnp.ndarray

# Commonly-used function/network signatures.
QNetwork = Callable[[Observation], QValues]
LSTMOutputs = Tuple[Tuple[Logits, Value], hk.LSTMState]
PolicyValueRNN = Callable[[Observation, hk.LSTMState], LSTMOutputs]
RecurrentQNetwork = Callable[[Observation, hk.LSTMState], Tuple[QValues, hk.LSTMState]]
SampleFn = Callable[[NetworkOutput, PRNGKey], Action]
LogProbFn = Callable[[NetworkOutput, Action], LogProb]


@dataclasses.dataclass
class FeedForwardNetwork:
    """Holds a pair of pure functions defining a feed-forward network.

    Attributes:
      init: A pure function: ``params = init(rng, *a, **k)``
      apply: A pure function: ``out = apply(params, rng, *a, **k)``
    """

    # Initializes and returns the networks parameters.
    init: Callable[..., Params]
    # Computes and returns the outputs of a forward pass.
    apply: Callable[..., NetworkOutput]


class ApplyFn(Protocol):
    def __call__(
        self,
        params: Params,
        observation: Observation,
        *args,
        is_training: bool,
        key: Optional[PRNGKey] = None,
        **kwargs,
    ) -> NetworkOutput:
        ...


@dataclasses.dataclass
class TypedFeedForwardNetwork:
    """FeedForwardNetwork with more specific types of the member functions.

    Attributes:
      init: A pure function. Initializes and returns the networks parameters.
      apply: A pure function. Computes and returns the outputs of a forward pass.
    """

    init: Callable[[PRNGKey], Params]
    apply: ApplyFn


def non_stochastic_network_to_typed(
    network: FeedForwardNetwork,
) -> TypedFeedForwardNetwork:
    """Converts non-stochastic FeedForwardNetwork to TypedFeedForwardNetwork.

    Non-stochastic network is the one that doesn't take a random key as an input
    for its `apply` method.

    Arguments:
      network: non-stochastic feed-forward network.

    Returns:
      corresponding TypedFeedForwardNetwork
    """

    def apply(
        params: Params,
        observation: Observation,
        *args,
        is_training: bool,
        key: Optional[PRNGKey] = None,
        **kwargs,
    ) -> NetworkOutput:
        del is_training, key
        return network.apply(params, observation, *args, **kwargs)

    return TypedFeedForwardNetwork(init=network.init, apply=apply)

corax/jax/networks/base.py
ethanluoyc/corax
# type: ignore
# 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.

"""Networks used in continuous control."""

from typing import Callable, Sequence

import haiku as hk
import jax
import jax.numpy as jnp

uniform_initializer = hk.initializers.UniformScaling(scale=0.333)


class NearZeroInitializedLinear(hk.Linear):
    """Simple linear layer, initialized at near zero weights and zero biases."""

    def __init__(self, output_size: int, scale: float = 1e-4):
        super().__init__(output_size, w_init=hk.initializers.VarianceScaling(scale))


class LayerNormMLP(hk.Module):
    """Simple feedforward MLP torso with initial layer-norm.

    This MLP's first linear layer is followed by a LayerNorm layer and a tanh
    non-linearity; subsequent layers use `activation`, which defaults to elu.

    Note! The default activation differs from the usual MLP default of ReLU for
    legacy reasons.
    """

    def __init__(
        self,
        layer_sizes: Sequence[int],
        w_init: hk.initializers.Initializer = uniform_initializer,
        activation: Callable[[jnp.ndarray], jnp.ndarray] = jax.nn.elu,
        activate_final: bool = False,
        name: str = "feedforward_mlp_torso",
    ):
        """Construct the MLP.

        Args:
          layer_sizes: a sequence of ints specifying the size of each layer.
          w_init: initializer for Linear layers.
          activation: nonlinearity to use in the MLP, defaults to elu.
            Note! The default activation differs from the usual MLP default of ReLU
            for legacy reasons.
          activate_final: whether or not to use the activation function on the final
            layer of the neural network.
          name: a name for the module.
        """
        super().__init__(name=name)

        self._network = hk.Sequential(
            [
                hk.Linear(layer_sizes[0], w_init=w_init),
                hk.LayerNorm(axis=-1, create_scale=True, create_offset=True),
                jax.lax.tanh,
                hk.nets.MLP(
                    layer_sizes[1:],
                    w_init=w_init,
                    activation=activation,
                    activate_final=activate_final,
                ),
            ]
        )

    def __call__(self, inputs: jnp.ndarray) -> jnp.ndarray:
        """Forwards the policy network."""
        return self._network(inputs)

corax/jax/networks/continuous.py
ethanluoyc/corax
# type: ignore
# 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.

"""Haiku modules that output tfd.Distributions."""

from typing import Any, Callable, List, Optional, Sequence, Union

import chex
import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np
import tensorflow_probability as tf_tfp
import tensorflow_probability.substrates.jax as tfp

hk_init = hk.initializers
tfd = tfp.distributions
_MIN_SCALE = 1e-4
Initializer = hk.initializers.Initializer


class CategoricalHead(hk.Module):
    """Module that produces a categorical distribution with the given number of values."""

    def __init__(
        self,
        num_values: Union[int, List[int]],
        dtype: Optional[Any] = jnp.int32,
        w_init: Optional[Initializer] = None,
        name: Optional[str] = None,
    ):
        super().__init__(name=name)
        self._dtype = dtype
        self._logit_shape = num_values
        self._linear = hk.Linear(np.prod(num_values), w_init=w_init)

    def __call__(self, inputs: jnp.ndarray) -> tfd.Distribution:
        logits = self._linear(inputs)
        if not isinstance(self._logit_shape, int):
            logits = hk.Reshape(self._logit_shape)(logits)
        return tfd.Categorical(logits=logits, dtype=self._dtype)


class GaussianMixture(hk.Module):
    """Module that outputs a Gaussian Mixture Distribution."""

    def __init__(
        self,
        num_dimensions: int,
        num_components: int,
        multivariate: bool,
        init_scale: Optional[float] = None,
        append_singleton_event_dim: bool = False,
        reinterpreted_batch_ndims: Optional[int] = None,
        transformation_fn: Optional[
            Callable[[tfd.Distribution], tfd.Distribution]
        ] = None,
        name: str = "GaussianMixture",
    ):
        """Initialization.

        Args:
          num_dimensions: dimensionality of the output distribution
          num_components: number of mixture components.
          multivariate: whether the resulting distribution is multivariate or not.
          init_scale: the initial scale for the Gaussian mixture components.
          append_singleton_event_dim: (univariate only) Whether to add an extra
            singleton dimension to the event shape.
          reinterpreted_batch_ndims: (univariate only) Number of batch dimensions to
            reinterpret as event dimensions.
          transformation_fn: Distribution transform such as TanhTransformation
            applied to individual components.
          name: name of the module passed to snt.Module parent class.
        """
        super().__init__(name=name)

        self._num_dimensions = num_dimensions
        self._num_components = num_components
        self._multivariate = multivariate
        self._append_singleton_event_dim = append_singleton_event_dim
        self._reinterpreted_batch_ndims = reinterpreted_batch_ndims

        if init_scale is not None:
            self._scale_factor = init_scale / jax.nn.softplus(0.0)
        else:
            self._scale_factor = 1.0  # Corresponds to init_scale = softplus(0).

        self._transformation_fn = transformation_fn

    def __call__(
        self, inputs: jnp.ndarray, low_noise_policy: bool = False
    ) -> tfd.Distribution:
        """Run the networks through inputs.

        Args:
          inputs: hidden activations of the policy network body.
          low_noise_policy: whether to set vanishingly small scales for each
            component. If this flag is set to True, the policy is effectively run
            without Gaussian noise.

        Returns:
          Mixture Gaussian distribution.
        """

        # Define the weight initializer.
        w_init = hk.initializers.VarianceScaling(scale=1e-5)

        # Create a layer that outputs the unnormalized log-weights.
        if self._multivariate:
            logits_size = self._num_components
        else:
            logits_size = self._num_dimensions * self._num_components
        logit_layer = hk.Linear(logits_size, w_init=w_init)

        # Create two layers that outputs a location and a scale, respectively, for
        # each dimension and each component.
        loc_layer = hk.Linear(
            self._num_dimensions * self._num_components, w_init=w_init
        )
        scale_layer = hk.Linear(
            self._num_dimensions * self._num_components, w_init=w_init
        )

        # Compute logits, locs, and scales if necessary.
        logits = logit_layer(inputs)
        locs = loc_layer(inputs)

        # When a low_noise_policy is requested, set the scales to its minimum value.
        if low_noise_policy:
            scales = jnp.full(locs.shape, _MIN_SCALE)
        else:
            scales = scale_layer(inputs)
            scales = self._scale_factor * jax.nn.softplus(scales) + _MIN_SCALE

        if self._multivariate:
            components_class = tfd.MultivariateNormalDiag
            shape = [-1, self._num_components, self._num_dimensions]  # [B, C, D]
            # In this case, no need to reshape logits as they are in the correct shape
            # already, namely [batch_size, num_components].
        else:
            components_class = tfd.Normal
            shape = [-1, self._num_dimensions, self._num_components]  # [B, D, C]
            if self._append_singleton_event_dim:
                shape.insert(2, 1)  # [B, D, 1, C]
            logits = logits.reshape(shape)

        # Reshape the mixture's location and scale parameters appropriately.
        locs = locs.reshape(shape)
        scales = scales.reshape(shape)

        if self._multivariate:
            components_distribution = components_class(loc=locs, scale_diag=scales)
        else:
            components_distribution = components_class(loc=locs, scale=scales)

        # Transformed the component distributions in the mixture.
        if self._transformation_fn:
            components_distribution = self._transformation_fn(components_distribution)

        # Create the mixture distribution.
        distribution = tfd.MixtureSameFamily(
            mixture_distribution=tfd.Categorical(logits=logits),
            components_distribution=components_distribution,
        )

        if not self._multivariate:
            distribution = tfd.Independent(
                distribution, reinterpreted_batch_ndims=self._reinterpreted_batch_ndims
            )

        return distribution


class TanhTransformedDistribution(tfd.TransformedDistribution):
    """Distribution followed by tanh."""

    def __init__(self, distribution, threshold=0.999, validate_args=False):
        """Initialize the distribution.

        Args:
          distribution: The distribution to transform.
          threshold: Clipping value of the action when computing the logprob.
          validate_args: Passed to super class.
        """
        super().__init__(
            distribution=distribution,
            bijector=tfp.bijectors.Tanh(),
            validate_args=validate_args,
        )
        # Computes the log of the average probability distribution outside the
        # clipping range, i.e. on the interval [-inf, -atanh(threshold)] for
        # log_prob_left and [atanh(threshold), inf] for log_prob_right.
        self._threshold = threshold
        inverse_threshold = self.bijector.inverse(threshold)
        # average(pdf) = p/epsilon
        # So log(average(pdf)) = log(p) - log(epsilon)
        log_epsilon = jnp.log(1.0 - threshold)
        # Those 2 values are differentiable w.r.t. model parameters, such that the
        # gradient is defined everywhere.
        self._log_prob_left = (
            self.distribution.log_cdf(-inverse_threshold) - log_epsilon
        )
        self._log_prob_right = (
            self.distribution.log_survival_function(inverse_threshold) - log_epsilon
        )

    def log_prob(self, event):
        # Without this clip there would be NaNs in the inner tf.where and that
        # causes issues for some reasons.
        event = jnp.clip(event, -self._threshold, self._threshold)
        # The inverse image of {threshold} is the interval [atanh(threshold), inf]
        # which has a probability of "log_prob_right" under the given distribution.
        return jnp.where(
            event <= -self._threshold,
            self._log_prob_left,
            jnp.where(
                event >= self._threshold, self._log_prob_right, super().log_prob(event)
            ),
        )

    def mode(self):
        return self.bijector.forward(self.distribution.mode())

    def entropy(self, seed=None):
        # We return an estimation using a single sample of the log_det_jacobian.
        # We can still do some backpropagation with this estimate.
        return self.distribution.entropy() + self.bijector.forward_log_det_jacobian(
            self.distribution.sample(seed=seed), event_ndims=0
        )

    @classmethod
    def _parameter_properties(cls, dtype: Optional[Any], num_classes=None):
        td_properties = super()._parameter_properties(dtype, num_classes=num_classes)
        del td_properties["bijector"]
        return td_properties


class NormalTanhDistribution(hk.Module):
    """Module that produces a TanhTransformedDistribution distribution."""

    def __init__(
        self,
        num_dimensions: int,
        min_scale: float = 1e-3,
        w_init: hk_init.Initializer = hk_init.VarianceScaling(1.0, "fan_in", "uniform"),
        b_init: hk_init.Initializer = hk_init.Constant(0.0),
    ):
        """Initialization.

        Args:
          num_dimensions: Number of dimensions of a distribution.
          min_scale: Minimum standard deviation.
          w_init: Initialization for linear layer weights.
          b_init: Initialization for linear layer biases.
        """
        super().__init__(name="Normal")
        self._min_scale = min_scale
        self._loc_layer = hk.Linear(num_dimensions, w_init=w_init, b_init=b_init)
        self._scale_layer = hk.Linear(num_dimensions, w_init=w_init, b_init=b_init)

    def __call__(self, inputs: jnp.ndarray) -> tfd.Distribution:
        loc = self._loc_layer(inputs)
        scale = self._scale_layer(inputs)
        scale = jax.nn.softplus(scale) + self._min_scale
        distribution = tfd.Normal(loc=loc, scale=scale)
        return tfd.Independent(
            TanhTransformedDistribution(distribution), reinterpreted_batch_ndims=1
        )


class MultivariateNormalDiagHead(hk.Module):
    """Module that produces a tfd.MultivariateNormalDiag distribution."""

    def __init__(
        self,
        num_dimensions: int,
        init_scale: float = 0.3,
        min_scale: float = 1e-6,
        w_init: hk_init.Initializer = hk_init.VarianceScaling(1e-4),
        b_init: hk_init.Initializer = hk_init.Constant(0.0),
    ):
        """Initialization.

        Args:
          num_dimensions: Number of dimensions of MVN distribution.
          init_scale: Initial standard deviation.
          min_scale: Minimum standard deviation.
          w_init: Initialization for linear layer weights.
          b_init: Initialization for linear layer biases.
        """
        super().__init__(name="MultivariateNormalDiagHead")
        self._min_scale = min_scale
        self._init_scale = init_scale
        self._loc_layer = hk.Linear(num_dimensions, w_init=w_init, b_init=b_init)
        self._scale_layer = hk.Linear(num_dimensions, w_init=w_init, b_init=b_init)

    def __call__(self, inputs: jnp.ndarray) -> tfd.Distribution:
        loc = self._loc_layer(inputs)
        scale = jax.nn.softplus(self._scale_layer(inputs))
        scale *= self._init_scale / jax.nn.softplus(0.0)
        scale += self._min_scale
        return tfd.MultivariateNormalDiag(loc=loc, scale_diag=scale)


class CategoricalValueHead(hk.Module):
    """Network head that produces a categorical distribution and value."""

    def __init__(
        self,
        num_values: int,
        name: Optional[str] = None,
    ):
        super().__init__(name=name)
        self._logit_layer = hk.Linear(num_values)
        self._value_layer = hk.Linear(1)

    def __call__(self, inputs: jnp.ndarray):
        logits = self._logit_layer(inputs)
        value = jnp.squeeze(self._value_layer(inputs), axis=-1)
        return (tfd.Categorical(logits=logits), value)


class DiscreteValued(hk.Module):
    """C51-style head.

    For each action, it produces the logits for a discrete distribution over
    atoms. Therefore, the returned logits represents several distributions, one
    for each action.
    """

    def __init__(
        self,
        num_actions: int,
        head_units: int = 512,
        num_atoms: int = 51,
        v_min: float = -1.0,
        v_max: float = 1.0,
    ):
        super().__init__("DiscreteValued")
        self._num_actions = num_actions
        self._num_atoms = num_atoms
        self._atoms = jnp.linspace(v_min, v_max, self._num_atoms)
        self._network = hk.nets.MLP([head_units, num_actions * num_atoms])

    def __call__(self, inputs: jnp.ndarray):
        q_logits = self._network(inputs)
        q_logits = jnp.reshape(q_logits, (-1, self._num_actions, self._num_atoms))
        q_dist = jax.nn.softmax(q_logits)
        q_values = jnp.sum(q_dist * self._atoms, axis=2)
        q_values = jax.lax.stop_gradient(q_values)
        return q_values, q_logits, self._atoms


class CategoricalCriticHead(hk.Module):
    """Critic head that uses a categorical to represent action values."""

    def __init__(
        self,
        num_bins: int = 601,
        vmax: Optional[float] = None,
        vmin: Optional[float] = None,
        w_init: hk_init.Initializer = hk_init.VarianceScaling(1e-5),
    ):
        super().__init__(name="categorical_critic_head")
        vmax = vmax if vmax is not None else 0.5 * (num_bins - 1)
        vmin = vmin if vmin is not None else -1.0 * vmax

        self._head = DiscreteValuedTfpHead(
            vmin=vmin, vmax=vmax, logits_shape=(1,), num_atoms=num_bins, w_init=w_init
        )

    def __call__(self, embedding: chex.Array) -> tfd.Distribution:
        output = self._head(embedding)
        return output


class DiscreteValuedTfpHead(hk.Module):
    """Represents a parameterized discrete valued distribution.

    The returned distribution is essentially a `tfd.Categorical` that knows its
    support and thus can compute the mean value.
    """

    def __init__(
        self,
        vmin: float,
        vmax: float,
        num_atoms: int,
        logits_shape: Optional[Sequence[int]] = None,
        w_init: Optional[Initializer] = None,
        b_init: Optional[Initializer] = None,
    ):
        """Initialization.

        If vmin and vmax have shape S, this will store the category values as a
        Tensor of shape (S*, num_atoms).

        Args:
          vmin: Minimum of the value range
          vmax: Maximum of the value range
          num_atoms: The atom values associated with each bin.
          logits_shape: The shape of the logits, excluding batch and num_atoms
            dimensions.
          w_init: Initialization for linear layer weights.
          b_init: Initialization for linear layer biases.
        """
        super().__init__(name="DiscreteValuedHead")
        self._values = np.linspace(vmin, vmax, num=num_atoms, axis=-1)
        if not logits_shape:
            logits_shape = ()
        self._logits_shape = logits_shape + (num_atoms,)
        self._w_init = w_init
        self._b_init = b_init

    def __call__(self, inputs: chex.Array) -> tfd.Distribution:
        net = hk.Linear(
            np.prod(self._logits_shape), w_init=self._w_init, b_init=self._b_init
        )
        logits = net(inputs)
        logits = hk.Reshape(self._logits_shape, preserve_dims=1)(logits)
        return DiscreteValuedTfpDistribution(values=self._values, logits=logits)


@tf_tfp.experimental.auto_composite_tensor
class DiscreteValuedTfpDistribution(tfd.Categorical):
    """This is a generalization of a categorical distribution.

    The support for the DiscreteValued distribution can be any real valued range,
    whereas the categorical distribution has support [0, n_categories - 1] or
    [1, n_categories]. This generalization allows us to take the mean of the
    distribution over its support.
    """

    def __init__(
        self,
        values: chex.Array,
        logits: Optional[chex.Array] = None,
        probs: Optional[chex.Array] = None,
        name: str = "DiscreteValuedDistribution",
    ):
        """Initialization.

        Args:
          values: Values making up support of the distribution. Should have a shape
            compatible with logits.
          logits: An N-D Tensor, N >= 1, representing the log probabilities of a set
            of Categorical distributions. The first N - 1 dimensions index into a
            batch of independent distributions and the last dimension indexes into
            the classes.
          probs: An N-D Tensor, N >= 1, representing the probabilities of a set of
            Categorical distributions. The first N - 1 dimensions index into a batch
            of independent distributions and the last dimension represents a vector
            of probabilities for each class. Only one of logits or probs should be
            passed in.
          name: Name of the distribution object.
        """
        parameters = dict(locals())
        self._values = np.asarray(values)

        if logits is not None:
            logits = jnp.asarray(logits)
            chex.assert_shape(logits, (..., *self._values.shape))

        if probs is not None:
            probs = jnp.asarray(probs)
            chex.assert_shape(probs, (..., *self._values.shape))

        super().__init__(logits=logits, probs=probs, name=name)

        self._parameters = parameters

    @property
    def values(self):
        return self._values

    @classmethod
    def _parameter_properties(cls, dtype, num_classes=None):
        return dict(
            values=tfp.util.ParameterProperties(
                event_ndims=None,
                shape_fn=lambda shape: (num_classes,),
                specifies_shape=True,
            ),
            logits=tfp.util.ParameterProperties(event_ndims=1),
            probs=tfp.util.ParameterProperties(event_ndims=1, is_preferred=False),
        )

    def _sample_n(self, key: chex.PRNGKey, n: int) -> chex.Array:
        indices = super()._sample_n(key=key, n=n)
        return jnp.take_along_axis(self._values, indices, axis=-1)

    def mean(self) -> chex.Array:
        """Overrides the Categorical mean by incorporating category values."""
        return jnp.sum(self.probs_parameter() * self._values, axis=-1)

    def variance(self) -> chex.Array:
        """Overrides the Categorical variance by incorporating category values."""
        dist_squared = jnp.square(jnp.expand_dims(self.mean(), -1) - self._values)
        return jnp.sum(self.probs_parameter() * dist_squared, axis=-1)

    def _event_shape(self):
        return jnp.zeros((), dtype=jnp.int32)

    def _event_shape_tensor(self):
        return []

corax/jax/networks/distributional.py
ethanluoyc/corax
# 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.

"""Rescaling layers (e.g. to match action specs)."""

import dataclasses

import jax.numpy as jnp
from jax import lax

from corax import specs


@dataclasses.dataclass
class ClipToSpec:
    """Clips inputs to within a BoundedArraySpec."""

    spec: specs.BoundedArray

    def __call__(self, inputs: jnp.ndarray) -> jnp.ndarray:
        return jnp.clip(inputs, self.spec.minimum, self.spec.maximum)


@dataclasses.dataclass
class RescaleToSpec:
    """Rescales inputs in [-1, 1] to match a BoundedArraySpec."""

    spec: specs.BoundedArray

    def __call__(self, inputs: jnp.ndarray) -> jnp.ndarray:
        scale = self.spec.maximum - self.spec.minimum
        offset = self.spec.minimum
        inputs = 0.5 * (inputs + 1.0)  # [0, 1]
        output = inputs * scale + offset  # [minimum, maximum]
        return output


@dataclasses.dataclass
class TanhToSpec:
    """Squashes real-valued inputs to match a BoundedArraySpec."""

    spec: specs.BoundedArray

    def __call__(self, inputs: jnp.ndarray) -> jnp.ndarray:
        scale = self.spec.maximum - self.spec.minimum
        offset = self.spec.minimum
        inputs = lax.tanh(inputs)  # [-1, 1]
        inputs = 0.5 * (inputs + 1.0)  # [0, 1]
        output = inputs * scale + offset  # [minimum, maximum]
        return output

corax/jax/networks/rescaling.py
ethanluoyc/corax
# type: ignore
# 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.

"""Utility functions to compute running statistics."""

import dataclasses
from typing import Any, Optional, Tuple, Union

import chex
import jax
import jax.numpy as jnp
import numpy as np
import tree

from corax import types
from corax.utils import tree_utils

Path = Tuple[Any, ...]
"""Path in a nested structure.

  A path is a tuple of indices (normally strings for maps and integers for
  arrays and tuples) that uniquely identifies a subtree in the nested structure.
  See
  https://tree.readthedocs.io/en/latest/api.html#tree.map_structure_with_path
  for more details.
"""


def _is_prefix(a: Path, b: Path) -> bool:
    """Returns whether `a` is a prefix of `b`."""
    return b[: len(a)] == a


def _zeros_like(nest: types.Nest, dtype=None) -> types.NestedArray:
    return jax.tree_map(lambda x: jnp.zeros(x.shape, dtype or x.dtype), nest)


def _ones_like(nest: types.Nest, dtype=None) -> types.NestedArray:
    return jax.tree_map(lambda x: jnp.ones(x.shape, dtype or x.dtype), nest)


@chex.dataclass(frozen=True)
class NestedMeanStd:
    """A container for running statistics (mean, std) of possibly nested data."""

    mean: types.NestedArray
    std: types.NestedArray


@chex.dataclass(frozen=True)
class RunningStatisticsState(NestedMeanStd):
    """Full state of running statistics computation."""

    count: Union[int, jnp.ndarray]
    summed_variance: types.NestedArray


@dataclasses.dataclass(frozen=True)
class NestStatisticsConfig:
    """Specifies how to compute statistics for Nests with the same structure.

    Attributes:
      paths: A sequence of Nest paths to compute statistics for. If there is a
        collision between paths (one is a prefix of the other), the shorter path
        takes precedence.
    """

    paths: Tuple[Path, ...] = ((),)


def _is_path_included(config: NestStatisticsConfig, path: Path) -> bool:
    """Returns whether the path is included in the config."""
    # A path is included in the config if it corresponds to a tree node that
    # belongs to a subtree rooted at the node corresponding to some path in
    # the config.
    return any(_is_prefix(config_path, path) for config_path in config.paths)


def init_state(nest: types.Nest) -> RunningStatisticsState:
    """Initializes the running statistics for the given nested structure."""
    dtype = jnp.float64 if jax.config.jax_enable_x64 else jnp.float32

    return RunningStatisticsState(  # pytype: disable=wrong-arg-types  # jax-ndarray
        count=0.0,
        mean=_zeros_like(nest, dtype=dtype),
        summed_variance=_zeros_like(nest, dtype=dtype),
        # Initialize with ones to make sure normalization works correctly
        # in the initial state.
        std=_ones_like(nest, dtype=dtype),
    )


def _validate_batch_shapes(
    batch: types.NestedArray,
    reference_sample: types.NestedArray,
    batch_dims: Tuple[int, ...],
) -> None:
    """Verifies shapes of the batch leaves against the reference sample.

    Checks that batch dimensions are the same in all leaves in the batch.
    Checks that non-batch dimensions for all leaves in the batch are the same
    as in the reference sample.

    Arguments:
      batch: the nested batch of data to be verified.
      reference_sample: the nested array to check non-batch dimensions.
      batch_dims: a Tuple of indices of batch dimensions in the batch shape.

    Returns:
      None.
    """

    def validate_node_shape(reference_sample: jnp.ndarray, batch: jnp.ndarray) -> None:
        expected_shape = batch_dims + reference_sample.shape
        assert batch.shape == expected_shape, f"{batch.shape} != {expected_shape}"

    tree_utils.fast_map_structure(validate_node_shape, reference_sample, batch)


def update(
    state: RunningStatisticsState,
    batch: types.NestedArray,
    *,
    config: NestStatisticsConfig = NestStatisticsConfig(),
    weights: Optional[jnp.ndarray] = None,
    std_min_value: float = 1e-6,
    std_max_value: float = 1e6,
    pmap_axis_name: Optional[str] = None,
    validate_shapes: bool = True,
) -> RunningStatisticsState:
    """Updates the running statistics with the given batch of data.

    Note: data batch and state elements (mean, etc.) must have the same structure.

    Note: by default will use int32 for counts and float32 for accumulated
    variance. This results in an integer overflow after 2^31 data points and
    degrading precision after 2^24 batch updates or even earlier if variance
    updates have large dynamic range.
    To improve precision, consider setting jax_enable_x64 to True, see
    https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#double-64bit-precision

    Arguments:
      state: The running statistics before the update.
      batch: The data to be used to update the running statistics.
      config: The config that specifies which leaves of the nested structure
        should the running statistics be computed for.
      weights: Weights of the batch data. Should match the batch dimensions.
        Passing a weight of 2. should be equivalent to updating on the
        corresponding data point twice.
      std_min_value: Minimum value for the standard deviation.
      std_max_value: Maximum value for the standard deviation.
      pmap_axis_name: Name of the pmapped axis, if any.
      validate_shapes: If true, the shapes of all leaves of the batch will be
        validated. Enabled by default. Doesn't impact performance when jitted.

    Returns:
      Updated running statistics.
    """
    # We require exactly the same structure to avoid issues when flattened
    # batch and state have different order of elements.
    tree.assert_same_structure(batch, state.mean)
    batch_shape = tree.flatten(batch)[0].shape
    # We assume the batch dimensions always go first.
    batch_dims = batch_shape[: len(batch_shape) - tree.flatten(state.mean)[0].ndim]
    batch_axis = range(len(batch_dims))
    if weights is None:
        step_increment = np.prod(batch_dims)
    else:
        step_increment = jnp.sum(weights)
    if pmap_axis_name is not None:
        step_increment = jax.lax.psum(step_increment, axis_name=pmap_axis_name)
    count = state.count + step_increment

    # Validation is important. If the shapes don't match exactly, but are
    # compatible, arrays will be silently broadcasted resulting in incorrect
    # statistics.
    if validate_shapes:
        if weights is not None:
            if weights.shape != batch_dims:
                raise ValueError(f"{weights.shape} != {batch_dims}")
        _validate_batch_shapes(batch, state.mean, batch_dims)

    def _compute_node_statistics(
        path: Path, mean: jnp.ndarray, summed_variance: jnp.ndarray, batch: jnp.ndarray
    ) -> Tuple[jnp.ndarray, jnp.ndarray]:
        assert isinstance(mean, jnp.ndarray), type(mean)
        assert isinstance(summed_variance, jnp.ndarray), type(summed_variance)
        if not _is_path_included(config, path):
            # Return unchanged.
            return mean, summed_variance
        # The mean and the sum of past variances are updated with Welford's
        # algorithm using batches (see https://stackoverflow.com/q/56402955).
        diff_to_old_mean = batch - mean
        if weights is not None:
            expanded_weights = jnp.reshape(
                weights, list(weights.shape) + [1] * (batch.ndim - weights.ndim)
            )
            diff_to_old_mean = diff_to_old_mean * expanded_weights
        mean_update = jnp.sum(diff_to_old_mean, axis=batch_axis) / count
        if pmap_axis_name is not None:
            mean_update = jax.lax.psum(mean_update, axis_name=pmap_axis_name)
        mean = mean + mean_update

        diff_to_new_mean = batch - mean
        variance_update = diff_to_old_mean * diff_to_new_mean
        variance_update = jnp.sum(variance_update, axis=batch_axis)
        if pmap_axis_name is not None:
            variance_update = jax.lax.psum(variance_update, axis_name=pmap_axis_name)
        summed_variance = summed_variance + variance_update
        return mean, summed_variance

    updated_stats = tree_utils.fast_map_structure_with_path(
        _compute_node_statistics, state.mean, state.summed_variance, batch
    )
    # map_structure_up_to is slow, so shortcut if we know the input is not
    # structured.
    if isinstance(state.mean, jnp.ndarray):
        mean, summed_variance = updated_stats
    else:
        # Reshape the updated stats from `nest(mean, summed_variance)` to
        # `nest(mean), nest(summed_variance)`.
        mean, summed_variance = [
            tree.map_structure_up_to(state.mean, lambda s, i=idx: s[i], updated_stats)
            for idx in range(2)
        ]

    def compute_std(
        path: Path, summed_variance: jnp.ndarray, std: jnp.ndarray
    ) -> jnp.ndarray:
        assert isinstance(summed_variance, jnp.ndarray)
        if not _is_path_included(config, path):
            return std
        # Summed variance can get negative due to rounding errors.
        summed_variance = jnp.maximum(summed_variance, 0)
        std = jnp.sqrt(summed_variance / count)
        std = jnp.clip(std, std_min_value, std_max_value)
        return std

    std = tree_utils.fast_map_structure_with_path(
        compute_std, summed_variance, state.std
    )

    return RunningStatisticsState(
        count=count, mean=mean, summed_variance=summed_variance, std=std
    )


def normalize(
    batch: types.NestedArray,
    mean_std: NestedMeanStd,
    max_abs_value: Optional[float] = None,
) -> types.NestedArray:
    """Normalizes data using running statistics."""

    def normalize_leaf(
        data: jnp.ndarray, mean: jnp.ndarray, std: jnp.ndarray
    ) -> jnp.ndarray:
        # Only normalize inexact types.
        if not jnp.issubdtype(data.dtype, jnp.inexact):
            return data
        data = (data - mean) / std
        if max_abs_value is not None:
            # TODO(b/124318564): remove pylint directive
            data = jnp.clip(
                data, -max_abs_value, +max_abs_value
            )  # pylint: disable=invalid-unary-operand-type
        return data

    return tree_utils.fast_map_structure(
        normalize_leaf, batch, mean_std.mean, mean_std.std
    )


def denormalize(batch: types.NestedArray, mean_std: NestedMeanStd) -> types.NestedArray:
    """Denormalizes values in a nested structure using the given mean/std.

    Only values of inexact types are denormalized.
    See https://numpy.org/doc/stable/_images/dtype-hierarchy.png for Numpy type
    hierarchy.

    Args:
      batch: a nested structure containing batch of data.
      mean_std: mean and standard deviation used for denormalization.

    Returns:
      Nested structure with denormalized values.
    """

    def denormalize_leaf(
        data: jnp.ndarray, mean: jnp.ndarray, std: jnp.ndarray
    ) -> jnp.ndarray:
        # Only denormalize inexact types.
        if not np.issubdtype(data.dtype, np.inexact):
            return data
        return data * std + mean

    return tree_utils.fast_map_structure(
        denormalize_leaf, batch, mean_std.mean, mean_std.std
    )


@dataclasses.dataclass(frozen=True)
class NestClippingConfig:
    """Specifies how to clip Nests with the same structure.

    Attributes:
      path_map: A map that specifies how to clip values in Nests with the same
        structure. Keys correspond to paths in the nest. Values are maximum
        absolute values to use for clipping. If there is a collision between paths
        (one path is a prefix of the other), the behavior is undefined.
    """

    path_map: Tuple[Tuple[Path, float], ...] = ()


def get_clip_config_for_path(
    config: NestClippingConfig, path: Path
) -> NestClippingConfig:
    """Returns the config for a subtree from the leaf defined by the path."""
    # Start with an empty config.
    path_map = []
    for map_path, max_abs_value in config.path_map:
        if _is_prefix(map_path, path):
            return NestClippingConfig(path_map=(((), max_abs_value),))
        if _is_prefix(path, map_path):
            path_map.append((map_path[len(path) :], max_abs_value))
    return NestClippingConfig(path_map=tuple(path_map))


def clip(
    batch: types.NestedArray, clipping_config: NestClippingConfig
) -> types.NestedArray:
    """Clips the batch."""

    def max_abs_value_for_path(path: Path, x: jnp.ndarray) -> Optional[float]:
        del x  # Unused, needed by interface.
        return next(
            (
                max_abs_value
                for clipping_path, max_abs_value in clipping_config.path_map
                if _is_prefix(clipping_path, path)
            ),
            None,
        )

    max_abs_values = tree_utils.fast_map_structure_with_path(
        max_abs_value_for_path, batch
    )

    def clip_leaf(data: jnp.ndarray, max_abs_value: Optional[float]) -> jnp.ndarray:
        if max_abs_value is not None:
            # TODO(b/124318564): remove pylint directive
            data = jnp.clip(
                data, -max_abs_value, +max_abs_value
            )  # pylint: disable=invalid-unary-operand-type
        return data

    return tree_utils.fast_map_structure(clip_leaf, batch, max_abs_values)


@dataclasses.dataclass(frozen=True)
class NestNormalizationConfig:
    """Specifies how to normalize Nests with the same structure.

    Attributes:
      stats_config: A config that defines how to compute running statistics to be
        used for normalization.
      clip_config: A config that defines how to clip normalized values.
    """

    stats_config: NestStatisticsConfig = NestStatisticsConfig()
    clip_config: NestClippingConfig = NestClippingConfig()

corax/jax/running_statistics.py
ethanluoyc/corax
# 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.

"""Tests for running statistics utilities."""

import functools
import math
from typing import NamedTuple

import jax
import jax.numpy as jnp
import numpy as np
import tree
from absl.testing import absltest
from jax.config import config as jax_config  # type: ignore

from corax import specs
from corax.jax import running_statistics

update_and_validate = functools.partial(running_statistics.update, validate_shapes=True)


class TestNestedSpec(NamedTuple):
    # Note: the fields are intentionally in reverse order to test ordering.
    a: specs.Array
    b: specs.Array


class RunningStatisticsTest(absltest.TestCase):
    def setUp(self):
        super().setUp()
        jax_config.update("jax_enable_x64", False)

    def assert_allclose(
        self, actual: jnp.ndarray, desired: jnp.ndarray, err_msg: str = ""
    ) -> None:
        np.testing.assert_allclose(
            actual, desired, atol=1e-5, rtol=1e-5, err_msg=err_msg
        )

    def test_normalize(self):
        state = running_statistics.init_state(specs.Array((5,), jnp.float32))

        x = jnp.arange(200, dtype=jnp.float32).reshape(20, 2, 5)
        x1, x2, x3, x4 = jnp.split(x, 4, axis=0)

        state = update_and_validate(state, x1)
        state = update_and_validate(state, x2)
        state = update_and_validate(state, x3)
        state = update_and_validate(state, x4)
        normalized = running_statistics.normalize(x, state)

        mean = jnp.mean(normalized)
        std = jnp.std(normalized)
        self.assert_allclose(mean, jnp.zeros_like(mean))
        self.assert_allclose(std, jnp.ones_like(std))

    def test_init_normalize(self):
        state = running_statistics.init_state(specs.Array((5,), jnp.float32))

        x = jnp.arange(200, dtype=jnp.float32).reshape(20, 2, 5)
        normalized = running_statistics.normalize(x, state)

        self.assert_allclose(normalized, x)

    def test_one_batch_dim(self):
        state = running_statistics.init_state(specs.Array((5,), jnp.float32))

        x = jnp.arange(10, dtype=jnp.float32).reshape(2, 5)

        state = update_and_validate(state, x)
        normalized = running_statistics.normalize(x, state)

        mean = jnp.mean(normalized, axis=0)
        std = jnp.std(normalized, axis=0)
        self.assert_allclose(mean, jnp.zeros_like(mean))
        self.assert_allclose(std, jnp.ones_like(std))

    def test_clip(self):
        state = running_statistics.init_state(specs.Array((), jnp.float32))

        x = jnp.arange(5, dtype=jnp.float32)

        state = update_and_validate(state, x)
        normalized = running_statistics.normalize(x, state, max_abs_value=1.0)

        mean = jnp.mean(normalized)
        std = jnp.std(normalized)
        self.assert_allclose(mean, jnp.zeros_like(mean))
        self.assert_allclose(std, jnp.ones_like(std) * math.sqrt(0.6))

    def test_nested_normalize(self):
        state = running_statistics.init_state(
            {"a": specs.Array((5,), jnp.float32), "b": specs.Array((2,), jnp.float32)}
        )

        x1 = {
            "a": jnp.arange(20, dtype=jnp.float32).reshape(2, 2, 5),
            "b": jnp.arange(8, dtype=jnp.float32).reshape(2, 2, 2),
        }
        x2 = {
            "a": jnp.arange(20, dtype=jnp.float32).reshape(2, 2, 5) + 20,
            "b": jnp.arange(8, dtype=jnp.float32).reshape(2, 2, 2) + 8,
        }
        x3 = {
            "a": jnp.arange(40, dtype=jnp.float32).reshape(4, 2, 5),
            "b": jnp.arange(16, dtype=jnp.float32).reshape(4, 2, 2),
        }

        state = update_and_validate(state, x1)
        state = update_and_validate(state, x2)
        state = update_and_validate(state, x3)
        normalized = running_statistics.normalize(x3, state)

        mean = tree.map_structure(lambda x: jnp.mean(x, axis=(0, 1)), normalized)
        std = tree.map_structure(lambda x: jnp.std(x, axis=(0, 1)), normalized)
        tree.map_structure(lambda x: self.assert_allclose(x, jnp.zeros_like(x)), mean)
        tree.map_structure(lambda x: self.assert_allclose(x, jnp.ones_like(x)), std)

    def test_validation(self):
        state = running_statistics.init_state(specs.Array((1, 2, 3), jnp.float32))

        x = jnp.arange(12, dtype=jnp.float32).reshape(2, 2, 3)
        with self.assertRaises(AssertionError):
            update_and_validate(state, x)

        x = jnp.arange(3, dtype=jnp.float32).reshape(1, 1, 3)
        with self.assertRaises(AssertionError):
            update_and_validate(state, x)

    def test_int_not_normalized(self):
        state = running_statistics.init_state(specs.Array((), jnp.int32))

        x = jnp.arange(5, dtype=jnp.int32)

        state = update_and_validate(state, x)
        normalized = running_statistics.normalize(x, state)

        np.testing.assert_array_equal(normalized, x)

    def test_pmap_update_nested(self):
        local_device_count = jax.local_device_count()
        state = running_statistics.init_state(
            {"a": specs.Array((5,), jnp.float32), "b": specs.Array((2,), jnp.float32)}
        )

        x = {
            "a": (jnp.arange(15 * local_device_count, dtype=jnp.float32)).reshape(
                local_device_count, 3, 5
            ),
            "b": (jnp.arange(6 * local_device_count, dtype=jnp.float32)).reshape(
                local_device_count, 3, 2
            ),
        }

        devices = jax.local_devices()
        state = jax.device_put_replicated(state, devices)
        pmap_axis_name = "i"
        state = jax.pmap(
            functools.partial(update_and_validate, pmap_axis_name=pmap_axis_name),
            pmap_axis_name,
        )(state, x)
        state = jax.pmap(
            functools.partial(update_and_validate, pmap_axis_name=pmap_axis_name),
            pmap_axis_name,
        )(state, x)
        normalized = jax.pmap(running_statistics.normalize)(x, state)

        mean = tree.map_structure(lambda x: jnp.mean(x, axis=(0, 1)), normalized)
        std = tree.map_structure(lambda x: jnp.std(x, axis=(0, 1)), normalized)
        tree.map_structure(lambda x: self.assert_allclose(x, jnp.zeros_like(x)), mean)
        tree.map_structure(lambda x: self.assert_allclose(x, jnp.ones_like(x)), std)

    def test_different_structure_normalize(self):
        spec = TestNestedSpec(
            a=specs.Array((5,), jnp.float32), b=specs.Array((2,), jnp.float32)
        )
        state = running_statistics.init_state(spec)

        x = {
            "a": jnp.arange(20, dtype=jnp.float32).reshape(2, 2, 5),
            "b": jnp.arange(8, dtype=jnp.float32).reshape(2, 2, 2),
        }

        with self.assertRaises(TypeError):
            state = update_and_validate(state, x)

    def test_weights(self):
        state = running_statistics.init_state(specs.Array((), jnp.float32))

        x = jnp.arange(5, dtype=jnp.float32)
        x_weights = jnp.ones_like(x)
        y = 2 * x + 5
        y_weights = 2 * x_weights
        z = jnp.concatenate([x, y])
        weights = jnp.concatenate([x_weights, y_weights])

        state = update_and_validate(state, z, weights=weights)

        self.assertEqual(state.mean, (jnp.mean(x) + 2 * jnp.mean(y)) / 3)
        big_z = jnp.concatenate([x, y, y])
        normalized = running_statistics.normalize(big_z, state)
        self.assertAlmostEqual(jnp.mean(normalized), 0.0, places=6)
        self.assertAlmostEqual(jnp.std(normalized), 1.0, places=6)

    def test_normalize_config(self):
        x = jnp.arange(200, dtype=jnp.float32).reshape(20, 2, 5)
        x_split = jnp.split(x, 5, axis=0)

        y = jnp.arange(160, dtype=jnp.float32).reshape(20, 2, 4)
        y_split = jnp.split(y, 5, axis=0)

        z = {"a": x, "b": y}

        z_split = [{"a": xx, "b": yy} for xx, yy in zip(x_split, y_split)]

        update = jax.jit(running_statistics.update, static_argnames=("config",))

        config = running_statistics.NestStatisticsConfig((("a",),))
        state = running_statistics.init_state(
            {"a": specs.Array((5,), jnp.float32), "b": specs.Array((4,), jnp.float32)}
        )
        # Test initialization from the first element.
        state = update(state, z_split[0], config=config)
        state = update(state, z_split[1], config=config)
        state = update(state, z_split[2], config=config)
        state = update(state, z_split[3], config=config)
        state = update(state, z_split[4], config=config)

        normalize = jax.jit(running_statistics.normalize)
        normalized = normalize(z, state)

        for key in normalized:
            mean = jnp.mean(normalized[key], axis=(0, 1))
            std = jnp.std(normalized[key], axis=(0, 1))
            if key == "a":
                self.assert_allclose(
                    mean,
                    jnp.zeros_like(mean),
                    err_msg=f"key:{key} mean:{mean} normalized:{normalized[key]}",
                )
                self.assert_allclose(
                    std,
                    jnp.ones_like(std),
                    err_msg=f"key:{key} std:{std} normalized:{normalized[key]}",
                )
            else:
                assert key == "b"
                np.testing.assert_array_equal(
                    normalized[key],
                    z[key],
                    err_msg=f"z:{z[key]} normalized:{normalized[key]}",
                )

    def test_clip_config(self):
        x = jnp.arange(10, dtype=jnp.float32) - 5
        y = jnp.arange(8, dtype=jnp.float32) - 4

        z = {"x": x, "y": y}

        max_abs_x = 2
        config = running_statistics.NestClippingConfig(((("x",), max_abs_x),))

        clipped_z = running_statistics.clip(z, config)

        clipped_x = jnp.clip(a=x, a_min=-max_abs_x, a_max=max_abs_x)
        np.testing.assert_array_equal(clipped_z["x"], clipped_x)

        np.testing.assert_array_equal(clipped_z["y"], z["y"])

    def test_denormalize(self):
        state = running_statistics.init_state(specs.Array((5,), jnp.float32))

        x = jnp.arange(100, dtype=jnp.float32).reshape(10, 2, 5)
        x1, x2 = jnp.split(x, 2, axis=0)

        state = update_and_validate(state, x1)
        state = update_and_validate(state, x2)
        normalized = running_statistics.normalize(x, state)

        mean = jnp.mean(normalized)
        std = jnp.std(normalized)
        self.assert_allclose(mean, jnp.zeros_like(mean))
        self.assert_allclose(std, jnp.ones_like(std))

        denormalized = running_statistics.denormalize(normalized, state)
        self.assert_allclose(denormalized, x)


if __name__ == "__main__":
    absltest.main()

corax/jax/running_statistics_test.py
ethanluoyc/corax
# 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.

"""Common JAX type definitions."""

from typing import Callable, Generic, Mapping, TypeVar

import chex
import dm_env
import jax
import jax.numpy as jnp

from corax import types

PRNGKey = jax.Array
Networks = TypeVar("Networks")
"""Container for all agent network components."""
Policy = TypeVar("Policy")
"""Function or container for agent policy functions."""
Sample = TypeVar("Sample")
"""Sample from the demonstrations or replay buffer."""
TrainingState = TypeVar("TrainingState")

TrainingMetrics = Mapping[str, jnp.ndarray]
"""Metrics returned by the training step.

Typically these are logged, so the values are expected to be scalars.
"""

Variables = Mapping[str, types.NestedArray]
"""Mapping of variable collections.

A mapping of variable collections, as defined by Learner.get_variables.
The keys are the collection names, the values are nested arrays representing
the values of the corresponding collection variables.
"""


@chex.dataclass(frozen=True, mappable_dataclass=False)
class TrainingStepOutput(Generic[TrainingState]):
    state: TrainingState
    metrics: TrainingMetrics


Seed = int
EnvironmentFactory = Callable[[Seed], dm_env.Environment]

corax/jax/types.py
ethanluoyc/corax
# type: ignore
# 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.

"""Utilities for JAX."""

import functools
import itertools
import queue
import threading
from typing import (
    Callable,
    Iterable,
    Iterator,
    NamedTuple,
    Optional,
    Sequence,
    Tuple,
    TypeVar,
)

import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np
import reverb
import tree
from absl import logging

from corax import core
from corax import types
from corax.jax import types as jax_types

F = TypeVar("F", bound=Callable)
N = TypeVar("N", bound=types.NestedArray)
T = TypeVar("T")


NUM_PREFETCH_THREADS = 1


def add_batch_dim(values: types.Nest) -> types.NestedArray:
    return jax.tree_map(lambda x: jnp.expand_dims(x, axis=0), values)


def _flatten(x: jnp.ndarray, num_batch_dims: int) -> jnp.ndarray:
    """Flattens the input, preserving the first ``num_batch_dims`` dimensions.

    If the input has fewer than ``num_batch_dims`` dimensions, it is returned
    unchanged.
    If the input has exactly ``num_batch_dims`` dimensions, an extra dimension
    is added. This is needed to handle batched scalars.

    Arguments:
      x: the input array to flatten.
      num_batch_dims: number of dimensions to preserve.

    Returns:
      flattened input.
    """
    # TODO(b/173492429): consider throwing an error instead.
    if x.ndim < num_batch_dims:
        return x
    return jnp.reshape(x, list(x.shape[:num_batch_dims]) + [-1])


def batch_concat(
    values: types.NestedArray,
    num_batch_dims: int = 1,
) -> jnp.ndarray:
    """Flatten and concatenate nested array structure, keeping batch dims."""
    flatten_fn = lambda x: _flatten(x, num_batch_dims)
    flat_leaves = tree.map_structure(flatten_fn, values)
    return jnp.concatenate(tree.flatten(flat_leaves), axis=-1)


def zeros_like(nest: types.Nest, dtype=None) -> types.NestedArray:
    return jax.tree_map(lambda x: jnp.zeros(x.shape, dtype or x.dtype), nest)


def ones_like(nest: types.Nest, dtype=None) -> types.NestedArray:
    return jax.tree_map(lambda x: jnp.ones(x.shape, dtype or x.dtype), nest)


def squeeze_batch_dim(nest: types.Nest) -> types.NestedArray:
    return jax.tree_map(lambda x: jnp.squeeze(x, axis=0), nest)


def to_numpy_squeeze(values: types.Nest) -> types.NestedArray:
    """Converts to numpy and squeezes out dummy batch dimension."""
    return jax.tree_map(lambda x: np.asarray(x).squeeze(axis=0), values)


def to_numpy(values: types.Nest) -> types.NestedArray:
    return jax.tree_map(np.asarray, values)


def fetch_devicearray(values: types.Nest) -> types.Nest:
    """Fetches and converts any DeviceArrays to np.ndarrays."""
    return tree.map_structure(_fetch_devicearray, values)


def _fetch_devicearray(x):
    if isinstance(x, jax.Array):
        return np.asarray(x)
    return x


def batch_to_sequence(values: types.Nest) -> types.NestedArray:
    return jax.tree_map(
        lambda x: jnp.transpose(x, axes=(1, 0, *range(2, len(x.shape)))), values
    )


def tile_array(array: jnp.ndarray, multiple: int) -> jnp.ndarray:
    """Tiles `multiple` copies of `array` along a new leading axis."""
    return jnp.stack([array] * multiple)


def tile_nested(inputs: types.Nest, multiple: int) -> types.Nest:
    """Tiles tensors in a nested structure along a new leading axis."""
    tile = functools.partial(tile_array, multiple=multiple)
    return jax.tree_map(tile, inputs)


def maybe_recover_lstm_type(state: types.NestedArray) -> types.NestedArray:
    """Recovers the type hk.LSTMState if LSTMState is in the type name.

    When the recurrent state of recurrent neural networks (RNN) is deserialized,
    for example when it is sampled from replay, it is sometimes repacked in a type
    that is identical to the source type but not the correct type itself. When
    using this state as the initial state in an hk.dynamic_unroll, this will
    cause hk.dynamic_unroll to raise an error as it requires its input and output
    states to be identical.

    Args:
      state: a nested structure of arrays representing the state of an RNN.

    Returns:
      Either the state unchanged if it is anything but an LSTMState, otherwise
      returns the state arrays properly contained in an hk.LSTMState.
    """
    return hk.LSTMState(*state) if type(state).__name__ == "LSTMState" else state


def prefetch(
    iterable: Iterable[T],
    buffer_size: int = 5,
    device: Optional[jax.Device] = None,
    num_threads: int = NUM_PREFETCH_THREADS,
) -> core.PrefetchingIterator[T]:
    """Returns prefetching iterator with additional 'ready' method."""

    return PrefetchIterator(iterable, buffer_size, device, num_threads)


class PrefetchingSplit(NamedTuple):
    host: types.NestedArray
    device: types.NestedArray


_SplitFunction = Callable[[types.NestedArray], PrefetchingSplit]


def keep_key_on_host(sample: reverb.ReplaySample) -> PrefetchingSplit:
    """Returns PrefetchingSplit which keeps uint64 reverb key on the host.

    We want to avoid truncation of the uint64 reverb key by JAX.

    Args:
      sample: a sample from a Reverb replay buffer.

    Returns:
      PrefetchingSplit with device having the reverb sample, and key on host.
    """
    return PrefetchingSplit(host=sample.info.key, device=sample)


def device_put(
    iterable: Iterable[types.NestedArray],
    device: jax.Device,
    split_fn: Optional[_SplitFunction] = None,
):
    """Returns iterator that samples an item and places it on the device."""

    return PutToDevicesIterable(
        iterable=iterable, pmapped_user=False, devices=[device], split_fn=split_fn
    )


def multi_device_put(
    iterable: Iterable[types.NestedArray],
    devices: Sequence[jax.Device],
    split_fn: Optional[_SplitFunction] = None,
):
    """Returns iterator that, per device, samples an item and places on device."""

    return PutToDevicesIterable(
        iterable=iterable, pmapped_user=True, devices=devices, split_fn=split_fn
    )


class PutToDevicesIterable(Iterable[types.NestedArray]):
    """Per device, samples an item from iterator and places on device.

    if pmapped_user:
      Items from the resulting generator are intended to be used in a pmapped
      function. Every element is a ShardedDeviceArray or (nested) Python container
      thereof. A single next() call to this iterator results in len(devices)
      calls to the underlying iterator. The returned items are put one on each
      device.
    if not pmapped_user:
      Places a sample from the iterator on the given device.

    Yields:
      If no split_fn is specified:
        DeviceArray/ShardedDeviceArray or (nested) Python container thereof
        representing the elements of shards stacked together, with each shard
        backed by physical device memory specified by the corresponding entry in
        devices.

      If split_fn is specified:
        PrefetchingSplit where the .host element is a stacked numpy array or
        (nested) Python contained thereof. The .device element is a
        DeviceArray/ShardedDeviceArray or (nested) Python container thereof.

    Raises:
      StopIteration: if there are not enough items left in the iterator to place
        one sample on each device.
      Any error thrown by the iterable_function. Note this is not raised inside
        the producer, but after it finishes executing.
    """

    def __init__(
        self,
        iterable: Iterable[types.NestedArray],
        pmapped_user: bool,
        devices: Sequence[jax.Device],
        split_fn: Optional[_SplitFunction] = None,
    ):
        """Constructs PutToDevicesIterable.

        Args:
          iterable: A python iterable. This is used to build the python prefetcher.
            Note that each iterable should only be passed to this function once as
            iterables aren't thread safe.
          pmapped_user: whether the user of data from this iterator is implemented
            using pmapping.
          devices: Devices used for prefecthing.
          split_fn: Optional function applied to every element from the iterable to
            split the parts of it that will be kept in the host and the parts that
            will sent to the device.

        Raises:
          ValueError: If devices list is empty, or if pmapped_use=False and more
            than 1 device is provided.
        """
        self.num_devices = len(devices)
        if self.num_devices == 0:
            raise ValueError("At least one device must be specified.")
        if (not pmapped_user) and (self.num_devices != 1):
            raise ValueError(
                "User is not implemented with pmapping but len(devices) "
                f"= {len(devices)} is not equal to 1! Devices given are:"
                f"\n{devices}"
            )

        self.iterable = iterable
        self.pmapped_user = pmapped_user
        self.split_fn = split_fn
        self.devices = devices
        self.iterator = iter(self.iterable)

    def __iter__(self) -> Iterator[types.NestedArray]:
        # It is important to structure the Iterable like this, because in
        # JustPrefetchIterator we must build a new iterable for each thread.
        # This is crucial if working with tensorflow datasets because tf.Graph
        # objects are thread local.
        self.iterator = iter(self.iterable)
        return self

    def __next__(self) -> types.NestedArray:
        try:
            if not self.pmapped_user:
                item = next(self.iterator)
                if self.split_fn is None:
                    return jax.device_put(item, self.devices[0])
                item_split = self.split_fn(item)
                return PrefetchingSplit(
                    host=item_split.host,
                    device=jax.device_put(item_split.device, self.devices[0]),
                )

            items = itertools.islice(self.iterator, self.num_devices)
            items = tuple(items)
            if len(items) < self.num_devices:
                raise StopIteration
            if self.split_fn is None:
                return jax.device_put_sharded(tuple(items), self.devices)
            else:
                # ((host: x1, device: y1), ..., (host: xN, device: yN)).
                items_split = (self.split_fn(item) for item in items)
                # (host: (x1, ..., xN), device: (y1, ..., yN)).
                split = tree.map_structure_up_to(
                    PrefetchingSplit(None, None), lambda *x: x, *items_split
                )

                return PrefetchingSplit(
                    host=np.stack(split.host),
                    device=jax.device_put_sharded(split.device, self.devices),
                )

        except StopIteration:
            raise

        except Exception:  # pylint: disable=broad-except
            logging.exception("Error for %s", self.iterable)
            raise


def sharded_prefetch(
    iterable: Iterable[types.NestedArray],
    buffer_size: int = 5,
    num_threads: int = 1,
    split_fn: Optional[_SplitFunction] = None,
    devices: Optional[Sequence[jax.Device]] = None,
) -> core.PrefetchingIterator:
    """Performs sharded prefetching from an iterable in separate threads.

    Elements from the resulting generator are intended to be used in a jax.pmap
    call. Every element is a sharded prefetched array with an additional replica
    dimension and corresponds to jax.local_device_count() elements from the
    original iterable.

    Args:
      iterable: A python iterable. This is used to build the python prefetcher.
        Note that each iterable should only be passed to this function once as
        iterables aren't thread safe.
      buffer_size (int): Number of elements to keep in the prefetch buffer.
      num_threads (int): Number of threads.
      split_fn: Optional function applied to every element from the iterable to
        split the parts of it that will be kept in the host and the parts that
        will sent to the device.
      devices: Devices used for prefecthing. Optional, jax.local_devices() by
        default.

    Returns:
      Prefetched elements from the original iterable with additional replica
      dimension.
    Raises:
      ValueError if the buffer_size <= 1.
      Any error thrown by the iterable_function. Note this is not raised inside
        the producer, but after it finishes executing.
    """

    devices = devices or jax.local_devices()

    iterable = PutToDevicesIterable(
        iterable=iterable, pmapped_user=True, devices=devices, split_fn=split_fn
    )

    return prefetch(iterable, buffer_size, device=None, num_threads=num_threads)


def replicate_in_all_devices(
    nest: N, devices: Optional[Sequence[jax.Device]] = None
) -> N:
    """Replicate array nest in all available devices."""
    devices = devices or jax.local_devices()
    return jax.device_put_sharded([nest] * len(devices), devices)


def get_from_first_device(nest: N, as_numpy: bool = True) -> N:
    """Gets the first array of a nest of `jax.Array`s.

    Args:
      nest: A nest of `jax.Array`s.
      as_numpy: If `True` then each `DeviceArray` that is retrieved is transformed
        (and copied if not on the host machine) into a `np.ndarray`.

    Returns:
      The first array of a nest of `jax.Array`s. Note that if
      `as_numpy=False` then the array will be a `DeviceArray` (which will live on
      the same device as the sharded device array). If `as_numpy=True` then the
      array will be copied to the host machine and converted into a `np.ndarray`.
    """
    zeroth_nest = jax.tree_map(lambda x: x[0], nest)
    return jax.device_get(zeroth_nest) if as_numpy else zeroth_nest


def mapreduce(
    f: F,
    reduce_fn: Optional[Callable[[jax.Array], jax.Array]] = None,
    **vmap_kwargs,
) -> F:
    """A simple decorator that transforms `f` into (`reduce_fn` o vmap o f).

    By default, we vmap over axis 0, and the `reduce_fn` is jnp.mean over axis 0.
    Note that the call signature of `f` is invariant under this transformation.

    If, for example, f has shape signature [H, W] -> [N], then mapreduce(f)
    (with the default arguments) will have shape signature [B, H, W] -> [N].

    Args:
      f: A pure function over examples.
      reduce_fn: A pure function that reduces DeviceArrays -> DeviceArrays.
      **vmap_kwargs: Keyword arguments to forward to `jax.vmap`.

    Returns:
      g: A pure function over batches of examples.
    """

    if reduce_fn is None:
        reduce_fn = lambda x: jnp.mean(x, axis=0)

    vmapped_f = jax.vmap(f, **vmap_kwargs)

    def g(*args, **kwargs):
        return jax.tree_map(reduce_fn, vmapped_f(*args, **kwargs))

    return g


_TrainingState = TypeVar("_TrainingState")
_TrainingData = TypeVar("_TrainingData")
_TrainingAux = TypeVar("_TrainingAux")


# TODO(b/192806089): migrate all callers to process_many_batches and remove this
# method.
def process_multiple_batches(
    process_one_batch: Callable[
        [_TrainingState, _TrainingData], Tuple[_TrainingState, _TrainingAux]
    ],
    num_batches: int,
    postprocess_aux: Optional[Callable[[_TrainingAux], _TrainingAux]] = None,
) -> Callable[[_TrainingState, _TrainingData], Tuple[_TrainingState, _TrainingAux]]:
    """Makes 'process_one_batch' process multiple batches at once.

    Args:
      process_one_batch: a function that takes 'state' and 'data', and returns
        'new_state' and 'aux' (for example 'metrics').
      num_batches: how many batches to process at once
      postprocess_aux: how to merge the extra information, defaults to taking the
        mean.

    Returns:
      A function with the same interface as 'process_one_batch' which processes
      multiple batches at once.
    """
    assert num_batches >= 1
    if num_batches == 1:
        if not postprocess_aux:
            return process_one_batch

        def _process_one_batch(state, data):
            state, aux = process_one_batch(state, data)
            return state, postprocess_aux(aux)

        return _process_one_batch

    if postprocess_aux is None:
        postprocess_aux = lambda x: jax.tree_map(jnp.mean, x)

    def _process_multiple_batches(state, data):
        data = jax.tree_map(
            lambda a: jnp.reshape(a, (num_batches, -1, *a.shape[1:])), data
        )

        state, aux = jax.lax.scan(process_one_batch, state, data, length=num_batches)
        return state, postprocess_aux(aux)

    return _process_multiple_batches


def process_many_batches(
    process_one_batch: Callable[
        [_TrainingState, _TrainingData], jax_types.TrainingStepOutput[_TrainingState]
    ],
    num_batches: int,
    postprocess_aux: Optional[
        Callable[[jax_types.TrainingMetrics], jax_types.TrainingMetrics]
    ] = None,
) -> Callable[
    [_TrainingState, _TrainingData], jax_types.TrainingStepOutput[_TrainingState]
]:
    """The version of 'process_multiple_batches' with stronger typing."""

    def _process_one_batch(
        state: _TrainingState, data: _TrainingData
    ) -> Tuple[_TrainingState, jax_types.TrainingMetrics]:
        result = process_one_batch(state, data)
        return result.state, result.metrics

    func = process_multiple_batches(_process_one_batch, num_batches, postprocess_aux)

    def _process_many_batches(
        state: _TrainingState, data: _TrainingData
    ) -> jax_types.TrainingStepOutput[_TrainingState]:
        state, aux = func(state, data)
        return jax_types.TrainingStepOutput(state, aux)

    return _process_many_batches


def weighted_softmax(x: jnp.ndarray, weights: jnp.ndarray, axis: int = 0):
    x = x - jnp.max(x, axis=axis)
    return (
        weights * jnp.exp(x) / jnp.sum(weights * jnp.exp(x), axis=axis, keepdims=True)
    )


def sample_uint32(random_key: jax_types.PRNGKey) -> int:
    """Returns an integer uniformly distributed in 0..2^32-1."""
    iinfo = jnp.iinfo(jnp.int32)
    # randint only accepts int32 values as min and max.
    jax_random = jax.random.randint(
        random_key, shape=(), minval=iinfo.min, maxval=iinfo.max, dtype=jnp.int32
    )
    return np.uint32(jax_random).item()


class PrefetchIterator(core.PrefetchingIterator):
    """Performs prefetching from an iterable in separate threads.

    Its interface is additionally extended with `ready` method which tells whether
    there is any data waiting for processing and a `retrieved_elements` method
    specifying number of elements retrieved from the iterator.

    Yields:
      Prefetched elements from the original iterable.

    Raises:
      ValueError: if the buffer_size < 1.
      StopIteration: If the iterable contains no more items.
      Any error thrown by the iterable_function. Note this is not raised inside
        the producer, but after it finishes executing.
    """

    def __init__(
        self,
        iterable: Iterable[types.NestedArray],
        buffer_size: int = 5,
        device: Optional[jax.Device] = None,
        num_threads: int = NUM_PREFETCH_THREADS,
    ):
        """Constructs PrefetchIterator.

        Args:
          iterable: A python iterable. This is used to build the python prefetcher.
            Note that each iterable should only be passed to this function once as
            iterables aren't thread safe.
          buffer_size (int): Number of elements to keep in the prefetch buffer.
          device (deprecated): Optionally place items from the iterable on the given
            device. If None, the items are returns as given by the iterable. This
            argument is deprecated and the recommended usage is to wrap the
            iterables using utils.device_put or utils.multi_device_put before using
            utils.prefetch.
          num_threads (int): Number of threads.
        """

        if buffer_size < 1:
            raise ValueError("the buffer_size should be >= 1")
        self.buffer = queue.Queue(maxsize=buffer_size)
        self.producer_error = []
        self.end = object()
        self.iterable = iterable
        self.device = device
        self.count = 0

        # Start producer threads.
        for _ in range(num_threads):
            threading.Thread(target=self.producer, daemon=True).start()

    def producer(self):
        """Enqueues items from `iterable` on a given thread."""
        try:
            # Build a new iterable for each thread. This is crucial if working with
            # tensorflow datasets because tf.Graph objects are thread local.
            for item in self.iterable:
                if self.device:
                    jax.device_put(item, self.device)
                self.buffer.put(item)
        except Exception as e:  # pylint: disable=broad-except
            logging.exception("Error in producer thread for %s", self.iterable)
            self.producer_error.append(e)
        finally:
            self.buffer.put(self.end)

    def __iter__(self):
        return self

    def ready(self):
        return not self.buffer.empty()

    def retrieved_elements(self):
        return self.count

    def __next__(self):
        value = self.buffer.get()
        if value is self.end:
            if self.producer_error:
                raise self.producer_error[0] from self.producer_error[0]
            raise StopIteration
        self.count += 1
        return value

corax/jax/utils.py
ethanluoyc/corax
# 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.

"""Tests for utils."""

import chex
import jax
import jax.numpy as jnp
import numpy as np
from absl.testing import absltest

from corax.jax import utils

chex.set_n_cpu_devices(4)


class JaxUtilsTest(absltest.TestCase):
    def test_batch_concat(self):
        batch_size = 32
        inputs = [
            jnp.zeros(shape=(batch_size, 2)),
            {"foo": jnp.zeros(shape=(batch_size, 5, 3))},
            [jnp.zeros(shape=(batch_size, 1))],
            jnp.zeros(shape=(batch_size,)),
        ]

        output_shape = utils.batch_concat(inputs).shape
        expected_shape = [batch_size, 2 + 5 * 3 + 1 + 1]
        self.assertSequenceEqual(output_shape, expected_shape)

    def test_mapreduce(self):
        @utils.mapreduce
        def f(y, x):
            return jnp.square(x + y)

        z = f(jnp.ones(shape=(32,)), jnp.ones(shape=(32,)))
        z = jax.device_get(z)
        self.assertEqual(z, 4)

    def test_get_from_first_device(self):
        sharded = {
            "a": jax.device_put_sharded(
                list(jnp.arange(16).reshape([jax.local_device_count(), 4])),
                jax.local_devices(),
            ),
            "b": jax.device_put_sharded(
                list(jnp.arange(8).reshape([jax.local_device_count(), 2])),
                jax.local_devices(),
            ),
        }

        want = {
            "a": jnp.arange(4),
            "b": jnp.arange(2),
        }

        # Get zeroth device content as DeviceArray.
        device_arrays = utils.get_from_first_device(sharded, as_numpy=False)
        jax.tree_map(lambda x: self.assertIsInstance(x, jax.Array), device_arrays)
        jax.tree_map(np.testing.assert_array_equal, want, device_arrays)

        # Get the zeroth device content as numpy arrays.
        numpy_arrays = utils.get_from_first_device(sharded, as_numpy=True)
        jax.tree_map(lambda x: self.assertIsInstance(x, np.ndarray), numpy_arrays)
        jax.tree_map(np.testing.assert_array_equal, want, numpy_arrays)


if __name__ == "__main__":
    absltest.main()

corax/jax/utils_test.py
ethanluoyc/corax
# type: ignore
# 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.

"""Variable utilities for JAX."""

import datetime
import time
from concurrent import futures
from typing import List, NamedTuple, Optional, Sequence, Union

import jax

from corax import core
from corax.jax import networks as network_types


class VariableReference(NamedTuple):
    variable_name: str


class ReferenceVariableSource(core.VariableSource):
    """Variable source which returns references instead of values.

    This is passed to each actor when using a centralized inference server. The
    actor uses this special variable source to get references rather than values.
    These references are then passed to calls to the inference server, which will
    dereference them to obtain the value of the corresponding variables at
    inference time. This avoids passing around copies of variables from each
    actor to the inference server.
    """

    def get_variables(self, names: Sequence[str]) -> List[VariableReference]:
        return [VariableReference(name) for name in names]


class VariableClient:
    """A variable client for updating variables from a remote source."""

    def __init__(
        self,
        client: core.VariableSource,
        key: Union[str, Sequence[str]],
        update_period: Union[int, datetime.timedelta] = 1,
        device: Optional[Union[str, jax.Device]] = None,
    ):
        """Initializes the variable client.

        Args:
          client: A variable source from which we fetch variables.
          key: Which variables to request. When multiple keys are used, params
            property will return a list of params.
          update_period: Interval between fetches, specified as either (int) a
            number of calls to update() between actual fetches or (timedelta) a time
            interval that has to pass since the last fetch.
          device: The name of a JAX device to put variables on. If None (default),
            VariableClient won't put params on any device.
        """
        self._update_period = update_period
        self._call_counter = 0
        self._last_call = time.time()
        self._client = client
        self._params: Sequence[network_types.Params] = None

        self._device = device
        if isinstance(self._device, str):
            self._device = jax.devices(device)[0]

        self._executor = futures.ThreadPoolExecutor(max_workers=1)

        if isinstance(key, str):
            key = [key]

        self._key = key
        self._request = lambda k=key: client.get_variables(k)
        self._future: Optional[futures.Future] = None  # pylint: disable=g-bare-generic
        self._async_request = lambda: self._executor.submit(self._request)

    def update(self, wait: bool = False) -> None:
        """Periodically updates the variables with the latest copy from the source.

        If wait is True, a blocking request is executed. Any active request will be
        cancelled.
        If wait is False, this method makes an asynchronous request for variables.

        Args:
          wait: Whether to execute asynchronous (False) or blocking updates (True).
            Defaults to False.
        """
        # Track calls (we only update periodically).
        self._call_counter += 1

        # Return if it's not time to fetch another update.
        if isinstance(self._update_period, datetime.timedelta):
            if self._update_period.total_seconds() + self._last_call > time.time():
                return
        else:
            if self._call_counter < self._update_period:
                return

        if wait:
            if self._future is not None:
                if self._future.running():
                    self._future.cancel()
                self._future = None
            self._call_counter = 0
            self._last_call = time.time()
            self.update_and_wait()
            return

        # Return early if we are still waiting for a previous request to come back.
        if self._future and not self._future.done():
            return

        # Get a future and add the copy function as a callback.
        self._call_counter = 0
        self._last_call = time.time()
        self._future = self._async_request()
        self._future.add_done_callback(lambda f: self._callback(f.result()))

    def update_and_wait(self):
        """Immediately update and block until we get the result."""
        self._callback(self._request())

    def _callback(self, params_list: List[network_types.Params]):
        if self._device and not isinstance(self._client, ReferenceVariableSource):
            # Move variables to a proper device.
            self._params = jax.device_put(params_list, self._device)
        else:
            self._params = params_list

    @property
    def device(self) -> Optional[jax.Device]:
        return self._device

    @property
    def params(self) -> Union[network_types.Params, List[network_types.Params]]:
        """Returns the first params for one key, otherwise the whole params list."""
        if self._params is None:
            self.update_and_wait()

        if len(self._params) == 1:
            return self._params[0]
        else:
            return self._params

corax/jax/variable_utils.py
ethanluoyc/corax
# 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.

"""Objects which specify the input/output spaces of an environment.

This module exposes the same spec classes as `dm_env` as well as providing an
additional `EnvironmentSpec` class which collects all of the specs for a given
environment. An `EnvironmentSpec` instance can be created directly or by using
the `make_environment_spec` helper given a `dm_env.Environment` instance.
"""

from typing import Any, NamedTuple

import dm_env
from dm_env import specs

Array = specs.Array
BoundedArray = specs.BoundedArray
DiscreteArray = specs.DiscreteArray


class EnvironmentSpec(NamedTuple):
    """Full specification of the domains used by a given environment."""

    observations: Any
    actions: Any
    rewards: Any
    discounts: Any


def make_environment_spec(environment: dm_env.Environment) -> EnvironmentSpec:
    """Returns an `EnvironmentSpec` describing values used by an environment."""
    return EnvironmentSpec(
        observations=environment.observation_spec(),
        actions=environment.action_spec(),
        rewards=environment.reward_spec(),
        discounts=environment.discount_spec(),
    )

corax/specs.py
ethanluoyc/corax
"""Testing helpers."""

corax/testing/__init__.py
ethanluoyc/corax
# type: ignore
# 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.

"""Fake (mock) components.

Minimal implementations of fake Acme components which can be instantiated in
order to test or interact with other components.
"""

import threading
from typing import Callable, Iterator, List, Mapping, Optional, Sequence

import dm_env
import numpy as np
import reverb
import tensorflow as tf
import tree
from rlds import rlds_types

from corax import core
from corax import specs
from corax import types
from corax import wrappers
from corax.utils import loggers


class Logger(loggers.Logger):
    """A fake logger for testing."""

    def __init__(self):
        self.data = []

    def write(self, data):
        self.data.append(data)

    @property
    def last_write(self):
        return self.data[-1]

    def close(self):
        pass


class Actor(core.Actor):
    """Fake actor which generates random actions and validates specs."""

    def __init__(self, spec: specs.EnvironmentSpec):
        self._spec = spec
        self.num_updates = 0

    def select_action(self, observation: types.NestedArray) -> types.NestedArray:
        _validate_spec(self._spec.observations, observation)
        return _generate_from_spec(self._spec.actions)

    def observe_first(self, timestep: dm_env.TimeStep):
        _validate_spec(self._spec.observations, timestep.observation)

    def observe(
        self,
        action: types.NestedArray,
        next_timestep: dm_env.TimeStep,
    ):
        _validate_spec(self._spec.actions, action)
        _validate_spec(self._spec.rewards, next_timestep.reward)
        _validate_spec(self._spec.discounts, next_timestep.discount)
        _validate_spec(self._spec.observations, next_timestep.observation)

    def update(self, wait: bool = False):
        self.num_updates += 1


class VariableSource(core.VariableSource):
    """Fake variable source."""

    def __init__(
        self,
        variables: Optional[types.NestedArray] = None,
        barrier: Optional[threading.Barrier] = None,
        use_default_key: bool = True,
    ):
        # Add dummy variables so we can expose them in get_variables.
        if use_default_key:
            self._variables = {"policy": [] if variables is None else variables}
        else:
            self._variables = variables
        self._barrier = barrier

    def get_variables(self, names: Sequence[str]) -> List[types.NestedArray]:
        if self._barrier is not None:
            self._barrier.wait()
        return [self._variables[name] for name in names]


class Learner(core.Learner, VariableSource):
    """Fake Learner."""

    def __init__(
        self,
        variables: Optional[types.NestedArray] = None,
        barrier: Optional[threading.Barrier] = None,
    ):
        super().__init__(variables=variables, barrier=barrier)
        self.step_counter = 0

    def step(self):
        self.step_counter += 1


class Environment(dm_env.Environment):
    """A fake environment with a given spec."""

    def __init__(
        self,
        spec: specs.EnvironmentSpec,
        *,
        episode_length: int = 25,
    ):
        # Assert that the discount spec is a BoundedArray with range [0, 1].
        def check_discount_spec(path, discount_spec):
            if (
                not isinstance(discount_spec, specs.BoundedArray)
                or not np.isclose(discount_spec.minimum, 0)
                or not np.isclose(discount_spec.maximum, 1)
            ):
                if path:
                    path_str = " " + "/".join(str(p) for p in path)
                else:
                    path_str = ""
                raise ValueError(
                    "discount_spec {}isn't a BoundedArray in [0, 1].".format(path_str)
                )

        tree.map_structure_with_path(check_discount_spec, spec.discounts)

        self._spec = spec
        self._episode_length = episode_length
        self._step = 0

    def _generate_fake_observation(self):
        return _generate_from_spec(self._spec.observations)

    def _generate_fake_reward(self):
        return _generate_from_spec(self._spec.rewards)

    def _generate_fake_discount(self):
        return _generate_from_spec(self._spec.discounts)

    def reset(self) -> dm_env.TimeStep:
        observation = self._generate_fake_observation()
        self._step = 1
        return dm_env.restart(observation)

    def step(self, action) -> dm_env.TimeStep:
        # Return a reset timestep if we haven't touched the environment yet.
        if not self._step:
            return self.reset()

        _validate_spec(self._spec.actions, action)

        observation = self._generate_fake_observation()
        reward = self._generate_fake_reward()
        discount = self._generate_fake_discount()

        if self._episode_length and (self._step == self._episode_length):
            self._step = 0
            # We can't use dm_env.termination directly because then the discount
            # wouldn't necessarily conform to the spec (if eg. we want float32).
            return dm_env.TimeStep(dm_env.StepType.LAST, reward, discount, observation)
        else:
            self._step += 1
            return dm_env.transition(
                reward=reward, observation=observation, discount=discount
            )

    def action_spec(self):
        return self._spec.actions

    def observation_spec(self):
        return self._spec.observations

    def reward_spec(self):
        return self._spec.rewards

    def discount_spec(self):
        return self._spec.discounts


class _BaseDiscreteEnvironment(Environment):
    """Discrete action fake environment."""

    def __init__(
        self,
        *,
        num_actions: int = 1,
        action_dtype=np.int32,
        observation_spec: types.NestedSpec,
        discount_spec: Optional[types.NestedSpec] = None,
        reward_spec: Optional[types.NestedSpec] = None,
        **kwargs,
    ):
        """Initialize the environment."""
        if reward_spec is None:
            reward_spec = specs.Array((), np.float32)

        if discount_spec is None:
            discount_spec = specs.BoundedArray((), np.float32, 0.0, 1.0)

        actions = specs.DiscreteArray(num_actions, dtype=action_dtype)

        super().__init__(
            spec=specs.EnvironmentSpec(
                observations=observation_spec,
                actions=actions,
                rewards=reward_spec,
                discounts=discount_spec,
            ),
            **kwargs,
        )


class DiscreteEnvironment(_BaseDiscreteEnvironment):
    """Discrete state and action fake environment."""

    def __init__(
        self,
        *,
        num_actions: int = 1,
        num_observations: int = 1,
        action_dtype=np.int32,
        obs_dtype=np.int32,
        obs_shape: Sequence[int] = (),
        discount_spec: Optional[types.NestedSpec] = None,
        reward_spec: Optional[types.NestedSpec] = None,
        **kwargs,
    ):
        """Initialize the environment."""
        observations_spec = specs.BoundedArray(
            shape=obs_shape,
            dtype=obs_dtype,
            minimum=obs_dtype(0),
            maximum=obs_dtype(num_observations - 1),
        )

        super().__init__(
            num_actions=num_actions,
            action_dtype=action_dtype,
            observation_spec=observations_spec,
            discount_spec=discount_spec,
            reward_spec=reward_spec,
            **kwargs,
        )


class NestedDiscreteEnvironment(_BaseDiscreteEnvironment):
    """Discrete action fake environment with nested discrete state."""

    def __init__(
        self,
        *,
        num_observations: Mapping[str, int],
        num_actions: int = 1,
        action_dtype=np.int32,
        obs_dtype=np.int32,
        obs_shape: Sequence[int] = (),
        discount_spec: Optional[types.NestedSpec] = None,
        reward_spec: Optional[types.NestedSpec] = None,
        **kwargs,
    ):
        """Initialize the environment."""

        observations_spec = {}
        for key in num_observations:
            observations_spec[key] = specs.BoundedArray(
                shape=obs_shape,
                dtype=obs_dtype,
                minimum=obs_dtype(0),
                maximum=obs_dtype(num_observations[key] - 1),
            )

        super().__init__(
            num_actions=num_actions,
            action_dtype=action_dtype,
            observation_spec=observations_spec,
            discount_spec=discount_spec,
            reward_spec=reward_spec,
            **kwargs,
        )


class ContinuousEnvironment(Environment):
    """Continuous state and action fake environment."""

    def __init__(
        self,
        *,
        action_dim: int = 1,
        observation_dim: int = 1,
        bounded: bool = False,
        dtype=np.float32,
        reward_dtype=np.float32,
        **kwargs,
    ):
        """Initialize the environment.

        Args:
          action_dim: number of action dimensions.
          observation_dim: number of observation dimensions.
          bounded: whether or not the actions are bounded in [-1, 1].
          dtype: dtype of the action and observation spaces.
          reward_dtype: dtype of the reward and discounts.
          **kwargs: additional kwargs passed to the Environment base class.
        """

        action_shape = () if action_dim == 0 else (action_dim,)
        observation_shape = () if observation_dim == 0 else (observation_dim,)

        observations = specs.Array(observation_shape, dtype)
        rewards = specs.Array((), reward_dtype)
        discounts = specs.BoundedArray((), reward_dtype, 0.0, 1.0)

        if bounded:
            actions = specs.BoundedArray(action_shape, dtype, -1.0, 1.0)
        else:
            actions = specs.Array(action_shape, dtype)

        super().__init__(
            spec=specs.EnvironmentSpec(
                observations=observations,
                actions=actions,
                rewards=rewards,
                discounts=discounts,
            ),
            **kwargs,
        )


def _validate_spec(spec: types.NestedSpec, value: types.NestedArray):
    """Validate a value from a potentially nested spec."""
    tree.assert_same_structure(value, spec)
    tree.map_structure(lambda s, v: s.validate(v), spec, value)


def _normalize_array(array: specs.Array) -> specs.Array:
    """Converts bounded arrays with (-inf,+inf) bounds to unbounded arrays.

    The returned array should be mostly equivalent to the input, except that
    `generate_value()` returns -infs on arrays bounded to (-inf,+inf) and zeros
    on unbounded arrays.

    Args:
      array: the array to be normalized.

    Returns:
      normalized array.
    """
    if isinstance(array, specs.DiscreteArray):
        return array
    if not isinstance(array, specs.BoundedArray):
        return array
    if not (array.minimum == float("-inf")).all():
        return array
    if not (array.maximum == float("+inf")).all():
        return array
    return specs.Array(array.shape, array.dtype, array.name)


def _generate_from_spec(spec: types.NestedSpec) -> types.NestedArray:
    """Generate a value from a potentially nested spec."""
    return tree.map_structure(lambda s: _normalize_array(s).generate_value(), spec)


def transition_dataset_from_spec(spec: specs.EnvironmentSpec) -> tf.data.Dataset:
    """Constructs fake dataset of Reverb N-step transition samples.

    Args:
      spec: Constructed fake transitions match the provided specification.

    Returns:
      tf.data.Dataset that produces the same fake N-step transition ReverbSample
      object indefinitely.
    """

    observation = _generate_from_spec(spec.observations)
    action = _generate_from_spec(spec.actions)
    reward = _generate_from_spec(spec.rewards)
    discount = _generate_from_spec(spec.discounts)
    data = types.Transition(observation, action, reward, discount, observation)

    info = tree.map_structure(
        lambda tf_dtype: tf.ones([], tf_dtype.as_numpy_dtype),
        reverb.SampleInfo.tf_dtypes(),
    )
    sample = reverb.ReplaySample(info=info, data=data)

    return tf.data.Dataset.from_tensors(sample).repeat()


def transition_dataset(environment: dm_env.Environment) -> tf.data.Dataset:
    """Constructs fake dataset of Reverb N-step transition samples.

    Args:
      environment: Constructed fake transitions will match the specification of
        this environment.

    Returns:
      tf.data.Dataset that produces the same fake N-step transition ReverbSample
      object indefinitely.
    """
    return transition_dataset_from_spec(specs.make_environment_spec(environment))


def transition_iterator_from_spec(
    spec: specs.EnvironmentSpec,
) -> Callable[[int], Iterator[types.Transition]]:
    """Constructs fake iterator of transitions.

    Args:
      spec: Constructed fake transitions match the provided specification..

    Returns:
      A callable that given a batch_size returns an iterator of transitions.
    """

    observation = _generate_from_spec(spec.observations)
    action = _generate_from_spec(spec.actions)
    reward = _generate_from_spec(spec.rewards)
    discount = _generate_from_spec(spec.discounts)
    data = types.Transition(observation, action, reward, discount, observation)

    dataset = tf.data.Dataset.from_tensors(data).repeat()

    return lambda batch_size: dataset.batch(batch_size).as_numpy_iterator()


def transition_iterator(
    environment: dm_env.Environment,
) -> Callable[[int], Iterator[types.Transition]]:
    """Constructs fake iterator of transitions.

    Args:
      environment: Constructed fake transitions will match the specification of
        this environment.

    Returns:
      A callable that given a batch_size returns an iterator of transitions.
    """
    return transition_iterator_from_spec(specs.make_environment_spec(environment))


def fake_atari_wrapped(
    episode_length: int = 10, oar_wrapper: bool = False
) -> dm_env.Environment:
    """Builds fake version of the environment to be used by tests.

    Args:
      episode_length: The length of episodes produced by this environment.
      oar_wrapper: Should ObservationActionRewardWrapper be applied.

    Returns:
      Fake version of the environment equivalent to the one returned by
      env_loader.load_atari_wrapped
    """
    env = DiscreteEnvironment(
        num_actions=18,
        num_observations=2,
        obs_shape=(84, 84, 4),
        obs_dtype=np.float32,
        episode_length=episode_length,
    )

    if oar_wrapper:
        env = wrappers.ObservationActionRewardWrapper(env)
    return env


def rlds_dataset_from_env_spec(
    spec: specs.EnvironmentSpec,
    *,
    episode_count: int = 10,
    episode_length: int = 25,
) -> tf.data.Dataset:
    """Constructs a fake RLDS dataset with the given spec.

    Args:
      spec: specification to use for generation of fake steps.
      episode_count: number of episodes in the dataset.
      episode_length: length of the episode in the dataset.

    Returns:
      a fake RLDS dataset.
    """

    fake_steps = {
        rlds_types.OBSERVATION: (
            [_generate_from_spec(spec.observations)] * episode_length
        ),
        rlds_types.ACTION: ([_generate_from_spec(spec.actions)] * episode_length),
        rlds_types.REWARD: ([_generate_from_spec(spec.rewards)] * episode_length),
        rlds_types.DISCOUNT: ([_generate_from_spec(spec.discounts)] * episode_length),
        rlds_types.IS_TERMINAL: [False] * (episode_length - 1) + [True],
        rlds_types.IS_FIRST: [True] + [False] * (episode_length - 1),
        rlds_types.IS_LAST: [False] * (episode_length - 1) + [True],
    }
    steps_dataset = tf.data.Dataset.from_tensor_slices(fake_steps)

    return tf.data.Dataset.from_tensor_slices(
        {rlds_types.STEPS: [steps_dataset] * episode_count}
    )

corax/testing/fakes.py
ethanluoyc/corax
"""Testing utilities."""

import sys
from typing import Optional

from absl import flags
from absl.testing import parameterized


class TestCase(parameterized.TestCase):
    """A custom TestCase which handles FLAG parsing for pytest compatibility."""

    def get_tempdir(self, name: Optional[str] = None) -> str:
        try:
            flags.FLAGS.test_tmpdir
        except flags.UnparsedFlagAccessError:
            # Need to initialize flags when running `pytest`.
            flags.FLAGS(sys.argv, known_only=True)
        return self.create_tempdir(name).full_path

corax/testing/test_utils.py
ethanluoyc/corax
# 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.

"""Common types used throughout Corax."""

from typing import Any, Callable, Iterable, Mapping, NamedTuple, Union

from corax import specs

NestedArray = Any
NestedTensor = Any

NestedSpec = Union[specs.Array, Iterable["NestedSpec"], Mapping[Any, "NestedSpec"]]

Nest = Union[NestedArray, NestedTensor, NestedSpec]

TensorTransformation = Callable[[NestedTensor], NestedTensor]
TensorValuedCallable = Callable[..., NestedTensor]


class Transition(NamedTuple):
    """Container for a transition."""

    observation: NestedArray
    action: NestedArray
    reward: NestedArray
    discount: NestedArray
    next_observation: NestedArray
    extras: NestedArray = ()

corax/types.py
ethanluoyc/corax
"""Corax utility functions."""

corax/utils/__init__.py
ethanluoyc/corax
# 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.

"""Utilities to use within loggers."""

import queue
import threading
from typing import Callable, Generic, TypeVar

from absl import logging

E = TypeVar("E")


class AsyncExecutor(Generic[E]):
    """Executes a blocking function asynchronously on a queue of items."""

    def __init__(
        self,
        fn: Callable[[E], None],
        queue_size: int = 1,
        interruptible_interval_secs: float = 1.0,
    ):
        """Buffers elements in a queue and runs `fn` asynchronously..

        NOTE: Once closed, `AsyncExecutor` will block until current `fn` finishes
          but is not guaranteed to dequeue all elements currently stored in
          the data queue. This is intentional so as to prevent a blocking `fn` call
          from preventing `AsyncExecutor` from closing.

        Args:
          fn: A callable to be executed upon dequeuing an element from data
            queue.
          queue_size: The maximum size of the synchronized buffer queue.
          interruptible_interval_secs: Timeout interval in seconds for blocking
            queue operations after which the background threads check for errors and
            if background threads should stop.
        """
        self._data = queue.Queue(maxsize=queue_size)
        self._should_stop = threading.Event()
        self._errors = queue.Queue()
        self._interruptible_interval_secs = interruptible_interval_secs

        def _dequeue() -> None:
            """Dequeue data from a queue and invoke blocking call."""
            while not self._should_stop.is_set():
                try:
                    element = self._data.get(timeout=self._interruptible_interval_secs)
                    # Execute fn upon dequeuing an element from the data queue.
                    fn(element)
                except queue.Empty:
                    # If queue is Empty for longer than the specified time interval,
                    # check again if should_stop has been requested and retry.
                    continue
                except Exception as e:
                    logging.error("AsyncExecuter thread terminated with error.")
                    logging.exception(e)
                    self._errors.put(e)
                    self._should_stop.set()
                    raise  # Never caught by anything, just terminates the thread.

        self._thread = threading.Thread(target=_dequeue, daemon=True)
        self._thread.start()

    def _raise_on_error(self) -> None:
        try:
            # Raise the error on the caller thread if an error has been raised in the
            # looper thread.
            raise self._errors.get_nowait()
        except queue.Empty:
            pass

    def close(self):
        self._should_stop.set()
        # Join all background threads.
        self._thread.join()
        # Raise errors produced by background threads.
        self._raise_on_error()

    def put(self, element: E) -> None:
        """Puts `element` asynchronuously onto the underlying data queue.

        The write call blocks if the underlying data_queue contains `queue_size`
          elements for over `self._interruptible_interval_secs` second, in which
          case we check if stop has been requested or if there has been an error
          raised on the looper thread. If neither happened, retry enqueue.

        Args:
          element: an element to be put into the underlying data queue and dequeued
            asynchronuously for `fn(element)` call.
        """
        while not self._should_stop.is_set():
            try:
                self._data.put(element, timeout=self._interruptible_interval_secs)
                break
            except queue.Full:
                continue
        else:
            # If `should_stop` has been set, then raises if any has been raised on
            # the background thread.
            self._raise_on_error()

corax/utils/async_utils.py
ethanluoyc/corax
# type: ignore
# 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.

"""A simple, hierarchical distributed counter."""

import threading
import time
from typing import Dict, Mapping, Optional, Union

from corax import core

Number = Union[int, float]


class Counter(core.Saveable):
    """A simple counter object that can periodically sync with a parent."""

    def __init__(
        self,
        parent: Optional["Counter"] = None,
        prefix: str = "",
        time_delta: float = 1.0,
        return_only_prefixed: bool = False,
    ):
        """Initialize the counter.

        Args:
          parent: a Counter object to cache locally (or None for no caching).
          prefix: string prefix to use for all local counts.
          time_delta: time difference in seconds between syncing with the parent
            counter.
          return_only_prefixed: if True, and if `prefix` isn't empty, return counts
            restricted to the given `prefix` on each call to `increment` and
            `get_counts`. The `prefix` is stripped from returned count names.
        """

        self._parent = parent
        self._prefix = prefix
        self._time_delta = time_delta

        # Hold local counts and we'll lock around that.
        # These are counts to be synced to the parent and the cache.
        self._counts = {}
        self._lock = threading.Lock()

        # We'll sync periodically (when the last sync was more than self._time_delta
        # seconds ago.)
        self._cache = {}
        self._last_sync_time = 0.0

        self._return_only_prefixed = return_only_prefixed

    def increment(self, **counts: Number) -> Dict[str, Number]:
        """Increment a set of counters.

        Args:
          **counts: keyword arguments specifying count increments.

        Returns:
          The [name, value] mapping of all counters stored, i.e. this will also
          include counts that were not updated by this call to increment.
        """
        with self._lock:
            for key, value in counts.items():
                self._counts.setdefault(key, 0)
                self._counts[key] += value
        return self.get_counts()

    def get_counts(self) -> Dict[str, Number]:
        """Return all counts tracked by this counter."""
        now = time.time()
        # TODO(b/144421838): use futures instead of blocking.
        if self._parent and (now - self._last_sync_time) > self._time_delta:
            with self._lock:
                counts = _prefix_keys(self._counts, self._prefix)
                # Reset the local counts, as they will be merged into the parent and the
                # cache.
                self._counts = {}
            self._cache = self._parent.increment(**counts)
            self._last_sync_time = now

        # Potentially prefix the keys in the counts dictionary.
        counts = _prefix_keys(self._counts, self._prefix)

        # If there's no prefix make a copy of the dictionary so we don't modify the
        # internal self._counts.
        if not self._prefix:
            counts = dict(counts)

        # Combine local counts with any parent counts.
        for key, value in self._cache.items():
            counts[key] = counts.get(key, 0) + value

        if self._prefix and self._return_only_prefixed:
            counts = dict(
                [
                    (key[len(self._prefix) + 1 :], value)
                    for key, value in counts.items()
                    if key.startswith(f"{self._prefix}_")
                ]
            )
        return counts

    def save(self) -> Mapping[str, Mapping[str, Number]]:
        return {"counts": self._counts, "cache": self._cache}

    def restore(self, state: Mapping[str, Mapping[str, Number]]):
        # Force a sync, if necessary, on the next get_counts call.
        self._last_sync_time = 0.0
        self._counts = state["counts"]
        self._cache = state["cache"]

    def get_steps_key(self) -> str:
        """Returns the key to use for steps by this counter."""
        if not self._prefix or self._return_only_prefixed:
            return "steps"
        return f"{self._prefix}_steps"


def _prefix_keys(dictionary: Dict[str, Number], prefix: str):
    """Return a dictionary with prefixed keys.

    Args:
      dictionary: dictionary to return a copy of.
      prefix: string to use as the prefix.

    Returns:
      Return a copy of the given dictionary whose keys are replaced by
      "{prefix}_{key}". If the prefix is the empty string it returns the given
      dictionary unchanged.
    """
    if prefix:
        dictionary = {f"{prefix}_{k}": v for k, v in dictionary.items()}
    return dictionary

corax/utils/counting.py
ethanluoyc/corax
# 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.

"""Tests for acme.utils.counting."""

import threading

from absl.testing import absltest

from corax.utils import counting


class Barrier:
    """Defines a simple barrier class to synchronize on a particular event."""

    def __init__(self, num_threads):
        """Constructor.

        Args:
          num_threads: int - how many threads will be syncronizing on this barrier
        """
        self._num_threads = num_threads
        self._count = 0
        self._cond = threading.Condition()

    def wait(self):
        """Waits on the barrier until all threads have called this method."""
        with self._cond:
            self._count += 1
            self._cond.notifyAll()
            while self._count < self._num_threads:
                self._cond.wait()


class CountingTest(absltest.TestCase):
    def test_counter_threading(self):
        counter = counting.Counter()
        num_threads = 10
        barrier = Barrier(num_threads)

        # Increment in every thread at the same time.
        def add_to_counter():
            barrier.wait()
            counter.increment(foo=1)

        # Run the threads.
        threads = []
        for _ in range(num_threads):
            t = threading.Thread(target=add_to_counter)
            t.start()
            threads.append(t)
        for t in threads:
            t.join()

        # Make sure the counter has been incremented once per thread.
        counts = counter.get_counts()
        self.assertEqual(counts["foo"], num_threads)

    def test_counter_caching(self):
        parent = counting.Counter()
        counter = counting.Counter(parent, time_delta=0.0)
        counter.increment(foo=12)
        self.assertEqual(parent.get_counts(), counter.get_counts())

    def test_shared_counts(self):
        # Two counters with shared parent should share counts (modulo namespacing).
        parent = counting.Counter()
        child1 = counting.Counter(parent, "child1")
        child2 = counting.Counter(parent, "child2")
        child1.increment(foo=1)
        result = child2.increment(foo=2)
        expected = {"child1_foo": 1, "child2_foo": 2}
        self.assertEqual(result, expected)

    def test_return_only_prefixed(self):
        parent = counting.Counter()
        child1 = counting.Counter(
            parent, "child1", time_delta=0.0, return_only_prefixed=False
        )
        child2 = counting.Counter(
            parent, "child2", time_delta=0.0, return_only_prefixed=True
        )
        child1.increment(foo=1)
        child2.increment(bar=1)
        self.assertEqual(child1.get_counts(), {"child1_foo": 1, "child2_bar": 1})
        self.assertEqual(child2.get_counts(), {"bar": 1})

    def test_get_steps_key(self):
        parent = counting.Counter()
        child1 = counting.Counter(
            parent, "child1", time_delta=0.0, return_only_prefixed=False
        )
        child2 = counting.Counter(
            parent, "child2", time_delta=0.0, return_only_prefixed=True
        )
        self.assertEqual(child1.get_steps_key(), "child1_steps")
        self.assertEqual(child2.get_steps_key(), "steps")
        child1.increment(steps=1)
        child2.increment(steps=2)
        self.assertEqual(child1.get_counts().get(child1.get_steps_key()), 1)
        self.assertEqual(child2.get_counts().get(child2.get_steps_key()), 2)

    def test_parent_prefix(self):
        parent = counting.Counter(prefix="parent")
        child = counting.Counter(parent, prefix="child", time_delta=0.0)
        self.assertEqual(child.get_steps_key(), "child_steps")


if __name__ == "__main__":
    absltest.main()

corax/utils/counting_test.py
ethanluoyc/corax
# Copyright 2020 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.

"""LazyLoader providing support for loading symbols upon first reference."""

import importlib
import sys
import typing

# Set to true in your local workspace to work around issues with lazy imports.
# Please report any cases of usage to stanczyk@, so we can address them.
_DISABLE_LAZY_IMPORTS = False


class LazyImports:
    """Turns all imports performed in its scope into lazy.

    It is used to avoid pulling in large dependencies.
    """

    def __init__(self, parent, add_to_globals=True):
        """Initializes LazyImports.

        Args:
          parent: Name of the parent module doing the imports.
          add_to_globals: Whether imported symbols should be added to importing
              module globals. Rule of thumb is to not add for __init__ modules,
              as it is important for these modules to resolve Lazy symbol upon its
              first access (not when accessing members of the Lazy import).
        """
        self._parent = sys.modules[parent]
        self._globals = vars(self._parent)
        self._add_to_globals = add_to_globals
        self._symbols = dict()

    def parent(self):
        return self._parent

    def _find_and_load(self, name, import_):
        del import_
        return _ModuleToLoad(name, self)

    def _handle_fromlist(self, module, fromlist, import_):
        del import_
        assert len(fromlist) == 1
        if not isinstance(module, _ModuleToLoad):
            return _ModuleToLoad(module.__name__, self)
        return module

    def __enter__(self):
        if typing.TYPE_CHECKING or _DISABLE_LAZY_IMPORTS:
            return
        # Start capturing import statements.
        self._org_find_and_load = importlib._bootstrap._find_and_load
        self._org_handle_fromlist = importlib._bootstrap._handle_fromlist
        importlib._bootstrap._find_and_load = self._find_and_load
        importlib._bootstrap._handle_fromlist = self._handle_fromlist

    def __exit__(self, type_, value, traceback):
        if typing.TYPE_CHECKING or _DISABLE_LAZY_IMPORTS:
            return
        # Stop capturing import statements.
        importlib._bootstrap._find_and_load = self._org_find_and_load
        importlib._bootstrap._handle_fromlist = self._org_handle_fromlist

        # If symbols are not added to globals then __getattr__ has to be patched.
        if not self._add_to_globals:
            self._parent.__getattr__ = self.__getattr__

        # Make sure lazy symbols know their local name within the importing module.
        for name in list(self._globals.keys()):
            member = self._globals[name]
            if isinstance(member, _MemberToLoad):
                member.set_export_name(name)
                if not self._add_to_globals:
                    del self._globals[name]
                    self._symbols[name] = member

    def __getattr__(self, name):
        if name in self._symbols:
            res = self._symbols[name].load()
            self._parent.__dict__[name] = res
            return res
        raise AttributeError(f"module {self._parent!r} has no attribute {name!r}")


class _MemberToLoad:
    """Module member to load in a lazy way."""

    def __init__(self, name, parent):
        self._name = name
        self._export_name = name
        self._parent = parent

    def __getattr__(self, item):
        return getattr(self.load(), item)

    def __call__(self, *args, **kwargs):
        return self.load()(*args, **kwargs)  # type: ignore

    def load(self):
        try:
            res = importlib.import_module(f"{self._parent.name()}.{self._name}")
        except ModuleNotFoundError:
            res = importlib.import_module(self._parent.name()).__dict__[self._name]

        self._parent.parent().parent().__dict__[self._export_name] = res
        return res

    def set_export_name(self, name):
        self._export_name = name


class _ModuleToLoad:
    """Module to load in a lazy way."""

    def __init__(self, name, parent):
        self._name = name
        self._parent = parent

    def name(self):
        return self._name

    def parent(self):
        return self._parent

    def __getattr__(self, item):
        return _MemberToLoad(item, self)

corax/utils/lazy_loader.py
ethanluoyc/corax
"""Loggers."""

from corax.utils.loggers.aggregators import Dispatcher
from corax.utils.loggers.asynchronous import AsyncLogger
from corax.utils.loggers.auto_close import AutoCloseLogger
from corax.utils.loggers.base import Logger
from corax.utils.loggers.base import LoggerFactory
from corax.utils.loggers.base import LoggerLabel
from corax.utils.loggers.base import LoggerStepsKey
from corax.utils.loggers.base import LoggingData
from corax.utils.loggers.base import NoOpLogger
from corax.utils.loggers.base import TaskInstance
from corax.utils.loggers.base import to_numpy
from corax.utils.loggers.csv import CSVLogger
from corax.utils.loggers.default import make_default_logger
from corax.utils.loggers.filters import GatedFilter
from corax.utils.loggers.filters import KeyFilter
from corax.utils.loggers.filters import NoneFilter
from corax.utils.loggers.filters import TimeFilter
from corax.utils.loggers.terminal import TerminalLogger
from corax.utils.loggers.wandb import WandbLogger

corax/utils/loggers/__init__.py
ethanluoyc/corax
# 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.

"""Utilities for aggregating to other loggers."""

from typing import Any, Callable, Optional, Sequence

from corax.utils.loggers import base


class Dispatcher(base.Logger):
    """Writes data to multiple `Logger` objects."""

    def __init__(
        self,
        to: Sequence[base.Logger],
        serialize_fn: Optional[Callable[[base.LoggingData], Any]] = None,
    ):
        """Initialize `Dispatcher` connected to several `Logger` objects."""
        self._to = to
        self._serialize_fn = serialize_fn

    def write(self, values: base.LoggingData):
        """Writes `values` to the underlying `Logger` objects."""
        if self._serialize_fn:
            values = self._serialize_fn(values)
        for logger in self._to:
            logger.write(values)

    def close(self):
        for logger in self._to:
            logger.close()

corax/utils/loggers/aggregators.py
ethanluoyc/corax
# 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.

"""Logger which makes another logger asynchronous."""

from typing import Any, Mapping

from corax.utils import async_utils
from corax.utils.loggers import base


class AsyncLogger(base.Logger):
    """Logger which makes the logging to another logger asyncronous."""

    def __init__(self, to: base.Logger):
        """Initializes the logger.

        Args:
          to: A `Logger` object to which the current object will forward its results
            when `write` is called.
        """
        self._to = to
        self._async_worker = async_utils.AsyncExecutor(self._to.write, queue_size=5)

    def write(self, values: Mapping[str, Any]):
        self._async_worker.put(values)

    def close(self):
        """Closes the logger, closing is synchronous."""
        self._async_worker.close()
        self._to.close()

corax/utils/loggers/asynchronous.py
ethanluoyc/corax
# 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.

"""Logger which self closes on exit if not closed yet."""

import weakref

from corax.utils.loggers import base


class AutoCloseLogger(base.Logger):
    """Logger which auto closes itself on exit if not already closed."""

    def __init__(self, logger: base.Logger):
        self._logger = logger
        # The finalizer "logger.close" is invoked in one of the following scenario:
        # 1) the current logger is GC
        # 2) from the python doc, when the program exits, each remaining live
        #    finalizer is called.
        # Note that in the normal flow, where "close" is explicitly called,
        # the finalizer is marked as dead using the detach function so that
        # the underlying logger is not closed twice (once explicitly and once
        # implicitly when the object is GC or when the program exits).
        self._finalizer = weakref.finalize(self, logger.close)

    def write(self, values: base.LoggingData):
        if self._logger is not None:
            self._logger.write(values)

    def close(self):
        if self._finalizer.detach():
            self._logger.close()  # type: ignore
        self._logger = None

corax/utils/loggers/auto_close.py
ethanluoyc/corax
# 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.

"""Base logger."""

import abc
from typing import Any, Mapping, Optional

import numpy as np
import tree
from typing_extensions import Protocol

LoggingData = Mapping[str, Any]


class Logger(abc.ABC):
    """A logger has a `write` method."""

    @abc.abstractmethod
    def write(self, data: LoggingData) -> None:
        """Writes `data` to destination (file, terminal, database, etc)."""

    @abc.abstractmethod
    def close(self) -> None:
        """Closes the logger, not expecting any further write."""


TaskInstance = int
LoggerLabel = str
LoggerStepsKey = str


class LoggerFactory(Protocol):
    def __call__(
        self,
        label: LoggerLabel,
        steps_key: Optional[LoggerStepsKey] = None,
        instance: Optional[TaskInstance] = None,
    ) -> Logger:
        ...


class NoOpLogger(Logger):
    """Simple Logger which does nothing and outputs no logs.

    This should be used sparingly, but it can prove useful if we want to quiet an
    individual component and have it produce no logging whatsoever.
    """

    def write(self, data: LoggingData):
        pass

    def close(self):
        pass


def tensor_to_numpy(value: Any):
    if hasattr(value, "numpy"):
        return value.numpy()  # tf.Tensor (TF2).
    if hasattr(value, "device_buffer"):
        return np.asarray(value)  # jnp.DeviceArray.
    return value


def to_numpy(values: Any):
    """Converts tensors in a nested structure to numpy.

    Converts tensors from TensorFlow to Numpy if needed without importing TF
    dependency.

    Args:
      values: nested structure with numpy and / or TF tensors.

    Returns:
      Same nested structure as values, but with numpy tensors.
    """
    return tree.map_structure(tensor_to_numpy, values)

corax/utils/loggers/base.py
ethanluoyc/corax
# 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.

"""A simple CSV logger.

Warning: Does not support preemption.
"""

import csv
import os
import time
from typing import TextIO, Union

from absl import logging

from corax.utils import paths
from corax.utils.loggers import base


class CSVLogger(base.Logger):
    """Standard CSV logger.

    The fields are inferred from the first call to write() and any additional
    fields afterwards are ignored.

    """

    _open = open

    def __init__(
        self,
        directory_or_file: Union[str, TextIO],
        label: str = "",
        time_delta: float = 0.0,
        add_uid: bool = True,
        flush_every: int = 30,
    ):
        """Instantiates the logger.

        Args:
          directory_or_file: Either a directory path as a string, or a file TextIO
            object.
          label: Extra label to add to logger. This is added as a suffix to the
            directory.
          time_delta: Interval in seconds between which writes are dropped to
            throttle throughput.
          add_uid: Whether to add a UID to the file path. See `paths.process_path`
            for details.
          flush_every: Interval (in writes) between flushes.
        """

        if flush_every <= 0:
            raise ValueError(
                f"`flush_every` must be a positive integer (got {flush_every})."
            )

        self._last_log_time = time.time() - time_delta
        self._time_delta = time_delta
        self._flush_every = flush_every
        self._add_uid = add_uid
        self._writer = None
        self._file_owner = False
        self._file = self._create_file(directory_or_file, label)
        self._writes = 0
        logging.info("Logging to %s", self.file_path)

    def _create_file(
        self,
        directory_or_file: Union[str, TextIO],
        label: str,
    ) -> TextIO:
        """Opens a file if input is a directory or use existing file."""
        if isinstance(directory_or_file, str):
            directory = paths.process_path(
                directory_or_file, "logs", label, add_uid=self._add_uid
            )
            file_path = os.path.join(directory, "logs.csv")
            self._file_owner = True
            return self._open(file_path, mode="a")  # type: ignore

        # TextIO instance.
        file = directory_or_file
        if label:
            logging.info("File, not directory, passed to CSVLogger; label not used.")
        if not file.mode.startswith("a"):
            raise ValueError(
                "File must be open in append mode; instead got " f'mode="{file.mode}".'
            )
        return file

    def write(self, data: base.LoggingData):
        """Writes a `data` into a row of comma-separated values."""
        # Only log if `time_delta` seconds have passed since last logging event.
        now = time.time()

        # TODO(b/192227744): Remove this in favour of filters.TimeFilter.
        elapsed = now - self._last_log_time
        if elapsed < self._time_delta:
            logging.debug(
                "Not due to log for another %.2f seconds, dropping data.",
                self._time_delta - elapsed,
            )
            return
        self._last_log_time = now

        # Append row to CSV.
        data = base.to_numpy(data)  # type: ignore
        # Use fields from initial `data` to create the header. If extra fields are
        # present in subsequent `data`, we ignore them.
        if not self._writer:
            fields = sorted(data.keys())
            self._writer = csv.DictWriter(
                self._file, fieldnames=fields, extrasaction="ignore"
            )
            # Write header only if the file is empty.
            if not self._file.tell():
                self._writer.writeheader()
        self._writer.writerow(data)

        # Flush every `flush_every` writes.
        if self._writes % self._flush_every == 0:
            self.flush()
        self._writes += 1

    def close(self):
        self.flush()
        if self._file_owner:
            self._file.close()

    def flush(self):
        self._file.flush()

    @property
    def file_path(self) -> str:
        return self._file.name

corax/utils/loggers/csv.py
ethanluoyc/corax
# 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.

"""Default logger."""

import logging
from typing import Any, Callable, Mapping, Optional

from corax.utils.loggers import aggregators
from corax.utils.loggers import asynchronous as async_logger
from corax.utils.loggers import base
from corax.utils.loggers import filters
from corax.utils.loggers import terminal


def make_default_logger(
    label: str,
    time_delta: float = 1.0,
    asynchronous: bool = False,
    print_fn: Optional[Callable[[str], None]] = None,
    serialize_fn: Optional[Callable[[Mapping[str, Any]], Any]] = base.to_numpy,
    steps_key: str = "steps",
) -> base.Logger:
    """Makes a default Acme logger.

    Args:
      label: Name to give to the logger.
      save_data: Whether to persist data.
      time_delta: Time (in seconds) between logging events.
      asynchronous: Whether the write function should block or not.
      print_fn: How to print to terminal (defaults to print).
      serialize_fn: An optional function to apply to the write inputs before
        passing them to the various loggers.
      steps_key: Ignored.

    Returns:
      A logger object that responds to logger.write(some_dict).
    """
    del steps_key
    if not print_fn:
        print_fn = logging.info
    terminal_logger = terminal.TerminalLogger(label=label, print_fn=print_fn)

    loggers = [terminal_logger]

    # Dispatch to all writers and filter Nones and by time.
    logger = aggregators.Dispatcher(loggers, serialize_fn)
    logger = filters.NoneFilter(logger)
    if asynchronous:
        logger = async_logger.AsyncLogger(logger)
    logger = filters.TimeFilter(logger, time_delta)

    return logger

corax/utils/loggers/default.py
ethanluoyc/corax
# 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.

"""Loggers which filter other loggers."""

import math
import time
from typing import Callable, Optional, Sequence

from corax.utils.loggers import base


class NoneFilter(base.Logger):
    """Logger which writes to another logger, filtering any `None` values."""

    def __init__(self, to: base.Logger):
        """Initializes the logger.

        Args:
          to: A `Logger` object to which the current object will forward its results
            when `write` is called.
        """
        self._to = to

    def write(self, values: base.LoggingData):
        values = {k: v for k, v in values.items() if v is not None}
        self._to.write(values)

    def close(self):
        self._to.close()


class TimeFilter(base.Logger):
    """Logger which writes to another logger at a given time interval."""

    def __init__(self, to: base.Logger, time_delta: float):
        """Initializes the logger.

        Args:
          to: A `Logger` object to which the current object will forward its results
            when `write` is called.
          time_delta: How often to write values out in seconds.
            Note that writes within `time_delta` are dropped.
        """
        self._to = to
        self._time = 0
        self._time_delta = time_delta
        if time_delta < 0:
            raise ValueError(f"time_delta must be greater than 0 (got {time_delta}).")

    def write(self, values: base.LoggingData):
        now = time.time()
        if (now - self._time) > self._time_delta:
            self._to.write(values)
            self._time = now

    def close(self):
        self._to.close()


class KeyFilter(base.Logger):
    """Logger which filters keys in logged data."""

    def __init__(
        self,
        to: base.Logger,
        *,
        keep: Optional[Sequence[str]] = None,
        drop: Optional[Sequence[str]] = None,
    ):
        """Creates the filter.

        Args:
          to: A `Logger` object to which the current object will forward its writes.
          keep: Keys that are kept by the filter. Note that `keep` and `drop` cannot
            be both set at once.
          drop: Keys that are dropped by the filter. Note that `keep` and `drop`
            cannot be both set at once.
        """
        if bool(keep) == bool(drop):
            raise ValueError("Exactly one of `keep` & `drop` arguments must be set.")
        self._to = to
        self._keep = keep
        self._drop = drop

    def write(self, data: base.LoggingData):
        if self._keep:
            data = {k: data[k] for k in self._keep}
        if self._drop:
            data = {k: v for k, v in data.items() if k not in self._drop}
        self._to.write(data)

    def close(self):
        self._to.close()


class GatedFilter(base.Logger):
    """Logger which writes to another logger based on a gating function.

    This logger tracks the number of times its `write` method is called, and uses
    a gating function on this number to decide when to write.
    """

    def __init__(self, to: base.Logger, gating_fn: Callable[[int], bool]):
        """Initialises the logger.

        Args:
          to: A `Logger` object to which the current object will forward its results
            when `write` is called.
          gating_fn: A function that takes an integer (number of calls) as input.
            For example, to log every tenth call: gating_fn=lambda t: t % 10 == 0.
        """
        self._to = to
        self._gating_fn = gating_fn
        self._calls = 0

    def write(self, values: base.LoggingData):
        if self._gating_fn(self._calls):
            self._to.write(values)
        self._calls += 1

    def close(self):
        self._to.close()

    @classmethod
    def logarithmic(cls, to: base.Logger, n: int = 10) -> "GatedFilter":
        """Builds a logger for writing at logarithmically-spaced intervals.

        This will log on a linear scale at each order of magnitude of `n`.
        For example, with n=10, this will log at times:
            [0, 1, 2, ..., 9, 10, 20, 30, ... 90, 100, 200, 300, ... 900, 1000]

        Args:
          to: The underlying logger to write to.
          n: Base (default 10) on which to operate.
        Returns:
          A GatedFilter logger, which gates logarithmically as described above.
        """

        def logarithmic_filter(t: int) -> bool:
            magnitude = math.floor(math.log10(max(t, 1)) / math.log10(n))
            return t % (n**magnitude) == 0

        return cls(to, gating_fn=logarithmic_filter)

    @classmethod
    def periodic(cls, to: base.Logger, interval: int = 10) -> "GatedFilter":
        """Builds a logger for writing at linearly-spaced intervals.

        Args:
          to: The underlying logger to write to.
          interval: The interval between writes.
        Returns:
          A GatedFilter logger, which gates periodically as described above.
        """
        return cls(to, gating_fn=lambda t: t % interval == 0)

corax/utils/loggers/filters.py
ethanluoyc/corax
# 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.

"""Tests for logging filters."""

import time

from absl.testing import absltest

from corax.utils.loggers import base
from corax.utils.loggers import filters


# TODO(jaslanides): extract this to test_utils, or similar, for re-use.
class FakeLogger(base.Logger):
    """A fake logger for testing."""

    def __init__(self):
        self.data = []

    def write(self, data):
        self.data.append(data)

    @property
    def last_write(self):
        return self.data[-1]

    def close(self):
        pass


class GatedFilterTest(absltest.TestCase):
    def test_logarithmic_filter(self):
        logger = FakeLogger()
        filtered = filters.GatedFilter.logarithmic(logger, n=10)
        for t in range(100):
            filtered.write({"t": t})
        rows = [row["t"] for row in logger.data]
        self.assertEqual(rows, [*range(10), *range(10, 100, 10)])

    def test_periodic_filter(self):
        logger = FakeLogger()
        filtered = filters.GatedFilter.periodic(logger, interval=10)
        for t in range(100):
            filtered.write({"t": t})
        rows = [row["t"] for row in logger.data]
        self.assertEqual(rows, list(range(0, 100, 10)))


class TimeFilterTest(absltest.TestCase):
    def test_delta(self):
        logger = FakeLogger()
        filtered = filters.TimeFilter(logger, time_delta=0.1)

        # Logged.
        filtered.write({"foo": 1})
        self.assertIn("foo", logger.last_write)

        # *Not* logged.
        filtered.write({"bar": 2})
        self.assertNotIn("bar", logger.last_write)

        # Wait out delta.
        time.sleep(0.11)

        # Logged.
        filtered.write({"baz": 3})
        self.assertIn("baz", logger.last_write)

        self.assertLen(logger.data, 2)


class KeyFilterTest(absltest.TestCase):
    def test_keep_filter(self):
        logger = FakeLogger()
        filtered = filters.KeyFilter(logger, keep=("foo",))
        filtered.write({"foo": "bar", "baz": 12})
        row, *_ = logger.data
        self.assertIn("foo", row)
        self.assertNotIn("baz", row)

    def test_drop_filter(self):
        logger = FakeLogger()
        filtered = filters.KeyFilter(logger, drop=("foo",))
        filtered.write({"foo": "bar", "baz": 12})
        row, *_ = logger.data
        self.assertIn("baz", row)
        self.assertNotIn("foo", row)

    def test_bad_arguments(self):
        with self.assertRaises(ValueError):
            filters.KeyFilter(FakeLogger())
        with self.assertRaises(ValueError):
            filters.KeyFilter(FakeLogger(), keep=("a",), drop=("b",))


if __name__ == "__main__":
    absltest.main()

corax/utils/loggers/filters_test.py
ethanluoyc/corax
# 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.

"""Utilities for logging to the terminal."""

import logging
import time
from typing import Any, Callable

import numpy as np

from corax.utils.loggers import base


def _format_key(key: str) -> str:
    """Internal function for formatting keys."""
    return key.replace("_", " ").title()


def _format_value(value: Any) -> str:
    """Internal function for formatting values."""
    value = base.to_numpy(value)
    if isinstance(value, (float, np.number)):
        return f"{value:0.3f}"
    return f"{value}"


def serialize(values: base.LoggingData) -> str:
    """Converts `values` to a pretty-printed string.

    This takes a dictionary `values` whose keys are strings and returns
    a formatted string such that each [key, value] pair is separated by ' = ' and
    each entry is separated by ' | '. The keys are sorted alphabetically to ensure
    a consistent order, and snake case is split into words.

    For example:

        values = {'a': 1, 'b' = 2.33333333, 'c': 'hello', 'big_value': 10}
        # Returns 'A = 1 | B = 2.333 | Big Value = 10 | C = hello'
        values_string = serialize(values)

    Args:
      values: A dictionary with string keys.

    Returns:
      A formatted string.
    """
    return " | ".join(
        f"{_format_key(k)} = {_format_value(v)}" for k, v in sorted(values.items())
    )


class TerminalLogger(base.Logger):
    """Logs to terminal."""

    def __init__(
        self,
        label: str = "",
        print_fn: Callable[[str], None] = logging.info,
        serialize_fn: Callable[[base.LoggingData], str] = serialize,
        time_delta: float = 0.0,
    ):
        """Initializes the logger.

        Args:
          label: label string to use when logging.
          print_fn: function to call which acts like print.
          serialize_fn: function to call which transforms values into a str.
          time_delta: How often (in seconds) to write values. This can be used to
            minimize terminal spam, but is 0 by default---ie everything is written.
        """

        self._print_fn = print_fn
        self._serialize_fn = serialize_fn
        self._label = label and f"[{_format_key(label)}] "
        self._time = time.time()
        self._time_delta = time_delta

    def write(self, values: base.LoggingData):
        now = time.time()
        if (now - self._time) > self._time_delta:
            self._print_fn(f"{self._label}{self._serialize_fn(values)}")
            self._time = now

    def close(self):
        pass

corax/utils/loggers/terminal.py
ethanluoyc/corax
# 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.

"""Tests for terminal logger."""

from absl.testing import absltest

from corax.utils.loggers import terminal


class LoggingTest(absltest.TestCase):
    def test_logging_output_format(self):
        inputs = {
            "c": "foo",
            "a": 1337,
            "b": 42.0001,
        }
        expected_outputs = "A = 1337 | B = 42.000 | C = foo"
        test_fn = lambda outputs: self.assertEqual(outputs, expected_outputs)

        logger = terminal.TerminalLogger(print_fn=test_fn)
        logger.write(inputs)

    def test_label(self):
        inputs = {"foo": "bar", "baz": 123}
        expected_outputs = "[Test] Baz = 123 | Foo = bar"
        test_fn = lambda outputs: self.assertEqual(outputs, expected_outputs)

        logger = terminal.TerminalLogger(print_fn=test_fn, label="test")
        logger.write(inputs)


if __name__ == "__main__":
    absltest.main()

corax/utils/loggers/terminal_test.py
ethanluoyc/corax
"""Logger that writes logs to WandB"""
from typing import Any, Optional

from absl import logging

from corax.utils.loggers import base

try:
    import wandb  # pytype: disable=import-error
except ImportError:
    wandb = None


class WandbLogger(base.Logger):
    """Logging results to weights and biases"""

    def __init__(
        self,
        label: Optional[str] = None,
        steps_key: Optional[str] = None,
        *,
        project: Optional[str] = None,
        entity: Optional[str] = None,
        dir: Optional[str] = None,  # pylint: disable=redefined-builtin
        name: Optional[str] = None,
        group: Optional[str] = None,
        config: Optional[Any] = None,
        run: Optional[Any] = None,
        **wandb_kwargs,
    ):
        if wandb is None:
            raise ImportError(
                "Logger not supported as `wandb` logger is not installed yet,"
                " install it with `pip install wandb`."
            )
        self._label = label
        self._iter = 0
        self._steps_key = steps_key
        if run is not None:
            self._run = run
        elif wandb.run is None:
            self._run = wandb.init(
                project=project,
                dir=dir,
                entity=entity,
                name=name,
                group=group,
                config=config,
                **wandb_kwargs,
            )
        else:
            self._run = wandb.run
        # define default x-axis (for latest wandb versions)
        if steps_key and getattr(self._run, "define_metric", None):
            prefix = f"{self._label}/*" if self._label else "*"
            self._run.define_metric(  # type: ignore
                prefix, step_metric=f"{self._label}/{self._steps_key}"
            )

    @property
    def run(self):
        """Return the current wandb run."""
        return self._run

    def write(self, data: base.LoggingData):
        data = base.to_numpy(data)  # type: ignore
        if self._steps_key is not None and self._steps_key not in data:
            logging.warn("steps key %s not found. Skip logging.", self._steps_key)
            return
        if self._label:
            stats = {f"{self._label}/{k}": v for k, v in data.items()}
        else:
            stats = data
        self._run.log(stats)  # type: ignore

        self._iter += 1

    def close(self):
        pass

corax/utils/loggers/wandb.py
ethanluoyc/corax
"""Corax observers."""

from corax.utils.observers.base import EnvLoopObserver
from corax.utils.observers.base import Number

corax/utils/observers/__init__.py
ethanluoyc/corax
"""Metrics observers."""

import abc
from typing import Dict, Union

import dm_env
import numpy as np

Number = Union[int, float]


class EnvLoopObserver(abc.ABC):
    """An interface for collecting metrics/counters in EnvironmentLoop."""

    @abc.abstractmethod
    def observe_first(self, env: dm_env.Environment, timestep: dm_env.TimeStep) -> None:
        """Observes the initial state."""

    @abc.abstractmethod
    def observe(
        self, env: dm_env.Environment, timestep: dm_env.TimeStep, action: np.ndarray
    ) -> None:
        """Records one environment step."""

    @abc.abstractmethod
    def get_metrics(self) -> Dict[str, Number]:
        """Returns metrics collected for the current episode."""

corax/utils/observers/base.py
ethanluoyc/corax
# 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.

"""Filesystem path helpers."""

import os
import os.path
import shutil
import time
from typing import Optional, Tuple

from absl import flags

ACME_ID = flags.DEFINE_string(
    "corax_id", None, "Experiment identifier to use for Acme."
)


def process_path(
    path: str,
    *subpaths: str,
    ttl_seconds: Optional[int] = None,
    backups: Optional[bool] = None,
    add_uid: bool = True,
) -> str:
    """Process the path string.

    This will process the path string by running `os.path.expanduser` to replace
    any initial "~". It will also append a unique string on the end of the path
    and create the directories leading to this path if necessary.

    Args:
      path: string defining the path to process and create.
      *subpaths: potential subpaths to include after uniqification.
      ttl_seconds: ignored.
      backups: ignored.
      add_uid: Whether to add a unique directory identifier between `path` and
        `subpaths`. If the `--acme_id` flag is set, will use that as the
        identifier.

    Returns:
      the processed, expanded path string.
    """
    del backups, ttl_seconds

    path = os.path.expanduser(path)
    if add_uid:
        path = os.path.join(path, *get_unique_id())
    path = os.path.join(path, *subpaths)
    os.makedirs(path, exist_ok=True)
    return path


_DATETIME = time.strftime("%Y%m%d-%H%M%S")


def get_unique_id() -> Tuple[str, ...]:
    """Makes a unique identifier for this process; override with --acme_id."""
    # By default we'll use the global id.
    identifier = _DATETIME

    # If the --acme_id flag is given prefer that; ignore if flag processing has
    # been skipped (this happens in colab or in tests).
    try:
        identifier = ACME_ID.value or identifier
    except flags.UnparsedFlagAccessError:
        pass

    # Return as a tuple (for future proofing).
    return (identifier,)


def rmdir(path: str):
    """Remove directory recursively."""
    shutil.rmtree(path)

corax/utils/paths.py
ethanluoyc/corax
import sys
from typing import Sequence, Tuple

import reverb


def disable_insert_blocking(
    tables: Sequence[reverb.Table],
) -> Tuple[Sequence[reverb.Table], Sequence[int]]:
    """Disables blocking of insert operations for a given collection of tables."""
    modified_tables = []
    sample_sizes = []
    for table in tables:
        rate_limiter_info = table.info.rate_limiter_info
        rate_limiter = reverb.rate_limiters.RateLimiter(
            samples_per_insert=rate_limiter_info.samples_per_insert,
            min_size_to_sample=rate_limiter_info.min_size_to_sample,
            min_diff=rate_limiter_info.min_diff,
            max_diff=sys.float_info.max,
        )
        modified_tables.append(table.replace(rate_limiter=rate_limiter))
        # Target the middle of the rate limiter's insert-sample balance window.
        sample_sizes.append(
            max(1, int((rate_limiter_info.max_diff - rate_limiter_info.min_diff) / 2))
        )
    return modified_tables, sample_sizes

corax/utils/reverb_utils.py
ethanluoyc/corax
# 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.

"""Tensor framework-agnostic utilities for manipulating nested structures."""

from typing import Any, List, Sequence, TypeVar

import numpy as np
import tree

ElementType = TypeVar("ElementType")


def fast_map_structure(func, *structure):
    """Faster map_structure implementation which skips some error checking."""
    flat_structure = (tree.flatten(s) for s in structure)
    entries = zip(*flat_structure)
    # Arbitrarily choose one of the structures of the original sequence (the last)
    # to match the structure for the flattened sequence.
    return tree.unflatten_as(structure[-1], [func(*x) for x in entries])


def fast_map_structure_with_path(func, *structure):
    """Faster map_structure_with_path implementation."""
    head_entries_with_path = tree.flatten_with_path(structure[0])
    if len(structure) > 1:
        tail_entries = (tree.flatten(s) for s in structure[1:])
        entries_with_path = [
            e[0] + e[1:] for e in zip(head_entries_with_path, *tail_entries)
        ]
    else:
        entries_with_path = head_entries_with_path
    # Arbitrarily choose one of the structures of the original sequence (the last)
    # to match the structure for the flattened sequence.
    return tree.unflatten_as(structure[-1], [func(*x) for x in entries_with_path])


def stack_sequence_fields(sequence: Sequence[ElementType]) -> ElementType:
    """Stacks a list of identically nested objects.

    This takes a sequence of identically nested objects and returns a single
    nested object whose ith leaf is a stacked numpy array of the corresponding
    ith leaf from each element of the sequence.

    For example, if `sequence` is:

    ```python
    [{
          'action': np.array([1.0]),
          'observation': (np.array([0.0, 1.0, 2.0]),),
          'reward': 1.0
     }, {
          'action': np.array([0.5]),
          'observation': (np.array([1.0, 2.0, 3.0]),),
          'reward': 0.0
     }, {
          'action': np.array([0.3]),1
          'observation': (np.array([2.0, 3.0, 4.0]),),
          'reward': 0.5
     }]
    ```

    Then this function will return:

    ```python
    {
        'action': np.array([....])         # array shape = [3 x 1]
        'observation': (np.array([...]),)  # array shape = [3 x 3]
        'reward': np.array([...])          # array shape = [3]
    }
    ```

    Note that the 'observation' entry in the above example has two levels of
    nesting, i.e it is a tuple of arrays.

    Args:
      sequence: a list of identically nested objects.

    Returns:
      A nested object with numpy.

    Raises:
      ValueError: If `sequence` is an empty sequence.
    """
    # Handle empty input sequences.
    if not sequence:
        raise ValueError("Input sequence must not be empty")

    # Default to asarray when arrays don't have the same shape to be compatible
    # with old behaviour.
    try:
        return fast_map_structure(lambda *values: np.stack(values), *sequence)  # type: ignore
    except ValueError:
        return fast_map_structure(
            lambda *values: np.asarray(values, dtype=object), *sequence  # type: ignore
        )


def unstack_sequence_fields(struct: ElementType, batch_size: int) -> List[ElementType]:
    """Converts a struct of batched arrays to a list of structs.

    This is effectively the inverse of `stack_sequence_fields`.

    Args:
      struct: An (arbitrarily nested) structure of arrays.
      batch_size: The length of the leading dimension of each array in the struct.
        This is assumed to be static and known.

    Returns:
      A list of structs with the same structure as `struct`, where each leaf node
       is an unbatched element of the original leaf node.
    """

    return [tree.map_structure(lambda s, i=i: s[i], struct) for i in range(batch_size)]  # type: ignore


def broadcast_structures(*args: Any) -> Any:
    """Returns versions of the arguments that give them the same nested structure.

    Any nested items in *args must have the same structure.

    Any non-nested item will be replaced with a nested version that shares that
    structure. The leaves will all be references to the same original non-nested
    item.

    If all *args are nested, or all *args are non-nested, this function will
    return *args unchanged.

    Example:
    ```
    a = ('a', 'b')
    b = 'c'
    tree_a, tree_b = broadcast_structure(a, b)
    tree_a
    > ('a', 'b')
    tree_b
    > ('c', 'c')
    ```

    Args:
      *args: A Sequence of nested or non-nested items.

    Returns:
      `*args`, except with all items sharing the same nest structure.
    """
    if not args:
        return

    reference_tree = None
    for arg in args:
        if tree.is_nested(arg):
            reference_tree = arg
            break

    # If reference_tree is None then none of args are nested and we can skip over
    # the rest of this function, which would be a no-op.
    if reference_tree is None:
        return args

    def mirror_structure(value, reference_tree):
        if tree.is_nested(value):
            # Use check_types=True so that the types of the trees we construct aren't
            # dependent on our arbitrary choice of which nested arg to use as the
            # reference_tree.
            tree.assert_same_structure(value, reference_tree, check_types=True)
            return value
        else:
            return tree.map_structure(lambda _: value, reference_tree)

    return tuple(mirror_structure(arg, reference_tree) for arg in args)


def tree_map(f):
    """Transforms `f` into a tree-mapped version."""

    def mapped_f(*structures):
        return tree.map_structure(f, *structures)

    return mapped_f

corax/utils/tree_utils.py
ethanluoyc/corax
from corax.utils import lazy_loader

with lazy_loader.LazyImports(__name__, False):
    from corax.wrappers.gym_wrapper import GymWrapper
    from corax.wrappers.gymnasium_wrapper import GymnasiumWrapper

del lazy_loader

corax/wrappers/__init__.py
ethanluoyc/corax
# 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.

"""Wraps an OpenAI Gym environment to be used as a dm_env environment."""

from typing import Any, Dict, Optional

import dm_env
import gym
import numpy as np
import tree
from gym import spaces

from corax import specs
from corax import types


class GymWrapper(dm_env.Environment):
    """Environment wrapper for OpenAI Gym environments.

    NOTE: This currently only supports GymV21 environments
    """

    # Note: we don't inherit from base.EnvironmentWrapper because that class
    # assumes that the wrapped environment is a dm_env.Environment.

    def __init__(self, environment: gym.Env):
        self._environment = environment
        self._reset_next_step = True
        self._last_info = None

        # Convert action and observation specs.
        obs_space = self._environment.observation_space
        act_space = self._environment.action_space
        self._observation_spec = _convert_to_spec(obs_space, name="observation")
        self._action_spec = _convert_to_spec(act_space, name="action")

    def reset(self) -> dm_env.TimeStep:
        """Resets the episode."""
        self._reset_next_step = False
        observation = self._environment.reset()
        # Reset the diagnostic information.
        self._last_info = None
        return dm_env.restart(observation)

    def step(self, action: types.NestedArray) -> dm_env.TimeStep:
        """Steps the environment."""
        if self._reset_next_step:
            return self.reset()

        observation, reward, done, info = self._environment.step(action)
        self._reset_next_step = done
        self._last_info = info

        # Convert the type of the reward based on the spec, respecting the scalar or
        # array property.
        reward = tree.map_structure(
            lambda x, t: (  # pylint: disable=g-long-lambda
                t.dtype.type(x) if np.isscalar(x) else np.asarray(x, dtype=t.dtype)
            ),
            reward,
            self.reward_spec(),
        )

        if done:
            truncated = info.get("TimeLimit.truncated", False)
            if truncated:
                return dm_env.truncation(reward, observation)
            return dm_env.termination(reward, observation)
        return dm_env.transition(reward, observation)

    def observation_spec(self) -> types.NestedSpec:
        return self._observation_spec

    def action_spec(self) -> types.NestedSpec:
        return self._action_spec

    def get_info(self) -> Optional[Dict[str, Any]]:
        """Returns the last info returned from env.step(action).

        Returns:
          info: dictionary of diagnostic information from the last environment step
        """
        return self._last_info

    @property
    def environment(self) -> gym.Env:
        """Returns the wrapped environment."""
        return self._environment

    def __getattr__(self, name: str):
        if name.startswith("__"):
            raise AttributeError(
                "attempted to get missing private attribute '{}'".format(name)
            )
        return getattr(self._environment, name)

    def close(self):
        self._environment.close()


def _convert_to_spec(space: gym.Space, name: Optional[str] = None) -> types.NestedSpec:
    """Converts an OpenAI Gym space to a dm_env spec or nested structure of specs.

    Box, MultiBinary and MultiDiscrete Gym spaces are converted to BoundedArray
    specs. Discrete OpenAI spaces are converted to DiscreteArray specs. Tuple and
    Dict spaces are recursively converted to tuples and dictionaries of specs.

    Args:
      space: The Gym space to convert.
      name: Optional name to apply to all return spec(s).

    Returns:
      A dm_env spec or nested structure of specs, corresponding to the input
      space.
    """
    if isinstance(space, spaces.Discrete):
        return specs.DiscreteArray(num_values=space.n, dtype=space.dtype, name=name)  # type: ignore

    elif isinstance(space, spaces.Box):
        return specs.BoundedArray(
            shape=space.shape,
            dtype=space.dtype,
            minimum=space.low,
            maximum=space.high,
            name=name,
        )

    elif isinstance(space, spaces.MultiBinary):
        return specs.BoundedArray(
            shape=space.shape, dtype=space.dtype, minimum=0.0, maximum=1.0, name=name
        )

    elif isinstance(space, spaces.MultiDiscrete):
        return specs.BoundedArray(
            shape=space.shape,
            dtype=space.dtype,
            minimum=np.zeros(space.shape),
            maximum=space.nvec - 1,
            name=name,
        )

    elif isinstance(space, spaces.Tuple):
        return tuple(_convert_to_spec(s, name) for s in space.spaces)

    elif isinstance(space, spaces.Dict):
        return {
            key: _convert_to_spec(value, key) for key, value in space.spaces.items()
        }

    else:
        raise ValueError("Unexpected gym space: {}".format(space))

corax/wrappers/gym_wrapper.py
ethanluoyc/corax
# type: ignore
"""Tests for gymnasium_wrapper."""

import gym
import numpy as np
from absl.testing import absltest
from dm_env import specs

from corax.wrappers import gym_wrapper


class GymWrapperTest(absltest.TestCase):
    def test_gym_cartpole(self):
        env = gym.GymWrapper(gym.make("CartPole-v0"))

        # Test converted observation spec.
        observation_spec: specs.BoundedArray = env.observation_spec()
        self.assertEqual(type(observation_spec), specs.BoundedArray)
        self.assertEqual(observation_spec.shape, (4,))
        self.assertEqual(observation_spec.minimum.shape, (4,))
        self.assertEqual(observation_spec.maximum.shape, (4,))
        self.assertEqual(observation_spec.dtype, np.dtype("float32"))

        # Test converted action spec.
        action_spec: specs.BoundedArray = env.action_spec()
        self.assertEqual(type(action_spec), specs.DiscreteArray)
        self.assertEqual(action_spec.shape, ())
        self.assertEqual(action_spec.minimum, 0)
        self.assertEqual(action_spec.maximum, 1)
        self.assertEqual(action_spec.num_values, 2)
        self.assertEqual(action_spec.dtype, np.dtype("int64"))

        # Test step.
        timestep = env.reset()
        self.assertTrue(timestep.first())
        timestep = env.step(1)
        self.assertEqual(timestep.reward, 1.0)
        self.assertTrue(np.isscalar(timestep.reward))
        self.assertEqual(timestep.observation.shape, (4,))
        env.close()

    def test_early_truncation(self):
        # Pendulum has no early termination condition. Recent versions of gym force
        # to use v1. We try both in case an earlier version is installed.
        try:
            gym_env = gym.make("Pendulum-v1")
        except:  # noqa
            gym_env = gym.make("Pendulum-v0")
        env = gym_wrapper.GymWrapper(gym_env)
        ts = env.reset()
        while not ts.last():
            ts = env.step(env.action_spec().generate_value())
        self.assertEqual(ts.discount, 1.0)
        self.assertTrue(np.isscalar(ts.reward))
        env.close()

    def test_multi_discrete(self):
        space = gym.spaces.MultiDiscrete([2, 3])
        spec = gym_wrapper._convert_to_spec(space)

        spec.validate([0, 0])
        spec.validate([1, 2])

        self.assertRaises(ValueError, spec.validate, [2, 2])
        self.assertRaises(ValueError, spec.validate, [1, 3])


if __name__ == "__main__":
    absltest.main()

corax/wrappers/gym_wrapper_test .py
ethanluoyc/corax
"""Wraps an Farama Gymnasium environment to be used as a dm_env environment."""

from typing import Any, Dict, Optional

import dm_env
import gymnasium as gym
import numpy as np
import tree
from gymnasium import spaces

from corax import specs
from corax import types


class GymnasiumWrapper(dm_env.Environment):
    """Environment wrapper for OpenAI Gym environments."""

    # Note: we don't inherit from base.EnvironmentWrapper because that class
    # assumes that the wrapped environment is a dm_env.Environment.

    def __init__(self, environment: gym.Env):
        self._environment = environment
        self._reset_next_step = True
        self._last_info = None

        # Convert action and observation specs.
        obs_space = self._environment.observation_space
        act_space = self._environment.action_space
        self._observation_spec = _convert_to_spec(obs_space, name="observation")
        self._action_spec = _convert_to_spec(act_space, name="action")

    def reset(self) -> dm_env.TimeStep:
        """Resets the episode."""
        self._reset_next_step = False
        observation, info = self._environment.reset()
        # Reset the diagnostic information.
        self._last_info = info
        return dm_env.restart(observation)

    def step(self, action: types.NestedArray) -> dm_env.TimeStep:
        """Steps the environment."""
        if self._reset_next_step:
            return self.reset()

        observation, reward, terminated, truncated, info = self._environment.step(
            action
        )

        self._reset_next_step = terminated or truncated
        self._last_info = info

        # Convert the type of the reward based on the spec, respecting the scalar or
        # array property.
        reward = tree.map_structure(
            lambda x, t: (  # pylint: disable=g-long-lambda
                t.dtype.type(x) if np.isscalar(x) else np.asarray(x, dtype=t.dtype)
            ),
            reward,
            self.reward_spec(),
        )

        if terminated or truncated:
            if truncated:
                return dm_env.truncation(reward, observation)
            return dm_env.termination(reward, observation)
        return dm_env.transition(reward, observation)

    def observation_spec(self) -> types.NestedSpec:
        return self._observation_spec

    def action_spec(self) -> types.NestedSpec:
        return self._action_spec

    def get_info(self) -> Optional[Dict[str, Any]]:
        """Returns the last info returned from env.step(action).

        Returns:
          info: dictionary of diagnostic information from the last environment step
        """
        return self._last_info

    @property
    def environment(self) -> gym.Env:
        """Returns the wrapped environment."""
        return self._environment

    def __getattr__(self, name: str):
        if name.startswith("__"):
            raise AttributeError(
                "attempted to get missing private attribute '{}'".format(name)
            )
        return getattr(self._environment, name)

    def close(self):
        self._environment.close()


def _convert_to_spec(space: gym.Space, name: Optional[str] = None) -> types.NestedSpec:
    """Converts an OpenAI Gym space to a dm_env spec or nested structure of specs.

    Box, MultiBinary and MultiDiscrete Gym spaces are converted to BoundedArray
    specs. Discrete OpenAI spaces are converted to DiscreteArray specs. Tuple and
    Dict spaces are recursively converted to tuples and dictionaries of specs.

    Args:
      space: The Gym space to convert.
      name: Optional name to apply to all return spec(s).

    Returns:
      A dm_env spec or nested structure of specs, corresponding to the input
      space.
    """
    if isinstance(space, spaces.Discrete):
        return specs.DiscreteArray(num_values=space.n, dtype=space.dtype, name=name)  # type: ignore

    elif isinstance(space, spaces.Box):
        return specs.BoundedArray(
            shape=space.shape,
            dtype=space.dtype,
            minimum=space.low,
            maximum=space.high,
            name=name,
        )

    elif isinstance(space, spaces.MultiBinary):
        return specs.BoundedArray(
            shape=space.shape, dtype=space.dtype, minimum=0.0, maximum=1.0, name=name
        )

    elif isinstance(space, spaces.MultiDiscrete):
        return specs.BoundedArray(
            shape=space.shape,
            dtype=space.dtype,
            minimum=np.zeros(space.shape),
            maximum=space.nvec - 1,
            name=name,
        )

    elif isinstance(space, spaces.Tuple):
        return tuple(_convert_to_spec(s, name) for s in space.spaces)

    elif isinstance(space, spaces.Dict):
        return {
            key: _convert_to_spec(value, key) for key, value in space.spaces.items()
        }

    else:
        raise ValueError("Unexpected gym space: {}".format(space))

corax/wrappers/gymnasium_wrapper.py
ethanluoyc/corax
# type: ignore
"""Tests for gymnasium_wrapper."""

import unittest

import numpy as np
from absl.testing import absltest
from dm_env import specs

SKIP_GYM_TESTS = False
SKIP_GYM_MESSAGE = "gym not installed."

try:
    # pylint: disable=g-import-not-at-top
    import gymnasium

    from corax.wrappers import gymnasium_wrapper

    # pylint: enable=g-import-not-at-top
except ModuleNotFoundError:
    SKIP_GYM_TESTS = True


@unittest.skipIf(SKIP_GYM_TESTS, SKIP_GYM_MESSAGE)
class GymWrapperTest(absltest.TestCase):
    def test_gym_cartpole(self):
        env = gymnasium_wrapper.GymnasiumWrapper(gymnasium.make("CartPole-v0"))

        # Test converted observation spec.
        observation_spec: specs.BoundedArray = env.observation_spec()
        self.assertEqual(type(observation_spec), specs.BoundedArray)
        self.assertEqual(observation_spec.shape, (4,))
        self.assertEqual(observation_spec.minimum.shape, (4,))
        self.assertEqual(observation_spec.maximum.shape, (4,))
        self.assertEqual(observation_spec.dtype, np.dtype("float32"))

        # Test converted action spec.
        action_spec: specs.BoundedArray = env.action_spec()
        self.assertEqual(type(action_spec), specs.DiscreteArray)
        self.assertEqual(action_spec.shape, ())
        self.assertEqual(action_spec.minimum, 0)
        self.assertEqual(action_spec.maximum, 1)
        self.assertEqual(action_spec.num_values, 2)
        self.assertEqual(action_spec.dtype, np.dtype("int64"))

        # Test step.
        timestep = env.reset()
        self.assertTrue(timestep.first())
        timestep = env.step(1)
        self.assertEqual(timestep.reward, 1.0)
        self.assertTrue(np.isscalar(timestep.reward))
        self.assertEqual(timestep.observation.shape, (4,))
        env.close()

    def test_early_truncation(self):
        # Pendulum has no early termination condition. Recent versions of gym force
        # to use v1. We try both in case an earlier version is installed.
        try:
            gym_env = gymnasium.make("Pendulum-v1")
        except:  # noqa
            gym_env = gymnasium.make("Pendulum-v0")
        env = gymnasium_wrapper.GymnasiumWrapper(gym_env)
        ts = env.reset()
        while not ts.last():
            ts = env.step(env.action_spec().generate_value())
        self.assertEqual(ts.discount, 1.0)
        self.assertTrue(np.isscalar(ts.reward))
        env.close()

    def test_multi_discrete(self):
        space = gymnasium.spaces.MultiDiscrete([2, 3])
        spec = gymnasium_wrapper._convert_to_spec(space)

        spec.validate([0, 0])
        spec.validate([1, 2])

        self.assertRaises(ValueError, spec.validate, [2, 2])
        self.assertRaises(ValueError, spec.validate, [1, 3])


if __name__ == "__main__":
    absltest.main()

corax/wrappers/gymnasium_wrapper_test.py
ethanluoyc/corax
# 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.

"""An environment wrapper to produce pixel observations from dm_control."""

import collections

import dm_env
import dm_env_wrappers
from dm_control.rl import control
from dm_control.suite.wrappers import pixels  # type: ignore


class MujocoPixelWrapper(dm_env_wrappers.EnvironmentWrapper):
    """Produces pixel observations from Mujoco environment observations."""

    def __init__(
        self,
        environment: control.Environment,
        *,
        height: int = 84,
        width: int = 84,
        camera_id: int = 0,
    ):
        render_kwargs = {"height": height, "width": width, "camera_id": camera_id}
        pixel_environment = pixels.Wrapper(
            environment, pixels_only=True, render_kwargs=render_kwargs
        )
        super().__init__(pixel_environment)

    def step(self, action) -> dm_env.TimeStep:
        return self._convert_timestep(self._environment.step(action))

    def reset(self) -> dm_env.TimeStep:
        return self._convert_timestep(self._environment.reset())

    def observation_spec(self):
        return self._environment.observation_spec()["pixels"]  # type: ignore

    def _convert_timestep(self, timestep: dm_env.TimeStep) -> dm_env.TimeStep:
        """Removes the pixel observation's OrderedDict wrapper."""
        observation: collections.OrderedDict = timestep.observation
        return timestep._replace(observation=observation["pixels"])

corax/wrappers/mujoco.py
ethanluoyc/corax
from ml_collections import config_dict


def get_config():
    config = config_dict.ConfigDict()
    config.dataset_name = "antmaze-medium-diverse-v2"
    config.seed = 0
    config.log_to_wandb = False

    config.policy_hidden_sizes = (256, 256)
    config.critic_hidden_sizes = (256, 256, 256, 256)

    config.policy_lr = 1e-4
    config.qf_lr = 3e-4
    config.batch_size = 256

    config.discount = 0.99

    config.reward_scale = 10.0
    config.reward_bias = -5

    config.initial_num_steps = 5000

    config.enable_calql = True
    config.cql_config = dict(
        cql_lagrange_threshold=0.8,
        cql_num_samples=10,
        max_target_backup=True,
        tau=5e-3,
    )

    # Offline training config
    config.offline_num_steps = int(1e6)
    config.offline_eval_every = int(5e4)
    config.offline_num_eval_episodes = 20

    # Offline training config
    config.mixing_ratio = 0.5
    config.online_utd_ratio = 1
    config.online_num_steps = int(1e6)
    config.online_eval_every = 2000
    config.online_num_eval_episodes = 20
    return config

projects/baselines/baselines/calql/configs/antmaze.py
ethanluoyc/corax
from typing import Iterator

import jax
import jax.numpy as jnp
import numpy as np
import optax
import reverb
import tensorflow as tf
import tree
from absl import app
from absl import flags
from ml_collections import config_flags

import corax
from baselines import d4rl_utils
from baselines import experiment_utils
from corax import specs
from corax import types
from corax.adders import reverb as adders_reverb
from corax.agents.jax import actor_core
from corax.agents.jax import actors
from corax.agents.jax import calql
from corax.agents.jax.calql.adder import SparseReward
from corax.datasets import reverb as reverb_datasets
from corax.datasets import tfds
from corax.jax import types as jax_types
from corax.jax import utils as jax_utils
from corax.jax import variable_utils
from corax.utils import counting
from corax.utils import loggers

_CONFIG = config_flags.DEFINE_config_file("config", None, lock_config=False)
flags.mark_flag_as_required("config")


@tf.function
def compute_return_to_go(rewards, discounts, gamma):
    rewards = tf.convert_to_tensor(rewards)
    discounts = tf.convert_to_tensor(discounts)

    def discounted_return_fn(acc, reward_discount):
        reward, discount = reward_discount
        return acc * discount * gamma + reward

    return tf.scan(
        fn=discounted_return_fn,
        elems=(rewards, discounts),
        reverse=True,
        initializer=tf.constant(0.0, dtype=rewards.dtype),
    )


@tf.function
def preprocess_episode(episode, negative_reward, reward_scale, reward_bias, gamma):
    steps = episode["steps"].batch(episode["steps"].cardinality()).get_single_element()

    observations = steps["observation"][:-1]
    next_observations = steps["observation"][1:]
    rewards = tf.cast(steps["reward"][:-1], tf.float64)
    actions = steps["action"][:-1]
    discounts = tf.cast(steps["discount"][:-1], tf.float64)

    rewards = rewards * reward_scale + reward_bias
    reward_negative = negative_reward * reward_scale + reward_bias
    gamma = tf.convert_to_tensor(gamma, dtype=rewards.dtype)

    if tf.reduce_all(rewards == reward_negative):
        return_to_go = tf.ones_like(rewards) * (reward_negative / (1 - gamma))  # type: ignore
    else:
        return_to_go = compute_return_to_go(rewards, discounts, gamma)

    return types.Transition(
        observation=observations,
        action=actions,
        discount=discounts,
        reward=tf.cast(rewards, tf.float32),
        next_observation=next_observations,
        extras={"mc_return": return_to_go},
    )


def get_transitions_dataset(
    d4rl_name, negative_reward, reward_scale, reward_bias, gamma
):
    tfds_name = d4rl_utils.get_tfds_name(d4rl_name)
    dataset = tfds.load_tfds_dataset(tfds_name)
    tf_transitions = []
    for episode in dataset:
        converted_transitions = preprocess_episode(
            episode, negative_reward, reward_scale, reward_bias, gamma
        )
        tf_transitions.append(converted_transitions)

    transitions = tf.data.Dataset.from_tensor_slices(
        tree.map_structure(lambda *x: tf.concat(x, axis=0), *tf_transitions)
    )
    return transitions


def main(_):
    config = _CONFIG.value

    logger_factory = experiment_utils.LoggerFactory(
        workdir=None, log_to_wandb=config.log_to_wandb, evaluator_time_delta=0.0
    )

    reward_config = SparseReward(
        reward_scale=config.reward_scale,
        reward_bias=config.reward_bias,
        positive_reward=1,
        negative_reward=0,
    )

    transitions = get_transitions_dataset(
        config.dataset_name,
        reward_config.negative_reward,
        reward_config.reward_scale,
        reward_config.reward_bias,
        config.discount,
    )

    env = d4rl_utils.load_d4rl_environment(config.dataset_name, config.seed)
    env_spec = corax.make_environment_spec(env)

    networks = calql.make_networks(
        env_spec,
        policy_hidden_sizes=config.policy_hidden_sizes,
        critic_hidden_sizes=config.critic_hidden_sizes,
    )

    eval_policy = actor_core.batched_feed_forward_to_actor_core(
        calql.apply_policy_and_sample(networks, eval_mode=True)
    )
    train_policy = actor_core.batched_feed_forward_to_actor_core(
        calql.apply_policy_and_sample(networks, eval_mode=False)
    )

    def make_offline_iterator(batch_size):
        return tfds.JaxInMemoryRandomSampleIterator(
            transitions, jax.random.PRNGKey(config.seed), batch_size
        )

    def make_offline_learner(
        networks: calql.CQLNetworks,
        random_key: jax_types.PRNGKey,
        dataset: Iterator[types.Transition],
        env_spec: specs.EnvironmentSpec,
        logger: loggers.Logger,
        counter: counting.Counter,
    ):
        target_entropy = -np.prod(env_spec.actions.shape).item()
        return calql.CalQLLearner(
            config.batch_size,
            networks,
            random_key,
            dataset,
            policy_optimizer=optax.adam(config.policy_lr),
            critic_optimizer=optax.adam(config.qf_lr),
            reward_scale=1.0,  # Reward scaled in dataset iterator
            target_entropy=target_entropy,
            discount=config.discount,
            num_bc_iters=0,
            use_calql=config.enable_calql,
            **config.cql_config,
            logger=logger,
            counter=counter,
        )

    def make_actor(policy, random_key, variable_source, adder=None):
        return actors.GenericActor(
            policy,
            random_key,
            variable_utils.VariableClient(
                variable_source, "policy", update_period=1, device="cpu"
            ),
            adder=adder,
            per_episode_update=True,
        )

    def make_online_iterator(
        replay_client: reverb.Client,
    ) -> Iterator[types.Transition]:
        # mix offline and online buffer
        assert config.mixing_ratio >= 0.0
        online_batch_size = int(config.mixing_ratio * config.batch_size)
        offline_batch_size = config.batch_size - online_batch_size
        offline_iterator = make_offline_iterator(offline_batch_size)

        online_dataset: Iterator[
            reverb.ReplaySample
        ] = reverb_datasets.make_reverb_dataset(
            server_address=replay_client.server_address,
            table=adders_reverb.DEFAULT_PRIORITY_TABLE,
            num_parallel_calls=4,
            batch_size=online_batch_size,
            prefetch_size=1,
        ).as_numpy_iterator()  # type: ignore

        while True:
            offline_batch = next(offline_iterator)
            offline_transitions: types.Transition = jax.device_put(offline_batch)
            online_transitions: types.Transition = jax.device_put(
                next(online_dataset).data
            )

            yield tree.map_structure(
                lambda x, y: jnp.concatenate([x, y]),
                offline_transitions,
                online_transitions,  # type: ignore
            )

    def make_evaluator(random_key: jax_types.PRNGKey, learner: corax.Learner):
        return d4rl_utils.D4RLEvaluator(
            lambda: env,
            make_actor(eval_policy, random_key, learner),
            logger=evaluator_logger,
            counter=evaluator_counter,
        )

    def make_replay_tables(env_spec: specs.EnvironmentSpec):
        return [
            reverb.Table(
                name=adders_reverb.DEFAULT_PRIORITY_TABLE,
                sampler=reverb.selectors.Uniform(),
                remover=reverb.selectors.Fifo(),
                # Do not limit the size of the table.
                max_size=config.online_num_steps,
                rate_limiter=reverb.rate_limiters.MinSize(1),
                signature=adders_reverb.NStepTransitionAdder.signature(
                    env_spec,
                    extras_spec={
                        "mc_return": tf.TensorSpec(shape=(), dtype=tf.float32)  # type: ignore
                    },
                ),
            )
        ]

    key = jax.random.PRNGKey(config.seed)

    parent_counter = counting.Counter(time_delta=0.0)
    learner_counter = counting.Counter(parent_counter, "learner", time_delta=0.0)
    learner_logger = logger_factory("learner", learner_counter.get_steps_key(), 0)

    learner_key, key = jax.random.split(key)

    offline_iterator = make_offline_iterator(config.batch_size)
    offline_iterator = jax_utils.prefetch(offline_iterator)

    offline_learner = make_offline_learner(
        networks,
        learner_key,
        offline_iterator,
        env_spec,
        learner_logger,
        learner_counter,
    )

    actor_key, key = jax.random.split(key)

    evaluator_counter = counting.Counter(parent_counter, "evaluator", time_delta=0.0)
    evaluator_logger = logger_factory("evaluator", evaluator_counter.get_steps_key(), 0)

    offline_evaluator = make_evaluator(actor_key, offline_learner)

    # Offline training
    max_num_offline_steps = config.offline_num_steps
    steps = 0
    eval_every = config.offline_eval_every
    while steps < max_num_offline_steps:
        learner_steps = min(eval_every, max_num_offline_steps - steps)
        for _ in range(learner_steps):
            offline_learner.step()
        offline_evaluator.run(config.offline_num_eval_episodes)
        steps += learner_steps

    reverb_tables = make_replay_tables(env_spec)
    replay_server = reverb.Server(reverb_tables, port=None)
    replay_client = replay_server.localhost_client()

    train_env = d4rl_utils.load_d4rl_environment(config.dataset_name, config.seed)
    online_counter = counting.Counter(parent_counter, "actor", time_delta=0.0)
    online_logger = logger_factory("actor", online_counter.get_steps_key(), 0)

    # Online fine-tuning
    online_dataset = make_online_iterator(replay_client)

    learner_key, key = jax.random.split(key)
    online_learner = make_offline_learner(
        networks,
        learner_key,
        online_dataset,
        env_spec,
        learner_logger,
        learner_counter,
    )
    online_learner.restore(offline_learner.save())

    del offline_learner, offline_iterator, offline_evaluator

    adder = calql.CalQLAdder(
        adders_reverb.NStepTransitionAdder(
            replay_client, n_step=1, discount=config.discount
        ),
        config.discount,
        reward_config=reward_config,
    )

    actor_key, eval_key = jax.random.split(key)
    online_actor = make_actor(train_policy, actor_key, online_learner, adder=adder)
    online_evaluator = make_evaluator(eval_key, online_learner)

    initial_num_steps = config.initial_num_steps
    eval_every = config.online_eval_every
    eval_episodes = config.online_num_eval_episodes
    num_steps = 0
    episode_length = 0
    episode_return = 0
    timestep = train_env.reset()
    online_actor.observe_first(timestep)
    online_actor.update()
    online_evaluator.run(num_episodes=eval_episodes)
    while True:
        action = online_actor.select_action(timestep.observation)
        next_timestep = train_env.step(action)
        num_steps += 1
        episode_return += next_timestep.reward
        episode_length += 1
        if num_steps >= int(config.online_num_steps):
            break
        if num_steps >= initial_num_steps:
            for _ in range(config.online_utd_ratio):
                online_learner.step()

        if num_steps % eval_every == 0:
            online_evaluator.run(num_episodes=eval_episodes)

        online_actor.observe(action, next_timestep)
        online_actor.update()

        if next_timestep.last():
            counts = online_counter.increment(episodes=1, steps=episode_length)
            online_logger.write({**counts, "episode_return": episode_return})
            episode_return = 0
            episode_length = 0
            timestep = train_env.reset()
            online_actor.observe_first(timestep)
        else:
            timestep = next_timestep


if __name__ == "__main__":
    app.run(main)

projects/baselines/baselines/calql/main.py
ethanluoyc/corax
import re
from typing import Callable, Tuple

import dm_env
import dm_env_wrappers
import numpy as np
from absl import logging
from d4rl import infos as d4rl_info

import corax
from corax.utils import counting
from corax.utils import loggers
from corax.utils import observers
from corax.wrappers import gym_wrapper


def load_d4rl_environment(name: str, seed: int):
    import d4rl  # noqa: F401
    import gym

    environment = gym.make(name)
    environment.seed(seed)
    environment = gym_wrapper.GymWrapper(environment)
    environment = dm_env_wrappers.SinglePrecisionWrapper(environment)

    return environment


def parse_d4rl_dataset_name(dataset_name: str) -> Tuple[str, str, str]:
    match = re.match(
        r"(?P<env>[a-z0-9]+)-(?P<dataset>[a-z\-]+)-(?P<version>v\d)", dataset_name
    )
    if not match:
        raise ValueError(f"Invalid D4RL dataset name: {dataset_name}")

    return match.group("env"), match.group("dataset"), match.group("version")


def get_tfds_name(d4rl_name: str) -> str:
    """Return the corresponding TFDS name for a given D4RL dataset name."""

    env, dataset, version = parse_d4rl_dataset_name(d4rl_name)
    if env in ["halfcheetah", "hopper", "walker2d", "ant"]:
        if version == "v0" and dataset == "medium-replay":
            return f"d4rl_mujoco_{env}/{version}-mixed"
        else:
            return f"d4rl_mujoco_{env}/{version}-{dataset}"
    elif env in ["antmaze"]:
        return f"d4rl_antmaze/{dataset}-{version}"
    elif env in ["pen", "door", "hammer", "relocate"]:
        return f"d4rl_adroit_{env}/{version}-{dataset}"
    else:
        raise ValueError(f"Unknown D4RL environment: {env}")


class D4RLScoreObserver(observers.EnvLoopObserver):
    def __init__(self, dataset_name: str) -> None:
        assert dataset_name in d4rl_info.REF_MAX_SCORE
        self._episode_return = 0
        self._dataset_name = dataset_name

    def observe_first(self, env, timestep) -> None:
        """Observes the initial state."""
        del env, timestep
        self._episode_return = 0

    def observe(self, env, timestep, action: np.ndarray) -> None:
        """Records one environment step."""
        del env, action
        self._episode_return += timestep.reward

    def get_metrics(self):
        """Returns metrics collected for the current episode."""
        score = self._episode_return
        ref_min_score = d4rl_info.REF_MIN_SCORE[self._dataset_name]
        ref_max_score = d4rl_info.REF_MAX_SCORE[self._dataset_name]
        normalized_score = (score - ref_min_score) / (ref_max_score - ref_min_score)
        return {"d4rl_normalized_score": normalized_score}


class D4RLEvaluator:
    def __init__(
        self,
        environment_factory: Callable[[], dm_env.Environment],
        actor: corax.Actor,
        counter: counting.Counter,
        logger: loggers.Logger,
    ):
        self._environment_factory = environment_factory
        self._actor = actor
        self._counter = counter
        self._logger = logger

    def run(self, num_episodes: int):
        environment = self._environment_factory()
        self._actor.update(wait=True)
        avg_reward = 0.0
        steps = 0
        for _ in range(num_episodes):
            timestep = environment.reset()
            self._actor.observe_first(timestep)
            while not timestep.last():
                action = self._actor.select_action(timestep.observation)
                timestep = environment.step(action)
                self._actor.observe(action, timestep)
                steps += 1
                avg_reward += timestep.reward

        counts = self._counter.increment(steps=steps, episodes=num_episodes)

        avg_reward /= num_episodes
        d4rl_score = environment.get_normalized_score(avg_reward)  # type: ignore
        results = {
            "average_episode_return": avg_reward,
            "average_normalized_score": d4rl_score,
        }
        results.update(counts)

        logging.info("---------------------------------------")
        logging.info(
            "Evaluation over %d episodes: %.3f (unormalized %.3f)",
            num_episodes,
            d4rl_score,
            avg_reward,
        )
        logging.info("---------------------------------------")
        self._logger.write(results)
        return results

projects/baselines/baselines/d4rl_utils.py
ethanluoyc/corax
import ml_collections


def get_config():
    config = ml_collections.ConfigDict()
    config.dataset = "medium"
    config.env = "hopper"
    config.mode = "normal"
    config.K = 20
    config.pct_traj = 1.0
    config.batch_size = 64
    config.network_config = dict(
        hidden_size=128,
        num_layers=3,
        num_heads=1,
        dropout_rate=0.1,
    )
    config.seed = 0
    config.learning_rate = 1e-4
    config.weight_decay = 1e-4
    config.warmup_steps = 10000
    config.num_eval_episodes = 10
    config.num_steps = int(1e5)
    config.eval_every = int(1e4)
    config.log_to_wandb = False
    config.max_ep_len = 1000
    config.scale = 1000.0
    return config

projects/baselines/baselines/decision_transformer/config.py
ethanluoyc/corax
import jax
import tensorflow as tf
from absl import app
from absl import flags
from ml_collections import config_flags

import corax
from baselines import d4rl_utils
from baselines import experiment_utils
from corax import environment_loop
from corax.agents.jax import decision_transformer
from corax.agents.jax.decision_transformer import dataset as dataset_lib
from corax.datasets import tfds
from corax.jax import utils as jax_utils

_CONFIG = config_flags.DEFINE_config_file("config", None)
flags.mark_flag_as_required("config")


def main(_):
    config = _CONFIG.value

    env_name = config["env"]
    max_ep_len = config.max_ep_len
    scale = config.scale  # normalization for rewards/returns
    seed = config.seed

    if env_name == "hopper":
        eval_env_name = "Hopper-v3"
        target_return = 3600  # evaluation conditioning targets
    elif env_name == "halfcheetah":
        eval_env_name = "HalfCheetah-v3"
        target_return = 12000
    elif env_name == "walker2d":
        eval_env_name = "Walker2d-v3"
        target_return = 5000
    else:
        raise NotImplementedError

    dataset_name = f"{env_name}-{config['dataset']}-v2"
    env = d4rl_utils.load_d4rl_environment(eval_env_name, seed)

    K = config.K
    batch_size = config.batch_size
    max_num_learner_steps = config.num_steps
    episode_dataset = tfds.load_tfds_dataset(d4rl_utils.get_tfds_name(dataset_name))
    observation_mean_std = dataset_lib.get_observation_mean_std(episode_dataset)

    def make_dataset_iterator(key):
        del key
        dataset = (
            dataset_lib.transform_decision_transformer_input(
                episode_dataset,
                sequence_length=K,
                scale=scale,
                observation_mean_std=observation_mean_std,
            )
            .shuffle(int(1e6))
            .repeat()
            .batch(batch_size)
            .as_numpy_iterator()
        )
        return jax_utils.device_put(dataset, jax.local_devices()[0])

    network_factory = lambda spec: decision_transformer.make_gym_networks(
        spec=spec, episode_length=max_ep_len, **config.network_config
    )

    dt_config = decision_transformer.DecisionTransformerConfig(
        context_length=K,
        target_return=target_return,
        return_scale=scale,
        mode=config.mode,
        learning_rate=config.learning_rate,
        warmup_steps=config.warmup_steps,
        grad_norm_clipping=0.25,
        weight_decay=config.weight_decay,
    )

    builder = decision_transformer.DecisionTransformerBuilder(
        dt_config,
        observation_mean_std,
        max_num_learner_steps=max_num_learner_steps,
        actor_device="gpu",
    )

    logger_factory = experiment_utils.LoggerFactory(log_to_wandb=config.log_to_wandb)
    random_key = jax.random.PRNGKey(config.seed)
    environment = env

    spec = corax.make_environment_spec(environment)
    dataset = make_dataset_iterator(random_key)

    networks = network_factory(spec)
    learner = builder.make_learner(
        random_key,
        networks,
        dataset,
        logger_factory,  # type: ignore
        spec,
    )
    learner.step()
    policy = builder.make_policy(networks, spec, True)
    actor = builder.make_actor(random_key, policy, spec, learner)

    eval_loop = environment_loop.EnvironmentLoop(environment, actor)  # type: ignore

    steps = 0
    while steps < config.num_steps:
        learner_steps = min(config.eval_every, max_num_learner_steps - steps)
        for _ in range(learner_steps):
            learner.step()
        if eval_loop:
            eval_loop.run(num_episodes=config.num_eval_episodes)
            steps += learner_steps


if __name__ == "__main__":
    tf.config.set_visible_devices([], "GPU")
    app.run(main)

projects/baselines/baselines/decision_transformer/main.py
ethanluoyc/corax
# 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.
"""Frame stacking utilities.
Note: Different from what's used in Acme
The aggregation is done via concatenate instead reshape.
This ensures that channels from the same image is contiguouse in the flattened representation
This ensures that channels from the same image is consecutive
Fill the stack with images from the first frame if empty.
This is consistent with the frame stacking used in DrQV2
"""

import collections

import dm_env
import dm_env_wrappers
import numpy as np
import tree
from dm_env import specs as dm_env_specs

from corax import types


class FrameStackingWrapper(dm_env_wrappers.EnvironmentWrapper):
    """Wrapper that stacks observations along a new final axis."""

    def __init__(
        self,
        environment: dm_env.Environment,
        num_frames: int = 4,
        flatten: bool = False,
    ):
        """Initializes a new FrameStackingWrapper.

        Args:
          environment: Environment.
          num_frames: Number frames to stack.
          flatten: Whether to flatten the channel and stack dimensions together.
        """
        super().__init__(environment)
        self._environment = environment
        original_spec = self._environment.observation_spec()
        self._stackers = tree.map_structure(
            lambda _: FrameStacker(num_frames=num_frames, flatten=flatten),
            self._environment.observation_spec(),
        )
        self._observation_spec = tree.map_structure(
            lambda stacker, spec: stacker.update_spec(spec),
            self._stackers,
            original_spec,
        )

    def _process_timestep(self, timestep: dm_env.TimeStep) -> dm_env.TimeStep:
        observation = tree.map_structure(
            lambda stacker, x: stacker.step(x), self._stackers, timestep.observation
        )
        return timestep._replace(observation=observation)

    def reset(self) -> dm_env.TimeStep:
        for stacker in tree.flatten(self._stackers):
            stacker.reset()
        return self._process_timestep(self._environment.reset())

    def step(self, action: int) -> dm_env.TimeStep:
        return self._process_timestep(self._environment.step(action))

    def observation_spec(self) -> types.NestedSpec:
        return self._observation_spec  # type: ignore


class FrameStacker:
    """Simple class for frame-stacking observations."""

    def __init__(self, num_frames: int, flatten: bool = False):
        self._num_frames = num_frames
        self._flatten = flatten
        self.reset()

    @property
    def num_frames(self) -> int:
        return self._num_frames

    def reset(self):
        self._stack = collections.deque(maxlen=self._num_frames)

    def step(self, frame: np.ndarray) -> np.ndarray:
        """Append frame to stack and return the stack."""
        if not self._stack:
            # Fill with the same frame if empty
            self._stack.extend([frame] * (self._num_frames - 1))
        self._stack.append(frame)

        if not self._flatten:
            return np.stack(self._stack, axis=-1)
        else:
            # Ensure that the a single frame is consecutive
            return np.concatenate(self._stack, axis=-1)

    def update_spec(self, spec: dm_env_specs.Array) -> dm_env_specs.Array:
        if not self._flatten:
            new_shape = (*spec.shape, self._num_frames)
        else:
            new_shape = spec.shape[:-1] + (self._num_frames * spec.shape[-1],)
        return dm_env_specs.Array(shape=new_shape, dtype=spec.dtype, name=spec.name)

projects/baselines/baselines/drq_bc/drq_frame_stacking.py
ethanluoyc/corax
import os

os.environ["MUJOCO_GL"] = "egl"

# ruff: noqa: E402
import dm_env_wrappers
import jax
import numpy as np
import optax
import reverb
import tensorflow as tf
import tensorflow_datasets as tfds
from absl import app
from dm_control import suite

import corax
from baselines.drq_bc import drq_frame_stacking
from baselines.drq_bc import vd4rl_preprocessor
from corax import types
from corax.agents.jax import actor_core as actor_core_lib
from corax.agents.jax import actors
from corax.agents.jax import drq_v2
from corax.agents.jax.drq_v2 import augmentations
from corax.jax import utils
from corax.jax import variable_utils
from corax.utils import counting
from corax.utils import loggers
from corax.wrappers import mujoco


def make_environment(domain: str, task: str):
    env = suite.load(domain, task, task_kwargs={"random": 0})
    env = mujoco.MujocoPixelWrapper(env, height=84, width=84, camera_id=0)
    env = dm_env_wrappers.ActionRepeatWrapper(env, 2)
    env = drq_frame_stacking.FrameStackingWrapper(env, 3, flatten=True)
    env = dm_env_wrappers.SinglePrecisionWrapper(env)
    return env


def get_dataset_iterator(domain, task):
    name = "medium_expert"
    dataset = tfds.load(f"vd4rl/main_{domain}_{task}_{name}_84px")["train"]  # type: ignore
    dataset = dataset.map(
        lambda episode: {"steps": vd4rl_preprocessor.process_data(episode["steps"], 3)},
        num_parallel_calls=tf.data.AUTOTUNE,
    )
    # TODO: use shuffling
    dataset = dataset.flat_map(
        lambda episode: vd4rl_preprocessor.tfds_get_n_step_transitions(
            episode["steps"], 3, 0.99
        )
    )

    iterator = dataset.cache().batch(256).repeat().prefetch(4).as_numpy_iterator()
    while True:
        batch = next(iterator)
        transitions = types.Transition(**batch)
        yield reverb.ReplaySample(data=transitions, info=None)  # type: ignore


def main(_):
    tf.config.set_visible_devices([], "GPU")
    np.random.seed(0)
    domain = "walker"
    task = "walk"
    dataset = utils.device_put(get_dataset_iterator(domain, task), jax.devices()[0])
    dataset = utils.prefetch(dataset, 2)
    environment = make_environment(domain, task)
    env_spec = corax.make_environment_spec(environment)
    key = jax.random.PRNGKey(0)
    learner_key, actor_key = jax.random.split(key)
    networks = drq_v2.make_networks(env_spec)
    parent_counter = counting.Counter(time_delta=0.0)

    learner = drq_v2.DrQV2Learner(
        learner_key,
        dataset,
        networks,
        optax.linear_schedule(0.5, 0.1, 250000),
        augmentation=augmentations.batched_random_shift_aug,
        policy_optimizer=optax.adam(1e-4),
        critic_optimizer=optax.adam(1e-4),
        encoder_optimizer=optax.adam(1e-4),
        critic_soft_update_rate=0.01,
        discount=0.99,
        bc_alpha=2.5,
        logger=loggers.make_default_logger("learner", asynchronous=True),
        counter=counting.Counter(parent_counter, "learner", time_delta=0.0),
    )

    device = "gpu"
    variable_client = variable_utils.VariableClient(learner, "policy", device=device)

    evaluator = actors.GenericActor(
        actor_core_lib.batched_feed_forward_to_actor_core(
            drq_v2.apply_policy_and_sample(networks, env_spec.actions, 0.0)
        ),
        actor_key,
        variable_client,
        backend=device,
    )

    eval_loop = corax.EnvironmentLoop(
        environment=environment,
        actor=evaluator,
        logger=loggers.make_default_logger("evaluator"),
        counter=counting.Counter(parent_counter, "evaluator", time_delta=0.0),
    )

    # Run the environment loop.
    max_steps = int(1e6)
    eval_every = 5000
    eval_episodes = 10
    steps = 0
    while steps < max_steps:
        for _ in range(eval_every):
            learner.step()
        steps += eval_every
        eval_loop.run(eval_episodes)


if __name__ == "__main__":
    app.run(main)

projects/baselines/baselines/drq_bc/main.py
ethanluoyc/corax
import numpy as np
import rlds
import tensorflow as tf
import tree


def stack_observations_np(observations, stack_size):
    def _repeat_first(sequence):
        reps = [1] * (sequence[0].ndim + 1)
        reps[0] = stack_size - 1
        repeated_first_step = np.tile(sequence[0], reps)
        return np.concatenate(
            [repeated_first_step, sequence], 0
        )  # [:-(stack_size - 1)]

    def _stack_observation(observation):
        stack = [np.roll(observation, i, axis=0) for i in range(stack_size)]
        stack.reverse()  # chronological order
        return np.concatenate(stack, axis=-1)

    observations = _repeat_first(observations)
    return _stack_observation(observations)[(stack_size - 1) :]


def get_n_step_transitions(episode_steps, nsteps, discount):
    """Get n_step transitions from an episode of steps"""
    episode_length = episode_steps["reward"].shape[0] - 1
    transitions = []

    for t in range(0, episode_length - nsteps + 1):
        observations = episode_steps["observation"][t : t + nsteps + 1]
        actions = episode_steps["action"][t : t + nsteps + 1]
        rewards = episode_steps["reward"][t : t + nsteps + 1]
        discounts = episode_steps["discount"][t : t + nsteps + 1]

        total_reward = rewards[0]
        total_discount = discounts[0]

        for n in range(1, len(discounts) - 1):
            total_discount *= discount
            total_reward += rewards[n] * total_discount
            total_discount *= discounts[n]

        transitions.append(
            {
                "observation": observations[0],
                "next_observation": observations[-1],
                "reward": total_reward,
                "discount": total_discount,
                "action": actions[0],
            }
        )
    # This can be faster
    return tree.map_structure(lambda *xs: np.stack(xs), *transitions)


def stack_observation_tfds(observations: tf.data.Dataset, stack_size: int):
    first_obs = observations.take(1)
    observations = first_obs.repeat(stack_size - 1).concatenate(observations)
    data = observations.window(stack_size, shift=1, drop_remainder=True).flat_map(
        lambda x: x.batch(stack_size, drop_remainder=True)
    )
    return data.map(
        lambda x: tf.concat(tf.unstack(x, axis=0), axis=-1),
        num_parallel_calls=tf.data.AUTOTUNE,
    )


def process_data(steps_dataset, stack_size):
    obs = steps_dataset.map(
        lambda step: step["observation"], num_parallel_calls=tf.data.AUTOTUNE
    )
    stacked_obs = stack_observation_tfds(obs, stack_size)
    zipped_ds = tf.data.Dataset.zip((stacked_obs, steps_dataset))

    def combine(stacked_obs, steps):
        steps["observation"] = stacked_obs
        return steps

    return zipped_ds.map(combine, num_parallel_calls=tf.data.AUTOTUNE)


@tf.function
def _sequence_to_transitions(sequence, discount):
    total_reward = sequence["reward"][0]
    total_discount = sequence["discount"][0]
    n_steps = len(sequence["reward"]) - 1

    for t in range(1, n_steps):
        total_discount *= discount
        total_reward += total_discount * sequence["reward"][t]
        total_discount *= sequence["discount"][t]

    return {
        "observation": sequence["observation"][0],
        "next_observation": sequence["observation"][-1],
        "action": sequence["action"][0],
        "reward": total_reward,
        "discount": total_discount,
    }


def tfds_get_n_step_transitions(steps_dataset, nsteps, discount):
    """Get n_step transitions from an episode of steps"""
    batched_steps = rlds.transformations.batch(
        steps_dataset, size=nsteps + 1, shift=1, drop_remainder=True
    )
    return batched_steps.map(
        lambda seq: _sequence_to_transitions(seq, discount),
        num_parallel_calls=tf.data.experimental.AUTOTUNE,
    )

projects/baselines/baselines/drq_bc/vd4rl_preprocessor.py
ethanluoyc/corax
import numpy as np
import tensorflow as tf
from absl.testing import absltest
from absl.testing import parameterized

from baselines.drq_bc import vd4rl_preprocessor

_TRANSITION_TEST_CASES = [
    {
        "testcase_name": "test_normal",
        "steps": {
            "observation": np.arange(4, dtype=np.float32),
            "reward": np.array([1, 2, 3, 4], dtype=np.float32),
            "discount": np.ones(4, dtype=np.float32),
            "action": np.arange(4, dtype=np.float32),
        },
        "expected_transitions": {
            "observation": np.array([0, 1]),
            "next_observation": np.array([2, 3]),
            "discount": np.array([0.99, 0.99]),
            "action": np.array([0, 1]),
            "reward": np.array([1 + 2 * 0.99, 2 + 3 * 0.99]),
        },
        "n_steps": 2,
        "discount": 0.99,
    },
    {
        "testcase_name": "test_normal_2",
        "steps": {
            "observation": np.array([0, 1, 2, 3, 4, 5], dtype=np.float32),
            "reward": np.array([1, 2, 3, 4, 5, 6], dtype=np.float32),
            "discount": np.array([1, 1, 1, 1, 1, 1], dtype=np.float32),
            "action": np.array([0, 1, 2, 3, 4, 5], dtype=np.float32),
        },
        "expected_transitions": {
            "observation": np.array([0, 1, 2]),
            "next_observation": np.array([3, 4, 5]),
            "discount": np.array([0.99 * 0.99, 0.99 * 0.99, 0.99 * 0.99]),
            "action": np.array([0, 1, 2]),
            "reward": np.array(
                [
                    1 + 2 * 0.99 + 3 * 0.99 * 0.99,
                    2 + 3 * 0.99 + 4 * 0.99 * 0.99,
                    3 + 4 * 0.99 + 5 * 0.99 * 0.99,
                ]
            ),
        },
        "n_steps": 3,
        "discount": 0.99,
    },
    {
        "testcase_name": "test_terminal",
        "steps": {
            "observation": np.array([0, 1, 2, 3, 4, 5], dtype=np.float32),
            "reward": np.array([1, 2, 3, 4, 5, 6], dtype=np.float32),
            "discount": np.array([1, 1, 1, 1, 0, 0], dtype=np.float32),
            "action": np.array([0, 1, 2, 3, 4, 5], dtype=np.float32),
        },
        "expected_transitions": {
            "observation": np.array([0, 1, 2]),
            "next_observation": np.array([3, 4, 5]),
            "discount": np.array([0.99 * 0.99, 0.99 * 0.99, 0]),
            "action": np.array([0, 1, 2]),
            "reward": np.array(
                [
                    1 + 2 * 0.99 + 3 * 0.99 * 0.99,
                    2 + 3 * 0.99 + 4 * 0.99 * 0.99,
                    # 3 + 4 * 0.99 + 5 * 0.99 * 0.99,
                    3 + 4 * 0.99 + 5 * 0.99 * 0.99,
                ]
            ),
        },
        "n_steps": 3,
        "discount": 0.99,
    },
    {
        "testcase_name": "test_terminal_2",
        "steps": {
            "observation": np.array([1, 2, 3, 4, 5], dtype=np.float32),
            "reward": np.array([2, 3, 5, 7, 0], dtype=np.float32),
            "discount": np.array([1, 1, 1, 0, 0], dtype=np.float32),
            "action": np.array([0, 0, 0, 0, 0], dtype=np.float32),
        },
        "expected_transitions": {
            "observation": np.array([1, 2]),
            "next_observation": np.array([4, 5]),
            "discount": np.array([0.5 * 0.5, 0]),
            "action": np.array([0, 0]),
            "reward": np.array(
                [
                    2 + 0.5 * 3 + 0.25 * 5,
                    3 + 0.5 * 5 + 0.25 * 7,
                ]
            ),
        },
        "n_steps": 3,
        "discount": 0.5,
    },
    {
        "testcase_name": "test_one_step",
        "steps": {
            "observation": np.array([1, 2, 3], dtype=np.float32),
            "reward": np.array([2, 3, -1], dtype=np.float32),
            "discount": np.array([1, 0, 0], dtype=np.float32),
            "action": np.array([0, 1, 0], dtype=np.float32),
        },
        "expected_transitions": {
            "observation": np.array([1, 2]),
            "next_observation": np.array([2, 3]),
            "discount": np.array([1, 0]),
            "action": np.array([0, 1]),
            "reward": np.array([2, 3]),
        },
        "n_steps": 1,
        "discount": 0.5,
    },
]


def _assert_transitions_equal(np_transitions, expected_transitions):
    np.testing.assert_allclose(
        np_transitions["observation"], expected_transitions["observation"]
    )
    np.testing.assert_allclose(
        np_transitions["next_observation"], expected_transitions["next_observation"]
    )
    np.testing.assert_allclose(np_transitions["reward"], expected_transitions["reward"])
    np.testing.assert_allclose(
        np_transitions["discount"], expected_transitions["discount"]
    )
    np.testing.assert_allclose(np_transitions["action"], expected_transitions["action"])


class VD4RLDatasetTest(parameterized.TestCase):
    def test_observation_stacking_np(self):
        test_obs = np.arange(4).reshape((4, 1))
        stacked_obs = vd4rl_preprocessor.stack_observations_np(test_obs, 2)
        expected_stacked_obs = np.array(
            [
                [0, 0],
                [0, 1],
                [1, 2],
                [2, 3],
            ]
        )
        np.testing.assert_equal(stacked_obs, expected_stacked_obs)

    def test_observation_stacking_tfds(self):
        test_obs = np.arange(4).reshape((4, 1))
        stacked_obs = next(
            vd4rl_preprocessor.stack_observation_tfds(
                tf.data.Dataset.from_tensor_slices(test_obs), 2
            )
            .batch(100)
            .as_numpy_iterator()
        )
        expected_stacked_obs = np.array(
            [
                [0, 0],
                [0, 1],
                [1, 2],
                [2, 3],
            ]
        )
        np.testing.assert_equal(stacked_obs, expected_stacked_obs)

    @parameterized.named_parameters(*_TRANSITION_TEST_CASES)
    def test_create_transitions_np(
        self, n_steps, steps, discount, expected_transitions
    ):
        np_transitions = vd4rl_preprocessor.get_n_step_transitions(
            steps, n_steps, discount
        )
        _assert_transitions_equal(np_transitions, expected_transitions)

    @parameterized.named_parameters(*_TRANSITION_TEST_CASES)
    def test_create_transitions_tfds(
        self, n_steps, steps, discount, expected_transitions
    ):
        steps_dataset = tf.data.Dataset.from_tensor_slices(steps)
        transition_ds = vd4rl_preprocessor.tfds_get_n_step_transitions(
            steps_dataset, n_steps, discount
        )
        tf_transitions = next(transition_ds.batch(100).as_numpy_iterator())
        _assert_transitions_equal(tf_transitions, expected_transitions)


if __name__ == "__main__":
    tf.config.set_visible_devices([], "GPU")
    absltest.main()

projects/baselines/baselines/drq_bc/vd4rl_preprocessor_test.py
ethanluoyc/corax
"""Utility definitions for experiments."""
import os
from typing import Any, Dict, List, Mapping, Optional

from absl import logging

from corax.jax import utils as jax_utils
from corax.utils import loggers as loggers_lib
from corax.utils.loggers import wandb as wandb_logger


def _get_time_delta(time_delta: Optional[float], default_time_delta: float):
    if time_delta is not None:
        return time_delta
    else:
        return default_time_delta


class LoggerFactory:
    """Factory for creating loggers used in magi RL experiments."""

    def __init__(
        self,
        workdir: Optional[str] = None,
        log_to_terminal: Optional[bool] = None,
        log_to_wandb: bool = False,
        wandb_kwargs: Optional[Mapping[str, Any]] = None,
        time_delta: float = 1.0,
        actor_time_delta: Optional[float] = None,
        learner_time_delta: Optional[float] = None,
        evaluator_time_delta: Optional[float] = None,
        async_learner_logger: bool = False,
        add_uid: bool = True,
    ):
        """Create a logger factory."""

        wandb_kwargs = wandb_kwargs or {}
        self._log_to_wandb = log_to_wandb

        if log_to_terminal is None:
            log_to_terminal = os.environ.get("JOB_ID", None) is None

        self._log_to_terminal = log_to_terminal
        self._run = None
        if log_to_wandb and self._run is None:
            import wandb

            wandb.require("service")
            self._run = wandb.init(**wandb_kwargs)
        if workdir is not None:
            os.makedirs(workdir, exist_ok=True)
        self._workdir = workdir
        self._time_delta = time_delta
        self._actor_time_delta = actor_time_delta
        self._learner_time_delta = learner_time_delta
        self._evaluator_time_delta = evaluator_time_delta
        self._async_learner_logger = async_learner_logger
        self._add_uid = add_uid

    @property
    def run(self):
        return self._run

    def __call__(
        self,
        label: str,
        steps_key: Optional[str] = None,
        instance: Optional[int] = None,
    ):
        """Create an experiment logger."""
        if steps_key is None:
            steps_key = f"{label}_steps"

        if instance is None:
            instance = 0

        # Binding the experiment logger factory ensures that
        # the wandb run associated with the launching process gets
        # serialized and passed to workers correctly.
        if label in ("learner",):
            return self.make_default_logger(
                label=label,
                log_to_terminal=self._log_to_terminal,
                asynchronous=self._async_learner_logger,
                time_delta=_get_time_delta(self._learner_time_delta, self._time_delta),
                serialize_fn=jax_utils.fetch_devicearray,
                workdir=self._workdir,
                log_to_wandb=self._log_to_wandb,
                steps_key=steps_key,
                wandb_run=self._run,
                add_uid=self._add_uid,
            )
        elif label in ("evaluator",):
            return self.make_default_logger(
                label=label,
                log_to_terminal=self._log_to_terminal,
                time_delta=_get_time_delta(
                    self._evaluator_time_delta, self._time_delta
                ),
                steps_key=steps_key,
                workdir=self._workdir,
                log_to_wandb=self._log_to_wandb,
                wandb_run=self._run,
                add_uid=self._add_uid,
            )
        elif label in ("actor",):
            return self.make_default_logger(
                label=label,
                log_to_terminal=self._log_to_terminal,
                save_data=instance == 0,
                time_delta=_get_time_delta(self._actor_time_delta, self._time_delta),
                steps_key=steps_key,
                workdir=self._workdir,
                log_to_wandb=self._log_to_wandb,
                wandb_run=self._run,
                add_uid=self._add_uid,
            )
        else:
            logging.warning("Unknown label %s. Fallback to default.", label)
            return self.make_default_logger(
                label=label,
                log_to_terminal=self._log_to_terminal,
                steps_key=steps_key,
                time_delta=self._time_delta,
                workdir=self._workdir,
                log_to_wandb=self._log_to_wandb,
                wandb_run=self._run,
                add_uid=self._add_uid,
            )

    @staticmethod
    def make_default_logger(
        label: str,
        log_to_terminal: bool = True,
        save_data: bool = True,
        time_delta: float = 1.0,
        asynchronous: bool = False,
        print_fn=None,
        workdir: Optional[str] = None,
        serialize_fn=loggers_lib.to_numpy,
        steps_key: str = "steps",
        log_to_wandb: bool = False,
        wandb_kwargs: Optional[Dict[str, Any]] = None,
        add_uid: bool = False,
        wandb_run=None,
    ) -> loggers_lib.Logger:
        """Make a default Acme logger.
        Args:
          label: Name to give to the logger.
          save_data: Whether to persist data.
          time_delta: Time (in seconds) between logging events.
          asynchronous: Whether the write function should block or not.
          print_fn: How to print to terminal (defaults to print).
          serialize_fn: An optional function to apply to the write inputs before
            passing them to the various loggers.
          steps_key: Ignored.

        Returns:
          A logger object that responds to logger.write(some_dict).
        """
        if not print_fn:
            print_fn = logging.info

        loggers: List[loggers_lib.Logger] = []
        if log_to_terminal:
            terminal_logger = loggers_lib.TerminalLogger(label=label, print_fn=print_fn)
            loggers.append(terminal_logger)
        if save_data and workdir is not None:
            loggers.append(loggers_lib.CSVLogger(workdir, label=label, add_uid=add_uid))
        if save_data and log_to_wandb:
            if wandb_kwargs is None:
                wandb_kwargs = {}
            loggers.append(
                wandb_logger.WandbLogger(
                    label=label, steps_key=steps_key, run=wandb_run, **wandb_kwargs
                )
            )

        # Dispatch to all writers and filter Nones and by time.
        logger = loggers_lib.Dispatcher(loggers, serialize_fn)
        logger = loggers_lib.NoneFilter(logger)
        if asynchronous:
            logger = loggers_lib.AsyncLogger(logger)
        logger = loggers_lib.TimeFilter(logger, time_delta)
        return loggers_lib.AutoCloseLogger(logger)

projects/baselines/baselines/experiment_utils.py
ethanluoyc/corax
import dataclasses
from typing import Optional, Sequence


@dataclasses.dataclass
class Config:
    env_name: str = "halfcheetah-medium-v2"
    num_episodes: Optional[int] = None
    batch_size: int = 256

    max_num_learner_steps: int = int(1e6)
    num_eval_episodes: int = 10
    eval_every: int = 5000
    seed: int = 0

    hidden_dims: Sequence[int] = (256, 256)

    use_cosine_decay: bool = True
    policy_lr: float = 3e-4
    critic_lr: float = 3e-4
    value_lr: float = 3e-4

    discount: float = 0.99
    tau: float = 5e-3
    expectile: float = 0.7
    temperature: float = 3.0

    log_to_wandb: bool = False
    num_sgd_steps_per_step: int = 1


def get_base_config():
    return Config()

projects/baselines/baselines/iql/base_config.py
ethanluoyc/corax
import itertools

from baselines.iql import base_config
from baselines.iql.configs import sweeps


def get_config():
    config = base_config.get_base_config()
    config.env_name = "antmaze-umaze-v2"

    config.eval_every = 50000
    config.num_eval_episodes = 100

    config.expectile = 0.9  # The actual tau for expectiles.
    config.temperature = 10.0

    config.tau = 0.005  # For soft target updates.

    return config


_NUM_SEEDS = 10


def sweep_datasets():
    parameters = []
    for seed, env_name in itertools.product(
        range(_NUM_SEEDS),
        sweeps.sweep_d4rl_antmaze(),
    ):
        parameters.append({"config.seed": seed, "config.env_name": env_name})
    return parameters

projects/baselines/baselines/iql/configs/antmaze.py
ethanluoyc/corax
import itertools

from baselines.iql import base_config
from baselines.iql.configs import sweeps

_NUM_SEEDS = 10


def get_config():
    return base_config.get_base_config()


def sweep_datasets():
    parameters = []
    for seed, env_name in itertools.product(
        range(_NUM_SEEDS),
        sweeps.sweep_d4rl_locomotion(),
    ):
        parameters.append({"config.seed": seed, "config.env_name": env_name})
    return parameters

projects/baselines/baselines/iql/configs/mujoco.py
ethanluoyc/corax
import itertools


def sweep_d4rl_locomotion(version="v2"):
    sweep = []
    envs = ["halfcheetah", "hopper", "walker2d"]
    datasets = ["medium", "medium-replay", "medium-expert"]
    for env, dataset in itertools.product(envs, datasets):
        sweep.append(f"{env}-{dataset}-{version}")
    return sweep


def sweep_d4rl_antmaze(version="v2"):
    envs = ["antmaze"]
    datasets = [
        "umaze",
        "umaze-diverse",
        "medium-play",
        "medium-diverse",
        "large-play",
        "large-diverse",
    ]
    sweep = []
    for env, dataset in itertools.product(envs, datasets):
        sweep.append(f"{env}-{dataset}-{version}")
    return sweep

projects/baselines/baselines/iql/configs/sweeps.py
ethanluoyc/corax
import jax
from absl import app
from absl import flags
from absl import logging
from ml_collections import config_flags

from baselines.iql import base_config
from baselines.iql import train

_CONFIG = config_flags.DEFINE_config_file("config", None)
flags.mark_flag_as_required("config")

FLAGS = flags.FLAGS


def main(_):
    config: base_config.Config = _CONFIG.value

    logging.info("Using config: %s", config_flags.get_config_filename(FLAGS["config"]))
    logging.info("Overrides: %s", config_flags.get_override_values(FLAGS["config"]))

    train.train_and_evaluate(config, workdir=None)


if __name__ == "__main__":
    jax.config.config_with_absl()
    app.run(main)

projects/baselines/baselines/iql/main.py
ethanluoyc/corax
import time
from typing import Optional, Sequence

import jax
import numpy as np
import optax
import rlds
import tensorflow as tf
from absl import logging

import corax
from baselines import d4rl_utils
from baselines import experiment_utils
from baselines import rlds_utils
from baselines.iql import base_config
from corax.agents.jax import actor_core
from corax.agents.jax import actors as actors_lib
from corax.agents.jax import iql
from corax.agents.jax import learners as learners_lib
from corax.datasets import tfds
from corax.jax import utils as jax_utils
from corax.jax import variable_utils
from corax.utils import counting
from corax.utils import loggers
from corax.utils import observers as observers_lib


def make_d4rl_transition_dataset(
    d4rl_name: str, batch_size: int, seed: int, num_episodes: Optional[int] = None
):
    dataset = tfds.load_tfds_dataset(
        d4rl_utils.get_tfds_name(d4rl_name), num_episodes=num_episodes, download=False
    )
    start_time = time.time()
    if "antmaze" in d4rl_name:
        reward_scale = 1.0
        reward_bias = -1.0
    else:
        episode_returns = dataset.map(
            lambda episode: rlds_utils.compute_episode_return(episode[rlds.STEPS]),
            num_parallel_calls=tf.data.AUTOTUNE,
        )
        episode_returns = (
            episode_returns.batch(int(dataset.cardinality())).get_single_element().numpy()  # type: ignore
        )
        max_episode_return = np.max(episode_returns)
        min_episode_return = np.min(episode_returns)

        reward_scale = 1000.0 / (max_episode_return - min_episode_return)
        reward_bias = 0.0

    dataset = rlds.transformations.map_nested_steps(
        dataset, lambda step: rlds_utils.rescale_reward(step, reward_scale, reward_bias)
    )
    dataset = dataset.map(
        rlds_utils.skip_truncated_last_step, num_parallel_calls=tf.data.AUTOTUNE
    )

    transitions = rlds_utils.episodes_to_transitions_dataset(
        dataset,
        cycle_length=16,
        block_length=16,
        num_parallel_calls=tf.data.AUTOTUNE,
        deterministic=True,
    )

    transitions = transitions.cache()

    iterator = tfds.JaxInMemoryRandomSampleIterator(
        transitions, jax.random.PRNGKey(seed), batch_size
    )

    elapsed_time = time.time() - start_time
    logging.info("Created dataset in %.2f seconds", elapsed_time)
    return iterator


def run_offline_experiment(
    environment_factory,
    dataset_factory,
    network_factory,
    learner_core_factory,
    policy_factory,
    *,
    seed: int,
    logger_factory: loggers.LoggerFactory,
    max_num_learner_steps: int,
    num_eval_episodes: int,
    eval_every: int,
    num_sgd_steps_per_step: int,
    observers: Sequence[observers_lib.EnvLoopObserver] = (),
):
    parent_counter = counting.Counter(time_delta=0.0)

    learner_counter = counting.Counter(parent_counter, prefix="learner", time_delta=0.0)

    environment = environment_factory(seed)
    environment_spec = corax.make_environment_spec(environment)

    random_key = jax.random.PRNGKey(seed)
    learner_key, dataset_key, evaluator_key = jax.random.split(random_key, 3)

    dataset = dataset_factory(dataset_key)

    networks = network_factory(environment_spec)
    learner_core = learner_core_factory(networks)
    learner = learners_lib.DefaultJaxLearner(
        learner_core,
        learner_key,
        dataset,
        counter=learner_counter,
        logger=logger_factory("learner", learner_counter.get_steps_key()),
        num_sgd_steps_per_step=num_sgd_steps_per_step,
    )

    eval_policy = policy_factory(networks, environment_spec, True)
    eval_counter = counting.Counter(parent_counter, prefix="evaluator", time_delta=0.0)
    eval_actor = actors_lib.GenericActor(
        eval_policy,
        evaluator_key,
        variable_utils.VariableClient(learner, "policy", device="cpu"),
        per_episode_update=True,
        backend="cpu",
    )

    eval_loop = corax.EnvironmentLoop(
        environment,
        eval_actor,
        logger=logger_factory("evaluator", eval_counter.get_steps_key()),
        observers=observers,
        counter=eval_counter,
    )

    learner_counter.increment(steps=0)
    eval_loop.run(num_episodes=num_eval_episodes)
    steps = 0
    while steps < max_num_learner_steps:
        learner_steps = min(eval_every, max_num_learner_steps - steps)
        for _ in range(learner_steps):
            learner.step()
        if eval_loop:
            eval_loop.run(num_episodes=num_eval_episodes)
            steps += learner_steps


def train_and_evaluate(config: base_config.Config, workdir: Optional[str] = None):
    del workdir

    if config.num_sgd_steps_per_step > 1:
        assert config.max_num_learner_steps % config.num_sgd_steps_per_step == 0
        assert config.eval_every % config.num_sgd_steps_per_step == 0

    logger_factory = experiment_utils.LoggerFactory(
        log_to_wandb=config.log_to_wandb,
        wandb_kwargs={"project": "corax", "config": config},
        learner_time_delta=1.0,
        async_learner_logger=True,
        evaluator_time_delta=0.0,
    )

    def environment_factory(seed: int):
        return d4rl_utils.load_d4rl_environment(config.env_name, seed=seed)

    def dataset_factory(key):
        return make_d4rl_transition_dataset(
            config.env_name,
            config.batch_size * config.num_sgd_steps_per_step,
            seed=jax_utils.sample_uint32(key),
            num_episodes=config.num_episodes,
        )

    def network_factory(spec):
        return iql.make_networks(spec, hidden_dims=config.hidden_dims)

    def learner_core_factory(networks):
        if config.use_cosine_decay:
            schedule_fn = optax.cosine_decay_schedule(
                -config.policy_lr, config.max_num_learner_steps
            )
            policy_optimizer = optax.chain(
                optax.scale_by_adam(),
                optax.scale_by_schedule(schedule_fn),
            )
        else:
            policy_optimizer = optax.adam(learning_rate=config.policy_lr)

        return iql.IQLLearnerCore(
            networks,
            policy_optimizer=policy_optimizer,
            critic_optimizer=optax.adam(config.critic_lr),
            value_optimizer=optax.adam(config.value_lr),
            discount=config.discount,
            tau=config.tau,
            expectile=config.expectile,
            temperature=config.temperature,
        )

    def policy_factory(networks, spec, evaluation):
        return actor_core.batched_feed_forward_to_actor_core(
            iql.apply_policy_and_sample(networks, spec.actions, eval_mode=evaluation)
        )

    run_offline_experiment(
        environment_factory,
        dataset_factory,
        network_factory,
        learner_core_factory,
        policy_factory,
        logger_factory=logger_factory,
        seed=config.seed,
        max_num_learner_steps=(
            config.max_num_learner_steps // config.num_sgd_steps_per_step
        ),
        eval_every=config.eval_every // config.num_sgd_steps_per_step,
        num_eval_episodes=config.num_eval_episodes,
        num_sgd_steps_per_step=config.num_sgd_steps_per_step,
        observers=[d4rl_utils.D4RLScoreObserver(config.env_name)],
    )

projects/baselines/baselines/iql/train.py
ethanluoyc/corax
from absl.testing import absltest

from baselines.iql import base_config
from baselines.iql import train


class TrainTest(absltest.TestCase):
    def test_train_and_evaluate(self):
        config = base_config.get_base_config()
        config.num_episodes = 2
        config.num_eval_episodes = 1
        config.batch_size = 1
        config.hidden_dims = (8, 8)
        config.max_num_learner_steps = 10
        train.train_and_evaluate(config, workdir=None)


if __name__ == "__main__":
    absltest.main()

projects/baselines/baselines/iql/train_test.py
ethanluoyc/corax
import dataclasses


@dataclasses.dataclass
class Config:
    env_name: str = "gymnasium:HalfCheetah-v4"
    max_num_actor_steps: int = int(1e6)
    num_eval_episodes: int = 10
    eval_every: int = 10000
    hidden_dims = (256, 256)
    seed: int = 0

    policy_lr: float = 3e-4
    critic_lr: float = 3e-4
    temperature_lr: float = 3e-4

    batch_size: int = 256
    max_replay_size: int = int(1e6)
    min_replay_size: int = 5000
    utd_ratio: int = 1
    discount: float = 0.99

    log_to_wandb: bool = False


def get_config():
    return Config()

projects/baselines/baselines/redq/config.py
ethanluoyc/corax
import dataclasses
from typing import Any, Callable, Iterator, Optional

import dm_env_wrappers
import jax
import optax
import reverb
from absl import app
from absl import flags
from ml_collections import config_flags

import corax
from baselines import experiment_utils
from baselines.redq.config import Config
from corax import adders as adders_lib
from corax import environment_loop
from corax import specs
from corax import types
from corax import wrappers
from corax.adders import reverb as adders
from corax.agents.jax import actor_core
from corax.agents.jax import actors as actors_lib
from corax.agents.jax import learners as learners_lib
from corax.agents.jax import local_layout
from corax.agents.jax import redq
from corax.datasets import reverb as reverb_datasets
from corax.jax import utils as jax_utils
from corax.jax import variable_utils
from corax.utils import counting
from corax.utils import reverb_utils

# TODO(yl): Move this to corax package


@dataclasses.dataclass
class ReverbReplay:
    server: reverb.Server
    adder: adders_lib.Adder
    data_iterator: Iterator[Any]
    client: Optional[reverb.Client] = None
    can_sample: Callable[[], bool] = lambda: True


def make_local_reverb_transition_replay(
    environment_spec: specs.EnvironmentSpec,
    extra_spec: types.NestedSpec = (),
    batch_size: int = 256,
    discount: float = 0.99,
    max_replay_size: int = int(1e6),
    min_replay_size: int = 1000,
    n_step: int = 1,
    prefetch_size: int = 1,
    samples_per_insert: float = 256,
    samples_per_insert_tolerance_rate: float = 0.1,
    replay_table_name: str = adders.DEFAULT_PRIORITY_TABLE,
):
    samples_per_insert_tolerance = (
        samples_per_insert_tolerance_rate * samples_per_insert
    )
    error_buffer = min_replay_size * samples_per_insert_tolerance
    limiter = reverb.rate_limiters.SampleToInsertRatio(
        min_size_to_sample=min_replay_size,
        samples_per_insert=samples_per_insert,
        error_buffer=max(error_buffer, 2 * samples_per_insert),
    )

    # Create a replay server to add data to. This uses no limiter behavior in
    # order to allow the Agent interface to handle it.
    replay_tables = [
        reverb.Table(
            name=replay_table_name,
            sampler=reverb.selectors.Uniform(),
            remover=reverb.selectors.Fifo(),
            max_size=max_replay_size,
            rate_limiter=limiter,
            signature=adders.NStepTransitionAdder.signature(
                environment_spec, extra_spec
            ),
        )
    ]

    replay_tables, sample_sizes = reverb_utils.disable_insert_blocking(replay_tables)

    server = reverb.Server(replay_tables, port=None)

    # The adder is used to insert observations into replay.
    client = server.localhost_client()
    adder = adders.NStepTransitionAdder(client, n_step=n_step, discount=discount)

    dataset = reverb_datasets.make_reverb_dataset(
        table=replay_table_name,
        server_address=client.server_address,
        batch_size=batch_size,
        prefetch_size=prefetch_size,
    )

    dataset = dataset.map(lambda x: x.data).as_numpy_iterator()

    def can_sample():
        for table, sample_size in zip(replay_tables, sample_sizes):
            if not table.can_sample(sample_size):
                return False
        return True

    return ReverbReplay(server, adder, dataset, client, can_sample=can_sample)  # type: ignore


def train(
    environment_factory,
    network_factory,
    replay_factory,
    learner_core_factory,
    policy_factory,
    logger_factory,
    seed: int,
    *,
    max_num_actor_steps: int,
    num_eval_episodes: int,
    eval_every: int,
):
    environment = environment_factory(seed)
    spec = corax.make_environment_spec(environment)
    replay = replay_factory(spec)
    networks = network_factory(spec)

    dataset = jax_utils.device_put(replay.data_iterator, jax.local_devices()[0])  # type: ignore
    dataset = jax_utils.prefetch(dataset, buffer_size=1)

    actor_key, learner_key, eval_key = jax.random.split(jax.random.PRNGKey(seed), 3)

    parent_counter = counting.Counter(time_delta=0.0)

    learner_core = learner_core_factory(spec, networks)

    learner_counter = counting.Counter(parent_counter, prefix="learner", time_delta=0.0)
    learner = learners_lib.DefaultJaxLearner(
        learner_core,
        learner_key,
        dataset,
        logger=logger_factory("learner", learner_counter.get_steps_key()),
        counter=learner_counter,
    )

    train_policy = policy_factory(networks, evaluation=False)
    eval_policy = policy_factory(networks, evaluation=False)

    train_actor = actors_lib.GenericActor(
        train_policy,
        random_key=actor_key,
        variable_client=variable_utils.VariableClient(learner, "policy", device="cpu"),
        adder=replay.adder,
    )

    agent = local_layout.LocalLayout(train_actor, learner, dataset, replay.can_sample)

    actor_counter = counting.Counter(parent_counter, prefix="actor", time_delta=0.0)
    train_loop = environment_loop.EnvironmentLoop(
        environment,
        agent,
        logger=logger_factory("actor", actor_counter.get_steps_key()),
        counter=actor_counter,
        should_update=True,
    )

    # Create the evaluation actor and loop.
    eval_counter = counting.Counter(parent_counter, prefix="evaluator", time_delta=0.0)
    eval_logger = logger_factory("evaluator", eval_counter.get_steps_key())

    eval_actor = actors_lib.GenericActor(
        eval_policy,
        random_key=eval_key,
        variable_client=variable_utils.VariableClient(learner, "policy", device="cpu"),
    )

    eval_loop = environment_loop.EnvironmentLoop(
        environment,
        eval_actor,
        counter=eval_counter,
        logger=eval_logger,
        should_update=True,
    )

    steps = 0
    while steps < max_num_actor_steps:
        eval_loop.run(num_episodes=num_eval_episodes)
        num_steps = min(eval_every, max_num_actor_steps - steps)
        steps += train_loop.run(num_steps=num_steps)

    eval_loop.run(num_episodes=num_eval_episodes)


def make_environment(suite, task, seed):
    if suite == "gymnasium":
        import gymnasium

        environment = gymnasium.make(task)
        environment.reset(seed=seed)
        environment = wrappers.GymnasiumWrapper(environment)
        environment = dm_env_wrappers.SinglePrecisionWrapper(environment)
        environment = dm_env_wrappers.CanonicalSpecWrapper(environment, clip=True)

    elif suite == "dmc":
        from dm_control import suite as dm_suite

        domain_name, task_name = task.split(":")
        environment = dm_suite.load(
            domain_name, task_name, task_kwargs={"random": seed}
        )
        environment = dm_env_wrappers.SinglePrecisionWrapper(environment)
        environment = dm_env_wrappers.ConcatObservationWrapper(environment)
        environment = dm_env_wrappers.CanonicalSpecWrapper(environment, clip=True)
    else:
        raise ValueError(f"Unknown environment suite: {suite}")
    return environment


_CONFIG = config_flags.DEFINE_config_file("config", None)
flags.mark_flag_as_required("config")


def main(_):
    config: Config = _CONFIG.value

    logger_factory = experiment_utils.LoggerFactory(
        log_to_wandb=config.log_to_wandb,
        wandb_kwargs={"project": "corax", "config": config},
        learner_time_delta=5.0,
        evaluator_time_delta=0.0,
    )

    def environment_factory(seed: int):
        suite, task = config.env_name.split(":", 1)
        return make_environment(suite, task, seed)

    def network_factory(spec):
        return redq.make_networks(spec, hidden_sizes=config.hidden_dims)

    def replay_factory(spec):
        return make_local_reverb_transition_replay(
            spec,
            batch_size=config.batch_size,
            discount=config.discount,
            max_replay_size=config.max_replay_size,
            min_replay_size=config.min_replay_size,
            samples_per_insert=config.batch_size * config.utd_ratio,
        )

    def learner_core_factory(spec, networks):
        return redq.REDQLearnerCore(
            networks,
            policy_optimizer=optax.adam(config.policy_lr),
            critic_optimizer=optax.adam(config.critic_lr),
            temperature_optimizer=optax.adam(config.temperature_lr),
            discount=config.discount,
            target_entropy=redq.target_entropy_from_spec(spec.actions),
            utd_ratio=config.utd_ratio,
        )

    def policy_factory(networks, evaluation):
        return actor_core.batched_feed_forward_to_actor_core(
            redq.apply_policy_and_sample(networks, evaluation=evaluation)
        )

    train(
        environment_factory,
        network_factory,
        replay_factory,
        learner_core_factory,
        policy_factory,
        logger_factory,
        seed=config.seed,
        max_num_actor_steps=config.max_num_actor_steps,
        num_eval_episodes=config.num_eval_episodes,
        eval_every=config.eval_every,
    )


if __name__ == "__main__":
    app.run(main)

projects/baselines/baselines/redq/main.py
ethanluoyc/corax
# type: ignore
from typing import Optional

import rlds
import tensorflow as tf

from corax import types


def is_truncated_last_step(step: rlds.Step) -> tf.Tensor:
    return tf.logical_and(step[rlds.IS_LAST], tf.logical_not(step[rlds.IS_TERMINAL]))


def skip_truncated_last_step(episode: rlds.Episode) -> rlds.Episode:
    filtered_steps = episode[rlds.STEPS].filter(
        lambda step: tf.logical_not(is_truncated_last_step(step))
    )
    return {**episode, rlds.STEPS: filtered_steps}


def episodes_to_transitions_dataset(
    episode_dataset: tf.data.Dataset,
    cycle_length: Optional[int] = None,
    block_length: Optional[int] = None,
    num_parallel_calls: Optional[int] = None,
    deterministic: Optional[bool] = None,
):
    def _batched_step_to_transition(step: rlds.Step) -> types.Transition:
        return types.Transition(
            observation=tf.nest.map_structure(lambda x: x[0], step[rlds.OBSERVATION]),
            action=tf.nest.map_structure(lambda x: x[0], step[rlds.ACTION]),
            reward=tf.nest.map_structure(lambda x: x[0], step[rlds.REWARD]),
            discount=1.0 - tf.cast(step[rlds.IS_TERMINAL][1], dtype=tf.float32),  # type: ignore
            next_observation=tf.nest.map_structure(
                lambda x: x[1], step[rlds.OBSERVATION]
            ),
        )

    def _batch_steps(episode: rlds.Episode) -> tf.data.Dataset:
        return rlds.transformations.batch(
            episode[rlds.STEPS], size=2, shift=1, drop_remainder=True
        )

    batched_steps = episode_dataset.interleave(
        _batch_steps,
        cycle_length=cycle_length,
        block_length=block_length,
        num_parallel_calls=num_parallel_calls,
        deterministic=deterministic,
    )
    transitions = rlds.transformations.map_steps(
        batched_steps, _batched_step_to_transition
    )
    return transitions


def compute_episode_return(steps_dataset: tf.data.Dataset):
    episode_length = tf.cast(
        rlds.transformations.episode_length(steps_dataset), tf.int64
    )
    steps = steps_dataset.batch(episode_length).get_single_element()
    episode_return = tf.reduce_sum(steps["reward"])
    return episode_return


def rescale_reward(step: rlds.Step, reward_scale, reward_bias):
    return {**step, rlds.REWARD: step[rlds.REWARD] * reward_scale + reward_bias}

projects/baselines/baselines/rlds_utils.py
ethanluoyc/corax
import ml_collections


def get_config():
    config = ml_collections.ConfigDict()
    config.actor_lr = 3e-4
    config.critic_lr = 3e-4

    config.hidden_dims = (256, 256)
    config.num_qs = 10
    config.num_min_qs = 2

    config.discount = 0.99

    config.tau = 0.005

    config.temp_lr = 3e-4

    config.init_temperature = 1.0
    config.backup_entropy = True
    config.critic_layer_norm = True

    config.log_to_wandb = False
    config.env_name = "halfcheetah-expert-v2"
    config.offline_ratio = 0.5
    config.seed = 42
    config.eval_episodes = 10
    config.eval_interval = 10000
    config.batch_size = 256
    config.max_steps = int(250000)
    config.utd_ratio = 20
    config.start_training = int(1e4)

    return config

projects/baselines/baselines/rlpd/config.py
ethanluoyc/corax
import jax
from absl import app
from absl import flags
from ml_collections import config_flags

from baselines import d4rl_utils
from baselines import experiment_utils
from corax.agents.jax import redq
from corax.datasets import tfds
from corax.jax import experiments

_CONFIG = config_flags.DEFINE_config_file(
    "config", None, "File path to the training configuration."
)
flags.mark_flag_as_required("config")

_WORKDIR = flags.DEFINE_string("workdir", None, "")


def _get_demonstration_dataset_factory(name, seed):
    def make_demonstrations(batch_size):
        tfds_name = d4rl_utils.get_tfds_name(name)
        transitions = tfds.get_tfds_dataset(tfds_name)
        with jax.default_device(jax.devices("cpu")[0]):  # type: ignore
            # NOTE(yl): The yield from is necessary to
            # for some reason returning the iterator directly causes
            # sampling to be much slower.
            yield from tfds.JaxInMemoryRandomSampleIterator(
                transitions, jax.random.PRNGKey(seed), batch_size
            )

    return make_demonstrations


def main(_):
    config = _CONFIG.value
    env_name = config.env_name

    seed = config.seed

    network_factory = lambda environment_spec: redq.make_networks(
        environment_spec,
        hidden_sizes=config.hidden_dims,
        num_qs=config.num_qs,
        num_min_qs=config.num_min_qs,
        critic_layer_norm=config.critic_layer_norm,
    )

    redq_config = redq.REDQConfig(
        actor_learning_rate=config.actor_lr,
        critic_learning_rate=config.critic_lr,
        temperature_learning_rate=config.temp_lr,
        init_temperature=config.init_temperature,
        backup_entropy=config.backup_entropy,
        discount=config.discount,
        n_step=1,
        target_entropy=None,  # Automatic entropy tuning.
        # Target smoothing coefficient.
        tau=config.tau,
        max_replay_size=config.max_steps,
        batch_size=config.batch_size,
        min_replay_size=config.start_training,
        # Convert from UTD to SPI
        # In the pure online setting, SPI = batch_size is equivalent to UTD = 1
        # For RLPD, SPI = online_batch_size * UTD = (1 - offline_ratio) * batch_size * UTD
        samples_per_insert=(
            config.utd_ratio * config.batch_size * (1 - config.offline_ratio)
        ),
        # Effectively equivalent to UTD
        num_sgd_steps_per_step=config.utd_ratio,
        offline_fraction=config.offline_ratio,
        reward_bias=-1 if "antmaze" in env_name else 0,
    )

    builder = redq.REDQBuilder(
        redq_config,
        make_demonstrations=_get_demonstration_dataset_factory(env_name, seed=seed),
    )

    environment_factory = lambda seed: d4rl_utils.load_d4rl_environment(env_name, seed)

    logger_factory = experiment_utils.LoggerFactory(
        workdir=_WORKDIR.value,
        log_to_wandb=config.log_to_wandb,
        wandb_kwargs={
            "project": "rlpd",
            "config": config,
        },
        evaluator_time_delta=0.01,
        add_uid=True,
    )

    experiment = experiments.ExperimentConfig(
        builder=builder,
        environment_factory=environment_factory,
        network_factory=network_factory,
        seed=seed,
        max_num_actor_steps=config.max_steps,
        logger_factory=logger_factory,  # type: ignore
        checkpointing=None,
    )
    experiments.run_experiment(
        experiment=experiment,
        eval_every=config.eval_interval,
        num_eval_episodes=config.eval_episodes,
    )


if __name__ == "__main__":
    app.run(main)

projects/baselines/baselines/rlpd/main.py
ethanluoyc/corax
from ml_collections import config_dict


def get_config():
    config = config_dict.ConfigDict()

    config.task = "quadruped-run"
    config.action_repeat = 4
    action_repeat = config.get_oneway_ref("action_repeat")
    config.discount = 0.99
    config.episode_length = 1000 // action_repeat
    config.train_steps = 500000 // action_repeat

    # planning
    config.iterations = 6
    config.num_samples = 512
    config.num_elites = 64
    config.mixture_coef = 0.05
    config.min_std = 0.05
    config.temperature = 0.5
    config.momentum = 0.1

    # learning
    config.batch_size = 512
    config.max_buffer_size = 1000000
    config.horizon = 5
    config.reward_coef = 0.5
    config.value_coef = 0.1
    config.consistency_coef = 2
    config.rho = 0.5
    config.kappa = 0.1
    config.lr = 1e-3
    schedule_steps = 25000
    config.std_schedule = dict(
        name="linear_schedule",
        kwargs=dict(
            init_value=0.5,
            end_value=config.get_oneway_ref("min_std"),
            transition_steps=schedule_steps,
        ),
    )
    # TODO(yl): Investigate if setting to episode_length works equally well.
    # i.e.
    # config.variable_update_period = config.get_oneway_ref('episode_length')
    config.variable_update_period = 1
    config.horizon_schedule = dict(
        name="linear_schedule",
        kwargs=dict(
            init_value=1,
            end_value=config.get_oneway_ref("horizon"),
            transition_steps=schedule_steps,
        ),
    )
    config.per_alpha = 0.6
    config.per_beta = 0.4
    config.grad_clip_norm = 10
    config.seed_steps = 5000
    config.update_freq = 2
    config.tau = 0.01

    # architecture
    config.enc_dim = 256
    config.mlp_dim = 512
    config.latent_dim = 50

    # wandb (insert your own)
    config.use_wandb = False
    config.wandb_project = None
    config.wandb_entity = None
    config.wandb_name = None

    # misc
    config.seed = 1
    config.exp_name = "default"
    config.eval_freq = 20000
    config.eval_episodes = 10
    config.save_video = False
    config.save_model = False
    return config


def get_sweep():
    params = []
    action_repat_overrides = {
        "walker-walk": 2,
        "walker-run": 2,
        "cartpole-swingup": 8,
        "finger-spin": 2,
    }
    for task in [
        "cartpole-swingup",
        "quadruped-run",
        "reacher-easy",
        "walker-walk",
        "walker-run",
        "finger-spin",
        "cheetah-run",
    ]:
        params.append(
            {
                "config.task": task,
                "config.action_repeat": action_repat_overrides.get(task, 4),
            }
        )
    return params

projects/baselines/baselines/tdmpc/configs/default.py
ethanluoyc/corax
from baselines.tdmpc.configs import default


def get_config():
    config = default.get_config()
    config.task = "dog-walk"

    config.lr = 3e-4
    config.action_repeat = 2
    config.train_steps = 5000000 // config.get_oneway_ref("action_repeat")
    config.iterations = 8
    config.latent_dim = 100
    config.batch_size = 2048

    return config

projects/baselines/baselines/tdmpc/configs/dog.py
ethanluoyc/corax
from baselines.tdmpc.configs import default


def get_config():
    config = default.get_config()
    config.task = "humanoid-walk"

    config.action_repeat = 2
    config.train_steps = 3000000 / config.get_oneway_ref("action_repeat")
    config.iterations = 12
    config.latent_dim = 100

    return config

projects/baselines/baselines/tdmpc/configs/humanoid.py
ethanluoyc/corax
from baselines.tdmpc.configs import default


def get_config():
    config = default.get_config()
    config.task = "walker-walk"
    config.action_repeat = 2

    return config

projects/baselines/baselines/tdmpc/configs/walker.py
ethanluoyc/corax
from baselines.tdmpc.configs import default


def get_config():
    config = default.get_config()
    config.task = "walker-walk"
    config.action_repeat = 2

    config.lr = 3e-4
    config.rewarder = dict(
        use_encoder=False,
        num_demos=10,
        update_period=10,
    )

    # config.reward_coef = 0.5
    # config.value_coef = 0.1
    # config.consistency_coef = 2

    return config

projects/baselines/baselines/tdmpc/configs/walker_il.py
ethanluoyc/corax
# type: ignore
"""Example running TD-MPC on dm_control.
See configs/ for configurations for other environments.

"""
import copy
import functools
import os

import dm_env_wrappers as wrappers
import optax
import rlax
import rlds
import tensorflow as tf
import tensorflow_datasets as tfds
import tree
from absl import app
from absl import flags
from absl import logging
from ml_collections import config_flags

from baselines import experiment_utils
from baselines.tdmpc import rewarder as rewarder_lib
from corax import types
from corax.agents.jax import tdmpc
from corax.jax import experiments
from corax.jax import variable_utils

os.environ["MUJOCO_GL"] = "egl"

_CONFIG = config_flags.DEFINE_config_file("config", None)
_WORKDIR = flags.DEFINE_string("workdir", None, "Where to store artifacts")
flags.mark_flag_as_required("config")


def make_logger_factory(config):
    wandb_kwargs = dict(
        name=config.wandb_name,
        entity=config.wandb_entity,
        project=config.wandb_project,
        config=config.to_dict(),
        tags=[config.task],
    )

    logger_factory = experiment_utils.LoggerFactory(
        log_to_wandb=config.get("use_wandb", False),
        workdir=_WORKDIR.value,
        learner_time_delta=10.0,
        wandb_kwargs=wandb_kwargs,
    )

    return logger_factory


def make_environment_factory(config):
    def environment_factory(seed):
        # pylint: disable=import-outside-toplevel
        from dm_control import suite

        domain, task = config.task.replace("-", "_").split("_", 1)
        domain = dict(cup="ball_in_cup").get(domain, domain)
        assert (domain, task) in suite.ALL_TASKS
        env = suite.load(
            domain, task, task_kwargs={"random": seed}, visualize_reward=False
        )
        env = wrappers.ConcatObservationWrapper(env)
        env = wrappers.ActionRepeatWrapper(env, config.action_repeat)
        env = wrappers.CanonicalSpecWrapper(env, clip=True)
        env = wrappers.SinglePrecisionWrapper(env)
        return env

    return environment_factory


def _make_schedule(config):
    return getattr(optax, config.name)(**config.kwargs)


# Create function that computes episode return from rlds dataset
EPISODE_RETURN = "episode_return"


@tf.function
def _add_episode_return(episode, output_key=EPISODE_RETURN):
    """Computes the return of the episode up to the 'placed' tag."""
    episode = copy.copy(episode)
    # Truncate the episode after the placed tag.
    episode[output_key] = rlds.transformations.sum_dataset(
        episode[rlds.STEPS], lambda step: step[rlds.REWARD]
    )
    return episode


def add_episode_return(episode_dataset):
    return episode_dataset.map(
        _add_episode_return, num_parallel_calls=tf.data.experimental.AUTOTUNE
    )


def filter_by_min_return(episode_dataset, min_return):
    return episode_dataset.filter(lambda episode: episode[EPISODE_RETURN] > min_return)


def concat_observations(episode_dataset):
    def _concat(step):
        flat_observations = tree.flatten(step[rlds.OBSERVATION])
        step[rlds.OBSERVATION] = tf.concat(flat_observations, axis=-1)
        return step

    return rlds.transformations.map_nested_steps(episode_dataset, _concat)


def transform_rlds_episode_to_otil_demo(episode_dataset: tf.data.Dataset):
    demos = []
    for episode in episode_dataset:
        steps_dataset = episode["steps"]
        episode_length = int(rlds.transformations.episode_length(steps_dataset))
        steps = steps_dataset.batch(episode_length).get_single_element()
        steps = tree.map_structure(lambda x: x.numpy(), steps)

        trajectory = []

        def get_nth_item(n, nest):
            return tree.map_structure(lambda x: x[n], nest)

        # Ignore terminal observation
        for t in range(episode_length):
            trajectory.append(
                types.Transition(
                    observation=get_nth_item(t, steps["observation"]),
                    action=get_nth_item(t, steps["action"]),
                    reward=(),
                    discount=(),
                    next_observation=(),
                )
            )
        demos.append(trajectory)
    return demos


def make_experiment_config(config):
    environment_factory = make_environment_factory(config)
    logger_factory = make_logger_factory(config)
    networks_factory = functools.partial(
        tdmpc.make_networks,
        latent_size=config.latent_dim,
        encoder_hidden_size=config.enc_dim,
        mlp_hidden_size=config.mlp_dim,
    )
    optimizer = optax.chain(
        optax.clip_by_global_norm(config.grad_clip_norm),
        optax.adam(config.lr),
    )

    def rewarder_factory(variable_source, spec):
        dataset = tfds.load(
            f"rlu_control_suite/{config.task.replace('-', '_')}", split="train"
        )
        min_return = 900
        transforms = [
            add_episode_return,
            functools.partial(filter_by_min_return, min_return=min_return),
            concat_observations,
        ]
        for transform in transforms:
            dataset = transform(dataset)

        rewarder_config = config.rewarder

        dataset = dataset.take(rewarder_config.num_demos)
        demonstrations = transform_rlds_episode_to_otil_demo(dataset)

        if rewarder_config.use_encoder:
            preprocessor = rewarder_lib.MeanStdPreprocessor(
                spec.observations, spec.actions, use_actions=False
            )
            networks = networks_factory(spec)

            def encoder_fn(params, obs):
                return networks.h(params, obs)

            preprocessor = rewarder_lib.EncoderPreprocessor(
                encoder_fn, use_actions=False
            )

            variable_client = variable_utils.VariableClient(variable_source, "policy")
            preprocessor_update_period = rewarder_config.update_period
        else:
            preprocessor = rewarder_lib.MeanStdPreprocessor(
                spec.observations, spec.actions, use_actions=False
            )
            preprocessor_update_period = 1
            variable_client = None

        rewarder = rewarder_lib.ROTRewarder(
            demonstrations=demonstrations,
            episode_length=config.episode_length,
            preprocessor=preprocessor,
            preprocessor_update_period=preprocessor_update_period,
            variable_client=variable_client,
        )

        return rewarder

    std_schedule = _make_schedule(config.std_schedule)
    horizon_schedule = _make_schedule(config.horizon_schedule)
    builder = tdmpc.TDMPCBuilder(
        tdmpc.TDMPCConfig(
            std_schedule=std_schedule,
            horizon_schedule=horizon_schedule,
            optimizer=optimizer,
            batch_size=config.batch_size,
            # One update per actor step.
            samples_per_insert=config.batch_size,
            samples_per_insert_tolerance_rate=0.1,
            max_replay_size=config.max_buffer_size,
            variable_update_period=config.variable_update_period,
            importance_sampling_exponent=config.per_alpha,
            priority_exponent=config.per_beta,
            discount=config.discount,
            num_trajectories=config.num_samples,
            min_std=config.min_std,
            temperature=config.temperature,
            momentum=config.momentum,
            num_elites=config.num_elites,
            num_iterations=config.iterations,
            critic_update_rate=config.tau,
            min_replay_size=config.seed_steps,
            policy_trajectory_fraction=config.mixture_coef,
            horizon=config.horizon,
            value_tx_pair=rlax.SIGNED_LOGP1_PAIR,
            consistency_loss_scale=config.consistency_coef,
            reward_loss_scale=config.reward_coef,
            value_loss_scale=config.value_coef,
            rho=config.rho,
        )
    )

    online_experiment = experiments.ExperimentConfig(
        builder=builder,
        network_factory=networks_factory,
        environment_factory=environment_factory,
        max_num_actor_steps=config.train_steps,
        seed=config.seed,
        logger_factory=logger_factory,
        checkpointing=None,
    )
    return experiments.ImitationExperimentConfig(online_experiment, rewarder_factory)


def main(_):
    tf.config.set_visible_devices([], "GPU")
    config = _CONFIG.value
    logging.info("Config:\n%s", config)
    experiment_config = make_experiment_config(config)
    experiments.run_imitation_experiment(
        experiment_config,
        eval_every=config.eval_freq,
        num_eval_episodes=config.eval_episodes,
    )


if __name__ == "__main__":
    app.run(main)

projects/baselines/baselines/tdmpc/imitation.py
ethanluoyc/corax
# type: ignore
"""Example running TD-MPC on dm_control.

See configs/ for configurations for other environments.

"""
import functools
import os

import dm_env_wrappers
import optax
import tensorflow as tf
from absl import app
from absl import flags
from absl import logging
from ml_collections import config_flags

from baselines import experiment_utils
from corax.agents.jax import tdmpc
from corax.jax import experiments

os.environ["MUJOCO_GL"] = "egl"  # noqa

_CONFIG = config_flags.DEFINE_config_file("config", None)
_WORKDIR = flags.DEFINE_string("workdir", None, "Where to store artifacts")
flags.mark_flag_as_required("config")


def make_logger_factory(config):
    wandb_kwargs = dict(
        name=config.wandb_name,
        entity=config.wandb_entity,
        project=config.wandb_project,
        config=config.to_dict(),
        tags=[config.task],
    )

    logger_factory = experiment_utils.LoggerFactory(
        log_to_wandb=config.get("use_wandb", False),
        workdir=_WORKDIR.value,
        learner_time_delta=10.0,
        wandb_kwargs=wandb_kwargs,
    )

    return logger_factory


def make_environment_factory(config):
    def environment_factory(seed):
        # pylint: disable=import-outside-toplevel
        from dm_control import suite

        domain, task = config.task.replace("-", "_").split("_", 1)
        domain = dict(cup="ball_in_cup").get(domain, domain)
        assert (domain, task) in suite.ALL_TASKS
        env = suite.load(
            domain, task, task_kwargs={"random": seed}, visualize_reward=False
        )
        # env = wrappers.ConcatObservationWrapper(env)
        env = dm_env_wrappers.ActionRepeatWrapper(env, config.action_repeat)
        env = dm_env_wrappers.CanonicalSpecWrapper(env, clip=True)
        env = dm_env_wrappers.SinglePrecisionWrapper(env)
        return env

    return environment_factory


def _make_schedule(config):
    return getattr(optax, config.name)(**config.kwargs)


def make_experiment_config(config):
    environment_factory = make_environment_factory(config)
    logger_factory = make_logger_factory(config)
    networks_factory = functools.partial(
        tdmpc.make_networks,
        latent_size=config.latent_dim,
        encoder_hidden_size=config.enc_dim,
        mlp_hidden_size=config.mlp_dim,
    )
    optimizer = optax.chain(
        optax.clip_by_global_norm(config.grad_clip_norm),
        optax.adam(config.lr),
    )
    std_schedule = _make_schedule(config.std_schedule)
    horizon_schedule = _make_schedule(config.horizon_schedule)
    builder = tdmpc.TDMPCBuilder(
        tdmpc.TDMPCConfig(
            std_schedule=std_schedule,
            horizon_schedule=horizon_schedule,
            optimizer=optimizer,
            batch_size=config.batch_size,
            # One update per actor step.
            samples_per_insert=config.batch_size,
            samples_per_insert_tolerance_rate=0.1,
            max_replay_size=config.max_buffer_size,
            variable_update_period=config.variable_update_period,
            importance_sampling_exponent=config.per_alpha,
            priority_exponent=config.per_beta,
            discount=config.discount,
            num_trajectories=config.num_samples,
            min_std=config.min_std,
            temperature=config.temperature,
            momentum=config.momentum,
            num_elites=config.num_elites,
            num_iterations=config.iterations,
            critic_update_rate=config.tau,
            min_replay_size=config.seed_steps,
            policy_trajectory_fraction=config.mixture_coef,
            horizon=config.horizon,
            consistency_loss_scale=config.consistency_coef,
            reward_loss_scale=config.reward_coef,
            value_loss_scale=config.value_coef,
            rho=config.rho,
        )
    )

    return experiments.ExperimentConfig(
        builder=builder,
        network_factory=networks_factory,
        environment_factory=environment_factory,
        max_num_actor_steps=config.train_steps,
        seed=config.seed,
        logger_factory=logger_factory,
        checkpointing=None,
    )


def main(_):
    tf.config.set_visible_devices([], "GPU")
    config = _CONFIG.value
    logging.info("Config:\n%s", config)
    experiment_config = make_experiment_config(config)
    experiments.run_experiment(
        experiment_config,
        eval_every=config.eval_freq,
        num_eval_episodes=config.eval_episodes,
    )


if __name__ == "__main__":
    app.run(main)

projects/baselines/baselines/tdmpc/main.py
ethanluoyc/corax
# type: ignore
"""Rewarder for computing rewards used in ROT."""
import threading
from typing import Any, Callable, Optional, Protocol, Sequence

import jax
import jax.numpy as jnp
import numpy as onp
import ott
import ott.geometry.costs
import tree
from ott.geometry import pointcloud
from ott.solvers.linear import sinkhorn

from corax import types
from corax.jax import networks as networks_lib
from corax.jax import running_statistics
from corax.jax import utils
from corax.jax import variable_utils

EncoderFn = Callable[[networks_lib.Params, networks_lib.Observation], types.NestedArray]
PreprocessorState = Any


class Preprocessor(Protocol):
    """Interface for preprocessor.

    The preprocessor is used for extracting informative state
    representation from the observations.

    Implementation of the Preprocessor protocol is expected
    to be stateless. This allows the preprocessor to be used
    in a JAX jitted function.
    """

    def init(self, expert_trajectories):
        ...

    def update(self, state: PreprocessorState, trajectory) -> PreprocessorState:
        ...

    def preprocess(
        self, params: networks_lib.Params, state: PreprocessorState, trajectory
    ):
        ...


class EncoderPreprocessor(Preprocessor):
    """Parametric encoder preprocessor."""

    def __init__(self, encoder_fn: EncoderFn, use_actions: bool = False):
        self._encoder_fn = encoder_fn
        self._use_actions = use_actions

    def init(self, expert_trajectories):
        del expert_trajectories
        return ()

    def update(self, state, trajectory):
        del state, trajectory
        return ()

    def preprocess(self, params, state, trajectory):
        observations, actions = trajectory
        encoded = self._encoder_fn(params, observations)  # pylint: disable=not-callable
        if self._use_actions:
            atoms = utils.batch_concat((encoded, actions))
        else:
            atoms = encoded
        return atoms


class MeanStdPreprocessor(Preprocessor):
    """Running mean/std preprocessor."""

    def __init__(
        self,
        spec: types.NestedArray,
        action_spec: types.NestedArray = None,
        partial_update: bool = False,
        use_actions: bool = False,
    ):
        if use_actions and action_spec is None:
            raise ValueError(
                "action_spec needs to be specified when use_actions is True."
            )
        self._observation_spec = spec
        self._action_spec = action_spec
        self._partial_update = partial_update
        self._use_actions = use_actions

    def init(self, expert_trajectories):
        del expert_trajectories
        if self._use_actions:
            return running_statistics.init_state(
                (self._observation_spec, self._action_spec)
            )
        else:
            return running_statistics.init_state(self._observation_spec)

    def update(self, state, trajectory):
        observations, actions = trajectory
        if self._use_actions:
            atoms = (observations, actions)
        else:
            atoms = observations
        if self._partial_update:
            # Partial update uses only statistics from the latest atoms
            state = running_statistics.init_state(self._observation_spec)
            new_state = running_statistics.update(state, atoms)
        else:
            new_state = running_statistics.update(state, atoms)
        new_state = new_state.replace(
            std=tree.map_structure(
                lambda s: (s < 1e-5).astype(jnp.float32) + s, new_state.std
            )
        )
        return new_state

    def preprocess(self, params, state, trajectory):
        del params
        observations, actions = trajectory
        if not self._use_actions:
            atoms = observations
        else:
            atoms = (observations, actions)
        normalized = running_statistics.normalize(atoms, state)
        return utils.batch_concat(normalized)


def _concat_trajectories(trajectories):
    observations = tree.map_structure(
        lambda *xs: onp.concatenate(xs, axis=0), *[traj[0] for traj in trajectories]
    )
    actions = tree.map_structure(
        lambda *xs: onp.concatenate(xs, axis=0), *[traj[1] for traj in trajectories]
    )
    return observations, actions


class FixedStatsPreprocessor(Preprocessor):
    """Running mean/std preprocessor."""

    def __init__(self, use_actions: bool = False):
        self._use_actions = use_actions

    def init(self, expert_trajectories):
        observations, actions = _concat_trajectories(expert_trajectories)
        if self._use_actions:
            atoms = (observations, actions)
        else:
            atoms = observations

        mean = tree.map_structure(lambda x: onp.mean(x, axis=0), atoms)
        std = tree.map_structure(lambda x: onp.std(x, axis=0, dtype="float64"), atoms)
        std = tree.map_structure(lambda s: (s < 1e-5) + s, std)
        return jax.device_put((mean, std))

    def update(self, state, trajectory):
        del trajectory
        return state

    def preprocess(self, params, state, trajectory):
        del params
        observations, actions = trajectory
        mean, std = state
        if not self._use_actions:
            atoms = observations
        else:
            atoms = (observations, actions)
        return utils.batch_concat(
            tree.map_structure(lambda a, m, s: (a - m) / s, atoms, mean, std)
        )


class AggregateFn(Protocol):
    """Function for combining pseudo-rewards from multiple expert demonstrations."""

    def __call__(self, rewards: jax.Array, **kwargs) -> jax.Array:
        ...


def aggregate_top_k(rewards, k=1):
    """Aggregate rewards from multiple expert demonstrations by mean of top-K demos."""
    scores = jnp.sum(rewards, axis=-1)
    _, indices = jax.lax.top_k(scores, k=k)
    return jnp.mean(rewards[indices], axis=0)


def aggregate_mean(rewards):
    """Aggregate rewards from multiple expert demonstrations by taking the mean"""
    return jnp.mean(rewards, axis=0)


class SquashingFn(Protocol):
    """Function for squashing pseudo-rewards computed from the OT computation."""

    def __call__(self, rewards: jax.Array, **kwargs) -> jax.Array:
        ...


def squashing_linear(rewards, alpha: float = 10.0):
    """Compute squashed rewards with alpha * rewards."""
    return alpha * rewards


def squashing_exponential(rewards, alpha: float = 5.0, beta: float = 5.0):
    """Compute squashed rewards with alpha * exp(beta * rewards)."""
    return alpha * jnp.exp(beta * rewards)


class ROTRewarder:
    """ROT rewarder.

    The rewarder measures similarities with the expert demonstration.
    """

    def __init__(
        self,
        demonstrations: Sequence[Sequence[types.Transition]],
        episode_length: int,
        preprocessor: Optional[Preprocessor] = None,
        aggregate_fn: AggregateFn = aggregate_top_k,
        max_iterations: float = 100,
        threshold: float = 1e-9,
        epsilon: float = 1e-2,
        preprocessor_update_period: int = 1,
        auto_reward_scale_factor: float = 10.0,
        cost_fn: str = "cosine",
        variable_client: Optional[variable_utils.VariableClient] = None,
    ):
        if cost_fn not in ["cosine", "euclidean"]:
            raise ValueError("cost_fn must be either cosine or euclidean")
        expert_trajectories = []

        for demo in demonstrations:
            expert_trajectory = self._pack_trajectory(demo)
            expert_trajectories.append(expert_trajectory)

        self._expert_trajectories = expert_trajectories
        self._cost_fn = cost_fn
        self._episode_length = episode_length
        self._aggregate_fn = jax.jit(aggregate_fn)
        self._preprocessor_update_period = preprocessor_update_period
        self._variable_client = variable_client
        self._params = variable_client.params if variable_client is not None else None
        self._num_episodes_seen = 0
        self._reward_scale = 1.0
        self._auto_reward_scale_factor = auto_reward_scale_factor

        self._preprocessor = preprocessor
        self._preprocessor_state = self._preprocessor.init(self._expert_trajectories)
        self._mutex = threading.Lock()

        def solve_ot(expert_atoms: jax.Array, agent_atoms: jax.Array):
            # agent_atom_mean = jnp.mean(agent_atoms, axis=0)
            # agent_atom_std = jnp.std(agent_atoms, axis=0)
            # expert_atoms = (expert_atoms - agent_atom_mean) / (agent_atom_std + 1e-5)
            # agent_atoms = (agent_atoms - agent_atom_mean) / (agent_atom_std + 1e-5)
            if self._cost_fn == "cosine":
                cost_fn = ott.geometry.costs.Cosine()
                geom = pointcloud.PointCloud(
                    agent_atoms,
                    expert_atoms,
                    cost_fn=cost_fn,
                    epsilon=epsilon,
                    scale_cost=True,
                )
            elif self._cost_fn == "euclidean":
                cost_fn = ott.geometry.costs.Euclidean()
                geom = pointcloud.PointCloud(
                    agent_atoms,
                    expert_atoms,
                    cost_fn=cost_fn,
                    epsilon=epsilon,
                    scale_cost=True,
                )

            problem = ott.problems.linear.linear_problem.LinearProblem(geom)
            solver = sinkhorn.Sinkhorn(
                threshold=threshold, max_iterations=max_iterations
            )
            solver_output = solver(problem)
            coupling_matrix = geom.transport_from_potentials(
                solver_output.f, solver_output.g
            )
            cost_matrix = cost_fn.all_pairs(agent_atoms, expert_atoms)
            # Eqn (2) of https://arxiv.org/pdf/2006.04678.pdf, Eqn (5) of OTIL
            ot_costs = jnp.einsum("ij,ij->i", coupling_matrix, cost_matrix)
            return ot_costs

        self._solve_ot = jax.jit(solve_ot)
        self._update_preprocessor = jax.jit(self._preprocessor.update)
        self._preprocess = jax.jit(self._preprocessor.preprocess)

    def _pack_trajectory(self, steps):
        observations = [step.observation for step in steps]
        actions = [step.action for step in steps]
        observations = tree.map_structure(lambda *xs: onp.stack(xs, 0), *observations)
        actions = tree.map_structure(lambda *xs: onp.stack(xs, 0), *actions)
        return observations, actions

    def _compute_offline_rewards(self, agent_steps, update: bool):
        """Compute rewards based on optimal transport."""

        agent_trajectory = self._pack_trajectory(agent_steps)

        # Update preprocessor state if necessary
        if update and self._num_episodes_seen % self._preprocessor_update_period == 0:
            # Maybe update parameters
            if self._variable_client is not None:
                self._variable_client.update_and_wait()
                self._params = self._variable_client.params
            self._preprocessor_state = self._update_preprocessor(
                self._preprocessor_state, agent_trajectory
            )

        agent_atoms = self._preprocess(
            self._params, self._preprocessor_state, agent_trajectory
        )

        rewards = []
        for expert_trajectory in self._expert_trajectories:
            # NOTE(yl): It may be possible to only recompute
            # expert atoms per preprocessor update, which saves some
            # computation, especially for an encoder preprocessors.
            # However, care must be taken to ensure there are
            # no concurrency issues when we move to distributed rewarder.
            expert_atoms = self._preprocess(
                self._params, self._preprocessor_state, expert_trajectory
            )
            ot_costs = self._solve_ot(expert_atoms, agent_atoms)
            pseudo_rewards = -ot_costs
            rewards.append(pseudo_rewards)

        rewards = jnp.stack(rewards)
        rewards = self._aggregate_fn(rewards)

        if update:
            if self._num_episodes_seen == 0:
                self._reward_scale = self._auto_reward_scale_factor / onp.sum(
                    onp.abs(rewards)
                )
            self._num_episodes_seen += 1

        rewards = rewards * self._reward_scale
        return rewards[:-1]

    def compute_offline_rewards(self, agent_steps, update: bool):
        """Compute rewards based on optimal transport."""
        with self._mutex:
            return self._compute_offline_rewards(agent_steps, update)

projects/baselines/baselines/tdmpc/rewarder.py
ethanluoyc/corax
import functools
import importlib
import os
import subprocess
import sys

from absl import app
from absl import flags
from absl import logging
from lxm3 import xm
from lxm3 import xm_cluster
from lxm3.contrib import ucl
from ml_collections import config_flags

_LAUNCH_ON_CLUSTER = flags.DEFINE_boolean(
    "launch_on_cluster", False, "Launch on cluster"
)
_SINGULARITY_CONTAINER = flags.DEFINE_string(
    "container", "corax-latest.sif", "Path to singularity container"
)
_EXP_NAME = flags.DEFINE_string("exp_name", None, "Name of experiment")
_ENTRYPOINT = flags.DEFINE_string("entrypoint", None, "Module of the entrypoint")
_SWEEP = flags.DEFINE_string("sweep", None, "Name of the sweep")
_TIMEOUT = flags.DEFINE_integer("timeout", 2, "Timeout in hours")
_WANDB_GROUP = flags.DEFINE_string("wandb_group", "{xid}_{name}", "wandb group")
_WANDB_PROJECT = flags.DEFINE_string("wandb_project", "corax", "wandb project")
_WANDB_ENTITY = flags.DEFINE_string("wandb_entity", "ethanluoyc", "wandb entity")
_WANDB_MODE = flags.DEFINE_string("wandb_mode", "online", "wandb mode")
_MAX_NUM_JOBS = flags.DEFINE_integer("max_num_jobs", None, "Maximum number of jobs")
config_flags.DEFINE_config_file("config", None, "Path to config")
flags.mark_flags_as_required(["config", "entrypoint"])

FLAGS = flags.FLAGS


@functools.lru_cache()
def _get_vcs_info():
    vcs = None
    try:
        import vcsinfo

        vcs_root = subprocess.check_output(
            ["git", "rev-parse", "--show-toplevel"], text=True
        ).strip()
        vcs = vcsinfo.detect_vcs(vcs_root)
    except subprocess.SubprocessError:
        logging.warn("Failed to detect VCS info")
    return vcs


def _get_wandb_env_vars(work_unit: xm.WorkUnit, experiment_name: str):
    xid = work_unit.experiment_id
    wid = work_unit.work_unit_id
    env_vars = {
        "WANDB_PROJECT": _WANDB_PROJECT.value,
        "WANDB_ENTITY": _WANDB_ENTITY.value,
        "WANDB_NAME": f"{experiment_name}_{xid}_{wid}",
        "WANDB_MODE": _WANDB_MODE.value,
        "WANDB_RUN_GROUP": _WANDB_GROUP.value.format(name=experiment_name, xid=xid),
    }
    vcs = _get_vcs_info()
    if vcs is not None:
        env_vars["WANDB_GIT_REMOTE_URL"] = vcs.upstream_repo
        env_vars["WANDB_GIT_COMMIT"] = vcs.id

    return env_vars


def _get_hyper():
    if _SWEEP.value is not None:
        sweep_file = config_flags.get_config_filename(FLAGS["config"])
        sys.path.insert(0, os.path.abspath(os.path.dirname(sweep_file)))
        sweep_module, _ = os.path.splitext(os.path.basename(sweep_file))
        m = importlib.import_module(sweep_module)
        sys.path.pop(0)
        sweep_fn_name = f"sweep_{_SWEEP.value}"
        logging.info(f"Running sweep {sweep_fn_name}")
        sweep_fn = getattr(m, sweep_fn_name, None)
        if sweep_fn is None:
            raise ValueError(f"Sweep {sweep_fn_name} does not exist in {sweep_file}")
        else:
            return sweep_fn()
    else:
        return [{}]


def main(_):
    exp_name = _EXP_NAME.value
    if exp_name is None:
        exp_name = _ENTRYPOINT.value.replace(".", "_")
    with xm_cluster.create_experiment(experiment_title=exp_name) as experiment:
        job_requirements = xm_cluster.JobRequirements(gpu=1, ram=8 * xm.GB)
        env_vars = {"XLA_PYTHON_CLIENT_PREALLOCATE": "false"}
        if _LAUNCH_ON_CLUSTER.value:
            # TODO: Make this configurable for non-UCL clusters
            tfds_data_dir = "/cluster/project0/offline_rl/tensorflow_datasets"
            d4rl_dataset_dir = "/cluster/project0/offline_rl/d4rl"
            executor = ucl.UclGridEngine(
                job_requirements,
                walltime=_TIMEOUT.value * xm.Hr,
                singularity_options=xm_cluster.SingularityOptions(
                    bind={
                        tfds_data_dir: tfds_data_dir,
                        d4rl_dataset_dir: d4rl_dataset_dir,
                    }
                ),
            )
            env_vars["TFDS_DATA_DIR"] = tfds_data_dir
            env_vars["D4RL_DATASET_DIR"] = d4rl_dataset_dir
        else:
            executor = xm_cluster.Local(job_requirements)

        config_resource = xm_cluster.Fileset(
            files={config_flags.get_config_filename(FLAGS["config"]): "config.py"}
        )

        spec = xm_cluster.PythonPackage(
            path="..",
            entrypoint=xm_cluster.ModuleName(_ENTRYPOINT.value),
            extra_packages=[os.path.join(os.path.dirname(__file__), "../../../")],
            resources=[config_resource],
        )

        singularity_container = _SINGULARITY_CONTAINER.value

        if singularity_container:
            spec = xm_cluster.SingularityContainer(
                spec, image_path=singularity_container
            )

        args = {"config": config_resource.get_path("config.py", executor.Spec())}  # type: ignore
        overrides = config_flags.get_override_values(FLAGS["config"])
        overrides = {f"config.{k}": v for k, v in overrides.items()}
        logging.info("Overrides: %r", overrides)
        args.update(overrides)

        sweep = list(_get_hyper())
        if _MAX_NUM_JOBS.value is not None:
            sweep = sweep[: _MAX_NUM_JOBS.value]

        logging.info("Will launch %d jobs", len(sweep))

        [executable] = experiment.package(
            [
                xm.Packageable(
                    spec, executor_spec=executor.Spec(), args=args, env_vars=env_vars
                )
            ]
        )

        async def make_job(work_unit: xm.WorkUnit, **args):
            job = xm.Job(
                executable,
                executor,
                args=args,
                env_vars={**_get_wandb_env_vars(work_unit, exp_name)},
            )

            work_unit.add(job)

        with experiment.batch():
            for parameters in sweep:
                experiment.add(make_job, parameters)


if __name__ == "__main__":
    app.run(main)

projects/baselines/cluster/launcher.py
ethanluoyc/corax