replay.py

# coding=utf-8
# Copyright 2022 The Google Research Authors.
#
# 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.

"""Replay buffer."""

import itertools as it
import pickle
import numpy as np


class Memory(object):
  """Container of episodes."""

  def __init__(self, data_keys=()):
    self.observations = []
    self.actions = []
    self.rewards = []
    self.data = {k: [] for k in data_keys}

  def log_init(self, obs):
    """Call this to begin logging a new episode."""
    self.observations.append([obs])
    self.actions.append([])
    self.rewards.append([])
    for value in self.data.values():
      value.append([])

  def log_experience(self, obs, act, reward, next_obs, data):
    """Add experience to the current episode."""
    assert (self.observations[-1][-1] == obs).all()
    self.observations[-1].append(next_obs)
    self.actions[-1].append(act)
    self.rewards[-1].append(reward)
    for key in data:
      self.data[key][-1].append(data[key])

  def entered_states(self):
    return np.array(list(it.chain.from_iterable(self.observations)))

  def exited_states(self):
    return np.array(list(it.chain.from_iterable(
        map(lambda obslist: obslist[0:-1], self.observations))))

  def attempted_actions(self):
    return np.array(list(it.chain.from_iterable(self.actions)))

  def observed_rewards(self):
    return np.array(list(it.chain.from_iterable(self.rewards)))

  def executed_state_action_pairs(self):
    return np.array(list(zip(self.exited_states(), self.attempted_actions())))

  def serialize(self):
    return pickle.dumps(
        (np.array(self.observations),
         np.array(self.actions),
         np.array(self.rewards),
         {key: np.array(value) for key, value in self.data.items()})
    )

  def unserialize(self, s):
    stuff = pickle.loads(s)
    obs, act, rew = stuff[:3]
    self.observations = obs
    self.actions = act
    self.rewards = rew
    if len(stuff) > 3:
      self.data = stuff[3]