dppo_math.py

import torch
import torch.nn as nn

class DashcamBuffer:
    """
    Records the internal (denoising steps) of the diffusion model 
    for the last K' active fine-tuning steps, alongside the environment rewards.
    """
    def __init__ (self,num_env_steps:int,K_prime:int,obs_dim:int,chunk_size:int,act_dim:int,device: str="cpu"):
        # num_env_steps: How many physical environment steps we take per rollout
        # K_prime: How many diffusion steps we are actively fine-tuning (e.g., the last 3 steps)
        self.total_capacity=num_env_steps*K_prime#total micro steps
        self.device=device
        self.ptr=0 # pointer to track where in the buffer


        #preallocate tensots for speed and memory efficiency
        self.states=torch.zeros((self.total_capacity,obs_dim),device=device)
        self.noisy_actions=torch.zeros((self.total_capacity,chunk_size,act_dim),device=device)
        self.committed_noises=torch.zeros((self.total_capacity,chunk_size,act_dim),device=device)
        self.k_steps=torch.zeros((self.total_capacity,),dtype=torch.long,device=device)
        self.log_probs=torch.zeros((self.total_capacity,),device=device)
        self.advantages=torch.zeros((self.total_capacity,),device=device)
        self.returns=torch.zeros((self.total_capacity,),device=device)

    def add_trajectory(self,states,noisy_actions,committed_noises,k_steps,log_probs,env_advantage,env_return):
        """
        Adds the K' micro-steps from ONE physical environment step into the buffer.
        Notice that env_advantage is a single scalar, but we assign it to all K' steps.

        Input Shapes (for a single environment step):
        - states: [K_prime, obs_dim]
        - noisy_actions: [K_prime, chunk_size, act_dim]  (x_k — input to actor)
        - committed_noises: [K_prime, chunk_size, act_dim]  (sampled epsilon — the RL action)
        - k_steps: [K_prime]
        - log_probs: [K_prime]
        - env_advantage: [Scalar]
        - env_return: [Scalar]
        """
        K_prime=states.shape[0]

        #determine the slice in our pre allocated tensors
        idx=slice(self.ptr,self.ptr+K_prime) # start,ending index

        self.states[idx]=states
        self.noisy_actions[idx]=noisy_actions
        self.committed_noises[idx]=committed_noises
        self.k_steps[idx]=k_steps
        self.log_probs[idx]=log_probs

        #broadcasting, copy the single physical advatange/return to all K' internal steps.
        self.advantages[idx]=env_advantage.expand(K_prime)
        self.returns[idx]=env_return.expand(K_prime)

        self.ptr+=K_prime

    def get_all(self):
        #returns the full buffer for the ppo update

        #normalize advantages across the whole batch to stablize ppo gradients
        adv=self.advantages
        normalized_adv=(adv-adv.mean())/(adv.std()+1e-8)

        return(self.states,self.noisy_actions,self.committed_noises,self.k_steps,self.log_probs,normalized_adv,self.returns)
    
    def clear(self):
        self.ptr=0

def calculate_gaussian_log_prob(predicted_noise, target_noise, log_variance):
    """
    Calculates the Log-Likelihood of the diffusion step transition.
    In DDIM, the reverse step is treated as a Gaussian distribution.
    
    predicted_noise Shape: [batch, chunk_size, act_dim] (Output of our MLP Actor)
    target_noise Shape: [batch, chunk_size, act_dim] (The actual noise we added during training)
    log_variance Shape: [batch, 1, 1] (From the diffusion scheduler)
    """
    # Math: log P(x) = -0.5 * ( (x - mu)^2 / var + log(var) + log(2*pi) )
    # Because diffusion predicts NOISE instead of the data directly,
    # we calculate the log prob of the noise matching.

    #calculate mean squared error between predicted and actual noise
    #shape[batch,chunk_size,act_dim]
    squared_error=(predicted_noise-target_noise)**2

    #divide by variance
    scaled_error=squared_error*torch.exp(-log_variance)

    #sum the log likelihoods across chunk and action dimension
    #so a single scalar log_prob per batch
    #log(2*pi) approx 1.837877
    #outputshape:[batch]

    log_prob=-0.5*(scaled_error+log_variance+1.837877).mean(dim=(1,2))

    return log_prob

def compute_ppo_objective(new_log_probs, old_log_probs, advantages, epsilon_clip=0.2):
    """
    The standard Clipped Surrogate Objective Function from PPO.
    
    new_log_probs Shape: [batch] (Calculated by the network RIGHT NOW)
    old_log_probs Shape: [batch] (Pulled from the Dashcam Buffer)
    advantages Shape: [batch] (Pulled from the Dashcam Buffer)
    #we dont use pi(a|s) because theres no way of mathmatically calculating this out
    #so we use the noise predicted for p(xk-1|xk,s) for the last k prime steps
    #batch size for training batch
    """
    
    #calculate the probability ratio new pi(a|s)/old pi(a|s)
    #exp(lognew-log old)=new/old
    ratio=torch.exp(new_log_probs-old_log_probs)

    #unclipped objective
    surr1=ratio*advantages

    # 3. Clipped Objective (Surrogate 2)
    # torch.clamp restricts the ratio to be between [0.8, 1.2]
    # This prevents the policy from updating too drastically in a single step.
    # Shape: [batch]

    clipped_ratio=torch.clamp(ratio,1-epsilon_clip,1+epsilon_clip)
    surr2=clipped_ratio*advantages

    ppo_loss=-torch.min(surr1,surr2).mean()

    return ppo_loss

def get_ddpm_log_variance(scheduler, k_steps, device):
    """
    Returns log(sigma_k^2) for each step in k_steps using the scheduler's beta schedule.
    sigma_k^2 = beta_k is the DDPM posterior variance.

    k_steps shape: [batch]
    output shape:  [batch, 1, 1]  — broadcasts over chunk_size and act_dim
    """
    betas = scheduler.betas.to(device)
    log_var = torch.log(betas[k_steps] + 1e-8)
    return log_var.view(-1, 1, 1)