2
\$\begingroup\$

I trained a DQN agent using tensorflow and OpenAI gym Atari environment called PongNoFrameskip-v4, but this code should be compatible with any gym environment that returns the state as an RGB frame. This has been originally the work described in those papers:

Description

A DQN, or Deep Q-Network, approximates a state-value function in a Q-Learning framework with a neural network. In the Atari Games case, they take in several frames of the game as an input and output state values for each action as an output.

It is usually used in conjunction with Experience Replay, for storing the episode steps in memory for off-policy learning, where samples are drawn from the replay memory at random. Additionally, the Q-Network is usually optimized towards a frozen target network that is periodically updated with the latest weights every k steps (where k is a hyperparameter). The latter makes training more stable by preventing short-term oscillations from a moving target. The former tackles autocorrelation that would occur from on-line learning, and having a replay memory makes the problem more like a supervised learning problem.

Result

Here's a demonstration of a game using a model I trained and note that training the agent, might require 1 - 2 GPU hours, so for the sake of demonstration, I'll include the necessary files to run the code without needing to train at your end.

PC - Agent

Pong improved model

You can download the model archive here ~= 20 mb, which you need to unzip in the same directory as the the 2 modules below utils.py and dqn.py, then just run:

python3 dqn.py

Things I would like to improve:

  • Model convergence speed

  • An issue with tensorflow 2 that does not use the GPU during the training unless I do:

      tf.compat.v1.disable_eager_execution() 
    
  • Inconsistent convergence speed: sometimes the model converges very fast and the target reward (set at 19 by default) is reached in less than 500,000 steps / ~= 200-250 games and sometimes it may need up to 700,000 steps to reach the same reward without changing hyperparameters.

  • Anything else that can be improved

dqn.py

import os
from collections import deque
from time import perf_counter

import cv2
import gym
import numpy as np
import wandb
from tensorflow.keras.layers import Conv2D, Dense, Flatten, Input
from tensorflow.keras.models import Model
from tensorflow.keras.optimizers import Adam

from utils import AtariPreprocessor


