DQN-2018-TAU/dqn_learn.py

"""
    This file is copied/apdated from https://github.com/berkeleydeeprlcourse/homework/tree/master/hw3
"""
import sys
import pickle
import numpy as np
from collections import namedtuple
from itertools import count
import random
import gym.spaces
from utils.saved_state import SavedState

import torch
import torch.optim as optim
import torch.nn as nn
import torch.nn.functional as F
import torch.autograd as autograd

from utils.replay_buffer import ReplayBuffer
from utils.gym import get_wrapper_by_name

USE_CUDA = torch.cuda.is_available()
print("USE_CUDA=", USE_CUDA)
print("tomer11")
dtype = torch.cuda.FloatTensor if USE_CUDA else torch.FloatTensor
longType = torch.cuda.LongTensor if USE_CUDA else torch.LongTensor

class Variable(autograd.Variable):
    def __init__(self, data, *args, **kwargs):
        if USE_CUDA:
            data = data.cuda()
        super(Variable, self).__init__(data, *args, **kwargs)

"""
    OptimizerSpec containing following attributes
        constructor: The optimizer constructor ex: RMSprop
        kwargs: {Dict} arguments for constructing optimizer
"""
OptimizerSpec = namedtuple("OptimizerSpec", ["constructor", "kwargs"])

