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
# * we are using has_weights=True to allow inputs to set weights
# * 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={'policy_loss': self.policy_loss},
has_weights=True)
self._policy_eval_model = policy_model(mode='eval')
policy_batch = next(self.policy_batches_stream())
self._policy_eval_model.init(policy_batch)
def __init__(self,
task,
joint_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 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.
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._collect_per_epoch = collect_per_epoch
self._max_slice_length = max_slice_length
self._policy_dist = distributions.create_distribution(
task.env.action_space)
# Inputs to the joint model are produced by self.batches_stream.
self._inputs = supervised.Inputs(
train_stream=lambda _: self.batches_stream())
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=joint_model,
optimizer=optimizer,
lr_schedule=lr_schedule,
loss_fn=self.joint_loss,
inputs=self._inputs,
output_dir=output_dir,
# TODO(lukaszkaiser): log policy and value losses too.
metrics={'joint_loss': self.joint_loss})
self._eval_model = 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(ExamplePolicyTrainer, self).__init__(task, 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)
# 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 CrossEntropyLoss
# This is because we are training a policy model, so targets are
# actions and they are integers -- CrossEntropyLoss will calculate
# the probability of each action in the state, 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.
# * we use id_to_mask=0
# This is because we reserved 0 for padding actions - so true actions
# start from 1 and we want to remove any loss on the 0 padding.
self._trainer = supervised.Trainer(
model=model, optimizer=optimizer, lr_schedule=lr_schedule,
loss_fn=tl.CrossEntropyLoss, inputs=self._inputs, output_dir=output_dir,
has_weights=True, id_to_mask=0)
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,
on_policy=True,
**kwargs): # Arguments of PolicyTrainer come here.
"""Configures the actor-critic Trainer."""
self._on_policy = on_policy
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)
# 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)
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)
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,
vocab_size=2,
**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.
vocab_size: Embedding vocabulary size (passed to `tl.Embedding`); used
only with discrete actions and when `q_value` is `True`.
**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
self._vocab_size = vocab_size
is_discrete = isinstance(self._task.action_space, gym.spaces.Discrete)
# TODO(henrykm) handle the case other than Discrete/Gaussian
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,
math.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,
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,
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 RLTrainer.
"""
super(PolicyTrainer, self).__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 = 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))
self._policy_eval_model = policy_model(
mode='eval') # 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,
**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)
**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
if scale_value_targets:
self._value_network_scale = 1 / (1 - self._task.gamma)
else:
self._value_network_scale = 1
self._value_eval_model = value_model(mode='eval')
self._value_eval_model.init(self._value_model_signature)
# 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()})
# Initialize policy training.
super(ActorCriticTrainer, self).__init__(task, **kwargs)
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)
