policies/consistency_policy.py

"""Consistency Policy for Fast Robot Control.

This module implements Consistency Policy, which distills a pretrained EDM
teacher into a fast 1-step or few-step action generator. Uses CTM (Consistency
Training Model) distillation with combined CTM + DSM losses.

Reference: "Consistency Policy: Accelerated Visuomotor Policies via
           Consistency Distillation" (RSS 2024)
Adapted from: https://github.com/Aaditya-Prasad/consistency-policy
"""

from __future__ import annotations

import copy
import logging
import math
from typing import Any, Dict, List, Optional, Tuple, Union

import torch
import torch.nn as nn
import torch.nn.functional as F
from omegaconf import DictConfig

from policies.base import BasePolicy
from networks.edm_scheduler import CTMScheduler, huber_loss
from networks.consistency_unet import ConsistencyUnet1D
from networks.unet import Unet

# Try importing vision encoder from diffusion_policy
try:
    from diffusion_policy.model.common.normalizer import LinearNormalizer
    from diffusion_policy.model.diffusion.mask_generator import LowdimMaskGenerator
    from diffusion_policy.common.robomimic_config_util import get_robomimic_config
    from diffusion_policy.common.pytorch_util import dict_apply, replace_submodules
    from robomimic.algo import algo_factory
    from robomimic.algo.algo import PolicyAlgo
    import robomimic.utils.obs_utils as ObsUtils
    import robomimic.models.base_nets as rmbn
    import diffusion_policy.model.vision.crop_randomizer as dmvc

    # Robustly locate CropRandomizer across robomimic versions
    try:
        from robomimic.models.obs_core import CropRandomizer as RM_CropRandomizer
    except Exception:
        try:
            from robomimic.models.base_nets import CropRandomizer as RM_CropRandomizer
        except Exception:
            RM_CropRandomizer = None

    HAS_ROBOMIMIC = True
except ImportError:
    HAS_ROBOMIMIC = False
    RM_CropRandomizer = None
    from diffusion_policy.model.common.normalizer import LinearNormalizer

logger = logging.getLogger(__name__)

__all__ = ["ConsistencyImagePolicy"]


def state_dict_to_model(state_dict: dict, pattern: str = r"model\.") -> dict:
    """Extract model state dict from checkpoint."""
    import re

    new_state_dict = {}
    prefix = re.compile(pattern)
    prefix_len = len(pattern.replace("\\", "").replace(".", "")) + 1  # +1 for the dot

    for k, v in state_dict["state_dicts"]["model"].items():
        if re.match(prefix, k):
            new_k = k[prefix_len:]
            new_state_dict[new_k] = v

    return new_state_dict


def state_dict_to_obs_encoder(state_dict: dict) -> dict:
    """Extract obs_encoder state dict from checkpoint."""
    new_state_dict = {}
    prefix = "obs_encoder."
    prefix_len = len(prefix)

    for k, v in state_dict["state_dicts"]["model"].items():
        if k.startswith(prefix):
            new_k = k[prefix_len:]
            new_state_dict[new_k] = v

    return new_state_dict