def dqn_learing(
    env,
    q_func,
    optimizer_spec,
    exploration,
    stopping_criterion=None,
    replay_buffer_size=1000000,
    batch_size=32,
    gamma=0.99,
    learning_starts=50000,
    learning_freq=4,
    frame_history_len=4,
    target_update_freq=10000,
    save_path=None,
    save_freq=100000,
    **kwargs
    ):

    """Run Deep Q-learning algorithm.

    You can specify your own convnet using q_func.

    All schedules are w.r.t. total number of steps taken in the environment.

    Parameters
    ----------
    env: gym.Env
        gym environment to train on.
    q_func: function
        Model to use for computing the q function. It should accept the
        following named arguments:
            input_channel: int
                number of channel of input.
            num_actions: int
                number of actions
    optimizer_spec: OptimizerSpec
        Specifying the constructor and kwargs, as well as learning rate schedule
        for the optimizer
    exploration: Schedule (defined in utils.schedule)
        schedule for probability of chosing random action.
    stopping_criterion: (env) -> bool
        should return true when it's ok for the RL algorithm to stop.
        takes in env and the number of steps executed so far.
    replay_buffer_size: int
        How many memories to store in the replay buffer.
    batch_size: int
        How many transitions to sample each time experience is replayed.
    gamma: float
        Discount Factor
    learning_starts: int
        After how many environment steps to start replaying experiences
    learning_freq: int
        How many steps of environment to take between every experience replay
    frame_history_len: int
        How many past frames to include as input to the model.
    target_update_freq: int
        How many experience replay rounds (not steps!) to perform between
        each update to the target Q network
    """
    assert type(env.observation_space) == gym.spaces.Box
    assert type(env.action_space)      == gym.spaces.Discrete

    ###############
    # BUILD MODEL #
    ###############

    if len(env.observation_space.shape) == 1:
        # This means we are running on low-dimensional observations (e.g. RAM)
        input_arg = env.observation_space.shape[0]
    else:
        img_h, img_w, img_c = env.observation_space.shape
        input_arg = frame_history_len * img_c
    num_actions = env.action_space.n

    def to_pytorch(obs, type=dtype, normalize=True):
        t = torch.from_numpy(obs).type(type)
        if normalize:
            return t / 255.0
        else:
            return t

    def to_pytorch_var(x, grad=False, type=dtype, normalize=True):
        return Variable(to_pytorch(x, type=type, normalize=normalize), requires_grad=grad)

    # Construct an epilson greedy policy with given exploration schedule
    def select_epsilon_greedy_action(model, obs, t):
        sample = random.random()
        eps_threshold = exploration.value(t)
        if sample > eps_threshold:
            obs = to_pytorch(obs).unsqueeze(0)
            # Use volatile = True if variable is only used in inference mode, i.e. don’t save the history
            return model(Variable(obs, volatile=True)).data.max(1)[1].cpu()
        else:
            return torch.IntTensor([[random.randrange(num_actions)]])

    # Initialize target q function and q function, i.e. build the model.
    ######

    # YOUR CODE HERE
    print("Input and output size of network:")
    print(input_arg,num_actions)
    Q = q_func(input_arg,num_actions)
    Q_target = q_func(input_arg,num_actions)

    if USE_CUDA:
        Q = Q.cuda()
        Q_target = Q_target.cuda()

    def update_Q_target():
        print("Updating Q_target")
        Q_target.load_state_dict(Q.state_dict())

    update_Q_target()

    ######

    # Construct Q network optimizer function
    optimizer = optimizer_spec.constructor(Q.parameters(), **optimizer_spec.kwargs)
    #scheduler = optim.lr_scheduler.ExponentialLR(optimizer, gamma_scheduler)


    replay_buffer = ReplayBuffer(replay_buffer_size, frame_history_len)

    start_step = 0
    Statistic = {
        "mean_episode_rewards": [],
        "best_mean_episode_rewards": []
    }
    mean_episode_reward = -float('nan')
    best_mean_episode_reward = -float('inf')
    if save_path is not None:
        try:
            print("Trying to load state from ", save_path)
            with open(save_path, 'rb') as f:
                saved_state = pickle.load(f)
                start_step = saved_state.timestep
                Q.load_state_dict(saved_state.state_dict)
                Q_target.load_state_dict(saved_state.state_dict)
                Statistic = saved_state.stats
                mean_episode_reward = Statistic["mean_episode_rewards"][-1][1]
                best_mean_episode_reward = Statistic["best_mean_episode_rewards"][-1][1]
        except:
            print("Saved state doesn't exist yet")

    def save_state(t):
        """
        Saves the current stable network weights, together with the current time step and statistics, for resuming later
        """
        if save_path is not None:
            print("Saving state")
            with open(save_path, 'wb') as f:
                pickle.dump(SavedState(Q_target.state_dict(),t,Statistic), f, pickle.HIGHEST_PROTOCOL)

    ###############
    # RUN ENV     #
    ###############
    num_param_updates = 0
    last_obs = env.reset()
    LOG_EVERY_N_STEPS = 1000

    for t in count(start=start_step):
        ### 1. Check stopping criterion
        if stopping_criterion is not None and stopping_criterion(env):
            break

        ### 2. Step the env and store the transition
        # At this point, "last_obs" contains the latest observation that was
        # recorded from the simulator. Here, your code needs to store this
        # observation and its outcome (reward, next observation, etc.) into
        # the replay buffer while stepping the simulator forward one step.
        # At the end of this block of code, the simulator should have been
        # advanced one step, and the replay buffer should contain one more
        # transition.
        # Specifically, last_obs must point to the new latest observation.
        # Useful functions you'll need to call:
        # obs, reward, done, info = env.step(action)
        # this steps the environment forward one step
        # obs = env.reset()
        # this resets the environment if you reached an episode boundary.
        # Don't forget to call env.reset() to get a new observation if done
        # is true!!
        # Note that you cannot use "last_obs" directly as input
        # into your network, since it needs to be processed to include context
        # from previous frames. You should check out the replay buffer
        # implementation in dqn_utils.py to see what functionality the replay
        # buffer exposes. The replay buffer has a function called
        # encode_recent_observation that will take the latest observation
        # that you pushed into the buffer and compute the corresponding
        # input that should be given to a Q network by appending some
        # previous frames.
        # Don't forget to include epsilon greedy exploration!
        # And remember that the first time you enter this loop, the model
        # may not yet have been initialized (but of course, the first step
        # might as well be random, since you haven't trained your net...)
        #####

        # YOUR CODE HERE
        last_frame_idx = replay_buffer.store_frame(last_obs)
        enc_last_obs = replay_buffer.encode_recent_observation()
        action = select_epsilon_greedy_action(Q, enc_last_obs, t)
        new_frame, r, done, _ = env.step(action)
        replay_buffer.store_effect(last_frame_idx, action, r, done)
        if done:
           last_obs = env.reset()
        else:
           last_obs = new_frame

        #####

        # at this point, the environment should have been advanced one step (and
        # reset if done was true), and last_obs should point to the new latest
        # observation

        ### 3. Perform experience replay and train the network.
        # Note that this is only done if the replay buffer contains enough samples
        # for us to learn something useful -- until then, the model will not be
        # initialized and random actions should be taken
        if (t > learning_starts and
                t % learning_freq == 0 and
                replay_buffer.can_sample(batch_size)):
            # Here, you should perform training. Training consists of four steps:
            # 3.a: use the replay buffer to sample a batch of transitions (see the
            # replay buffer code for function definition, each batch that you sample
            # should consist of current observations, current actions, rewards,
            # next observations, and done indicator).
            # Note: Move the variables to the GPU if avialable
            # 3.b: fill in your own code to compute the Bellman error. This requires
            # evaluating the current and next Q-values and constructing the corresponding error.
            # Note: don't forget to clip the error between [-1,1], multiply is by -1 (since pytorch minimizes) and
            #       maskout post terminal status Q-values (see ReplayBuffer code).
            # 3.c: train the model. To do this, use the bellman error you calculated perviously.
            # Pytorch will differentiate this error for you, to backward the error use the following API:
            #       current.backward(d_error.data.unsqueeze(1))
            # Where "current" is the variable holding current Q Values and d_error is the clipped bellman error.
            # Your code should produce one scalar-valued tensor.
            # Note: don't forget to call optimizer.zero_grad() before the backward call and
            #       optimizer.step() after the backward call.
            # 3.d: periodically update the target network by loading the current Q network weights into the
            #      target_Q network. see state_dict() and load_state_dict() methods.
            #      you should update every target_update_freq steps, and you may find the
            #      variable num_param_updates useful for this (it was initialized to 0)
            #####

            # YOUR CODE HERE
            optimizer.zero_grad()
            # TODO: Add a learning rate scheduler step here

            obs_batch,act_batch,r_batch,next_obs_batch,done_mask = replay_buffer.sample(batch_size)

            Q_val_batch = Q(to_pytorch_var(obs_batch))
            Q_target_val_batch = Q_target(to_pytorch_var(next_obs_batch)).detach()
            # The following code will take only one cell from each vector of the Q_val_batch tensor.
            # Each vector corresponds to a single output of the Q net, and each cell corresponds to a single action.
            # This means we take only the cells of the actions that we actually took, since all others are irrelevant
            # when calculating the loss.
            act_batch_var = to_pytorch_var(act_batch, type=longType, normalize=False).unsqueeze(1)
            vals_of_actions_taken = Q_val_batch.gather(1, act_batch_var)
            Q_target_val_max, _ = Q_target_val_batch.max(1)
            Q_target_val_max = Q_target_val_max.unsqueeze(1)
            reverse_done_mask = 1 - to_pytorch_var(done_mask, normalize=False).unsqueeze(1)
            Q_target_masked_val_max = (reverse_done_mask * Q_target_val_max)
            Q_target_discounted = (gamma * Q_target_masked_val_max)
            r_batch_var = to_pytorch_var(r_batch, normalize=False).unsqueeze(1)
            Q_masked_target = r_batch_var + Q_target_discounted
            bellman_error = Q_masked_target - vals_of_actions_taken
            clipped_error = bellman_error.clamp(-1,1)
            vals_of_actions_taken.backward(-clipped_error)

            optimizer.step()
            #scheduler.step()

            num_param_updates += 1
            if num_param_updates % target_update_freq == 0:
                update_Q_target()
            #####

        ### 4. Log progress and keep track of statistics
        episode_rewards = get_wrapper_by_name(env, "Monitor").get_episode_rewards()
        if len(episode_rewards) > 0:
            mean_episode_reward = np.mean(episode_rewards[-100:])
        if len(episode_rewards) > 100:
            best_mean_episode_reward = max(best_mean_episode_reward, mean_episode_reward)

        Statistic["mean_episode_rewards"].append((t,mean_episode_reward))
        Statistic["best_mean_episode_rewards"].append((t,best_mean_episode_reward))

        if t % LOG_EVERY_N_STEPS == 0 and t > learning_starts and t > start_step:
            print("Timestep %d" % (t,))
            print("mean reward (100 episodes) %f" % mean_episode_reward)
            print("best mean reward %f" % best_mean_episode_reward)
            print("episodes %d" % len(episode_rewards))
            print("exploration %f" % exploration.value(t))
            sys.stdout.flush()

        if t % save_freq == 0 and t > start_step:
            save_state(t)