research/efficient-hrl/agent.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.
# ==============================================================================
"""A UVF agent.
"""
import tensorflow as tf
import gin.tf
from agents import ddpg_agent
# pylint: disable=unused-import
import cond_fn
from utils import utils as uvf_utils
from context import gin_imports
# pylint: enable=unused-import
slim = tf.contrib.slim
@gin.configurable
class UvfAgentCore(object):
"""Defines basic functions for UVF agent. Must be inherited with an RL agent.
Used as lower-level agent.
"""
def __init__(self,
observation_spec,
action_spec,
tf_env,
tf_context,
step_cond_fn=cond_fn.env_transition,
reset_episode_cond_fn=cond_fn.env_restart,
reset_env_cond_fn=cond_fn.false_fn,
metrics=None,
**base_agent_kwargs):
"""Constructs a UVF agent.
Args:
observation_spec: A TensorSpec defining the observations.
action_spec: A BoundedTensorSpec defining the actions.
tf_env: A Tensorflow environment object.
tf_context: A Context class.
step_cond_fn: A function indicating whether to increment the num of steps.
reset_episode_cond_fn: A function indicating whether to restart the
episode, resampling the context.
reset_env_cond_fn: A function indicating whether to perform a manual reset
of the environment.
metrics: A list of functions that evaluate metrics of the agent.
**base_agent_kwargs: A dictionary of parameters for base RL Agent.
Raises:
ValueError: If 'dqda_clipping' is < 0.
"""
self._step_cond_fn = step_cond_fn
self._reset_episode_cond_fn = reset_episode_cond_fn
self._reset_env_cond_fn = reset_env_cond_fn
self.metrics = metrics
# expose tf_context methods
self.tf_context = tf_context(tf_env=tf_env)
self.set_replay = self.tf_context.set_replay
self.sample_contexts = self.tf_context.sample_contexts
self.compute_rewards = self.tf_context.compute_rewards
self.gamma_index = self.tf_context.gamma_index
self.context_specs = self.tf_context.context_specs
self.context_as_action_specs = self.tf_context.context_as_action_specs
self.init_context_vars = self.tf_context.create_vars
self.env_observation_spec = observation_spec[0]
merged_observation_spec = (uvf_utils.merge_specs(
(self.env_observation_spec,) + self.context_specs),)
self._context_vars = dict()
self._action_vars = dict()
self.BASE_AGENT_CLASS.__init__(
self,
observation_spec=merged_observation_spec,
action_spec=action_spec,
**base_agent_kwargs
)
def set_meta_agent(self, agent=None):
self._meta_agent = agent
@property
def meta_agent(self):
return self._meta_agent
def actor_loss(self, states, actions, rewards, discounts,
next_states):
"""Returns the next action for the state.
Args:
state: A [num_state_dims] tensor representing a state.
context: A list of [num_context_dims] tensor representing a context.
Returns:
A [num_action_dims] tensor representing the action.
"""
return self.BASE_AGENT_CLASS.actor_loss(self, states)
def action(self, state, context=None):
"""Returns the next action for the state.
Args:
state: A [num_state_dims] tensor representing a state.
context: A list of [num_context_dims] tensor representing a context.
Returns:
A [num_action_dims] tensor representing the action.
"""
merged_state = self.merged_state(state, context)
return self.BASE_AGENT_CLASS.action(self, merged_state)
def actions(self, state, context=None):
"""Returns the next action for the state.
Args:
state: A [-1, num_state_dims] tensor representing a state.
context: A list of [-1, num_context_dims] tensor representing a context.
Returns:
A [-1, num_action_dims] tensor representing the action.
"""
merged_states = self.merged_states(state, context)
return self.BASE_AGENT_CLASS.actor_net(self, merged_states)
def log_probs(self, states, actions, state_reprs, contexts=None):
assert contexts is not None
batch_dims = [tf.shape(states)[0], tf.shape(states)[1]]
contexts = self.tf_context.context_multi_transition_fn(
contexts, states=tf.to_float(state_reprs))
flat_states = tf.reshape(states,
[batch_dims[0] * batch_dims[1], states.shape[-1]])
flat_contexts = [tf.reshape(tf.cast(context, states.dtype),
[batch_dims[0] * batch_dims[1], context.shape[-1]])
for context in contexts]
flat_pred_actions = self.actions(flat_states, flat_contexts)
pred_actions = tf.reshape(flat_pred_actions,
batch_dims + [flat_pred_actions.shape[-1]])
error = tf.square(actions - pred_actions)
spec_range = (self._action_spec.maximum - self._action_spec.minimum) / 2
normalized_error = tf.cast(error, tf.float64) / tf.constant(spec_range) ** 2
return -normalized_error
@gin.configurable('uvf_add_noise_fn')
def add_noise_fn(self, action_fn, stddev=1.0, debug=False,
clip=True, global_step=None):
"""Returns the action_fn with additive Gaussian noise.
Args:
action_fn: A callable(`state`, `context`) which returns a
[num_action_dims] tensor representing a action.
stddev: stddev for the Ornstein-Uhlenbeck noise.
debug: Print debug messages.
Returns:
A [num_action_dims] action tensor.
"""
if global_step is not None:
stddev *= tf.maximum( # Decay exploration during training.
tf.train.exponential_decay(1.0, global_step, 1e6, 0.8), 0.5)
def noisy_action_fn(state, context=None):
"""Noisy action fn."""
action = action_fn(state, context)
if debug:
action = uvf_utils.tf_print(
action, [action],
message='[add_noise_fn] pre-noise action',
first_n=100)
noise_dist = tf.distributions.Normal(tf.zeros_like(action),
tf.ones_like(action) * stddev)
noise = noise_dist.sample()
action += noise
if debug:
action = uvf_utils.tf_print(
action, [action],
message='[add_noise_fn] post-noise action',
first_n=100)
if clip:
action = uvf_utils.clip_to_spec(action, self._action_spec)
return action
return noisy_action_fn
def merged_state(self, state, context=None):
"""Returns the merged state from the environment state and contexts.
Args:
state: A [num_state_dims] tensor representing a state.
context: A list of [num_context_dims] tensor representing a context.
If None, use the internal context.
Returns:
A [num_merged_state_dims] tensor representing the merged state.
"""
if context is None:
context = list(self.context_vars)
state = tf.concat([state,] + context, axis=-1)
self._validate_states(self._batch_state(state))
return state
def merged_states(self, states, contexts=None):
"""Returns the batch merged state from the batch env state and contexts.
Args:
states: A [batch_size, num_state_dims] tensor representing a batch
of states.
contexts: A list of [batch_size, num_context_dims] tensor
representing a batch of contexts. If None,
use the internal context.
Returns:
A [batch_size, num_merged_state_dims] tensor representing the batch
of merged states.
"""
if contexts is None:
contexts = [tf.tile(tf.expand_dims(context, axis=0),
(tf.shape(states)[0], 1)) for
context in self.context_vars]
states = tf.concat([states,] + contexts, axis=-1)
self._validate_states(states)
return states
def unmerged_states(self, merged_states):
"""Returns the batch state and contexts from the batch merged state.
Args:
merged_states: A [batch_size, num_merged_state_dims] tensor
representing a batch of merged states.
Returns:
A [batch_size, num_state_dims] tensor and a list of
[batch_size, num_context_dims] tensors representing the batch state
and contexts respectively.
"""
self._validate_states(merged_states)
num_state_dims = self.env_observation_spec.shape.as_list()[0]
num_context_dims_list = [c.shape.as_list()[0] for c in self.context_specs]
states = merged_states[:, :num_state_dims]
contexts = []
i = num_state_dims
for num_context_dims in num_context_dims_list:
contexts.append(merged_states[:, i: i+num_context_dims])
i += num_context_dims
return states, contexts
def sample_random_actions(self, batch_size=1):
"""Return random actions.
Args:
batch_size: Batch size.
Returns:
A [batch_size, num_action_dims] tensor representing the batch of actions.
"""
actions = tf.concat(
[
tf.random_uniform(
shape=(batch_size, 1),
minval=self._action_spec.minimum[i],
maxval=self._action_spec.maximum[i])
for i in range(self._action_spec.shape[0].value)
],
axis=1)
return actions
def clip_actions(self, actions):
"""Clip actions to spec.
Args:
actions: A [batch_size, num_action_dims] tensor representing
the batch of actions.
Returns:
A [batch_size, num_action_dims] tensor representing the batch
of clipped actions.
"""
actions = tf.concat(
[
tf.clip_by_value(
actions[:, i:i+1],
self._action_spec.minimum[i],
self._action_spec.maximum[i])
for i in range(self._action_spec.shape[0].value)
],
axis=1)
return actions
def mix_contexts(self, contexts, insert_contexts, indices):
"""Mix two contexts based on indices.
Args:
contexts: A list of [batch_size, num_context_dims] tensor representing
the batch of contexts.
insert_contexts: A list of [batch_size, num_context_dims] tensor
representing the batch of contexts to be inserted.
indices: A list of a list of integers denoting indices to replace.
Returns:
A list of resulting contexts.
"""
if indices is None: indices = [[]] * len(contexts)
assert len(contexts) == len(indices)
assert all([spec.shape.ndims == 1 for spec in self.context_specs])
mix_contexts = []
for contexts_, insert_contexts_, indices_, spec in zip(
contexts, insert_contexts, indices, self.context_specs):
mix_contexts.append(
tf.concat(
[
insert_contexts_[:, i:i + 1] if i in indices_ else
contexts_[:, i:i + 1] for i in range(spec.shape.as_list()[0])
],
axis=1))
return mix_contexts
def begin_episode_ops(self, mode, action_fn=None, state=None):
"""Returns ops that reset agent at beginning of episodes.
Args:
mode: a string representing the mode=[train, explore, eval].
Returns:
A list of ops.
"""
all_ops = []
for _, action_var in sorted(self._action_vars.items()):
sample_action = self.sample_random_actions(1)[0]
all_ops.append(tf.assign(action_var, sample_action))
all_ops += self.tf_context.reset(mode=mode, agent=self._meta_agent,
action_fn=action_fn, state=state)
return all_ops
def cond_begin_episode_op(self, cond, input_vars, mode, meta_action_fn):
"""Returns op that resets agent at beginning of episodes.
A new episode is begun if the cond op evalues to `False`.
Args:
cond: a Boolean tensor variable.
input_vars: A list of tensor variables.
mode: a string representing the mode=[train, explore, eval].
Returns:
Conditional begin op.
"""
(state, action, reward, next_state,
state_repr, next_state_repr) = input_vars
def continue_fn():
"""Continue op fn."""
items = [state, action, reward, next_state,
state_repr, next_state_repr] + list(self.context_vars)
batch_items = [tf.expand_dims(item, 0) for item in items]
(states, actions, rewards, next_states,
state_reprs, next_state_reprs) = batch_items[:6]
context_reward = self.compute_rewards(
mode, state_reprs, actions, rewards, next_state_reprs,
batch_items[6:])[0][0]
context_reward = tf.cast(context_reward, dtype=reward.dtype)
if self.meta_agent is not None:
meta_action = tf.concat(self.context_vars, -1)
items = [state, meta_action, reward, next_state,
state_repr, next_state_repr] + list(self.meta_agent.context_vars)
batch_items = [tf.expand_dims(item, 0) for item in items]
(states, meta_actions, rewards, next_states,
state_reprs, next_state_reprs) = batch_items[:6]
meta_reward = self.meta_agent.compute_rewards(
mode, states, meta_actions, rewards,
next_states, batch_items[6:])[0][0]
meta_reward = tf.cast(meta_reward, dtype=reward.dtype)
else:
meta_reward = tf.constant(0, dtype=reward.dtype)
with tf.control_dependencies([context_reward, meta_reward]):
step_ops = self.tf_context.step(mode=mode, agent=self._meta_agent,
state=state,
next_state=next_state,
state_repr=state_repr,
next_state_repr=next_state_repr,
action_fn=meta_action_fn)
with tf.control_dependencies(step_ops):
context_reward, meta_reward = map(tf.identity, [context_reward, meta_reward])
return context_reward, meta_reward
def begin_episode_fn():
"""Begin op fn."""
begin_ops = self.begin_episode_ops(mode=mode, action_fn=meta_action_fn, state=state)
with tf.control_dependencies(begin_ops):
return tf.zeros_like(reward), tf.zeros_like(reward)
with tf.control_dependencies(input_vars):
cond_begin_episode_op = tf.cond(cond, continue_fn, begin_episode_fn)
return cond_begin_episode_op
def get_env_base_wrapper(self, env_base, **begin_kwargs):
"""Create a wrapper around env_base, with agent-specific begin/end_episode.
Args:
env_base: A python environment base.
**begin_kwargs: Keyword args for begin_episode_ops.
Returns:
An object with begin_episode() and end_episode().
"""
begin_ops = self.begin_episode_ops(**begin_kwargs)
return uvf_utils.get_contextual_env_base(env_base, begin_ops)
def init_action_vars(self, name, i=None):
"""Create and return a tensorflow Variable holding an action.
Args:
name: Name of the variables.
i: Integer id.
Returns:
A [num_action_dims] tensor.
"""
if i is not None:
name += '_%d' % i
assert name not in self._action_vars, ('Conflict! %s is already '
'initialized.') % name
self._action_vars[name] = tf.Variable(
self.sample_random_actions(1)[0], name='%s_action' % (name))
self._validate_actions(tf.expand_dims(self._action_vars[name], 0))
return self._action_vars[name]
@gin.configurable('uvf_critic_function')
def critic_function(self, critic_vals, states, critic_fn=None):
"""Computes q values based on outputs from the critic net.
Args:
critic_vals: A tf.float32 [batch_size, ...] tensor representing outputs
from the critic net.
states: A [batch_size, num_state_dims] tensor representing a batch
of states.
critic_fn: A callable that process outputs from critic_net and
outputs a [batch_size] tensor representing q values.
Returns:
A tf.float32 [batch_size] tensor representing q values.
"""
if critic_fn is not None:
env_states, contexts = self.unmerged_states(states)
critic_vals = critic_fn(critic_vals, env_states, contexts)
critic_vals.shape.assert_has_rank(1)
return critic_vals
def get_action_vars(self, key):
return self._action_vars[key]
def get_context_vars(self, key):
return self.tf_context.context_vars[key]
def step_cond_fn(self, *args):
return self._step_cond_fn(self, *args)
def reset_episode_cond_fn(self, *args):
return self._reset_episode_cond_fn(self, *args)
def reset_env_cond_fn(self, *args):
return self._reset_env_cond_fn(self, *args)
@property
def context_vars(self):
return self.tf_context.vars
@gin.configurable
class MetaAgentCore(UvfAgentCore):
"""Defines basic functions for UVF Meta-agent. Must be inherited with an RL agent.
Used as higher-level agent.
"""
def __init__(self,
observation_spec,
action_spec,
tf_env,
tf_context,
sub_context,
step_cond_fn=cond_fn.env_transition,
reset_episode_cond_fn=cond_fn.env_restart,
reset_env_cond_fn=cond_fn.false_fn,
metrics=None,
actions_reg=0.,
k=2,
**base_agent_kwargs):
"""Constructs a Meta agent.
Args:
observation_spec: A TensorSpec defining the observations.
action_spec: A BoundedTensorSpec defining the actions.
tf_env: A Tensorflow environment object.
tf_context: A Context class.
step_cond_fn: A function indicating whether to increment the num of steps.
reset_episode_cond_fn: A function indicating whether to restart the
episode, resampling the context.
reset_env_cond_fn: A function indicating whether to perform a manual reset
of the environment.
metrics: A list of functions that evaluate metrics of the agent.
**base_agent_kwargs: A dictionary of parameters for base RL Agent.
Raises:
ValueError: If 'dqda_clipping' is < 0.
"""
self._step_cond_fn = step_cond_fn
self._reset_episode_cond_fn = reset_episode_cond_fn
self._reset_env_cond_fn = reset_env_cond_fn
self.metrics = metrics
self._actions_reg = actions_reg
self._k = k
# expose tf_context methods
self.tf_context = tf_context(tf_env=tf_env)
self.sub_context = sub_context(tf_env=tf_env)
self.set_replay = self.tf_context.set_replay
self.sample_contexts = self.tf_context.sample_contexts
self.compute_rewards = self.tf_context.compute_rewards
self.gamma_index = self.tf_context.gamma_index
self.context_specs = self.tf_context.context_specs
self.context_as_action_specs = self.tf_context.context_as_action_specs
self.sub_context_as_action_specs = self.sub_context.context_as_action_specs
self.init_context_vars = self.tf_context.create_vars
self.env_observation_spec = observation_spec[0]
merged_observation_spec = (uvf_utils.merge_specs(
(self.env_observation_spec,) + self.context_specs),)
self._context_vars = dict()
self._action_vars = dict()
assert len(self.context_as_action_specs) == 1
self.BASE_AGENT_CLASS.__init__(
self,
observation_spec=merged_observation_spec,
action_spec=self.sub_context_as_action_specs,
**base_agent_kwargs
)
@gin.configurable('meta_add_noise_fn')
def add_noise_fn(self, action_fn, stddev=1.0, debug=False,
global_step=None):
noisy_action_fn = super(MetaAgentCore, self).add_noise_fn(
action_fn, stddev,
clip=True, global_step=global_step)
return noisy_action_fn
def actor_loss(self, states, actions, rewards, discounts,
next_states):
"""Returns the next action for the state.
Args:
state: A [num_state_dims] tensor representing a state.
context: A list of [num_context_dims] tensor representing a context.
Returns:
A [num_action_dims] tensor representing the action.
"""
actions = self.actor_net(states, stop_gradients=False)
regularizer = self._actions_reg * tf.reduce_mean(
tf.reduce_sum(tf.abs(actions[:, self._k:]), -1), 0)
loss = self.BASE_AGENT_CLASS.actor_loss(self, states)
return regularizer + loss
@gin.configurable
class UvfAgent(UvfAgentCore, ddpg_agent.TD3Agent):
"""A DDPG agent with UVF.
"""
BASE_AGENT_CLASS = ddpg_agent.TD3Agent
ACTION_TYPE = 'continuous'
def __init__(self, *args, **kwargs):
UvfAgentCore.__init__(self, *args, **kwargs)
@gin.configurable
class MetaAgent(MetaAgentCore, ddpg_agent.TD3Agent):
"""A DDPG meta-agent.
"""
BASE_AGENT_CLASS = ddpg_agent.TD3Agent
ACTION_TYPE = 'continuous'
def __init__(self, *args, **kwargs):
MetaAgentCore.__init__(self, *args, **kwargs)
@gin.configurable()
def state_preprocess_net(
states,
num_output_dims=2,
states_hidden_layers=(100,),
normalizer_fn=None,
activation_fn=tf.nn.relu,
zero_time=True,
images=False):
"""Creates a simple feed forward net for embedding states.
"""
with slim.arg_scope(
[slim.fully_connected],
activation_fn=activation_fn,
normalizer_fn=normalizer_fn,
weights_initializer=slim.variance_scaling_initializer(
factor=1.0/3.0, mode='FAN_IN', uniform=True)):
states_shape = tf.shape(states)
states_dtype = states.dtype
states = tf.to_float(states)
if images: # Zero-out x-y
states *= tf.constant([0.] * 2 + [1.] * (states.shape[-1] - 2), dtype=states.dtype)
if zero_time:
states *= tf.constant([1.] * (states.shape[-1] - 1) + [0.], dtype=states.dtype)
orig_states = states
embed = states
if states_hidden_layers:
embed = slim.stack(embed, slim.fully_connected, states_hidden_layers,
scope='states')
with slim.arg_scope([slim.fully_connected],
weights_regularizer=None,
weights_initializer=tf.random_uniform_initializer(
minval=-0.003, maxval=0.003)):
embed = slim.fully_connected(embed, num_output_dims,
activation_fn=None,
normalizer_fn=None,
scope='value')
output = embed
output = tf.cast(output, states_dtype)
return output
@gin.configurable()
def action_embed_net(
actions,
states=None,
num_output_dims=2,
hidden_layers=(400, 300),
normalizer_fn=None,
activation_fn=tf.nn.relu,
zero_time=True,
images=False):
"""Creates a simple feed forward net for embedding actions.
"""
with slim.arg_scope(
[slim.fully_connected],
activation_fn=activation_fn,
normalizer_fn=normalizer_fn,
weights_initializer=slim.variance_scaling_initializer(
factor=1.0/3.0, mode='FAN_IN', uniform=True)):
actions = tf.to_float(actions)
if states is not None:
if images: # Zero-out x-y
states *= tf.constant([0.] * 2 + [1.] * (states.shape[-1] - 2), dtype=states.dtype)
if zero_time:
states *= tf.constant([1.] * (states.shape[-1] - 1) + [0.], dtype=states.dtype)
actions = tf.concat([actions, tf.to_float(states)], -1)
embed = actions
if hidden_layers:
embed = slim.stack(embed, slim.fully_connected, hidden_layers,
scope='hidden')
with slim.arg_scope([slim.fully_connected],
weights_regularizer=None,
weights_initializer=tf.random_uniform_initializer(
minval=-0.003, maxval=0.003)):
embed = slim.fully_connected(embed, num_output_dims,
activation_fn=None,
normalizer_fn=None,
scope='value')
if num_output_dims == 1:
return embed[:, 0, ...]
else:
return embed
def huber(x, kappa=0.1):
return (0.5 * tf.square(x) * tf.to_float(tf.abs(x) <= kappa) +
kappa * (tf.abs(x) - 0.5 * kappa) * tf.to_float(tf.abs(x) > kappa)
) / kappa
@gin.configurable()
class StatePreprocess(object):
STATE_PREPROCESS_NET_SCOPE = 'state_process_net'
ACTION_EMBED_NET_SCOPE = 'action_embed_net'
def __init__(self, trainable=False,
state_preprocess_net=lambda states: states,
action_embed_net=lambda actions, *args, **kwargs: actions,
ndims=None):
self.trainable = trainable
self._scope = tf.get_variable_scope().name
self._ndims = ndims
self._state_preprocess_net = tf.make_template(
self.STATE_PREPROCESS_NET_SCOPE, state_preprocess_net,
create_scope_now_=True)
self._action_embed_net = tf.make_template(
self.ACTION_EMBED_NET_SCOPE, action_embed_net,
create_scope_now_=True)
def __call__(self, states):
batched = states.get_shape().ndims != 1
if not batched:
states = tf.expand_dims(states, 0)
embedded = self._state_preprocess_net(states)
if self._ndims is not None:
embedded = embedded[..., :self._ndims]
if not batched:
return embedded[0]
return embedded
def loss(self, states, next_states, low_actions, low_states):
batch_size = tf.shape(states)[0]
d = int(low_states.shape[1])
# Sample indices into meta-transition to train on.
probs = 0.99 ** tf.range(d, dtype=tf.float32)
probs *= tf.constant([1.0] * (d - 1) + [1.0 / (1 - 0.99)],
dtype=tf.float32)
probs /= tf.reduce_sum(probs)
index_dist = tf.distributions.Categorical(probs=probs, dtype=tf.int64)
indices = index_dist.sample(batch_size)
batch_size = tf.cast(batch_size, tf.int64)
next_indices = tf.concat(
[tf.range(batch_size, dtype=tf.int64)[:, None],
(1 + indices[:, None]) % d], -1)
new_next_states = tf.where(indices < d - 1,
tf.gather_nd(low_states, next_indices),
next_states)
next_states = new_next_states
embed1 = tf.to_float(self._state_preprocess_net(states))
embed2 = tf.to_float(self._state_preprocess_net(next_states))
action_embed = self._action_embed_net(
tf.layers.flatten(low_actions), states=states)
tau = 2.0
fn = lambda z: tau * tf.reduce_sum(huber(z), -1)
all_embed = tf.get_variable('all_embed', [1024, int(embed1.shape[-1])],
initializer=tf.zeros_initializer())
upd = all_embed.assign(tf.concat([all_embed[batch_size:], embed2], 0))
with tf.control_dependencies([upd]):
close = 1 * tf.reduce_mean(fn(embed1 + action_embed - embed2))
prior_log_probs = tf.reduce_logsumexp(
-fn((embed1 + action_embed)[:, None, :] - all_embed[None, :, :]),
axis=-1) - tf.log(tf.to_float(all_embed.shape[0]))
far = tf.reduce_mean(tf.exp(-fn((embed1 + action_embed)[1:] - embed2[:-1])
- tf.stop_gradient(prior_log_probs[1:])))
repr_log_probs = tf.stop_gradient(
-fn(embed1 + action_embed - embed2) - prior_log_probs) / tau
return close + far, repr_log_probs, indices
def get_trainable_vars(self):
return (
slim.get_trainable_variables(
uvf_utils.join_scope(self._scope, self.STATE_PREPROCESS_NET_SCOPE)) +
slim.get_trainable_variables(
uvf_utils.join_scope(self._scope, self.ACTION_EMBED_NET_SCOPE)))
@gin.configurable()
class InverseDynamics(object):
INVERSE_DYNAMICS_NET_SCOPE = 'inverse_dynamics'
def __init__(self, spec):
self._spec = spec
def sample(self, states, next_states, num_samples, orig_goals, sc=0.5):
goal_dim = orig_goals.shape[-1]
spec_range = (self._spec.maximum - self._spec.minimum) / 2 * tf.ones([goal_dim])
loc = tf.cast(next_states - states, tf.float32)[:, :goal_dim]
scale = sc * tf.tile(tf.reshape(spec_range, [1, goal_dim]),
[tf.shape(states)[0], 1])
dist = tf.distributions.Normal(loc, scale)
if num_samples == 1:
return dist.sample()
samples = tf.concat([dist.sample(num_samples - 2),
tf.expand_dims(loc, 0),
tf.expand_dims(orig_goals, 0)], 0)
return uvf_utils.clip_to_spec(samples, self._spec)