Merge branch 'master' of https://github.com/tensorflow/models into cifar_keras

research/efficient-hrl/context/configs/ant_maze.gin

+#-*-Python-*-
+create_maze_env.env_name = "AntMaze"
+context_range = (%CONTEXT_RANGE_MIN, %CONTEXT_RANGE_MAX)
+meta_context_range = ((-4, -4), (20, 20))
+
+RESET_EPISODE_PERIOD = 500
+RESET_ENV_PERIOD = 1
+# End episode every N steps
+UvfAgent.reset_episode_cond_fn = @every_n_steps
+every_n_steps.n = %RESET_EPISODE_PERIOD
+train_uvf.max_steps_per_episode = %RESET_EPISODE_PERIOD
+# Do a manual reset every N episodes
+UvfAgent.reset_env_cond_fn = @every_n_episodes
+every_n_episodes.n = %RESET_ENV_PERIOD
+every_n_episodes.steps_per_episode = %RESET_EPISODE_PERIOD
+
+## Config defaults
+EVAL_MODES = ["eval1", "eval2", "eval3"]
+
+## Config agent
+CONTEXT = @agent/Context
+META_CONTEXT = @meta/Context
+
+## Config agent context
+agent/Context.context_ranges = [%context_range]
+agent/Context.context_shapes = [%SUBGOAL_DIM]
+agent/Context.meta_action_every_n = 10
+agent/Context.samplers = {
+    "train": [@train/DirectionSampler],
+    "explore": [@train/DirectionSampler],
+}
+
+agent/Context.context_transition_fn = @relative_context_transition_fn
+agent/Context.context_multi_transition_fn = @relative_context_multi_transition_fn
+
+agent/Context.reward_fn = @uvf/negative_distance
+
+## Config meta context
+meta/Context.context_ranges = [%meta_context_range]
+meta/Context.context_shapes = [2]
+meta/Context.samplers = {
+    "train": [@train/RandomSampler],
+    "explore": [@train/RandomSampler],
+    "eval1": [@eval1/ConstantSampler],
+    "eval2": [@eval2/ConstantSampler],
+    "eval3": [@eval3/ConstantSampler],
+}
+meta/Context.reward_fn = @task/negative_distance
+
+## Config rewards
+task/negative_distance.state_indices = [0, 1]
+task/negative_distance.relative_context = False
+task/negative_distance.diff = False
+task/negative_distance.offset = 0.0
+
+## Config samplers
+train/RandomSampler.context_range = %meta_context_range
+train/DirectionSampler.context_range = %context_range
+train/DirectionSampler.k = %SUBGOAL_DIM
+relative_context_transition_fn.k = %SUBGOAL_DIM
+relative_context_multi_transition_fn.k = %SUBGOAL_DIM
+MetaAgent.k = %SUBGOAL_DIM
+
+eval1/ConstantSampler.value = [16, 0]
+eval2/ConstantSampler.value = [16, 16]
+eval3/ConstantSampler.value = [0, 16]

research/efficient-hrl/context/configs/ant_maze_img.gin

+#-*-Python-*-
+create_maze_env.env_name = "AntMaze"
+IMAGES = True
+
+context_range = (%CONTEXT_RANGE_MIN, %CONTEXT_RANGE_MAX)
+meta_context_range = ((-4, -4), (20, 20))
+
+RESET_EPISODE_PERIOD = 500
+RESET_ENV_PERIOD = 1
+# End episode every N steps
+UvfAgent.reset_episode_cond_fn = @every_n_steps
+every_n_steps.n = %RESET_EPISODE_PERIOD
+train_uvf.max_steps_per_episode = %RESET_EPISODE_PERIOD
+# Do a manual reset every N episodes
+UvfAgent.reset_env_cond_fn = @every_n_episodes
+every_n_episodes.n = %RESET_ENV_PERIOD
+every_n_episodes.steps_per_episode = %RESET_EPISODE_PERIOD
+
+## Config defaults
+EVAL_MODES = ["eval1", "eval2", "eval3"]
+
+## Config agent
+CONTEXT = @agent/Context
+META_CONTEXT = @meta/Context
+
+## Config agent context
+agent/Context.context_ranges = [%context_range]
+agent/Context.context_shapes = [%SUBGOAL_DIM]
+agent/Context.meta_action_every_n = 10
+agent/Context.samplers = {
+    "train": [@train/DirectionSampler],
+    "explore": [@train/DirectionSampler],
+}
+
+agent/Context.context_transition_fn = @relative_context_transition_fn
+agent/Context.context_multi_transition_fn = @relative_context_multi_transition_fn
+
+agent/Context.reward_fn = @uvf/negative_distance
+
+## Config meta context
+meta/Context.context_ranges = [%meta_context_range]
+meta/Context.context_shapes = [2]
+meta/Context.samplers = {
+    "train": [@train/RandomSampler],
+    "explore": [@train/RandomSampler],
+    "eval1": [@eval1/ConstantSampler],
+    "eval2": [@eval2/ConstantSampler],
+    "eval3": [@eval3/ConstantSampler],
+}
+meta/Context.context_transition_fn = @task/relative_context_transition_fn
+meta/Context.context_multi_transition_fn = @task/relative_context_multi_transition_fn
+meta/Context.reward_fn = @task/negative_distance
+
+## Config rewards
+task/negative_distance.state_indices = [0, 1]
+task/negative_distance.relative_context = True
+task/negative_distance.diff = False
+task/negative_distance.offset = 0.0
+
+## Config samplers
+train/RandomSampler.context_range = %meta_context_range
+train/DirectionSampler.context_range = %context_range
+train/DirectionSampler.k = %SUBGOAL_DIM
+relative_context_transition_fn.k = %SUBGOAL_DIM
+relative_context_multi_transition_fn.k = %SUBGOAL_DIM
+task/relative_context_transition_fn.k = 2
+task/relative_context_multi_transition_fn.k = 2
+MetaAgent.k = %SUBGOAL_DIM
+
+eval1/ConstantSampler.value = [16, 0]
+eval2/ConstantSampler.value = [16, 16]
+eval3/ConstantSampler.value = [0, 16]

research/efficient-hrl/context/configs/ant_push_multi.gin

+#-*-Python-*-
+create_maze_env.env_name = "AntPush"
+context_range = (%CONTEXT_RANGE_MIN, %CONTEXT_RANGE_MAX)
+meta_context_range = ((-16, -4), (16, 20))
+
+RESET_EPISODE_PERIOD = 500
+RESET_ENV_PERIOD = 1
+# End episode every N steps
+UvfAgent.reset_episode_cond_fn = @every_n_steps
+every_n_steps.n = %RESET_EPISODE_PERIOD
+train_uvf.max_steps_per_episode = %RESET_EPISODE_PERIOD
+# Do a manual reset every N episodes
+UvfAgent.reset_env_cond_fn = @every_n_episodes
+every_n_episodes.n = %RESET_ENV_PERIOD
+every_n_episodes.steps_per_episode = %RESET_EPISODE_PERIOD
+
+## Config defaults
+EVAL_MODES = ["eval2"]
+
+## Config agent
+CONTEXT = @agent/Context
+META_CONTEXT = @meta/Context
+
+## Config agent context
+agent/Context.context_ranges = [%context_range]
+agent/Context.context_shapes = [%SUBGOAL_DIM]
+agent/Context.meta_action_every_n = 10
+agent/Context.samplers = {
+    "train": [@train/DirectionSampler],
+    "explore": [@train/DirectionSampler],
+}
+
+agent/Context.context_transition_fn = @relative_context_transition_fn
+agent/Context.context_multi_transition_fn = @relative_context_multi_transition_fn
+
+agent/Context.reward_fn = @uvf/negative_distance
+
+## Config meta context
+meta/Context.context_ranges = [%meta_context_range]
+meta/Context.context_shapes = [2]
+meta/Context.samplers = {
+    "train": [@train/RandomSampler],
+    "explore": [@train/RandomSampler],
+    "eval2": [@eval2/ConstantSampler],
+}
+meta/Context.reward_fn = @task/negative_distance
+
+## Config rewards
+task/negative_distance.state_indices = [0, 1]
+task/negative_distance.relative_context = False
+task/negative_distance.diff = False
+task/negative_distance.offset = 0.0
+
+## Config samplers
+train/RandomSampler.context_range = %meta_context_range
+train/DirectionSampler.context_range = %context_range
+train/DirectionSampler.k = %SUBGOAL_DIM
+relative_context_transition_fn.k = %SUBGOAL_DIM
+relative_context_multi_transition_fn.k = %SUBGOAL_DIM
+MetaAgent.k = %SUBGOAL_DIM
+
+eval2/ConstantSampler.value = [0, 19]

research/efficient-hrl/context/configs/ant_push_multi_img.gin

+#-*-Python-*-
+create_maze_env.env_name = "AntPush"
+IMAGES = True
+
+context_range = (%CONTEXT_RANGE_MIN, %CONTEXT_RANGE_MAX)
+meta_context_range = ((-16, -4), (16, 20))
+
+RESET_EPISODE_PERIOD = 500
+RESET_ENV_PERIOD = 1
+# End episode every N steps
+UvfAgent.reset_episode_cond_fn = @every_n_steps
+every_n_steps.n = %RESET_EPISODE_PERIOD
+train_uvf.max_steps_per_episode = %RESET_EPISODE_PERIOD
+# Do a manual reset every N episodes
+UvfAgent.reset_env_cond_fn = @every_n_episodes
+every_n_episodes.n = %RESET_ENV_PERIOD
+every_n_episodes.steps_per_episode = %RESET_EPISODE_PERIOD
+
+## Config defaults
+EVAL_MODES = ["eval2"]
+
+## Config agent
+CONTEXT = @agent/Context
+META_CONTEXT = @meta/Context
+
+## Config agent context
+agent/Context.context_ranges = [%context_range]
+agent/Context.context_shapes = [%SUBGOAL_DIM]
+agent/Context.meta_action_every_n = 10
+agent/Context.samplers = {
+    "train": [@train/DirectionSampler],
+    "explore": [@train/DirectionSampler],
+}
+
+agent/Context.context_transition_fn = @relative_context_transition_fn
+agent/Context.context_multi_transition_fn = @relative_context_multi_transition_fn
+
+agent/Context.reward_fn = @uvf/negative_distance
+
+## Config meta context
+meta/Context.context_ranges = [%meta_context_range]
+meta/Context.context_shapes = [2]
+meta/Context.samplers = {
+    "train": [@train/RandomSampler],
+    "explore": [@train/RandomSampler],
+    "eval2": [@eval2/ConstantSampler],
+}
+meta/Context.context_transition_fn = @task/relative_context_transition_fn
+meta/Context.context_multi_transition_fn = @task/relative_context_multi_transition_fn
+meta/Context.reward_fn = @task/negative_distance
+
+## Config rewards
+task/negative_distance.state_indices = [0, 1]
+task/negative_distance.relative_context = True
+task/negative_distance.diff = False
+task/negative_distance.offset = 0.0
+
+## Config samplers
+train/RandomSampler.context_range = %meta_context_range
+train/DirectionSampler.context_range = %context_range
+train/DirectionSampler.k = %SUBGOAL_DIM
+relative_context_transition_fn.k = %SUBGOAL_DIM
+relative_context_multi_transition_fn.k = %SUBGOAL_DIM
+task/relative_context_transition_fn.k = 2
+task/relative_context_multi_transition_fn.k = 2
+MetaAgent.k = %SUBGOAL_DIM
+
+eval2/ConstantSampler.value = [0, 19]

research/efficient-hrl/context/configs/ant_push_single.gin

