Transformer-augmented reinforcement learning in the Unity Food Collector environment

A course research project (USF, spring 2026) that asked one question: does bolting a transformer onto a model-free reinforcement-learning agent beat the plain multilayer-perceptron (MLP) version of the same algorithm? A three-person team split the question by where the transformer sits. I built the PPO variant with a transformer actor and a separate MLP critic, and I also ran the memoryless PPO baseline everyone was measured against. One teammate, Zaima, built a PPO variant with a transformer encoder shared by actor and critic; the other, Raniero, built the SAC baseline and a SAC variant with a transformer critic. All five models were evaluated on the same Unity ML-Agents Food Collector binary under the same protocol.

The short version: on a held-out evaluation my transformer-actor agent scored a mean episodic return of 61.84 against 22.33 for the PPO baseline and 49.54 for the SAC baseline. The two designs that gave the critic temporal context did better still (85.33 and 86.95). Most of the work, and most of what I learned, was in getting a 38M-parameter policy to train stably under PPO at all.

Code and report: repository, my final submission folder, and the final report (PDF).

Problem

The environment is Unity ML-Agents Food Collector: several agents share an arena, earn +1 for each green food pellet collected and −1 for each red one, and try to maximise the discounted return at γ = 0.99. Each agent sees the world through an ego-centric grid sensor, a 5 × 40 × 40 tensor whose five channels are semantic masks for food, poison, walls, other active agents, and frozen agents. Flattened, that is an 8,000-dimensional observation per step. Each of the 4 parallel arenas holds 3 agents, so every environment step yields 12 observations. Episodes truncate after 1,000 steps and the wrapper handles respawns.

The action space is hybrid: three continuous controls in [−1, 1] (forward/back, strafe, rotate) plus an optional binary laser branch that freezes other agents. The headline run enables the laser so the comparison against my teammates' models, which all use the hybrid space, is apples to apples; every intra-project ablation used the no-laser sibling.

The hypothesis going in was the textbook one. Food appears and disappears, agents move, and a single frame does not tell you where a pellet was two seconds ago. A transformer encoder over a sliding window of past frames, with positional encoding to preserve their order, should let the policy learn a content-dependent weighting over its recent history. An MLP has neither the memory nor the sense of when something was seen.

HYPOTHESIS: Self-attention over a window of past observations will recover hidden state that a memoryless PPO policy cannot, and raise episodic return.

Constraints

• No dataset, only rollouts. PPO is on-policy: each iteration collects 10,240 transitions across the parallel agents, updates, and throws the buffer away. There is no train/validation/test split, so evaluation had to be defined three ways: the smoothed training reward curve, a held-out run of the peak checkpoint (20 episodes, seed 42, deterministic mean actions, gradients off), and the cross-team comparison on that same protocol.
• Tiny batches during collection. During rollout the policy runs one step at a time per agent, so a forward pass sees a batch of 3. That is far too small for BatchNorm statistics, which mattered once the visual encoder arrived.
• The Unity binary is not redistributable. Reproduction requires building the Food Collector scene from ML-Agents 1.1.x or supplying a binary compiled for three agents. The repo pins everything else: Python 3.10.12, the wrapper, and one YAML config per experiment.
• Single GPU, single seed. Every run in the sequence, including the headline, used seed 42. One replication under seed 43 (v5.4) is reported honestly below; a proper multi-seed study did not fit the timeline.
• Shared measurement, separate code. The three architectures share the environment wrapper and the evaluation metrics but not the training code, so cross-team differences also carry algorithm, encoder, and entropy-regularisation differences.

Approach

The comparison is controlled where it can be. The baseline and the transformer variant share the same action parameterisation (diagonal Gaussian with a learned per-dimension log σ), the same GAE advantage estimator (γ = 0.99, λ = 0.95), the same clipped surrogate with ε = 0.2, and the same value and entropy losses. Any reward delta between them is meant to come from the architecture, not from a different RL algorithm.

The project ran as a numbered sequence of experiments, each with its own frozen YAML, and four of them define the story:

1. baseline_001: PPO with a two-hidden-layer MLP actor-critic (256 units), no memory, no CNN.
2. v3: transformer actor over a 32-frame window, the 8,000-dim frame projected by a single linear layer, no CNN.
3. v5.1-rerun: the same transformer actor with the window collapsed to a single frame, trained under a cosine learning-rate schedule with warmup.
4. v6.2 (headline): v5.1 plus a CoordConv-ResNet visual encoder in front of the transformer.

Between those milestones sit the runs that failed and taught something: v2 (dropout inside the encoder), v4 (a small plain CNN), v5.0 and v5.3 (linear schedule; sequence length 3), v5.4 and v5.5 (seed and head/layer sweeps), and v6.1 (the CoordConv encoder with hyperparameters ported unchanged from v5).

