def __init__(self,
task,
policy_model=None,
policy_optimizer=None,
policy_lr_schedule=lr.MultifactorSchedule,
policy_batch_size=64,
policy_train_steps_per_epoch=500,
policy_evals_per_epoch=1,
policy_eval_steps=1,
collect_per_epoch=50,
max_slice_length=1,
output_dir=None):
"""Configures the policy trainer.
Args:
task: RLTask instance, which defines the environment to train on.
policy_model: Trax layer, representing the policy model.
functions and eval functions (a.k.a. metrics) are considered to be
outside the core model, taking core model output and data labels as
their two inputs.
policy_optimizer: the optimizer to use to train the policy model.
policy_lr_schedule: learning rate schedule to use to train the policy.
policy_batch_size: batch size used to train the policy model.
policy_train_steps_per_epoch: how long to train policy in each RL epoch.
policy_evals_per_epoch: number of policy trainer evaluations per RL epoch
- only affects metric reporting.
policy_eval_steps: number of policy trainer steps per evaluation - only
affects metric reporting.
collect_per_epoch: how many trajectories to collect per epoch.
max_slice_length: the maximum length of trajectory slices to use.
output_dir: Path telling where to save outputs (evals and checkpoints).
"""
super(PolicyTrainer,
self).__init__(task,
collect_per_epoch=collect_per_epoch,
output_dir=output_dir)
self._policy_batch_size = policy_batch_size
self._policy_train_steps_per_epoch = policy_train_steps_per_epoch
self._policy_evals_per_epoch = policy_evals_per_epoch
self._policy_eval_steps = policy_eval_steps
self._collect_per_epoch = collect_per_epoch
self._max_slice_length = max_slice_length
self._policy_dist = distributions.create_distribution(
task.action_space)
# Inputs to the policy model are produced by self._policy_batches_stream.
self._policy_inputs = supervised.Inputs(
train_stream=lambda _: self.policy_batches_stream())
policy_model = functools.partial(
policy_model,
policy_distribution=self._policy_dist,
)
# This is the policy Trainer that will be used to train the policy model.
# * inputs to the trainer come from self.policy_batches_stream
# * outputs, targets and weights are passed to self.policy_loss
self._policy_trainer = supervised.Trainer(
model=policy_model,
optimizer=policy_optimizer,
lr_schedule=policy_lr_schedule,
loss_fn=self.policy_loss,
inputs=self._policy_inputs,
output_dir=output_dir,
metrics=self.policy_metrics,
)
self._policy_collect_model = policy_model(mode='collect')
policy_batch = next(self.policy_batches_stream())
self._policy_collect_model.init(shapes.signature(policy_batch))
if self._task._initial_trajectories == 0:
self._task.remove_epoch(0)
self.task.collect_trajectories(self.policy,
self._collect_per_epoch,
self._epoch)
def __init__(self,
task,
joint_model=None,
optimizer=None,
lr_schedule=lr.multifactor,
batch_size=64,
train_steps_per_epoch=500,
supervised_evals_per_epoch=1,
supervised_eval_steps=1,
n_trajectories_per_epoch=50,
max_slice_length=1,
normalize_advantages=True,
output_dir=None,
n_replay_epochs=1):
"""Configures the joint trainer.
Args:
task: RLTask instance, which defines the environment to train on.
joint_model: Trax layer, representing the joint policy and value model.
optimizer: the optimizer to use to train the joint model.
lr_schedule: learning rate schedule to use to train the joint model/.
batch_size: batch size used to train the joint model.
train_steps_per_epoch: how long to train the joint model in each RL epoch.
supervised_evals_per_epoch: number of value trainer evaluations per RL
epoch - only affects metric reporting.
supervised_eval_steps: number of value trainer steps per evaluation -
only affects metric reporting.
n_trajectories_per_epoch: how many trajectories to collect per epoch.
max_slice_length: the maximum length of trajectory slices to use.
normalize_advantages: if True, then normalize advantages - currently
implemented only in PPO.
output_dir: Path telling where to save outputs (evals and checkpoints).
n_replay_epochs: how many last epochs to take into the replay buffer;
> 1 only makes sense for off-policy algorithms.
"""
super(ActorCriticJointTrainer, self).__init__(
task,
n_trajectories_per_epoch=n_trajectories_per_epoch,
output_dir=output_dir,
)
self._batch_size = batch_size
self._train_steps_per_epoch = train_steps_per_epoch
self._supervised_evals_per_epoch = supervised_evals_per_epoch
self._supervised_eval_steps = supervised_eval_steps
self._n_trajectories_per_epoch = n_trajectories_per_epoch
self._max_slice_length = max_slice_length
self._policy_dist = distributions.create_distribution(
task.action_space)
self._lr_schedule = lr_schedule()
self._optimizer = optimizer
self._normalize_advantages = normalize_advantages
self._n_replay_epochs = n_replay_epochs
self._task.set_n_replay_epochs(n_replay_epochs)
# Inputs to the joint model are produced by self.batches_stream.
self._inputs = supervised.Inputs(
train_stream=lambda _: self.batches_stream())
self._joint_model = functools.partial(
joint_model,
policy_distribution=self._policy_dist,
)
# This is the joint Trainer that will be used to train the policy model.
# * inputs to the trainer come from self.batches_stream
# * outputs are passed to self._joint_loss
self._trainer = supervised.Trainer(model=self._joint_model,
optimizer=self._optimizer,
lr_schedule=self._lr_schedule,
loss_fn=self.joint_loss,
inputs=self._inputs,
output_dir=output_dir,
metrics={
'joint_loss':
self.joint_loss,
'advantage_mean':
self.advantage_mean,
'advantage_norm':
self.advantage_norm,
'value_loss':
self.value_loss,
'explained_variance':
self.explained_variance,
'log_probs_mean':
self.log_probs_mean,
'preferred_move':
self.preferred_move
})
self._eval_model = tl.Accelerate(self._joint_model(mode='eval'),
n_devices=1)
example_batch = next(self.batches_stream())
self._eval_model.init(example_batch)
def __init__(self,
task,
joint_model=None,
optimizer=None,
lr_schedule=lr.MultifactorSchedule,
batch_size=64,
train_steps_per_epoch=500,
supervised_evals_per_epoch=1,
supervised_eval_steps=1,
collect_per_epoch=50,
max_slice_length=1,
output_dir=None):
"""Configures the joint trainer.
Args:
task: RLTask instance, which defines the environment to train on.
joint_model: Trax layer, representing the joint policy and value model.
optimizer: the optimizer to use to train the joint model.
lr_schedule: learning rate schedule to use to train the joint model/.
batch_size: batch size used to train the joint model.
train_steps_per_epoch: how long to train the joint model in each RL epoch.
supervised_evals_per_epoch: number of value trainer evaluations per RL
epoch - only affects metric reporting.
supervised_eval_steps: number of value trainer steps per evaluation -
only affects metric reporting.
collect_per_epoch: how many trajectories to collect per epoch.
max_slice_length: the maximum length of trajectory slices to use.
output_dir: Path telling where to save outputs (evals and checkpoints).
"""
super(ActorCriticJointTrainer,
self).__init__(task,
collect_per_epoch=collect_per_epoch,
output_dir=output_dir)
self._batch_size = batch_size
self._train_steps_per_epoch = train_steps_per_epoch
self._supervised_evals_per_epoch = supervised_evals_per_epoch
self._supervised_eval_steps = supervised_eval_steps
self._collect_per_epoch = collect_per_epoch
self._max_slice_length = max_slice_length
self._policy_dist = distributions.create_distribution(
task.action_space)
self._total_samples = 0
self._lr_schedule = lr_schedule
self._optimizer = optimizer
# Inputs to the joint model are produced by self.batches_stream.
self._inputs = supervised.Inputs(
train_stream=lambda _: self.batches_stream())
self._joint_model = functools.partial(
joint_model,
policy_distribution=self._policy_dist,
)
# This is the joint Trainer that will be used to train the policy model.
# * inputs to the trainer come from self.batches_stream
# * outputs are passed to self._joint_loss
self._trainer = supervised.Trainer(model=self._joint_model,
optimizer=self._optimizer,
lr_schedule=self._lr_schedule,
loss_fn=self.joint_loss,
inputs=self._inputs,
output_dir=output_dir,
metrics={
'joint_loss':
self.joint_loss,
'advantage_mean':
self.advantage_mean,
'value_loss':
self.value_loss,
'log_probs_mean':
self.log_probs_mean,
'preferred_move':
self.preferred_move
})
self._eval_model = self._joint_model(mode='eval')
example_batch = next(self.batches_stream())
self._eval_model.init(example_batch)
def __init__(self, task, model=None, optimizer=None,
lr_schedule=lr.MultifactorSchedule, batch_size=64,
train_steps_per_epoch=500, collect_per_epoch=50,
max_slice_length=1, output_dir=None):
"""Configures the Reinforce loop.
Args:
task: RLTask instance, which defines the environment to train on.
model: Trax layer, representing the policy model.
functions and eval functions (a.k.a. metrics) are considered to be
outside the core model, taking core model output and data labels as
their two inputs.
optimizer: the optimizer to use to train the model.
lr_schedule: learning rate schedule to use to train the model.
batch_size: batch size used to train the model.
train_steps_per_epoch: how long to train in each RL epoch.
collect_per_epoch: how many trajectories to collect per epoch.
max_slice_length: the maximum length of trajectory slices to use.
output_dir: Path telling where to save outputs (evals and checkpoints).
Can be None if both `eval_task` and `checkpoint_at` are None.
"""
super(PolicyGradientTrainer, self).__init__(
task, collect_per_epoch=collect_per_epoch, output_dir=output_dir)
self._batch_size = batch_size
self._train_steps_per_epoch = train_steps_per_epoch
self._collect_per_epoch = collect_per_epoch
self._max_slice_length = max_slice_length
self._epoch = 0
self._eval_model = model(mode='eval')
example_batch = next(self._batches_stream())
self._eval_model.init(example_batch)
self._policy_dist = distributions.create_distribution(task.env.action_space)
# Inputs to the policy model are produced by self._batches_stream.
# As you can see below, the stream returns (observation, action, return)
# from the RLTask, which the model uses as (inputs, targets, loss weights).
self._inputs = supervised.Inputs(
train_stream=lambda _: self._batches_stream())
# This is the main Trainer that will be used to train the policy using
# a policy gradient loss. Note a few of the choices here:
#
# * this is a policy trainer, so:
# inputs are states and targets are actions + loss weights (see below)
# * we are using LogLoss
# This is because we are training a policy model, so targets are
# actions and they are points sampled from the policy distribution --
# LogLoss will calculate
# the log probability of each action in the state, log pi(s, a).
# * we are using has_weights=True
# We set has_weights = True because pi(s, a) will be multiplied by
# a number -- a factor that can change depending on which policy
# gradient algorithms you use; here, we just use the return from
# from this state and action, but many other variants can be tried.
loss = functools.partial(
distributions.LogLoss, distribution=self._policy_dist
)
self._trainer = supervised.Trainer(
model=model, optimizer=optimizer, lr_schedule=lr_schedule,
loss_fn=loss, inputs=self._inputs, output_dir=output_dir,
metrics={'loss': loss}, has_weights=True)
def __init__(self, task,
value_model=None,
value_optimizer=None,
value_lr_schedule=lr.MultifactorSchedule,
value_batch_size=64,
value_train_steps_per_epoch=500,
value_evals_per_epoch=1,
value_eval_steps=1,
n_shared_layers=0,
added_policy_slice_length=0,
n_replay_epochs=1,
scale_value_targets=False,
q_value=False,
q_value_n_samples=1,
**kwargs): # Arguments of PolicyTrainer come here.
"""Configures the actor-critic Trainer.
Args:
task: RLTask instance to use
value_model: the model to use for the value function
value_optimizer: the optimizer to train the value model
value_lr_schedule: lr schedule for value model training
value_batch_size: batch size for value model training
value_train_steps_per_epoch: how many steps are we using to
train the value model in each epoch
value_evals_per_epoch: number of value trainer evaluations per RL
epoch - only affects metric reporting.
value_eval_steps: number of value trainer steps per evaluation -
only affects metric reporting.
n_shared_layers: how many layers to share between value and
policy models
added_policy_slice_length: how much longer should slices of
trajectories be for policy than for value training; this
is useful for TD calculations and only affect the length
of elements produced for policy batches; value batches
have maximum length set by max_slice_length in **kwargs
n_replay_epochs: how many last epochs to take into the replay buffer;
only makes sense for off-policy algorithms
scale_value_targets: whether to scale targets for the value function by
1 / (1 - gamma)
q_value: whether to use Q-values as baselines
q_value_n_samples: number of samples to average over when calculating
baselines based on Q-values
**kwargs: arguments for PolicyTrainer super-class
"""
self._n_shared_layers = n_shared_layers
self._value_batch_size = value_batch_size
self._value_train_steps_per_epoch = value_train_steps_per_epoch
self._value_evals_per_epoch = value_evals_per_epoch
self._value_eval_steps = value_eval_steps
# The 2 below will be initalized in super.__init__ anyway, but are needed
# to construct value batches which are needed before PolicyTrainer init
# since policy input creation calls the value model -- hence this code.
self._task = task
self._max_slice_length = kwargs.get('max_slice_length', 1)
self._added_policy_slice_length = added_policy_slice_length
self._n_replay_epochs = n_replay_epochs
task.set_n_replay_epochs(n_replay_epochs)
if scale_value_targets:
self._value_network_scale = 1 / (1 - self._task.gamma)
else:
self._value_network_scale = 1
self._q_value = q_value
self._q_value_n_samples = q_value_n_samples
if q_value:
value_model = functools.partial(value_model, inject_actions=True)
self._value_eval_model = value_model(mode='eval')
self._value_eval_model.init(self._value_model_signature)
# Initialize policy training.
super(ActorCriticTrainer, self).__init__(task, **kwargs)
# Initialize training of the value function.
value_output_dir = kwargs.get('output_dir', None)
if value_output_dir is not None:
value_output_dir = os.path.join(value_output_dir, 'value')
# If needed, create value_output_dir and missing parent directories.
if not tf.io.gfile.isdir(value_output_dir):
tf.io.gfile.makedirs(value_output_dir)
self._value_inputs = supervised.Inputs(
train_stream=lambda _: self.value_batches_stream())
self._value_trainer = supervised.Trainer(
model=value_model,
optimizer=value_optimizer,
lr_schedule=value_lr_schedule,
loss_fn=tl.L2Loss(),
inputs=self._value_inputs,
output_dir=value_output_dir,
metrics={'value_loss': tl.L2Loss()})
def __init__(self, task,
value_body=None,
value_optimizer=None,
value_lr_schedule=lr.multifactor,
value_batch_size=64,
value_train_steps_per_epoch=500,
value_evals_per_epoch=1,
value_eval_steps=1,
exploration_rate=functools.partial(
lr.multifactor,
factors='constant * decay_every',
constant=1., # pylint: disable=redefined-outer-name
decay_factor=0.99,
steps_per_decay=1,
minimum=0.1),
n_eval_episodes=0,
only_eval=False,
n_replay_epochs=1,
max_slice_length=1,
sync_freq=1000,
scale_value_targets=True,
output_dir=None,
**kwargs):
"""Configures the value trainer.
Args:
task: RLTask instance, which defines the environment to train on.
value_body: Trax layer, representing the body of the value model.
functions and eval functions (a.k.a. metrics) are considered to be
outside the core model, taking core model output and data labels as
their two inputs.
value_optimizer: the optimizer to use to train the policy model.
value_lr_schedule: learning rate schedule to use to train the policy.
value_batch_size: batch size used to train the policy model.
value_train_steps_per_epoch: how long to train policy in each RL epoch.
value_evals_per_epoch: number of policy trainer evaluations per RL epoch
- only affects metric reporting.
value_eval_steps: number of policy trainer steps per evaluation - only
affects metric reporting.
exploration_rate: exploration rate schedule - used in the policy method.
n_eval_episodes: number of episodes to play with policy at
temperature 0 in each epoch -- used for evaluation only
only_eval: If set to True, then trajectories are collected only for
for evaluation purposes, but they are not recorded.
n_replay_epochs: Number of last epochs to take into the replay buffer;
only makes sense for off-policy algorithms.
max_slice_length: the maximum length of trajectory slices to use; it is
the second dimenions of the value network output:
(batch, max_slice_length, number of actions)
Higher max_slice_length implies that the network has to predict more
values into the future.
sync_freq: frequency when to synchronize the target
network with the trained network. This is necessary for training the
network on bootstrapped targets, e.g. using n-step returns.
scale_value_targets: If `True`, scale value function targets by
`1 / (1 - gamma)`. We are trying to fix the problem with very large
returns in some games in a way which does not introduce an additional
hyperparameters.
output_dir: Path telling where to save outputs (evals and checkpoints).
**kwargs: arguments for the superclass RLTrainer.
"""
super(ValueAgent, self).__init__(
task,
n_eval_episodes=n_eval_episodes,
output_dir=output_dir,
**kwargs
)
self._value_batch_size = value_batch_size
self._value_train_steps_per_epoch = value_train_steps_per_epoch
self._value_evals_per_epoch = value_evals_per_epoch
self._value_eval_steps = value_eval_steps
self._only_eval = only_eval
self._max_slice_length = max_slice_length
self._policy_dist = distributions.create_distribution(task.action_space)
self._n_replay_epochs = n_replay_epochs
self._exploration_rate = exploration_rate()
self._sync_at = (lambda step: step % sync_freq == 0)
if scale_value_targets:
self._value_network_scale = 1 / (1 - self._task.gamma)
else:
self._value_network_scale = 1
value_model = functools.partial(
models.Quality,
body=value_body,
n_actions=self.task.action_space.n)
self._value_eval_model = value_model(mode='eval')
self._value_eval_model.init(self._value_model_signature)
self._value_eval_jit = tl.jit_forward(
self._value_eval_model.pure_fn, fastmath.device_count(), do_mean=False)
# Inputs to the value model are produced by self._values_batches_stream.
self._inputs = data.inputs.Inputs(
train_stream=lambda _: self.value_batches_stream())
# This is the value Trainer that will be used to train the value model.
# * inputs to the trainer come from self.value_batches_stream
# * outputs, targets and weights are passed to self.value_loss
self._value_trainer = supervised.Trainer(
model=value_model,
optimizer=value_optimizer,
lr_schedule=value_lr_schedule(),
loss_fn=self.value_loss,
inputs=self._inputs,
output_dir=output_dir,
metrics={'value_loss': self.value_loss,
'value_mean': self.value_mean,
'returns_mean': self.returns_mean}
)
value_batch = next(self.value_batches_stream())
self._eval_model = tl.Accelerate(
value_model(mode='collect'), n_devices=1)
self._eval_model.init(shapes.signature(value_batch))
if self._task._initial_trajectories == 0:
self._task.remove_epoch(0)
self._collect_trajectories()
def __init__(self, task, policy_model=None, policy_optimizer=None,
policy_lr_schedule=lr.multifactor, policy_batch_size=64,
policy_train_steps_per_epoch=500, policy_evals_per_epoch=1,
policy_eval_steps=1, n_eval_episodes=0,
only_eval=False, max_slice_length=1, output_dir=None, **kwargs):
"""Configures the policy trainer.
Args:
task: RLTask instance, which defines the environment to train on.
policy_model: Trax layer, representing the policy model.
functions and eval functions (a.k.a. metrics) are considered to be
outside the core model, taking core model output and data labels as
their two inputs.
policy_optimizer: the optimizer to use to train the policy model.
policy_lr_schedule: learning rate schedule to use to train the policy.
policy_batch_size: batch size used to train the policy model.
policy_train_steps_per_epoch: how long to train policy in each RL epoch.
policy_evals_per_epoch: number of policy trainer evaluations per RL epoch
- only affects metric reporting.
policy_eval_steps: number of policy trainer steps per evaluation - only
affects metric reporting.
n_eval_episodes: number of episodes to play with policy at
temperature 0 in each epoch -- used for evaluation only
only_eval: If set to True, then trajectories are collected only for
for evaluation purposes, but they are not recorded.
max_slice_length: the maximum length of trajectory slices to use.
output_dir: Path telling where to save outputs (evals and checkpoints).
**kwargs: arguments for the superclass Agent.
"""
super().__init__(
task,
n_eval_episodes=n_eval_episodes,
output_dir=output_dir,
**kwargs
)
self._policy_batch_size = policy_batch_size
self._policy_train_steps_per_epoch = policy_train_steps_per_epoch
self._policy_evals_per_epoch = policy_evals_per_epoch
self._policy_eval_steps = policy_eval_steps
self._only_eval = only_eval
self._max_slice_length = max_slice_length
self._policy_dist = distributions.create_distribution(task.action_space)
# Inputs to the policy model are produced by self._policy_batches_stream.
self._policy_inputs = data.inputs.Inputs(
train_stream=lambda _: self.policy_batches_stream())
policy_model = functools.partial(
policy_model,
policy_distribution=self._policy_dist,
)
# This is the policy Trainer that will be used to train the policy model.
# * inputs to the trainer come from self.policy_batches_stream
# * outputs, targets and weights are passed to self.policy_loss
self._policy_trainer = supervised.Trainer(
model=policy_model,
optimizer=policy_optimizer,
lr_schedule=policy_lr_schedule(),
loss_fn=self.policy_loss,
inputs=self._policy_inputs,
output_dir=output_dir,
metrics=self.policy_metrics,
)
self._policy_collect_model = tl.Accelerate(
policy_model(mode='collect'), n_devices=1)
policy_batch = next(self.policy_batches_stream())
self._policy_collect_model.init(shapes.signature(policy_batch))
self._policy_eval_model = tl.Accelerate(
policy_model(mode='eval'), n_devices=1) # Not collecting stats
self._policy_eval_model.init(shapes.signature(policy_batch))
if self._task._initial_trajectories == 0:
self._task.remove_epoch(0)
self._collect_trajectories()
def __init__(self, task,
value_model=None,
value_optimizer=None,
value_lr_schedule=lr.MultifactorSchedule,
value_batch_size=64,
value_train_steps_per_epoch=500,
value_evals_per_epoch=1,
value_eval_steps=1,
n_shared_layers=0,
added_policy_slice_length=0,
n_replay_epochs=1,
scale_value_targets=False,
q_value=False,
q_value_aggregate_max=True,
q_value_n_samples=1,
**kwargs): # Arguments of PolicyTrainer come here.
"""Configures the actor-critic trainer.
Args:
task: `RLTask` instance to use.
value_model: Model to use for the value function.
value_optimizer: Optimizer to train the value model.
value_lr_schedule: lr schedule for value model training.
value_batch_size: Batch size for value model training.
value_train_steps_per_epoch: Number of steps are we using to train the
value model in each epoch.
value_evals_per_epoch: Number of value trainer evaluations per RL epoch;
only affects metric reporting.
value_eval_steps: Number of value trainer steps per evaluation; only
affects metric reporting.
n_shared_layers: Number of layers to share between value and policy
models.
added_policy_slice_length: How much longer should slices of
trajectories be for policy than for value training; this
is useful for TD calculations and only affect the length
of elements produced for policy batches; value batches
have maximum length set by `max_slice_length` in `**kwargs`.
n_replay_epochs: Number of last epochs to take into the replay buffer;
only makes sense for off-policy algorithms.
scale_value_targets: If `True`, scale value function targets by
`1 / (1 - gamma)`.
q_value: If `True`, use Q-values as baselines.
q_value_aggregate_max: If `True`, aggregate Q-values with max (or mean).
q_value_n_samples: Number of samples to average over when calculating
baselines based on Q-values.
**kwargs: Arguments for `PolicyTrainer` superclass.
"""
self._n_shared_layers = n_shared_layers
self._value_batch_size = value_batch_size
self._value_train_steps_per_epoch = value_train_steps_per_epoch
self._value_evals_per_epoch = value_evals_per_epoch
self._value_eval_steps = value_eval_steps
# The 2 below will be initalized in super.__init__ anyway, but are needed
# to construct value batches which are needed before PolicyTrainer init
# since policy input creation calls the value model -- hence this code.
self._task = task
self._max_slice_length = kwargs.get('max_slice_length', 1)
self._added_policy_slice_length = added_policy_slice_length
self._n_replay_epochs = n_replay_epochs
task.set_n_replay_epochs(n_replay_epochs)
if scale_value_targets:
self._value_network_scale = 1 / (1 - self._task.gamma)
else:
self._value_network_scale = 1
self._q_value = q_value
self._q_value_aggregate_max = q_value_aggregate_max
self._q_value_n_samples = q_value_n_samples
is_discrete = isinstance(self._task.action_space, gym.spaces.Discrete)
self._is_discrete = is_discrete
self._vocab_size = None
self._sample_all_discrete_actions = False
if q_value and is_discrete:
self._vocab_size = self.task.action_space.n
# TODO(lukaszkaiser): the code below is specific to AWR, move it.
# If n_samples = n_actions, we'll take them all in actor and reweight.
if self._q_value_n_samples == self._vocab_size:
# TODO(lukaszkaiser): set this explicitly once it's in AWR Trainer.
self._sample_all_discrete_actions = True
if q_value:
value_model = functools.partial(value_model,
inject_actions=True,
is_discrete=is_discrete,
vocab_size=self._vocab_size)
self._value_eval_model = value_model(mode='eval')
self._value_eval_model.init(self._value_model_signature)
self._value_eval_jit = tl.jit_forward(
self._value_eval_model.pure_fn, fastmath.device_count(), do_mean=False)
# Initialize policy training.
super(ActorCriticTrainer, self).__init__(task, **kwargs)
# Initialize training of the value function.
value_output_dir = kwargs.get('output_dir', None)
if value_output_dir is not None:
value_output_dir = os.path.join(value_output_dir, 'value')
# If needed, create value_output_dir and missing parent directories.
if not tf.io.gfile.isdir(value_output_dir):
tf.io.gfile.makedirs(value_output_dir)
self._value_inputs = supervised.Inputs(
train_stream=lambda _: self.value_batches_stream())
self._value_trainer = supervised.Trainer(
model=value_model,
optimizer=value_optimizer,
lr_schedule=value_lr_schedule,
loss_fn=tl.L2Loss(),
inputs=self._value_inputs,
output_dir=value_output_dir,
metrics={'value_loss': tl.L2Loss()})
def __init__(self,
task,
value_model=None,
value_optimizer=None,
value_lr_schedule=lr.MultifactorSchedule,
value_batch_size=64,
value_train_steps_per_epoch=500,
n_shared_layers=0,
added_policy_slice_length=0,
**kwargs): # Arguments of PolicyTrainer come here.
"""Configures the actor-critic Trainer.
Args:
task: RLTask instance to use
value_model: the model to use for the value function
value_optimizer: the optimizer to train the value model
value_lr_schedule: lr schedule for value model training
value_batch_size: batch size for value model training
value_train_steps_per_epoch: how many steps are we using to
train the value model in each epoch
n_shared_layers: how many layers to share between value and
policy models
added_policy_slice_length: how much longer should slices of
trajectories be for policy than for value training; this
is useful for TD calculations and only affect the length
of elements produced for policy batches; value batches
have maximum length set by max_slice_length in **kwargs
**kwargs: arguments for PolicyTrainer super-class
"""
self._n_shared_layers = n_shared_layers
self._value_batch_size = value_batch_size
self._value_train_steps_per_epoch = value_train_steps_per_epoch
# The 2 below will be initalized in super.__init__ anyway, but are needed
# to construct value batches which are needed before PolicyTrainer init
# since policy input creation calls the value model -- hence this code.
self._task = task
self._max_slice_length = kwargs.get('max_slice_length', None)
self._added_policy_slice_length = added_policy_slice_length
# Initialize training of the value function.
value_output_dir = kwargs.get('output_dir', None)
if value_output_dir is not None:
value_output_dir = os.path.join(value_output_dir, '/value')
self._value_inputs = supervised.Inputs(
train_stream=lambda _: self.value_batches_stream())
self._value_trainer = supervised.Trainer(
model=value_model,
optimizer=value_optimizer,
lr_schedule=value_lr_schedule,
loss_fn=tl.L2Loss,
inputs=self._value_inputs,
output_dir=value_output_dir,
metrics={'value_loss': tl.L2Loss},
has_weights=True)
self._value_eval_model = value_model(mode='eval')
value_batch = next(self.value_batches_stream())
self._value_eval_model.init(value_batch)
# Initialize policy training.
super(ActorCriticTrainer, self).__init__(task, **kwargs)