+#-*-Python-*-
+create_maze_env.env_name = "AntPush"
+context_range = (%CONTEXT_RANGE_MIN, %CONTEXT_RANGE_MAX)
+meta_context_range = ((-16, -4), (16, 20))
+
+RESET_EPISODE_PERIOD = 500
+RESET_ENV_PERIOD = 1
+# End episode every N steps
+UvfAgent.reset_episode_cond_fn = @every_n_steps
+every_n_steps.n = %RESET_EPISODE_PERIOD
+train_uvf.max_steps_per_episode = %RESET_EPISODE_PERIOD
+# Do a manual reset every N episodes
+UvfAgent.reset_env_cond_fn = @every_n_episodes
+every_n_episodes.n = %RESET_ENV_PERIOD
+every_n_episodes.steps_per_episode = %RESET_EPISODE_PERIOD
+
+## Config defaults
+EVAL_MODES = ["eval2"]
+
+## Config agent
+CONTEXT = @agent/Context
+META_CONTEXT = @meta/Context
+
+## Config agent context
+agent/Context.context_ranges = [%context_range]
+agent/Context.context_shapes = [%SUBGOAL_DIM]
+agent/Context.meta_action_every_n = 10
+agent/Context.samplers = {
+    "train": [@train/DirectionSampler],
+    "explore": [@train/DirectionSampler],
+}
+
+agent/Context.context_transition_fn = @relative_context_transition_fn
+agent/Context.context_multi_transition_fn = @relative_context_multi_transition_fn
+
+agent/Context.reward_fn = @uvf/negative_distance
+
+## Config meta context
+meta/Context.context_ranges = [%meta_context_range]
+meta/Context.context_shapes = [2]
+meta/Context.samplers = {
+    "train": [@eval2/ConstantSampler],
+    "explore": [@eval2/ConstantSampler],
+    "eval2": [@eval2/ConstantSampler],
+}
+meta/Context.reward_fn = @task/negative_distance
+
+## Config rewards
+task/negative_distance.state_indices = [0, 1]
+task/negative_distance.relative_context = False
+task/negative_distance.diff = False
+task/negative_distance.offset = 0.0
+
+## Config samplers
+train/RandomSampler.context_range = %meta_context_range
+train/DirectionSampler.context_range = %context_range
+train/DirectionSampler.k = %SUBGOAL_DIM
+relative_context_transition_fn.k = %SUBGOAL_DIM
+relative_context_multi_transition_fn.k = %SUBGOAL_DIM
+MetaAgent.k = %SUBGOAL_DIM
+
+eval2/ConstantSampler.value = [0, 19]

research/efficient-hrl/context/configs/default.gin

+#-*-Python-*-
+ENV_CONTEXT = None
+EVAL_MODES = ["eval"]
+TARGET_Q_CLIPPING = None
+RESET_EPISODE_PERIOD = None
+ZERO_OBS = False
+CONTEXT_RANGE_MIN = -10
+CONTEXT_RANGE_MAX = 10
+SUBGOAL_DIM = 2
+
+uvf/negative_distance.summarize = False
+uvf/negative_distance.relative_context = True

research/efficient-hrl/context/configs/hiro_orig.gin

+#-*-Python-*-
+ENV_CONTEXT = None
+EVAL_MODES = ["eval"]
+TARGET_Q_CLIPPING = None
+RESET_EPISODE_PERIOD = None
+ZERO_OBS = True
+IMAGES = False
+CONTEXT_RANGE_MIN = (-10, -10, -0.5, -1, -1, -1, -1, -0.5, -0.3, -0.5, -0.3, -0.5, -0.3, -0.5, -0.3)
+CONTEXT_RANGE_MAX = ( 10,  10,  0.5,  1,  1,  1,  1,  0.5,  0.3,  0.5,  0.3,  0.5,  0.3,  0.5,  0.3)
+SUBGOAL_DIM = 15
+META_EXPLORE_NOISE = 1.0
+
+uvf/negative_distance.summarize = False
+uvf/negative_distance.relative_context = True

research/efficient-hrl/context/configs/hiro_repr.gin

+#-*-Python-*-
+ENV_CONTEXT = None
+EVAL_MODES = ["eval"]
+TARGET_Q_CLIPPING = None
+RESET_EPISODE_PERIOD = None
+ZERO_OBS = False
+IMAGES = False
+CONTEXT_RANGE_MIN = -10
+CONTEXT_RANGE_MAX = 10
+SUBGOAL_DIM = 2
+META_EXPLORE_NOISE = 5.0
+
+StatePreprocess.trainable = True
+StatePreprocess.state_preprocess_net = @state_preprocess_net
+StatePreprocess.action_embed_net = @action_embed_net
+
+uvf/negative_distance.summarize = False
+uvf/negative_distance.relative_context = True

research/efficient-hrl/context/configs/hiro_xy.gin

+#-*-Python-*-
+ENV_CONTEXT = None
+EVAL_MODES = ["eval"]
+TARGET_Q_CLIPPING = None
+RESET_EPISODE_PERIOD = None
+ZERO_OBS = False
+IMAGES = False
+CONTEXT_RANGE_MIN = -10
+CONTEXT_RANGE_MAX = 10
+SUBGOAL_DIM = 2
+META_EXPLORE_NOISE = 1.0
+
+uvf/negative_distance.summarize = False
+uvf/negative_distance.relative_context = True

research/efficient-hrl/context/configs/point_maze.gin

+#-*-Python-*-
+# NOTE: For best training, low-level exploration (uvf_add_noise_fn.stddev)
+# should be reduced to around 0.1.
+create_maze_env.env_name = "PointMaze"
+context_range_min = -10
+context_range_max = 10
+context_range = (%context_range_min, %context_range_max)
+meta_context_range = ((-2, -2), (10, 10))
+
+RESET_EPISODE_PERIOD = 500
+RESET_ENV_PERIOD = 1
+# End episode every N steps
+UvfAgent.reset_episode_cond_fn = @every_n_steps
+every_n_steps.n = %RESET_EPISODE_PERIOD
+train_uvf.max_steps_per_episode = %RESET_EPISODE_PERIOD
+# Do a manual reset every N episodes
+UvfAgent.reset_env_cond_fn = @every_n_episodes
+every_n_episodes.n = %RESET_ENV_PERIOD
+every_n_episodes.steps_per_episode = %RESET_EPISODE_PERIOD
+
+## Config defaults
+EVAL_MODES = ["eval1", "eval2", "eval3"]
+
+## Config agent
+CONTEXT = @agent/Context
+META_CONTEXT = @meta/Context
+
+## Config agent context
+agent/Context.context_ranges = [%context_range]
+agent/Context.context_shapes = [%SUBGOAL_DIM]
+agent/Context.meta_action_every_n = 10
+agent/Context.samplers = {
+    "train": [@train/DirectionSampler],
+    "explore": [@train/DirectionSampler],
+    "eval1": [@uvf_eval1/ConstantSampler],
+    "eval2": [@uvf_eval2/ConstantSampler],
+    "eval3": [@uvf_eval3/ConstantSampler],
+}
+
+agent/Context.context_transition_fn = @relative_context_transition_fn
+agent/Context.context_multi_transition_fn = @relative_context_multi_transition_fn
+
+agent/Context.reward_fn = @uvf/negative_distance
+
+## Config meta context
+meta/Context.context_ranges = [%meta_context_range]
+meta/Context.context_shapes = [2]
+meta/Context.samplers = {
+    "train": [@train/RandomSampler],
+    "explore": [@train/RandomSampler],
+    "eval1": [@eval1/ConstantSampler],
+    "eval2": [@eval2/ConstantSampler],
+    "eval3": [@eval3/ConstantSampler],
+}
+meta/Context.reward_fn = @task/negative_distance
+
+## Config rewards
+task/negative_distance.state_indices = [0, 1]
+task/negative_distance.relative_context = False
+task/negative_distance.diff = False
+task/negative_distance.offset = 0.0
+
+## Config samplers
+train/RandomSampler.context_range = %meta_context_range
+train/DirectionSampler.context_range = %context_range
+train/DirectionSampler.k = %SUBGOAL_DIM
+relative_context_transition_fn.k = %SUBGOAL_DIM
+relative_context_multi_transition_fn.k = %SUBGOAL_DIM
+MetaAgent.k = %SUBGOAL_DIM
+
+eval1/ConstantSampler.value = [8, 0]
+eval2/ConstantSampler.value = [8, 8]
+eval3/ConstantSampler.value = [0, 8]

research/efficient-hrl/context/context.py