Everything is PyTorch. Rollout buffers, GAE, and PPO minibatching all live on CUDA tensors so no host-device copies happen between collection and update. Metrics go to Aim: mean, min, max, and std of episode reward, approximate KL, clip fraction, gradient norm, entropy, explained variance, learning rate, policy and value loss, and the per-dimension action standard deviations. The headline run is Aim hash c3cf3c51.

Architecture

The v6.2 policy has 37,761,031 parameters: roughly 31.1M in the visual encoder and 6.7M in the transformer body and heads. The matched no-CNN v5 policy has 2.6M, so the 14.5× growth is almost entirely the encoder. It has four parts.

Visual encoder

Each frame is reshaped to (5, 40, 40) and passed through a CoordConv-ResNet. A CoordConv layer prepends two constant channels holding the normalised x and y coordinate of every cell before an ordinary convolution, which breaks translation equivariance on purpose: in Food Collector the four walls and corners are not interchangeable, and a plain CNN cannot tell south-west from north-east. Four residual blocks (CoordConv → BatchNorm → ReLU → CoordConv → BatchNorm, plus a 1 × 1 projected skip where shapes change) step the map through (5, 40, 40) → (64, 40, 40) → (128, 20, 20) → (256, 10, 10) → (512, 10, 10). A flatten and a linear layer compress the 51,200-dim feature map to a 512-dim per-frame embedding. The block design was ported from Zaima's view processor so that the cross-team comparison isolates the transformer placement rather than the input pipeline.

Temporal encoder

The trainer keeps a sliding window of the last L encoded frames per agent. At episode start the window is filled with L copies of the first frame, so the transformer always sees a full sequence and never needs a padding mask; on termination it is reset and refilled from the new starting frame, which is the leakage guard between episodes. Each embedding is projected to d_model = 512, fixed sinusoidal positional encodings are added, and two standard TransformerEncoderLayer blocks (2 heads, feed-forward width 1,024, dropout 0) process the sequence. The outputs are mean-pooled over the sequence dimension into a 512-dim context vector.

The headline uses L = 1. That is a finding about the task, not a shortcut, and it is discussed under results.

Action heads

The pooled context feeds two heads in parallel. The continuous head is a two-layer MLP (512 → 256 → 3) producing the Gaussian mean μ over forward, strafe, and rotate, paired with a learned log σ vector clamped to [−5, 2]; sampled actions are tanh-squashed and hard-clamped to [−1, 1] before they reach Unity. The laser head is a single 512 → 2 linear layer emitting categorical logits over off/on. The heads are independent given the context, and PPO trains on the joint log-probability log π(a | s) = log π_cont(a_c | s) + log π_disc(a_d | s). The laser branch adds only 1,026 parameters.

A deliberately asymmetric critic

The critic does not read the encoder. It is a small MLP (8,000 → 256 → 1) over the raw flattened current frame. Early in training the CoordConv-ResNet's weights are random and its outputs are noise; feeding that noise to the value function poisoned its early targets and corrupted the GAE advantages in earlier iterations. Decoupling the critic removed that failure mode at the cost of one redundant input pathway.

DECISION: Actor on encoded, attended features; critic on the raw frame. Random CNN features are worse than no features for a value function that has to be useful from step zero.

Training settings for the headline run

• PPO clip 0.2, γ 0.99, GAE λ 0.95, value-loss coefficient 0.5, entropy coefficient 0.002, max gradient norm 0.5.
• Adam with a cosine schedule from a peak of 5 × 10⁻⁵ to 5 × 10⁻⁶ and a 10,000-step linear warmup from zero.
• 10,240 transitions per rollout, 5 epochs per rollout, minibatch 512, advantages normalised per minibatch.
• Guardrails: log σ clamped to [−5, 2]; the log importance ratio clamped to [−10, 10] before exponentiation; KL early-stop disabled (it was used at a 0.05 target in v3 through v5).
• BatchNorm handling: policy.eval() during rollout collection so the 3-sample batches use the running averages, policy.train() during the 512-sample PPO updates, which are what actually move those averages.
• Seed 42 for Python, NumPy, PyTorch (CPU and CUDA), and the Unity side channel.

The baseline differs in the ways you would expect for a first run: a constant learning rate of 1 × 10⁻⁴, 3 epochs, a 9,000-transition buffer, entropy 0.005, and none of the clamps.

What broke

The failure catalogue is the most useful part of the report, so it belongs here.

