Module atari_cnn 

Nature-DQN-scale CNN policies (actor-critic + Q-network) for the Atari (ALE) workload. Burn-backend Nature-DQN-scale CNN policy for the Atari (ALE) workload (Epic #306, Phase 3 — issue #327).

Implements the classic Nature-DQN convolutional stack (Mnih et al., Human-level control through deep reinforcement learning, 2015) as two Burn modules that share the same conv trunk:

• crate::policy::atari_cnn::NatureDqnBurnPolicy — actor-critic variant (policy + value heads), consumed by crate::train::ppo::trainer::PPOTrainerBurn.
• crate::policy::atari_cnn::NatureDqnQNetwork — single-Q-head variant, consumed by crate::train::dqn::DQNTrainerBurn (with a copy_params_from target-net sync, mirroring crate::policy::q_network::QNetworkBurn).

Architecture

obs [B, 4, 84, 84]
  → conv1 (32 ch, 8x8, stride 4)  → ReLU   → [B, 32, 20, 20]
  → conv2 (64 ch, 4x4, stride 2)  → ReLU   → [B, 64,  9,  9]
  → conv3 (64 ch, 3x3, stride 1)  → ReLU   → [B, 64,  7,  7]
  → flatten (64*7*7 = 3136)
  → fc_common (3136 -> 512)       → ReLU
  → heads:
      actor-critic: policy_head (512 -> A) [logits], value_head (512 -> 1)
      q-network:    q_head      (512 -> A) [Q(s, a)]

Convolutions use Burn’s default PaddingConfig2d::Valid (no padding), matching the Nature-DQN spec; the spatial reductions are therefore 84 → 20 → 9 → 7, giving a cached flat_size of 64 * 7 * 7 = 3136.

Input contract

• Layout: NCHW [batch, channels, height, width] — same convention as crate::policy::snake_cnn::SnakeCnnBurnPolicy and Burn’s Conv2d.
• Channels: 4 (frame-stack dimension, produced by the preprocessor — not this module).
• Spatial: 84 × 84.
• Dtype: f32, pixel-scaled to 0.0–1.0 (uint8 ÷ 255). The network is scale-agnostic, but this is the expected convention.
• No batch-size constraint.

Trainer integration (closure-based, flat rollout buffers)

Both Burn trainers are closure-based, not trait-based; the only module bounds are AutodiffModule<B> + Clone, satisfied automatically by #[derive(Module, Debug)]. The rollout buffers hand the closure a flat observation tensor [B, C*H*W], so the closure must reshape to [B, C, H, W] before calling forward/evaluate_actions — the same pattern used by examples/games/snake/train_snake_multi_v2.rs (lines 239–253):

// PPO (actor-critic):
let evaluate_fn = |p: &NatureDqnBurnPolicy<B>, o_flat: Tensor<B, 2>, acts: Tensor<B, 1, Int>| {
    let b = o_flat.dims()[0];
    let o4 = o_flat.reshape([b, 4, 84, 84]);
    p.evaluate_actions(o4, acts) // (log_probs [B], entropy [B], values [B])
};

// DQN (Q-network):
let forward_fn = |q: &NatureDqnQNetwork<B>, o_flat: Tensor<B, 2>| {
    let b = o_flat.dims()[0];
    q.forward(o_flat.reshape([b, 4, 84, 84])) // Q-values [B, A]
};

Seeded initialization

Seeded construction (see crate::policy::atari_cnn::NatureDqnConfig) drives the three FC layers from deterministically-derived host-side RNG streams via the shared mlp.rs helpers (derive_layer_seed / seeded_layer_weights / linear_from_weights), so two constructions with the same seed produce bit-identical FC weights. Fixed per-variant layer indices:

• NatureDqnBurnPolicy: 0 = fc_common, 1 = policy_head, 2 = value_head
• NatureDqnQNetwork: 0 = fc_common, 1 = q_head

Conv layers are intentionally unseeded. Burn’s Conv2dConfig — like LinearConfig — exposes no seed parameter, so the seeded path cannot reach the convolutions. This is a deliberate, second-order concern: the conv parameters total ~78K versus ~1.6M for fc_common alone, so the FC layers dominate reproducibility. The seeded path therefore covers only the three FC layers; the unseeded (seed: None) path routes every layer through Burn’s stock Initializer.

Structs

NatureDqnBurnPolicy

Nature-DQN-scale actor-critic CNN policy on Burn.

NatureDqnBurnPolicyRecord

The record type for the module.

NatureDqnBurnPolicyRecordItem

The record item type for the module.

NatureDqnConfig

Configuration for the Nature-DQN policies.

NatureDqnQNetwork

Nature-DQN-scale Q-network CNN on Burn.

NatureDqnQNetworkRecord

The record type for the module.

NatureDqnQNetworkRecordItem

The record item type for the module.