+# Copyright 2018 The TensorFlow Authors 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.
+# ==============================================================================
+
+"""Context for Universal Value Function agents.
+
+A context specifies a list of contextual variables, each with
+  own sampling and reward computation methods.
+
+Examples of contextual variables include
+  goal states, reward combination vectors, etc.
+"""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+import numpy as np
+import tensorflow as tf
+from tf_agents import specs
+import gin.tf
+from utils import utils as uvf_utils
+
+
+@gin.configurable
+class Context(object):
+  """Base context."""
+  VAR_NAME = 'action'
+
+  def __init__(self,
+               tf_env,
+               context_ranges=None,
+               context_shapes=None,
+               state_indices=None,
+               variable_indices=None,
+               gamma_index=None,
+               settable_context=False,
+               timers=None,
+               samplers=None,
+               reward_weights=None,
+               reward_fn=None,
+               random_sampler_mode='random',
+               normalizers=None,
+               context_transition_fn=None,
+               context_multi_transition_fn=None,
+               meta_action_every_n=None):
+    self._tf_env = tf_env
+    self.variable_indices = variable_indices
+    self.gamma_index = gamma_index
+    self._settable_context = settable_context
+    self.timers = timers
+    self._context_transition_fn = context_transition_fn
+    self._context_multi_transition_fn = context_multi_transition_fn
+    self._random_sampler_mode = random_sampler_mode
+
+    # assign specs
+    self._obs_spec = self._tf_env.observation_spec()
+    self._context_shapes = tuple([
+        shape if shape is not None else self._obs_spec.shape
+        for shape in context_shapes
+    ])
+    self.context_specs = tuple([
+        specs.TensorSpec(dtype=self._obs_spec.dtype, shape=shape)
+        for shape in self._context_shapes
+    ])
+    if context_ranges is not None:
+      self.context_ranges = context_ranges
+    else:
+      self.context_ranges = [None] * len(self._context_shapes)
+
+    self.context_as_action_specs = tuple([
+        specs.BoundedTensorSpec(
+            shape=shape,
+            dtype=(tf.float32 if self._obs_spec.dtype in
+                   [tf.float32, tf.float64] else self._obs_spec.dtype),
+            minimum=context_range[0],
+            maximum=context_range[-1])
+        for shape, context_range in zip(self._context_shapes, self.context_ranges)
+    ])
+
+    if state_indices is not None:
+      self.state_indices = state_indices
+    else:
+      self.state_indices = [None] * len(self._context_shapes)
+    if self.variable_indices is not None and self.n != len(
+        self.variable_indices):
+      raise ValueError(
+          'variable_indices (%s) must have the same length as contexts (%s).' %
+          (self.variable_indices, self.context_specs))
+    assert self.n == len(self.context_ranges)
+    assert self.n == len(self.state_indices)
+
+    # assign reward/sampler fns
+    self._sampler_fns = dict()
+    self._samplers = dict()
+    self._reward_fns = dict()
+
+    # assign reward fns
+    self._add_custom_reward_fns()
+    reward_weights = reward_weights or None
+    self._reward_fn = self._make_reward_fn(reward_fn, reward_weights)
+
+    # assign samplers
+    self._add_custom_sampler_fns()
+    for mode, sampler_fns in samplers.items():
+      self._make_sampler_fn(sampler_fns, mode)
+
+    # create normalizers
+    if normalizers is None:
+      self._normalizers = [None] * len(self.context_specs)
+    else:
+      self._normalizers = [
+          normalizer(tf.zeros(shape=spec.shape, dtype=spec.dtype))
+          if normalizer is not None else None
+          for normalizer, spec in zip(normalizers, self.context_specs)
+      ]
+    assert self.n == len(self._normalizers)
+
+    self.meta_action_every_n = meta_action_every_n
+
+    # create vars
+    self.context_vars = {}
+    self.timer_vars = {}
+    self.create_vars(self.VAR_NAME)
+    self.t = tf.Variable(
+        tf.zeros(shape=(), dtype=tf.int32), name='num_timer_steps')
+
+  def _add_custom_reward_fns(self):
+    pass
+
+  def _add_custom_sampler_fns(self):
+    pass
+
+  def sample_random_contexts(self, batch_size):
+    """Sample random batch contexts."""
+    assert self._random_sampler_mode is not None
+    return self.sample_contexts(self._random_sampler_mode, batch_size)[0]
+
+  def sample_contexts(self, mode, batch_size, state=None, next_state=None,
+                      **kwargs):
+    """Sample a batch of contexts.
+
+    Args:
+      mode: A string representing the mode [`train`, `explore`, `eval`].
+      batch_size: Batch size.
+    Returns:
+      Two lists of [batch_size, num_context_dims] contexts.
+    """
+    contexts, next_contexts = self._sampler_fns[mode](
+        batch_size, state=state, next_state=next_state,
+        **kwargs)
+    self._validate_contexts(contexts)
+    self._validate_contexts(next_contexts)
+    return contexts, next_contexts
+
+  def compute_rewards(self, mode, states, actions, rewards, next_states,
+                      contexts):
+    """Compute context-based rewards.
+
+    Args:
+      mode: A string representing the mode ['uvf', 'task'].
+      states: A [batch_size, num_state_dims] tensor.
+      actions: A [batch_size, num_action_dims] tensor.
+      rewards: A [batch_size] tensor representing unmodified rewards.
+      next_states: A [batch_size, num_state_dims] tensor.
+      contexts: A list of [batch_size, num_context_dims] tensors.
+    Returns:
+      A [batch_size] tensor representing rewards.
+    """
+    return self._reward_fn(states, actions, rewards, next_states,
+                           contexts)
+
+  def _make_reward_fn(self, reward_fns_list, reward_weights):
+    """Returns a fn that computes rewards.
+
+    Args:
+      reward_fns_list: A fn or a list of reward fns.
+      mode: A string representing the operating mode.
+      reward_weights: A list of reward weights.
+    """
+    if not isinstance(reward_fns_list, (list, tuple)):
+      reward_fns_list = [reward_fns_list]
+    if reward_weights is None:
+      reward_weights = [1.0] * len(reward_fns_list)
+    assert len(reward_fns_list) == len(reward_weights)
+
+    reward_fns_list = [
+        self._custom_reward_fns[fn] if isinstance(fn, (str,)) else fn
+        for fn in reward_fns_list
+    ]
+
+    def reward_fn(*args, **kwargs):
+      """Returns rewards, discounts."""
+      reward_tuples = [
+          reward_fn(*args, **kwargs) for reward_fn in reward_fns_list
+      ]
+      rewards_list = [reward_tuple[0] for reward_tuple in reward_tuples]
+      discounts_list = [reward_tuple[1] for reward_tuple in reward_tuples]
+      ndims = max([r.shape.ndims for r in rewards_list])
+      if ndims > 1:  # expand reward shapes to allow broadcasting
+        for i in range(len(rewards_list)):
+          for _ in range(rewards_list[i].shape.ndims - ndims):
+            rewards_list[i] = tf.expand_dims(rewards_list[i], axis=-1)
+          for _ in range(discounts_list[i].shape.ndims - ndims):
+            discounts_list[i] = tf.expand_dims(discounts_list[i], axis=-1)
+      rewards = tf.add_n(
+          [r * tf.to_float(w) for r, w in zip(rewards_list, reward_weights)])
+      discounts = discounts_list[0]
+      for d in discounts_list[1:]:
+        discounts *= d
+
+      return rewards, discounts
+
+    return reward_fn
+
+  def _make_sampler_fn(self, sampler_cls_list, mode):
+    """Returns a fn that samples a list of context vars.
+
+    Args:
+      sampler_cls_list: A list of sampler classes.
+      mode: A string representing the operating mode.
+    """
+    if not isinstance(sampler_cls_list, (list, tuple)):
+      sampler_cls_list = [sampler_cls_list]
+
+    self._samplers[mode] = []
+    sampler_fns = []
+    for spec, sampler in zip(self.context_specs, sampler_cls_list):
+      if isinstance(sampler, (str,)):
+        sampler_fn = self._custom_sampler_fns[sampler]
+      else:
+        sampler_fn = sampler(context_spec=spec)
+        self._samplers[mode].append(sampler_fn)
+      sampler_fns.append(sampler_fn)
+
+    def batch_sampler_fn(batch_size, state=None, next_state=None, **kwargs):
+      """Sampler fn."""
+      contexts_tuples = [
+          sampler(batch_size, state=state, next_state=next_state, **kwargs)
+          for sampler in sampler_fns]
+      contexts = [c[0] for c in contexts_tuples]
+      next_contexts = [c[1] for c in contexts_tuples]
+      contexts = [
+          normalizer.update_apply(c) if normalizer is not None else c
+          for normalizer, c in zip(self._normalizers, contexts)
+      ]
+      next_contexts = [
+          normalizer.apply(c) if normalizer is not None else c
+          for normalizer, c in zip(self._normalizers, next_contexts)
+      ]
+      return contexts, next_contexts
+
+    self._sampler_fns[mode] = batch_sampler_fn
+
+  def set_env_context_op(self, context, disable_unnormalizer=False):
+    """Returns a TensorFlow op that sets the environment context.
+
+    Args:
+      context: A list of context Tensor variables.
+      disable_unnormalizer: Disable unnormalization.
+    Returns:
+      A TensorFlow op that sets the environment context.
+    """
+    ret_val = np.array(1.0, dtype=np.float32)
+    if not self._settable_context:
+      return tf.identity(ret_val)
+
+    if not disable_unnormalizer:
+      context = [
+          normalizer.unapply(tf.expand_dims(c, 0))[0]
+          if normalizer is not None else c
+          for normalizer, c in zip(self._normalizers, context)
+      ]
+
+    def set_context_func(*env_context_values):
+      tf.logging.info('[set_env_context_op] Setting gym environment context.')
+      # pylint: disable=protected-access
+      self.gym_env.set_context(*env_context_values)
+      return ret_val
+      # pylint: enable=protected-access
+
+    with tf.name_scope('set_env_context'):
+      set_op = tf.py_func(set_context_func, context, tf.float32,
+                          name='set_env_context_py_func')
+      set_op.set_shape([])
+    return set_op
+
+  def set_replay(self, replay):
+    """Set replay buffer for samplers.
+
+    Args:
+      replay: A replay buffer.
+    """
+    for _, samplers in self._samplers.items():
+      for sampler in samplers:
+        sampler.set_replay(replay)
+
+  def get_clip_fns(self):
+    """Returns a list of clip fns for contexts.
+
+    Returns:
+      A list of fns that clip context tensors.
+    """
+    clip_fns = []
+    for context_range in self.context_ranges:
+      def clip_fn(var_, range_=context_range):
+        """Clip a tensor."""
+        if range_ is None:
+          clipped_var = tf.identity(var_)
+        elif isinstance(range_[0], (int, long, float, list, np.ndarray)):
+          clipped_var = tf.clip_by_value(
+              var_,
+              range_[0],
+              range_[1],)
+        else: raise NotImplementedError(range_)
+        return clipped_var
+      clip_fns.append(clip_fn)
+    return clip_fns
+
+  def _validate_contexts(self, contexts):
+    """Validate if contexts have right specs.
+
+    Args:
+      contexts: A list of [batch_size, num_context_dim] tensors.
+    Raises:
+      ValueError: If shape or dtype mismatches that of spec.
+    """
+    for i, (context, spec) in enumerate(zip(contexts, self.context_specs)):
+      if context[0].shape != spec.shape:
+        raise ValueError('contexts[%d] has invalid shape %s wrt spec shape %s' %
+                         (i, context[0].shape, spec.shape))
+      if context.dtype != spec.dtype:
+        raise ValueError('contexts[%d] has invalid dtype %s wrt spec dtype %s' %
+                         (i, context.dtype, spec.dtype))
+
+  def context_multi_transition_fn(self, contexts, **kwargs):
+    """Returns multiple future contexts starting from a batch."""
+    assert self._context_multi_transition_fn
+    return self._context_multi_transition_fn(contexts, None, None, **kwargs)
+
+  def step(self, mode, agent=None, action_fn=None, **kwargs):
+    """Returns [next_contexts..., next_timer] list of ops.
+
+    Args:
+      mode: a string representing the mode=[train, explore, eval].
+      **kwargs: kwargs for context_transition_fn.
+    Returns:
+      a list of ops that set the context.
+    """
+    if agent is None:
+      ops = []
+      if self._context_transition_fn is not None:
+        def sampler_fn():
+          samples = self.sample_contexts(mode, 1)[0]
+          return [s[0] for s in samples]
+        values = self._context_transition_fn(self.vars, self.t, sampler_fn, **kwargs)
+        ops += [tf.assign(var, value) for var, value in zip(self.vars, values)]
+      ops.append(tf.assign_add(self.t, 1))  # increment timer
+      return ops
+    else:
+      ops = agent.tf_context.step(mode, **kwargs)
+      state = kwargs['state']
+      next_state = kwargs['next_state']
+      state_repr = kwargs['state_repr']
+      next_state_repr = kwargs['next_state_repr']
+      with tf.control_dependencies(ops):  # Step high level context before computing low level one.
+        # Get the context transition function output.
+        values = self._context_transition_fn(self.vars, self.t, None,
+                                             state=state_repr,
+                                             next_state=next_state_repr)
+        # Select a new goal every C steps, otherwise use context transition.
+        low_level_context = [
+            tf.cond(tf.equal(self.t % self.meta_action_every_n, 0),
+                    lambda: tf.cast(action_fn(next_state, context=None), tf.float32),
+                    lambda: values)]
+        ops = [tf.assign(var, value)
+               for var, value in zip(self.vars, low_level_context)]
+        with tf.control_dependencies(ops):
+          return [tf.assign_add(self.t, 1)]  # increment timer
+        return ops
+
+  def reset(self, mode, agent=None, action_fn=None, state=None):
+    """Returns ops that reset the context.
+
+    Args:
+      mode: a string representing the mode=[train, explore, eval].
+    Returns:
+      a list of ops that reset the context.
+    """
+    if agent is None:
+      values = self.sample_contexts(mode=mode, batch_size=1)[0]
+      if values is None:
+        return []
+      values = [value[0] for value in values]
+      values[0] = uvf_utils.tf_print(
+          values[0],
+          values,
+          message='context:reset, mode=%s' % mode,
+          first_n=10,
+          name='context:reset:%s' % mode)
+      all_ops = []
+      for _, context_vars in sorted(self.context_vars.items()):
+        ops = [tf.assign(var, value) for var, value in zip(context_vars, values)]
+      all_ops += ops
+      all_ops.append(self.set_env_context_op(values))
+      all_ops.append(tf.assign(self.t, 0))  # reset timer
+      return all_ops
+    else:
+      ops = agent.tf_context.reset(mode)
+      # NOTE: The code is currently written in such a way that the higher level
+      # policy does not provide a low-level context until the second
+      # observation.  Insead, we just zero-out low-level contexts.
+      for key, context_vars in sorted(self.context_vars.items()):
+        ops += [tf.assign(var, tf.zeros_like(var)) for var, meta_var in
+                zip(context_vars, agent.tf_context.context_vars[key])]
+
+      ops.append(tf.assign(self.t, 0))  # reset timer
+      return ops
+
+  def create_vars(self, name, agent=None):
+    """Create tf variables for contexts.
+
+    Args:
+      name: Name of the variables.
+    Returns:
+      A list of [num_context_dims] tensors.
+    """
+    if agent is not None:
+      meta_vars = agent.create_vars(name)
+    else:
+      meta_vars = {}
+    assert name not in self.context_vars, ('Conflict! %s is already '
+                                           'initialized.') % name
+    self.context_vars[name] = tuple([
+        tf.Variable(
+            tf.zeros(shape=spec.shape, dtype=spec.dtype),
+            name='%s_context_%d' % (name, i))
+        for i, spec in enumerate(self.context_specs)
+    ])
+    return self.context_vars[name], meta_vars
+
+  @property
+  def n(self):
+    return len(self.context_specs)
+
+  @property
+  def vars(self):
+    return self.context_vars[self.VAR_NAME]
+
+  # pylint: disable=protected-access
+  @property
+  def gym_env(self):
+    return self._tf_env.pyenv._gym_env
+
+  @property
+  def tf_env(self):
+    return self._tf_env
+  # pylint: enable=protected-access

