utils.py

import re
from pathlib import Path

import hydra
import numpy as np
import os
import random
import torch
import torch.nn.functional as F
import wandb

from omegaconf import DictConfig
from typing import Tuple


def seed_all(seed):
    np.random.seed(seed)
    torch.manual_seed(seed)
    random.seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)

def mask_random_obs(
    images: torch.Tensor,
    p: float = 0.15
) -> torch.Tensor:
    if not (0 <= p <= 1):
        raise ValueError(f"Probability p must be between 0 and 1, got {p}")

    B_N, T, H, W, C = images.shape
    device = images.device

    mask = torch.bernoulli(torch.full((B_N, T, 1, 1, 1), 1 - p, device=device))
    masked_images = images * mask

    return masked_images

def random_shift(images: torch.Tensor, padding: int = 4) -> torch.Tensor:
    """
    Applies a random shift to each image in the batch.

    Args:
        images (torch.Tensor): Tensor of shape (B * N, T, H, W, C).
        padding (int, optional): Number of pixels to pad on each side. Defaults to 4.

    Returns:
        torch.Tensor: Shifted images with the same shape as input.
    """
    B_N, T, H, W, C = images.shape

    # Permute to (B * N * T, C, H, W) for processing
    images = images.permute(0, 4, 1, 2, 3).reshape(B_N * T, C, H, W)

    # Generate random shifts in pixels within [-padding, padding]
    shift_x = torch.randint(-padding, padding + 1, (B_N * T,), device=images.device).float()
    shift_y = torch.randint(-padding, padding + 1, (B_N * T,), device=images.device).float()

    # Normalize shifts to [-1, 1] based on image dimensions
    shift_x_norm = shift_x / W
    shift_y_norm = shift_y / H

    # Create affine transformation matrices for each image
    theta = torch.zeros((B_N * T, 2, 3), device=images.device)
    theta[:, 0, 0] = 1
    theta[:, 1, 1] = 1
    theta[:, 0, 2] = shift_x_norm
    theta[:, 1, 2] = shift_y_norm

    # Generate grids for sampling
    grid = F.affine_grid(theta, images.size(), align_corners=False)

    # Apply grid sampling with reflection padding to handle borders
    shifted_images = F.grid_sample(images, grid, padding_mode='reflection', align_corners=False)

    # Reshape back to (B * N, T, H, W, C)
    shifted_images = shifted_images.view(B_N, T, C, H, W).permute(0, 1, 3, 4, 2)

    return shifted_images

def random_crop(images: torch.Tensor, crop_size: Tuple[int, int] = None) -> torch.Tensor:
    """
    Applies a random crop to each image and resizes it back to the original size.

    Args:
        images (torch.Tensor): Tensor of shape (B * N, T, H, W, C).
        crop_size (Tuple[int, int], optional): Desired crop size (H_crop, W_crop).
                                              If None, defaults to original size. Defaults to None.

    Returns:
        torch.Tensor: Cropped and resized images with the same shape as input.
    """
    B_N, T, H, W, C = images.shape

    if crop_size is None:
        crop_size = (H, W)
    H_crop, W_crop = crop_size

    # Permute to (B * N * T, C, H, W) for processing
    images = images.permute(0, 4, 1, 2, 3).reshape(B_N * T, C, H, W)

    # Determine scaling factors
    scale_h = H_crop / H
    scale_w = W_crop / W

    # Generate random translations in normalized coordinates
    translations_h = (torch.randint(0, H - H_crop + 1, (B_N * T,), device=images.device).float() / (H / 2)) - 1.0
    translations_w = (torch.randint(0, W - W_crop + 1, (B_N * T,), device=images.device).float() / (W / 2)) - 1.0

    # Create affine transformation matrices for cropping and resizing
    theta = torch.zeros((B_N * T, 2, 3), device=images.device)
    theta[:, 0, 0] = scale_w
    theta[:, 1, 1] = scale_h
    theta[:, 0, 2] = translations_w
    theta[:, 1, 2] = translations_h

    # Generate grids for sampling
    grid = F.affine_grid(theta, images.size(), align_corners=False)

    # Apply grid sampling with reflection padding to handle borders
    cropped_resized_images = F.grid_sample(images, grid, padding_mode='reflection', align_corners=False)

    # Reshape back to (B * N, T, H, W, C)
    cropped_resized_images = cropped_resized_images.view(B_N, T, C, H, W).permute(0, 1, 3, 4, 2)

    return cropped_resized_images