class DQN:
    def __init__(
        self,
        env,
        replay_buffer_size=10000,
        batch_size=32,
        checkpoint=None,
        reward_buffer_size=100,
        epsilon_start=1.0,
        epsilon_end=0.01,
        frame_skips=4,
        resize_shape=(84, 84),
        state_buffer_size=2,
    ):
        """
        Initialize agent settings.
        Args:
            env: gym environment that returns states as atari frames.
            replay_buffer_size: Size of the replay buffer that will hold record the
                last n observations in the form of (state, action, reward, done, new state)
            batch_size: Training batch size.
            checkpoint: Path to .tf filename under which the trained model will be saved.
            reward_buffer_size: Size of the reward buffer that will hold the last n total
                rewards which will be used for calculating the mean reward.
            epsilon_start: Start value of epsilon that regulates exploration during training.
            epsilon_end: End value of epsilon which represents the minimum value of epsilon
                which will not be decayed further when reached.
            frame_skips: Number of frame skips to use per environment step.
            resize_shape: (m, n) dimensions for the frame preprocessor
            state_buffer_size: Size of the state buffer used by the frame preprocessor.
        """
        self.env = gym.make(env)
        self.env = AtariPreprocessor(
            self.env, frame_skips, resize_shape, state_buffer_size
        )
        self.input_shape = self.env.observation_space.shape
        self.main_model = self.create_model()
        self.target_model = self.create_model()
        self.buffer_size = replay_buffer_size
        self.buffer = deque(maxlen=replay_buffer_size)
        self.batch_size = batch_size
        self.checkpoint_path = checkpoint
        self.total_rewards = deque(maxlen=reward_buffer_size)
        self.best_reward = -float('inf')
        self.mean_reward = -float('inf')
        self.state = self.env.reset()
        self.steps = 0
        self.frame_speed = 0
        self.last_reset_frame = 0
        self.epsilon_start = self.epsilon = epsilon_start
        self.epsilon_end = epsilon_end
        self.games = 0

    def create_model(self):
        """
        Create model that will be used for the main and target models.
        Returns:
            None
        """
        x0 = Input(self.input_shape)
        x = Conv2D(32, 8, 4, activation='relu')(x0)
        x = Conv2D(64, 4, 2, activation='relu')(x)
        x = Conv2D(64, 3, 1, activation='relu')(x)
        x = Flatten()(x)
        x = Dense(512, 'relu')(x)
        x = Dense(self.env.action_space.n)(x)
        return Model(x0, x)

    def get_action(self, training=True):
        """
        Generate action following an epsilon-greedy policy.
        Args:
            training: If False, no use of randomness will apply.

        Returns:
            A random action or Q argmax.
        """
        if training and np.random.random() < self.epsilon:
            return self.env.action_space.sample()
        q_values = self.main_model.predict(np.expand_dims(self.state, 0))
        return np.argmax(q_values)

    def get_buffer_sample(self):
        """
        Get a sample of the replay buffer.
        Returns:
            A batch of observations in the form of
            [[states], [actions], [rewards], [dones], [next states]]
        """
        indices = np.random.choice(len(self.buffer), self.batch_size, replace=False)
        memories = [self.buffer[i] for i in indices]
        batch = [np.array(item) for item in zip(*memories)]
        return batch

    def update(self, batch, gamma):
        """
        Update gradients on a given a batch.
        Args:
            batch: A batch of observations in the form of
                [[states], [actions], [rewards], [dones], [next states]]
            gamma: Discount factor.

        Returns:
            None
        """
        states, actions, rewards, dones, new_states = batch
        q_states = self.main_model.predict(states)
        new_state_values = self.target_model.predict(new_states).max(1)
        new_state_values[dones] = 0
        target_values = np.copy(q_states)
        target_values[np.arange(self.batch_size), actions] = (
            new_state_values * gamma + rewards
        )
        self.main_model.fit(states, target_values, verbose=0)

    def checkpoint(self):
        """
        Save model weights if improved.
        Returns:
            None
        """
        if self.best_reward < self.mean_reward:
            print(f'Best reward updated: {self.best_reward} -> {self.mean_reward}')
            if self.checkpoint_path:
                self.main_model.save_weights(self.checkpoint_path)
        self.best_reward = max(self.mean_reward, self.best_reward)

    def display_metrics(self):
        """
        Display progress metrics to the console.
        Returns:
            None
        """
        display_titles = (
            'frame',
            'games',
            'mean reward',
            'best_reward',
            'epsilon',
            'speed',
        )
        display_values = (
            self.steps,
            self.games,
            self.mean_reward,
            self.best_reward,
            np.around(self.epsilon, 2),
            f'{round(self.frame_speed)} steps/s',
        )
        display = (
            f'{title}: {value}' for title, value in zip(display_titles, display_values)
        )
        print(', '.join(display))

    def update_metrics(self, episode_reward, start_time):
        """
        Update progress metrics.
        Args:
            episode_reward: Total reward per a single episode (game).
            start_time: Episode start time, used for calculating fps.

        Returns:
            None
        """
        self.games += 1
        self.checkpoint()
        self.total_rewards.append(episode_reward)
        self.frame_speed = (self.steps - self.last_reset_frame) / (
            perf_counter() - start_time
        )
        self.last_reset_frame = self.steps
        self.mean_reward = np.around(np.mean(self.total_rewards), 2)
        self.display_metrics()

    def fit(
        self,
        decay_n_steps=150000,
        learning_rate=1e-4,
        gamma=0.99,
        update_target_steps=1000,
        monitor_session=None,
        weights=None,
        max_steps=None,
        target_reward=19,
    ):
        """
        Train agent on a supported environment
        Args:
            decay_n_steps: Maximum steps that determine epsilon decay rate.
            learning_rate: Model learning rate shared by both main and target networks.
            gamma: Discount factor used for gradient updates.
            update_target_steps: Update target model every n steps.
            monitor_session: Session name to use for monitoring the training with wandb.
            weights: Path to .tf trained model weights to continue training.
            max_steps: Maximum number of steps, if reached the training will stop.
            target_reward: Target reward, if achieved, the training will stop

        Returns:
            None
        """
        if monitor_session:
            wandb.init(name=monitor_session)
        episode_reward = 0
        start_time = perf_counter()
        optimizer = Adam(learning_rate)
        if weights:
            self.main_model.load_weights(weights)
            self.target_model.load_weights(weights)
        self.main_model.compile(optimizer, loss='mse')
        self.target_model.compile(optimizer, loss='mse')
        while True:
            self.steps += 1
            self.epsilon = max(
                self.epsilon_end, self.epsilon_start - self.steps / decay_n_steps
            )
            action = self.get_action()
            new_state, reward, done, info = self.env.step(action)
            episode_reward += reward
            self.buffer.append((self.state, action, reward, done, new_state))
            self.state = new_state
            if done:
                if self.mean_reward >= target_reward:
                    print(f'Reward achieved in {self.steps} steps!')
                    break
                if max_steps and self.steps >= max_steps:
                    print(f'Maximum steps exceeded')
                    break
                self.update_metrics(episode_reward, start_time)
                start_time = perf_counter()
                episode_reward = 0
                self.state = self.env.reset()
            if len(self.buffer) < self.buffer_size:
                continue
            batch = self.get_buffer_sample()
            self.update(batch, gamma)
            if self.steps % update_target_steps == 0:
                self.target_model.set_weights(self.main_model.get_weights())

    def play(self, weights=None, video_dir=None, render=False, frame_dir=None):
        """
        Play and display a game.
        Args:
            weights: Path to trained weights, if not specified, the most recent
                model weights will be used.
            video_dir: Path to directory to save the resulting game video.
            render: If True, the game will be displayed.
            frame_dir: Path to directory to save game frames.

        Returns:
            None
        """
        if weights:
            self.main_model.load_weights(weights)
        if video_dir:
            self.env = gym.wrappers.Monitor(self.env, video_dir)
        self.state = self.env.reset()
        steps = 0
        for dir_name in (video_dir, frame_dir):
            os.makedirs(dir_name or '.', exist_ok=True)
        while True:
            if render:
                self.env.render()
            if frame_dir:
                frame = self.env.render(mode='rgb_array')
                cv2.imwrite(os.path.join(frame_dir, f'{steps:05d}.jpg'), frame)
            action = self.get_action(False)
            self.state, reward, done, info = self.env.step(action)
            if done:
                break
            steps += 1


