def main(argv): del argv # Create the task environment. test_w = [float(x) for x in FLAGS.test_w] env_config = configs.get_fig5_task_config(test_w) env = scavenger.Scavenger(**env_config) env = environment_wrappers.EnvironmentWithLogging(env) # Create the flat agent. agent = dqn_agent.Agent(obs_spec=env.observation_spec(), action_spec=env.action_spec(), network_kwargs=dict( output_sizes=(64, 128), activate_final=True, ), epsilon=0.1, additional_discount=0.9, batch_size=10, optimizer_name="AdamOptimizer", optimizer_kwargs=dict(learning_rate=3e-4, )) _, ema_returns = experiment.run(env, agent, num_episodes=FLAGS.num_episodes, report_every=FLAGS.report_every) if FLAGS.output_path: experiment.write_returns_to_file(FLAGS.output_path, ema_returns)
def main(argv): del argv # Load the keyboard. keyboard = smart_module.SmartModuleImport(hub.Module(FLAGS.keyboard_path)) # Create the task environment. base_env_config = configs.get_fig4_task_config() base_env = scavenger.Scavenger(**base_env_config) base_env = environment_wrappers.EnvironmentWithLogging(base_env) # Wrap the task environment with the keyboard. additional_discount = 0.9 env = environment_wrappers.EnvironmentWithKeyboardDirect( env=base_env, keyboard=keyboard, keyboard_ckpt_path=None, additional_discount=additional_discount, call_and_return=False) # Create the player agent. agent = regressed_agent.Agent( batch_size=10, optimizer_name="AdamOptimizer", optimizer_kwargs=dict(learning_rate=3e-2, ), init_w=np.random.normal(size=keyboard.num_cumulants) * 0.1, ) _, ema_returns = experiment.run(env, agent, num_episodes=FLAGS.num_episodes, report_every=FLAGS.report_every, num_eval_reps=100) if FLAGS.output_path: experiment.write_returns_to_file(FLAGS.output_path, ema_returns)
def create_and_train_keyboard_with_phi(num_episodes, phi_model_path, policy_weights, export_path=None): """Train an option keyboard.""" env_config = configs.get_pretrain_config() env = scavenger.Scavenger(**env_config) env = environment_wrappers.EnvironmentWithLogging(env) env = environment_wrappers.EnvironmentWithLearnedPhi(env, phi_model_path) agent = keyboard_agent.Agent(obs_spec=env.observation_spec(), action_spec=env.action_spec(), policy_weights=policy_weights, network_kwargs=dict( output_sizes=(64, 128), activate_final=True, ), epsilon=0.1, additional_discount=0.9, batch_size=10, optimizer_name="AdamOptimizer", optimizer_kwargs=dict(learning_rate=3e-4, )) if num_episodes: experiment.run(env, agent, num_episodes=num_episodes) agent.export(export_path) return agent
def _train_keyboard(num_episodes): """Train an option keyboard.""" env_config = configs.get_pretrain_config() env = scavenger.Scavenger(**env_config) env = environment_wrappers.EnvironmentWithLogging(env) agent = keyboard_agent.Agent(obs_spec=env.observation_spec(), action_spec=env.action_spec(), policy_weights=np.array([ [1.0, 0.0], [0.0, 1.0], ]), network_kwargs=dict( output_sizes=(64, 128), activate_final=True, ), epsilon=0.1, additional_discount=0.9, batch_size=10, optimizer_name="AdamOptimizer", optimizer_kwargs=dict(learning_rate=3e-4, )) experiment.run(env, agent, num_episodes=num_episodes) return agent
def main(argv): del argv # Load the keyboard. keyboard = smart_module.SmartModuleImport(hub.Module(FLAGS.keyboard_path)) # Create the task environment. base_env_config = configs.get_fig4_task_config() base_env = scavenger.Scavenger(**base_env_config) base_env = environment_wrappers.EnvironmentWithLogging(base_env) # Wrap the task environment with the keyboard. additional_discount = 0.9 env = environment_wrappers.EnvironmentWithKeyboardDirect( env=base_env, keyboard=keyboard, keyboard_ckpt_path=None, additional_discount=additional_discount, call_and_return=False) # Create the player agent. agent = regressed_agent.Agent( batch_size=10, optimizer_name="AdamOptimizer", # Disable training. optimizer_kwargs=dict(learning_rate=0.0,), init_w=[1., -1.]) returns = [] for _ in range(FLAGS.num_episodes): returns.append(experiment.run_episode(env, agent)) tf.logging.info("#" * 80) tf.logging.info( f"Avg. return over {FLAGS.num_episodes} episodes is {np.mean(returns)}") tf.logging.info("#" * 80)
def main(argv): del argv # Pretrain the keyboard and save a checkpoint. if FLAGS.keyboard_path: keyboard_path = FLAGS.keyboard_path else: with tf.Graph().as_default(): export_path = "/tmp/option_keyboard/keyboard" _ = keyboard_utils.create_and_train_keyboard( num_episodes=FLAGS.num_pretrain_episodes, export_path=export_path) keyboard_path = os.path.join(export_path, "tfhub") # Load the keyboard. keyboard = smart_module.SmartModuleImport(hub.Module(keyboard_path)) # Create the task environment. base_env_config = configs.get_task_config() base_env = scavenger.Scavenger(**base_env_config) base_env = environment_wrappers.EnvironmentWithLogging(base_env) # Wrap the task environment with the keyboard. additional_discount = 0.9 env = environment_wrappers.EnvironmentWithKeyboard( env=base_env, keyboard=keyboard, keyboard_ckpt_path=None, n_actions_per_dim=3, additional_discount=additional_discount, call_and_return=False) # Create the player agent. agent = dqn_agent.Agent(obs_spec=env.observation_spec(), action_spec=env.action_spec(), network_kwargs=dict( output_sizes=(64, 128), activate_final=True, ), epsilon=0.1, additional_discount=additional_discount, batch_size=10, optimizer_name="AdamOptimizer", optimizer_kwargs=dict(learning_rate=3e-4, )) _, ema_returns = experiment.run(env, agent, num_episodes=FLAGS.num_episodes, report_every=FLAGS.report_every) if FLAGS.output_path: experiment.write_returns_to_file(FLAGS.output_path, ema_returns)
def evaluate_keyboard(keyboard_path): """Evaluate a keyboard.""" angles_to_sweep = np.deg2rad(np.linspace(-90, 180, num=19, endpoint=True)) weights_to_sweep = np.stack( [np.cos(angles_to_sweep), np.sin(angles_to_sweep)], axis=-1) weights_to_sweep /= np.sum( np.maximum(weights_to_sweep, 0.0), axis=-1, keepdims=True) weights_to_sweep = np.clip(weights_to_sweep, -1000, 1000) tf.logging.info(weights_to_sweep) # Load the keyboard. keyboard = smart_module.SmartModuleImport(hub.Module(keyboard_path)) # Create the task environment. all_returns = [] for w_to_sweep in weights_to_sweep.tolist(): base_env_config = configs.get_fig5_task_config(w_to_sweep) base_env = scavenger.Scavenger(**base_env_config) base_env = environment_wrappers.EnvironmentWithLogging(base_env) # Wrap the task environment with the keyboard. with tf.variable_scope(None, default_name="inner_loop"): additional_discount = 0.9 env = environment_wrappers.EnvironmentWithKeyboardDirect( env=base_env, keyboard=keyboard, keyboard_ckpt_path=None, additional_discount=additional_discount, call_and_return=False) # Create the player agent. agent = regressed_agent.Agent( batch_size=10, optimizer_name="AdamOptimizer", # Disable training. optimizer_kwargs=dict(learning_rate=0.0,), init_w=w_to_sweep) returns = [] for _ in range(FLAGS.num_episodes): returns.append(experiment.run_episode(env, agent)) tf.logging.info(f"Task: {w_to_sweep}, mean returns over " f"{FLAGS.num_episodes} episodes is {np.mean(returns)}") all_returns.append(returns) return all_returns, weights_to_sweep
def main(argv): del argv # Load the keyboard. keyboard = smart_module.SmartModuleImport(hub.Module(FLAGS.keyboard_path)) # Create the task environment. base_env_config = configs.get_task_config() base_env = scavenger.Scavenger(**base_env_config) base_env = environment_wrappers.EnvironmentWithLogging(base_env) # Wrap the task environment with the keyboard. additional_discount = 0.9 env = environment_wrappers.EnvironmentWithKeyboardDirect( env=base_env, keyboard=keyboard, keyboard_ckpt_path=None, additional_discount=additional_discount, call_and_return=False) # Create the player agent. agent = regressed_agent.Agent( batch_size=10, optimizer_name="AdamOptimizer", # Disable training. optimizer_kwargs=dict(learning_rate=0.0, ), init_w=[float(x) for x in FLAGS.test_w]) returns = [] for _ in range(FLAGS.num_episodes): returns.append(experiment.run_episode(env, agent)) tf.logging.info("#" * 80) tf.logging.info( f"Avg. return over {FLAGS.num_episodes} episodes is {np.mean(returns)}" ) tf.logging.info("#" * 80) if FLAGS.output_path: with gfile.GFile(FLAGS.output_path, "w") as file: writer = csv.writer(file, delimiter=" ", quoting=csv.QUOTE_MINIMAL) writer.writerow(["return"]) for val in returns: writer.writerow([val])
def main(argv): del argv # Create the task environment. env_config = configs.get_fig4_task_config() env = scavenger.Scavenger(**env_config) env = environment_wrappers.EnvironmentWithLogging(env) # Create the flat agent. agent = dqn_agent.Agent(obs_spec=env.observation_spec(), action_spec=env.action_spec(), network_kwargs=dict( output_sizes=(64, 128), activate_final=True, ), epsilon=0.1, additional_discount=0.9, batch_size=10, optimizer_name="AdamOptimizer", optimizer_kwargs=dict(learning_rate=3e-4, )) experiment.run(env, agent, num_episodes=FLAGS.num_episodes)
def main(argv): del argv # Pretrain the keyboard and save a checkpoint. pretrain_agent = _train_keyboard(num_episodes=FLAGS.num_pretrain_episodes) keyboard_ckpt_path = "/tmp/option_keyboard/keyboard.ckpt" pretrain_agent.export(keyboard_ckpt_path) # Create the task environment. base_env_config = configs.get_task_config() base_env = scavenger.Scavenger(**base_env_config) base_env = environment_wrappers.EnvironmentWithLogging(base_env) # Wrap the task environment with the keyboard. additional_discount = 0.9 env = environment_wrappers.EnvironmentWithKeyboard( env=base_env, keyboard=pretrain_agent.keyboard, keyboard_ckpt_path=keyboard_ckpt_path, n_actions_per_dim=3, additional_discount=additional_discount, call_and_return=True) # Create the player agent. agent = dqn_agent.Agent(obs_spec=env.observation_spec(), action_spec=env.action_spec(), network_kwargs=dict( output_sizes=(64, 128), activate_final=True, ), epsilon=0.1, additional_discount=additional_discount, batch_size=10, optimizer_name="AdamOptimizer", optimizer_kwargs=dict(learning_rate=3e-4, )) experiment.run(env, agent, num_episodes=FLAGS.num_episodes)
def main(argv): del argv if FLAGS.use_random_tasks: tasks = np.random.normal(size=(8, 2)) else: tasks = [ [1.0, 0.0], [0.0, 1.0], [-1.0, 0.0], [0.0, -1.0], [0.7, 0.3], [-0.3, -0.7], ] if FLAGS.normalisation == "L1": tasks /= np.sum(np.abs(tasks), axis=-1, keepdims=True) elif FLAGS.normalisation == "L2": tasks /= np.linalg.norm(tasks, axis=-1, keepdims=True) else: raise ValueError("Unknown normlisation_method {}".format( FLAGS.normalisation)) logging.info("Tasks: %s", tasks) env_config = dict(arena_size=11, num_channels=2, max_num_steps=100, num_init_objects=10, object_priors=[1.0, 1.0], egocentric=True, default_w=None, aux_tasks_w=tasks) env = scavenger.Scavenger(**env_config) num_actions = env.action_spec().maximum + 1 model_config = dict( n_actions=num_actions, n_phis=FLAGS.num_phis, network_kwargs=dict( output_sizes=(64, 128), activate_final=True, ), ) model = smart_module.SmartModuleExport(lambda: PhiModel(**model_config)) dummy_steps = collect_experience(env, num_episodes=10, verbose=True) num_rewards = dummy_steps.rewards.shape[-1] # Placeholders steps_ph = tree.map_structure(create_ph, dummy_steps) phis = model(steps_ph.obs, steps_ph.actions) phis_to_rewards = snt.Linear(num_rewards, initializers=dict(w=tf.zeros), use_bias=False) preds = phis_to_rewards(phis) loss_per_batch = tf.square(preds - steps_ph.rewards) loss_op = tf.reduce_mean(loss_per_batch) replay = [] # Optimizer and train op. with tf.variable_scope("optimizer"): optimizer = tf.train.AdamOptimizer(learning_rate=FLAGS.learning_rate) train_op = optimizer.minimize(loss_op) # Add normalisation of weights in phis_to_rewards if FLAGS.normalisation == "L1": w_norm = tf.reduce_sum(tf.abs(phis_to_rewards.w), axis=0, keepdims=True) elif FLAGS.normalisation == "L2": w_norm = tf.norm(phis_to_rewards.w, axis=0, keepdims=True) else: raise ValueError("Unknown normlisation_method {}".format( FLAGS.normalisation)) normalise_w = tf.assign(phis_to_rewards.w, phis_to_rewards.w / tf.maximum(w_norm, 1e-6)) def filter_steps(steps): mask = np.sum(np.abs(steps.rewards), axis=-1) > 0.1 nonzero_inds = np.where(mask)[0] zero_inds = np.where(np.logical_not(mask))[0] zero_inds = np.random.choice(zero_inds, size=len(nonzero_inds), replace=False) selected_inds = np.concatenate([nonzero_inds, zero_inds]) selected_steps = tree.map_structure(lambda x: x[selected_inds], steps) return selected_steps, selected_inds with tf.Session() as sess: sess.run(tf.global_variables_initializer()) step = 0 while step < FLAGS.num_train_steps: step += 1 steps_output = collect_experience(env, num_episodes=10) selected_step_outputs, selected_inds = filter_steps(steps_output) if len(replay) > FLAGS.min_replay_size: # Do training. for _ in range(FLAGS.num_train_repeats): train_samples = random.choices(replay, k=128) train_samples = tree.map_structure( lambda *x: np.stack(x, axis=0), *train_samples) train_samples = tree.unflatten_as(steps_ph, train_samples) feed_dict = dict( zip(tree.flatten(steps_ph), tree.flatten(train_samples))) _, train_loss = sess.run([train_op, loss_op], feed_dict=feed_dict) sess.run(normalise_w) # Do evaluation. if step % 50 == 0: feed_dict = dict( zip(tree.flatten(steps_ph), tree.flatten(selected_step_outputs))) eval_loss = sess.run(loss_op, feed_dict=feed_dict) logging.info( "Step %d, train loss %f, eval loss %f, replay %s", step, train_loss, eval_loss, len(replay)) print(sess.run(phis_to_rewards.get_variables())[0].T) values = dict(step=step, train_loss=train_loss, eval_loss=eval_loss) logging.info(values) # Add to replay. if step <= FLAGS.num_replay_steps: def select_fn(ind): return lambda x: x[ind] for idx in range(len(selected_inds)): replay.append( tree.flatten( tree.map_structure(select_fn(idx), selected_step_outputs))) # Export trained model. if FLAGS.export_path: model.export(FLAGS.export_path, sess, overwrite=True)