<jupyter_start><jupyter_text>DQN (Deep Q Network)In previous Labs, we have learned to use Pytorch to build deep learning models. In this lab, we will apply deep learning as function approximations in reinforcement learning. Reference: DQN https://arxiv.org/abs/1312.5602In tabular Q-learning, we maintain a table of state-action pairs $(s,a)$ and save one action value for each entry $Q(s,a),\forall (s,a)$. At each time step $t$, we are in state $s_t$, then we choose action based on $\epsilon-$greedy strategy. With prob $\epsilon$, choose action uniformly random; with prob $1-\epsilon$, choose action based on $$a_t = \arg\max_a Q(s_t,a)$$ We then get the instant reward $r_t$, update the Q-table using the following rule$$Q(s_t,a_t) \leftarrow (1-\alpha)Q(s_t,a_t) + \alpha (r_t + \max_a \gamma Q(s_{t+1},a))$$where $\alpha \in (0,1)$ is learning rate. The algorithm is shown to converge in tabular cases. However, in cases where we cannot keep a table for state and action, we need function approximation. Consider using neural network with parameter $\theta$, the network takes as input state $s$ and action $a$. (*there are alternative parameterizations here*). Let $Q_\theta(s,a)$ be the output of the network, to estimate the optimal action value function in state $s$ and take action $a$ (and follow optimal policy thereafter). $$Q_\theta(s,a) \approx Q^\ast(s,a)$$  Bellman optimality equationWe will use Bellman optimality equation to find $\theta$ such that the above approximation holds better. Recall that for optimal Q function $Q^\ast(s,a)$ the following holds for all $(s,a)$$$Q^\ast(s_t,a_t) = \mathbb{E}\big[r_t + \gamma \max_a Q^\ast(s_{t+1},a)\big]$$where the expectation is wrt the random reward $r_t$ and transition to the next state $s_{t+1}$. A natural objective to consider is $$\min_\theta\  (Q_\theta(s_t,a_t) -\mathbb{E}\big[r_t + \gamma  \max_a  Q_{\hat \theta}(s_{t+1},a)\big])^2$$at the current or previous $\hat \theta$.  Building the DQN model The first step is to build a neural network with parameters $\theta$ that predicts $Q_\theta(s,a)$ for any $(s,a)$. You can either build a network that * (in case of small number $K$ of discrete actions) takes as input a  representation of state $s$ and outputs a $K$-dimensional vector giving scores $Q(s,a), a=1,\ldots, K$ for all actionsor * takes as input a concatenated representation of state and action $(s,a)$ and output one dimensional score $Q_\theta(s,a)$,Below we have provided a skeleton code (incomplete) for defining and training the Q-function. **You need to fill in the DNN model definition and loss function definition**. Refer to regression lab (lab 2) for help.<jupyter_code># define neural net Q_\theta(s,a) as a class

