updates

hw2/model.py

@@ -2,6 +2,10 @@ from hydra.utils import instantiate
 import torch
 import torch.nn as nn
 
+# additional imports for extra credit
+import math
+import torch.nn.functional as F
+
 
 class QNetwork(nn.Module):
     def __init__(self, state_size, action_size, hidden_size, activation):
@@ -49,5 +53,55 @@ class DuelingQNetwork(nn.Module):
         ############################
         return Qs
 
+# Extra credit: implementing Noisy DQN
+class NoisyLinear(nn.Linear):
+    
+    # code reference from:
+    # (1) https://github.com/PacktPublishing/Deep-Reinforcement-Learning-Hands-On/blob/baa9d013596ea8ea8ed6826b9de6679d98b897ca/Chapter07/lib/dqn_model.py#L9
+    # (2) https://github.com/thomashirtz/noisy-networks/blob/main/noisynetworks.py
+    
+    def __init__(self, in_features, out_features, sigma_init=0.5, bias=True):
+        super().__init__(in_features, out_features, bias=bias)
+        # assume noise is gaussian, set sigma as learnable parameters
+        self.sigma_weight = nn.Parameter(torch.full((out_features, in_features), sigma_init))
+        self.register_buffer('epsilon_weight', torch.full((out_features, in_features), sigma_init))
+        if bias:
+            self.sigma_bias = nn.Parameter(torch.full((out_features,), sigma_init))
+            self.register_buffer('epsilon_bias', torch.full((out_features,), sigma_init))
 
-    
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        """
+        Reset the weights and bias of the noisy linear layer to a uniform distribution with std dev of sqrt(3 / in_features)
+        """
+        std = math.sqrt(3 / self.in_features)
+        self.weight.data.uniform_(-std, std)
+        self.bias.data.uniform_(-std, std)
+    
+    def forward(self, input):
+        """
+        Forward pass of noisy linear layer, adding gaussian noise to the weight and bias
+        """
+        self.epsilon_weight.normal_()
+        weight = self.weight + self.sigma_weight * self.epsilon_weight.data
+        bias = self.bias
+        if bias is not None:
+            self.epsilon_bias.normal_()
+            bias = bias + self.sigma_bias * self.epsilon_bias.data
+        return F.linear(input, weight, bias)
+
+class NoisyQNetwork(nn.Module):
+    def __init__(self, state_size, action_size, hidden_size, activation, sigma_init=0.5):
+        super(NoisyQNetwork, self).__init__()
+        self.q_head = nn.Sequential(
+            NoisyLinear(state_size, hidden_size, sigma_init=sigma_init),
+            instantiate(activation),
+            NoisyLinear(hidden_size, hidden_size, sigma_init=sigma_init),
+            instantiate(activation),
+            NoisyLinear(hidden_size, action_size, sigma_init=sigma_init)
+        )
+
+    def forward(self, state):
+        Qs = self.q_head(state)
+        return Qs