research/efficient-hrl/context/context_transition_functions.py

+# Copyright 2018 The TensorFlow Authors 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.
+# ==============================================================================
+
+"""Context functions.
+
+Given the current contexts, timer and context sampler, returns new contexts
+  after an environment step. This can be used to define a high-level policy
+  that controls contexts as its actions.
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow as tf
+import gin.tf
+import utils as uvf_utils
+
+
+@gin.configurable
+def periodic_context_fn(contexts, timer, sampler_fn, period=1):
+  """Periodically samples contexts.
+
+  Args:
+    contexts: a list of [num_context_dims] tensor variables representing
+      current contexts.
+    timer: a scalar integer tensor variable holding the current time step.
+    sampler_fn: a sampler function that samples a list of [num_context_dims]
+      tensors.
+    period: (integer) period of update.
+  Returns:
+    a list of [num_context_dims] tensors.
+  """
+  contexts = list(contexts[:])  # create copy
+  return tf.cond(tf.mod(timer, period) == 0, sampler_fn, lambda: contexts)
+
+
+@gin.configurable
+def timer_context_fn(contexts,
+                     timer,
+                     sampler_fn,
+                     period=1,
+                     timer_index=-1,
+                     debug=False):
+  """Samples contexts based on timer in contexts.
+
+  Args:
+    contexts: a list of [num_context_dims] tensor variables representing
+      current contexts.
+    timer: a scalar integer tensor variable holding the current time step.
+    sampler_fn: a sampler function that samples a list of [num_context_dims]
+      tensors.
+    period: (integer) period of update; actual period = `period` + 1.
+    timer_index: (integer) Index of context list that present timer.
+    debug: (boolean) Print debug messages.
+  Returns:
+    a list of [num_context_dims] tensors.
+  """
+  contexts = list(contexts[:])  # create copy
+  cond = tf.equal(contexts[timer_index][0], 0)
+  def reset():
+    """Sample context and reset the timer."""
+    new_contexts = sampler_fn()
+    new_contexts[timer_index] = tf.zeros_like(
+        contexts[timer_index]) + period
+    return new_contexts
+  def update():
+    """Decrement the timer."""
+    contexts[timer_index] -= 1
+    return contexts
+  values = tf.cond(cond, reset, update)
+  if debug:
+    values[0] = uvf_utils.tf_print(
+        values[0],
+        values + [timer],
+        'timer_context_fn',
+        first_n=200,
+        name='timer_context_fn:contexts')
+  return values
+
+
+@gin.configurable
+def relative_context_transition_fn(
+    contexts, timer, sampler_fn,
+    k=2, state=None, next_state=None,
+    **kwargs):
+  """Contexts updated to be relative to next state.
+  """
+  contexts = list(contexts[:])  # create copy
+  assert len(contexts) == 1
+  new_contexts = [
+      tf.concat(
+          [contexts[0][:k] + state[:k] - next_state[:k],
+           contexts[0][k:]], -1)]
+  return new_contexts
+
+
+@gin.configurable
+def relative_context_multi_transition_fn(
+    contexts, timer, sampler_fn,
+    k=2, states=None,
+    **kwargs):
+  """Given contexts at first state and sequence of states, derives sequence of all contexts.
+  """
+  contexts = list(contexts[:])  # create copy
+  assert len(contexts) == 1
+  contexts = [
+      tf.concat(
+          [tf.expand_dims(contexts[0][:, :k] + states[:, 0, :k], 1) - states[:, :, :k],
+           contexts[0][:, None, k:] * tf.ones_like(states[:, :, :1])], -1)]
+  return contexts

research/astronet/astronet/astro_cnn_model/__init__.py → research/efficient-hrl/context/gin_imports.py

-# Copyright 2018 The TensorFlow Authors.
+# Copyright 2018 The TensorFlow Authors 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.
@@ -11,4 +11,15 @@
 # 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.
+# ==============================================================================
 
+"""Import gin configurable modules.
+"""
+
+# pylint: disable=unused-import
+from context import context
+from context import context_transition_functions
+from context import gin_utils
+from context import rewards_functions
+from context import samplers
+# pylint: disable=unused-import

research/astronet/light_curve_util/cc/python/median_filter_test.py → research/efficient-hrl/context/gin_utils.py

-# Copyright 2018 The TensorFlow Authors.
+# Copyright 2018 The TensorFlow Authors 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.
@@ -11,38 +11,35 @@
 # 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 the Python wrapping of the median_filter library."""
+"""Gin configurable utility functions.
+"""
 
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-from absl.testing import absltest
 import numpy as np
+import gin.tf
 
-from light_curve_util.cc.python import median_filter
 
+@gin.configurable
+def gin_sparse_array(size, values, indices, fill_value=0):
+  arr = np.zeros(size)
+  arr.fill(fill_value)
+  arr[indices] = values
+  return arr
 
-class MedianFilterTest(absltest.TestCase):
 
-  def testError(self):
-    x = [2, 0, 1]
-    y = [1, 2, 3]
-    with self.assertRaises(ValueError):
-      median_filter.median_filter(
-          x, y, num_bins=2, bin_width=1, x_min=0, x_max=2)
+@gin.configurable
+def gin_sum(values):
+  result = values[0]
+  for value in values[1:]:
+    result += value
+  return result
 
-  def testMedianFilter(self):
-    x = np.arange(-6, 7)
-    y = np.arange(1, 14)
 
-    result = median_filter.median_filter(
-        x, y, num_bins=5, bin_width=2, x_min=-5, x_max=5)
-
-    expected = [2.5, 4.5, 6.5, 8.5, 10.5]
-    np.testing.assert_almost_equal(result, expected)
-
-
-if __name__ == "__main__":
-  absltest.main()
+@gin.configurable
+def gin_range(n):
+  return range(n)

research/efficient-hrl/context/rewards_functions.py

