def collect_n_transitions(tf_env, policy, data, n, log_freq=10000): """Adds desired number of transitions to dataset.""" collector = train_eval_utils.DataCollector(tf_env, policy, data) time_st = time.time() timed_at_step = 0 steps_collected = 0 while steps_collected < n: count = collector.collect_transition() steps_collected += count if (steps_collected % log_freq == 0 or steps_collected == n) and count > 0: steps_per_sec = ((steps_collected - timed_at_step) / (time.time() - time_st)) timed_at_step = steps_collected time_st = time.time() logging.info('(%d/%d) steps collected at %.4g steps/s.', steps_collected, n, steps_per_sec)
def train_eval_online( # Basic args. log_dir, agent_module, env_name='HalfCheetah-v2', # Train and eval args. total_train_steps=int(1e6), summary_freq=100, print_freq=1000, save_freq=int(1e8), eval_freq=5000, n_eval_episodes=20, # For saving a partially trained policy. eval_target=None, # Target return value to stop training. eval_target_n=2, # Stop after n consecutive evals above eval_target. # Agent train args. initial_explore_steps=10000, replay_buffer_size=int(1e6), model_params=(((200, 200),), 2), optimizers=(('adam', 0.001),), batch_size=256, weight_decays=(0.0,), update_freq=1, update_rate=0.005, discount=0.99, ): """Training a policy with online interaction.""" # Create tf_env to get specs. tf_env = train_eval_utils.env_factory(env_name) tf_env_test = train_eval_utils.env_factory(env_name) observation_spec = tf_env.observation_spec() action_spec = tf_env.action_spec() # Initialize dataset. with tf.device('/cpu:0'): train_data = dataset.Dataset( observation_spec, action_spec, replay_buffer_size, circular=True, ) data_ckpt = tf.train.Checkpoint(data=train_data) data_ckpt_name = os.path.join(log_dir, 'replay') time_st_total = time.time() time_st = time.time() timed_at_step = 0 # Collect data from random policy. explore_policy = policies.ContinuousRandomPolicy(action_spec) steps_collected = 0 log_freq = 5000 logging.info('Collecting data ...') collector = train_eval_utils.DataCollector(tf_env, explore_policy, train_data) while steps_collected < initial_explore_steps: count = collector.collect_transition() steps_collected += count if (steps_collected % log_freq == 0 or steps_collected == initial_explore_steps) and count > 0: steps_per_sec = ((steps_collected - timed_at_step) / (time.time() - time_st)) timed_at_step = steps_collected time_st = time.time() logging.info('(%d/%d) steps collected at %.4g steps/s.', steps_collected, initial_explore_steps, steps_per_sec) # Construct agent. agent_flags = utils.Flags( action_spec=action_spec, model_params=model_params, optimizers=optimizers, batch_size=batch_size, weight_decays=weight_decays, update_freq=update_freq, update_rate=update_rate, discount=discount, train_data=train_data) agent_args = agent_module.Config(agent_flags).agent_args agent = agent_module.Agent(**vars(agent_args)) # Prepare savers for models and results. train_summary_dir = os.path.join(log_dir, 'train') eval_summary_dir = os.path.join(log_dir, 'eval') train_summary_writer = tf.compat.v2.summary.create_file_writer( train_summary_dir) eval_summary_writers = collections.OrderedDict() for policy_key in agent.test_policies.keys(): eval_summary_writer = tf.compat.v2.summary.create_file_writer( os.path.join(eval_summary_dir, policy_key)) eval_summary_writers[policy_key] = eval_summary_writer agent_ckpt_name = os.path.join(log_dir, 'agent') eval_results = [] # Train agent. logging.info('Start training ....') time_st = time.time() timed_at_step = 0 target_partial_policy_saved = False collector = train_eval_utils.DataCollector( tf_env, agent.online_policy, train_data) for _ in range(total_train_steps): collector.collect_transition() agent.train_step() step = agent.global_step if step % summary_freq == 0 or step == total_train_steps: agent.write_train_summary(train_summary_writer) if step % print_freq == 0 or step == total_train_steps: agent.print_train_info() if step % eval_freq == 0 or step == total_train_steps: time_ed = time.time() time_cost = time_ed - time_st logging.info( 'Training at %.4g steps/s.', (step - timed_at_step) / time_cost) eval_result, eval_infos = train_eval_utils.eval_policies( tf_env_test, agent.test_policies, n_eval_episodes) eval_results.append([step] + eval_result) # Cecide whether to save a partially trained policy based on current model # performance. if (eval_target is not None and len(eval_results) >= eval_target_n and not target_partial_policy_saved): evals_ = list([eval_results[-(i + 1)][1] for i in range(eval_target_n)]) evals_ = np.array(evals_) if np.min(evals_) >= eval_target: agent.save(agent_ckpt_name + '_partial_target') dataset.save_copy(train_data, data_ckpt_name + '_partial_target') logging.info('A partially trained policy was saved at step %d,' ' with episodic return %.4g.', step, evals_[-1]) target_partial_policy_saved = True logging.info('Testing at step %d:', step) for policy_key, policy_info in eval_infos.items(): logging.info(utils.get_summary_str( step=None, info=policy_info, prefix=policy_key + ': ')) utils.write_summary(eval_summary_writers[policy_key], step, policy_info) time_st = time.time() timed_at_step = step if step % save_freq == 0: agent.save(agent_ckpt_name + '-' + str(step)) # Final save after training. agent.save(agent_ckpt_name + '_final') data_ckpt.write(data_ckpt_name + '_final') time_cost = time.time() - time_st_total logging.info('Training finished, time cost %.4gs.', time_cost) return np.array(eval_results)