Fix 7 bugs found in deep code review: models, env wrappers, train loop

[Copilot]

Deep review of the RL pipeline and model architectures surfaced several correctness bugs and a few code quality issues.

Critical bugs

• DartPolicy.value_net final nn.ReLU() — clamped all critic outputs to ≥ 0, preventing the value function from representing negative returns. Removed.

• AttentionMLP attention was a no-op — bmm(q[B,1,d], k[B,d,1]) produces a (B,1,1) scalar; softmax of a single value is always 1.0, so every attention weight was constant and attended_x == x. The Q/K layers had zero effect. Replaced with element-wise per-feature attention:

  # Before (broken): single scalar softmax, always 1.0
  scores = torch.bmm(q.unsqueeze(1), k.unsqueeze(2)) / sqrt(d)  # (B,1,1)
  weights = torch.softmax(scores, dim=-1)                        # always 1.0
  
  # After: per-feature weights that actually gate the input
  scores = (q * k) / sqrt(d)          # (B, num_features)
  weights = torch.softmax(scores, dim=-1)
  attended_x = weights * x

• CustomMLPExtractor.forward() ran shared_net twice — delegated to forward_actor then forward_critic, each independently running shared_net. Fixed by computing shared once in forward() and passing it directly to both heads.

• WarpFrameDict non-grayscale path returned HWC — the np.transpose((2,0,1)) converting back to CHW after resize was commented out, while observation_space declared CHW. Uncommented.

Minor bugs

• HybridMambaMLPExtractor dead gru_layers param — copied from GRUMlpExtractor, never used inside the Mamba _build_temporal_model. Removed.

Code quality

• isMLP() — replaced long or-chain with a frozenset for O(1) lookup and easier maintenance.
• train.py — removed deprecated mode='human' kwarg from env.render() (dropped in Gymnasium ≥ 0.26).