+# Copyright 2018 The TensorFlow Authors 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.
+# ==============================================================================
+
+"""Reward shaping functions used by Contexts.
+
+  Each reward function should take the following inputs and return new rewards,
+    and discounts.
+
+  new_rewards, discounts = reward_fn(states, actions, rewards,
+    next_states, contexts)
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow as tf
+import gin.tf
+
+
+def summarize_stats(stats):
+  """Summarize a dictionary of variables.
+
+  Args:
+    stats: a dictionary of {name: tensor} to compute stats over.
+  """
+  for name, stat in stats.items():
+    mean = tf.reduce_mean(stat)
+    tf.summary.scalar('mean_%s' % name, mean)
+    tf.summary.scalar('max_%s' % name, tf.reduce_max(stat))
+    tf.summary.scalar('min_%s' % name, tf.reduce_min(stat))
+    std = tf.sqrt(tf.reduce_mean(tf.square(stat)) - tf.square(mean) + 1e-10)
+    tf.summary.scalar('std_%s' % name, std)
+    tf.summary.histogram(name, stat)
+
+
+def index_states(states, indices):
+  """Return indexed states.
+
+  Args:
+    states: A [batch_size, num_state_dims] Tensor representing a batch
+        of states.
+    indices: (a list of Numpy integer array) Indices of states dimensions
+      to be mapped.
+  Returns:
+    A [batch_size, num_indices] Tensor representing the batch of indexed states.
+  """
+  if indices is None:
+    return states
+  indices = tf.constant(indices, dtype=tf.int32)
+  return tf.gather(states, indices=indices, axis=1)
+
+
+def record_tensor(tensor, indices, stats, name='states'):
+  """Record specified tensor dimensions into stats.
+
+  Args:
+    tensor: A [batch_size, num_dims] Tensor.
+    indices: (a list of integers) Indices of dimensions to record.
+    stats: A dictionary holding stats.
+    name: (string) Name of tensor.
+  """
+  if indices is None:
+    indices = range(tensor.shape.as_list()[1])
+  for index in indices:
+    stats['%s_%02d' % (name, index)] = tensor[:, index]
+
+
+@gin.configurable
+def potential_rewards(states,
+                      actions,
+                      rewards,
+                      next_states,
+                      contexts,
+                      gamma=1.0,
+                      reward_fn=None):
+  """Return the potential-based rewards.
+
+  Args:
+    states: A [batch_size, num_state_dims] Tensor representing a batch
+        of states.
+    actions: A [batch_size, num_action_dims] Tensor representing a batch
+      of actions.
+    rewards: A [batch_size] Tensor representing a batch of rewards.
+    next_states: A [batch_size, num_state_dims] Tensor representing a batch
+      of next states.
+    contexts: A list of [batch_size, num_context_dims] Tensor representing
+      a batch of contexts.
+    gamma: Reward discount.
+    reward_fn: A reward function.
+  Returns:
+    A new tf.float32 [batch_size] rewards Tensor, and
+      tf.float32 [batch_size] discounts tensor.
+  """
+  del actions  # unused args
+  gamma = tf.to_float(gamma)
+  rewards_tp1, discounts = reward_fn(None, None, rewards, next_states, contexts)
+  rewards, _ = reward_fn(None, None, rewards, states, contexts)
+  return -rewards + gamma * rewards_tp1, discounts
+
+
+@gin.configurable
+def timed_rewards(states,
+                  actions,
+                  rewards,
+                  next_states,
+                  contexts,
+                  reward_fn=None,
+                  dense=False,
+                  timer_index=-1):
+  """Return the timed rewards.
+
+  Args:
+    states: A [batch_size, num_state_dims] Tensor representing a batch
+        of states.
+    actions: A [batch_size, num_action_dims] Tensor representing a batch
+      of actions.
+    rewards: A [batch_size] Tensor representing a batch of rewards.
+    next_states: A [batch_size, num_state_dims] Tensor representing a batch
+      of next states.
+    contexts: A list of [batch_size, num_context_dims] Tensor representing
+      a batch of contexts.
+    reward_fn: A reward function.
+    dense: (boolean) Provide dense rewards or sparse rewards at time = 0.
+    timer_index: (integer) The context list index that specifies timer.
+  Returns:
+    A new tf.float32 [batch_size] rewards Tensor, and
+      tf.float32 [batch_size] discounts tensor.
+  """
+  assert contexts[timer_index].get_shape().as_list()[1] == 1
+  timers = contexts[timer_index][:, 0]
+  rewards, discounts = reward_fn(states, actions, rewards, next_states,
+                                 contexts)
+  terminates = tf.to_float(timers <= 0)  # if terminate set 1, else set 0
+  for _ in range(rewards.shape.ndims - 1):
+    terminates = tf.expand_dims(terminates, axis=-1)
+  if not dense:
+    rewards *= terminates  # if terminate, return rewards, else return 0
+  discounts *= (tf.to_float(1.0) - terminates)
+  return rewards, discounts
+
+
+@gin.configurable
+def reset_rewards(states,
+                  actions,
+                  rewards,
+                  next_states,
+                  contexts,
+                  reset_index=0,
+                  reset_state=None,
+                  reset_reward_function=None,
+                  include_forward_rewards=True,
+                  include_reset_rewards=True):
+  """Returns the rewards for a forward/reset agent.
+
+  Args:
+    states: A [batch_size, num_state_dims] Tensor representing a batch
+        of states.
+    actions: A [batch_size, num_action_dims] Tensor representing a batch
+      of actions.
+    rewards: A [batch_size] Tensor representing a batch of rewards.
+    next_states: A [batch_size, num_state_dims] Tensor representing a batch
+      of next states.
+    contexts: A list of [batch_size, num_context_dims] Tensor representing
+      a batch of contexts.
+    reset_index: (integer) The context list index that specifies reset.
+    reset_state: Reset state.
+    reset_reward_function: Reward function for reset step.
+    include_forward_rewards: Include the rewards from the forward pass.
+    include_reset_rewards: Include the rewards from the reset pass.
+
+  Returns:
+    A new tf.float32 [batch_size] rewards Tensor, and
+      tf.float32 [batch_size] discounts tensor.
+  """
+  reset_state = tf.constant(
+      reset_state, dtype=next_states.dtype, shape=next_states.shape)
+  reset_states = tf.expand_dims(reset_state, 0)
+
+  def true_fn():
+    if include_reset_rewards:
+      return reset_reward_function(states, actions, rewards, next_states,
+                                   [reset_states] + contexts[1:])
+    else:
+      return tf.zeros_like(rewards), tf.ones_like(rewards)
+
+  def false_fn():
+    if include_forward_rewards:
+      return plain_rewards(states, actions, rewards, next_states, contexts)
+    else:
+      return tf.zeros_like(rewards), tf.ones_like(rewards)
+
+  rewards, discounts = tf.cond(
+      tf.cast(contexts[reset_index][0, 0], dtype=tf.bool), true_fn, false_fn)
+  return rewards, discounts
+
+
+@gin.configurable
+def tanh_similarity(states,
+                    actions,
+                    rewards,
+                    next_states,
+                    contexts,
+                    mse_scale=1.0,
+                    state_scales=1.0,
+                    goal_scales=1.0,
+                    summarize=False):
+  """Returns the similarity between next_states and contexts using tanh and mse.
+
+  Args:
+    states: A [batch_size, num_state_dims] Tensor representing a batch
+        of states.
+    actions: A [batch_size, num_action_dims] Tensor representing a batch
+      of actions.
+    rewards: A [batch_size] Tensor representing a batch of rewards.
+    next_states: A [batch_size, num_state_dims] Tensor representing a batch
+      of next states.
+    contexts: A list of [batch_size, num_context_dims] Tensor representing
+      a batch of contexts.
+    mse_scale: A float, to scale mse before tanh.
+    state_scales: multiplicative scale for (next) states. A scalar or 1D tensor,
+      must be broadcastable to number of state dimensions.
+    goal_scales: multiplicative scale for contexts. A scalar or 1D tensor,
+      must be broadcastable to number of goal dimensions.
+    summarize: (boolean) enable summary ops.
+
+
+  Returns:
+    A new tf.float32 [batch_size] rewards Tensor, and
+      tf.float32 [batch_size] discounts tensor.
+  """
+  del states, actions, rewards  # Unused
+  mse = tf.reduce_mean(tf.squared_difference(next_states * state_scales,
+                                             contexts[0] * goal_scales), -1)
+  tanh = tf.tanh(mse_scale * mse)
+  if summarize:
+    with tf.name_scope('RewardFn/'):
+      tf.summary.scalar('mean_mse', tf.reduce_mean(mse))
+      tf.summary.histogram('mse', mse)
+      tf.summary.scalar('mean_tanh', tf.reduce_mean(tanh))
+      tf.summary.histogram('tanh', tanh)
+  rewards = tf.to_float(1 - tanh)
+  return rewards, tf.ones_like(rewards)
+
+
+@gin.configurable
+def negative_mse(states,
+                 actions,
+                 rewards,
+                 next_states,
+                 contexts,
+                 state_scales=1.0,
+                 goal_scales=1.0,
+                 summarize=False):
+  """Returns the negative mean square error between next_states and contexts.
+
+  Args:
+    states: A [batch_size, num_state_dims] Tensor representing a batch
+        of states.
+    actions: A [batch_size, num_action_dims] Tensor representing a batch
+      of actions.
+    rewards: A [batch_size] Tensor representing a batch of rewards.
+    next_states: A [batch_size, num_state_dims] Tensor representing a batch
+      of next states.
+    contexts: A list of [batch_size, num_context_dims] Tensor representing
+      a batch of contexts.
+    state_scales: multiplicative scale for (next) states. A scalar or 1D tensor,
+      must be broadcastable to number of state dimensions.
+    goal_scales: multiplicative scale for contexts. A scalar or 1D tensor,
+      must be broadcastable to number of goal dimensions.
+    summarize: (boolean) enable summary ops.
+
+  Returns:
+    A new tf.float32 [batch_size] rewards Tensor, and
+      tf.float32 [batch_size] discounts tensor.
+  """
+  del states, actions, rewards  # Unused
+  mse = tf.reduce_mean(tf.squared_difference(next_states * state_scales,
+                                             contexts[0] * goal_scales), -1)
+  if summarize:
+    with tf.name_scope('RewardFn/'):
+      tf.summary.scalar('mean_mse', tf.reduce_mean(mse))
+      tf.summary.histogram('mse', mse)
+  rewards = tf.to_float(-mse)
+  return rewards, tf.ones_like(rewards)
+
+
+@gin.configurable
+def negative_distance(states,
+                      actions,
+                      rewards,
+                      next_states,
+                      contexts,
+                      state_scales=1.0,
+                      goal_scales=1.0,
+                      reward_scales=1.0,
+                      weight_index=None,
+                      weight_vector=None,
+                      summarize=False,
+                      termination_epsilon=1e-4,
+                      state_indices=None,
+                      goal_indices=None,
+                      vectorize=False,
+                      relative_context=False,
+                      diff=False,
+                      norm='L2',
+                      epsilon=1e-10,
+                      bonus_epsilon=0., #5.,
+                      offset=0.0):
+  """Returns the negative euclidean distance between next_states and contexts.
+
+  Args:
+    states: A [batch_size, num_state_dims] Tensor representing a batch
+        of states.
+    actions: A [batch_size, num_action_dims] Tensor representing a batch
+      of actions.
+    rewards: A [batch_size] Tensor representing a batch of rewards.
+    next_states: A [batch_size, num_state_dims] Tensor representing a batch
+      of next states.
+    contexts: A list of [batch_size, num_context_dims] Tensor representing
+      a batch of contexts.
+    state_scales: multiplicative scale for (next) states. A scalar or 1D tensor,
+      must be broadcastable to number of state dimensions.
+    goal_scales: multiplicative scale for goals. A scalar or 1D tensor,
+      must be broadcastable to number of goal dimensions.
+    reward_scales: multiplicative scale for rewards. A scalar or 1D tensor,
+      must be broadcastable to number of reward dimensions.
+    weight_index: (integer) The context list index that specifies weight.
+    weight_vector: (a number or a list or Numpy array) The weighting vector,
+      broadcastable to `next_states`.
+    summarize: (boolean) enable summary ops.
+    termination_epsilon: terminate if dist is less than this quantity.
+    state_indices: (a list of integers) list of state indices to select.
+    goal_indices: (a list of integers) list of goal indices to select.
+    vectorize: Return a vectorized form.
+    norm: L1 or L2.
+    epsilon: small offset to ensure non-negative/zero distance.
+
+  Returns:
+    A new tf.float32 [batch_size] rewards Tensor, and
+      tf.float32 [batch_size] discounts tensor.
+  """
+  del actions, rewards  # Unused
+  stats = {}
+  record_tensor(next_states, state_indices, stats, 'next_states')
+  states = index_states(states, state_indices)
+  next_states = index_states(next_states, state_indices)
+  goals = index_states(contexts[0], goal_indices)
+  if relative_context:
+    goals = states + goals
+  sq_dists = tf.squared_difference(next_states * state_scales,
+                                   goals * goal_scales)
+  old_sq_dists = tf.squared_difference(states * state_scales,
+                                       goals * goal_scales)
+  record_tensor(sq_dists, None, stats, 'sq_dists')
+  if weight_vector is not None:
+    sq_dists *= tf.convert_to_tensor(weight_vector, dtype=next_states.dtype)
+    old_sq_dists *= tf.convert_to_tensor(weight_vector, dtype=next_states.dtype)
+  if weight_index is not None:
+    #sq_dists *= contexts[weight_index]
+    weights = tf.abs(index_states(contexts[0], weight_index))
+    #weights /= tf.reduce_sum(weights, -1, keepdims=True)
+    sq_dists *= weights
+    old_sq_dists *= weights
+  if norm == 'L1':
+    dist = tf.sqrt(sq_dists + epsilon)
+    old_dist = tf.sqrt(old_sq_dists + epsilon)
+    if not vectorize:
+      dist = tf.reduce_sum(dist, -1)
+      old_dist = tf.reduce_sum(old_dist, -1)
+  elif norm == 'L2':
+    if vectorize:
+      dist = sq_dists
+      old_dist = old_sq_dists
+    else:
+      dist = tf.reduce_sum(sq_dists, -1)
+      old_dist = tf.reduce_sum(old_sq_dists, -1)
+    dist = tf.sqrt(dist + epsilon)  # tf.gradients fails when tf.sqrt(-0.0)
+    old_dist = tf.sqrt(old_dist + epsilon)  # tf.gradients fails when tf.sqrt(-0.0)
+  else:
+    raise NotImplementedError(norm)
+  discounts = dist > termination_epsilon
+  if summarize:
+    with tf.name_scope('RewardFn/'):
+      tf.summary.scalar('mean_dist', tf.reduce_mean(dist))
+      tf.summary.histogram('dist', dist)
+      summarize_stats(stats)
+  bonus = tf.to_float(dist < bonus_epsilon)
+  dist *= reward_scales
+  old_dist *= reward_scales
+  if diff:
+    return bonus + offset + tf.to_float(old_dist - dist), tf.to_float(discounts)
+  return bonus + offset + tf.to_float(-dist), tf.to_float(discounts)
+
+
+@gin.configurable
+def cosine_similarity(states,
+                      actions,
+                      rewards,
+                      next_states,
+                      contexts,
+                      state_scales=1.0,
+                      goal_scales=1.0,
+                      reward_scales=1.0,
+                      normalize_states=True,
+                      normalize_goals=True,
+                      weight_index=None,
+                      weight_vector=None,
+                      summarize=False,
+                      state_indices=None,
+                      goal_indices=None,
+                      offset=0.0):
+  """Returns the cosine similarity between next_states - states and contexts.
+
+  Args:
+    states: A [batch_size, num_state_dims] Tensor representing a batch
+        of states.
+    actions: A [batch_size, num_action_dims] Tensor representing a batch
+      of actions.
+    rewards: A [batch_size] Tensor representing a batch of rewards.
+    next_states: A [batch_size, num_state_dims] Tensor representing a batch
+      of next states.
+    contexts: A list of [batch_size, num_context_dims] Tensor representing
+      a batch of contexts.
+    state_scales: multiplicative scale for (next) states. A scalar or 1D tensor,
+      must be broadcastable to number of state dimensions.
+    goal_scales: multiplicative scale for goals. A scalar or 1D tensor,
+      must be broadcastable to number of goal dimensions.
+    reward_scales: multiplicative scale for rewards. A scalar or 1D tensor,
+      must be broadcastable to number of reward dimensions.
+    weight_index: (integer) The context list index that specifies weight.
+    weight_vector: (a number or a list or Numpy array) The weighting vector,
+      broadcastable to `next_states`.
+    summarize: (boolean) enable summary ops.
+    termination_epsilon: terminate if dist is less than this quantity.
+    state_indices: (a list of integers) list of state indices to select.
+    goal_indices: (a list of integers) list of goal indices to select.
+    vectorize: Return a vectorized form.
+    norm: L1 or L2.
+    epsilon: small offset to ensure non-negative/zero distance.
+
+  Returns:
+    A new tf.float32 [batch_size] rewards Tensor, and
+      tf.float32 [batch_size] discounts tensor.
+  """
+  del actions, rewards  # Unused
+  stats = {}
+  record_tensor(next_states, state_indices, stats, 'next_states')
+  states = index_states(states, state_indices)
+  next_states = index_states(next_states, state_indices)
+  goals = index_states(contexts[0], goal_indices)
+
+  if weight_vector is not None:
+    goals *= tf.convert_to_tensor(weight_vector, dtype=next_states.dtype)
+  if weight_index is not None:
+    weights = tf.abs(index_states(contexts[0], weight_index))
+    goals *= weights
+
+  direction_vec = next_states - states
+  if normalize_states:
+    direction_vec = tf.nn.l2_normalize(direction_vec, -1)
+  goal_vec = goals
+  if normalize_goals:
+    goal_vec = tf.nn.l2_normalize(goal_vec, -1)
+
+  similarity = tf.reduce_sum(goal_vec * direction_vec, -1)
+  discounts = tf.ones_like(similarity)
+  return offset + tf.to_float(similarity), tf.to_float(discounts)
+
+
+@gin.configurable
+def diff_distance(states,
+                  actions,
+                  rewards,
+                  next_states,
+                  contexts,
+                  state_scales=1.0,
+                  goal_scales=1.0,
+                  reward_scales=1.0,
+                  weight_index=None,
+                  weight_vector=None,
+                  summarize=False,
+                  termination_epsilon=1e-4,
+                  state_indices=None,
+                  goal_indices=None,
+                  norm='L2',
+                  epsilon=1e-10):
+  """Returns the difference in euclidean distance between states/next_states and contexts.
+
+  Args:
+    states: A [batch_size, num_state_dims] Tensor representing a batch
+        of states.
+    actions: A [batch_size, num_action_dims] Tensor representing a batch
+      of actions.
+    rewards: A [batch_size] Tensor representing a batch of rewards.
+    next_states: A [batch_size, num_state_dims] Tensor representing a batch
+      of next states.
+    contexts: A list of [batch_size, num_context_dims] Tensor representing
+      a batch of contexts.
+    state_scales: multiplicative scale for (next) states. A scalar or 1D tensor,
+      must be broadcastable to number of state dimensions.
+    goal_scales: multiplicative scale for goals. A scalar or 1D tensor,
+      must be broadcastable to number of goal dimensions.
+    reward_scales: multiplicative scale for rewards. A scalar or 1D tensor,
+      must be broadcastable to number of reward dimensions.
+    weight_index: (integer) The context list index that specifies weight.
+    weight_vector: (a number or a list or Numpy array) The weighting vector,
+      broadcastable to `next_states`.
+    summarize: (boolean) enable summary ops.
+    termination_epsilon: terminate if dist is less than this quantity.
+    state_indices: (a list of integers) list of state indices to select.
+    goal_indices: (a list of integers) list of goal indices to select.
+    vectorize: Return a vectorized form.
+    norm: L1 or L2.
+    epsilon: small offset to ensure non-negative/zero distance.
+
+  Returns:
+    A new tf.float32 [batch_size] rewards Tensor, and
+      tf.float32 [batch_size] discounts tensor.
+  """
+  del actions, rewards  # Unused
+  stats = {}
+  record_tensor(next_states, state_indices, stats, 'next_states')
+  next_states = index_states(next_states, state_indices)
+  states = index_states(states, state_indices)
+  goals = index_states(contexts[0], goal_indices)
+  next_sq_dists = tf.squared_difference(next_states * state_scales,
+                                        goals * goal_scales)
+  sq_dists = tf.squared_difference(states * state_scales,
+                                   goals * goal_scales)
+  record_tensor(sq_dists, None, stats, 'sq_dists')
+  if weight_vector is not None:
+    next_sq_dists *= tf.convert_to_tensor(weight_vector, dtype=next_states.dtype)
+    sq_dists *= tf.convert_to_tensor(weight_vector, dtype=next_states.dtype)
+  if weight_index is not None:
+    next_sq_dists *= contexts[weight_index]
+    sq_dists *= contexts[weight_index]
+  if norm == 'L1':
+    next_dist = tf.sqrt(next_sq_dists + epsilon)
+    dist = tf.sqrt(sq_dists + epsilon)
+    next_dist = tf.reduce_sum(next_dist, -1)
+    dist = tf.reduce_sum(dist, -1)
+  elif norm == 'L2':
+    next_dist = tf.reduce_sum(next_sq_dists, -1)
+    next_dist = tf.sqrt(next_dist + epsilon)  # tf.gradients fails when tf.sqrt(-0.0)
+    dist = tf.reduce_sum(sq_dists, -1)
+    dist = tf.sqrt(dist + epsilon)  # tf.gradients fails when tf.sqrt(-0.0)
+  else:
+    raise NotImplementedError(norm)
+  discounts = next_dist > termination_epsilon
+  if summarize:
+    with tf.name_scope('RewardFn/'):
+      tf.summary.scalar('mean_dist', tf.reduce_mean(dist))
+      tf.summary.histogram('dist', dist)
+      summarize_stats(stats)
+  diff = dist - next_dist
+  diff *= reward_scales
+  return tf.to_float(diff), tf.to_float(discounts)
+
+
+@gin.configurable
+def binary_indicator(states,
+                     actions,
+                     rewards,
+                     next_states,
+                     contexts,
+                     termination_epsilon=1e-4,
+                     offset=0,
+                     epsilon=1e-10,
+                     state_indices=None,
+                     summarize=False):
+  """Returns 0/1 by checking if next_states and contexts overlap.
+
+  Args:
+    states: A [batch_size, num_state_dims] Tensor representing a batch
+        of states.
+    actions: A [batch_size, num_action_dims] Tensor representing a batch
+      of actions.
+    rewards: A [batch_size] Tensor representing a batch of rewards.
+    next_states: A [batch_size, num_state_dims] Tensor representing a batch
+      of next states.
+    contexts: A list of [batch_size, num_context_dims] Tensor representing
+      a batch of contexts.
+    termination_epsilon: terminate if dist is less than this quantity.
+    offset: Offset the rewards.
+    epsilon: small offset to ensure non-negative/zero distance.
+
+  Returns:
+    A new tf.float32 [batch_size] rewards Tensor, and
+      tf.float32 [batch_size] discounts tensor.
+  """
+  del states, actions  # unused args
+  next_states = index_states(next_states, state_indices)
+  dist = tf.reduce_sum(tf.squared_difference(next_states, contexts[0]), -1)
+  dist = tf.sqrt(dist + epsilon)
+  discounts = dist > termination_epsilon
+  rewards = tf.logical_not(discounts)
+  rewards = tf.to_float(rewards) + offset
+  return tf.to_float(rewards), tf.ones_like(tf.to_float(discounts)) #tf.to_float(discounts)
+
+
+@gin.configurable
+def plain_rewards(states, actions, rewards, next_states, contexts):
+  """Returns the given rewards.
+
+  Args:
+    states: A [batch_size, num_state_dims] Tensor representing a batch
+        of states.
+    actions: A [batch_size, num_action_dims] Tensor representing a batch
+      of actions.
+    rewards: A [batch_size] Tensor representing a batch of rewards.
+    next_states: A [batch_size, num_state_dims] Tensor representing a batch
+      of next states.
+    contexts: A list of [batch_size, num_context_dims] Tensor representing
+      a batch of contexts.
+
+  Returns:
+    A new tf.float32 [batch_size] rewards Tensor, and
+      tf.float32 [batch_size] discounts tensor.
+  """
+  del states, actions, next_states, contexts  # Unused
+  return rewards, tf.ones_like(rewards)
+
+
+@gin.configurable
+def ctrl_rewards(states,
+                 actions,
+                 rewards,
+                 next_states,
+                 contexts,
+                 reward_scales=1.0):
+  """Returns the negative control cost.
+
+  Args:
+    states: A [batch_size, num_state_dims] Tensor representing a batch
+        of states.
+    actions: A [batch_size, num_action_dims] Tensor representing a batch
+      of actions.
+    rewards: A [batch_size] Tensor representing a batch of rewards.
+    next_states: A [batch_size, num_state_dims] Tensor representing a batch
+      of next states.
+    contexts: A list of [batch_size, num_context_dims] Tensor representing
+      a batch of contexts.
+    reward_scales: multiplicative scale for rewards. A scalar or 1D tensor,
+      must be broadcastable to number of reward dimensions.
+
+  Returns:
+    A new tf.float32 [batch_size] rewards Tensor, and
+      tf.float32 [batch_size] discounts tensor.
+  """
+  del states, rewards, contexts  # Unused
+  if actions is None:
+    rewards = tf.to_float(tf.zeros(shape=next_states.shape[:1]))
+  else:
+    rewards = -tf.reduce_sum(tf.square(actions), axis=1)
+    rewards *= reward_scales
+    rewards = tf.to_float(rewards)
+  return rewards, tf.ones_like(rewards)
+
+
+@gin.configurable
+def diff_rewards(
+    states,
+    actions,
+    rewards,
+    next_states,
+    contexts,
+    state_indices=None,
+    goal_index=0,):
+  """Returns (next_states - goals) as a batched vector reward."""
+  del states, rewards, actions  # Unused
+  if state_indices is not None:
+    next_states = index_states(next_states, state_indices)
+  rewards = tf.to_float(next_states - contexts[goal_index])
+  return rewards, tf.ones_like(rewards)
+
+
+@gin.configurable
+def state_rewards(states,
+                  actions,
+                  rewards,
+                  next_states,
+                  contexts,
+                  weight_index=None,
+                  state_indices=None,
+                  weight_vector=1.0,
+                  offset_vector=0.0,
+                  summarize=False):
+  """Returns the rewards that are linear mapping of next_states.
+
+  Args:
+    states: A [batch_size, num_state_dims] Tensor representing a batch
+        of states.
+    actions: A [batch_size, num_action_dims] Tensor representing a batch
+      of actions.
+    rewards: A [batch_size] Tensor representing a batch of rewards.
+    next_states: A [batch_size, num_state_dims] Tensor representing a batch
+      of next states.
+    contexts: A list of [batch_size, num_context_dims] Tensor representing
+      a batch of contexts.
+    weight_index: (integer) Index of contexts lists that specify weighting.
+    state_indices: (a list of Numpy integer array) Indices of states dimensions
+      to be mapped.
+    weight_vector: (a number or a list or Numpy array) The weighting vector,
+      broadcastable to `next_states`.
+    offset_vector: (a number or a list of Numpy array) The off vector.
+    summarize: (boolean) enable summary ops.
+
+  Returns:
+    A new tf.float32 [batch_size] rewards Tensor, and
+      tf.float32 [batch_size] discounts tensor.
+  """
+  del states, actions, rewards  # unused args
+  stats = {}
+  record_tensor(next_states, state_indices, stats)
+  next_states = index_states(next_states, state_indices)
+  weight = tf.constant(
+      weight_vector, dtype=next_states.dtype, shape=next_states[0].shape)
+  weights = tf.expand_dims(weight, 0)
+  offset = tf.constant(
+      offset_vector, dtype=next_states.dtype, shape=next_states[0].shape)
+  offsets = tf.expand_dims(offset, 0)
+  if weight_index is not None:
+    weights *= contexts[weight_index]
+  rewards = tf.to_float(tf.reduce_sum(weights * (next_states+offsets), axis=1))
+  if summarize:
+    with tf.name_scope('RewardFn/'):
+      summarize_stats(stats)
+  return rewards, tf.ones_like(rewards)

research/efficient-hrl/context/samplers.py

+# Copyright 2018 The TensorFlow Authors 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.
+# ==============================================================================
+
+"""Samplers for Contexts.
+
+  Each sampler class should define __call__(batch_size).
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+import tensorflow as tf
+slim = tf.contrib.slim
+import gin.tf
+
+
+@gin.configurable
+class BaseSampler(object):
+  """Base sampler."""
+
+  def __init__(self, context_spec, context_range=None, k=2, scope='sampler'):
+    """Construct a base sampler.
+
+    Args:
+      context_spec: A context spec.
+      context_range: A tuple of (minval, max), where minval, maxval are floats
+        or Numpy arrays with the same shape as the context.
+      scope: A string denoting scope.
+    """
+    self._context_spec = context_spec
+    self._context_range = context_range
+    self._k = k
+    self._scope = scope
+
+  def __call__(self, batch_size, **kwargs):
+    raise NotImplementedError
+
+  def set_replay(self, replay=None):
+    pass
+
+  def _validate_contexts(self, contexts):
+    """Validate if contexts have right spec.
+
+    Args:
+      contexts: A [batch_size, num_contexts_dim] tensor.
+    Raises:
+      ValueError: If shape or dtype mismatches that of spec.
+    """
+    if contexts[0].shape != self._context_spec.shape:
+      raise ValueError('contexts has invalid shape %s wrt spec shape %s' %
+                       (contexts[0].shape, self._context_spec.shape))
+    if contexts.dtype != self._context_spec.dtype:
+      raise ValueError('contexts has invalid dtype %s wrt spec dtype %s' %
+                       (contexts.dtype, self._context_spec.dtype))
+
+
+@gin.configurable
+class ZeroSampler(BaseSampler):
+  """Zero sampler."""
+
+  def __call__(self, batch_size, **kwargs):
+    """Sample a batch of context.
+
+    Args:
+      batch_size: Batch size.
+    Returns:
+      Two [batch_size, num_context_dims] tensors.
+    """
+    contexts = tf.zeros(
+        dtype=self._context_spec.dtype,
+        shape=[
+            batch_size,
+        ] + self._context_spec.shape.as_list())
+    return contexts, contexts
+
+
+@gin.configurable
+class BinarySampler(BaseSampler):
+  """Binary sampler."""
+
+  def __init__(self, probs=0.5, *args, **kwargs):
+    """Constructor."""
+    super(BinarySampler, self).__init__(*args, **kwargs)
+    self._probs = probs
+
+  def __call__(self, batch_size, **kwargs):
+    """Sample a batch of context."""
+    spec = self._context_spec
+    contexts = tf.random_uniform(
+        shape=[
+            batch_size,
+        ] + spec.shape.as_list(), dtype=tf.float32)
+    contexts = tf.cast(tf.greater(contexts, self._probs), dtype=spec.dtype)
+    return contexts, contexts
+
+
+@gin.configurable
+class RandomSampler(BaseSampler):
+  """Random sampler."""
+
+  def __call__(self, batch_size, **kwargs):
+    """Sample a batch of context.
+
+    Args:
+      batch_size: Batch size.
+    Returns:
+      Two [batch_size, num_context_dims] tensors.
+    """
+    spec = self._context_spec
+    context_range = self._context_range
+    if isinstance(context_range[0], (int, float)):
+      contexts = tf.random_uniform(
+          shape=[
+              batch_size,
+          ] + spec.shape.as_list(),
+          minval=context_range[0],
+          maxval=context_range[1],
+          dtype=spec.dtype)
+    elif isinstance(context_range[0], (list, tuple, np.ndarray)):
+      assert len(spec.shape.as_list()) == 1
+      assert spec.shape.as_list()[0] == len(context_range[0])
+      assert spec.shape.as_list()[0] == len(context_range[1])
+      contexts = tf.concat(
+          [
+              tf.random_uniform(
+                  shape=[
+                      batch_size, 1,
+                  ] + spec.shape.as_list()[1:],
+                  minval=context_range[0][i],
+                  maxval=context_range[1][i],
+                  dtype=spec.dtype) for i in range(spec.shape.as_list()[0])
+          ],
+          axis=1)
+    else: raise NotImplementedError(context_range)
+    self._validate_contexts(contexts)
+    state, next_state = kwargs['state'], kwargs['next_state']
+    if state is not None and next_state is not None:
+      pass
+      #contexts = tf.concat(
+      #    [tf.random_normal(tf.shape(state[:, :self._k]), dtype=tf.float64) +
+      #     tf.random_shuffle(state[:, :self._k]),
+      #     contexts[:, self._k:]], 1)
+
+    return contexts, contexts
+
+
+@gin.configurable
+class ScheduledSampler(BaseSampler):
+  """Scheduled sampler."""
+
+  def __init__(self,
+               scope='default',
+               values=None,
+               scheduler='cycle',
+               scheduler_params=None,
+               *args, **kwargs):
+    """Construct sampler.
+
+    Args:
+      scope: Scope name.
+      values: A list of numbers or [num_context_dim] Numpy arrays
+        representing the values to cycle.
+      scheduler: scheduler type.
+      scheduler_params: scheduler parameters.
+      *args: arguments.
+      **kwargs: keyword arguments.
+    """
+    super(ScheduledSampler, self).__init__(*args, **kwargs)
+    self._scope = scope
+    self._values = values
+    self._scheduler = scheduler
+    self._scheduler_params = scheduler_params or {}
+    assert self._values is not None and len(
+        self._values), 'must provide non-empty values.'
+    self._n = len(self._values)
+    # TODO(shanegu): move variable creation outside. resolve tf.cond problem.
+    self._count = 0
+    self._i = tf.Variable(
+        tf.zeros(shape=(), dtype=tf.int32),
+        name='%s-scheduled_sampler_%d' % (self._scope, self._count))
+    self._values = tf.constant(self._values, dtype=self._context_spec.dtype)
+
+  def __call__(self, batch_size, **kwargs):
+    """Sample a batch of context.
+
+    Args:
+      batch_size: Batch size.
+    Returns:
+      Two [batch_size, num_context_dims] tensors.
+    """
+    spec = self._context_spec
+    next_op = self._next(self._i)
+    with tf.control_dependencies([next_op]):
+      value = self._values[self._i]
+      if value.get_shape().as_list():
+        values = tf.tile(
+            tf.expand_dims(value, 0), (batch_size,) + (1,) * spec.shape.ndims)
+      else:
+        values = value + tf.zeros(
+            shape=[
+                batch_size,
+            ] + spec.shape.as_list(), dtype=spec.dtype)
+    self._validate_contexts(values)
+    self._count += 1
+    return values, values
+
+  def _next(self, i):
+    """Return op that increments pointer to next value.
+
+    Args:
+      i: A tensorflow integer variable.
+    Returns:
+      Op that increments pointer.
+    """
+    if self._scheduler == 'cycle':
+      inc = ('inc' in self._scheduler_params and
+             self._scheduler_params['inc']) or 1
+      return tf.assign(i, tf.mod(i+inc, self._n))
+    else:
+      raise NotImplementedError(self._scheduler)
+
+
+@gin.configurable
+class ReplaySampler(BaseSampler):
+  """Replay sampler."""
+
+  def __init__(self,
+               prefetch_queue_capacity=2,
+               override_indices=None,
+               state_indices=None,
+               *args,
+               **kwargs):
+    """Construct sampler.
+
+    Args:
+      prefetch_queue_capacity: Capacity for prefetch queue.
+      override_indices: Override indices.
+      state_indices: Select certain indices from state dimension.
+      *args: arguments.
+      **kwargs: keyword arguments.
+    """
+    super(ReplaySampler, self).__init__(*args, **kwargs)
+    self._prefetch_queue_capacity = prefetch_queue_capacity
+    self._override_indices = override_indices
+    self._state_indices = state_indices
+
+  def set_replay(self, replay):
+    """Set replay.
+
+    Args:
+      replay: A replay buffer.
+    """
+    self._replay = replay
+
+  def __call__(self, batch_size, **kwargs):
+    """Sample a batch of context.
+
+    Args:
+      batch_size: Batch size.
+    Returns:
+      Two [batch_size, num_context_dims] tensors.
+    """
+    batch = self._replay.GetRandomBatch(batch_size)
+    next_states = batch[4]
+    if self._prefetch_queue_capacity > 0:
+      batch_queue = slim.prefetch_queue.prefetch_queue(
+          [next_states],
+          capacity=self._prefetch_queue_capacity,
+          name='%s/batch_context_queue' % self._scope)
+      next_states = batch_queue.dequeue()
+    if self._override_indices is not None:
+      assert self._context_range is not None and isinstance(
+          self._context_range[0], (int, long, float))
+      next_states = tf.concat(
+          [
+              tf.random_uniform(
+                  shape=next_states[:, :1].shape,
+                  minval=self._context_range[0],
+                  maxval=self._context_range[1],
+                  dtype=next_states.dtype)
+              if i in self._override_indices else next_states[:, i:i + 1]
+              for i in range(self._context_spec.shape.as_list()[0])
+          ],
+          axis=1)
+    if self._state_indices is not None:
+      next_states = tf.concat(
+          [
+              next_states[:, i:i + 1]
+              for i in range(self._context_spec.shape.as_list()[0])
+          ],
+          axis=1)
+    self._validate_contexts(next_states)
+    return next_states, next_states
+
+
+@gin.configurable
+class TimeSampler(BaseSampler):
+  """Time Sampler."""
+
+  def __init__(self, minval=0, maxval=1, timestep=-1, *args, **kwargs):
+    """Construct sampler.
+
+    Args:
+      minval: Min value integer.
+      maxval: Max value integer.
+      timestep: Time step between states and next_states.
+      *args: arguments.
+      **kwargs: keyword arguments.
+    """
+    super(TimeSampler, self).__init__(*args, **kwargs)
+    assert self._context_spec.shape.as_list() == [1]
+    self._minval = minval
+    self._maxval = maxval
+    self._timestep = timestep
+
+  def __call__(self, batch_size, **kwargs):
+    """Sample a batch of context.
+
+    Args:
+      batch_size: Batch size.
+    Returns:
+      Two [batch_size, num_context_dims] tensors.
+    """
+    if self._maxval == self._minval:
+      contexts = tf.constant(
+          self._maxval, shape=[batch_size, 1], dtype=tf.int32)
+    else:
+      contexts = tf.random_uniform(
+          shape=[batch_size, 1],
+          dtype=tf.int32,
+          maxval=self._maxval,
+          minval=self._minval)
+    next_contexts = tf.maximum(contexts + self._timestep, 0)
+
+    return tf.cast(
+        contexts, dtype=self._context_spec.dtype), tf.cast(
+            next_contexts, dtype=self._context_spec.dtype)
+
+
+@gin.configurable
+class ConstantSampler(BaseSampler):
+  """Constant sampler."""
+
+  def __init__(self, value=None, *args, **kwargs):
+    """Construct sampler.
+
+    Args:
+      value: A list or Numpy array for values of the constant.
+      *args: arguments.
+      **kwargs: keyword arguments.
+    """
+    super(ConstantSampler, self).__init__(*args, **kwargs)
+    self._value = value
+
+  def __call__(self, batch_size, **kwargs):
+    """Sample a batch of context.
+
+    Args:
+      batch_size: Batch size.
+    Returns:
+      Two [batch_size, num_context_dims] tensors.
+    """
+    spec = self._context_spec
+    value_ = tf.constant(self._value, shape=spec.shape, dtype=spec.dtype)
+    values = tf.tile(
+        tf.expand_dims(value_, 0), (batch_size,) + (1,) * spec.shape.ndims)
+    self._validate_contexts(values)
+    return values, values
+
+
+@gin.configurable
+class DirectionSampler(RandomSampler):
+  """Direction sampler."""
+
+  def __call__(self, batch_size, **kwargs):
+    """Sample a batch of context.
+
+    Args:
+      batch_size: Batch size.
+    Returns:
+      Two [batch_size, num_context_dims] tensors.
+    """
+    spec = self._context_spec
+    context_range = self._context_range
+    if isinstance(context_range[0], (int, float)):
+      contexts = tf.random_uniform(
+          shape=[
+              batch_size,
+          ] + spec.shape.as_list(),
+          minval=context_range[0],
+          maxval=context_range[1],
+          dtype=spec.dtype)
+    elif isinstance(context_range[0], (list, tuple, np.ndarray)):
+      assert len(spec.shape.as_list()) == 1
+      assert spec.shape.as_list()[0] == len(context_range[0])
+      assert spec.shape.as_list()[0] == len(context_range[1])
+      contexts = tf.concat(
+          [
+              tf.random_uniform(
+                  shape=[
+                      batch_size, 1,
+                  ] + spec.shape.as_list()[1:],
+                  minval=context_range[0][i],
+                  maxval=context_range[1][i],
+                  dtype=spec.dtype) for i in range(spec.shape.as_list()[0])
+          ],
+          axis=1)
+    else: raise NotImplementedError(context_range)
+    self._validate_contexts(contexts)
+    if 'sampler_fn' in kwargs:
+      other_contexts = kwargs['sampler_fn']()
+    else:
+      other_contexts = contexts
+    state, next_state = kwargs['state'], kwargs['next_state']
+    if state is not None and next_state is not None:
+      my_context_range = (np.array(context_range[1]) - np.array(context_range[0])) / 2 * np.ones(spec.shape.as_list())
+      contexts = tf.concat(
+          [0.1 * my_context_range[:self._k] *
+           tf.random_normal(tf.shape(state[:, :self._k]), dtype=state.dtype) +
+           tf.random_shuffle(state[:, :self._k]) - state[:, :self._k],
+           other_contexts[:, self._k:]], 1)
+      #contexts = tf.Print(contexts,
+      #                    [contexts, tf.reduce_max(contexts, 0),
+      #                     tf.reduce_min(state, 0), tf.reduce_max(state, 0)], 'contexts', summarize=15)
+      next_contexts = tf.concat( #LALA
+          [state[:, :self._k] + contexts[:, :self._k] - next_state[:, :self._k],
+           other_contexts[:, self._k:]], 1)
+      next_contexts = contexts  #LALA cosine
+    else:
+      next_contexts = contexts
+    return tf.stop_gradient(contexts), tf.stop_gradient(next_contexts)

