deepmind-research/tvt/rma.py

# Lint as: python2, python3
# pylint: disable=g-bad-file-header
# Copyright 2019 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.
# ============================================================================
"""RMA agent."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections

import numpy as np
from six.moves import range
from six.moves import zip
import sonnet as snt
import tensorflow.compat.v1 as tf
import trfl

from tvt import losses
from tvt import memory as memory_module
from tensorflow.contrib import framework as contrib_framework

nest = contrib_framework.nest

PolicyOutputs = collections.namedtuple(
    'PolicyOutputs', ['policy', 'action', 'baseline'])

StepOutput = collections.namedtuple(
    'StepOutput', ['action', 'baseline', 'read_info'])

AgentState = collections.namedtuple(
    'AgentState', ['core_state', 'prev_action'])

Observation = collections.namedtuple(
    'Observation', ['image', 'last_action', 'last_reward'])

RNNStateNoMem = collections.namedtuple(
    'RNNStateNoMem', ['controller_outputs', 'h_controller'])

RNNState = collections.namedtuple(
    'RNNState',
    list(RNNStateNoMem._fields) + ['memory', 'mem_reads', 'h_mem_writer'])

CoreOutputs = collections.namedtuple(
    'CoreOutputs', ['action', 'policy', 'baseline', 'z', 'read_info'])


def rnn_inputs_to_static_rnn_inputs(inputs):
  """Converts time major tensors to timestep lists."""
  # Inputs to core build method are expected to be a tensor or tuple of tensors.
  if isinstance(inputs, tuple):
    num_timesteps = inputs[0].shape.as_list()[0]
    converted_inputs = [tf.unstack(input_, num_timesteps) for input_ in inputs]
    return list(zip(*converted_inputs))
  else:
    return tf.unstack(inputs)


def static_rnn_outputs_to_core_outputs(outputs):
  """Convert from length T list of nests to nest of tensors with first dim T."""
  list_of_flats = [nest.flatten(n) for n in outputs]
  new_outputs = list()
  for i in range(len(list_of_flats[0])):
    new_outputs.append(tf.stack([flat_nest[i] for flat_nest in list_of_flats]))
  return nest.pack_sequence_as(structure=outputs[0], flat_sequence=new_outputs)


def unroll(core, initial_state, inputs, dtype=tf.float32):
  """Perform a static unroll of the core."""
  static_rnn_inputs = rnn_inputs_to_static_rnn_inputs(inputs)
  static_outputs, _ = tf.nn.static_rnn(
      core,
      inputs=static_rnn_inputs,
      initial_state=initial_state,
      dtype=dtype)
  core_outputs = static_rnn_outputs_to_core_outputs(static_outputs)
  return core_outputs


class ImageEncoderDecoder(snt.AbstractModule):
  """Image Encoder/Decoder module."""

  def __init__(
      self,
      image_code_size,
      name='image_encoder_decoder'):
    """Initialize the image encoder/decoder."""
    super(ImageEncoderDecoder, self).__init__(name=name)

    # This is set by a call to `encode`. `decode` will fail before this is set.
    self._convnet_output_shape = None

    with self._enter_variable_scope():
      self._convnet = snt.nets.ConvNet2D(
          output_channels=(16, 32),
          kernel_shapes=(3, 3),
          strides=(1, 1),
          paddings=('SAME',))
      self._post_convnet_layer = snt.Linear(image_code_size, name='final_layer')

  @snt.reuse_variables
  def encode(self, image):
    """Encode the image observation."""
    convnet_output = self._convnet(image)

    # Store unflattened convnet output shape for use in decoder.
    self._convnet_output_shape = convnet_output.shape[1:]

    # Flatten convnet outputs and pass through final layer to get image code.
    return self._post_convnet_layer(snt.BatchFlatten()(convnet_output))

  @snt.reuse_variables
  def decode(self, code):
    """Decode the image observation from a latent code."""
    if self._convnet_output_shape is None:
      raise ValueError('Must call `encode` before `decode`.')
    transpose_convnet_in_flat = snt.Linear(
        self._convnet_output_shape.num_elements(),
        name='decode_initial_linear')(
            code)
    transpose_convnet_in_flat = tf.nn.relu(transpose_convnet_in_flat)
    transpose_convnet_in = snt.BatchReshape(
        self._convnet_output_shape.as_list())(transpose_convnet_in_flat)
    return self._convnet.transpose(None)(transpose_convnet_in)

  def _build(self, *args):  # Unused. Use encode/decode instead.
    raise NotImplementedError('Use encode/decode methods instead of __call__.')


class Policy(snt.AbstractModule):
  """A policy module possibly containing a read-only DNC."""

  def __init__(self,
               num_actions,
               num_policy_hiddens=(),
               num_baseline_hiddens=(),
               activation=tf.nn.tanh,
               policy_clip_abs_value=10.0,
               name='Policy'):
    """Construct a policy module possibly containing a read-only DNC.

    Args:
      num_actions: Number of discrete actions to choose from.
      num_policy_hiddens: Tuple or List, sizes of policy MLP hidden layers.
      num_baseline_hiddens: Tuple or List, sizes of baseline MLP hidden layers.
          An empty tuple/list results in a linear layer instead of an MLP.
      activation: Callable, e.g. tf.nn.tanh.
      policy_clip_abs_value: float, Policy gradient clip value.
      name: A string, the module's name
    """
    super(Policy, self).__init__(name=name)

    self._num_actions = num_actions
    self._policy_layers = tuple(num_policy_hiddens) + (num_actions,)
    self._baseline_layers = tuple(num_baseline_hiddens) + (1,)
    self._policy_clip_abs_value = policy_clip_abs_value
    self._activation = activation

  def _build(self, inputs):
    (shared_inputs, extra_policy_inputs) = inputs
    policy_in = tf.concat([shared_inputs, extra_policy_inputs], axis=1)

    policy = snt.nets.MLP(
        output_sizes=self._policy_layers,
        activation=self._activation,
        name='policy_mlp')(
            policy_in)

    # Sample an action from the policy logits.
    action = tf.multinomial(policy, num_samples=1, output_dtype=tf.int32)
    action = tf.squeeze(action, 1)  # [B, 1] -> [B]

    if self._policy_clip_abs_value > 0:
      policy = snt.clip_gradient(
          net=policy,
          clip_value_min=-self._policy_clip_abs_value,
          clip_value_max=self._policy_clip_abs_value)

    baseline_in = tf.concat([shared_inputs, tf.stop_gradient(policy)], axis=1)
    baseline = snt.nets.MLP(
        self._baseline_layers,
        activation=self._activation,
        name='baseline_mlp')(
            baseline_in)
    baseline = tf.squeeze(baseline, axis=-1)  # [B, 1] -> [B]

    if self._policy_clip_abs_value > 0:
      baseline = snt.clip_gradient(
          net=baseline,
          clip_value_min=-self._policy_clip_abs_value,
          clip_value_max=self._policy_clip_abs_value)

    outputs = PolicyOutputs(
        policy=policy,
        action=action,
        baseline=baseline)

    return outputs


class _RMACore(snt.RNNCore):
  """RMA RNN Core."""

  def __init__(self,
               num_actions,
               with_memory=True,
               name='rma_core'):
    super(_RMACore, self).__init__(name=name)

    # MLP activation as callable.
    mlp_activation = tf.nn.tanh

    # Size of latent code written to memory (if using it) and used to
    # reconstruct from (if including reconstructions).
    num_latents = 200

    # Value function decode settings.
    baseline_mlp_num_hiddens = (200,)

    # Policy settings.
    num_policy_hiddens = (200,)  # Only used for non-recurrent core.

    # Controller settings.
    control_hidden_size = 256
    control_num_layers = 2

    # Memory settings (only used if with_memory=True).
    memory_size = 1000
    memory_num_reads = 3
    memory_top_k = 50

    self._with_memory = with_memory

    with self._enter_variable_scope():
      # Construct the features -> latent encoder.
      self._z_encoder_mlp = snt.nets.MLP(
          output_sizes=(2 * num_latents, num_latents),
          activation=mlp_activation,
          activate_final=False,
          name='z_encoder_mlp')

      # Construct controller.
      rnn_cores = [snt.LSTM(control_hidden_size)
                   for _ in range(control_num_layers)]
      self._controller = snt.DeepRNN(
          rnn_cores, skip_connections=True, name='controller')

      # Construct memory.
      if self._with_memory:
        memory_dim = num_latents  # Each write to memory is of size memory_dim.
        self._mem_shape = (memory_size, memory_dim)
        self._memory_reader = memory_module.MemoryReader(
            memory_word_size=memory_dim,
            num_read_heads=memory_num_reads,
            top_k=memory_top_k,
            memory_size=memory_size)
        self._memory_writer = memory_module.MemoryWriter(
            mem_shape=self._mem_shape)

      # Construct policy, starting with policy_core and policy_action_head.
      # `extra_inputs` in this case will be mem_out from current time step (note
      # that mem_out is just the controller output if with_memory=False).
      self._policy = Policy(
          num_policy_hiddens=num_policy_hiddens,
          num_actions=num_actions,
          num_baseline_hiddens=baseline_mlp_num_hiddens,
          activation=mlp_activation,
          policy_clip_abs_value=10.0,)

    # Set state_size and output_size.
    controller_out_size = self._controller.output_size
    controller_state_size = self._controller.state_size
    self._state_size = RNNStateNoMem(controller_outputs=controller_out_size,
                                     h_controller=controller_state_size)
    read_info_size = ()
    if self._with_memory:
      mem_reads_size, read_info_size = self._memory_reader.output_size
      mem_writer_state_size = self._memory_writer.state_size
      self._state_size = RNNState(memory=tf.TensorShape(self._mem_shape),
                                  mem_reads=mem_reads_size,
                                  h_mem_writer=mem_writer_state_size,
                                  **self._state_size._asdict())

    z_size = num_latents
    self._output_size = CoreOutputs(
        action=tf.TensorShape([]),  # Scalar tensor shapes must be explicit.
        policy=num_actions,
        baseline=tf.TensorShape([]),  # Scalar tensor shapes must be explicit.
        z=z_size,
        read_info=read_info_size)

  def _build(self, inputs, h_prev):
    features = inputs

    z_net_inputs = [features, h_prev.controller_outputs]
    if self._with_memory:
      z_net_inputs.append(h_prev.mem_reads)
    z_net_inputs_concat = tf.concat(z_net_inputs, axis=1)
    z = self._z_encoder_mlp(z_net_inputs_concat)

    controller_out, h_controller = self._controller(z, h_prev.h_controller)

    read_info = ()
    if self._with_memory:
      # Perform a memory read/write step before generating the policy_modules.
      mem_reads, read_info = self._memory_reader((controller_out,
                                                  h_prev.memory))
      memory, h_mem_writer = self._memory_writer((z, h_prev.memory),
                                                 h_prev.h_mem_writer)
      policy_extra_input = tf.concat([controller_out, mem_reads], axis=1)
    else:
      policy_extra_input = controller_out

    # Get policy, action and (possible empty) baseline from policy module.
    policy_inputs = (z, policy_extra_input)
    policy_outputs = self._policy(policy_inputs)
    core_outputs = CoreOutputs(
        z=z,
        read_info=read_info,
        **policy_outputs._asdict())

    h_next = RNNStateNoMem(controller_outputs=controller_out,
                           h_controller=h_controller)
    if self._with_memory:
      h_next = RNNState(memory=memory,
                        mem_reads=mem_reads,
                        h_mem_writer=h_mem_writer,
                        **h_next._asdict())

    return core_outputs, h_next

  def initial_state(self, batch_size):
    """Use initial state for RNN modules, otherwise use zero state."""
    zero_state = self.zero_state(batch_size, dtype=tf.float32)
    controller_out = zero_state.controller_outputs
    h_controller = self._controller.initial_state(batch_size)

    state = RNNStateNoMem(controller_outputs=controller_out,
                          h_controller=h_controller)
    if self._with_memory:
      memory = zero_state.memory
      mem_reads = zero_state.mem_reads
      h_mem_writer = self._memory_writer.initial_state(batch_size)
      state = RNNState(memory=memory,
                       mem_reads=mem_reads,
                       h_mem_writer=h_mem_writer,
                       **state._asdict())
    return state

  @property
  def state_size(self):
    return self._state_size

  @property
  def output_size(self):
    return self._output_size


class Agent(snt.AbstractModule):
  """Myriad RMA agent.

  `latents` here refers to a purely deterministic encoding of the inputs, rather
  than VAE-like latents in e.g. the MERLIN agent.
  """

  def __init__(self,
               batch_size,
               with_reconstructions=True,
               with_memory=True,
               image_code_size=500,
               image_cost_weight=50.,
               num_actions=None,
               observation_shape=None,
               entropy_cost=0.01,
               return_cost_weight=0.4,
               gamma=0.96,
               read_strength_cost=5e-5,
               read_strength_tolerance=2.,
               name='rma_agent'):
    super(Agent, self).__init__(name=name)

    self._batch_size = batch_size
    self._with_reconstructions = with_reconstructions
    self._image_cost_weight = image_cost_weight
    self._image_code_size = image_code_size
    self._entropy_cost = entropy_cost
    self._return_cost_weight = return_cost_weight
    self._gamma = gamma
    self._read_strength_cost = read_strength_cost
    self._read_strength_tolerance = read_strength_tolerance
    self._num_actions = num_actions
    self._name = name
    self._logged_values = {}

    # Store total number of pixels across channels (for image loss scaling).
    self._total_num_pixels = np.prod(observation_shape)

    with self._enter_variable_scope():

      # Construct image encoder/decoder.
      self._image_encoder_decoder = ImageEncoderDecoder(
          image_code_size=image_code_size)

      self._core = _RMACore(
          num_actions=self._num_actions,
          with_memory=with_memory)

  def initial_state(self, batch_size):
    with tf.name_scope(self._name + '/initial_state'):
      return AgentState(
          core_state=self._core.initial_state(batch_size),
          prev_action=tf.zeros(shape=(batch_size,), dtype=tf.int32))

  def _prepare_observations(self, observation, last_reward, last_action):
    image = observation

    # Make sure the entries are in [0, 1) range.
    if image.dtype.is_integer:
      image = tf.cast(image, tf.float32) / 255.

    if last_reward is None:
      # For some envs, in the first timestep the last_reward can be None.
      batch_size = observation.shape[0]
      last_reward = tf.zeros((batch_size,), dtype=tf.float32)

    return Observation(
        image=image,
        last_action=last_action,
        last_reward=last_reward)

  @snt.reuse_variables
  def _encode(self, observation, last_reward, last_action):
    inputs = self._prepare_observations(observation, last_reward, last_action)

    # Encode image observation.
    obs_code = self._image_encoder_decoder.encode(inputs.image)

    # Encode last action.
    action_code = tf.one_hot(inputs.last_action, self._num_actions)

    # Encode last reward.
    reward_code = tf.expand_dims(inputs.last_reward, -1)

    features = tf.concat([obs_code, action_code, reward_code], axis=1)

    return inputs, features

  @snt.reuse_variables
  def _decode(self, z):
    # Decode image.
    image_recon = self._image_encoder_decoder.decode(z)

    # Decode action.
    action_recon = snt.Linear(self._num_actions, name='action_recon_linear')(z)

    # Decode reward.
    reward_recon = snt.Linear(1, name='reward_recon_linear')(z)

    # Full reconstructions.
    recons = Observation(
        image=image_recon,
        last_reward=reward_recon,
        last_action=action_recon)

    return recons

  def step(self, reward, observation, prev_state):
    with tf.name_scope(self._name + '/step'):
      _, features = self._encode(observation, reward, prev_state.prev_action)

      core_outputs, next_core_state = self._core(
          features, prev_state.core_state)

      action = core_outputs.action

    step_output = StepOutput(
        action=action,
        baseline=core_outputs.baseline,
        read_info=core_outputs.read_info)
    agent_state = AgentState(
        core_state=next_core_state,
        prev_action=action)
    return step_output, agent_state

  @snt.reuse_variables
  def loss(self, observations, rewards, actions, additional_rewards=None):
    """Compute the loss."""
    dummy_zeroth_step_actions = tf.zeros_like(actions[:1])
    all_actions = tf.concat([dummy_zeroth_step_actions, actions], axis=0)
    inputs, features = snt.BatchApply(self._encode)(
        observations, rewards, all_actions)

    rewards = rewards[1:]  # Zeroth step reward not correlated to actions.
    if additional_rewards is not None:
      # Additional rewards are not passed to the encoder (above) in order to be
      # consistent with the step, nor to the recon loss so that recons are
      # consistent with the observations. Thus, additional rewards only affect
      # the returns used to learn the value function.
      rewards += additional_rewards

    initial_state = self._core.initial_state(self._batch_size)

    rnn_inputs = features
    core_outputs = unroll(self._core, initial_state, rnn_inputs)

    # Remove final timestep of outputs.
    core_outputs = nest.map_structure(lambda t: t[:-1], core_outputs)

    if self._with_reconstructions:
      recons = snt.BatchApply(self._decode)(core_outputs.z)
      recon_targets = nest.map_structure(lambda t: t[:-1], inputs)
      recon_loss, recon_logged_values = losses.reconstruction_losses(
          recons=recons,
          targets=recon_targets,
          image_cost=self._image_cost_weight / self._total_num_pixels,
          action_cost=1.,
          reward_cost=1.)
    else:
      recon_loss = tf.constant(0.0)
      recon_logged_values = dict()

    if core_outputs.read_info is not tuple():
      read_reg_loss, read_reg_logged_values = (
          losses.read_regularization_loss(
              read_info=core_outputs.read_info,
              strength_cost=self._read_strength_cost,
              strength_tolerance=self._read_strength_tolerance,
              strength_reg_mode='L1',
              key_norm_cost=0.,
              key_norm_tolerance=1.))
    else:
      read_reg_loss = tf.constant(0.0)
      read_reg_logged_values = dict()

    # Bootstrap value is at end of episode so is zero.
    bootstrap_value = tf.zeros(shape=(self._batch_size,), dtype=tf.float32)

    discounts = self._gamma * tf.ones_like(rewards)

    a2c_loss, a2c_loss_extra = trfl.sequence_advantage_actor_critic_loss(
        policy_logits=core_outputs.policy,
        baseline_values=core_outputs.baseline,
        actions=actions,
        rewards=rewards,
        pcontinues=discounts,
        bootstrap_value=bootstrap_value,
        lambda_=self._gamma,
        entropy_cost=self._entropy_cost,
        baseline_cost=self._return_cost_weight,
        name='SequenceA2CLoss')

    a2c_loss = tf.reduce_mean(a2c_loss)  # Average over batch.

    total_loss = a2c_loss + recon_loss + read_reg_loss

    a2c_loss_logged_values = dict(
        pg_loss=tf.reduce_mean(a2c_loss_extra.policy_gradient_loss),
        baseline_loss=tf.reduce_mean(a2c_loss_extra.baseline_loss),
        entropy_loss=tf.reduce_mean(a2c_loss_extra.entropy_loss))
    agent_loss_log = losses.combine_logged_values(
        a2c_loss_logged_values,
        recon_logged_values,
        read_reg_logged_values)
    agent_loss_log['total_loss'] = total_loss

    return total_loss, agent_loss_log

  def _build(self, *args):  # Unused.
    # pylint: disable=no-value-for-parameter
    return self.step(*args)
    # pylint: enable=no-value-for-parameter