ModelTrain/module/policy.py

import torch.nn as nn
from torch.nn import functional as F
import torchvision.transforms as transforms
import torch, numpy as np
from ModelTrain.detr.main import build_ACT_model_and_optimizer, build_CNNMLP_model_and_optimizer
import IPython
e = IPython.embed
from robomimic.models.base_nets import ResNet18Conv, SpatialSoftmax
from robomimic.algo.diffusion_policy import replace_bn_with_gn, ConditionalUnet1D
from diffusers.schedulers.scheduling_ddim import DDIMScheduler
from diffusers.training_utils import EMAModel

class DiffusionPolicy(nn.Module):

    def __init__(self, args_override):
        super().__init__()
        self.camera_names = args_override["camera_names"]
        self.observation_horizon = args_override["observation_horizon"]
        self.action_horizon = args_override["action_horizon"]
        self.prediction_horizon = args_override["prediction_horizon"]
        self.num_inference_timesteps = args_override["num_inference_timesteps"]
        self.ema_power = args_override["ema_power"]
        self.lr = args_override["lr"]
        self.weight_decay = 0
        self.num_kp = 32
        self.feature_dimension = 64
        self.ac_dim = args_override["action_dim"]
        self.obs_dim = self.feature_dimension * len(self.camera_names) + 14
        backbones = []
        pools = []
        linears = []
        for _ in self.camera_names:
            backbones.append(ResNet18Conv(input_channel=3, pretrained=False, input_coord_conv=False))
            pools.append(SpatialSoftmax(input_shape=[512, 15, 20], num_kp=self.num_kp, temperature=1.0, learnable_temperature=False, noise_std=0.0))
            linears.append(torch.nn.Linear(int(np.prod([self.num_kp, 2])), self.feature_dimension))
        else:
            backbones = nn.ModuleList(backbones)
            pools = nn.ModuleList(pools)
            linears = nn.ModuleList(linears)
            backbones = replace_bn_with_gn(backbones)
            noise_pred_net = ConditionalUnet1D(input_dim=(self.ac_dim),
              global_cond_dim=(self.obs_dim * self.observation_horizon))
            nets = nn.ModuleDict({"policy": (nn.ModuleDict({
                        'backbones': backbones, 
                        'pools': pools, 
                        'linears': linears, 
                        'noise_pred_net': noise_pred_net}))})
            nets = nets.float().cuda()
            ENABLE_EMA = True
            if ENABLE_EMA:
                ema = EMAModel(model=nets, power=(self.ema_power))
            else:
                ema = None
            self.nets = nets
            self.ema = ema
            self.noise_scheduler = DDIMScheduler(num_train_timesteps=50,
              beta_schedule="squaredcos_cap_v2",
              clip_sample=True,
              set_alpha_to_one=True,
              steps_offset=0,
              prediction_type="epsilon")
            n_parameters = sum((p.numel() for p in self.nets.parameters()))
            print("number of parameters: %.2fM" % (n_parameters / 1000000.0,))

    def configure_optimizers(self):
        optimizer = torch.optim.AdamW((self.nets.parameters()), lr=(self.lr), weight_decay=(self.weight_decay))
        return optimizer

    def __call__(self, qpos, image, actions=None, is_pad=None):
        B = qpos.shape[0]
        if actions is not None:
            nets = self.nets
            all_features = []
            for cam_id in range(len(self.camera_names)):
                cam_image = image[:, cam_id]
                cam_features = nets["policy"]["backbones"][cam_id](cam_image)
                pool_features = nets["policy"]["pools"][cam_id](cam_features)
                pool_features = torch.flatten(pool_features, start_dim=1)
                out_features = nets["policy"]["linears"][cam_id](pool_features)
                all_features.append(out_features)
            else:
                obs_cond = torch.cat((all_features + [qpos]), dim=1)
                noise = torch.randn((actions.shape), device=(obs_cond.device))
                timesteps = torch.randint(0,
                  (self.noise_scheduler.config.num_train_timesteps), (
                 B,),
                  device=(obs_cond.device)).long()
                noisy_actions = self.noise_scheduler.add_noise(actions, noise, timesteps)
                noise_pred = nets["policy"]["noise_pred_net"](noisy_actions, timesteps, global_cond=obs_cond)
                all_l2 = F.mse_loss(noise_pred, noise, reduction="none")
                loss = (all_l2 * ~is_pad.unsqueeze(-1)).mean()
                loss_dict = {}
                loss_dict["l2_loss"] = loss
                loss_dict["loss"] = loss
                if self.training:
                    if self.ema is not None:
                        self.ema.step(nets)
                return loss_dict

        To = self.observation_horizon
        Ta = self.action_horizon
        Tp = self.prediction_horizon
        action_dim = self.ac_dim
        nets = self.nets
        if self.ema is not None:
            nets = self.ema.averaged_model
        all_features = []
        for cam_id in range(len(self.camera_names)):
            cam_image = image[:, cam_id]
            cam_features = nets["policy"]["backbones"][cam_id](cam_image)
            pool_features = nets["policy"]["pools"][cam_id](cam_features)
            pool_features = torch.flatten(pool_features, start_dim=1)
            out_features = nets["policy"]["linears"][cam_id](pool_features)
            all_features.append(out_features)
        else:
            obs_cond = torch.cat((all_features + [qpos]), dim=1)
            noisy_action = torch.randn((
             B, Tp, action_dim),
              device=(obs_cond.device))
            naction = noisy_action
            self.noise_scheduler.set_timesteps(self.num_inference_timesteps)
            for k in self.noise_scheduler.timesteps:
                noise_pred = nets["policy"]["noise_pred_net"](sample=naction,
                  timestep=k,
                  global_cond=obs_cond)
                naction = self.noise_scheduler.step(model_output=noise_pred,
                  timestep=k,
                  sample=naction).prev_sample
            else:
                return naction

    def serialize(self):
        return {'nets':(self.nets.state_dict)(),  'ema':self.ema.averaged_model.state_dict() if (self.ema is not None) else None}

    def deserialize(self, model_dict):
        status = self.nets.load_state_dict(model_dict["nets"])
        print("Loaded model")
        if model_dict.get("ema", None) is not None:
            print("Loaded EMA")
            status_ema = self.ema.averaged_model.load_state_dict(model_dict["ema"])
            status = [status, status_ema]
        return status