research/efficient-hrl/environments/__init__.py

-# Copyright 2018 The TensorFlow Authors 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.
-# ==============================================================================
 
-"""Random policy on an environment."""
-
-import tensorflow as tf
-import numpy as np
-import random
-
-import create_maze_env
-
-app = tf.app
-flags = tf.flags
-logging = tf.logging
-
-FLAGS = flags.FLAGS
-
-flags.DEFINE_string('env', 'AntMaze', 'environment name: AntMaze, AntPush, or AntFall')
-flags.DEFINE_integer('episode_length', 500, 'episode length')
-flags.DEFINE_integer('num_episodes', 50, 'number of episodes')
-
-
-def get_goal_sample_fn(env_name):
-  if env_name == 'AntMaze':
-    # NOTE: When evaluating (i.e. the metrics shown in the paper,
-    # we use the commented out goal sampling function.  The uncommented
-    # one is only used for training.
-    #return lambda: np.array([0., 16.])
-    return lambda: np.random.uniform((-4, -4), (20, 20))
-  elif env_name == 'AntPush':
-    return lambda: np.array([0., 19.])
-  elif env_name == 'AntFall':
-    return lambda: np.array([0., 27., 4.5])
-  else:
-    assert False, 'Unknown env'
-
-
-def get_reward_fn(env_name):
-  if env_name == 'AntMaze':
-    return lambda obs, goal: -np.sum(np.square(obs[:2] - goal)) ** 0.5
-  elif env_name == 'AntPush':
-    return lambda obs, goal: -np.sum(np.square(obs[:2] - goal)) ** 0.5
-  elif env_name == 'AntFall':
-    return lambda obs, goal: -np.sum(np.square(obs[:3] - goal)) ** 0.5
-  else:
-    assert False, 'Unknown env'
-
-
-def success_fn(last_reward):
-  return last_reward > -5.0
-
-
-class EnvWithGoal(object):
-
-  def __init__(self, base_env, env_name):
-    self.base_env = base_env
-    self.goal_sample_fn = get_goal_sample_fn(env_name)
-    self.reward_fn = get_reward_fn(env_name)
-    self.goal = None
-
-  def reset(self):
-    obs = self.base_env.reset()
-    self.goal = self.goal_sample_fn()
-    return np.concatenate([obs, self.goal])
-
-  def step(self, a):
-    obs, _, done, info = self.base_env.step(a)
-    reward = self.reward_fn(obs, self.goal)
-    return np.concatenate([obs, self.goal]), reward, done, info
-
-  @property
-  def action_space(self):
-    return self.base_env.action_space
-
-
-def run_environment(env_name, episode_length, num_episodes):
-  env = EnvWithGoal(
-      create_maze_env.create_maze_env(env_name),
-      env_name)
-
-  def action_fn(obs):
-    action_space = env.action_space
-    action_space_mean = (action_space.low + action_space.high) / 2.0
-    action_space_magn = (action_space.high - action_space.low) / 2.0
-    random_action = (action_space_mean +
-                     action_space_magn *
-                     np.random.uniform(low=-1.0, high=1.0,
-                                       size=action_space.shape))
-    return random_action
-
-  rewards = []
-  successes = []
-  for ep in range(num_episodes):
-    rewards.append(0.0)
-    successes.append(False)
-    obs = env.reset()
-    for _ in range(episode_length):
-      obs, reward, done, _ = env.step(action_fn(obs))
-      rewards[-1] += reward
-      successes[-1] = success_fn(reward)
-      if done:
-        break
-    logging.info('Episode %d reward: %.2f, Success: %d', ep + 1, rewards[-1], successes[-1])
-
-  logging.info('Average Reward over %d episodes: %.2f',
-               num_episodes, np.mean(rewards))
-  logging.info('Average Success over %d episodes: %.2f',
-               num_episodes, np.mean(successes))
-
-
-def main(unused_argv):
-  logging.set_verbosity(logging.INFO)
-  run_environment(FLAGS.env, FLAGS.episode_length, FLAGS.num_episodes)
-
-
-if __name__ == '__main__':
-  app.run()