def compute_discounted_returns(gamma, rewards):
        batch_size, trajectory_len = rewards.shape
        returns = torch.empty(
            (batch_size, trajectory_len), 
            device=rewards.device,
            dtype=torch.float32
        )
        returns[:, -1] = rewards[:, -1]
        for t in reversed(range(trajectory_len-1)):
            returns[:, t] = rewards[:, t] + gamma * returns[:, t+1]
        return returns

def compute_gae_advantages(rewards, values, gamma=0.96, lmbda=0.95):
    # Compute deltas
    with torch.no_grad():
        deltas = rewards[:, :-1] + gamma * values[:, 1:] - values[:, :-1]
        
        advantages = torch.empty_like(deltas)
        advantage = 0
        for t in reversed(range(deltas.size(1))):
            advantages[:, t] = advantage = deltas[:, t] + gamma * lmbda * advantage
        
        # Normalize advantages
        # advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
    
    return advantages

def instantiate_agent(cfg: DictConfig):
    assert cfg.use_transformer is True, "Only Transformer is supported for now."
    
    encoder = hydra.utils.instantiate(cfg.encoder)
    actor = hydra.utils.instantiate(
        cfg.mlp_model,
        encoder=encoder,
    )
    
    if not cfg.share_encoder:
        encoder_c = hydra.utils.instantiate(cfg.encoder)
        critic = hydra.utils.instantiate(
            cfg.linear_model, 
            encoder=encoder_c,
        )
    else:
        critic = hydra.utils.instantiate(
            cfg.linear_model,
            encoder=encoder,
        )

    encoder_t = hydra.utils.instantiate(cfg.encoder)
    target = hydra.utils.instantiate(
        cfg.linear_model, 
        encoder=encoder_t,
    )

    ss_module = hydra.utils.instantiate(
        cfg.ss_module,
        encoder=encoder,
    )

    target.load_state_dict(critic.state_dict())

    optimizer_actor = hydra.utils.instantiate(cfg.optimizer_actor, params=actor.parameters())
    optimizer_critic = hydra.utils.instantiate(cfg.optimizer_critic, params=critic.parameters())
    optimizer_ss = hydra.utils.instantiate(cfg.optimizer_ss, params=ss_module.parameters())

    agent = hydra.utils.instantiate(
        cfg.agent,
        actor=actor,
        critic=critic,
        target=target,
        ss_module=ss_module,
        critic_optimizer=optimizer_critic,
        actor_optimizer=optimizer_actor,
        ss_optimizer=optimizer_ss,
    )
    return agent

def update_target_network(target, source, tau=0.005):
    for target_param, param in zip(target.parameters(), source.parameters()):
        target_param.data.copy_(tau*param.data + (1.0-tau)*target_param.data)

