def train_dyke_agent(train_env: TFPyEnvironment, eval_env: TFPyEnvironment, agent: DqnAgent, train_steps: int, steps_per_episode: int, eval_episodes: int) -> Dict[str, Any]: """ Trains the DQN agent on the dyke maintenance task. :param train_env: The training environment. :param eval_env: The environment for testing agent performance. :param agent: The agent. :param train_steps: The number of training steps to use. :param steps_per_episode: The number of time steps that can be taken in a single dyke environment episode. :param eval_episodes: The number of episodes to use per evaluation. :return: A mapping to various metrics pertaining to the training's results. """ losses: np.ndarray = np.zeros(shape=(train_steps, steps_per_episode)) evaluations: np.ndarray = np.zeros(shape=(train_steps, eval_episodes)) train_metrics: Tuple = (AverageReturnMetric, ) train_metric_results: np.ndarray = np.zeros(shape=(len(train_metrics), train_steps, steps_per_episode)) for step in range(train_steps): # we uniformly sample experiences (single time steps) from one episode per train step print('STEP %d/%d' % (step + 1, train_steps)) train_env.reset() rep_buf = _dyke_replay_buffer(train_env, agent, steps_per_episode) train_metric_inst: Tuple = tuple( [metric() for metric in train_metrics]) # instantiate the metrics obs: Tuple = (rep_buf.add_batch, ) + train_metric_inst _ = DynamicStepDriver( env=train_env, policy=agent.collect_policy, observers=obs, num_steps=steps_per_episode ).run( ) # experience a single episode using the agent's current configuration dataset: tf.data.Dataset = rep_buf.as_dataset( sample_batch_size=_REP_BUF_BATCH_SIZE, num_steps=_REP_BUF_NUM_STEPS) iterator = iter(dataset) for tr in range(steps_per_episode): trajectories, _ = next(iterator) losses[step, tr] = agent.train(experience=trajectories).loss for met in range(len(train_metrics)): train_metric_results[ met, step, tr] = train_metric_inst[met].result().numpy() evaluations[step, :] = _evaluate_dyke_agent(eval_env, agent, eval_episodes) return { 'loss': losses, 'eval': evaluations, 'train-metrics': train_metric_results }
def create_pong_agent(train_environment: TFEnvironment, dense_layer_sizes: Sequence[int], learning_rate: float) -> (DqnAgent, QNetwork): optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate) global_step = tf.compat.v1.train.get_or_create_global_step() q_net = QNetwork(input_tensor_spec=train_environment.observation_spec(), action_spec=train_environment.action_spec(), fc_layer_params=dense_layer_sizes) agent = DqnAgent(time_step_spec=train_environment.time_step_spec(), action_spec=train_environment.action_spec(), q_network=q_net, optimizer=optimizer, td_errors_loss_fn=element_wise_squared_loss, train_step_counter=global_step) agent.initialize() agent.train = common.function(agent.train) agent.train_step_counter.assign(0) return agent
decay_steps=epsilon_decay_steps, end_learning_rate=epsilon_final) # final ε agent = DqnAgent( tf_env.time_step_spec(), tf_env.action_spec(), q_network=q_net, optimizer=optimizer, target_update_period=target_update_period, td_errors_loss_fn=keras.losses.Huber(reduction="none"), gamma=discount_factor, # discount factor train_step_counter=train_step, epsilon_greedy=lambda: epsilon_fn(train_step)) agent.initialize() # Speed up as tensorflow function agent.train = function(agent.train) ## ------------------------------------------------------------------------------ ## ------------------------------------------------------------------------------ ## ------------------------------------------------------------------------------ replay_buffer = tf_uniform_replay_buffer.TFUniformReplayBuffer( # Determines the data spec type data_spec=agent.collect_data_spec, # The number of trajectories added at each step batch_size=tf_env.batch_size, # This can store 4 million trajectories (note: requires a lot of RAM) max_length=n_iterations) # Create the observer that adds trajectories to the replay buffer replay_buffer_observer = replay_buffer.add_batch
# execute the random policy in the environment for a few steps # and record the data (observations, actions, rewards etc) in the replay buffer print("Collecting initial random steps") random_policy = RandomTFPolicy(train_env.time_step_spec(), train_env.action_spec()) for _ in range(initial_collect_steps): collect_step(train_env, random_policy) dataset = replay_buffer.as_dataset(num_parallel_calls=3, sample_batch_size=batch_size, num_steps=2).prefetch(3) iterator = iter(dataset) # train the agent print("Training the agent") tf_agent.train = common.function(tf_agent.train) # Reset the train step tf_agent.train_step_counter.assign(0) def compute_avg_reward(environment, policy, num_episodes=10): total_reward = 0.0 for episode in range(num_episodes): print(f"Computing reward: game {episode + 1}/{num_episodes}") time_step = environment.reset() episode_reward = 0.0 while not time_step.is_last(): action_step = policy.action(time_step)
class DQNAgent: def __init__(self) -> None: """ A class for training a TF-agent based on https://www.tensorflow.org/agents/tutorials/1_dqn_tutorial """ self.train_env = None # Training environment self.agent = None # The algorithm used to solve an RL problem is represented by a TF-Agent self.replay_buffer = None # The replay buffer keeps track of data collected from the environment self.dataset = None # The agent needs access to the replay buffer via an iterable tf.data.Dataset self.iterator = None # The iterator of self.dataset def compile(self, X_train: np.ndarray, y_train: np.ndarray, lr: float, epsilon: float, gamma: float, imb_ratio: float, replay_buffer_max_length: int, layers: dict) -> None: """ Create the Q-network, agent and policy Args: X_train: A np.ndarray for training samples. y_train: A np.ndarray for the class labels of the training samples. lr: learn rate for the optimizer (default Adam) epsilon: Used for the default epsilon greedy policy for choosing a random action. gamma: The discount factor for learning Q-values imb_ratio: ratio of imbalance. Used to specifiy reward in the environment replay_buffer_max_length: Maximum lenght of replay memory. layers: A dict containing the layers of the Q-Network (eg, conv, dense, rnn, dropout). """ dense_layers = layers.get("dense") conv_layers = layers.get("conv") dropout_layers = layers.get("dropout") self.train_env = TFPyEnvironment(ClassifyEnv(X_train, y_train, imb_ratio)) # create a custom environment q_net = QNetwork(self.train_env.observation_spec(), self.train_env.action_spec(), conv_layer_params=conv_layers, fc_layer_params=dense_layers, dropout_layer_params=dropout_layers) optimizer = tf.compat.v1.train.AdamOptimizer(learning_rate=lr) train_step_counter = tf.Variable(0) self.agent = DqnAgent( self.train_env.time_step_spec(), self.train_env.action_spec(), q_network=q_net, optimizer=optimizer, td_errors_loss_fn=common.element_wise_squared_loss, train_step_counter=train_step_counter, gamma=gamma, epsilon_greedy=epsilon, ) self.agent.initialize() self.replay_buffer = TFUniformReplayBuffer( data_spec=self.agent.collect_data_spec, batch_size=self.train_env.batch_size, max_length=replay_buffer_max_length) def fit(self, X_train: np.ndarray, y_train: np.ndarray, epochs: int, batch_size: int, eval_step: int, log_step: int, collect_steps_per_episode: int) -> None: """ Starts the training of the Agent. Args: X_train: A np.ndarray for training samples. y_train: A np.ndarray for the class labels of the training samples. epochs: Number of epochs to train Agent batch_size: The Batch Size eval_step: Evaluate Model each 'eval_step' log_step: Monitor results of model each 'log_step' collect_steps_per_episode: Collect a few steps using collect_policy and save to the replay buffer. """ self.dataset = self.replay_buffer.as_dataset( num_parallel_calls=3, sample_batch_size=batch_size, num_steps=2).prefetch(3) self.iterator = iter(self.dataset) def collect_step(environment, policy, buffer): time_step = environment.current_time_step() action_step = policy.action(time_step) next_time_step = environment.step(action_step.action) traj = trajectory.from_transition(time_step, action_step, next_time_step) # Add trajectory to the replay buffer buffer.add_batch(traj) def collect_data(env, policy, buffer, steps): for _ in range(steps): collect_step(env, policy, buffer) # (Optional) Optimize by wrapping some of the code in a graph using TF function. self.agent.train = common.function(self.agent.train) # Reset the train step self.agent.train_step_counter.assign(0) for _ in range(epochs): #print("epoch: ", _) # Collect a few steps using collect_policy and save to the replay buffer. collect_data(self.train_env, self.agent.collect_policy, self.replay_buffer, collect_steps_per_episode) # Sample a batch of data from the buffer and update the agent's network. experience, _ = next(self.iterator) train_loss = self.agent.train(experience).loss step = self.agent.train_step_counter.numpy() if step % log_step == 0: print('step = {0}: loss = {1}'.format(step, train_loss)) if step % eval_step == 0: metrics = self.compute_metrics(X_train, y_train) print(metrics) def compute_metrics(self, X: np.ndarray, y_true: list) -> dict: """Compute Metrics for Evaluation""" # TODO: apply softmax layer for q logits? q, _ = self.agent._target_q_network (X, training=False) # y_scores = np.max(q.numpy(), axis=1) # predicted scores (Q-Values) y_pred = np.argmax(q.numpy(), axis=1) # predicted class label metrics = custom_metrics(y_true, y_pred) return metrics def evaluate(self, X: np.ndarray, y: list, X_train=None, y_train=None) -> dict: """ Evaluation of trained Q-network """ metrics = self.compute_metrics(X, y) print("evaluation: ", metrics) return metrics
def main(_): # Environment env_name = "Breakout-v4" train_num_parallel_environments = 5 max_steps_per_episode = 1000 # Replay buffer replay_buffer_capacity = 50000 init_replay_buffer = 500 # Driver collect_steps_per_iteration = 1 * train_num_parallel_environments # Training train_batch_size = 32 train_iterations = 100000 train_summary_interval = 200 train_checkpoint_interval = 200 # Evaluation eval_num_parallel_environments = 5 eval_summary_interval = 500 eval_num_episodes = 20 # File paths path = pathlib.Path(__file__) parent_dir = path.parent.resolve() folder_name = path.stem + time.strftime("_%Y%m%d_%H%M%S") train_checkpoint_dir = str(parent_dir / folder_name / "train_checkpoint") train_summary_dir = str(parent_dir / folder_name / "train_summary") eval_summary_dir = str(parent_dir / folder_name / "eval_summary") # Parallel training environment tf_env = TFPyEnvironment( ParallelPyEnvironment([ lambda: suite_atari.load( env_name, env_wrappers= [lambda env: TimeLimit(env, duration=max_steps_per_episode)], gym_env_wrappers=[AtariPreprocessing, FrameStack4], ) ] * train_num_parallel_environments)) tf_env.seed([42] * tf_env.batch_size) tf_env.reset() # Parallel evaluation environment eval_tf_env = TFPyEnvironment( ParallelPyEnvironment([ lambda: suite_atari.load( env_name, env_wrappers= [lambda env: TimeLimit(env, duration=max_steps_per_episode)], gym_env_wrappers=[AtariPreprocessing, FrameStack4], ) ] * eval_num_parallel_environments)) eval_tf_env.seed([42] * eval_tf_env.batch_size) eval_tf_env.reset() # Creating the Deep Q-Network preprocessing_layer = keras.layers.Lambda( lambda obs: tf.cast(obs, np.float32) / 255.) conv_layer_params = [(32, (8, 8), 4), (64, (4, 4), 2), (64, (3, 3), 1)] fc_layer_params = [512] q_net = QNetwork(tf_env.observation_spec(), tf_env.action_spec(), preprocessing_layers=preprocessing_layer, conv_layer_params=conv_layer_params, fc_layer_params=fc_layer_params) # Creating the DQN Agent optimizer = keras.optimizers.RMSprop(lr=2.5e-4, rho=0.95, momentum=0.0, epsilon=0.00001, centered=True) epsilon_fn = keras.optimizers.schedules.PolynomialDecay( initial_learning_rate=1.0, # initial ε decay_steps=2500000, end_learning_rate=0.01) # final ε global_step = tf.compat.v1.train.get_or_create_global_step() agent = DqnAgent( tf_env.time_step_spec(), tf_env.action_spec(), q_network=q_net, optimizer=optimizer, target_update_period=200, td_errors_loss_fn=keras.losses.Huber(reduction="none"), gamma=0.99, # discount factor train_step_counter=global_step, epsilon_greedy=lambda: epsilon_fn(global_step)) agent.initialize() # Creating the Replay Buffer replay_buffer = tf_uniform_replay_buffer.TFUniformReplayBuffer( data_spec=agent.collect_data_spec, batch_size=tf_env.batch_size, max_length=replay_buffer_capacity) # Observer: Replay Buffer Observer replay_buffer_observer = replay_buffer.add_batch # Observer: Training Metrics train_metrics = [ tf_metrics.NumberOfEpisodes(), tf_metrics.EnvironmentSteps(), tf_metrics.AverageReturnMetric(batch_size=tf_env.batch_size), tf_metrics.AverageEpisodeLengthMetric(batch_size=tf_env.batch_size), ] # Creating the Collect Driver collect_driver = DynamicStepDriver(tf_env, agent.collect_policy, observers=[replay_buffer_observer] + train_metrics, num_steps=collect_steps_per_iteration) # Initialize replay buffer initial_collect_policy = RandomTFPolicy(tf_env.time_step_spec(), tf_env.action_spec()) init_driver = DynamicStepDriver( tf_env, initial_collect_policy, observers=[replay_buffer_observer, ShowProgress()], num_steps=init_replay_buffer) final_time_step, final_policy_state = init_driver.run() # Creating the Dataset dataset = replay_buffer.as_dataset(sample_batch_size=train_batch_size, num_steps=2, num_parallel_calls=3).prefetch(3) # Optimize by wrapping some of the code in a graph using TF function. collect_driver.run = function(collect_driver.run) agent.train = function(agent.train) print("\n\n++++++++++++++++++++++++++++++++++\n") # Create checkpoint train_checkpointer = Checkpointer( ckpt_dir=train_checkpoint_dir, max_to_keep=1, agent=agent, # replay_buffer=replay_buffer, global_step=global_step, # metrics=metric_utils.MetricsGroup(train_metrics, 'train_metrics') ) # Restore checkpoint # train_checkpointer.initialize_or_restore() # Summary writers and metrics train_summary_writer = tf.summary.create_file_writer(train_summary_dir) eval_summary_writer = tf.summary.create_file_writer(eval_summary_dir) eval_metrics = [ tf_metrics.NumberOfEpisodes(), tf_metrics.EnvironmentSteps(), tf_metrics.AverageReturnMetric(batch_size=eval_tf_env.batch_size, buffer_size=eval_num_episodes), tf_metrics.AverageEpisodeLengthMetric( batch_size=eval_tf_env.batch_size, buffer_size=eval_num_episodes) ] # Create evaluate callback function eval_callback = evaluate(eval_metrics=eval_metrics, eval_tf_env=eval_tf_env, eval_policy=agent.policy, eval_num_episodes=eval_num_episodes, train_step=global_step, eval_summary_writer=eval_summary_writer) # Train agent train_agent(tf_env=tf_env, train_iterations=train_iterations, global_step=global_step, agent=agent, dataset=dataset, collect_driver=collect_driver, train_metrics=train_metrics, train_checkpointer=train_checkpointer, train_checkpoint_interval=train_checkpoint_interval, train_summary_writer=train_summary_writer, train_summary_interval=train_summary_interval, eval_summary_interval=eval_summary_interval, eval_callback=eval_callback) print("\n\n++++++++++ END OF TF_AGENTS RL TRAINING ++++++++++\n\n")
num_parallel_calls=3, sample_batch_size=BATCH_SIZE, num_steps=2).prefetch(3) iterator = iter(dataset) agent.train_step_counter.assign(0) avg_return = compute_avg_return(eval_env, agent.policy) returns = [avg_return] # Pre-populate replay buffer for _ in range(PRETRAIN_LEN): driver.run() # Train # Optimize agent.train = common.function(agent.train) start_time = time.time() for _ in range(NUM_ITERATIONS): driver.run() experience, unused_info = next(iterator) train_loss = agent.train(experience).loss step = agent.train_step_counter.numpy() if step % 200 == 0: print(f'Step {step}: loss = {train_loss}') if step % 1000 == 0: avg_return = compute_avg_return(eval_env, agent.policy) current_time = time.time() elapsed_time = current_time - start_time