if __name__ == '__main__':
    agn = DQN('PongNoFrameskip-v4')
    agn.play('model/pong_test.tf', render=True)

utils.py

from collections import deque

import cv2
import gym
import numpy as np


class AtariPreprocessor(gym.Wrapper):
    """
    gym wrapper for preprocessing atari frames.
    """

    def __init__(self, env, frame_skips=4, resize_shape=(84, 84), state_buffer_size=2):
        """
        Initialize preprocessing settings.
        Args:
            env: gym environment that returns states as atari frames.
            frame_skips: Number of frame skips to use per environment step.
            resize_shape: (m, n) output frame size.
            state_buffer_size: State buffer for max pooling.
        """
        super(AtariPreprocessor, self).__init__(env)
        self.skips = frame_skips
        self.frame_shape = resize_shape
        self.observation_space.shape = (*resize_shape, 1)
        self.observation_buffer = deque(maxlen=state_buffer_size)

    def process_frame(self, frame):
        """
        Resize and convert atari frame to grayscale.
        Args:
            frame: Image as numpy.ndarray

        Returns:
            Processed frame.
        """
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        frame = cv2.resize(frame, self.frame_shape) / 255
        return np.expand_dims(frame, -1)

    def step(self, action: int):
        """
        Step respective to self.skips.
        Args:
            action: Action supported by self.env

        Returns:
            (state, reward, done, info)
        """
        total_reward = 0
        state, done, info = 3 * [None]
        for _ in range(self.skips):
            state, reward, done, info = self.env.step(action)
            total_reward += reward
            self.observation_buffer.append(state)
            if done:
                break
        max_frame = np.max(np.stack(self.observation_buffer), axis=0)
        return self.process_frame(max_frame), total_reward, done, info

    def reset(self, **kwargs):
        """
        Reset self.env
        Args:
            **kwargs: kwargs passed to self.env.reset()

        Returns:
            Processed atari frame.
        """
        self.observation_buffer.clear()
        observation = self.env.reset(**kwargs)
        self.observation_buffer.append(observation)
        return self.process_frame(observation)

References:

\$\endgroup\$

0

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.