Baseline σ-collapse (baseline_001). Training reward climbed to 12.2 around step 504k and then the run destroyed itself. The learned log σ drifted steadily more negative, entropy fell from 4.25 at initialisation to 3.0 at the peak and kept falling, and once the Gaussian became a narrow spike, actions sampled under the old policy landed in the tails of the new one. PPO's importance ratio π_new / π_old overflowed and the policy loss hit 7.5 × 10²⁴ at step 828k. Three guardrails were added afterwards: the log σ clamp, the log-ratio clamp, and a KL early-stop. They prevented every later explosion but did not recover the run, which plateaued at 7.3.

Transformer KL stall (v3). The first transformer configuration at a learning rate of 1 × 10⁻⁴ peaked at 19.1 near step 348k, then froze. The KL early-stop target was set too tight at 0.02; from step 562k it fired after every single minibatch, epochs completed per rollout fell from 5 to 1, the clip fraction approached 1.0, and the policy loss saturated at zero because every sample sat outside the clip window where the gradient is zero. The actor stopped learning while the critic kept training. Raising the target to 0.05 unstuck it too late; halving the learning rate to 5 × 10⁻⁵ in v4 restored full epochs and reward growth. I later dropped the early-stop entirely once the cosine schedule removed the reason it existed.

LESSON: A KL early-stop that fires on every minibatch is not a safety net. It is a frozen actor with a healthy-looking loss curve. Log the per-update KL against the target so you can see it happening.

Trust-region collapse on the scaled-up network (v6.1). Porting the CoordConv-ResNet with v5's hyperparameters unchanged took the parameter count from 2.6M to 37.8M. The run peaked at 14.80 at step 133k, then approximate KL jumped from 0.08 to about 125 (roughly 1,500× the old threshold), the clip fraction reached 0.91, the gradient norm peaked at 1,649, and reward fell into the 1 to 3 range. It never recovered and was killed at 327k. Three things compounded: random CNN initialisation, BatchNorm running statistics that need warm-up, and uninformative early features feeding the transformer and therefore the critic. The fix in v6.2 was the 10,000-step warmup plus the lower peak rate and cosine decay. The warmup was the decisive change.

Dropout noise (v2). Dropout inside the encoder made approximate KL noisy because collection ran in eval mode and updates in train mode, producing log-probability differences unrelated to any real policy change. Removing it in v3 eliminated the noise, which is why the headline runs with dropout 0.

Two behavioural failures were seen in rollout video but are not yet quantified: the agent barrels into red pellets when they cluster next to green ones, treating the cluster as a reward gradient, and it occasionally pins itself against a wall when boxed in by geometry and a frozen partner, oscillating until the episode times out.

Results

The run sequence

Peak smoothed training reward, single seed:

• baseline_001, PPO + MLP: 12.20
• v3, transformer actor, 32-frame window, no CNN: 19.10
• v5.0, transformer actor, 1-frame window, linear schedule: 28.60
• v5.1-rerun, same, cosine schedule with warmup: 43.30
• v6.1, CoordConv-ResNet with naive hyperparameters: 14.80 at the collapse step
• v6.2 (headline), CoordConv-ResNet with cosine and warmup: 53.20 at step 1,320,444

The headline run continued to step 1,760,592 without exceeding that peak and was stopped manually. Against the baseline that is a 4.4× improvement; against the architecture-matched no-CNN transformer it is 23%, and that last contrast holds the schedule, warmup, critic, and seed constant while changing only the visual front-end.

RESULT: The visual encoder alone was worth 23% over the best no-CNN transformer, but only after the learning-rate schedule made a 38M-parameter policy trainable at all. Either piece without the other underperformed: the CNN with the naive schedule collapsed, and the schedule without the CNN topped out at 43.30.

Temporal memory did not help

This is the negative result, and the project reports it rather than burying it. Shrinking the attention window from 32 frames to 1 raised reward, from 19.10 (v3) to 28.60 (v5.0) and then 43.30 under the corrected schedule. A window of 3 (v5.3) fell to 18.70. Two explanations fit the data and cannot be separated by it: the 5 × 40 × 40 grid sensor covers so much of the arena that there is little hidden state for attention to recover, or the 32× larger input plus the positional table expands the optimisation surface beyond what the budget can navigate. The report leans on the first because of the sensor's geometric coverage. Either way, on this task the transformer earned its keep as a position-aware per-frame feature learner, not as memory.

RESULT: On a near-fully-observable grid sensor, a sequence length of 1 beat a sequence length of 32. The original hypothesis, as stated, was not supported.

Sample efficiency

v6.2 first crosses a mean reward of 10 at step 112,596, 20 at 204,720, 30 at 409,440, 40 at 481,092, and the v5.1-rerun peak of 43.30 at 552,744, against v5.1-rerun's own peak near 2.4M steps. The same no-CNN ceiling is reached roughly 4× faster once the encoder and schedule are right. The curve trails v5.1 for the first 200k steps or so, the warmup and stabilisation regime, then pulls decisively ahead once CNN features start informing the actor.