class Qfunction(object):
    
    def __init__(self, obssize, actsize, lr):
        """
        obssize: dimension of state space
        actsize: dimension of action space
        sess: sess to execute this Qfunction
        optimizer: 
        """
        # DEFINE THE MODEL
        self.model = torch.nn.Sequential(
                    #TODO      
                    #input layer
                    torch.nn.Linear(obssize, 128),
                    torch.nn.ReLU(),
                    torch.nn.Linear(128, 256),
                    torch.nn.ReLU(),
                    torch.nn.Linear(256, actsize)
                )
        
        # DEFINE THE OPTIMIZER
        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=lr)
        
        # RECORD HYPER-PARAMS
        self.obssize = obssize
        self.actsize = actsize
        
    def _to_one_hot(self, y, num_classes):
        """
        convert an integer vector y into one-hot representation
        """
        scatter_dim = len(y.size())
        y_tensor = y.view(*y.size(), -1)
        zeros = torch.zeros(*y.size(), num_classes, dtype=y.dtype)
        return zeros.scatter(scatter_dim, y_tensor, 1)

    
    def compute_Qvalues(self, states, actions):
        """
        input: list of numsamples state-action pairs
        output: List of Q values for each input (s,a). The output will have size [numsamples, 1] 
        """
        #Below is example code when neural network is set to take as input state and output Q-value for all actions.
        #This will be different for neural network that takes as input a state-action pair
        
        states = torch.FloatTensor(states)
        q_preds = self.model(states)
        action_onehot = self._to_one_hot(actions, actsize)
        q_preds_selected = torch.sum(q_preds * action_onehot, axis=-1)

        return q_preds_selected
        
    def compute_maxQvalues(self, states):
        """
        input: a list of numsamples states 
        output: max_a Q(s,a) values for every input state s in states. The output will have size numsamples
        """
        #Below is example code when neural network is set to take as input state and output Q-value for all actions.
        #if the neural takes as input a state-action pair, then the code will need to loop over all actions to compute all values
        
        states = torch.FloatTensor(states)
        Qvalues = self.model(states).cpu().data.numpy()
        q_preds_greedy = np.max(Qvalues,1) 
        
        return q_preds_greedy
        
    def compute_argmaxQ(self, state):
        """
        input: one state s
        output: arg max_a Q(self.model(states).cpu().data.numpy()s,a) values for the input state s. The output will have size 1
        """
        #Below is example code when neural network is set to take as input state and output Q-value for all actions.
        #if the neural takes as input a state-action pair, then the code will need to loop over all actions to compute all values
        
        state = torch.FloatTensor(state)
        Qvalue = self.model(state).cpu().data.numpy()
        greedy_action = np.argmax(Qvalue.flatten())
        
        return greedy_action
        
        
    def train(self, states, actions, targets):
        """
        states: numpy array as input to compute loss (s)
        actions: numpy array as input to compute loss (a)
        targets: numpy array as input to compute loss (Q targets)
        """
        states = torch.FloatTensor(states)
        actions = torch.LongTensor(actions)
        targets = torch.FloatTensor(targets)
        
        # COMPUTE Q PREDICTIONS for all state-action pairs
        q_preds_selected = self.compute_Qvalues(states, actions)

                
        # LOSS
        #print(q_preds_selected.shape, targets.shape)
        loss = torch.mean((q_preds_selected - targets)**2)

        # BACKWARD PASS
        self.optimizer.zero_grad()
        loss.backward()

        # UPDATE
        self.optimizer.step()
            
        return loss.detach().cpu().data.numpy()<jupyter_output><empty_output><jupyter_text>At this point you can skip ahead to implementing the basic Q-learning that at every step $t$ in the environment* given state $s_t$, computes greedy actions from Q-values (using compute_argmaxQ function above) and uses $\epsilon$-greedy select an action $a_t$, * makes observation of reward $r_t$ and next state $s_{t+1}$* using compute_maxQvalues() function, computes target  $$r_t + \gamma \max_a Q_\theta(s_{t+1},a)$$and then retrains the Q-function using train() function above (with numsamples=1)However, for improved performance you may want to consider ideas like batch training (numsamples>1 is the batch size) with experience replay buffer and target-networks.  **Replay Buffer**Maintain a buffer $R$ to store trainsition tuples $(s_t,a_t,r_t,s_{t+1})$, when we minimize the Bellman error. When optimizing the Bellman error loss, we sample batches from the replay buffer and compute gradients for update on these batches. In particular, in each update, we sample $N$ tuples from buffer $(s_i,a_i,r_i,s_{i}') \sim R$ and then computetargets$$d_i=r_i + \max_a \gamma Q_{\theta}(s_i^\prime,a)$$for all $i$. Use the above training function train() with input as list $(s_i, a_i, d_i)_{i=1}^N$  to update parameters using backprop.**Target Network**Maintain a target network in addition to the original pricipal network. The target network is just a copy of the original network but the parameters are not updated by gradients. The target network $\theta^-$ is copied from the principal network every $\tau$ time steps. Target network is used to compute the targets for update$$d_i =  r_t + \gamma \max_a Q_{\theta^{-}}(s_{i}^\prime,a)$$the targets are used in the loss function to update the principal network parameters. This slowly updated target network ensures that the targets come from a relatively stationary distribution and hence stabilize learning. Hence several critical parts of the complete pseudocode for DQN is as follows:**Initialization.**principal network $Q_\theta(s,a)$, target network $Q_{\theta^{-}}(s,a)$. Replay buffer $R = \{\}$ (empty). **At each time step $t.$**The agent executes action using $\epsilon-$greedy based on the principal network $Q_\theta(s,a)$. To update $\theta$: sample $N$ tuples $(s_i,a_i,r_i,s_i^\prime) \sim R$, compute empirical loss $$\frac{1}{N} \sum_{i=1}^N (Q_\theta(s_i,a_i) - (r_i + \gamma \max_a Q_{\theta^{-}}(s_i^\prime,a))^2$$and update parameter $\theta$ using backprop (just take one gradient step).**Update target network.**Every $\tau$ time steps, update target network by copying $\theta_{\text{target}} \leftarrow \theta$.**Bellman target.**Above, we have defined the target values as being computed from a target net with parameter $\theta^-$ $$r_i + \gamma \max_a Q_{\theta^{-}}(s_i^\prime,a)$$It is worth thinking about what happens if we are at the end of an episode, that is, what if $s_i^\prime$ here is a terminal state. In this case, should the Bellman error be defined exactly the same as above? Do we need some modifications? Think carefully about this as this will greatly impact the algorithmic performance.  Implementation of replay buffer<jupyter_code># Implement replay buffer
import random
class ReplayBuffer(object):
    
    def __init__(self, maxlength):
        """
        maxlength: max number of tuples to store in the buffer
        if there are more tuples than maxlength, pop out the oldest tuples
        """
        self.buffer = deque()
        self.number = 0
        self.maxlength = maxlength
    
    def append(self, experience):
        """
        this function implements appending new experience tuple
        experience: a tuple of the form (s,a,r,s^\prime)
        """
        self.buffer.append(experience)
        self.number += 1
        
    def pop(self):
        """
        pop out the oldest tuples if self.number > self.maxlength
        """
        while self.number > self.maxlength:
            self.buffer.popleft()
            self.number -= 1
    
    def sample(self, batchsize):
        """
        this function samples 'batchsize' experience tuples
        batchsize: size of the minibatch to be sampled
        return: a list of tuples of form (s,a,r,s^\prime)
        """
        minibatch = random.sample(self.buffer,batchsize)
        return minibatch<jupyter_output><empty_output><jupyter_text>Code snippet for copying target networkYou may use th following to update target network i.e. to copy from principal network to target network. We need to use tensorflow scope to distinguish the computational graphs of target and principal networks. The following function builds a tensorflow operation that does the copying $\theta^- \leftarrow \theta$<jupyter_code>def run_target_update(Qprincipal, Qtarget):
    for v,v_ in zip(Qprincipal.model.parameters(), Qtarget.model.parameters()):
        v_.data.copy_(v.data)<jupyter_output><empty_output><jupyter_text>Main code for DQN Now that we have all the ingredients for DQN, we can write the main procedure to train DQN on a given environment. The implementation is straightforward if you follow the pseudocode pdf.<jupyter_code>%%wandb
#remove above line if you do not want to see inline plots from wandb

# hyper-parameters
lr = 1e-3  # learning rate for gradient update
batchsize = 128  # batchsize for buffer sampling
maxlength = 500  # max number of tuples held by buffer
envname = "CartPole-v0"  # environment name
tau = 100  # time steps for target update
episodes = 500  # number of episodes to run
initialsize = 100  # initial time steps before start training
epsilon = .2  # constant for exploration
gamma = .99  # discount

explore_rate = 1.0
min_explore_rate = 0.05  
max_explore_rate = 1.0
explore_decay_rate = 0.005


# initialize environment
env = gym.make(envname)
obssize = env.observation_space.low.size
actsize = env.action_space.n

# initialize Q-function networks (princpal and target)
Qprincipal = Qfunction(obssize, actsize, lr)
Qtarget = Qfunction(obssize, actsize, lr)

# initialization of graph and buffer
buffer = ReplayBuffer(maxlength)

# main iteration
rrecord = []
totalstep = 0

for episode in range(episodes):
    
    obs = env.reset()
    done = False
    rsum = 0 
    while not done:
        
        explore_rate = min_explore_rate + \
        (max_explore_rate - min_explore_rate) * np.exp(-explore_decay_rate*(totalstep))
        
        # epsilon greedy for exploration
        random_num = random.uniform(0, 1) 
        if random_num <= explore_rate:
            action = env.action_space.sample()
        else:
            action = Qprincipal.compute_argmaxQ(np.expand_dims(obs,0))
      
        newobs, r, done, _ = env.step(action)
        done_ = 1 if done else 0
        e = (obs, action, r, done_, newobs) 
        # e = (array([ 0.0297525 ,  0.0352823 , -0.03689473, -0.03157034]), 
        # 1, 1.0, 0, array([ 0.03045814,  0.23091338, -0.03752614, -0.33566185]))
     
        
        #IF NOT USING BUFFER:
        #use single sample (obs, action, r, done_, newobs) with Qtarget to compute target and train Qprincipal
        
        # ELSE IF USING REPLAY BUFFER
        # append experiences e to buffer
        
        buffer.append(e)
        while buffer.number > maxlength:
            buffer.pop()
            
    
        # every few episodes (decide the frequency) sample a minibatch from buffer
        # compute targets in batch using Qtarget and train Qprincipal 
        if episode >= initialsize:
            if buffer.number >= maxlength:
                batch_sample = buffer.sample(batchsize)

                b_states =  list(np.array(batch_sample)[:, 0])
                b_actions = list(np.array(batch_sample)[:, 1])
                b_reward  = list(np.array(batch_sample)[:, 2])
                b_next_state = list(np.array(batch_sample)[:, 4])
                b_done = list(np.array(batch_sample)[:, 3])
                
                next_state_q_vals = np.zeros(batchsize)
                next_state_q_vals = Qtarget.compute_maxQvalues(b_next_state)
                
                # deal with done states 
                for idx in range(batchsize):
                    if b_done[idx]:
                        next_state_q_vals[idx] = b_reward[idx]
                     
                #targets = Qtarget.compute_maxQvalues(b_states)
                targets = b_reward + gamma * next_state_q_vals

                Qprincipal.train(b_states, b_actions, targets)
       
        #UPDATE target network 
        #every tau steps update copy the principal network to the target network
        if totalstep % tau == 0:
            run_target_update(Qprincipal, Qtarget)
 
        # update
        totalstep += 1
        rsum += r
        obs = newobs 
        if done_: break 
    
        

    #The code below is for printing and debugging at the end of episode
    rrecord.append(rsum)
    
    # printing functions for debugging purposes. Feel free to add more
    if episode % 10 == 0:
        print('buffersize {}'.format(buffer.number))
        print('episode {} ave training returns {}'.format(episode, np.mean(rrecord[-10:])))
    
    #printing moving averages for smoothed visualization. 
    fixedWindow=100
    movingAverage=0
    if len(rrecord) >= fixedWindow:
        movingAverage=np.mean(rrecord[len(rrecord)-fixedWindow:len(rrecord)-1])
    wandb.log({ "training reward" : rsum, "train reward moving average" : movingAverage})<jupyter_output><empty_output><jupyter_text>Finally, we evaluate the performance of the trained agent. We will evaluate the performance of the greedy policy wrt learned Q-function. The evaluation will be run 10 times, each for eval_epsiodes and print out the average performance across these episodes. Please **do not** change the code below.<jupyter_code>### DO NOT CHANGE
def evaluate(Q, env, episodes):
    # main iteration
    score = 0.0
    for episode in range(episodes):
        
        obs = env.reset()
        done = False
        rsum = 0
        
        while not done:
            # always greedy
            action = Q.compute_argmaxQ(np.expand_dims(obs,0))
            

            # mdp stepping forward
            newobs, r, done, _ = env.step(action)

            # update data
            rsum += r
            obs = newobs        

        
        wandb.log({"eval reward" : rsum})
        score = score + rsum
    score = score/episodes
    
    return score
# DO NOT CHANGE CODE HERE
# after training, we will evaluate the performance of the agent
# on a target environment
env_test = gym.make(envname)
eval_episodes = 1000
score = evaluate(Qprincipal, env_test, eval_episodes)
wandb.run.summary["score"]=score 

print("eval performance of DQN agent: {}".format(score))
run.finish()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.09315286624203821, 'fraction_numerical': 0.010805277525022748, 'mean_word_length': 2.8005186946185434, 'pattern_counts': {'":': 0, '<': 22, '<?xml version=': 0, '>': 28, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '393306', 'license': 'MIT', 'n_tokens_mistral': 5219, 'n_tokens_neox': 4964, 'n_words': 1859, 'path': 'labs/lab3_DQN.ipynb', 'repo_name': 'EricaWei053/Reinforcement-Learning', 'size': '17584'}