research/efficient-hrl/environments/ant.py

@@ -21,8 +21,21 @@ from gym import utils
 from gym.envs.mujoco import mujoco_env
 
 
+def q_inv(a):
+  return [a[0], -a[1], -a[2], -a[3]]
+
+
+def q_mult(a, b): # multiply two quaternion
+  w = a[0] * b[0] - a[1] * b[1] - a[2] * b[2] - a[3] * b[3]
+  i = a[0] * b[1] + a[1] * b[0] + a[2] * b[3] - a[3] * b[2]
+  j = a[0] * b[2] - a[1] * b[3] + a[2] * b[0] + a[3] * b[1]
+  k = a[0] * b[3] + a[1] * b[2] - a[2] * b[1] + a[3] * b[0]
+  return [w, i, j, k]
+
+
 class AntEnv(mujoco_env.MujocoEnv, utils.EzPickle):
   FILE = "ant.xml"
+  ORI_IND = 3
 
   def __init__(self, file_path=None, expose_all_qpos=True,
                expose_body_coms=None, expose_body_comvels=None):
@@ -101,3 +114,21 @@ class AntEnv(mujoco_env.MujocoEnv, utils.EzPickle):
 
   def viewer_setup(self):
     self.viewer.cam.distance = self.model.stat.extent * 0.5