Seed variance

Only one replication exists: v5.4 repeated v5.0's configuration under seed 43 and peaked at 20.70 against 28.60. That spread is large enough that the single-seed headline numbers should be read as one draw, not an expectation.

Held-out comparison across the team

Each model's best checkpoint was loaded in inference mode and run for 20 episodes under seed 42 in the same binary. Three metrics: J_mean averages reward over all agents and episodes; J_min averages the worst agent per episode, a robustness measure; J_max averages the best agent per episode, a ceiling measure.

• This work, PPO transformer actor + MLP critic: J_mean 61.84, J_min 30.60, J_max 80.80
• PPO baseline (MLP): 22.33, 13.00, 37.00
• SAC baseline: 49.54, 30.45, 67.20
• PPO shared transformer encoder (Zaima): 85.33, 72.45, 190.66
• SAC with transformer critic (Raniero): 86.95, 75.05, 100.45

My design beats both baselines on every metric. The two teammates' designs beat it on J_mean and beat it decisively on J_min, and that J_min row is the most interesting number in the project. The two architectures whose critic sees a single frame, my raw-observation MLP critic and the SAC baseline's one-step Q-network, cluster at a worst-agent score of about 30. The two whose critic gets a temporal window cluster at about 73. A plausible mechanism is credit assignment: a critic with history can attribute a delayed pellet to the actions that led to it instead of relying on a one-step bootstrap to chain the credit. Confirming that needs a run that holds the actor and algorithm fixed and varies only the critic's context.

The 190.66 J_max for the shared-encoder PPO is more than double its own mean, which suggests one agent per episode collecting most of the reward, the predator dynamic the laser enables. The SAC transformer critic's tighter gap (100.45 versus 86.95) at a comparable mean points to a more uniformly competent population.

The other two placements

Shared transformer encoder for actor and critic (PPO, Zaima). The same CoordConv-ResNet front-end, but the encoded frame is kept as a 10 × 10 grid of tokens rather than flattened. A spatial attention block turns each frame's 100 tokens into a single class-token embedding, and a temporal attention block over a 5-frame history turns those into one history representation. Actor and critic both read that shared context, with a two-layer MLP critic head. Her branch also carries five ablations: no temporal transformer, no spatial transformer, class tokens only, mean-pooled spatial tokens, and no positional tokens. Branch: zaima-trasnformer-ppo.

Transformer critic (SAC, Raniero). Both SAC entries share one hybrid actor: a plain CNN backbone, a Gaussian movement head, a Gumbel-Softmax laser head, and separately learned entropy temperatures for the continuous and discrete branches. They differ only in the critic. The baseline uses twin one-step Q-networks over CNN features concatenated with both actions; the variant replaces that with a causal transformer critic that ingests a chunk of actions for multi-step Bellman targets. SAC's higher floor even at baseline mostly comes from being off-policy: reusing each transition many times gives it more training signal per frame than PPO gets. Branch: SAC_Base.

Lessons and next questions

• Trust-region stability dominates capacity once the encoder scales. Hyperparameters tuned on a 2.6M-parameter network did not transfer to one 14× larger. A fixed-step warmup mattered more than the peak learning rate.
• Window length is task-dependent, and the task decides. Measure observability before assuming memory will help.
• Put the guardrails in from run one. The log σ clamp, the log-ratio clamp, and a well-monitored KL check should exist before the first explosion, not after it.
• Every knob moves other knobs. Each architecture change required re-tuning several other settings, and the changes were never uniform. Freezing one YAML per experiment was what made the sequence legible.

Open questions the data cannot answer yet: what the seed variance of v6.2 actually is; whether giving my critic temporal context closes the J_min gap while holding the actor fixed; whether GroupNorm would remove the train/eval toggle cleanly; and what the red-pellet collision rate and wall-pin duration look like as first-class metrics, since the environment already exposes the signals needed to count them.

What is in the repository

The main branch holds ryanpowers_final, my submission folder: the report, an environment README documenting the observation tensor and action space, the PPO baseline package (ppo_baseline) and the transformer package (ppo_transformer, with models, training, and utils), pinned requirements, and four configs: the headline v6.2 config, its laser-enabled sibling, the frozen evaluation config for the headline checkpoint, and the baseline. The folder is hashed so graders can verify it was not modified after submission. Teammates' work lives on their own branches, and the Environment branch carries the Gymnasium wrapper around the Unity binary.

Both training entry points take a YAML config and a run name, accept --resume, and log to Aim. PyTorch is installed separately because the wheel depends on the CUDA version.