class ConsistencyImagePolicy(BasePolicy):
    """Consistency Policy for fast image-based robot control.

    This policy is distilled from a pretrained EDM teacher using CTM
    (Consistency Training Model) distillation. It can generate actions
    in 1 step (fastest) or use multi-step chaining for better quality.

    Args:
        shape_meta: Dictionary containing observation and action shapes.
        horizon: Total planning horizon (number of timesteps).
        n_action_steps: Number of action steps to predict.
        n_obs_steps: Number of observation steps to condition on.
        noise_scheduler: CTMScheduler instance (or config dict).
        num_inference_steps: Number of inference steps (1 for fastest).
        chaining_times: Timesteps for multi-step chaining.
        teacher_path: Path to pretrained EDM teacher checkpoint.
        losses: Loss configuration dict (ctm_weight, dsm_weight).
        dsm_weights: DSM loss weighting scheme.
        dropout_rate: Dropout rate for consistency training.
        delta: Huber loss delta parameter.
        crop_shape: Image crop size for augmentation.
        diffusion_step_embed_dim: Dimension of diffusion time embedding.
        down_dims: U-Net channel dimensions.
        kernel_size: Convolution kernel size.
        n_groups: GroupNorm groups.
        cond_predict_scale: Use FiLM scale+shift conditioning.
        obs_encoder_group_norm: Use GroupNorm in vision encoder.
        eval_fixed_crop: Use center crop during evaluation.
        obs_as_global_cond: Use observation as global conditioning.
        initial_ema_decay: EMA decay rate for student model.
        action_space: Gym action space (for clipping).
        device: Device string.
        dtype: Data type string.
        clip_actions: Whether to clip actions.
    """

    def __init__(
        self,
        shape_meta: dict,
        horizon: int = 16,
        n_action_steps: int = 8,
        n_obs_steps: int = 2,
        # Scheduler
        noise_scheduler: Optional[Union[CTMScheduler, dict]] = None,
        # Inference
        num_inference_steps: int = 1,
        chaining_times: Optional[List[Union[str, int]]] = None,
        # Teacher
        teacher_path: Optional[str] = None,
        inference_mode: bool = False,
        # Loss configuration
        losses: Optional[Dict[str, float]] = None,
        dsm_weights: str = "none",
        # Training
        dropout_rate: float = 0.2,
        delta: float = 0.0,
        initial_ema_decay: float = 0.9,
        # Vision encoder
        crop_shape: Tuple[int, int] = (76, 76),
        obs_encoder_group_norm: bool = False,
        eval_fixed_crop: bool = False,
        obs_as_global_cond: bool = True,
        # U-Net architecture
        diffusion_step_embed_dim: int = 128,
        down_dims: Tuple[int, ...] = (256, 512, 1024),
        kernel_size: int = 5,
        n_groups: int = 8,
        cond_predict_scale: bool = True,
        # Base policy args
        action_space: Any = None,
        device: str = "cpu",
        dtype: str = "float32",
        clip_actions: bool = True,
        **kwargs,
    ):
        super().__init__(
            action_space=action_space,
            device=device,
            dtype=dtype,
            clip_actions=clip_actions,
        )

        # Default loss configuration
        if losses is None:
            losses = {"ctm": 1.0, "dsm": 0.5}

        # Default chaining times (for 3-step inference)
        if chaining_times is None:
            chaining_times = ["D", 27, 54]

        # Parse shape_meta
        action_shape = shape_meta["action"]["shape"]
        assert len(action_shape) == 1
        action_dim = action_shape[0]
        obs_shape_meta = shape_meta["obs"]

        # Build observation config
        obs_config = {"low_dim": [], "rgb": [], "depth": [], "scan": []}
        obs_key_shapes = {}

        for key, attr in obs_shape_meta.items():
            shape = attr["shape"]
            obs_key_shapes[key] = list(shape)
            obs_type = attr.get("type", "low_dim")

            if obs_type == "rgb":
                obs_config["rgb"].append(key)
            elif obs_type == "low_dim":
                obs_config["low_dim"].append(key)
            else:
                raise RuntimeError(f"Unsupported obs type: {obs_type}")

        # Build vision encoder using robomimic
        if HAS_ROBOMIMIC:
            config = get_robomimic_config(
                algo_name="bc_rnn",
                hdf5_type="image",
                task_name="square",
                dataset_type="ph",
            )

            with config.unlocked():
                config.observation.modalities.obs = obs_config

                if crop_shape is None:
                    for key, modality in config.observation.encoder.items():
                        if modality.obs_randomizer_class == "CropRandomizer":
                            modality["obs_randomizer_class"] = None
                else:
                    ch, cw = crop_shape
                    for key, modality in config.observation.encoder.items():
                        if modality.obs_randomizer_class == "CropRandomizer":
                            modality.obs_randomizer_kwargs.crop_height = ch
                            modality.obs_randomizer_kwargs.crop_width = cw

            ObsUtils.initialize_obs_utils_with_config(config)

            policy: PolicyAlgo = algo_factory(
                algo_name=config.algo_name,
                config=config,
                obs_key_shapes=obs_key_shapes,
                ac_dim=action_dim,
                device="cpu",
            )

            obs_encoder = policy.nets["policy"].nets["encoder"].nets["obs"]

            if obs_encoder_group_norm:
                replace_submodules(
                    root_module=obs_encoder,
                    predicate=lambda x: isinstance(x, nn.BatchNorm2d),
                    func=lambda x: nn.GroupNorm(
                        num_groups=x.num_features // 16, num_channels=x.num_features
                    ),
                )

            if eval_fixed_crop and RM_CropRandomizer is not None:
                replace_submodules(
                    root_module=obs_encoder,
                    predicate=lambda x: isinstance(x, RM_CropRandomizer),
                    func=lambda x: dmvc.CropRandomizer(
                        input_shape=x.input_shape,
                        crop_height=x.crop_height,
                        crop_width=x.crop_width,
                        num_crops=x.num_crops,
                        pos_enc=x.pos_enc,
                    ),
                )
            elif eval_fixed_crop and RM_CropRandomizer is None:
                logger.warning(
                    "eval_fixed_crop=True but could not locate robomimic CropRandomizer; skipping."
                )

            obs_feature_dim = obs_encoder.output_shape()[0]
        else:
            obs_feature_dim = sum(
                sum(s) if isinstance(s, (list, tuple)) else s
                for s in obs_key_shapes.values()
            )
            obs_encoder = nn.Identity()
            logger.warning(
                "robomimic not available, using identity obs encoder."
            )

        # Create student model (ConsistencyUnet1D with stop-time)
        input_dim = action_dim
        global_cond_dim = None
        if obs_as_global_cond:
            global_cond_dim = obs_feature_dim * n_obs_steps

        model = ConsistencyUnet1D(
            input_dim=input_dim,
            global_cond_dim=global_cond_dim,
            diffusion_step_embed_dim=diffusion_step_embed_dim,
            down_dims=list(down_dims),
            kernel_size=kernel_size,
            n_groups=n_groups,
            cond_predict_scale=cond_predict_scale,
            dropout_rate=dropout_rate,
        )

        # Create EMA model for consistency training
        model_ema = copy.deepcopy(model)
        model_ema.requires_grad_(False)

        # Prepare dropout generators
        model.prepare_drop_generators()
        model_ema.prepare_drop_generators()

        # Create teacher model (standard Unet, loaded from checkpoint)
        teacher = Unet(
            input_dim=input_dim,
            cond_dim=global_cond_dim,
            diffusion_step_embed_dim=diffusion_step_embed_dim,
            down_dims=list(down_dims),
            kernel_size=kernel_size,
            n_groups=n_groups,
            cond_predict_scale=cond_predict_scale,
        )

        # Load teacher weights if provided
        if not inference_mode and teacher_path is not None:
            checkpoint = torch.load(teacher_path, map_location="cpu")
            # Load teacher Unet weights
            unet_state_dict = state_dict_to_model(checkpoint)
            teacher.load_state_dict(unet_state_dict)
            teacher.eval()
            teacher.requires_grad_(False)
            # Load obs_encoder weights from teacher (critical for matching embeddings)
            obs_encoder_state_dict = state_dict_to_obs_encoder(checkpoint)
            obs_encoder.load_state_dict(obs_encoder_state_dict)
            logger.info(f"Loaded teacher Unet and obs_encoder from: {teacher_path}")

        # Create noise scheduler
        if noise_scheduler is None:
            noise_scheduler = CTMScheduler(ode_steps_max=1)
        elif isinstance(noise_scheduler, (dict, DictConfig)):
            noise_scheduler = CTMScheduler(**noise_scheduler)

        # Store components
        self.obs_encoder = obs_encoder
        self.model = model
        self.model_ema = model_ema
        self.teacher = teacher
        self.noise_scheduler = noise_scheduler
        self.normalizer = LinearNormalizer()

        self.horizon = horizon
        self.obs_feature_dim = obs_feature_dim
        self.action_dim = action_dim
        self.n_action_steps = n_action_steps
        self.n_obs_steps = n_obs_steps
        self.obs_as_global_cond = obs_as_global_cond

        # Training parameters
        self.losses = losses
        self.dsm_weights = dsm_weights
        self.delta = delta
        self.ema_decay = initial_ema_decay

        # Inference parameters
        self.num_inference_steps = num_inference_steps
        self.chaining_times = chaining_times
        self.chain = False  # Enable via enable_chaining()

        # Mask generator
        self.mask_generator = LowdimMaskGenerator(
            action_dim=action_dim,
            obs_dim=0 if obs_as_global_cond else obs_feature_dim,
            max_n_obs_steps=n_obs_steps,
            fix_obs_steps=True,
            action_visible=False,
        )

        logger.info(
            f"ConsistencyImagePolicy: student_params={sum(p.numel() for p in self.model.parameters())}, "
            f"teacher_params={sum(p.numel() for p in self.teacher.parameters())}, "
            f"vision_params={sum(p.numel() for p in self.obs_encoder.parameters())}"
        )
        logger.info(f"Using losses: {self.losses}")

    def set_normalizer(self, normalizer: LinearNormalizer) -> None:
        """Load normalizer from fitted instance."""
        self.normalizer.load_state_dict(normalizer.state_dict())

    def drop_teacher(self) -> None:
        """Remove teacher model to save memory during inference."""
        self.teacher = None

    def enable_chaining(self) -> None:
        """Enable multi-step chaining for inference."""
        if self.chaining_times is not None:
            self.chain = True
            logger.info(f"Chaining enabled with times: {self.chaining_times}")
        else:
            raise ValueError("Chaining times not set")

    def disable_chaining(self) -> None:
        """Disable multi-step chaining."""
        self.chain = False

    # ==================== FORWARD PASS ====================

    def _forward(
        self,
        model: nn.Module,
        sample: torch.Tensor,
        timestep: torch.Tensor,
        stop_time: torch.Tensor,
        local_cond: Optional[torch.Tensor] = None,
        global_cond: Optional[torch.Tensor] = None,
        clamp: bool = False,
    ) -> torch.Tensor:
        """Forward pass through CTM model with stop-time conditioning."""
        denoise = lambda x, t, s: model(
            x, t, s, local_cond=local_cond, global_cond=global_cond
        )
        return self.noise_scheduler.ctm_calc_out(
            denoise, sample, timestep, stop_time, clamp=clamp
        )

    # ==================== INFERENCE ====================

    def conditional_sample(
        self,
        condition_data: torch.Tensor,
        condition_mask: torch.Tensor,
        local_cond: Optional[torch.Tensor] = None,
        global_cond: Optional[torch.Tensor] = None,
        generator: Optional[torch.Generator] = None,
    ) -> torch.Tensor:
        """Sample trajectory using Consistency Model.

        Single-step: Directly denoise from time_max to time_min.
        Multi-step: Use chaining for better quality.
        """
        scheduler = self.noise_scheduler

        # Sample initial position (reduced variance)
        trajectory = scheduler.sample_initial_position(condition_data, generator=generator)

        t = torch.tensor([scheduler.time_max], device=condition_data.device)
        s = torch.tensor([scheduler.time_min], device=condition_data.device)

        # Apply conditioning
        trajectory[condition_mask] = condition_data[condition_mask]

        # Single-step generation: directly to time 0
        out = self._forward(
            self.model,
            trajectory,
            t,
            s,
            local_cond=local_cond,
            global_cond=global_cond,
            clamp=True,
        )

        out[condition_mask] = condition_data[condition_mask]

        if not self.chain:
            return out

        # Multi-step chaining
        for chain_t in self.chaining_times[1:]:
            t = torch.tensor([float(chain_t)], device=condition_data.device)

            if self.chaining_times[0] == "C":
                # Convert from bin index to time
                t = scheduler.timesteps_to_times(t)

            s = torch.tensor([scheduler.time_min], device=condition_data.device)

            # Re-noise to intermediate time
            trajectory = scheduler.add_noise(out, t)

            # Denoise to time 0
            out = self._forward(
                self.model,
                trajectory,
                t,
                s,
                local_cond=local_cond,
                global_cond=global_cond,
                clamp=True,
            )

        return out

    def predict_action(
        self, obs_dict: Dict[str, torch.Tensor]
    ) -> Dict[str, torch.Tensor]:
        """Predict action sequence from observations."""
        assert "past_action" not in obs_dict

        # Normalize input
        nobs = self.normalizer.normalize(obs_dict)
        value = next(iter(nobs.values()))
        B, To = value.shape[:2]
        T = self.horizon
        Da = self.action_dim
        Do = self.obs_feature_dim

        device = self.device
        dtype = self.dtype

        # Build conditioning
        local_cond = None
        global_cond = None

        if self.obs_as_global_cond:
            this_nobs = dict_apply(
                nobs, lambda x: x[:, : self.n_obs_steps, ...].reshape(-1, *x.shape[2:])
            )
            nobs_features = self.obs_encoder(this_nobs)
            global_cond = nobs_features.reshape(B, -1)

            cond_data = torch.zeros(size=(B, T, Da), device=device, dtype=dtype)
            cond_mask = torch.zeros_like(cond_data, dtype=torch.bool)
        else:
            this_nobs = dict_apply(
                nobs, lambda x: x[:, : self.n_obs_steps, ...].reshape(-1, *x.shape[2:])
            )
            nobs_features = self.obs_encoder(this_nobs)
            nobs_features = nobs_features.reshape(B, self.n_obs_steps, -1)

            cond_data = torch.zeros(size=(B, T, Da + Do), device=device, dtype=dtype)
            cond_mask = torch.zeros_like(cond_data, dtype=torch.bool)
            cond_data[:, : self.n_obs_steps, Da:] = nobs_features
            cond_mask[:, : self.n_obs_steps, Da:] = True

        # Run sampling
        nsample = self.conditional_sample(
            cond_data, cond_mask, local_cond=local_cond, global_cond=global_cond
        )

        # Unnormalize prediction
        naction_pred = nsample[..., :Da]
        action_pred = self.normalizer["action"].unnormalize(naction_pred)

        # Extract action window
        start = self.n_obs_steps - 1
        end = start + self.n_action_steps
        action = action_pred[:, start:end]

        return {"action": action, "action_pred": action_pred}

    def forward(
        self, obs: Any, *, deterministic: bool = False, **kwargs: Any
    ) -> torch.Tensor:
        """Generate actions from observations."""
        input_obs = obs if isinstance(obs, dict) else {"obs": obs}

        for key in input_obs:
            if isinstance(input_obs[key], torch.Tensor):
                input_obs[key] = input_obs[key].to(self.device, self.dtype)
            else:
                input_obs[key] = torch.as_tensor(
                    input_obs[key], device=self.device, dtype=self.dtype
                )

        result = self.predict_action(input_obs)
        action = result["action"]

        if action.shape[0] == 1:
            action = action[0]
        return action

    def reset(self) -> None:
        """Reset policy state between episodes."""
        pass

    # ==================== TRAINING ====================

    def compute_loss(self, batch: Any) -> Dict[str, torch.Tensor]:
        """Compute CTM + DSM training loss.

        Returns a dictionary of losses for logging.
        """
        total_loss = {}

        assert "valid_mask" not in batch

        # Normalize
        nobs = self.normalizer.normalize(batch["obs"])
        nactions = self.normalizer["action"].normalize(batch["action"])
        batch_size = nactions.shape[0]
        horizon = nactions.shape[1]

        # Build conditioning
        local_cond = None
        global_cond = None
        trajectory = nactions
        cond_data = trajectory

        if self.obs_as_global_cond:
            this_nobs = dict_apply(
                nobs, lambda x: x[:, : self.n_obs_steps, ...].reshape(-1, *x.shape[2:])
            )
            nobs_features = self.obs_encoder(this_nobs)
            global_cond = nobs_features.reshape(batch_size, -1)
        else:
            this_nobs = dict_apply(nobs, lambda x: x.reshape(-1, *x.shape[2:]))
            nobs_features = self.obs_encoder(this_nobs)
            nobs_features = nobs_features.reshape(batch_size, horizon, -1)
            cond_data = torch.cat([nactions, nobs_features], dim=-1)
            trajectory = cond_data.detach()

        # Generate mask
        condition_mask = self.mask_generator(trajectory.shape)

        # ==================== CTM LOSS ====================
        if "ctm" in self.losses:
            # Sample t, s, u (as bin indices)
            t, s, u = self.noise_scheduler.sample_times_ctm(trajectory)

            times = self.noise_scheduler.timesteps_to_times(t)
            stops = self.noise_scheduler.timesteps_to_times(s)
            u_times = self.noise_scheduler.timesteps_to_times(u)

            # Add noise at time t
            noise_traj = self.noise_scheduler.add_noise(trajectory, times)

            # Use teacher to denoise from t to u
            denoise = lambda x, t: self.teacher(
                x, t, cond=global_cond, local_cond=local_cond
            )

            u_noise_traj = noise_traj
            distances = u - t
            max_d = torch.max(distances)

            for d in range(self.noise_scheduler.ode_steps_max):
                ct = torch.stack(
                    [
                        (t_i + d).clamp(int(t_i.item()), int(u_i.item()))
                        for t_i, u_i in zip(t, u)
                    ]
                )
                nt = torch.stack(
                    [
                        (t_i + d + 1).clamp(int(t_i.item()), int(u_i.item()))
                        for t_i, u_i in zip(t, u)
                    ]
                )

                current_times = self.noise_scheduler.timesteps_to_times(ct)
                next_times = self.noise_scheduler.timesteps_to_times(nt)

                u_noise_traj = self.noise_scheduler.step(
                    denoise, u_noise_traj, current_times, next_times, clamp=False
                )

            # Student: t -> s
            pred = self._forward(
                self.model,
                noise_traj,
                times,
                stops,
                local_cond=local_cond,
                global_cond=global_cond,
            )

            # EMA Student: u -> s
            target = self._forward(
                self.model_ema,
                u_noise_traj,
                u_times,
                stops,
                local_cond=local_cond,
                global_cond=global_cond,
            )

            # Both back to time 0
            start = torch.tensor(
                [self.noise_scheduler.time_min], device=trajectory.device
            ).expand(times.shape)

            pred = self._forward(
                self.model_ema,
                pred,
                stops,
                start,
                local_cond=local_cond,
                global_cond=global_cond,
            )

            target = self._forward(
                self.model_ema,
                target,
                stops,
                start,
                local_cond=local_cond,
                global_cond=global_cond,
            )

            loss = huber_loss(pred, target, delta=self.delta, weights=None)
            total_loss["ctm"] = loss * self.losses["ctm"]

        # ==================== DSM LOSS ====================
        if "dsm" in self.losses:
            # Sample times for DSM
            times, _ = self.noise_scheduler.sample_times(
                trajectory, time_sampler="ctm_dsm"
            )
            weights = self.noise_scheduler.get_weights(times, None, self.dsm_weights)

            # Add noise
            noisy_trajectory = self.noise_scheduler.add_noise(trajectory, times)

            # Predict clean trajectory (stop at time_min)
            stop = torch.tensor(
                [self.noise_scheduler.time_min], device=trajectory.device
            ).expand(times.shape)

            pred = self._forward(
                self.model,
                noisy_trajectory,
                times,
                stop,
                local_cond=local_cond,
                global_cond=global_cond,
                clamp=False,
            )

            target = trajectory

            loss = huber_loss(pred, target, delta=self.delta, weights=weights)
            total_loss["dsm"] = loss * self.losses["dsm"]

        return total_loss

    @torch.no_grad()
    def ema_update(self) -> None:
        """Update EMA model parameters."""
        param = [p.data for p in self.model.parameters()]
        param_ema = [p.data for p in self.model_ema.parameters()]

        torch._foreach_mul_(param_ema, self.ema_decay)
        torch._foreach_add_(param_ema, param, alpha=1 - self.ema_decay)