+
+  def get_ori(self):
+    ori = [0, 1, 0, 0]
+    rot = self.model.data.qpos[self.__class__.ORI_IND:self.__class__.ORI_IND + 4]  # take the quaternion
+    ori = q_mult(q_mult(rot, ori), q_inv(rot))[1:3]  # project onto x-y plane
+    ori = math.atan2(ori[1], ori[0])
+    return ori
+
+  def set_xy(self, xy):
+    qpos = np.copy(self.physics.data.qpos)
+    qpos[0] = xy[0]
+    qpos[1] = xy[1]
+
+    qvel = self.physics.data.qvel
+    self.set_state(qpos, qvel)
+
+  def get_xy(self):
+    return self.physics.data.qpos[:2]

research/efficient-hrl/environments/ant_maze_env.py

@@ -13,8 +13,8 @@
 # limitations under the License.
 # ==============================================================================
 
-from maze_env import MazeEnv
-from ant import AntEnv
+from environments.maze_env import MazeEnv
+from environments.ant import AntEnv
 
 
 class AntMazeEnv(MazeEnv):

research/efficient-hrl/environments/create_maze_env.py

@@ -13,18 +13,85 @@
 # limitations under the License.
 # ==============================================================================
 
-from ant_maze_env import AntMazeEnv
+from environments.ant_maze_env import AntMazeEnv
+from environments.point_maze_env import PointMazeEnv
 
+import tensorflow as tf
+import gin.tf
+from tf_agents.environments import gym_wrapper
+from tf_agents.environments import tf_py_environment
+
+
+@gin.configurable
+def create_maze_env(env_name=None, top_down_view=False):
+  n_bins = 0
+  manual_collision = False
+  if env_name.startswith('Ego'):
+    n_bins = 8
+    env_name = env_name[3:]
+  if env_name.startswith('Ant'):
+    cls = AntMazeEnv
+    env_name = env_name[3:]
+    maze_size_scaling = 8
+  elif env_name.startswith('Point'):
+    cls = PointMazeEnv
+    manual_collision = True
+    env_name = env_name[5:]
+    maze_size_scaling = 4
+  else:
+    assert False, 'unknown env %s' % env_name
 
-def create_maze_env(env_name=None):
   maze_id = None
-  if env_name.startswith('AntMaze'):
+  observe_blocks = False
+  put_spin_near_agent = False
+  if env_name == 'Maze':
     maze_id = 'Maze'
-  elif env_name.startswith('AntPush'):
+  elif env_name == 'Push':
     maze_id = 'Push'
-  elif env_name.startswith('AntFall'):
+  elif env_name == 'Fall':
     maze_id = 'Fall'
+  elif env_name == 'Block':
+    maze_id = 'Block'
+    put_spin_near_agent = True
+    observe_blocks = True
+  elif env_name == 'BlockMaze':
+    maze_id = 'BlockMaze'
+    put_spin_near_agent = True
+    observe_blocks = True
   else:
     raise ValueError('Unknown maze environment %s' % env_name)
 
-  return AntMazeEnv(maze_id=maze_id)
+  gym_mujoco_kwargs = {
+      'maze_id': maze_id,
+      'n_bins': n_bins,
+      'observe_blocks': observe_blocks,
+      'put_spin_near_agent': put_spin_near_agent,
+      'top_down_view': top_down_view,
+      'manual_collision': manual_collision,
+      'maze_size_scaling': maze_size_scaling
+  }
+  gym_env = cls(**gym_mujoco_kwargs)
+  gym_env.reset()
+  wrapped_env = gym_wrapper.GymWrapper(gym_env)
+  return wrapped_env
+
+
+class TFPyEnvironment(tf_py_environment.TFPyEnvironment):
+
+  def __init__(self, *args, **kwargs):
+    super(TFPyEnvironment, self).__init__(*args, **kwargs)
+
+  def start_collect(self):
+    pass
+
+  def current_obs(self):
+    time_step = self.current_time_step()
+    return time_step.observation[0]  # For some reason, there is an extra dim.
+
+  def step(self, actions):
+    actions = tf.expand_dims(actions, 0)
+    next_step = super(TFPyEnvironment, self).step(actions)
+    return next_step.is_last()[0], next_step.reward[0], next_step.discount[0]
+
+  def reset(self):
+    return super(TFPyEnvironment, self).reset()