def save_checkpoint(agent, replay_buffer, agent_replay_buffer, step, checkpoint_dir, save_agent_replay_buffer=False, also_clean_old=False):
    # Get the W&B run ID
    run_id = wandb.run.id

    from train import ReplayBuffer, GPUMemoryFriendlyAgentReplayBuffer
    replay_buffer: ReplayBuffer
    agent_replay_buffer: GPUMemoryFriendlyAgentReplayBuffer

    # Create a folder with the run ID if it doesn't exist
    run_folder = Path(checkpoint_dir) / run_id
    os.makedirs(run_folder, exist_ok=True)

    # Generate a unique checkpoint directory for the current step
    checkpoint_path = run_folder / ('step_'+str(step))
    os.makedirs(checkpoint_path, exist_ok=True)

    agent_checkpoint = agent.state_dict()
    torch.save(agent_checkpoint, checkpoint_path / 'agent.pt')

    if replay_buffer is not None and save_agent_replay_buffer:
        if replay_buffer.use_memmap:
            # When using memmap, only save the metadata
            replay_buffer_state_dict = {
                'number_of_added_trajectories': replay_buffer.number_of_added_trajectories,
                'marker': replay_buffer.marker,
                'fm_marker': replay_buffer.fm_marker
            }
            torch.save(replay_buffer_state_dict, checkpoint_path / 'replay_buffer_meta.pt')
        else:
            # If not using memmap, save the entire data
            torch.save(replay_buffer.trajectory_batch.data, checkpoint_path / 'replay_buffer_data.pt')
    if agent_replay_buffer is not None and save_agent_replay_buffer:
        if agent_replay_buffer.use_memmap:
            # When using memmap, only save the metadata
            agent_replay_buffer_state_dict = {
                'num_added_agents': agent_replay_buffer.num_added_agents,
                'marker': agent_replay_buffer.marker
            }
            torch.save(agent_replay_buffer_state_dict, checkpoint_path / 'agent_replay_buffer_meta.pt')
        else:
            # If not using memmap, save the entire data
            torch.save({'params': agent_replay_buffer.agents_batched_state_dicts,
                        'num_added_agents': agent_replay_buffer.num_added_agents},
                       checkpoint_path / 'agent_replay_buffer.pt')

    if also_clean_old:
        for d in run_folder.iterdir():
            if re.match(f'step_{str(step)}', d.name):
                continue
            elif re.match(r'step_\d+', d.name):
                # make sure we don't delete the current checkpoint
                d = Path(d)
                assert d != checkpoint_path
                # remove the replay buffer and agent replay buffer if they exist
                if (d / 'replay_buffer_data.pt').exists():
                    os.remove(d / 'replay_buffer_data.pt')
                if (d / 'agent_replay_buffer.pt').exists():
                    os.remove(d / 'agent_replay_buffer.pt')
                # Remove the metadata files as well
                if (d / 'replay_buffer_meta.pt').exists():
                    os.remove(d / 'replay_buffer_meta.pt')
                if (d / 'agent_replay_buffer_meta.pt').exists():
                    os.remove(d / 'agent_replay_buffer_meta.pt')

    print(f"(@@-o)Agent Checkpoint saved at: {checkpoint_path}")


def wandb_stats(trajectory, B, N):
    stats = {}
    stats['Average reward'] = trajectory.data['rewards'].mean().detach()
    stats['Average sum rewards'] = torch.sum(
        trajectory.data['rewards'].reshape(B, N, -1), dim=1
    ).unsqueeze(1).repeat(1, N, 1).mean().detach()
    return stats


from omegaconf import OmegaConf


def flatten_native_dict(cfg, parent_key='', sep='.'):
    """
    Recursively flattens a Python dictionary (or list) into a list of key-value pairs.
    The keys will be in dot notation.

    Args:
        cfg: Python dict or list to flatten.
        parent_key: Current key prefix (used in recursion).
        sep: Separator between keys (default is a dot).

    Returns:
        A list of tuples containing flattened key-value pairs.
    """
    items = []

    if isinstance(cfg, dict):
        for k, v in cfg.items():
            new_key = f"{parent_key}{sep}{k}" if parent_key else k
            if isinstance(v, (dict, list)):
                items.extend(flatten_native_dict(v, new_key, sep=sep))
            else:
                items.append((new_key, v))
    elif isinstance(cfg, list):
        for i, v in enumerate(cfg):
            new_key = f"{parent_key}{sep}{i}" if parent_key else str(i)
            if isinstance(v, (dict, list)):
                items.extend(flatten_native_dict(v, new_key, sep=sep))
            else:
                items.append((new_key, v))

    return items