diff --git a/scripts/env_utils.py b/scripts/env_utils.py
index 674a6c9..6d6e4c2 100644
--- a/scripts/env_utils.py
+++ b/scripts/env_utils.py
@@ -11,10 +11,20 @@
 from utils import resolve_clip_reward, resolve_sticky_action_settings
 
 
+_MLP_POLICIES = frozenset({
+    'MlpPolicy',
+    'MlpDropoutPolicy',
+    'CombinedPolicy',
+    'AttentionMLPPolicy',
+    'EntityAttentionPolicy',
+    'HockeyMultiHeadPolicy',
+    'HybridMambaPolicy',
+    'GRUMlpPolicy',
+})
+
+
 def isMLP(name):
-    return name == 'MlpPolicy' or name == 'MlpDropoutPolicy' or name == 'CombinedPolicy' \
-          or name == 'AttentionMLPPolicy' or name == 'EntityAttentionPolicy' or name == 'HockeyMultiHeadPolicy' \
-          or name == 'HybridMambaPolicy' or name == 'GRUMlpPolicy'
+    return name in _MLP_POLICIES
 
 
 def make_retro(
diff --git a/scripts/env_wrappers.py b/scripts/env_wrappers.py
index fae14ae..96d3c84 100644
--- a/scripts/env_wrappers.py
+++ b/scripts/env_wrappers.py
@@ -69,7 +69,7 @@ def observation(self, obs):
         else:
             img = cv2.resize(img, (self.width, self.height), interpolation=cv2.INTER_AREA)
             # Convert back to CHW
-            #img = np.transpose(img, (2, 0, 1))
+            img = np.transpose(img, (2, 0, 1))
 
         img = img.astype(np.uint8)
 
diff --git a/scripts/models.py b/scripts/models.py
index 4d17117..e0cb77f 100644
--- a/scripts/models.py
+++ b/scripts/models.py
@@ -278,21 +278,18 @@ def __init__(self, observation_space, features_dim=64, hidden_dim=128):
     def forward(self, x):
         if x.dim() == 1:
             x = x.unsqueeze(0)
-        # Self-attention
+        # Self-attention: compute element-wise attention weights across all features
         q = self.query(x)  # [batch_size, num_features]
         k = self.key(x)    # [batch_size, num_features]
 
-        # Compute attention scores
-        # q.unsqueeze(1): [batch_size, 1, num_features]
-        # k.unsqueeze(2): [batch_size, num_features, 1]
-        # bmm result: [batch_size, 1, 1]
-        attention_scores = torch.bmm(q.unsqueeze(1), k.unsqueeze(2)) / (self.num_features ** 0.5)
+        # Element-wise product gives a per-feature attention score, then
+        # softmax normalises across the feature dimension so that the weights
+        # sum to 1 and every feature dimension gets its own importance score.
+        attention_scores = (q * k) / (self.num_features ** 0.5)  # [batch_size, num_features]
+        weights = torch.softmax(attention_scores, dim=-1)         # [batch_size, num_features]
 
-        # Apply softmax - need to specify dimension
-        weights = torch.softmax(attention_scores, dim=-1)  # Apply along last dimension
-
-        # Apply attention weights
-        attended_x = weights.squeeze(-1) * x  # [batch_size, num_features]
+        # Scale input features by their attention weights
+        attended_x = weights * x  # [batch_size, num_features]
 
         # MLP
         return self.mlp(attended_x)
@@ -392,7 +389,6 @@ def __init__(
         d_state=64,
         d_conv=4,
         expand=2,
-        gru_layers=2,
     ):
         self.d_state = d_state
         self.d_conv = d_conv
@@ -575,12 +571,12 @@ def __init__(self, observation_space, action_space, lr_schedule, net_arch=None,
             nn.ReLU()
         )
 
-        # Value head should output a single value
+        # Value head should output a single value (no activation after final linear
+        # — value estimates can be negative)
         self.value_net = nn.Sequential(
             nn.Linear(latent_dim_vf, latent_dim_vf),
             nn.ReLU(),
-            nn.Linear(latent_dim_vf, 1),  # Output single value
-            nn.ReLU()
+            nn.Linear(latent_dim_vf, 1),
         )
 
     def forward(self, obs, deterministic=False):
@@ -824,8 +820,9 @@ def __init__(self, feature_dim, net_arch, dropout_prob=0.5):
         )
 
     def forward(self, features):
-        # Return policy and value features for the base policy's use
-        return self.forward_actor(features), self.forward_critic(features)
+        # Run shared layers once and reuse for both heads
+        shared = self.shared_net(features)
+        return self.policy_net(shared), self.value_net(shared)
 
     def forward_actor(self, features):
         shared = self.shared_net(features)
diff --git a/scripts/train.py b/scripts/train.py
index 984aaa5..eb4e505 100644
--- a/scripts/train.py
+++ b/scripts/train.py
@@ -102,7 +102,7 @@ def play(self, args, continuous=True):
 
         state = self.env.reset()
         while True:
-            self.env.render(mode='human')
+            self.env.render()
 
             p1_actions = self.p1_model.predict(state, deterministic=args.deterministic)