def main(unused_argv):
# Save the string flags now as we modify them later.
string_flags = FLAGS.flags_into_string()
gin.parse_config_files_and_bindings(
[FLAGS.gin_config] if FLAGS.gin_config else [],
# Gin uses slashes to denote scopes but XM doesn't allow slashes in
# parameter names so we use __ instead and convert it to slashes here.
[s.replace('__', '/') for s in FLAGS.gin_bindings])
gym_kwargs = {}
if FLAGS.mujoco_model:
local_mujoco_model = tempfile.mkstemp(prefix='mujoco_model',
suffix='.xml')[1]
logging.info('Copying remote model %s to local file %s',
FLAGS.mujoco_model, local_mujoco_model)
tf.io.gfile.copy(FLAGS.mujoco_model,
local_mujoco_model,
overwrite=True)
gym_kwargs['model_path'] = local_mujoco_model
create_environment = lambda task, config: env.create_environment(
env_name=config.env_name,
discretization='none',
n_actions_per_dim=11,
action_ratio=30,
gym_kwargs=gym_kwargs)
if FLAGS.run_mode == 'actor':
actor.actor_loop(create_environment)
elif FLAGS.run_mode == 'learner':
logdir = FLAGS.logdir
settings = utils.init_learner_multi_host(FLAGS.num_training_tpus)
learner.learner_loop(
create_environment,
create_agent,
create_optimizer,
learner_flags.training_config_from_flags(),
settings,
action_distribution_config=continuous_action_config())
with tf.io.gfile.GFile(os.path.join(logdir, 'learner_flags.txt'),
'w') as f:
f.write(string_flags)
with tf.io.gfile.GFile(os.path.join(logdir, 'learner.gin'), 'w') as f:
f.write(gin.operative_config_str())
else:
raise ValueError('Unsupported run mode {}'.format(FLAGS.run_mode))
def learner_loop(create_env_fn, create_agent_fn, create_optimizer_fn):
"""Main learner loop.
Args:
create_env_fn: Callable that must return a newly created environment. The
callable takes the task ID as argument - an arbitrary task ID of 0 will be
passed by the learner. The returned environment should follow GYM's API.
It is only used for infering tensor shapes. This environment will not be
used to generate experience.
create_agent_fn: Function that must create a new tf.Module with the neural
network that outputs actions and new agent state given the environment
observations and previous agent state. See dmlab.agents.ImpalaDeep for an
example. The factory function takes as input the environment action and
observation spaces and a parametric distribution over actions.
create_optimizer_fn: Function that takes the final iteration as argument
and must return a tf.keras.optimizers.Optimizer and a
tf.keras.optimizers.schedules.LearningRateSchedule.
"""
logging.info('Starting learner loop')
validate_config()
settings = utils.init_learner_multi_host(FLAGS.num_training_tpus)
strategy, hosts, training_strategy, encode, decode = settings
env = create_env_fn(0, FLAGS)
parametric_action_distribution = get_parametric_distribution_for_action_space(
env.action_space)
env_output_specs = utils.EnvOutput(
tf.TensorSpec([], tf.float32, 'reward'),
tf.TensorSpec([], tf.bool, 'done'),
tf.TensorSpec(env.observation_space.shape, env.observation_space.dtype,
'observation'),
tf.TensorSpec([], tf.bool, 'abandoned'),
tf.TensorSpec([], tf.int32, 'episode_step'),
)
action_specs = tf.TensorSpec(env.action_space.shape,
env.action_space.dtype, 'action')
agent_input_specs = (action_specs, env_output_specs)
# Initialize agent and variables.
agent = create_agent_fn(env.action_space, env.observation_space,
parametric_action_distribution)
initial_agent_state = agent.initial_state(1)
agent_state_specs = tf.nest.map_structure(
lambda t: tf.TensorSpec(t.shape[1:], t.dtype), initial_agent_state)
unroll_specs = [None] # Lazy initialization.
input_ = tf.nest.map_structure(
lambda s: tf.zeros([1] + list(s.shape), s.dtype), agent_input_specs)
input_ = encode(input_)
with strategy.scope():
@tf.function
def create_variables(*args):
return agent.get_action(*decode(args))
initial_agent_output, _ = create_variables(*input_,
initial_agent_state)
if not hasattr(agent, 'entropy_cost'):
mul = FLAGS.entropy_cost_adjustment_speed
agent.entropy_cost_param = tf.Variable(
tf.math.log(FLAGS.entropy_cost) / mul,
# Without the constraint, the param gradient may get rounded to 0
# for very small values.
constraint=lambda v: tf.clip_by_value(v, -20 / mul, 20 / mul),
trainable=True,
dtype=tf.float32)
agent.entropy_cost = lambda: tf.exp(mul * agent.entropy_cost_param)
# Create optimizer.
iter_frame_ratio = (FLAGS.batch_size * FLAGS.unroll_length *
FLAGS.num_action_repeats)
final_iteration = int(
math.ceil(FLAGS.total_environment_frames / iter_frame_ratio))
optimizer, learning_rate_fn = create_optimizer_fn(final_iteration)
iterations = optimizer.iterations
optimizer._create_hypers()
optimizer._create_slots(agent.trainable_variables)
# ON_READ causes the replicated variable to act as independent variables for
# each replica.
temp_grads = [
tf.Variable(tf.zeros_like(v),
trainable=False,
synchronization=tf.VariableSynchronization.ON_READ)
for v in agent.trainable_variables
]
@tf.function
def minimize(iterator):
data = next(iterator)
def compute_gradients(args):
args = tf.nest.pack_sequence_as(unroll_specs[0],
decode(args, data))
with tf.GradientTape() as tape:
loss, logs = compute_loss(logger,
parametric_action_distribution,
agent, *args)
grads = tape.gradient(loss, agent.trainable_variables)
for t, g in zip(temp_grads, grads):
t.assign(g)
return loss, logs
loss, logs = training_strategy.run(compute_gradients, (data, ))
loss = training_strategy.experimental_local_results(loss)[0]
def apply_gradients(_):
optimizer.apply_gradients(
zip(temp_grads, agent.trainable_variables))
strategy.run(apply_gradients, (loss, ))
getattr(
agent, 'end_of_training_step_callback',
lambda: logging.info('end_of_training_step_callback not found'))()
logger.step_end(logs, training_strategy, iter_frame_ratio)
agent_output_specs = tf.nest.map_structure(
lambda t: tf.TensorSpec(t.shape[1:], t.dtype), initial_agent_output)
# Setup checkpointing and restore checkpoint.
ckpt = tf.train.Checkpoint(agent=agent, optimizer=optimizer)
if FLAGS.init_checkpoint is not None:
tf.print('Loading initial checkpoint from %s...' %
FLAGS.init_checkpoint)
ckpt.restore(FLAGS.init_checkpoint).assert_consumed()
manager = tf.train.CheckpointManager(ckpt,
FLAGS.logdir,
max_to_keep=1,
keep_checkpoint_every_n_hours=6)
last_ckpt_time = 0 # Force checkpointing of the initial model.
if manager.latest_checkpoint:
logging.info('Restoring checkpoint: %s', manager.latest_checkpoint)
ckpt.restore(manager.latest_checkpoint).assert_consumed()
last_ckpt_time = time.time()
# Logging.
summary_writer = tf.summary.create_file_writer(FLAGS.logdir,
flush_millis=20000,
max_queue=1000)
logger = utils.ProgressLogger(summary_writer=summary_writer,
starting_step=iterations * iter_frame_ratio)
servers = []
unroll_queues = []
info_specs = (
tf.TensorSpec([], tf.int64, 'episode_num_frames'),
tf.TensorSpec([], tf.float32, 'episode_returns'),
tf.TensorSpec([], tf.float32, 'episode_raw_returns'),
)
info_queue = utils.StructuredFIFOQueue(-1, info_specs)
def create_host(i, host, inference_devices):
with tf.device(host):
server = grpc.Server([FLAGS.server_address])
store = utils.UnrollStore(
FLAGS.num_envs, FLAGS.unroll_length,
(action_specs, env_output_specs, agent_output_specs))
env_run_ids = utils.Aggregator(
FLAGS.num_envs, tf.TensorSpec([], tf.int64, 'run_ids'))
env_infos = utils.Aggregator(FLAGS.num_envs, info_specs,
'env_infos')
# First agent state in an unroll.
first_agent_states = utils.Aggregator(FLAGS.num_envs,
agent_state_specs,
'first_agent_states')
# Current agent state and action.
agent_states = utils.Aggregator(FLAGS.num_envs, agent_state_specs,
'agent_states')
actions = utils.Aggregator(FLAGS.num_envs, action_specs, 'actions')
unroll_specs[0] = Unroll(agent_state_specs, *store.unroll_specs)
unroll_queue = utils.StructuredFIFOQueue(1, unroll_specs[0])
def add_batch_size(ts):
return tf.TensorSpec([FLAGS.inference_batch_size] +
list(ts.shape), ts.dtype, ts.name)
inference_specs = (
tf.TensorSpec([], tf.int32, 'env_id'),
tf.TensorSpec([], tf.int64, 'run_id'),
env_output_specs,
tf.TensorSpec([], tf.float32, 'raw_reward'),
)
inference_specs = tf.nest.map_structure(add_batch_size,
inference_specs)
def create_inference_fn(inference_device):
@tf.function(input_signature=inference_specs)
def inference(env_ids, run_ids, env_outputs, raw_rewards):
# Reset the environments that had their first run or crashed.
previous_run_ids = env_run_ids.read(env_ids)
env_run_ids.replace(env_ids, run_ids)
reset_indices = tf.where(
tf.not_equal(previous_run_ids, run_ids))[:, 0]
envs_needing_reset = tf.gather(env_ids, reset_indices)
if tf.not_equal(tf.shape(envs_needing_reset)[0], 0):
tf.print('Environment ids needing reset:',
envs_needing_reset)
env_infos.reset(envs_needing_reset)
store.reset(envs_needing_reset)
initial_agent_states = agent.initial_state(
tf.shape(envs_needing_reset)[0])
first_agent_states.replace(envs_needing_reset,
initial_agent_states)
agent_states.replace(envs_needing_reset,
initial_agent_states)
actions.reset(envs_needing_reset)
tf.debugging.assert_non_positive(
tf.cast(env_outputs.abandoned, tf.int32),
'Abandoned done states are not supported in VTRACE.')
# Update steps and return.
env_infos.add(env_ids,
(0, env_outputs.reward, raw_rewards))
done_ids = tf.gather(env_ids,
tf.where(env_outputs.done)[:, 0])
if i == 0:
info_queue.enqueue_many(env_infos.read(done_ids))
env_infos.reset(done_ids)
env_infos.add(env_ids, (FLAGS.num_action_repeats, 0., 0.))
# Inference.
prev_actions = actions.read(env_ids)
input_ = encode((prev_actions, env_outputs))
prev_agent_states = agent_states.read(env_ids)
with tf.device(inference_device):
@tf.function
def agent_inference(*args):
return agent(*decode(args), is_training=False)
agent_outputs, curr_agent_states = agent_inference(
*input_, prev_agent_states)
# Append the latest outputs to the unroll and insert completed unrolls
# in queue.
completed_ids, unrolls = store.append(
env_ids, (prev_actions, env_outputs, agent_outputs))
unrolls = Unroll(first_agent_states.read(completed_ids),
*unrolls)
unroll_queue.enqueue_many(unrolls)
first_agent_states.replace(
completed_ids, agent_states.read(completed_ids))
# Update current state.
agent_states.replace(env_ids, curr_agent_states)
actions.replace(env_ids, agent_outputs.action)
# Return environment actions to environments.
return agent_outputs.action
return inference
with strategy.scope():
server.bind(
[create_inference_fn(d) for d in inference_devices])
server.start()
unroll_queues.append(unroll_queue)
servers.append(server)
for i, (host, inference_devices) in enumerate(hosts):
create_host(i, host, inference_devices)
def dequeue(ctx):
# Create batch (time major).
env_outputs = tf.nest.map_structure(
lambda *args: tf.stack(args), *[
unroll_queues[ctx.input_pipeline_id].dequeue() for i in range(
ctx.get_per_replica_batch_size(FLAGS.batch_size))
])
env_outputs = env_outputs._replace(
prev_actions=utils.make_time_major(env_outputs.prev_actions),
env_outputs=utils.make_time_major(env_outputs.env_outputs),
agent_outputs=utils.make_time_major(env_outputs.agent_outputs))
env_outputs = env_outputs._replace(
env_outputs=encode(env_outputs.env_outputs))
# tf.data.Dataset treats list leafs as tensors, so we need to flatten and
# repack.
return tf.nest.flatten(env_outputs)
def dataset_fn(ctx):
dataset = tf.data.Dataset.from_tensors(0).repeat(None)
def _dequeue(_):
return dequeue(ctx)
return dataset.map(_dequeue,
num_parallel_calls=ctx.num_replicas_in_sync //
len(hosts))
dataset = training_strategy.experimental_distribute_datasets_from_function(
dataset_fn)
it = iter(dataset)
def additional_logs():
tf.summary.scalar('learning_rate', learning_rate_fn(iterations))
n_episodes = info_queue.size()
n_episodes -= n_episodes % FLAGS.log_episode_frequency
if tf.not_equal(n_episodes, 0):
episode_stats = info_queue.dequeue_many(n_episodes)
episode_keys = [
'episode_num_frames', 'episode_return', 'episode_raw_return'
]
for key, values in zip(episode_keys, episode_stats):
for value in tf.split(
values,
values.shape[0] // FLAGS.log_episode_frequency):
tf.summary.scalar(key, tf.reduce_mean(value))
for (frames, ep_return, raw_return) in zip(*episode_stats):
logging.info('Return: %f Raw return: %f Frames: %i', ep_return,
raw_return, frames)
logger.start(additional_logs)
# Execute learning.
while iterations < final_iteration:
# Save checkpoint.
current_time = time.time()
if current_time - last_ckpt_time >= FLAGS.save_checkpoint_secs:
manager.save()
# Apart from checkpointing, we also save the full model (including
# the graph). This way we can load it after the code/parameters changed.
tf.saved_model.save(agent, os.path.join(FLAGS.logdir,
'saved_model'))
last_ckpt_time = current_time
minimize(it)
logger.shutdown()
manager.save()
tf.saved_model.save(agent, os.path.join(FLAGS.logdir, 'saved_model'))
for server in servers:
server.shutdown()
for unroll_queue in unroll_queues:
unroll_queue.close()
def visualize(create_env_fn, create_agent_fn, create_optimizer_fn):
print('Visualization launched...')
settings = utils.init_learner_multi_host(1)
strategy, hosts, training_strategy, encode, decode = settings
env = create_env_fn(0)
parametric_action_distribution = get_parametric_distribution_for_action_space(
env.action_space)
agent = create_agent_fn(env.action_space, env.observation_space,
parametric_action_distribution)
optimizer, learning_rate_fn = create_optimizer_fn(1e9)
env_output_specs = utils.EnvOutput(
tf.TensorSpec([], tf.float32, 'reward'),
tf.TensorSpec([], tf.bool, 'done'),
tf.TensorSpec(env.observation_space.shape, env.observation_space.dtype,
'observation'),
tf.TensorSpec([], tf.bool, 'abandoned'),
tf.TensorSpec([], tf.int32, 'episode_step'),
)
action_specs = tf.TensorSpec(env.action_space.shape,
env.action_space.dtype, 'action')
agent_input_specs = (action_specs, env_output_specs)
# Initialize agent and variables.
agent = create_agent_fn(env.action_space, env.observation_space,
parametric_action_distribution)
initial_agent_state = agent.initial_state(1)
agent_state_specs = tf.nest.map_structure(
lambda t: tf.TensorSpec(t.shape[1:], t.dtype), initial_agent_state)
unroll_specs = [None] # Lazy initialization.
input_ = tf.nest.map_structure(
lambda s: tf.zeros([1] + list(s.shape), s.dtype), agent_input_specs)
input_ = encode(input_)
with strategy.scope():
@tf.function
def create_variables(*args):
return agent.get_action(*decode(args))
initial_agent_output, _ = create_variables(*input_, initial_agent_state)
if not hasattr(agent, 'entropy_cost'):
mul = FLAGS.entropy_cost_adjustment_speed
agent.entropy_cost_param = tf.Variable(
tf.math.log(FLAGS.entropy_cost) / mul,
# Without the constraint, the param gradient may get rounded to 0
# for very small values.
constraint=lambda v: tf.clip_by_value(v, -20 / mul, 20 / mul),
trainable=True,
dtype=tf.float32)
agent.entropy_cost = lambda: tf.exp(mul * agent.entropy_cost_param)
# Create optimizer.
iter_frame_ratio = (
FLAGS.batch_size * FLAGS.unroll_length * FLAGS.num_action_repeats)
final_iteration = int(
math.ceil(FLAGS.total_environment_frames / iter_frame_ratio))
optimizer, learning_rate_fn = create_optimizer_fn(final_iteration)
iterations = optimizer.iterations
optimizer._create_hypers()
optimizer._create_slots(agent.trainable_variables)
# ON_READ causes the replicated variable to act as independent variables for
# each replica.
temp_grads = [
tf.Variable(tf.zeros_like(v), trainable=False,
synchronization=tf.VariableSynchronization.ON_READ)
for v in agent.trainable_variables
]
agent_output_specs = tf.nest.map_structure(
lambda t: tf.TensorSpec(t.shape[1:], t.dtype), initial_agent_output)
if True:
ckpt = tf.train.Checkpoint(agent=agent, optimizer=optimizer)
ckpt.restore('seed_rl/checkpoints/agent_good_3m/ckpt-9').assert_consumed()
def get_agent_action(obs):
initial_agent_state = agent.initial_state(1)
shaped_obs = tf.reshape(tf.convert_to_tensor(obs), shape=(1,)+env.observation_space.shape)
initial_env_output = (tf.constant([1.]), tf.constant([False]), shaped_obs,
tf.constant([False]), tf.constant([1], dtype=tf.float32),)
agent_out = agent(tf.zeros([0], dtype=tf.float32), initial_env_output,
initial_agent_state)
return agent_out
def run_episode(steps):
mode = None
obs = env.reset()
rewards = []
for _ in range(steps):
agent_out, state = get_agent_action(obs)
action = agent_out.action.numpy()[0]
obs, rew, done, info = env.step(action)
rewards.append(rew)
if done:
break
reward = np.sum(rewards)
print('reward: {0}'.format(reward))
return reward
all_rewards = []
iter = 0
while True:
all_rewards.append(run_episode(250))
if len(all_rewards) > 1000:
all_rewards = all_rewards[-1000:]
print('mean cum reward: {0}'.format(np.mean(all_rewards)))
if iter % 10 == 0:
env.save_replay()
print('\n REPLAY SAVED\n')
iter += 1
print('Graceful termination')
sys.exit(0)