class ACTPolicy(nn.Module):

    def __init__(self, args_override):
        super().__init__()
        model, optimizer = build_ACT_model_and_optimizer(args_override)
        self.model = model
        self.optimizer = optimizer
        self.kl_weight = args_override["kl_weight"]
        self.vq = args_override["vq"]
        print(f"KL Weight {self.kl_weight}")

    def __call__(self, qpos, image, actions=None, is_pad=None, vq_sample=None):
        env_state = None
        normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        image = normalize(image)
        if actions is not None:
            actions = actions[:, :self.model.num_queries]
            is_pad = is_pad[:, :self.model.num_queries]
            loss_dict = dict()
            a_hat, is_pad_hat, (mu, logvar), probs, binaries = self.model(qpos, image, env_state, actions, is_pad, vq_sample)
            if self.vq or self.model.encoder is None:
                total_kld = [
                 torch.tensor(0.0)]
            else:
                total_kld, dim_wise_kld, mean_kld = kl_divergence(mu, logvar)
            if self.vq:
                loss_dict["vq_discrepancy"] = F.l1_loss(probs, binaries, reduction="mean")
            all_l1 = F.l1_loss(actions, a_hat, reduction="none")
            l1 = (all_l1 * ~is_pad.unsqueeze(-1)).mean()
            loss_dict["l1"] = l1
            loss_dict["kl"] = total_kld[0]
            loss_dict["loss"] = loss_dict["l1"] + loss_dict["kl"] * self.kl_weight
            return loss_dict
        a_hat, _, (_, _), _, _ = self.model(qpos, image, env_state, vq_sample=vq_sample)
        return a_hat

    def configure_optimizers(self):
        return self.optimizer

    @torch.no_grad()
    def vq_encode(self, qpos, actions, is_pad):
        actions = actions[:, :self.model.num_queries]
        is_pad = is_pad[:, :self.model.num_queries]
        _, _, binaries, _, _ = self.model.encode(qpos, actions, is_pad)
        return binaries

    def serialize(self):
        return self.state_dict()

    def deserialize(self, model_dict):
        return self.load_state_dict(model_dict)


class CNNMLPPolicy(nn.Module):

    def __init__(self, args_override):
        super().__init__()
        model, optimizer = build_CNNMLP_model_and_optimizer(args_override)
        self.model = model
        self.optimizer = optimizer

    def __call__(self, qpos, image, actions=None, is_pad=None):
        env_state = None
        normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[
         0.229, 0.224, 0.225])
        image = normalize(image)
        if actions is not None:
            actions = actions[:, 0]
            a_hat = self.model(qpos, image, env_state, actions)
            mse = F.mse_loss(actions, a_hat)
            loss_dict = dict()
            loss_dict["mse"] = mse
            loss_dict["loss"] = loss_dict["mse"]
            return loss_dict
        a_hat = self.model(qpos, image, env_state)
        return a_hat

    def configure_optimizers(self):
        return self.optimizer


def kl_divergence(mu, logvar):
    batch_size = mu.size(0)
    assert batch_size != 0
    if mu.data.ndimension() == 4:
        mu = mu.view(mu.size(0), mu.size(1))
    if logvar.data.ndimension() == 4:
        logvar = logvar.view(logvar.size(0), logvar.size(1))
    klds = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp())
    total_kld = klds.sum(1).mean(0, True)
    dimension_wise_kld = klds.mean(0)
    mean_kld = klds.mean(1).mean(0, True)
    return (
     total_kld, dimension_wise_kld, mean_kld)

# okay decompiling policy.pyc