normalization.py

"""
Value normalization utilities for PPO training stability
"""

import torch
import numpy as np
from typing import Union, Tuple


class RunningMeanStd:
    """
    Tracks running mean and standard deviation using recursive formulas.
    Provides numerically stable value normalization for RL training.
    """
    
    def __init__(self, epsilon: float = 1e-4, shape: Tuple[int, ...] = (), device: str = 'cpu'):
        self.mean = torch.zeros(shape, dtype=torch.float32, device=device)
        self.var = torch.ones(shape, dtype=torch.float32, device=device)
        self.count = epsilon
        self.device = device
    
    def update(self, x: torch.Tensor):
        """Update running statistics with new batch of data"""
        if x.numel() == 0:
            return
            
        batch_mean = torch.mean(x, dim=0)
        batch_var = torch.var(x, dim=0, unbiased=False)
        batch_count = x.shape[0]
        
        self.update_from_moments(batch_mean, batch_var, batch_count)
    
    def update_from_moments(self, batch_mean: torch.Tensor, batch_var: torch.Tensor, batch_count: int):
        """Update from batch statistics (more numerically stable)"""
        delta = batch_mean - self.mean
        tot_count = self.count + batch_count
        
        new_mean = self.mean + delta * batch_count / tot_count
        m_a = self.var * self.count
        m_b = batch_var * batch_count
        M2 = m_a + m_b + torch.square(delta) * self.count * batch_count / tot_count
        new_var = M2 / tot_count
        
        self.mean = new_mean
        self.var = new_var
        self.count = tot_count
    
    @property
    def std(self) -> torch.Tensor:
        """Return current standard deviation"""
        return torch.sqrt(self.var + 1e-8)
    
    def normalize(self, x: torch.Tensor) -> torch.Tensor:
        """Normalize input using current statistics"""
        return (x - self.mean) / self.std
    
    def denormalize(self, x: torch.Tensor) -> torch.Tensor:
        """Denormalize input using current statistics"""
        return x * self.std + self.mean


class ValueNormalizer:
    """
    Normalizes value function targets for stable training.
    Uses exponential moving average for efficient online updates.
    """
    
    def __init__(self, beta: float = 0.99, device: str = 'cpu'):
        self.beta = beta
        self.device = device
        self.running_mean = 0.0
        self.running_var = 1.0
        self.count = 0
    
    def update(self, values: torch.Tensor):
        """Update normalization statistics"""
        if values.numel() == 0:
            return
            
        batch_mean = torch.mean(values).item()
        batch_var = torch.var(values).item()
        
        if self.count == 0:
            self.running_mean = batch_mean
            self.running_var = batch_var
        else:
            self.running_mean = self.beta * self.running_mean + (1 - self.beta) * batch_mean
            self.running_var = self.beta * self.running_var + (1 - self.beta) * batch_var
            
        self.count += 1
    
    def normalize(self, values: torch.Tensor) -> torch.Tensor:
        """Normalize values"""
        std = np.sqrt(self.running_var + 1e-8)
        return (values - self.running_mean) / std
    
    @property
    def std(self) -> float:
        return np.sqrt(self.running_var + 1e-8)


class RewardClipper:
    """
    Clips and normalizes rewards for training stability.
    Implements both hard clipping and adaptive scaling.
    
    For Breakout:
    - Individual rewards: 1-7 points per brick
    - Typical episode: 100-400 points
    - Clipping at 10.0 preserves reward signal while preventing instability
    """
    
    def __init__(self, clip_range: float = 10.0, adaptive: bool = True, device: str = 'cpu'):
        # Breakout-specific: 10.0 preserves max single-brick reward (7) with safety margin
        # Alternative values: 5.0 (more conservative), 15.0 (less conservative)
        self.clip_range = clip_range
        self.adaptive = adaptive
        self.device = device
        
        if adaptive:
            self.reward_rms = RunningMeanStd(shape=(), device=device)
    
    def process_rewards(self, rewards: torch.Tensor) -> torch.Tensor:
        """Process rewards with clipping and optional normalization"""
        if self.adaptive and rewards.numel() > 0:
            # Update running statistics
            self.reward_rms.update(rewards)
            
            # Normalize by running std (but not mean - preserve reward signal)
            processed_rewards = rewards / self.reward_rms.std
        else:
            processed_rewards = rewards
        
        # Apply hard clipping
        processed_rewards = torch.clamp(processed_rewards, -self.clip_range, self.clip_range)
        
        return processed_rewards


class ObservationNormalizer:
    """
    Normalizes observations for consistent input scaling.
    """
    
    def __init__(self, observation_shape: Tuple[int, ...], device: str = 'cpu'):
        self.obs_rms = RunningMeanStd(shape=observation_shape, device=device)
        self.device = device
    
    def update(self, observations: torch.Tensor):
        """Update normalization statistics"""
        # Flatten batch dimensions for update
        flat_obs = observations.reshape(-1, *observations.shape[-len(self.obs_rms.mean.shape):])
        self.obs_rms.update(flat_obs)
    
    def normalize(self, observations: torch.Tensor) -> torch.Tensor:
        """Normalize observations"""
        return self.obs_rms.normalize(observations)
    
    def denormalize(self, observations: torch.Tensor) -> torch.Tensor:
        """Denormalize observations"""
        return self.obs_rms.denormalize(observations)