Overview The DOOM Neuron project uses Proximal Policy Optimization (PPO) to train biological neurons through reinforcement learning. The encoder learns to generate optimal stimulation patterns, while the decoder maps neural responses (spikes) to game actions.
PPO Algorithm Core Concept PPO is an on-policy reinforcement learning algorithm that:
Collects experience rollouts by interacting with the environment
Estimates advantages using Generalized Advantage Estimation (GAE)
Updates the policy using clipped surrogate objective
Constrains updates to prevent catastrophic policy collapse
Training Loop (training_server.py) for episode in range (max_episodes): # Phase 1: Collect rollout (2048 steps) for step in range (steps_per_update): # Encoder: obs → stimulation parameters frequencies, amplitudes, enc_logprob, enc_entropy = policy.sample_encoder(obs) # Apply stimulation to neurons policy.apply_stimulation(stim_socket, frequencies, amplitudes) # Collect neural responses spike_counts = policy.collect_spikes(spike_socket) # Decoder: spikes → actions actions, dec_logprob, dec_entropy = policy.decode_spikes_to_action(spike_counts) # Execute action in environment next_obs, reward, done, info = env.step(actions) # Store transition buffer.store(obs, actions, reward, logprob, value, spike_counts, frequencies, amplitudes) # Phase 2: Compute advantages using GAE advantages, returns = compute_gae_returns(rewards, values, gamma = 0.997 , gae_lambda = 0.95 ) # Phase 3: PPO update (4 epochs) for epoch in range (num_epochs): for batch in minibatches( batch_size = 256 ): # Re-evaluate actions under current policy logprobs, values, entropy, enc_logprob, enc_entropy = policy.evaluate_actions( batch.spike_features, batch.actions, batch.obs, batch.stim_frequencies, batch.stim_amplitudes ) # PPO clipped loss policy_loss = compute_policy_loss(logprobs, batch.old_logprobs, advantages) value_loss = compute_value_loss(values, returns) entropy_bonus = entropy.mean() * entropy_coef encoder_entropy_penalty = enc_entropy.mean() * encoder_entropy_coef loss = policy_loss + value_loss_coef * value_loss - entropy_bonus + encoder_entropy_penalty optimizer.zero_grad() loss.backward() nn.utils.clip_grad_norm_(policy.parameters(), max_grad_norm) optimizer.step()
Policy Gradients Encoder Gradients The encoder network learns to generate stimulation parameters via policy gradients through Beta distributions :
# From EncoderNetwork.sample() in ppo_doom.py:495-515 class EncoderNetwork ( nn . Module ): def sample ( self , obs , deterministic = False ): # Compute Beta distribution parameters features = self ._combine_features(obs) # Process obs with CNN + MLP freq_alpha = F.softplus( self .freq_alpha_head(features)) + 1.0 freq_beta = F.softplus( self .freq_beta_head(features)) + 1.0 amp_alpha = F.softplus( self .amp_alpha_head(features)) + 1.0 amp_beta = F.softplus( self .amp_beta_head(features)) + 1.0 # Create Beta distributions freq_dist = Beta(freq_alpha, freq_beta) amp_dist = Beta(amp_alpha, amp_beta) # Sample and scale to physical ranges freq_u = freq_dist.rsample() # Reparameterization trick amp_u = amp_dist.rsample() frequencies = min_frequency + freq_u * (max_frequency - min_frequency) # 4-40 Hz amplitudes = min_amplitude + amp_u * (max_amplitude - min_amplitude) # 1.0-2.5 μA # Compute log probabilities with Jacobian correction freq_log_prob = freq_dist.log_prob(freq_u) - log(freq_range) amp_log_prob = amp_dist.log_prob(amp_u) - log(amp_range) return frequencies, amplitudes, freq_log_prob, amp_log_prob, freq_entropy, amp_entropy
The encoder uses Beta distributions instead of Gaussian because stimulation parameters are naturally bounded (frequency: 4-40 Hz, amplitude: 1.0-2.5 μA). Beta distributions provide better sample efficiency in bounded spaces.
Decoder Gradients The decoder receives gradients through its action logits:
# From DecoderNetwork.forward() in training_server.py:627-629 class DecoderNetwork ( nn . Module ): def forward ( self , spike_features ): # Linear readout: spikes → action logits head_input = self .shared(spike_features) if self .use_mlp else spike_features joint_logits = self .joint_head(head_input) # (batch, 54) for combinatorial actions return joint_logits
Action probabilities flow through the decoder:
spike_counts → decoder → action_logits → Categorical → sampled_action ↑ ↓ └───── policy gradient ──────┘
By setting decoder_zero_bias=True (default), the decoder is forced to rely entirely on neural spike activity rather than learning a bias-driven policy. This ensures biological neurons are genuinely controlling the agent.
Generalized Advantage Estimation (GAE) Advantage Computation GAE computes advantage estimates that balance bias and variance:
# Pseudocode from PPO implementation def compute_gae ( rewards , values , gamma = 0.997 , gae_lambda = 0.95 ): """ Compute GAE advantages and returns. Args: rewards: List of rewards [r_0, r_1, ..., r_T] values: Value estimates [V(s_0), V(s_1), ..., V(s_T)] gamma: Discount factor (0.997) gae_lambda: GAE lambda parameter (0.95) """ advantages = [] gae = 0 # Backward pass through episode for t in reversed ( range ( len (rewards))): # TD error: δ_t = r_t + γ V(s_{t+1}) - V(s_t) delta = rewards[t] + gamma * values[t + 1 ] - values[t] # GAE: A_t = δ_t + (γλ) δ_{t+1} + (γλ)² δ_{t+2} + ... gae = delta + gamma * gae_lambda * gae advantages.insert( 0 , gae) # Returns: R_t = A_t + V(s_t) returns = [adv + val for adv, val in zip (advantages, values)] return advantages, returns
Hyperparameters
γ (gamma) = 0.997 : Discount factor for long-term rewards
Higher than typical 0.99 to encourage survival in DOOM
Balances immediate kills vs. staying alive
λ (lambda) = 0.95 : GAE smoothing parameter
Controls bias-variance tradeoff
λ=0: high bias, low variance (TD learning)
λ=1: low bias, high variance (Monte Carlo)
Why GAE? Standard advantage estimation faces a dilemma:
TD errors (1-step): Low variance but biased by value estimatesMonte Carlo returns (full episode): Unbiased but high variance GAE creates a smooth interpolation: A^(1) = δ_t # 1-step (high bias) A^(2) = δ_t + γλ δ_{t+1} # 2-step A^(3) = δ_t + γλ δ_{t+1} + (γλ)² δ_{t+2} # 3-step ... A^(∞) = Σ (γλ)^k δ_{t+k} # Full return (low bias)
Lambda=0.95 provides a sweet spot for DOOM, where:
Value function learns to predict episode outcomes
Advantages capture meaningful state-action quality differences
Updates remain stable despite sparse rewards
Loss Functions Policy Loss (Clipped Surrogate Objective) def compute_policy_loss ( logprobs , old_logprobs , advantages , clip_epsilon = 0.2 ): """ PPO clipped surrogate objective. Args: logprobs: Log probabilities under current policy π_θ old_logprobs: Log probabilities under old policy π_{θ_old} advantages: GAE advantage estimates clip_epsilon: Clipping range (0.2) """ # Importance sampling ratio: r_t(θ) = π_θ(a_t|s_t) / π_{θ_old}(a_t|s_t) ratio = torch.exp(logprobs - old_logprobs) # Unclipped objective: r_t(θ) * A_t unclipped_obj = ratio * advantages # Clipped objective: clip(r_t(θ), 1-ε, 1+ε) * A_t clipped_obj = torch.clamp(ratio, 1 - clip_epsilon, 1 + clip_epsilon) * advantages # Take minimum (pessimistic bound) policy_loss = - torch.min(unclipped_obj, clipped_obj).mean() return policy_loss
The clipping mechanism prevents large policy updates that could break the neural interface. With biological neurons, catastrophic policy changes could overstimulate channels and require retraining from scratch.
Value Loss (MSE with optional normalization) def compute_value_loss ( values , returns , normalize_returns = False ): """ Mean squared error between value predictions and returns. Args: values: V(s) predictions from critic returns: GAE returns (advantages + values) normalize_returns: Whether to normalize returns (default: False) """ if normalize_returns: # Normalize for stable training returns = (returns - returns.mean()) / (returns.std() + 1e-8 ) value_loss = F.mse_loss(values, returns) return value_loss
Configuration: normalize_returns=False (disabled per Doom Initial Report)
Entropy Bonus Entropy regularization encourages exploration:
# From PPOConfig entropy_coef = 0.02 # Decoder entropy bonus encoder_entropy_coef = - 0.10 # Encoder entropy penalty # Total entropy term entropy_term = ( entropy_coef * decoder_entropy.mean() + # Encourage decoder exploration encoder_entropy_coef * encoder_entropy.mean() # Penalize encoder entropy ) total_loss = policy_loss + value_loss_coef * value_loss - entropy_term
The negative encoder entropy coefficient (-0.10) penalizes high entropy in stimulation parameters. This encourages the encoder to be more deterministic once it finds effective stimulation patterns, reducing noise in the neural interface.
Reward Shaping Base Rewards from VizDoom # From VizDoomEnv configuration game.set_kill_reward( 200 ) # +200 for each enemy killed
Event-Based Rewards The environment tracks game events and computes shaped rewards:
# From VizDoomEnv.step() reward computation reward_components = { 'kill_reward' : info[ 'event_enemy_kill' ] * 200 , 'armor_bonus' : info[ 'event_armor_pickup' ] * 50 , 'damage_penalty' : info[ 'event_took_damage' ] * ( - 10 ), 'ammo_waste_penalty' : info[ 'event_ammo_waste' ] * ( - 5 ), 'approach_bonus' : info[ 'event_move_closer' ] * 1.0 , 'retreat_penalty' : info[ 'event_move_farther' ] * ( - 0.5 ), } total_reward = sum (reward_components.values())
Simplified Reward Mode With simplified_reward=True (default), only core events contribute:
if config.simplified_reward: # Focus on kills and survival reward = ( kill_reward + damage_penalty + armor_bonus ) else : # Full shaped rewards including aim alignment and velocity reward = ( kill_reward + damage_penalty + armor_bonus + aim_alignment_reward + # Bonus for facing enemies velocity_reward + # Encourage movement ammo_waste_penalty + approach_bonus + retreat_penalty )
Simplified reward was enabled based on the “Doom Initial Report” findings. The full shaped reward system can provide denser learning signals but may introduce confounding factors when analyzing biological learning.
Training Hyperparameters PPO Configuration # From PPOConfig in training_server.py:120-128 learning_rate = 3e-4 # Adam optimizer learning rate gamma = 0.997 # Discount factor (higher for survival) gae_lambda = 0.95 # GAE lambda parameter clip_epsilon = 0.2 # PPO clipping range value_loss_coef = 0.3 # Value loss weight entropy_coef = 0.02 # Decoder entropy bonus max_grad_norm = 3.0 # Gradient clipping threshold normalize_returns = False # Return normalization disabled
Rollout Parameters num_envs = 1 # Single environment (no parallelization) steps_per_update = 2048 # Steps collected before each update batch_size = 256 # Minibatch size for SGD num_epochs = 4 # PPO epochs per rollout max_episodes = 2000 # Total training episodes
Network Architecture hidden_size = 128 # MLP hidden layer size encoder_cnn_channels = 64 # CNN base channels (increased from 16) encoder_trainable = True # Learn encoder via policy gradients encoder_entropy_coef = - 0.10 # Penalize encoder entropy decoder_use_mlp = False # Use linear readout (not MLP) decoder_zero_bias = True # Zero decoder biases decoder_enforce_nonnegative = False # Allow negative weights decoder_mlp_hidden = 256 # MLP hidden size (if enabled)
Gradient Clipping # Prevent gradient explosion nn.utils.clip_grad_norm_(policy.parameters(), max_norm = max_grad_norm)
max_grad_norm = 3.0: Conservative clipping to protect neural interfaceBiological neurons may be sensitive to extreme stimulation changes
Lower values (0.5-1.0) tested but found too restrictive
Combinatorial Action Space The decoder outputs logits over 54 discrete joint actions:
# From PPOPolicy.__init__() in training_server.py:752-769 forward_options = [ 'none' , 'forward' , 'backward' ] # 3 options strafe_options = [ 'none' , 'left' , 'right' ] # 3 options camera_options = [ 'none' , 'turn_left' , 'turn_right' ] # 3 options attack_options = [ 'idle' , 'attack' ] # 2 options speed_options = [ 'off' ] # 1 option (removed) # Total: 3 × 3 × 3 × 2 × 1 = 54 actions for fwd in forward_options: for strafe in strafe_options: for turn in camera_options: for attack in attack_options: combinatorial_action_defs.append({ 'forward' : fwd, 'strafe' : strafe, 'turn' : turn, 'attack' : attack })
Why Combinatorial Actions?
Earlier versions used factored action spaces (separate distributions for movement/turning/attack), but this caused issues:
Independence assumption violated : Movement and attack should be coordinatedCredit assignment difficulty : Hard to tell which action component caused rewardExploration challenges : Random independent actions rarely produce coherent behavior The combinatorial action space:
Models realistic action combinations (e.g., “strafe left while shooting”)
Simplifies credit assignment to single action selection
Reduces decoder to single softmax head (54 logits)
Allows natural exploration through single categorical distribution
Training Dynamics Exploration Strategy # Sample from action distribution during training joint_dist = Categorical( logits = joint_logits) joint_action = joint_dist.sample() # Stochastic sampling # Deterministic evaluation joint_action = joint_logits.argmax( dim =- 1 ) # Greedy selection
Exploration is managed by:
Entropy bonus (0.02): Encourages trying diverse actionsStochastic sampling : Naturally explores via softmax distributionEncoder exploration : Beta distributions add stimulation variability
Value Function Bootstrap The value network learns to predict episode returns:
# Value targets from GAE returns = advantages + values # R_t = A_t + V(s_t) # Value loss value_loss = (value_predictions - returns)²
This helps with:
Sparse rewards : Predict long-term value even without immediate rewardCredit assignment : Understand which states lead to eventual successAdvantage estimation : Reduce variance in policy gradients
Biological Learning Considerations Encoder Adaptation The encoder must learn stimulation patterns that:
Evoke informative spike patterns from biological neuronsDifferentiate game states through neural activityRemain within safe stimulation bounds (1-2.5 μA, 4-40 Hz)
Decoder Constraints # Linear readout without bias if decoder_zero_bias: decoder.bias.data.zero_() decoder.bias.requires_grad = False # Optional non-negative weights if decoder_enforce_nonnegative: weights = F.softplus(decoder.weight) # Ensure w ≥ 0
These constraints ensure:
Actions depend on spike counts, not learned offsets
Decoder interpretability (positive weights = excitatory)
Biological plausibility of the read-out mechanism
Monitor the ratio of decoder weight magnitude to bias magnitude during training. If bias dominates, the biological neurons aren’t contributing meaningfully to decisions.
Logging and Metrics Key metrics tracked during training:
metrics = { 'Policy/loss' : policy_loss.item(), 'Policy/entropy' : entropy.mean().item(), 'Value/loss' : value_loss.item(), 'Value/mean' : values.mean().item(), 'Encoder/entropy' : encoder_entropy.mean().item(), 'Encoder/freq_mean' : frequencies.mean().item(), 'Encoder/amp_mean' : amplitudes.mean().item(), 'Decoder/weight_l2' : decoder.weight.pow( 2 ).sum().item(), 'Episode/reward' : episode_reward, 'Episode/length' : episode_length, 'Episode/kills' : kill_count, }
These are logged to TensorBoard for monitoring convergence and diagnosing issues.